From 8c6f9d8b29a569e5b908bee4153b3bf4d80b14ee Mon Sep 17 00:00:00 2001
From: raySavignone <ramon@zipline.inc>
Date: Wed, 2 Sep 2026 17:49:05 -0700
Subject: [PATCH] Fix: make StreamableHTTP#close and Native::Client#stop
 exception-safe

Wrap teardown in ensure so cleanup always runs when session termination
raises, while the original error still propagates. Prevents leaked SSE
threads, HTTPX clients, and per-client HumanInTheLoopRegistry entries in
long-lived processes that build short-lived StreamableHTTP clients.
---
 lib/ruby_llm/mcp/native/client.rb               |  1 +
 .../mcp/native/transports/streamable_http.rb    |  1 +
 spec/ruby_llm/mcp/native/client_spec.rb         | 17 +++++++++++++++++
 .../native/transports/streamable_http_spec.rb   | 14 ++++++++++++++
 4 files changed, 33 insertions(+)

diff --git a/lib/ruby_llm/mcp/native/client.rb b/lib/ruby_llm/mcp/native/client.rb
index da4d3c48..850638f1 100644
--- a/lib/ruby_llm/mcp/native/client.rb
+++ b/lib/ruby_llm/mcp/native/client.rb
@@ -124,6 +124,7 @@ def start
 
         def stop
           @transport&.close
+        ensure
           @capabilities = nil
           @transport = nil
           @task_registry = Native::TaskRegistry.new
diff --git a/lib/ruby_llm/mcp/native/transports/streamable_http.rb b/lib/ruby_llm/mcp/native/transports/streamable_http.rb
index b43b225c..bc31fcc0 100644
--- a/lib/ruby_llm/mcp/native/transports/streamable_http.rb
+++ b/lib/ruby_llm/mcp/native/transports/streamable_http.rb
@@ -135,6 +135,7 @@ def alive?
 
           def close
             terminate_session
+          ensure
             cleanup_sse_resources
             cleanup_connection
           end
diff --git a/spec/ruby_llm/mcp/native/client_spec.rb b/spec/ruby_llm/mcp/native/client_spec.rb
index 3397d1d7..5a7eaf43 100644
--- a/spec/ruby_llm/mcp/native/client_spec.rb
+++ b/spec/ruby_llm/mcp/native/client_spec.rb
@@ -23,6 +23,23 @@ def build_client(human_callback: nil, **options)
     )
   end
 
+  describe "#stop" do
+    it "releases the per-client human-in-the-loop registry even when transport teardown fails" do
+      client = build_client
+      owner_id = client.registry_owner_id
+      transport = instance_double(RubyLLM::MCP::Native::Transport)
+      allow(transport).to receive(:close).and_raise(StandardError, "teardown failed")
+      client.instance_variable_set(:@transport, transport)
+
+      expect(RubyLLM::MCP::Handlers::HumanInTheLoopRegistry)
+        .to receive(:release).with(owner_id).and_call_original
+
+      expect { client.stop }.to raise_error(StandardError, /teardown failed/)
+
+      expect(client.instance_variable_get(:@transport)).to be_nil
+    end
+  end
+
   describe "#execute_tool" do
     it "fails closed when callback returns non-normalized decision" do
       client = build_client(human_callback: ->(_name, _params) { true })
diff --git a/spec/ruby_llm/mcp/native/transports/streamable_http_spec.rb b/spec/ruby_llm/mcp/native/transports/streamable_http_spec.rb
index bb1e5a53..931e01fe 100644
--- a/spec/ruby_llm/mcp/native/transports/streamable_http_spec.rb
+++ b/spec/ruby_llm/mcp/native/transports/streamable_http_spec.rb
@@ -33,6 +33,20 @@
     allow(logger).to receive(:info)
   end
 
+  describe "#close" do
+    it "runs the remaining cleanup steps and re-raises when session termination fails" do
+      allow(transport).to receive(:terminate_session)
+        .and_raise(RubyLLM::MCP::Errors::TransportError.new(message: "session termination failed", code: nil))
+      expect(transport).to receive(:cleanup_sse_resources).and_call_original
+      expect(transport).to receive(:cleanup_connection).and_call_original
+
+      expect { transport.close }
+        .to raise_error(RubyLLM::MCP::Errors::TransportError, /session termination failed/)
+
+      expect(transport.instance_variable_get(:@connection)).to be_nil
+    end
+  end
+
   describe "#handle_oauth_authorization_error" do
     let(:response) { instance_double(HTTPX::Response) }