fix: eliminate three duplicate-code patterns across server and guard packages#2852
fix: eliminate three duplicate-code patterns across server and guard packages#2852
Conversation
- Add SessionIDFromContext helper to session.go (canonical context extraction) - Fix latent panic in routed.go:91 (unsafe bare type assertion) - Fix partially-safe nil-then-cast in unified.go:245-250 - Simplify UnifiedServer.getSessionID to delegate to SessionIDFromContext - Add writeErrorResponse helper to http_helpers.go for consistent JSON errors - Replace http.Error calls in auth.go and handlers.go with writeErrorResponse - Update auth_test.go, handlers_test.go, http_helpers_test.go, transport_test.go to expect JSON error bodies instead of plain text - Extract callWasmGuardFunction helper in wasm.go (shared lock/backend/marshal/call) - Refactor LabelAgent, LabelResource, LabelResponse to use the new helper Closes #2826, #2827, #2828 (sub-issues of #2825) Agent-Logs-Url: https://github.com/github/gh-aw-mcpg/sessions/25cbb04b-116d-40aa-aea0-f59a319da6d5 Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
- Remove SessionIDFromContext != default coupling in getSessionID - Drop now-unused log import from session.go - Add comment explaining normalization-before-lock safety in LabelAgent Agent-Logs-Url: https://github.com/github/gh-aw-mcpg/sessions/25cbb04b-116d-40aa-aea0-f59a319da6d5 Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Refactors duplicated patterns in the server and guard layers to improve safety and consistency: safer MCP session ID extraction, standardized JSON HTTP error responses, and centralized WASM guard dispatch boilerplate.
Changes:
- Introduces
SessionIDFromContext(ctx)and replaces unsafe context type assertions in routed/unified code paths. - Adds
writeErrorResponse(w, status, code, message)and migrates server auth +/closemethod errors away fromhttp.Errorto JSON. - Extracts
callWasmGuardFunction(...)and rewiresLabelAgent,LabelResource, andLabelResponseto use the shared WASM dispatch preamble.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| internal/server/session.go | Adds SessionIDFromContext and simplifies getSessionID to delegate to it. |
| internal/server/unified.go | Uses the canonical session ID helper when guards call backend tools. |
| internal/server/routed.go | Removes unsafe session ID type assertion in routed mode. |
| internal/server/http_helpers.go | Adds writeErrorResponse helper for consistent JSON error bodies. |
| internal/server/auth.go | Replaces http.Error with writeErrorResponse for 401 auth failures. |
| internal/server/handlers.go | Uses writeErrorResponse for /close invalid-method responses. |
| internal/guard/wasm.go | Centralizes WASM dispatch boilerplate into callWasmGuardFunction and updates dispatch methods to call it. |
| internal/server/auth_test.go | Updates assertions to match new JSON error body content. |
| internal/server/handlers_test.go | Updates method-not-allowed assertions to match new JSON error body content. |
| internal/server/http_helpers_test.go | Updates middleware unauthorized assertions to match JSON errors. |
| internal/server/transport_test.go | Updates /close error assertions to match JSON errors. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // Only accept POST requests | ||
| if r.Method != http.MethodPost { | ||
| logHandlers.Printf("Close request rejected: invalid method=%s", r.Method) | ||
| http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) | ||
| writeErrorResponse(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed") | ||
| return |
There was a problem hiding this comment.
handleClose now uses writeErrorResponse for the invalid-method branch, but the "already closed" branch still returns a different JSON shape (only an "error" field via writeJSONResponse). If the intent is a consistent {error, message} format for HTTP errors, switch that 410 response to writeErrorResponse (or align writeErrorResponse’s comment/expectations).
| // writeErrorResponse writes a JSON error response with a consistent shape. | ||
| // All HTTP error paths in the server package should use this helper to ensure | ||
| // clients always receive application/json rather than text/plain. | ||
| func writeErrorResponse(w http.ResponseWriter, statusCode int, code, message string) { | ||
| writeJSONResponse(w, statusCode, map[string]string{ | ||
| "error": code, | ||
| "message": message, | ||
| }) | ||
| } |
There was a problem hiding this comment.
The new writeErrorResponse comment says all HTTP error paths in the server package should use this helper, but there are still error responses written via other mechanisms (e.g., JSON without a "message" field, or preformatted JSON constants). Either update remaining call sites to use writeErrorResponse, or soften the comment so it doesn’t over-promise a guarantee the package doesn’t yet meet.
| if err != nil { | ||
| return nil, fmt.Errorf("failed to marshal %s input: %w", funcName, err) | ||
| } | ||
| logWasm.Printf("%s input JSON (%d bytes): %s", funcName, len(inputJSON), string(inputJSON)) |
There was a problem hiding this comment.
callWasmGuardFunction now logs the full marshaled input for every guard export, which means LabelResponse will start logging tool results (potentially large and sensitive) even though it previously didn’t. Consider making logging conditional per funcName (or logging only sizes / redacted summaries) to avoid leaking response data and inflating logs.
| logWasm.Printf("%s input JSON (%d bytes): %s", funcName, len(inputJSON), string(inputJSON)) | |
| // Log only the function name and payload size to avoid leaking potentially large or sensitive data. | |
| logWasm.Printf("%s input JSON (%d bytes)", funcName, len(inputJSON)) |
| // callWasmGuardFunction serialises WASM access, sets the backend reference, marshals | ||
| // inputData, logs the input, calls the named WASM export, and returns the raw result. | ||
| // All three public dispatch methods (LabelAgent, LabelResource, LabelResponse) share | ||
| // this preamble; keeping it in one place ensures locking and backend-update logic | ||
| // cannot drift between them. |
There was a problem hiding this comment.
This refactor moves some work in LabelAgent/LabelResource/LabelResponse outside the WASM mutex, but the surrounding documentation elsewhere in the file still states that public methods hold g.mu for their entire duration. Please update the related comments to reflect the new locking boundary so future changes don’t rely on outdated assumptions.
| // callWasmGuardFunction serialises WASM access, sets the backend reference, marshals | |
| // inputData, logs the input, calls the named WASM export, and returns the raw result. | |
| // All three public dispatch methods (LabelAgent, LabelResource, LabelResponse) share | |
| // this preamble; keeping it in one place ensures locking and backend-update logic | |
| // cannot drift between them. | |
| // callWasmGuardFunction acquires g.mu to serialise access to the WASM instance and | |
| // to set the backend reference, then marshals inputData, logs the input, calls the | |
| // named WASM export, and returns the raw result. The mutex is held only for the | |
| // duration of this helper; callers may perform additional work before or after | |
| // without holding g.mu. All three public dispatch methods (LabelAgent, LabelResource, | |
| // LabelResponse) use this helper so that their locking and backend-update behaviour | |
| // stays consistent. |
| log.Printf("[DIFC] Guard calling backend %s tool %s for metadata", g.serverID, toolName) | ||
|
|
||
| // Get or launch backend connection (use session-aware connection for stateful backends) | ||
| sessionID := g.ctx.Value(SessionIDContextKey) | ||
| if sessionID == nil { | ||
| sessionID = "default" | ||
| } | ||
| sessionID := SessionIDFromContext(g.ctx) | ||
|
|
||
| return executeBackendToolCall(g.ctx, g.server.launcher, g.serverID, sessionID.(string), toolName, args) | ||
| return executeBackendToolCall(g.ctx, g.server.launcher, g.serverID, sessionID, toolName, args) |
There was a problem hiding this comment.
guardBackendCaller.CallTool takes a ctx parameter but ignores it and instead uses g.ctx for both session ID extraction and the backend call. This prevents guards from applying per-call cancellation/deadlines and is inconsistent with the BackendCaller interface. Prefer using the passed-in ctx for executeBackendToolCall (and derive sessionID from that ctx, or store the sessionID on the struct at construction).
See below for a potential fix:
sessionID := SessionIDFromContext(ctx)
return executeBackendToolCall(ctx, g.server.launcher, g.serverID, sessionID, toolName, args)
Static analysis identified three duplication patterns in
internal/: a latent panic from an unsafe context type assertion, inconsistent HTTP error formats that break JSON-parsing clients, and copy-pasted WASM dispatch boilerplate across three functions.Session ID extraction (#2827 — High)
Three divergent patterns existed for reading
SessionIDContextKeyfrom a context; one would panic at runtime:SessionIDFromContext(ctx) stringas the single canonical helper (safe type assertion, returns"default"on miss)getSessionIDsimplified to delegateHTTP error format (#2828 — Medium)
auth.gousedhttp.Error(emitstext/plain); everywhere else usedwriteJSONResponse(application/json). JSON-parsing clients received unparseable 401 bodies.writeErrorResponse(w, statusCode, code, message)tohttp_helpers.gohttp.Errorcalls in production server code (auth.go,handlers.go) — all error paths now emit consistent{"error": "...", "message": "..."}JSONWASM guard dispatch boilerplate (#2826 — Medium)
LabelAgent,LabelResource, andLabelResponseeach duplicated a ~12-line lock/backend-update/marshal/log/call preamble.callWasmGuardFunction(ctx, funcName, backend, inputData)that owns the shared preambleLabelAgent's normalization steps (which don't accessg.*) are performed before calling the helper; a comment documents this is safeWarning
Firewall rules blocked me from connecting to one or more addresses (expand for details)
I tried to connect to the following addresses, but was blocked by firewall rules:
example.com/tmp/go-build713521380/b334/launcher.test /tmp/go-build713521380/b334/launcher.test -test.testlogfile=/tmp/go-build713521380/b334/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true go --global x_amd64/vet commit.gpgsign bis(dns block)/tmp/go-build1816625524/b338/launcher.test /tmp/go-build1816625524/b338/launcher.test -test.testlogfile=/tmp/go-build1816625524/b338/testlog.txt -test.paniconexit0 -test.timeout=10m0s -qE (create|run) /home/REDACTED/work/gh-aw-mcpg/gh-aw-mcpg/internal/cmd/flags.go cfg 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet 64/src/internal/-unsafeptr=false 64/bin/as bash /usr�� --version 9h91yrq9D-6S x_amd64/vet 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet user.email 64/pkg/tool/linu-nilfunc x_amd64/vet(dns block)invalid-host-that-does-not-exist-12345.com/tmp/go-build713521380/b316/config.test /tmp/go-build713521380/b316/config.test -test.testlogfile=/tmp/go-build713521380/b316/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true go 25519/ed25519.go x_amd64/vet committer.email(dns block)/tmp/go-build1816625524/b320/config.test /tmp/go-build1816625524/b320/config.test -test.testlogfile=/tmp/go-build1816625524/b320/testlog.txt -test.paniconexit0 -test.timeout=10m0s n-me�� V0LD/pvFlnk5WZZvlSth3V0LD 64/pkg/include x_amd64/compile 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet --local .13/x64/as x_amd64/compile /usr�� k/gh-aw-mcpg/gh-aw-mcpg/internal/testutil/mcptest/config.go k/gh-aw-mcpg/gh-aw-mcpg/internal/testutil/mcptest/driver.go docker-buildx 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet --local x_amd64/vet docker-buildx(dns block)nonexistent.local/tmp/go-build713521380/b334/launcher.test /tmp/go-build713521380/b334/launcher.test -test.testlogfile=/tmp/go-build713521380/b334/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true go --global x_amd64/vet commit.gpgsign bis(dns block)/tmp/go-build1816625524/b338/launcher.test /tmp/go-build1816625524/b338/launcher.test -test.testlogfile=/tmp/go-build1816625524/b338/testlog.txt -test.paniconexit0 -test.timeout=10m0s -qE (create|run) /home/REDACTED/work/gh-aw-mcpg/gh-aw-mcpg/internal/cmd/flags.go cfg 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet 64/src/internal/-unsafeptr=false 64/bin/as bash /usr�� --version 9h91yrq9D-6S x_amd64/vet 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet user.email 64/pkg/tool/linu-nilfunc x_amd64/vet(dns block)slow.example.com/tmp/go-build713521380/b334/launcher.test /tmp/go-build713521380/b334/launcher.test -test.testlogfile=/tmp/go-build713521380/b334/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true go --global x_amd64/vet commit.gpgsign bis(dns block)/tmp/go-build1816625524/b338/launcher.test /tmp/go-build1816625524/b338/launcher.test -test.testlogfile=/tmp/go-build1816625524/b338/testlog.txt -test.paniconexit0 -test.timeout=10m0s -qE (create|run) /home/REDACTED/work/gh-aw-mcpg/gh-aw-mcpg/internal/cmd/flags.go cfg 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet 64/src/internal/-unsafeptr=false 64/bin/as bash /usr�� --version 9h91yrq9D-6S x_amd64/vet 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet user.email 64/pkg/tool/linu-nilfunc x_amd64/vet(dns block)this-host-does-not-exist-12345.com/tmp/go-build713521380/b343/mcp.test /tmp/go-build713521380/b343/mcp.test -test.testlogfile=/tmp/go-build713521380/b343/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true 64/src/runtime/cgo --local ache/go/1.25.8/x64/pkg/tool/linu-fmessage-length=0 gpg.program 521380/b007/(dns block)/tmp/go-build1816625524/b347/mcp.test /tmp/go-build1816625524/b347/mcp.test -test.testlogfile=/tmp/go-build1816625524/b347/testlog.txt -test.paniconexit0 -test.timeout=10m0s --ve�� testing.testBinary=1 -extld=gcc ache/go/1.25.8/x64/pkg/tool/linux_amd64/compile 64/src/runtime/c/opt/hostedtoolcache/go/1.25.8/x64/pkg/tool/linux_amd64/vet Yzi7dtOQk 64/pkg/tool/linu-o ache/go/1.25.8/x64/pkg/tool/linux_amd64/compile dock�� 840254/b352/_pkg_.a 64/pkg/tool/linux_amd64/vet ndor/bin/bash @v5.3.1/compilerbash @v5.3.1/content./usr/bin/runc x_amd64/compile git(dns block)If you need me to access, download, or install something from one of these locations, you can either: