Skip to content

refactor: split oversized files by responsibility (mcp/connection, proxy/handler, httputil) #4701

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lpcox merged 3 commits into from
Apr 28, 2026
Merged

refactor: split oversized files by responsibility (mcp/connection, proxy/handler, httputil) #4701

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e514972
Initial plan
Copilot Apr 28, 2026
2ddc6ee
refactor: split connection.go, extract response transforms, split htt…
Copilot Apr 28, 2026
28ed63f
Update internal/mcp/connection_methods.go
lpcox Apr 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions internal/httputil/github_http.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
35 changes: 0 additions & 35 deletions internal/httputil/httputil.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,6 @@ package httputil
import (
"encoding/json"
"net/http"
"strconv"
"strings"
"time"

"github.com/github/gh-aw-mcpg/internal/logger"
)
Expand Down Expand Up @@ -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 {
Expand Down
279 changes: 0 additions & 279 deletions internal/mcp/connection.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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).
Expand All @@ -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, &params); 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)
Expand Down
Loading
Loading