fix: coroutine leaks from connection pool

Author: stainless-app[bot]

lib/openai/internal/transport/pooled_net_requester.rb

@@ -134,9 +134,9 @@ def execute(request)
 
           # rubocop:disable Metrics/BlockLength
           enum = Enumerator.new do |y|
-            with_pool(url, deadline: deadline) do |conn|
-              next if finished
+            next if finished
 
+            with_pool(url, deadline: deadline) do |conn|
               req, closing = self.class.build_request(request) do
                 self.class.calibrate_socket_timeout(conn, deadline)
               end
@@ -149,7 +149,7 @@ def execute(request)
 
               self.class.calibrate_socket_timeout(conn, deadline)
               conn.request(req) do |rsp|
-                y << [conn, req, rsp]
+                y << [req, rsp]
                 break if finished
 
                 rsp.read_body do |bytes|
@@ -160,6 +160,8 @@ def execute(request)
                 end
                 eof = true
               end
+            ensure
+              conn.finish if !eof && conn&.started?
             end
           rescue Timeout::Error
             raise OpenAI::Errors::APITimeoutError.new(url: url, request: req)
@@ -168,16 +170,11 @@ def execute(request)
           end
           # rubocop:enable Metrics/BlockLength
 
-          conn, _, response = enum.next
+          _, response = enum.next
           body = OpenAI::Internal::Util.fused_enum(enum, external: true) do
             finished = true
-            tap do
-              enum.next
-            rescue StopIteration
-              nil
-            end
+            loop { enum.next }
           ensure
-            conn.finish if !eof && conn&.started?
             closing&.call
           end
           [Integer(response.code), response, body]

lib/openai/internal/type/base_stream.rb

@@ -13,21 +13,6 @@ module Type
       module BaseStream
         include Enumerable
 
-        class << self
-          # Attempt to close the underlying transport when the stream itself is garbage
-          # collected.
-          #
-          # This should not be relied upon for resource clean up, as the garbage collector
-          # is not guaranteed to run.
-          #
-          # @param stream [Enumerable<Object>]
-          #
-          # @return [Proc]
-          #
-          # @see https://rubyapi.org/3.2/o/objectspace#method-c-define_finalizer
-          def defer_closing(stream) = ->(_id) { OpenAI::Internal::Util.close_fused!(stream) }
-        end
-
         # @return [Integer]
         attr_reader :status
 
@@ -82,8 +67,6 @@ def initialize(model:, url:, status:, headers:, response:, unwrap:, stream:)
           @unwrap = unwrap
           @stream = stream
           @iterator = iterator
-
-          ObjectSpace.define_finalizer(self, OpenAI::Internal::Type::BaseStream.defer_closing(@stream))
         end
 
         # @api private

rbi/openai/internal/type/base_stream.rbi

@@ -12,21 +12,6 @@ module OpenAI
         Message = type_member(:in)
         Elem = type_member(:out)
 
-        class << self
-          # Attempt to close the underlying transport when the stream itself is garbage
-          # collected.
-          #
-          # This should not be relied upon for resource clean up, as the garbage collector
-          # is not guaranteed to run.
-          sig do
-            params(stream: T::Enumerable[T.anything]).returns(
-              T.proc.params(arg0: Integer).void
-            )
-          end
-          def defer_closing(stream)
-          end
-        end
-
         sig { returns(Integer) }
         attr_reader :status
 

sig/openai/internal/type/base_stream.rbs

@@ -4,10 +4,6 @@ module OpenAI
       module BaseStream[Message, Elem]
         include Enumerable[Elem]
 
-        def self.defer_closing: (
-          Enumerable[top] stream
-        ) -> (^(Integer arg0) -> void)
-
         attr_reader status: Integer
 
         attr_reader headers: ::Hash[String, String]

test/openai/internal/util_test.rb

@@ -310,6 +310,31 @@ def test_copy_write
 end
 
 class OpenAI::Test::UtilFusedEnumTest < Minitest::Test
+  def test_rewind_closing
+    touched = false
+    once = 0
+    steps = 0
+    enum = Enumerator.new do |y|
+      next if touched
+
+      10.times do
+        steps = _1
+        y << _1
+      end
+    ensure
+      once = once.succ
+    end
+
+    fused = OpenAI::Internal::Util.fused_enum(enum, external: true) do
+      touched = true
+      loop { enum.next }
+    end
+    OpenAI::Internal::Util.close_fused!(fused)
+
+    assert_equal(1, once)
+    assert_equal(0, steps)
+  end
+
   def test_closing
     arr = [1, 2, 3]
     once = 0