Fix logging feature consuming response body and preventing streaming

The logging feature's wrap_response eagerly called response.body.to_s to
log the body, which consumed the stream and made subsequent
response.body.each calls raise StateError.

Replace eager consumption with a tee-style BodyLogger stream wrapper
that logs each chunk as it flows through readpartial. Also generalize
Body's stream detection from is_a?(Inflater) to respond_to?(:connection)
so any stream wrapper is supported.

Closes #785.

Author: sferik

CHANGELOG.md

@@ -23,6 +23,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `LocalJumpError` when using instrumentation with instrumenters that
   unconditionally yield in `#instrument` (e.g., `ActiveSupport::Notifications`)
   (#673)
+- Logging feature no longer eagerly consumes the response body at debug level;
+  body chunks are now logged as they are streamed, preserving
+  `response.body.each` (#785)
 
 ### Changed
 

lib/http/features/logging.rb

@@ -72,19 +72,95 @@ def wrap_request(request)
       # @api public
       def wrap_response(response)
         logger.info { "< #{response.status}" }
-        logger.debug { "#{stringify_headers(response.headers)}\n\n#{response.body}" }
 
-        response
+        return log_response_body_inline(response) unless response.body.is_a?(Response::Body)
+
+        logger.debug { stringify_headers(response.headers) }
+        return response unless logger.debug?
+
+        Response.new(logged_response_options(response))
       end
 
       private
 
+      # Log response with body inline (for non-streaming string bodies)
+      # @return [HTTP::Response]
+      # @api private
+      def log_response_body_inline(response)
+        logger.debug { "#{stringify_headers(response.headers)}\n\n#{response.body}" }
+        response
+      end
+
+      # Build options hash for a response with body logging
+      # @return [Hash]
+      # @api private
+      def logged_response_options(response)
+        {
+          status:        response.status,
+          version:       response.version,
+          headers:       response.headers,
+          proxy_headers: response.proxy_headers,
+          connection:    response.connection,
+          body:          logged_body(response.body),
+          request:       response.request
+        }
+      end
+
+      # Wrap a response body with a logging stream
+      # @return [HTTP::Response::Body]
+      # @api private
+      def logged_body(body)
+        stream = BodyLogger.new(body.instance_variable_get(:@stream), logger)
+        Response::Body.new(stream, encoding: body.encoding)
+      end
+
       # Convert headers to a string representation
       # @return [String]
       # @api private
       def stringify_headers(headers)
         headers.map { |name, value| "#{name}: #{value}" }.join("\n")
       end
+
+      # Stream wrapper that logs each chunk as it flows through readpartial
+      class BodyLogger
+        # The underlying connection
+        #
+        # @example
+        #   body_logger.connection
+        #
+        # @return [HTTP::Connection] the underlying connection
+        # @api public
+        attr_reader :connection
+
+        # Create a new BodyLogger wrapping a stream
+        #
+        # @example
+        #   BodyLogger.new(stream, logger)
+        #
+        # @param stream [#readpartial] the stream to wrap
+        # @param logger [#debug] the logger instance
+        # @return [BodyLogger]
+        # @api public
+        def initialize(stream, logger)
+          @stream = stream
+          @connection = stream.respond_to?(:connection) ? stream.connection : stream
+          @logger = logger
+        end
+
+        # Read a chunk from the underlying stream and log it
+        #
+        # @example
+        #   body_logger.readpartial # => "chunk"
+        #
+        # @return [String] the chunk read from the stream
+        # @raise [EOFError] when no more data left
+        # @api public
+        def readpartial(*)
+          chunk = @stream.readpartial(*)
+          @logger.debug { chunk }
+          chunk
+        end
+      end
     end
   end
 end

lib/http/response/body.rb

@@ -41,7 +41,7 @@ class Body
       # @api public
       def initialize(stream, encoding: Encoding::BINARY)
         @stream     = stream
-        @connection = stream.is_a?(Inflater) ? stream.connection : stream
+        @connection = stream.respond_to?(:connection) ? stream.connection : stream
         @streaming  = nil
         @contents   = nil
         @encoding   = find_encoding(encoding)

sig/http.rbs

@@ -322,7 +322,17 @@ module HTTP
 
       private
 
+      def log_response_body_inline: (untyped response) -> untyped
+      def logged_response_options: (untyped response) -> Hash[Symbol, untyped]
+      def logged_body: (Response::Body body) -> Response::Body
       def stringify_headers: (Headers headers) -> String
+
+      class BodyLogger
+        attr_reader connection: untyped
+
+        def initialize: (untyped stream, untyped logger) -> void
+        def readpartial: (*untyped) -> String
+      end
     end
 
     class NormalizeUri < Feature

test/http/features/logging_test.rb

@@ -58,28 +58,147 @@
   end
 
   describe "logging the response" do
-    let(:response) do
-      HTTP::Response.new(
-        version: "1.1",
-        status:  200,
-        headers: { content_type: "application/json" },
-        body:    '{"success": true}',
-        request: HTTP::Request.new(verb: :get, uri: "https://example.com")
+    context "with a string body" do
+      let(:response) do
+        HTTP::Response.new(
+          version: "1.1",
+          status:  200,
+          headers: { content_type: "application/json" },
+          body:    '{"success": true}',
+          request: HTTP::Request.new(verb: :get, uri: "https://example.com")
+        )
+      end
+
+      it "logs the response with body" do
+        feature.wrap_response(response)
+
+        expected = <<~OUTPUT
+          ** INFO **
+          < 200 OK
+          ** DEBUG **
+          Content-Type: application/json
+
+          {"success": true}
+        OUTPUT
+        assert_equal expected, logdev.string
+      end
+    end
+
+    context "with a streaming body" do
+      let(:chunks) { %w[{"suc cess" :true}] }
+      let(:stream) do
+        fake(
+          readpartial: proc { chunks.shift or raise EOFError },
+          close:       nil,
+          closed?:     true
+        )
+      end
+      let(:body) { HTTP::Response::Body.new(stream) }
+      let(:response) do
+        HTTP::Response.new(
+          version: "1.1",
+          status:  200,
+          headers: { content_type: "application/json" },
+          body:    body,
+          request: HTTP::Request.new(verb: :get, uri: "https://example.com")
+        )
+      end
+
+      it "does not consume the body" do
+        wrapped = feature.wrap_response(response)
+
+        assert_nil wrapped.body.instance_variable_get(:@streaming)
+      end
+
+      it "logs body chunks as they are streamed" do
+        wrapped = feature.wrap_response(response)
+        wrapped.body.to_s
+
+        assert_includes logdev.string, '{"suc'
+        assert_includes logdev.string, 'cess"'
+        assert_includes logdev.string, ":true}"
+      end
+
+      it "preserves the full body content" do
+        wrapped = feature.wrap_response(response)
+
+        assert_equal '{"success":true}', wrapped.body.to_s
+      end
+    end
+
+    context "when logger level is above debug" do
+      let(:feature) do
+        logger           = Logger.new(logdev)
+        logger.formatter = ->(severity, _, _, message) { format("** %s **\n%s\n", severity, message) }
+        logger.level     = Logger::INFO
+
+        HTTP::Features::Logging.new(logger: logger)
+      end
+
+      let(:stream) do
+        fake(
+          readpartial: proc { raise EOFError },
+          close:       nil,
+          closed?:     true
+        )
+      end
+      let(:body) { HTTP::Response::Body.new(stream) }
+      let(:response) do
+        HTTP::Response.new(
+          version: "1.1",
+          status:  200,
+          headers: { content_type: "text/plain" },
+          body:    body,
+          request: HTTP::Request.new(verb: :get, uri: "https://example.com")
+        )
+      end
+
+      it "does not wrap the body" do
+        wrapped = feature.wrap_response(response)
+
+        assert_same response, wrapped
+      end
+    end
+  end
+
+  describe "BodyLogger" do
+    let(:logger) do
+      logger           = Logger.new(logdev)
+      logger.formatter = ->(severity, _, _, message) { format("** %s **\n%s\n", severity, message) }
+      logger
+    end
+
+    it "passes through chunks and logs them" do
+      chunks = %w[hello world]
+      stream = fake(readpartial: proc { chunks.shift or raise EOFError })
+
+      body_logger = HTTP::Features::Logging::BodyLogger.new(stream, logger)
+
+      assert_equal "hello", body_logger.readpartial
+      assert_equal "world", body_logger.readpartial
+      assert_raises(EOFError) { body_logger.readpartial }
+      assert_includes logdev.string, "hello"
+      assert_includes logdev.string, "world"
+    end
+
+    it "exposes the underlying connection" do
+      connection = Object.new
+      stream = fake(
+        readpartial: proc { raise EOFError },
+        connection:  connection
       )
+
+      body_logger = HTTP::Features::Logging::BodyLogger.new(stream, logger)
+
+      assert_same connection, body_logger.connection
     end
 
-    it "logs the response" do
-      feature.wrap_response(response)
+    it "uses the stream as connection when stream has no connection method" do
+      stream = fake(readpartial: proc { raise EOFError })
 
-      expected = <<~OUTPUT
-        ** INFO **
-        < 200 OK
-        ** DEBUG **
-        Content-Type: application/json
+      body_logger = HTTP::Features::Logging::BodyLogger.new(stream, logger)
 
-        {"success": true}
-      OUTPUT
-      assert_equal expected, logdev.string
+      assert_same stream, body_logger.connection
     end
   end
 end