diff --git a/lib/braintrust/api/datasets.rb b/lib/braintrust/api/datasets.rb index 120d0f7..a04a949 100644 --- a/lib/braintrust/api/datasets.rb +++ b/lib/braintrust/api/datasets.rb @@ -4,6 +4,7 @@ require "json" require "uri" require_relative "../logger" +require_relative "../internal/http" module Braintrust class API @@ -111,6 +112,7 @@ def fetch(id:, limit: 1000, cursor: nil, version: nil) payload[:version] = version if version response = http_post_json_raw("/btql", payload) + Braintrust::Internal::Http.decompress_response!(response) # Parse JSONL response records = response.body.lines @@ -158,9 +160,7 @@ def http_request(method, path, params: {}, payload: nil, base_url: nil, parse_js start_time = Time.now Log.debug("[API] #{method.upcase} #{uri}") - http = Net::HTTP.new(uri.host, uri.port) - http.use_ssl = (uri.scheme == "https") - response = http.request(request) + response = Braintrust::Internal::Http.with_redirects(uri, request) duration_ms = ((Time.now - start_time) * 1000).round(2) Log.debug("[API] #{method.upcase} #{uri} -> #{response.code} (#{duration_ms}ms, #{response.body.bytesize} bytes)") diff --git a/lib/braintrust/api/functions.rb b/lib/braintrust/api/functions.rb index 1d569d9..2173bd6 100644 --- a/lib/braintrust/api/functions.rb +++ b/lib/braintrust/api/functions.rb @@ -4,6 +4,7 @@ require "json" require "uri" require_relative "../logger" +require_relative "../internal/http" module Braintrust class API @@ -242,9 +243,7 @@ def http_request(method, path, params: {}, payload: nil, parse_json: true) start_time = Time.now Log.debug("[API] #{method.upcase} #{uri}") - http = Net::HTTP.new(uri.host, uri.port) - http.use_ssl = (uri.scheme == "https") - response = http.request(request) + response = Braintrust::Internal::Http.with_redirects(uri, request) duration_ms = ((Time.now - start_time) * 1000).round(2) Log.debug("[API] #{method.upcase} #{uri} -> #{response.code} (#{duration_ms}ms, #{response.body.bytesize} bytes)") diff --git a/lib/braintrust/api/internal/auth.rb b/lib/braintrust/api/internal/auth.rb index 457c4d3..c642d8a 100644 --- a/lib/braintrust/api/internal/auth.rb +++ b/lib/braintrust/api/internal/auth.rb @@ -4,6 +4,7 @@ require "json" require "uri" require_relative "../../logger" +require_relative "../../internal/http" module Braintrust class API @@ -44,12 +45,7 @@ def self.login(api_key:, app_url:, org_name: nil) request = Net::HTTP::Post.new(uri) request["Authorization"] = "Bearer #{api_key}" - http = Net::HTTP.new(uri.hostname, uri.port) - http.use_ssl = true if uri.scheme == "https" - - response = http.start do |http_session| - http_session.request(request) - end + response = Braintrust::Internal::Http.with_redirects(uri, request) Log.debug("Login: received response [#{response.code}]") diff --git a/lib/braintrust/api/internal/experiments.rb b/lib/braintrust/api/internal/experiments.rb index b2e2ea5..55830b4 100644 --- a/lib/braintrust/api/internal/experiments.rb +++ b/lib/braintrust/api/internal/experiments.rb @@ -3,6 +3,7 @@ require "net/http" require "json" require "uri" +require_relative "../../internal/http" module Braintrust class API @@ -41,9 +42,7 @@ def create(name:, project_id:, ensure_new: true, tags: nil, metadata: nil, request["Authorization"] = "Bearer #{@state.api_key}" request.body = JSON.dump(payload) - http = Net::HTTP.new(uri.host, uri.port) - http.use_ssl = (uri.scheme == "https") - response = http.request(request) + response = Braintrust::Internal::Http.with_redirects(uri, request) unless response.is_a?(Net::HTTPSuccess) raise Error, "HTTP #{response.code} for POST #{uri}: #{response.body}" diff --git a/lib/braintrust/api/internal/projects.rb b/lib/braintrust/api/internal/projects.rb index 4995a68..317aa32 100644 --- a/lib/braintrust/api/internal/projects.rb +++ b/lib/braintrust/api/internal/projects.rb @@ -3,6 +3,7 @@ require "net/http" require "json" require "uri" +require_relative "../../internal/http" module Braintrust class API @@ -26,9 +27,7 @@ def create(name:) request["Authorization"] = "Bearer #{@state.api_key}" request.body = JSON.dump({name: name}) - http = Net::HTTP.new(uri.host, uri.port) - http.use_ssl = (uri.scheme == "https") - response = http.request(request) + response = Braintrust::Internal::Http.with_redirects(uri, request) unless response.is_a?(Net::HTTPSuccess) raise Error, "HTTP #{response.code} for POST #{uri}: #{response.body}" diff --git a/lib/braintrust/internal/http.rb b/lib/braintrust/internal/http.rb new file mode 100644 index 0000000..b569ff8 --- /dev/null +++ b/lib/braintrust/internal/http.rb @@ -0,0 +1,97 @@ +# frozen_string_literal: true + +require "net/http" +require "uri" +require "zlib" +require "stringio" +require_relative "../logger" + +module Braintrust + module Internal + # HTTP utilities for redirect following and response decompression. + # Drop-in enhancement for raw Net::HTTP request calls throughout the SDK. + module Http + DEFAULT_MAX_REDIRECTS = 5 + + # Execute an HTTP request, following redirects as needed. + # + # @param uri [URI] The request URI + # @param request [Net::HTTPRequest] The prepared request object + # @param max_redirects [Integer] Maximum number of redirects to follow + # @return [Net::HTTPResponse] The final response + # @raise [Braintrust::Error] On too many redirects or missing Location header + def self.with_redirects(uri, request, max_redirects: DEFAULT_MAX_REDIRECTS) + response = perform_request(uri, request) + + redirects = 0 + original_request = request + + while response.is_a?(Net::HTTPRedirection) + redirects += 1 + if redirects > max_redirects + raise Error, "Too many redirects (max #{max_redirects})" + end + + location = response["location"] + unless location + raise Error, "Redirect response #{response.code} without Location header" + end + + redirect_uri = URI(location) + redirect_uri = uri + redirect_uri unless redirect_uri.host + + Log.debug("[HTTP] Following #{response.code} redirect to #{redirect_uri}") + + request = build_redirect_request(response, redirect_uri, original_request, uri) + uri = redirect_uri + response = perform_request(uri, request) + end + + response + end + + # Decompress an HTTP response body in place based on Content-Encoding. + # No-op if the response has no recognized encoding. + # + # @param response [Net::HTTPResponse] The response to decompress + # @return [void] + def self.decompress_response!(response) + encoding = response["content-encoding"]&.downcase + case encoding + when "gzip", "x-gzip" + gz = Zlib::GzipReader.new(StringIO.new(response.body)) + response.body.replace(gz.read) + gz.close + response.delete("content-encoding") + end + end + + def self.perform_request(uri, request) + http = Net::HTTP.new(uri.host, uri.port) + http.use_ssl = (uri.scheme == "https") + http.request(request) + end + private_class_method :perform_request + + def self.build_redirect_request(response, redirect_uri, original_request, original_uri) + if response.code == "307" || response.code == "308" + request = original_request.class.new(redirect_uri) + request.body = original_request.body + request["Content-Type"] = original_request["Content-Type"] if original_request["Content-Type"] + else + # 301, 302, 303: follow with GET, no body + request = Net::HTTP::Get.new(redirect_uri) + end + + # Strip Authorization when redirecting to a different host (e.g. S3) + if original_uri.host == redirect_uri.host + auth = original_request["Authorization"] + request["Authorization"] = auth if auth + end + + request + end + private_class_method :build_redirect_request + end + end +end diff --git a/lib/braintrust/trace/attachment.rb b/lib/braintrust/trace/attachment.rb index 51a98d4..53f9ceb 100644 --- a/lib/braintrust/trace/attachment.rb +++ b/lib/braintrust/trace/attachment.rb @@ -2,6 +2,7 @@ require "net/http" require_relative "../internal/encoding" +require_relative "../internal/http" require "uri" module Braintrust @@ -91,7 +92,8 @@ def self.from_file(content_type, path) # att = Braintrust::Trace::Attachment.from_url("https://example.com/image.png") def self.from_url(url) uri = URI.parse(url) - response = Net::HTTP.get_response(uri) + request = Net::HTTP::Get.new(uri) + response = Braintrust::Internal::Http.with_redirects(uri, request) unless response.is_a?(Net::HTTPSuccess) raise StandardError, "Failed to fetch URL: #{response.code} #{response.message}" diff --git a/test/braintrust/api/internal/auth_test.rb b/test/braintrust/api/internal/auth_test.rb index 70b1bb1..22e45ca 100644 --- a/test/braintrust/api/internal/auth_test.rb +++ b/test/braintrust/api/internal/auth_test.rb @@ -91,7 +91,7 @@ def test_login_server_error def test_login_unexpected_response stub = stub_request(:post, "https://www.braintrust.dev/api/apikey/login") - .to_return(status: 301, body: "", headers: {}) + .to_return(status: 100, body: "", headers: {}) error = assert_raises(Braintrust::Error) do Braintrust::API::Internal::Auth.login( @@ -101,7 +101,7 @@ def test_login_unexpected_response end assert_match(/unexpected response/i, error.message) - assert_match(/301/, error.message) + assert_match(/100/, error.message) ensure remove_request_stub(stub) end diff --git a/test/braintrust/internal/http_test.rb b/test/braintrust/internal/http_test.rb new file mode 100644 index 0000000..641b718 --- /dev/null +++ b/test/braintrust/internal/http_test.rb @@ -0,0 +1,199 @@ +# frozen_string_literal: true + +require "test_helper" +require "braintrust/internal/http" + +class Braintrust::Internal::HttpTest < Minitest::Test + # -- with_redirects -- + + def test_with_redirects_returns_response_on_success + stub = stub_request(:get, "https://api.example.com/v1/dataset") + .to_return(status: 200, body: '{"ok":true}') + + uri = URI("https://api.example.com/v1/dataset") + request = Net::HTTP::Get.new(uri) + + response = Braintrust::Internal::Http.with_redirects(uri, request) + + assert_equal "200", response.code + assert_equal '{"ok":true}', response.body + ensure + remove_request_stub(stub) + end + + def test_with_redirects_follows_303_to_different_host_with_get + api_stub = stub_request(:post, "https://api.example.com/btql") + .to_return( + status: 303, + headers: {"Location" => "https://s3.amazonaws.com/bucket/response.jsonl"} + ) + + # Stub only accepts GET with no Authorization - verifies method conversion and auth stripping + s3_stub = stub_request(:get, "https://s3.amazonaws.com/bucket/response.jsonl") + .with { |req| req.headers["Authorization"].nil? } + .to_return(status: 200, body: '{"id":"1"}') + + uri = URI("https://api.example.com/btql") + request = Net::HTTP::Post.new(uri) + request["Authorization"] = "Bearer secret-token" + request["Content-Type"] = "application/json" + request.body = '{"query":"test"}' + + response = Braintrust::Internal::Http.with_redirects(uri, request) + + assert_equal "200", response.code + assert_equal '{"id":"1"}', response.body + assert_requested(s3_stub, times: 1) + ensure + remove_request_stub(api_stub) + remove_request_stub(s3_stub) + end + + def test_with_redirects_preserves_auth_for_same_host + redirect_stub = stub_request(:get, "https://api.example.com/old") + .to_return( + status: 303, + headers: {"Location" => "https://api.example.com/new"} + ) + + # Stub requires Authorization header - verifies auth is preserved for same host + target_stub = stub_request(:get, "https://api.example.com/new") + .with(headers: {"Authorization" => "Bearer secret-token"}) + .to_return(status: 200, body: "ok") + + uri = URI("https://api.example.com/old") + request = Net::HTTP::Get.new(uri) + request["Authorization"] = "Bearer secret-token" + + Braintrust::Internal::Http.with_redirects(uri, request) + + assert_requested(target_stub, times: 1) + ensure + remove_request_stub(redirect_stub) + remove_request_stub(target_stub) + end + + def test_with_redirects_307_preserves_method_and_body + redirect_stub = stub_request(:post, "https://api.example.com/old") + .to_return( + status: 307, + headers: {"Location" => "https://api.example.com/new"} + ) + + # Stub requires POST with matching body and content-type - verifies method/body preservation + target_stub = stub_request(:post, "https://api.example.com/new") + .with(body: '{"data":1}', headers: {"Content-Type" => "application/json"}) + .to_return(status: 200, body: "ok") + + uri = URI("https://api.example.com/old") + request = Net::HTTP::Post.new(uri) + request["Content-Type"] = "application/json" + request["Authorization"] = "Bearer secret-token" + request.body = '{"data":1}' + + Braintrust::Internal::Http.with_redirects(uri, request) + + assert_requested(target_stub, times: 1) + ensure + remove_request_stub(redirect_stub) + remove_request_stub(target_stub) + end + + def test_with_redirects_raises_on_too_many_redirects + stubs = (0..5).map { |i| + stub_request(:get, "https://api.example.com/r#{i}") + .to_return( + status: 302, + headers: {"Location" => "https://api.example.com/r#{i + 1}"} + ) + } + + uri = URI("https://api.example.com/r0") + request = Net::HTTP::Get.new(uri) + + error = assert_raises(Braintrust::Error) do + Braintrust::Internal::Http.with_redirects(uri, request) + end + + assert_match(/too many redirects/i, error.message) + ensure + stubs.each { |s| remove_request_stub(s) } + end + + def test_with_redirects_respects_custom_max_redirects + stubs = (0..2).map { |i| + stub_request(:get, "https://api.example.com/r#{i}") + .to_return( + status: 302, + headers: {"Location" => "https://api.example.com/r#{i + 1}"} + ) + } + + uri = URI("https://api.example.com/r0") + request = Net::HTTP::Get.new(uri) + + error = assert_raises(Braintrust::Error) do + Braintrust::Internal::Http.with_redirects(uri, request, max_redirects: 2) + end + + assert_match(/too many redirects.*max 2/i, error.message) + ensure + stubs.each { |s| remove_request_stub(s) } + end + + def test_with_redirects_raises_on_missing_location_header + stub = stub_request(:get, "https://api.example.com/missing") + .to_return(status: 303) + + uri = URI("https://api.example.com/missing") + request = Net::HTTP::Get.new(uri) + + error = assert_raises(Braintrust::Error) do + Braintrust::Internal::Http.with_redirects(uri, request) + end + + assert_match(/without location header/i, error.message) + ensure + remove_request_stub(stub) + end + + # -- decompress_response! -- + + def test_decompress_response_handles_gzip + original = "hello world" + gzipped = gzip_string(original) + + stub = stub_request(:get, "https://s3.example.com/data.jsonl.gz") + .to_return( + status: 200, + body: gzipped, + headers: {"Content-Encoding" => "gzip"} + ) + + uri = URI("https://s3.example.com/data.jsonl.gz") + request = Net::HTTP::Get.new(uri) + response = Braintrust::Internal::Http.with_redirects(uri, request) + + Braintrust::Internal::Http.decompress_response!(response) + + assert_equal original, response.body + assert_nil response["content-encoding"] + ensure + remove_request_stub(stub) + end + + def test_decompress_response_noop_without_encoding + stub = stub_request(:get, "https://api.example.com/data") + .to_return(status: 200, body: '{"ok":true}') + + uri = URI("https://api.example.com/data") + request = Net::HTTP::Get.new(uri) + response = Braintrust::Internal::Http.with_redirects(uri, request) + + Braintrust::Internal::Http.decompress_response!(response) + + assert_equal '{"ok":true}', response.body + ensure + remove_request_stub(stub) + end +end diff --git a/test/braintrust/trace/attachment_test.rb b/test/braintrust/trace/attachment_test.rb index 8674c37..b220d88 100644 --- a/test/braintrust/trace/attachment_test.rb +++ b/test/braintrust/trace/attachment_test.rb @@ -96,62 +96,47 @@ def test_attachment_is_reusable end def test_from_url_fetches_remote_image - # Mock HTTP response for testing - mock_response = Minitest::Mock.new - mock_response.expect(:is_a?, true, [Net::HTTPSuccess]) - mock_response.expect(:content_type, "image/png") - mock_response.expect(:body, @test_png_data) + stub = stub_request(:get, "https://example.com/image.png") + .to_return(status: 200, body: @test_png_data, headers: {"Content-Type" => "image/png"}) - Net::HTTP.stub(:get_response, mock_response) do - att = Braintrust::Trace::Attachment.from_url("https://example.com/image.png") + att = Braintrust::Trace::Attachment.from_url("https://example.com/image.png") - refute_nil att - assert_instance_of Braintrust::Trace::Attachment, att - - # Should be able to convert to data URL - data_url = att.to_data_url - assert_match(/^data:image\/png;base64,/, data_url) - end + refute_nil att + assert_instance_of Braintrust::Trace::Attachment, att - mock_response.verify + # Should be able to convert to data URL + data_url = att.to_data_url + assert_match(/^data:image\/png;base64,/, data_url) + ensure + remove_request_stub(stub) end def test_from_url_handles_content_type_from_response - # Mock HTTP response with JPEG content type - mock_response = Minitest::Mock.new - mock_response.expect(:is_a?, true, [Net::HTTPSuccess]) - mock_response.expect(:content_type, "image/jpeg") - mock_response.expect(:body, "fake jpeg data") + stub = stub_request(:get, "https://example.com/photo.jpg") + .to_return(status: 200, body: "fake jpeg data", headers: {"Content-Type" => "image/jpeg"}) - Net::HTTP.stub(:get_response, mock_response) do - att = Braintrust::Trace::Attachment.from_url("https://example.com/photo.jpg") + att = Braintrust::Trace::Attachment.from_url("https://example.com/photo.jpg") - refute_nil att - data_url = att.to_data_url - - # Should detect JPEG content type - assert_match(/^data:image\/jpeg;base64,/, data_url) - end + refute_nil att + data_url = att.to_data_url - mock_response.verify + # Should detect JPEG content type + assert_match(/^data:image\/jpeg;base64,/, data_url) + ensure + remove_request_stub(stub) end def test_from_url_raises_on_network_error - # Mock HTTP error response - mock_response = Minitest::Mock.new - mock_response.expect(:is_a?, false, [Net::HTTPSuccess]) - mock_response.expect(:code, "404") - mock_response.expect(:message, "Not Found") - - Net::HTTP.stub(:get_response, mock_response) do - error = assert_raises(StandardError) do - Braintrust::Trace::Attachment.from_url("https://example.com/nonexistent.png") - end - - assert_match(/Failed to fetch URL: 404 Not Found/, error.message) + stub = stub_request(:get, "https://example.com/nonexistent.png") + .to_return(status: 404) + + error = assert_raises(StandardError) do + Braintrust::Trace::Attachment.from_url("https://example.com/nonexistent.png") end - mock_response.verify + assert_match(/Failed to fetch URL: 404/, error.message) + ensure + remove_request_stub(stub) end def test_different_content_types diff --git a/test/support/compression_helper.rb b/test/support/compression_helper.rb new file mode 100644 index 0000000..eadb22b --- /dev/null +++ b/test/support/compression_helper.rb @@ -0,0 +1,19 @@ +# frozen_string_literal: true + +require "zlib" +require "stringio" + +module Test + module Support + module CompressionHelper + def gzip_string(str) + io = StringIO.new + io.set_encoding("ASCII-8BIT") + gz = Zlib::GzipWriter.new(io) + gz.write(str) + gz.close + io.string + end + end + end +end diff --git a/test/test_helper.rb b/test/test_helper.rb index 2f84d74..dad4740 100644 --- a/test/test_helper.rb +++ b/test/test_helper.rb @@ -70,6 +70,7 @@ require_relative "support/assert_helper" require_relative "support/braintrust_helper" +require_relative "support/compression_helper" require_relative "support/fixture_helper" require_relative "support/log_helper" require_relative "support/mock_helper" @@ -80,6 +81,7 @@ class Minitest::Test include ::Test::Support::AssertHelper include ::Test::Support::BraintrustHelper + include ::Test::Support::CompressionHelper include ::Test::Support::FixtureHelper include ::Test::Support::LogHelper include ::Test::Support::MockHelper