feat: handle thread interrupts in the core HTTP client

Author: stainless-app[bot]

lib/openai/internal/transport/pooled_net_requester.rb

@@ -128,40 +128,48 @@ def execute(request)
           url, deadline = request.fetch_values(:url, :deadline)
 
           req = nil
-          eof = false
           finished = false
-          closing = nil
 
           # rubocop:disable Metrics/BlockLength
           enum = Enumerator.new do |y|
             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
-
-              self.class.calibrate_socket_timeout(conn, deadline)
-              unless conn.started?
-                conn.keep_alive_timeout = self.class::KEEP_ALIVE_TIMEOUT
-                conn.start
-              end
+              eof = false
+              closing = nil
+              ::Thread.handle_interrupt(Object => :never) do
+                ::Thread.handle_interrupt(Object => :immediate) do
+                  req, closing = self.class.build_request(request) do
+                    self.class.calibrate_socket_timeout(conn, deadline)
+                  end
 
-              self.class.calibrate_socket_timeout(conn, deadline)
-              conn.request(req) do |rsp|
-                y << [req, rsp]
-                break if finished
-
-                rsp.read_body do |bytes|
-                  y << bytes.force_encoding(Encoding::BINARY)
-                  break if finished
+                  self.class.calibrate_socket_timeout(conn, deadline)
+                  unless conn.started?
+                    conn.keep_alive_timeout = self.class::KEEP_ALIVE_TIMEOUT
+                    conn.start
+                  end
 
                   self.class.calibrate_socket_timeout(conn, deadline)
+                  conn.request(req) do |rsp|
+                    y << [req, rsp]
+                    break if finished
+
+                    rsp.read_body do |bytes|
+                      y << bytes.force_encoding(Encoding::BINARY)
+                      break if finished
+
+                      self.class.calibrate_socket_timeout(conn, deadline)
+                    end
+                    eof = true
+                  end
+                end
+              ensure
+                begin
+                  conn.finish if !eof && conn&.started?
+                ensure
+                  closing&.call
                 end
-                eof = true
               end
-            ensure
-              conn.finish if !eof && conn&.started?
             end
           rescue Timeout::Error
             raise OpenAI::Errors::APITimeoutError.new(url: url, request: req)
@@ -174,8 +182,6 @@ def execute(request)
           body = OpenAI::Internal::Util.fused_enum(enum, external: true) do
             finished = true
             loop { enum.next }
-          ensure
-            closing&.call
           end
           [Integer(response.code), response, body]
         end

test/openai/internal/util_test.rb

@@ -343,6 +343,29 @@ def test_rewind_closing
     assert_equal(0, steps)
   end
 
+  def test_thread_interrupts
+    once = 0
+    que = Queue.new
+    enum = Enumerator.new do |y|
+      10.times { y << _1 }
+    ensure
+      once = once.succ
+    end
+
+    fused_1 = OpenAI::Internal::Util.fused_enum(enum, external: true) { loop { enum.next } }
+    fused_2 = OpenAI::Internal::Util.chain_fused(fused_1) { fused_1.each(&_1) }
+    fused_3 = OpenAI::Internal::Util.chain_fused(fused_2) { fused_2.each(&_1) }
+
+    th = ::Thread.new do
+      que << "🐶"
+      fused_3.each { sleep(10) }
+    end
+
+    assert_equal("🐶", que.pop)
+    th.kill.join
+    assert_equal(1, once)
+  end
+
   def test_closing
     arr = [1, 2, 3]
     once = 0