fix: add mutex protection to toolCache and provide thread-safe accessors in ClientSession

Author: guglielmo-san

mcp/client.go

@@ -325,6 +325,8 @@ type ClientSession struct {
 	pendingElicitationsMu sync.Mutex
 	pendingElicitations   map[string]chan struct{}
 
+	// toolCacheMu guards toolCache.
+	toolCacheMu sync.RWMutex
 	// toolCache stores tool definitions keyed by name.
 	// It is used to look up x-mcp-header annotations when
 	// constructing Mcp-Param-* headers for tools/call requests.
@@ -375,6 +377,8 @@ func (cs *ClientSession) Wait() error {
 }
 
 func (cs *ClientSession) cacheTools(tools []*Tool) {
+	cs.toolCacheMu.Lock()
+	defer cs.toolCacheMu.Unlock()
 	if cs.toolCache == nil {
 		cs.toolCache = make(map[string]*Tool, len(tools))
 	}
@@ -383,6 +387,12 @@ func (cs *ClientSession) cacheTools(tools []*Tool) {
 	}
 }
 
+func (cs *ClientSession) getCachedTool(name string) *Tool {
+	cs.toolCacheMu.RLock()
+	defer cs.toolCacheMu.RUnlock()
+	return cs.toolCache[name]
+}
+
 // registerElicitationWaiter registers a waiter for an elicitation complete
 // notification with the given elicitation ID. It returns two functions: an await
 // function that waits for the notification or context cancellation, and a cleanup
@@ -1021,7 +1031,7 @@ func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) (
 		// Avoid sending nil over the wire.
 		params.Arguments = map[string]any{}
 	}
-	if tool := cs.toolCache[params.Name]; tool != nil {
+	if tool := cs.getCachedTool(params.Name); tool != nil {
 		ctx = context.WithValue(ctx, toolContextKey, tool)
 	}
 	return handleSend[*CallToolResult](ctx, methodCallTool, newClientRequest(cs, orZero[Params](params)))

mcp/client_test.go

@@ -440,6 +440,80 @@ func TestClientCapabilities(t *testing.T) {
 	}
 }
 
+func TestToolCache(t *testing.T) {
+	tool1 := &Tool{Name: "tool1", Description: "first"}
+	tool2 := &Tool{Name: "tool2", Description: "second"}
+	tool1Updated := &Tool{Name: "tool1", Description: "updated"}
+
+	testCases := []struct {
+		name         string
+		cacheBatches [][]*Tool
+		lookup       string
+		want         *Tool
+	}{
+		{
+			name:   "empty cache",
+			lookup: "tool1",
+			want:   nil,
+		},
+		{
+			name:         "single tool found",
+			cacheBatches: [][]*Tool{{tool1}},
+			lookup:       "tool1",
+			want:         tool1,
+		},
+		{
+			name:         "unknown tool",
+			cacheBatches: [][]*Tool{{tool1}},
+			lookup:       "nonexistent",
+			want:         nil,
+		},
+		{
+			name:         "multiple tools single batch",
+			cacheBatches: [][]*Tool{{tool1, tool2}},
+			lookup:       "tool2",
+			want:         tool2,
+		},
+		{
+			name:         "additive first tool retained",
+			cacheBatches: [][]*Tool{{tool1}, {tool2}},
+			lookup:       "tool1",
+			want:         tool1,
+		},
+		{
+			name:         "additive second tool added",
+			cacheBatches: [][]*Tool{{tool1}, {tool2}},
+			lookup:       "tool2",
+			want:         tool2,
+		},
+		{
+			name:         "overwrite existing entry",
+			cacheBatches: [][]*Tool{{tool1}, {tool1Updated}},
+			lookup:       "tool1",
+			want:         tool1Updated,
+		},
+		{
+			name:         "empty batch no-op",
+			cacheBatches: [][]*Tool{{}},
+			lookup:       "tool1",
+			want:         nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			cs := &ClientSession{}
+			for _, batch := range tc.cacheBatches {
+				cs.cacheTools(batch)
+			}
+			got := cs.getCachedTool(tc.lookup)
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("getCachedTool(%q) mismatch (-want +got):\n%s", tc.lookup, diff)
+			}
+		})
+	}
+}
+
 func TestClientCapabilitiesOverWire(t *testing.T) {
 	testCases := []struct {
 		name             string

mcp/streamable_headers.go

@@ -27,9 +27,6 @@ const (
 	mcpHeaderExtension           = "x-mcp-header"
 )
 
-// ---------------------------------------------------------------------------
-// Shared helpers (used by both client and server)
-// ---------------------------------------------------------------------------
 
 func extractName(method string, params json.RawMessage) (string, bool) {
 	switch method {
@@ -122,10 +119,6 @@ func unmarshalPrimitive(raw json.RawMessage) any {
 	}
 }
 
-// ---------------------------------------------------------------------------
-// Client-side helpers
-// ---------------------------------------------------------------------------
-
 // setStandardHeaders populates standard MCP headers.
 // It requires the protocol version header to be set.
 func setStandardHeaders(header http.Header, msg jsonrpc.Message) {
@@ -279,10 +272,6 @@ func validateHeaderName(name string) error {
 	return nil
 }
 
-// ---------------------------------------------------------------------------
-// Server-side helpers
-// ---------------------------------------------------------------------------
-
 func validateMcpHeaders(header http.Header, msg jsonrpc.Message, tool *Tool) error {
 	protocolVersion := header.Get(protocolVersionHeader)
 	if protocolVersion == "" || protocolVersion < minVersionForStandardHeaders {