From a3aca5db18c21f0cd54ef3c753bbb2db4784ad4c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 5 Apr 2026 22:15:18 +0000 Subject: [PATCH 1/3] Initial plan From e4881876ab848c6a25e1581f38148edebecb1e62 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 5 Apr 2026 23:20:11 +0000 Subject: [PATCH 2/3] go-sdk: extract registerToolWithoutValidation, use ParseToolArguments in mcptest, add filteredServerCache max-size - Extract registerToolWithoutValidation helper in tool_registry.go to deduplicate the server.AddTool bypass pattern (with long comment) from tool_registry.go and routed.go - Use mcp.ParseToolArguments in mcptest/server.go instead of inline json.Unmarshal - Add maxSize field (default 1000) to filteredServerCache with LRU eviction to prevent unbounded memory growth - Add tests: TestFilteredServerCache_MaxSize, TestFilteredServerCache_TTLEviction, TestRegisterToolWithoutValidation Agent-Logs-Url: https://github.com/github/gh-aw-mcpg/sessions/8fb7f3cd-ebb2-4acd-b60e-77bb0297dc9f Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com> --- internal/server/routed.go | 46 +++++++---- internal/server/routed_test.go | 123 ++++++++++++++++++++++++++++ internal/server/tool_registry.go | 34 ++++---- internal/testutil/mcptest/server.go | 22 +++-- 4 files changed, 180 insertions(+), 45 deletions(-) diff --git a/internal/server/routed.go b/internal/server/routed.go index e95e4a04..773242d9 100644 --- a/internal/server/routed.go +++ b/internal/server/routed.go @@ -31,12 +31,18 @@ func rejectIfShutdown(unifiedServer *UnifiedServer, next http.Handler, logNamesp }) } +// filteredServerCacheMaxSize is the maximum number of entries the filteredServerCache +// will hold. When the cache is full, the least-recently-used entry is evicted to make room. +const filteredServerCacheMaxSize = 1000 + // filteredServerCache caches filtered server instances per (backend, session) key. // Entries are evicted after the configured TTL to prevent unbounded memory growth -// in long-running deployments with many sessions. +// in long-running deployments with many sessions. A max-size cap provides an additional +// safety guard against an unbounded number of unique sessions. type filteredServerCache struct { servers map[string]*filteredServerEntry ttl time.Duration + maxSize int mu sync.RWMutex } @@ -50,11 +56,13 @@ func newFilteredServerCache(ttl time.Duration) *filteredServerCache { return &filteredServerCache{ servers: make(map[string]*filteredServerEntry), ttl: ttl, + maxSize: filteredServerCacheMaxSize, } } // getOrCreate returns a cached server or creates a new one. -// Expired entries are lazily evicted on each call. +// Expired entries are lazily evicted on each call. When the cache has reached its +// maximum size, the least-recently-used entry is evicted to make room. func (c *filteredServerCache) getOrCreate(backendID, sessionID string, creator func() *sdk.Server) *sdk.Server { key := fmt.Sprintf("%s/%s", backendID, sessionID) now := time.Now() @@ -75,6 +83,20 @@ func (c *filteredServerCache) getOrCreate(backendID, sessionID string, creator f return entry.server } + // Enforce max-size limit: evict the least-recently-used entry when at capacity. + if len(c.servers) >= c.maxSize { + var lruKey string + var lruTime time.Time + for k, entry := range c.servers { + if lruKey == "" || entry.lastUsed.Before(lruTime) { + lruKey = k + lruTime = entry.lastUsed + } + } + logRouted.Printf("[CACHE] Max size reached (%d), evicting LRU entry: key=%s", c.maxSize, lruKey) + delete(c.servers, lruKey) + } + logRouted.Printf("[CACHE] Creating new filtered server: backend=%s, session=%s", backendID, sessionID) server := creator() c.servers[key] = &filteredServerEntry{server: server, lastUsed: now} @@ -172,22 +194,16 @@ func createFilteredServer(unifiedServer *UnifiedServer, backendID string) *sdk.S continue } - // Use Server.AddTool method (not sdk.AddTool function) to avoid schema validation - // This allows including InputSchema from backends using different JSON Schema versions - // Wrap the typed handler to match the simple ToolHandler signature - wrappedHandler := func(ctx context.Context, req *sdk.CallToolRequest) (*sdk.CallToolResult, error) { - // Call the unified server's handler directly - // This ensures we go through the same session and connection pool - log.Printf("[ROUTED] Calling unified handler for: %s", toolNameCopy) - result, _, err := handler(ctx, req, nil) - return result, err - } - - server.AddTool(&sdk.Tool{ + // Use registerToolWithoutValidation to bypass JSON Schema validation, allowing + // InputSchema from backends using different JSON Schema versions (e.g., draft-07). + registerToolWithoutValidation(server, &sdk.Tool{ Name: toolInfo.Name, // Without prefix for the client Description: toolInfo.Description, InputSchema: toolInfo.InputSchema, // Include schema for clients - }, wrappedHandler) + }, func(ctx context.Context, req *sdk.CallToolRequest, _ interface{}) (*sdk.CallToolResult, interface{}, error) { + log.Printf("[ROUTED] Calling unified handler for: %s", toolNameCopy) + return handler(ctx, req, nil) + }) } return server diff --git a/internal/server/routed_test.go b/internal/server/routed_test.go index 8f61dac7..19b59192 100644 --- a/internal/server/routed_test.go +++ b/internal/server/routed_test.go @@ -6,6 +6,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -556,6 +557,128 @@ func TestCreateFilteredServer_EdgeCases(t *testing.T) { }) } +// TestFilteredServerCache_MaxSize verifies that the cache enforces its max-size limit +// by evicting the least-recently-used entry when the cache is full. +func TestFilteredServerCache_MaxSize(t *testing.T) { + assert := assert.New(t) + + ttl := time.Hour + cache := newFilteredServerCache(ttl) + cache.maxSize = 3 // Use a small max for the test + + callCount := 0 + creator := func() *sdk.Server { + callCount++ + return sdk.NewServer(&sdk.Implementation{Name: "test", Version: "1.0"}, &sdk.ServerOptions{}) + } + + // Fill the cache to max capacity + s1 := cache.getOrCreate("backend", "session1", creator) + s2 := cache.getOrCreate("backend", "session2", creator) + s3 := cache.getOrCreate("backend", "session3", creator) + assert.Equal(3, callCount, "Should have created 3 servers") + assert.NotNil(s1) + assert.NotNil(s2) + assert.NotNil(s3) + assert.Equal(3, len(cache.servers), "Cache should have 3 entries") + + // Touch session1 and session3 so session2 becomes the LRU + cache.getOrCreate("backend", "session1", creator) + cache.getOrCreate("backend", "session3", creator) + assert.Equal(3, callCount, "Cache hits should not create new servers") + + // Adding a fourth entry should evict the LRU (session2) + s4 := cache.getOrCreate("backend", "session4", creator) + assert.Equal(4, callCount, "Should have created a 4th server") + assert.NotNil(s4) + assert.Equal(3, len(cache.servers), "Cache should still be at max size (3)") + + // session2 should have been evicted + _, session2Exists := cache.servers["backend/session2"] + assert.False(session2Exists, "session2 should have been evicted as LRU") + + // session1, session3, and session4 should still be present + _, session1Exists := cache.servers["backend/session1"] + assert.True(session1Exists, "session1 should still be cached") + _, session3Exists := cache.servers["backend/session3"] + assert.True(session3Exists, "session3 should still be cached") + _, session4Exists := cache.servers["backend/session4"] + assert.True(session4Exists, "session4 should still be cached") +} + +// TestFilteredServerCache_TTLEviction verifies that expired entries are evicted. +func TestFilteredServerCache_TTLEviction(t *testing.T) { + assert := assert.New(t) + + ttl := 10 * time.Millisecond + cache := newFilteredServerCache(ttl) + + callCount := 0 + creator := func() *sdk.Server { + callCount++ + return sdk.NewServer(&sdk.Implementation{Name: "test", Version: "1.0"}, &sdk.ServerOptions{}) + } + + // Add an entry + cache.getOrCreate("backend", "session1", creator) + assert.Equal(1, callCount) + assert.Equal(1, len(cache.servers)) + + // Wait for TTL to expire + time.Sleep(20 * time.Millisecond) + + // Next call should evict the expired entry and create a new one + cache.getOrCreate("backend", "session2", creator) + assert.Equal(2, callCount, "Should have created a new server after TTL eviction") + + // session1 should have been evicted during the lazy eviction scan + _, session1Exists := cache.servers["backend/session1"] + assert.False(session1Exists, "Expired session1 should have been evicted") +} + +// TestRegisterToolWithoutValidation verifies that tools are registered on the server +// and that the wrapped handler forwards calls correctly via in-memory transport. +func TestRegisterToolWithoutValidation(t *testing.T) { + assert := assert.New(t) + require := require.New(t) + + server := sdk.NewServer(&sdk.Implementation{Name: "test", Version: "1.0"}, &sdk.ServerOptions{}) + + var handlerCalled bool + handler := func(ctx context.Context, req *sdk.CallToolRequest, state interface{}) (*sdk.CallToolResult, interface{}, error) { + handlerCalled = true + return &sdk.CallToolResult{IsError: false}, nil, nil + } + + registerToolWithoutValidation(server, &sdk.Tool{ + Name: "test_tool", + Description: "A test tool", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + }, + }, handler) + + // Use in-memory transports to connect a client to the server and invoke the tool + serverTransport, clientTransport := sdk.NewInMemoryTransports() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + go func() { + _ = server.Run(ctx, serverTransport) + }() + + client := sdk.NewClient(&sdk.Implementation{Name: "test-client", Version: "1.0"}, &sdk.ClientOptions{}) + clientSession, err := client.Connect(ctx, clientTransport, nil) + require.NoError(err) + defer clientSession.Close() + + result, err := clientSession.CallTool(ctx, &sdk.CallToolParams{Name: "test_tool"}) + require.NoError(err) + assert.False(result.IsError) + assert.True(handlerCalled, "Handler should have been called") +} + // TestCreateHTTPServerForRoutedMode_OAuth tests OAuth discovery endpoint in routed mode func TestCreateHTTPServerForRoutedMode_OAuth(t *testing.T) { tests := []struct { diff --git a/internal/server/tool_registry.go b/internal/server/tool_registry.go index 3ba14eda..423d7eee 100644 --- a/internal/server/tool_registry.go +++ b/internal/server/tool_registry.go @@ -24,6 +24,20 @@ type launchResult struct { duration time.Duration } +// registerToolWithoutValidation registers a tool with the SDK server using the Server.AddTool +// method (not the sdk.AddTool function) to bypass JSON Schema validation. This allows including +// InputSchema from backends that use different JSON Schema versions (e.g., draft-07) without +// validation errors, which is critical for clients to understand tool parameters. +// +// The handler's third parameter (pre-validated input) is passed as nil since argument +// unmarshaling is handled inside the handler itself. +func registerToolWithoutValidation(server *sdk.Server, tool *sdk.Tool, handler func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error)) { + server.AddTool(tool, func(ctx context.Context, req *sdk.CallToolRequest) (*sdk.CallToolResult, error) { + result, _, err := handler(ctx, req, nil) + return result, err + }) +} + // registerAllTools fetches and registers tools from all backend servers func (us *UnifiedServer) registerAllTools() error { log.Println("Registering tools from all backends...") @@ -235,28 +249,12 @@ func (us *UnifiedServer) registerToolsFromBackend(serverID string) error { us.tools[prefixedName].Handler = finalHandler us.toolsMu.Unlock() - // Register the tool with the SDK using the Server.AddTool method (not sdk.AddTool function) - // The method version does NOT perform schema validation, allowing us to include - // InputSchema from backends that use different JSON Schema versions (e.g., draft-07) - // without validation errors. This is critical for clients to understand tool parameters. - // - // We need to wrap our typed handler to match the simpler ToolHandler signature. - // The typed handler signature: func(context.Context, *CallToolRequest, interface{}) (*CallToolResult, interface{}, error) - // The simple handler signature: func(context.Context, *CallToolRequest) (*CallToolResult, error) - wrappedHandler := func(ctx context.Context, req *sdk.CallToolRequest) (*sdk.CallToolResult, error) { - // Call the final handler (which may include middleware wrapping) - // The third parameter would be the pre-unmarshaled/validated input if using sdk.AddTool, - // but we handle unmarshaling ourselves in the handler, so we pass nil - result, _, err := finalHandler(ctx, req, nil) - return result, err - } - - us.server.AddTool(&sdk.Tool{ + registerToolWithoutValidation(us.server, &sdk.Tool{ Name: prefixedName, Description: toolDesc, InputSchema: normalizedSchema, // Include the schema for clients to understand parameters Annotations: tool.Annotations, - }, wrappedHandler) + }, finalHandler) log.Printf("Registered tool: %s", logName) } diff --git a/internal/testutil/mcptest/server.go b/internal/testutil/mcptest/server.go index a8491918..198132b3 100644 --- a/internal/testutil/mcptest/server.go +++ b/internal/testutil/mcptest/server.go @@ -2,10 +2,10 @@ package mcptest import ( "context" - "encoding/json" "fmt" "log" + "github.com/github/gh-aw-mcpg/internal/mcp" sdk "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -48,18 +48,16 @@ func (s *Server) Start() error { Description: tool.Description, InputSchema: tool.InputSchema, }, func(ctx context.Context, req *sdk.CallToolRequest) (*sdk.CallToolResult, error) { - var args map[string]interface{} - if len(req.Params.Arguments) > 0 { - if err := json.Unmarshal(req.Params.Arguments, &args); err != nil { - return &sdk.CallToolResult{ - IsError: true, - Content: []sdk.Content{ - &sdk.TextContent{ - Text: fmt.Sprintf("Failed to parse arguments: %v", err), - }, + args, err := mcp.ParseToolArguments(req) + if err != nil { + return &sdk.CallToolResult{ + IsError: true, + Content: []sdk.Content{ + &sdk.TextContent{ + Text: fmt.Sprintf("Failed to parse arguments: %v", err), }, - }, nil - } + }, + }, nil } content, err := tool.Handler(args) From 7b0dd7e268d62100b9c54a1e26616a324451f4bb Mon Sep 17 00:00:00 2001 From: Landon Cox Date: Sun, 5 Apr 2026 17:00:20 -0700 Subject: [PATCH 3/3] Address PR review feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix duplicate error wrapping in mcptest: change 'Failed to parse arguments' to 'Failed to parse tool arguments' since ParseToolArguments already wraps with 'failed to parse arguments' - Remove LRU eviction of active cache entries: routed mode relies on reusing the same filtered server instance per (backend, session), so evicting non-expired entries could break StreamableHTTP session semantics. Instead log a warning and allow cache growth until TTL eviction. - Truncate session IDs in cache log messages using auth.TruncateSessionID to avoid leaking secrets into debug logs - Increase TTL test margins (10ms→100ms TTL, 20ms→200ms sleep) to reduce flakiness under CI load - Fix search_repositories URL encoding (url.QueryEscape) for query param Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- internal/proxy/proxy.go | 2 +- internal/server/routed.go | 21 ++++++++------------- internal/server/routed_test.go | 27 ++++++++++----------------- internal/testutil/mcptest/server.go | 2 +- 4 files changed, 20 insertions(+), 32 deletions(-) diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 581e1619..eea9a2fe 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -324,7 +324,7 @@ func (r *restBackendCaller) CallTool(ctx context.Context, toolName string, args if pp, ok := argsMap["perPage"].(float64); ok { perPage = fmt.Sprintf("%d", int(pp)) } - apiPath = fmt.Sprintf("/search/repositories?q=%s&per_page=%s", query, perPage) + apiPath = fmt.Sprintf("/search/repositories?q=%s&per_page=%s", url.QueryEscape(query), perPage) case "get_collaborator_permission": owner, _ := argsMap["owner"].(string) diff --git a/internal/server/routed.go b/internal/server/routed.go index 773242d9..fd24a144 100644 --- a/internal/server/routed.go +++ b/internal/server/routed.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "github.com/github/gh-aw-mcpg/internal/auth" "github.com/github/gh-aw-mcpg/internal/httputil" "github.com/github/gh-aw-mcpg/internal/logger" "github.com/github/gh-aw-mcpg/internal/version" @@ -73,7 +74,7 @@ func (c *filteredServerCache) getOrCreate(backendID, sessionID string, creator f // Lazy eviction of expired entries for k, entry := range c.servers { if now.Sub(entry.lastUsed) > c.ttl { - logRouted.Printf("[CACHE] Evicting expired server: key=%s (idle %s)", k, now.Sub(entry.lastUsed).Round(time.Second)) + logRouted.Printf("[CACHE] Evicting expired server: key=%s (idle %s)", auth.TruncateSessionID(k), now.Sub(entry.lastUsed).Round(time.Second)) delete(c.servers, k) } } @@ -83,21 +84,15 @@ func (c *filteredServerCache) getOrCreate(backendID, sessionID string, creator f return entry.server } - // Enforce max-size limit: evict the least-recently-used entry when at capacity. + // Safety bound: if at capacity after TTL eviction, log a warning but do not + // evict non-expired entries. Routed mode relies on reusing the same filtered + // server instance for a given (backend, session), and evicting an active entry + // would recreate that server mid-session, breaking StreamableHTTP semantics. if len(c.servers) >= c.maxSize { - var lruKey string - var lruTime time.Time - for k, entry := range c.servers { - if lruKey == "" || entry.lastUsed.Before(lruTime) { - lruKey = k - lruTime = entry.lastUsed - } - } - logRouted.Printf("[CACHE] Max size reached (%d), evicting LRU entry: key=%s", c.maxSize, lruKey) - delete(c.servers, lruKey) + logRouted.Printf("[CACHE] Max size reached (%d), retaining active entries until TTL eviction", c.maxSize) } - logRouted.Printf("[CACHE] Creating new filtered server: backend=%s, session=%s", backendID, sessionID) + logRouted.Printf("[CACHE] Creating new filtered server: backend=%s, session=%s", backendID, auth.TruncateSessionID(sessionID)) server := creator() c.servers[key] = &filteredServerEntry{server: server, lastUsed: now} return server diff --git a/internal/server/routed_test.go b/internal/server/routed_test.go index 19b59192..131cd979 100644 --- a/internal/server/routed_test.go +++ b/internal/server/routed_test.go @@ -557,8 +557,8 @@ func TestCreateFilteredServer_EdgeCases(t *testing.T) { }) } -// TestFilteredServerCache_MaxSize verifies that the cache enforces its max-size limit -// by evicting the least-recently-used entry when the cache is full. +// TestFilteredServerCache_MaxSize verifies that the cache allows growth beyond maxSize +// when all entries are still active (non-expired), to avoid disrupting sessions. func TestFilteredServerCache_MaxSize(t *testing.T) { assert := assert.New(t) @@ -582,24 +582,17 @@ func TestFilteredServerCache_MaxSize(t *testing.T) { assert.NotNil(s3) assert.Equal(3, len(cache.servers), "Cache should have 3 entries") - // Touch session1 and session3 so session2 becomes the LRU - cache.getOrCreate("backend", "session1", creator) - cache.getOrCreate("backend", "session3", creator) - assert.Equal(3, callCount, "Cache hits should not create new servers") - - // Adding a fourth entry should evict the LRU (session2) + // Adding a fourth entry should be allowed (no LRU eviction of active sessions) s4 := cache.getOrCreate("backend", "session4", creator) assert.Equal(4, callCount, "Should have created a 4th server") assert.NotNil(s4) - assert.Equal(3, len(cache.servers), "Cache should still be at max size (3)") + assert.Equal(4, len(cache.servers), "Cache should grow beyond maxSize for active sessions") - // session2 should have been evicted - _, session2Exists := cache.servers["backend/session2"] - assert.False(session2Exists, "session2 should have been evicted as LRU") - - // session1, session3, and session4 should still be present + // All sessions should still be present _, session1Exists := cache.servers["backend/session1"] assert.True(session1Exists, "session1 should still be cached") + _, session2Exists := cache.servers["backend/session2"] + assert.True(session2Exists, "session2 should still be cached") _, session3Exists := cache.servers["backend/session3"] assert.True(session3Exists, "session3 should still be cached") _, session4Exists := cache.servers["backend/session4"] @@ -610,7 +603,7 @@ func TestFilteredServerCache_MaxSize(t *testing.T) { func TestFilteredServerCache_TTLEviction(t *testing.T) { assert := assert.New(t) - ttl := 10 * time.Millisecond + ttl := 100 * time.Millisecond cache := newFilteredServerCache(ttl) callCount := 0 @@ -624,8 +617,8 @@ func TestFilteredServerCache_TTLEviction(t *testing.T) { assert.Equal(1, callCount) assert.Equal(1, len(cache.servers)) - // Wait for TTL to expire - time.Sleep(20 * time.Millisecond) + // Wait for TTL to expire (use generous margin to avoid CI flakiness) + time.Sleep(200 * time.Millisecond) // Next call should evict the expired entry and create a new one cache.getOrCreate("backend", "session2", creator) diff --git a/internal/testutil/mcptest/server.go b/internal/testutil/mcptest/server.go index 198132b3..b3e5c51c 100644 --- a/internal/testutil/mcptest/server.go +++ b/internal/testutil/mcptest/server.go @@ -54,7 +54,7 @@ func (s *Server) Start() error { IsError: true, Content: []sdk.Content{ &sdk.TextContent{ - Text: fmt.Sprintf("Failed to parse arguments: %v", err), + Text: fmt.Sprintf("Failed to parse tool arguments: %v", err), }, }, }, nil