From e5149725684679180c424e496ba984baa77151f4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 28 Apr 2026 01:53:28 +0000 Subject: [PATCH 1/3] Initial plan From 2ddc6ee716e552c00f837e0be81a2c4f41512854 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 28 Apr 2026 02:02:30 +0000 Subject: [PATCH 2/3] refactor: split connection.go, extract response transforms, split httputil Agent-Logs-Url: https://github.com/github/gh-aw-mcpg/sessions/157ad9ad-2fb1-4217-89cc-7854435a6d50 Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com> --- internal/httputil/github_http.go | 40 ++++ internal/httputil/httputil.go | 35 ---- internal/mcp/connection.go | 279 --------------------------- internal/mcp/connection_logging.go | 54 ++++++ internal/mcp/connection_methods.go | 116 +++++++++++ internal/mcp/pagination.go | 129 +++++++++++++ internal/proxy/handler.go | 142 -------------- internal/proxy/response_transform.go | 145 ++++++++++++++ 8 files changed, 484 insertions(+), 456 deletions(-) create mode 100644 internal/httputil/github_http.go create mode 100644 internal/mcp/connection_logging.go create mode 100644 internal/mcp/connection_methods.go create mode 100644 internal/mcp/pagination.go create mode 100644 internal/proxy/response_transform.go diff --git a/internal/httputil/github_http.go b/internal/httputil/github_http.go new file mode 100644 index 00000000..ae6d4b87 --- /dev/null +++ b/internal/httputil/github_http.go @@ -0,0 +1,40 @@ +package httputil + +import ( + "net/http" + "strconv" + "strings" + "time" +) + +// GitHubUserAgent is the User-Agent header value sent on all GitHub API requests. +const GitHubUserAgent = "awmg/1.0" + +// ApplyGitHubAPIHeaders sets the standard GitHub API request headers on req. +// authHeader should be the full Authorization header value (e.g. "token xyz" or +// "Bearer xyz"). When authHeader is empty no Authorization header is set, which +// is appropriate when the caller has already decided that no auth is available. +func ApplyGitHubAPIHeaders(req *http.Request, authHeader string) { + if authHeader != "" { + req.Header.Set("Authorization", authHeader) + } + req.Header.Set("Accept", "application/vnd.github+json") + req.Header.Set("User-Agent", GitHubUserAgent) +} + +// ParseRateLimitResetHeader parses the Unix-timestamp value of the +// X-RateLimit-Reset HTTP header into a time.Time. +// Returns zero time when the header value is absent or malformed. +func ParseRateLimitResetHeader(value string) time.Time { + if value == "" { + return time.Time{} + } + unix, err := strconv.ParseInt(strings.TrimSpace(value), 10, 64) + if err != nil { + logHTTP.Printf("Failed to parse X-RateLimit-Reset header value=%q: %v", value, err) + return time.Time{} + } + reset := time.Unix(unix, 0) + logHTTP.Printf("Parsed X-RateLimit-Reset: resetAt=%s", reset.UTC().Format(time.RFC3339)) + return reset +} diff --git a/internal/httputil/httputil.go b/internal/httputil/httputil.go index befe0f28..2c55c2a9 100644 --- a/internal/httputil/httputil.go +++ b/internal/httputil/httputil.go @@ -5,9 +5,6 @@ package httputil import ( "encoding/json" "net/http" - "strconv" - "strings" - "time" "github.com/github/gh-aw-mcpg/internal/logger" ) @@ -36,38 +33,6 @@ func WriteJSONResponse(w http.ResponseWriter, statusCode int, body interface{}) } } -// ParseRateLimitResetHeader parses the Unix-timestamp value of the -// X-RateLimit-Reset HTTP header into a time.Time. -// Returns zero time when the header value is absent or malformed. -func ParseRateLimitResetHeader(value string) time.Time { - if value == "" { - return time.Time{} - } - unix, err := strconv.ParseInt(strings.TrimSpace(value), 10, 64) - if err != nil { - logHTTP.Printf("Failed to parse X-RateLimit-Reset header value=%q: %v", value, err) - return time.Time{} - } - reset := time.Unix(unix, 0) - logHTTP.Printf("Parsed X-RateLimit-Reset: resetAt=%s", reset.UTC().Format(time.RFC3339)) - return reset -} - -// GitHubUserAgent is the User-Agent header value sent on all GitHub API requests. -const GitHubUserAgent = "awmg/1.0" - -// ApplyGitHubAPIHeaders sets the standard GitHub API request headers on req. -// authHeader should be the full Authorization header value (e.g. "token xyz" or -// "Bearer xyz"). When authHeader is empty no Authorization header is set, which -// is appropriate when the caller has already decided that no auth is available. -func ApplyGitHubAPIHeaders(req *http.Request, authHeader string) { - if authHeader != "" { - req.Header.Set("Authorization", authHeader) - } - req.Header.Set("Accept", "application/vnd.github+json") - req.Header.Set("User-Agent", GitHubUserAgent) -} - // IsTransientHTTPError returns true for status codes that indicate a temporary // server-side condition (rate-limiting or transient failure) worth retrying. func IsTransientHTTPError(statusCode int) bool { diff --git a/internal/mcp/connection.go b/internal/mcp/connection.go index bb273ec7..89f05fe5 100644 --- a/internal/mcp/connection.go +++ b/internal/mcp/connection.go @@ -314,21 +314,6 @@ func (c *Connection) ServerInfo() (name, version string) { return initResult.ServerInfo.Name, initResult.ServerInfo.Version } -// logReconnectStart emits the structured log warning that is common to all reconnect paths. -func (c *Connection) logReconnectStart() { - logger.LogWarn("backend", "MCP session expired for %s, attempting to reconnect...", c.serverID) -} - -// logReconnectResult emits the structured log entry that signals whether the reconnect -// succeeded or failed. It is the common success/failure telemetry shared by all reconnect paths. -func (c *Connection) logReconnectResult(err error) { - if err != nil { - logger.LogError("backend", "Session reconnect failed for %s: %v", c.serverID, err) - } else { - logger.LogInfo("backend", "Session successfully reconnected for %s", c.serverID) - } -} - // reconnectPlainJSON re-initialises the plain JSON-RPC session with the HTTP backend. // It is safe for concurrent callers: only one reconnect runs at a time, and the updated // session ID is available to all callers once the lock is released. @@ -422,38 +407,6 @@ func (c *Connection) callSDKMethodWithReconnect(method string, params interface{ return result, err } -// logOutboundRPCRequest logs an outbound RPC request, optionally attaching agent DIFC tag snapshots. -// When shouldAttachTags is true, snapshot must be non-nil. -func logOutboundRPCRequest(serverID string, method string, payload []byte, shouldAttachTags bool, snapshot *AgentTagsSnapshot) { - if shouldAttachTags { - logger.LogRPCRequestWithAgentSnapshot(logger.RPCDirectionOutbound, serverID, method, payload, snapshot.Secrecy, snapshot.Integrity) - } else { - logger.LogRPCRequest(logger.RPCDirectionOutbound, serverID, method, payload) - } -} - -// logInboundRPCResponse logs an inbound RPC response, optionally attaching agent DIFC tag snapshots. -// When shouldAttachTags is true, snapshot must be non-nil. -func logInboundRPCResponse(serverID string, payload []byte, err error, shouldAttachTags bool, snapshot *AgentTagsSnapshot) { - if shouldAttachTags { - logger.LogRPCResponseWithAgentSnapshot(logger.RPCDirectionInbound, serverID, payload, err, snapshot.Secrecy, snapshot.Integrity) - } else { - logger.LogRPCResponse(logger.RPCDirectionInbound, serverID, payload, err) - } -} - -// logInboundRPCResponseFromResult attempts to marshal a response payload for logging, -// silently ignores marshal failures, logs the inbound response, and returns the -// original result and error unchanged. -func logInboundRPCResponseFromResult(serverID string, result *Response, err error, shouldAttachTags bool, snapshot *AgentTagsSnapshot) (*Response, error) { - var responsePayload []byte - if result != nil { - responsePayload, _ = json.Marshal(result) - } - logInboundRPCResponse(serverID, responsePayload, err, shouldAttachTags, snapshot) - return result, err -} - // SendRequest sends a JSON-RPC request and waits for the response // The serverID parameter is used for logging to associate the request with a backend server func (c *Connection) SendRequest(method string, params interface{}) (*Response, error) { @@ -494,49 +447,6 @@ func (c *Connection) SendRequestWithServerID(ctx context.Context, method string, return logInboundRPCResponseFromResult(serverID, result, err, shouldAttachAgentTags, snapshot) } -// callSDKMethod calls the appropriate SDK method based on the method name -// This centralizes the method dispatch logic used by both HTTP SDK transports and stdio -func (c *Connection) callSDKMethod(method string, params interface{}) (*Response, error) { - logConn.Printf("Dispatching SDK method: %s, serverID=%s", method, c.serverID) - switch method { - case "tools/list": - return c.listTools() - case "tools/call": - return c.callTool(params) - case "resources/list": - return c.listResources() - case "resources/read": - return c.readResource(params) - case "prompts/list": - return c.listPrompts() - case "prompts/get": - return c.getPrompt(params) - default: - logConn.Printf("Unsupported method: %s", method) - return nil, fmt.Errorf("unsupported method: %s", method) - } -} - -// marshalToResponse marshals an SDK result into a Response object. -// This helper reduces code duplication across all MCP method wrappers. -// -// The ID field is set to a static placeholder (1) because this Response is only -// constructed after the SDK's session.XXX() call has already resolved the -// request–response correlation internally. The gateway never uses this ID for -// matching; it is present solely to satisfy the JSON-RPC 2.0 structure. -func marshalToResponse(result interface{}) (*Response, error) { - resultJSON, err := json.Marshal(result) - if err != nil { - return nil, fmt.Errorf("failed to marshal result: %w", err) - } - - return &Response{ - JSONRPC: "2.0", - ID: 1, // Placeholder – see function comment for safety rationale - Result: resultJSON, - }, nil -} - // requireSession validates that a session is available for SDK operations. // This helper centralizes session validation logic across all MCP method wrappers. // Returns an error if the session is nil (e.g., for plain JSON-RPC transport). @@ -547,195 +457,6 @@ func (c *Connection) requireSession() error { return nil } -// unmarshalParams converts generic interface{} params to a specific struct type. -// This helper reduces code duplication across MCP method wrappers and ensures -// consistent error handling for parameter conversion. It uses marshal/unmarshal -// to maintain JSON schema validation benefits. -func unmarshalParams(params interface{}, target interface{}) error { - paramsJSON, err := json.Marshal(params) - if err != nil { - return fmt.Errorf("failed to marshal params: %w", err) - } - if err := json.Unmarshal(paramsJSON, target); err != nil { - return fmt.Errorf("invalid params: %w", err) - } - return nil -} - -// callParamMethod is a generic helper for SDK operations that require typed parameters. -// It handles the common pattern of: requireSession → unmarshalParams → fn(params) → marshalToResponse. -// P is the type of the parameter struct to unmarshal into. -func callParamMethod[P any](c *Connection, rawParams interface{}, fn func(P) (interface{}, error)) (*Response, error) { - if err := c.requireSession(); err != nil { - return nil, err - } - var params P - if err := unmarshalParams(rawParams, ¶ms); err != nil { - return nil, err - } - result, err := fn(params) - if err != nil { - return nil, err - } - return marshalToResponse(result) -} - -// paginatedPage holds a single page of results from a paginated SDK list call. -type paginatedPage[T any] struct { - Items []T - NextCursor string -} - -// paginateAllMaxPages is the maximum number of pages that paginateAll will fetch. -// This guards against misbehaving or adversarial backends that return an unbounded -// sequence of pages, which would otherwise consume unbounded memory and time. -const paginateAllMaxPages = 100 - -// paginateAll collects all items across paginated SDK list calls. -// It returns an error if the backend returns more than paginateAllMaxPages pages, -// protecting against runaway backends. -func paginateAll[T any]( - serverID string, - itemKind string, - fetch func(cursor string) (paginatedPage[T], error), -) ([]T, error) { - first, err := fetch("") - if err != nil { - return nil, err - } - all := make([]T, len(first.Items), max(len(first.Items), 1)) - copy(all, first.Items) - logConn.Printf("list%s: received page of %d %s from serverID=%s", itemKind, len(first.Items), itemKind, serverID) - - cursor := first.NextCursor - seenCursors := make(map[string]struct{}) - for pageCount := 1; cursor != ""; pageCount++ { - if pageCount >= paginateAllMaxPages { - return nil, fmt.Errorf("list%s: backend serverID=%s returned more than %d pages; aborting to prevent unbounded memory growth", itemKind, serverID, paginateAllMaxPages) - } - if _, seen := seenCursors[cursor]; seen { - return nil, fmt.Errorf("list%s: backend serverID=%s returned cyclical cursor %q", itemKind, serverID, cursor) - } - seenCursors[cursor] = struct{}{} - page, err := fetch(cursor) - if err != nil { - return nil, err - } - all = append(all, page.Items...) - logConn.Printf("list%s: received page of %d %s (total so far: %d) from serverID=%s", itemKind, len(page.Items), itemKind, len(all), serverID) - cursor = page.NextCursor - } - logConn.Printf("list%s: received %d %s total from serverID=%s", itemKind, len(all), itemKind, serverID) - return all, nil -} - -// listMCPItems is a generic helper for the list* family of MCP operations. -// It handles session validation, logging, pagination, and response marshalling, -// eliminating the boilerplate that was previously duplicated across listTools, -// listResources, and listPrompts. -func listMCPItems[Item any, Result any]( - c *Connection, - kind string, - fetchPage func(cursor string) (paginatedPage[Item], error), - buildResult func([]Item) Result, -) (*Response, error) { - if err := c.requireSession(); err != nil { - return nil, err - } - logConn.Printf("list%s: requesting %s list from backend serverID=%s", kind, kind, c.serverID) - items, err := paginateAll(c.serverID, kind, fetchPage) - if err != nil { - return nil, err - } - return marshalToResponse(buildResult(items)) -} - -func (c *Connection) listTools() (*Response, error) { - return listMCPItems(c, "tools", - func(cursor string) (paginatedPage[*sdk.Tool], error) { - result, err := c.getSDKSession().ListTools(c.ctx, &sdk.ListToolsParams{Cursor: cursor}) - if err != nil { - return paginatedPage[*sdk.Tool]{}, err - } - return paginatedPage[*sdk.Tool]{Items: result.Tools, NextCursor: result.NextCursor}, nil - }, - func(items []*sdk.Tool) *sdk.ListToolsResult { - return &sdk.ListToolsResult{Tools: items} - }, - ) -} - -func (c *Connection) callTool(params interface{}) (*Response, error) { - return callParamMethod(c, params, func(p CallToolParams) (interface{}, error) { - // Ensure arguments is never nil - default to empty map - // This is required by the MCP protocol which expects arguments to always be present - if p.Arguments == nil { - p.Arguments = make(map[string]interface{}) - } - logConn.Printf("callTool: parsed name=%s, arguments=%+v", p.Name, p.Arguments) - return c.getSDKSession().CallTool(c.ctx, &sdk.CallToolParams{ - Name: p.Name, - Arguments: p.Arguments, - }) - }) -} - -func (c *Connection) listResources() (*Response, error) { - return listMCPItems(c, "resources", - func(cursor string) (paginatedPage[*sdk.Resource], error) { - result, err := c.getSDKSession().ListResources(c.ctx, &sdk.ListResourcesParams{Cursor: cursor}) - if err != nil { - return paginatedPage[*sdk.Resource]{}, err - } - return paginatedPage[*sdk.Resource]{Items: result.Resources, NextCursor: result.NextCursor}, nil - }, - func(items []*sdk.Resource) *sdk.ListResourcesResult { - return &sdk.ListResourcesResult{Resources: items} - }, - ) -} - -func (c *Connection) readResource(params interface{}) (*Response, error) { - type readResourceParams struct { - URI string `json:"uri"` - } - return callParamMethod(c, params, func(p readResourceParams) (interface{}, error) { - logConn.Printf("readResource: reading resource uri=%s from serverID=%s", p.URI, c.serverID) - return c.getSDKSession().ReadResource(c.ctx, &sdk.ReadResourceParams{ - URI: p.URI, - }) - }) -} - -func (c *Connection) listPrompts() (*Response, error) { - return listMCPItems(c, "prompts", - func(cursor string) (paginatedPage[*sdk.Prompt], error) { - result, err := c.getSDKSession().ListPrompts(c.ctx, &sdk.ListPromptsParams{Cursor: cursor}) - if err != nil { - return paginatedPage[*sdk.Prompt]{}, err - } - return paginatedPage[*sdk.Prompt]{Items: result.Prompts, NextCursor: result.NextCursor}, nil - }, - func(items []*sdk.Prompt) *sdk.ListPromptsResult { - return &sdk.ListPromptsResult{Prompts: items} - }, - ) -} - -func (c *Connection) getPrompt(params interface{}) (*Response, error) { - type getPromptParams struct { - Name string `json:"name"` - Arguments map[string]string `json:"arguments"` - } - return callParamMethod(c, params, func(p getPromptParams) (interface{}, error) { - logConn.Printf("getPrompt: getting prompt name=%s from serverID=%s", p.Name, c.serverID) - return c.getSDKSession().GetPrompt(c.ctx, &sdk.GetPromptParams{ - Name: p.Name, - Arguments: p.Arguments, - }) - }) -} - // Close closes the connection func (c *Connection) Close() error { logConn.Printf("Closing connection: serverID=%s, isHTTP=%v", c.serverID, c.isHTTP) diff --git a/internal/mcp/connection_logging.go b/internal/mcp/connection_logging.go new file mode 100644 index 00000000..90dc17c1 --- /dev/null +++ b/internal/mcp/connection_logging.go @@ -0,0 +1,54 @@ +package mcp + +import ( + "encoding/json" + + "github.com/github/gh-aw-mcpg/internal/logger" +) + +// logReconnectStart emits the structured log warning that is common to all reconnect paths. +func (c *Connection) logReconnectStart() { + logger.LogWarn("backend", "MCP session expired for %s, attempting to reconnect...", c.serverID) +} + +// logReconnectResult emits the structured log entry that signals whether the reconnect +// succeeded or failed. It is the common success/failure telemetry shared by all reconnect paths. +func (c *Connection) logReconnectResult(err error) { + if err != nil { + logger.LogError("backend", "Session reconnect failed for %s: %v", c.serverID, err) + } else { + logger.LogInfo("backend", "Session successfully reconnected for %s", c.serverID) + } +} + +// logOutboundRPCRequest logs an outbound RPC request, optionally attaching agent DIFC tag snapshots. +// When shouldAttachTags is true, snapshot must be non-nil. +func logOutboundRPCRequest(serverID string, method string, payload []byte, shouldAttachTags bool, snapshot *AgentTagsSnapshot) { + if shouldAttachTags { + logger.LogRPCRequestWithAgentSnapshot(logger.RPCDirectionOutbound, serverID, method, payload, snapshot.Secrecy, snapshot.Integrity) + } else { + logger.LogRPCRequest(logger.RPCDirectionOutbound, serverID, method, payload) + } +} + +// logInboundRPCResponse logs an inbound RPC response, optionally attaching agent DIFC tag snapshots. +// When shouldAttachTags is true, snapshot must be non-nil. +func logInboundRPCResponse(serverID string, payload []byte, err error, shouldAttachTags bool, snapshot *AgentTagsSnapshot) { + if shouldAttachTags { + logger.LogRPCResponseWithAgentSnapshot(logger.RPCDirectionInbound, serverID, payload, err, snapshot.Secrecy, snapshot.Integrity) + } else { + logger.LogRPCResponse(logger.RPCDirectionInbound, serverID, payload, err) + } +} + +// logInboundRPCResponseFromResult attempts to marshal a response payload for logging, +// silently ignores marshal failures, logs the inbound response, and returns the +// original result and error unchanged. +func logInboundRPCResponseFromResult(serverID string, result *Response, err error, shouldAttachTags bool, snapshot *AgentTagsSnapshot) (*Response, error) { + var responsePayload []byte + if result != nil { + responsePayload, _ = json.Marshal(result) + } + logInboundRPCResponse(serverID, responsePayload, err, shouldAttachTags, snapshot) + return result, err +} diff --git a/internal/mcp/connection_methods.go b/internal/mcp/connection_methods.go new file mode 100644 index 00000000..c1fe18ac --- /dev/null +++ b/internal/mcp/connection_methods.go @@ -0,0 +1,116 @@ +package mcp + +import ( + "fmt" + + sdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// callSDKMethod calls the appropriate SDK method based on the method name. +// This centralizes the method dispatch logic used by both HTTP SDK transports and stdio. +func (c *Connection) callSDKMethod(method string, params interface{}) (*Response, error) { + logConn.Printf("Dispatching SDK method: %s, serverID=%s", method, c.serverID) + switch method { + case "tools/list": + return c.listTools() + case "tools/call": + return c.callTool(params) + case "resources/list": + return c.listResources() + case "resources/read": + return c.readResource(params) + case "prompts/list": + return c.listPrompts() + case "prompts/get": + return c.getPrompt(params) + default: + logConn.Printf("Unsupported method: %s", method) + return nil, fmt.Errorf("unsupported method: %s", method) + } +} + +func (c *Connection) listTools() (*Response, error) { + return listMCPItems(c, "tools", + func(cursor string) (paginatedPage[*sdk.Tool], error) { + result, err := c.getSDKSession().ListTools(c.ctx, &sdk.ListToolsParams{Cursor: cursor}) + if err != nil { + return paginatedPage[*sdk.Tool]{}, err + } + return paginatedPage[*sdk.Tool]{Items: result.Tools, NextCursor: result.NextCursor}, nil + }, + func(items []*sdk.Tool) *sdk.ListToolsResult { + return &sdk.ListToolsResult{Tools: items} + }, + ) +} + +func (c *Connection) callTool(params interface{}) (*Response, error) { + return callParamMethod(c, params, func(p CallToolParams) (interface{}, error) { + // Ensure arguments is never nil - default to empty map + // This is required by the MCP protocol which expects arguments to always be present + if p.Arguments == nil { + p.Arguments = make(map[string]interface{}) + } + logConn.Printf("callTool: parsed name=%s, arguments=%+v", p.Name, p.Arguments) + return c.getSDKSession().CallTool(c.ctx, &sdk.CallToolParams{ + Name: p.Name, + Arguments: p.Arguments, + }) + }) +} + +func (c *Connection) listResources() (*Response, error) { + return listMCPItems(c, "resources", + func(cursor string) (paginatedPage[*sdk.Resource], error) { + result, err := c.getSDKSession().ListResources(c.ctx, &sdk.ListResourcesParams{Cursor: cursor}) + if err != nil { + return paginatedPage[*sdk.Resource]{}, err + } + return paginatedPage[*sdk.Resource]{Items: result.Resources, NextCursor: result.NextCursor}, nil + }, + func(items []*sdk.Resource) *sdk.ListResourcesResult { + return &sdk.ListResourcesResult{Resources: items} + }, + ) +} + +func (c *Connection) readResource(params interface{}) (*Response, error) { + type readResourceParams struct { + URI string `json:"uri"` + } + return callParamMethod(c, params, func(p readResourceParams) (interface{}, error) { + logConn.Printf("readResource: reading resource uri=%s from serverID=%s", p.URI, c.serverID) + return c.getSDKSession().ReadResource(c.ctx, &sdk.ReadResourceParams{ + URI: p.URI, + }) + }) +} + +func (c *Connection) listPrompts() (*Response, error) { + return listMCPItems(c, "prompts", + func(cursor string) (paginatedPage[*sdk.Prompt], error) { + result, err := c.getSDKSession().ListPrompts(c.ctx, &sdk.ListPromptsParams{Cursor: cursor}) + if err != nil { + return paginatedPage[*sdk.Prompt]{}, err + } + return paginatedPage[*sdk.Prompt]{Items: result.Prompts, NextCursor: result.NextCursor}, nil + }, + func(items []*sdk.Prompt) *sdk.ListPromptsResult { + return &sdk.ListPromptsResult{Prompts: items} + }, + ) +} + +func (c *Connection) getPrompt(params interface{}) (*Response, error) { + type getPromptParams struct { + Name string `json:"name"` + Arguments map[string]string `json:"arguments"` + } + return callParamMethod(c, params, func(p getPromptParams) (interface{}, error) { + logConn.Printf("getPrompt: getting prompt name=%s from serverID=%s", p.Name, c.serverID) + return c.getSDKSession().GetPrompt(c.ctx, &sdk.GetPromptParams{ + Name: p.Name, + Arguments: p.Arguments, + }) + }) +} diff --git a/internal/mcp/pagination.go b/internal/mcp/pagination.go new file mode 100644 index 00000000..d040d2f5 --- /dev/null +++ b/internal/mcp/pagination.go @@ -0,0 +1,129 @@ +package mcp + +import ( + "encoding/json" + "fmt" +) + +// marshalToResponse marshals an SDK result into a Response object. +// This helper reduces code duplication across all MCP method wrappers. +// +// The ID field is set to a static placeholder (1) because this Response is only +// constructed after the SDK's session.XXX() call has already resolved the +// request–response correlation internally. The gateway never uses this ID for +// matching; it is present solely to satisfy the JSON-RPC 2.0 structure. +func marshalToResponse(result interface{}) (*Response, error) { + resultJSON, err := json.Marshal(result) + if err != nil { + return nil, fmt.Errorf("failed to marshal result: %w", err) + } + + return &Response{ + JSONRPC: "2.0", + ID: 1, // Placeholder – see function comment for safety rationale + Result: resultJSON, + }, nil +} + +// unmarshalParams converts generic interface{} params to a specific struct type. +// This helper reduces code duplication across MCP method wrappers and ensures +// consistent error handling for parameter conversion. It uses marshal/unmarshal +// to maintain JSON schema validation benefits. +func unmarshalParams(params interface{}, target interface{}) error { + paramsJSON, err := json.Marshal(params) + if err != nil { + return fmt.Errorf("failed to marshal params: %w", err) + } + if err := json.Unmarshal(paramsJSON, target); err != nil { + return fmt.Errorf("invalid params: %w", err) + } + return nil +} + +// callParamMethod is a generic helper for SDK operations that require typed parameters. +// It handles the common pattern of: requireSession → unmarshalParams → fn(params) → marshalToResponse. +// P is the type of the parameter struct to unmarshal into. +func callParamMethod[P any](c *Connection, rawParams interface{}, fn func(P) (interface{}, error)) (*Response, error) { + if err := c.requireSession(); err != nil { + return nil, err + } + var params P + if err := unmarshalParams(rawParams, ¶ms); err != nil { + return nil, err + } + result, err := fn(params) + if err != nil { + return nil, err + } + return marshalToResponse(result) +} + +// paginatedPage holds a single page of results from a paginated SDK list call. +type paginatedPage[T any] struct { + Items []T + NextCursor string +} + +// paginateAllMaxPages is the maximum number of pages that paginateAll will fetch. +// This guards against misbehaving or adversarial backends that return an unbounded +// sequence of pages, which would otherwise consume unbounded memory and time. +const paginateAllMaxPages = 100 + +// paginateAll collects all items across paginated SDK list calls. +// It returns an error if the backend returns more than paginateAllMaxPages pages, +// protecting against runaway backends. +func paginateAll[T any]( + serverID string, + itemKind string, + fetch func(cursor string) (paginatedPage[T], error), +) ([]T, error) { + first, err := fetch("") + if err != nil { + return nil, err + } + all := make([]T, len(first.Items), max(len(first.Items), 1)) + copy(all, first.Items) + logConn.Printf("list%s: received page of %d %s from serverID=%s", itemKind, len(first.Items), itemKind, serverID) + + cursor := first.NextCursor + seenCursors := make(map[string]struct{}) + for pageCount := 1; cursor != ""; pageCount++ { + if pageCount >= paginateAllMaxPages { + return nil, fmt.Errorf("list%s: backend serverID=%s returned more than %d pages; aborting to prevent unbounded memory growth", itemKind, serverID, paginateAllMaxPages) + } + if _, seen := seenCursors[cursor]; seen { + return nil, fmt.Errorf("list%s: backend serverID=%s returned cyclical cursor %q", itemKind, serverID, cursor) + } + seenCursors[cursor] = struct{}{} + page, err := fetch(cursor) + if err != nil { + return nil, err + } + all = append(all, page.Items...) + logConn.Printf("list%s: received page of %d %s (total so far: %d) from serverID=%s", itemKind, len(page.Items), itemKind, len(all), serverID) + cursor = page.NextCursor + } + logConn.Printf("list%s: received %d %s total from serverID=%s", itemKind, len(all), itemKind, serverID) + return all, nil +} + +// listMCPItems is a generic helper for the list* family of MCP operations. +// It handles session validation, logging, pagination, and response marshalling, +// eliminating the boilerplate that was previously duplicated across listTools, +// listResources, and listPrompts. +func listMCPItems[Item any, Result any]( + c *Connection, + kind string, + fetchPage func(cursor string) (paginatedPage[Item], error), + buildResult func([]Item) Result, +) (*Response, error) { + if err := c.requireSession(); err != nil { + return nil, err + } + logConn.Printf("list%s: requesting %s list from backend serverID=%s", kind, kind, c.serverID) + items, err := paginateAll(c.serverID, kind, fetchPage) + if err != nil { + return nil, err + } + return marshalToResponse(buildResult(items)) +} diff --git a/internal/proxy/handler.go b/internal/proxy/handler.go index d31778d0..eb0478bb 100644 --- a/internal/proxy/handler.go +++ b/internal/proxy/handler.go @@ -474,145 +474,3 @@ func computeRetryAfter(resetAt time.Time) int { } return secs } - -// rewrapSearchResponse re-wraps filtered items into the original search response -// envelope. GitHub search endpoints return {"total_count": N, "items": [...]}; -// ToResult() returns a bare array, so we rebuild the wrapper. -func rewrapSearchResponse(originalData interface{}, filteredItems interface{}) interface{} { - original, ok := originalData.(map[string]interface{}) - if !ok { - return filteredItems - } - // Detect search response wrapper (has total_count + items/repositories) - if _, hasTotalCount := original["total_count"]; !hasTotalCount { - return filteredItems - } - items, ok := filteredItems.([]interface{}) - if !ok { - return filteredItems - } - // Rebuild the search wrapper with filtered items - result := make(map[string]interface{}) - for k, v := range original { - result[k] = v - } - // Replace items key — search can use "items", "repositories", etc. - for _, key := range []string{"items", "repositories"} { - if _, ok := original[key]; ok { - result[key] = items - break - } - } - result["total_count"] = float64(len(items)) - result["incomplete_results"] = false - return result -} - -// unwrapSingleObject preserves the original response shape for single-object endpoints. -// When the guard wraps a single object in a collection, ToResult() returns [obj]. -// This unwraps it back to obj when the original response was a single object -// (e.g., get_file_contents, get_commit, issue_read). -func unwrapSingleObject(originalData interface{}, filteredData interface{}) interface{} { - original, isMap := originalData.(map[string]interface{}) - if !isMap { - return filteredData - } - // Don't unwrap search envelopes (handled by rewrapSearchResponse) - if _, hasTotalCount := original["total_count"]; hasTotalCount { - return filteredData - } - // Don't unwrap GraphQL responses (handled separately) - if _, hasData := original["data"]; hasData { - return filteredData - } - // If filtered result is a single-element array, unwrap to match original shape - if arr, ok := filteredData.([]interface{}); ok && len(arr) == 1 { - return arr[0] - } - return filteredData -} - -// rebuildGraphQLResponse reconstructs a GraphQL response with only accessible -// items, preserving the {"data": {...}} envelope that clients expect. -func rebuildGraphQLResponse(originalData interface{}, filtered *difc.FilteredCollectionLabeledData) interface{} { - original, ok := originalData.(map[string]interface{}) - if !ok { - return map[string]interface{}{"data": nil} - } - if _, ok := original["data"]; !ok { - return map[string]interface{}{"data": nil} - } - - // If all items were filtered out, return {"data": null} to avoid leaking - // the original response through non-collection fields (e.g., viewer). - if filtered.GetAccessibleCount() == 0 { - return map[string]interface{}{"data": nil} - } - - // Deep-clone the original data structure - cloned := deepCloneJSON(original) - - // Build accessible items set - accessibleItems := make([]interface{}, 0, len(filtered.Accessible)) - for _, item := range filtered.Accessible { - accessibleItems = append(accessibleItems, item.Data) - } - - // Walk the cloned structure and replace nodes/edges arrays. - // If no nodes/edges found, return {"data": null} to prevent leaking - // non-collection data (e.g., viewer { login }). - if clonedMap, ok := cloned.(map[string]interface{}); ok { - if clonedData, ok := clonedMap["data"]; ok { - if !replaceNodesArray(clonedData, accessibleItems) { - return map[string]interface{}{"data": nil} - } - } - } - - return cloned -} - -// replaceNodesArray walks a JSON tree and replaces the first "nodes" or "edges" -// array with the given items, and updates any adjacent "totalCount". -func replaceNodesArray(v interface{}, items []interface{}) bool { - obj, ok := v.(map[string]interface{}) - if !ok { - return false - } - for _, key := range []string{"nodes", "edges"} { - if _, ok := obj[key]; ok { - obj[key] = items - if _, ok := obj["totalCount"]; ok { - obj["totalCount"] = float64(len(items)) - } - return true - } - } - // Recurse into child objects - for _, child := range obj { - if replaceNodesArray(child, items) { - return true - } - } - return false -} - -// deepCloneJSON creates a deep copy of a JSON-compatible value. -func deepCloneJSON(v interface{}) interface{} { - switch val := v.(type) { - case map[string]interface{}: - clone := make(map[string]interface{}, len(val)) - for k, v := range val { - clone[k] = deepCloneJSON(v) - } - return clone - case []interface{}: - clone := make([]interface{}, len(val)) - for i, v := range val { - clone[i] = deepCloneJSON(v) - } - return clone - default: - return v - } -} diff --git a/internal/proxy/response_transform.go b/internal/proxy/response_transform.go new file mode 100644 index 00000000..50414236 --- /dev/null +++ b/internal/proxy/response_transform.go @@ -0,0 +1,145 @@ +package proxy + +import "github.com/github/gh-aw-mcpg/internal/difc" + +// rewrapSearchResponse re-wraps filtered items into the original search response +// envelope. GitHub search endpoints return {"total_count": N, "items": [...]}; +// ToResult() returns a bare array, so we rebuild the wrapper. +func rewrapSearchResponse(originalData interface{}, filteredItems interface{}) interface{} { + original, ok := originalData.(map[string]interface{}) + if !ok { + return filteredItems + } + // Detect search response wrapper (has total_count + items/repositories) + if _, hasTotalCount := original["total_count"]; !hasTotalCount { + return filteredItems + } + items, ok := filteredItems.([]interface{}) + if !ok { + return filteredItems + } + // Rebuild the search wrapper with filtered items + result := make(map[string]interface{}) + for k, v := range original { + result[k] = v + } + // Replace items key — search can use "items", "repositories", etc. + for _, key := range []string{"items", "repositories"} { + if _, ok := original[key]; ok { + result[key] = items + break + } + } + result["total_count"] = float64(len(items)) + result["incomplete_results"] = false + return result +} + +// unwrapSingleObject preserves the original response shape for single-object endpoints. +// When the guard wraps a single object in a collection, ToResult() returns [obj]. +// This unwraps it back to obj when the original response was a single object +// (e.g., get_file_contents, get_commit, issue_read). +func unwrapSingleObject(originalData interface{}, filteredData interface{}) interface{} { + original, isMap := originalData.(map[string]interface{}) + if !isMap { + return filteredData + } + // Don't unwrap search envelopes (handled by rewrapSearchResponse) + if _, hasTotalCount := original["total_count"]; hasTotalCount { + return filteredData + } + // Don't unwrap GraphQL responses (handled separately) + if _, hasData := original["data"]; hasData { + return filteredData + } + // If filtered result is a single-element array, unwrap to match original shape + if arr, ok := filteredData.([]interface{}); ok && len(arr) == 1 { + return arr[0] + } + return filteredData +} + +// rebuildGraphQLResponse reconstructs a GraphQL response with only accessible +// items, preserving the {"data": {...}} envelope that clients expect. +func rebuildGraphQLResponse(originalData interface{}, filtered *difc.FilteredCollectionLabeledData) interface{} { + original, ok := originalData.(map[string]interface{}) + if !ok { + return map[string]interface{}{"data": nil} + } + if _, ok := original["data"]; !ok { + return map[string]interface{}{"data": nil} + } + + // If all items were filtered out, return {"data": null} to avoid leaking + // the original response through non-collection fields (e.g., viewer). + if filtered.GetAccessibleCount() == 0 { + return map[string]interface{}{"data": nil} + } + + // Deep-clone the original data structure + cloned := deepCloneJSON(original) + + // Build accessible items set + accessibleItems := make([]interface{}, 0, len(filtered.Accessible)) + for _, item := range filtered.Accessible { + accessibleItems = append(accessibleItems, item.Data) + } + + // Walk the cloned structure and replace nodes/edges arrays. + // If no nodes/edges found, return {"data": null} to prevent leaking + // non-collection data (e.g., viewer { login }). + if clonedMap, ok := cloned.(map[string]interface{}); ok { + if clonedData, ok := clonedMap["data"]; ok { + if !replaceNodesArray(clonedData, accessibleItems) { + return map[string]interface{}{"data": nil} + } + } + } + + return cloned +} + +// replaceNodesArray walks a JSON tree and replaces the first "nodes" or "edges" +// array with the given items, and updates any adjacent "totalCount". +func replaceNodesArray(v interface{}, items []interface{}) bool { + obj, ok := v.(map[string]interface{}) + if !ok { + return false + } + for _, key := range []string{"nodes", "edges"} { + if _, ok := obj[key]; ok { + obj[key] = items + if _, ok := obj["totalCount"]; ok { + obj["totalCount"] = float64(len(items)) + } + return true + } + } + // Recurse into child objects + for _, child := range obj { + if replaceNodesArray(child, items) { + return true + } + } + return false +} + +// deepCloneJSON creates a deep copy of a JSON-compatible value. +func deepCloneJSON(v interface{}) interface{} { + switch val := v.(type) { + case map[string]interface{}: + clone := make(map[string]interface{}, len(val)) + for k, v := range val { + clone[k] = deepCloneJSON(v) + } + return clone + case []interface{}: + clone := make([]interface{}, len(val)) + for i, v := range val { + clone[i] = deepCloneJSON(v) + } + return clone + default: + return v + } +} From 28ed63f130f9fc0abd985e78893285985da0f058 Mon Sep 17 00:00:00 2001 From: Landon Cox Date: Mon, 27 Apr 2026 20:21:21 -0700 Subject: [PATCH 3/3] Update internal/mcp/connection_methods.go Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- internal/mcp/connection_methods.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/mcp/connection_methods.go b/internal/mcp/connection_methods.go index c1fe18ac..82b5b66d 100644 --- a/internal/mcp/connection_methods.go +++ b/internal/mcp/connection_methods.go @@ -51,7 +51,7 @@ func (c *Connection) callTool(params interface{}) (*Response, error) { if p.Arguments == nil { p.Arguments = make(map[string]interface{}) } - logConn.Printf("callTool: parsed name=%s, arguments=%+v", p.Name, p.Arguments) + logConn.Printf("callTool: parsed name=%s, argumentCount=%d", p.Name, len(p.Arguments)) return c.getSDKSession().CallTool(c.ctx, &sdk.CallToolParams{ Name: p.Name, Arguments: p.Arguments,