diff --git a/enterprise/lib/captain/tools/http_tool.rb b/enterprise/lib/captain/tools/http_tool.rb index b37d3b815..84725033a 100644 --- a/enterprise/lib/captain/tools/http_tool.rb +++ b/enterprise/lib/captain/tools/http_tool.rb @@ -2,6 +2,8 @@ require 'agents' class Captain::Tools::HttpTool < Agents::Tool MAX_RESPONSE_SIZE = 1.megabyte + OPEN_TIMEOUT = 10 + READ_TIMEOUT = 30 def initialize(assistant, custom_tool) @assistant = assistant @@ -34,6 +36,8 @@ class Captain::Tools::HttpTool < Agents::Tool method: @custom_tool.http_method, body: body, max_bytes: MAX_RESPONSE_SIZE, + open_timeout: OPEN_TIMEOUT, + read_timeout: READ_TIMEOUT, headers: request_headers(tool_context, body), http_basic_authentication: @custom_tool.build_basic_auth_credentials, validate_content_type: false diff --git a/lib/safe_fetch.rb b/lib/safe_fetch.rb index 6a0c26542..a7d96ac14 100644 --- a/lib/safe_fetch.rb +++ b/lib/safe_fetch.rb @@ -51,6 +51,8 @@ module SafeFetch method: :get, body: nil, max_bytes: nil, + open_timeout: DEFAULT_OPEN_TIMEOUT, + read_timeout: DEFAULT_READ_TIMEOUT, headers: nil, http_basic_authentication: nil, allowed_content_type_prefixes: DEFAULT_ALLOWED_CONTENT_TYPE_PREFIXES, @@ -69,6 +71,8 @@ module SafeFetch body, tempfile, effective_max_bytes, + open_timeout, + read_timeout, headers, http_basic_authentication, allowed_content_type_prefixes, @@ -98,7 +102,7 @@ module SafeFetch private # rubocop:disable Metrics/MethodLength, Metrics/ParameterLists - def stream_to_tempfile(url, method, body, tempfile, max_bytes, headers, http_basic_authentication, + def stream_to_tempfile(url, method, body, tempfile, max_bytes, open_timeout, read_timeout, headers, http_basic_authentication, allowed_content_type_prefixes, allowed_content_types, validate_content_type) response = nil bytes_written = 0 @@ -112,7 +116,7 @@ module SafeFetch body: body, request_proc: request_proc(http_basic_authentication), sensitive_headers: sensitive_headers(headers), - http_options: { open_timeout: DEFAULT_OPEN_TIMEOUT, read_timeout: DEFAULT_READ_TIMEOUT } + http_options: { open_timeout: open_timeout, read_timeout: read_timeout } ) do |res| response = res next unless res.is_a?(Net::HTTPSuccess) diff --git a/spec/enterprise/lib/captain/tools/http_tool_spec.rb b/spec/enterprise/lib/captain/tools/http_tool_spec.rb index 6a2b1aa5e..fefaf3d4b 100644 --- a/spec/enterprise/lib/captain/tools/http_tool_spec.rb +++ b/spec/enterprise/lib/captain/tools/http_tool_spec.rb @@ -1,4 +1,5 @@ require 'rails_helper' +require 'stringio' RSpec.describe Captain::Tools::HttpTool, type: :model do let(:account) { create(:account) } @@ -45,6 +46,19 @@ RSpec.describe Captain::Tools::HttpTool, type: :model do expect(result).to eq('{"status": "success"}') expect(WebMock).to have_requested(:get, 'https://example.com/orders/123') end + + it 'preserves the previous timeout budget when fetching' do + allow(SafeFetch).to receive(:fetch) do |fetched_url, **options, &block| + expect(fetched_url).to eq('https://example.com/orders/123') + expect(options[:open_timeout]).to eq(10) + expect(options[:read_timeout]).to eq(30) + block.call(OpenStruct.new(tempfile: StringIO.new('{"status": "success"}'))) + end + + result = tool.perform(tool_context) + + expect(result).to eq('{"status": "success"}') + end end context 'with POST request' do