fix(mcptoolset): allow callers to close MCP sessions (#1713)

Implement io.Closer on MCP toolsets to release session resources when they are no longer needed. Keep closure final so in-flight failures cannot trigger reconnection.

Add focused lifecycle tests and use the public close method in the registry integration test.

Refs #1712

agentregistry/factory.go

@@ -343,6 +343,8 @@ func toA2ASkills(skills []Skill) []a2a.AgentSkill {
 // default, requests to *.googleapis.com endpoints are authenticated with the
 // registry's Application Default Credentials; use [WithMCPHTTPClient] and/or
 // [WithMCPHeaders] to override or augment egress.
+//
+// Callers must close the returned toolset as described in [mcptoolset.New].
 func (c *Client) MCPToolset(ctx context.Context, name string, opts ...MCPToolsetOption) (tool.Toolset, error) {
 	server, err := c.GetMCPServer(ctx, name)
 	if err != nil {

agentregistry/mcptoolset_e2e_test.go

@@ -17,6 +17,7 @@ package agentregistry
 import (
 	"context"
 	"encoding/json"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -43,9 +44,6 @@ func TestMCPToolset_E2E(t *testing.T) {
 	mcpSrv := httptest.NewServer(mcp.NewStreamableHTTPHandler(
 		func(*http.Request) *mcp.Server { return mcpServer }, nil))
 	defer mcpSrv.Close()
-	// The toolset holds a persistent streamable-HTTP connection with no public
-	// close; force-close it (runs before Close) so the server can shut down.
-	defer mcpSrv.CloseClientConnections()
 
 	regSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		rec := MCPServer{
@@ -63,6 +61,11 @@ func TestMCPToolset_E2E(t *testing.T) {
 	if err != nil {
 		t.Fatalf("MCPToolset() error = %v", err)
 	}
+	defer func() {
+		if err := ts.(io.Closer).Close(); err != nil {
+			t.Errorf("Close() error = %v", err)
+		}
+	}()
 
 	rctx := icontext.NewReadonlyContext(icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}))
 	tools, err := ts.Tools(rctx)

tool/mcptoolset/client.go

@@ -42,6 +42,10 @@ type connectionRefresher struct {
 
 	mu      sync.Mutex
 	session *mcp.ClientSession
+	closed  bool
+
+	closeOnce sync.Once
+	closeErr  error
 }
 
 // refreshableErrors is a list of errors that should trigger a connection refresh.
@@ -64,6 +68,22 @@ func newConnectionRefresher(client *mcp.Client, transport mcp.Transport) *connec
 	}
 }
 
+func (c *connectionRefresher) Close() error {
+	c.closeOnce.Do(func() {
+		c.mu.Lock()
+		// Closing is final so failed in-flight calls cannot reconnect.
+		c.closed = true
+		session := c.session
+		c.session = nil
+		c.mu.Unlock()
+
+		if session != nil {
+			c.closeErr = session.Close()
+		}
+	})
+	return c.closeErr
+}
+
 // CallTool calls a tool on the MCP server, automatically reconnecting if needed.
 func (c *connectionRefresher) CallTool(ctx context.Context, params *mcp.CallToolParams) (*mcp.CallToolResult, error) {
 	result, _, err := withRetry(ctx, c, func(session *mcp.ClientSession) (*mcp.CallToolResult, error) {
@@ -149,6 +169,9 @@ func (c *connectionRefresher) getSession(ctx context.Context) (*mcp.ClientSessio
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	if c.closed {
+		return nil, mcp.ErrConnectionClosed
+	}
 	if c.session != nil {
 		return c.session, nil
 	}
@@ -166,6 +189,10 @@ func (c *connectionRefresher) refreshConnection(ctx context.Context) (*mcp.Clien
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	if c.closed {
+		return nil, mcp.ErrConnectionClosed
+	}
+
 	// Ping to verify the connection is actually dead before reconnecting.
 	// This handles the case where another goroutine already reconnected.
 	if c.session != nil {

tool/mcptoolset/client_test.go

@@ -0,0 +1,61 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package mcptoolset
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+func TestClosePreventsRetry(t *testing.T) {
+	server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "1"}, nil)
+	connects := 0
+	transport := connectFunc(func(ctx context.Context) (mcp.Connection, error) {
+		connects++
+		ct, st := mcp.NewInMemoryTransports()
+		session, err := server.Connect(ctx, st, nil)
+		if err != nil {
+			return nil, err
+		}
+		t.Cleanup(func() { _ = session.Close() })
+		return ct.Connect(ctx)
+	})
+	client := newConnectionRefresher(nil, transport)
+	t.Cleanup(func() { _ = client.Close() })
+	calls := 0
+	_, _, err := withRetry(t.Context(), client, func(*mcp.ClientSession) (*mcp.CallToolResult, error) {
+		calls++
+		// The toolset closes after an operation starts, before its error is handled.
+		if err := client.Close(); err != nil {
+			t.Fatal(err)
+		}
+		return nil, mcp.ErrConnectionClosed
+	})
+	if !errors.Is(err, mcp.ErrConnectionClosed) {
+		t.Errorf("withRetry() error = %v, want ErrConnectionClosed", err)
+	}
+	if calls != 1 || connects != 1 {
+		t.Errorf("after Close(): %d calls and %d connections, want 1 each", calls, connects)
+	}
+}
+
+type connectFunc func(context.Context) (mcp.Connection, error)
+
+func (f connectFunc) Connect(ctx context.Context) (mcp.Connection, error) {
+	return f(ctx)
+}

tool/mcptoolset/set.go

@@ -17,6 +17,7 @@ package mcptoolset
 
 import (
 	"fmt"
+	"io"
 	"net/http"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
@@ -33,6 +34,13 @@ import (
 // MCP session is created lazily on the first request to LLM.
 // New returns an error if neither Transport nor Endpoint is provided.
 //
+// The returned toolset implements [io.Closer]. Callers must close it when it
+// is no longer in use. Close waits for ongoing connection setup and MCP
+// operations to finish, and is safe to call concurrently or more than once.
+// Further MCP operations fail with [mcp.ErrConnectionClosed]; closing does not
+// affect other sessions created with Config.Client.
+// When using [tool.FilterToolset], retain the original toolset to close it.
+//
 // Usage: create MCP ToolSet with mcptoolset.New() and provide it to the
 // LLMAgent in the llmagent.Config.
 //
@@ -146,12 +154,18 @@ type Config struct {
 }
 
 type set struct {
-	mcpClient                   MCPClient
+	mcpClient                   *connectionRefresher
 	toolFilter                  tool.Predicate
 	requireConfirmation         bool
 	requireConfirmationProvider tool.ConfirmationProvider
 }
 
+func (s *set) Close() error {
+	return s.mcpClient.Close()
+}
+
+var _ io.Closer = (*set)(nil)
+
 func (*set) Name() string {
 	return "mcp_tool_set"
 }

tool/mcptoolset/set_test.go

@@ -23,7 +23,9 @@ import (
 	"net/http"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
@@ -256,6 +258,77 @@ func TestToolFilter(t *testing.T) {
 	}
 }
 
+func TestMCPToolSetClose(t *testing.T) {
+	for _, connected := range []bool{false, true} {
+		t.Run(fmt.Sprintf("connected=%t", connected), func(t *testing.T) {
+			server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "1"}, nil)
+			mcp.AddTool(server, &mcp.Tool{Name: "get_weather"}, weatherFunc)
+			transport := &spyTransport{Transport: &reconnectableTransport{server: server}}
+			ts, err := mcptoolset.New(mcptoolset.Config{Transport: transport})
+			if err != nil {
+				t.Fatal(err)
+			}
+			t.Cleanup(func() {
+				for session := range server.Sessions() {
+					_ = session.Close()
+				}
+			})
+			closer, ok := ts.(io.Closer)
+			if !ok {
+				t.Fatal("MCP toolset does not implement io.Closer")
+			}
+			invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{})
+			ctx := icontext.NewReadonlyContext(invCtx)
+			var tools []tool.Tool
+			var sessionClosed <-chan struct{}
+			if connected {
+				tools, err = ts.Tools(ctx)
+				if err != nil {
+					t.Fatal(err)
+				}
+				for session := range server.Sessions() {
+					done := make(chan struct{})
+					sessionClosed = done
+					go func() {
+						_ = session.Wait()
+						close(done)
+					}()
+				}
+			}
+			var wg sync.WaitGroup
+			for range 2 {
+				wg.Go(func() {
+					if err := closer.Close(); err != nil {
+						t.Errorf("Close() failed: %v", err)
+					}
+				})
+			}
+			wg.Wait()
+			if connected {
+				select {
+				case <-sessionClosed:
+				case <-time.After(5 * time.Second):
+					t.Fatal("Close() left the MCP session open")
+				}
+				fnTool := tools[0].(toolinternal.FunctionTool)
+				if _, err := fnTool.Run(agent.NewToolContext(invCtx, "", nil, nil), map[string]any{"city": "Paris"}); !errors.Is(err, mcp.ErrConnectionClosed) {
+					t.Errorf("Run() after Close() error = %v, want ErrConnectionClosed", err)
+				}
+			}
+			if _, err := ts.Tools(ctx); !errors.Is(err, mcp.ErrConnectionClosed) {
+				t.Errorf("Tools() after Close() error = %v, want ErrConnectionClosed", err)
+			}
+			wantConnections := 0
+			if connected {
+				wantConnections = 1
+			}
+			if transport.connectCount != wantConnections {
+				t.Errorf("Connect() called %d times, want %d", transport.connectCount, wantConnections)
+			}
+		})
+	}
+}
+
 func TestListToolsReconnection(t *testing.T) {
 	server := mcp.NewServer(&mcp.Implementation{Name: "test_server", Version: "v1.0.0"}, nil)
 	mcp.AddTool(server, &mcp.Tool{Name: "get_weather", Description: "returns weather in the given city"}, weatherFunc)