Skip to content

Commit d4ac02b

Browse files
lpcoxlpcox
authored
refactor: relocate misplaced functions to their natural homes (#5006)
Four findings from semantic function clustering analysis: session lifecycle helpers in the wrong file, feature-specific constructors buried in `unified.go`, a zero-logic passthrough wrapper, and two tiny tracing files that should be one. ## Changes - **Finding 1 — Session helpers → `session.go`** Moved `extractAndValidateSession`, `injectSessionContext`, `setupSessionCallback`, `logHTTPRequestBody` out of `http_helpers.go` into `session.go`. `http_helpers.go` shrinks ~100 lines; all session lifecycle logic is now in one file. - **Finding 2 — Feature constructors → their owning files** - `buildCircuitBreakers`: `unified.go` → `circuit_breaker.go` - `buildAllowedToolSets` + `hasWildcard`: `unified.go` → `tool_registry.go` Added `config` import to both destination files. - **Finding 3 — Remove `writeErrorResponse` wrapper** Deleted the one-line passthrough and replaced all call sites with `httputil.WriteErrorResponse` directly: ```go // before func writeErrorResponse(w http.ResponseWriter, statusCode int, code, message string) { httputil.WriteErrorResponse(w, statusCode, code, message) } // after — call sites use httputil directly httputil.WriteErrorResponse(w, status, code, msg) ``` - **Finding 4 — Merge tracing CLI files** `cmd/tracing_helpers.go` + `cmd/flags_tracing.go` (75 lines combined) → `cmd/tracing.go`. > [!WARNING] > > <details> > <summary>Firewall rules blocked me from connecting to one or more addresses (expand for details)</summary> > > #### I tried to connect to the following addresses, but was blocked by firewall rules: > > - `example.com` > - Triggering command: `/tmp/go-build3872797525/b513/launcher.test /tmp/go-build3872797525/b513/launcher.test -test.testlogfile=/tmp/go-build3872797525/b513/testlog.txt -test.paniconexit0 -test.timeout=10m0s 8043�� olang.org/grpc@v1.80.0/internal/idle/idle.go cfg x_amd64/vet . --gdwarf2 --64 x_amd64/vet -I 8043057/b418/_pkg_.a -I x_amd64/vet --gdwarf-5 g/grpc/internal//usr/bin/runc -o x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3399420040/b509/launcher.test /tmp/go-build3399420040/b509/launcher.test -test.testlogfile=/tmp/go-build3399420040/b509/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true /tmp/go-build2509960937/b173/vet/run/containerd/io.containerd.runtime.v2.task/moby/db2aa2a60d29ftail 64/pkg/tool/linu-buildtags bash g_.a -trimpath &#34;REQUESTS_CA_B--root /opt/hostedtoolc/var/run/docker/runtime-runc/moby -uns�� y x_amd64/vet e-handler @v1.43.0/propaga/usr/libexec/docker/cli-plugins/docker-compose @v1.43.0/propagadocker-cli-plugin-metadata x_amd64/vet e-handler` (dns block) > - `invalid-host-that-does-not-exist-12345.com` > - Triggering command: `/tmp/go-build3872797525/b495/config.test /tmp/go-build3872797525/b495/config.test -test.testlogfile=/tmp/go-build3872797525/b495/testlog.txt -test.paniconexit0 -test.timeout=10m0s -I g_.a NbAn/6jVna5QQyS18CrolNbAn x_amd64/vet --gdwarf-5 resolver/delegat/tmp/go-build2509960937/b300/vet.cfg -o x_amd64/vet 8043�� qzhGZjM5u cfg 64/pkg/tool/linux_amd64/vet (compatibility issue with Go 1.25.0). Continuing with other checks...&#34;; \ elif command -v golan -imultiarch` (dns block) > - Triggering command: `/tmp/go-build3399420040/b491/config.test /tmp/go-build3399420040/b491/config.test -test.testlogfile=/tmp/go-build3399420040/b491/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true 88635552ec0a846ca0962f57bea95f938934c2423396c1fa129 x_amd64/compile ateway_with_pipe.sh cdc8770c07c851b4aae2b041e8136bb0e9e/log.json pc_health_v1/hea--prefix=/net/ipv4/conf/vethe00a9b0 x_amd64/compile ateway_with_pipe--prefix=/net/ipv6/conf/vethe00a9b0 -uns�� 3c2c91bb777a34fa x_amd64/compile by/db2aa2a60d29f83ac692768d36051d3e54950f8b5453b528136881dc8dce1cff/log.json g_.a -trimpath x_amd64/vet dCQFD9yY_0oeD/ar--prefix=/net/ipv6/conf/vethdaded9f` (dns block) > - `nonexistent.local` > - Triggering command: `/tmp/go-build3872797525/b513/launcher.test /tmp/go-build3872797525/b513/launcher.test -test.testlogfile=/tmp/go-build3872797525/b513/testlog.txt -test.paniconexit0 -test.timeout=10m0s 8043�� olang.org/grpc@v1.80.0/internal/idle/idle.go cfg x_amd64/vet . --gdwarf2 --64 x_amd64/vet -I 8043057/b418/_pkg_.a -I x_amd64/vet --gdwarf-5 g/grpc/internal//usr/bin/runc -o x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3399420040/b509/launcher.test /tmp/go-build3399420040/b509/launcher.test -test.testlogfile=/tmp/go-build3399420040/b509/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true /tmp/go-build2509960937/b173/vet/run/containerd/io.containerd.runtime.v2.task/moby/db2aa2a60d29ftail 64/pkg/tool/linu-buildtags bash g_.a -trimpath &#34;REQUESTS_CA_B--root /opt/hostedtoolc/var/run/docker/runtime-runc/moby -uns�� y x_amd64/vet e-handler @v1.43.0/propaga/usr/libexec/docker/cli-plugins/docker-compose @v1.43.0/propagadocker-cli-plugin-metadata x_amd64/vet e-handler` (dns block) > - `slow.example.com` > - Triggering command: `/tmp/go-build3872797525/b513/launcher.test /tmp/go-build3872797525/b513/launcher.test -test.testlogfile=/tmp/go-build3872797525/b513/testlog.txt -test.paniconexit0 -test.timeout=10m0s 8043�� olang.org/grpc@v1.80.0/internal/idle/idle.go cfg x_amd64/vet . --gdwarf2 --64 x_amd64/vet -I 8043057/b418/_pkg_.a -I x_amd64/vet --gdwarf-5 g/grpc/internal//usr/bin/runc -o x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3399420040/b509/launcher.test /tmp/go-build3399420040/b509/launcher.test -test.testlogfile=/tmp/go-build3399420040/b509/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true /tmp/go-build2509960937/b173/vet/run/containerd/io.containerd.runtime.v2.task/moby/db2aa2a60d29ftail 64/pkg/tool/linu-buildtags bash g_.a -trimpath &#34;REQUESTS_CA_B--root /opt/hostedtoolc/var/run/docker/runtime-runc/moby -uns�� y x_amd64/vet e-handler @v1.43.0/propaga/usr/libexec/docker/cli-plugins/docker-compose @v1.43.0/propagadocker-cli-plugin-metadata x_amd64/vet e-handler` (dns block) > - `this-host-does-not-exist-12345.com` > - Triggering command: `/tmp/go-build3872797525/b522/mcp.test /tmp/go-build3872797525/b522/mcp.test -test.testlogfile=/tmp/go-build3872797525/b522/testlog.txt -test.paniconexit0 -test.timeout=10m0s -I 8043057/b414/_pkg_.a -fPIC x_amd64/vet -pthread g/grpc/internal//usr/bin/runc -fmessage-length--version x_amd64/vet` (dns block) > - Triggering command: `/tmp/go-build3399420040/b518/mcp.test /tmp/go-build3399420040/b518/mcp.test -test.testlogfile=/tmp/go-build3399420040/b518/testlog.txt -test.paniconexit0 -test.timeout=10m0s -test.v=true -I by/eca09d10600f3--log-format ker/docker-init by/eca09d10600f3/usr/libexec/docker/cli-plugins/docker-compose b1a7041eaee94398docker-cli-plugin-metadata x_amd64/vet 91e/log.json /tmp�� tes.crt x_amd64/vet /usr/sbin/bash by/b7068f2c73a9cbash /interface.go x_amd64/vet bash` (dns block) > > If you need me to access, download, or install something from one of these locations, you can either: > > - Configure [Actions setup steps](https://gh.io/copilot/actions-setup-steps) to set up my environment, which run before the firewall is enabled > - Add the appropriate URLs or hosts to the custom allowlist in this repository's [Copilot coding agent settings](https://github.com/github/gh-aw-mcpg/settings/copilot/coding_agent) (admins only) > > </details>
2 parents 19ddcfc + ad12b78 commit d4ac02b

9 files changed

Lines changed: 178 additions & 180 deletions

File tree

‎internal/cmd/flags_tracing.go‎

Lines changed: 0 additions & 23 deletions
This file was deleted.

internal/cmd/tracing_helpers.go renamed to internal/cmd/tracing.go

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,35 @@
11
package cmd
22

3+
// Tracing-related flags and helpers for OpenTelemetry OTLP trace export.
4+
35
import (
46
"context"
57
"time"
68

9+
"github.com/spf13/cobra"
10+
"github.com/spf13/pflag"
11+
712
"github.com/github/gh-aw-mcpg/internal/config"
813
"github.com/github/gh-aw-mcpg/internal/envutil"
914
"github.com/github/gh-aw-mcpg/internal/tracing"
10-
"github.com/spf13/pflag"
1115
)
1216

17+
// Tracing flag variables
18+
var (
19+
otlpEndpoint string
20+
otlpServiceName string
21+
otlpSampleRate float64
22+
)
23+
24+
func init() {
25+
RegisterFlag(func(cmd *cobra.Command) {
26+
registerTracingFlags(cmd.Flags(), &otlpEndpoint, &otlpServiceName, &otlpSampleRate,
27+
"OTLP HTTP endpoint for trace export (e.g. http://localhost:4318). Defaults from OTEL_EXPORTER_OTLP_ENDPOINT when set. Tracing is disabled when empty.",
28+
"Service name reported in traces. Defaults from OTEL_SERVICE_NAME when set.",
29+
"Fraction of traces to sample and export (0.0–1.0). Default 1.0 samples everything.")
30+
})
31+
}
32+
1333
func registerTracingFlags(flags *pflag.FlagSet, endpoint *string, serviceName *string, sampleRate *float64, endpointUsage string, serviceUsage string, sampleUsage string) {
1434
flags.StringVar(endpoint, "otlp-endpoint", envutil.GetEnvString("OTEL_EXPORTER_OTLP_ENDPOINT", ""),
1535
endpointUsage)

‎internal/server/circuit_breaker.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"sync"
88
"time"
99

10+
"github.com/github/gh-aw-mcpg/internal/config"
1011
"github.com/github/gh-aw-mcpg/internal/logger"
1112
"github.com/github/gh-aw-mcpg/internal/strutil"
1213
)
@@ -233,6 +234,22 @@ func formatResetAt(t time.Time) string {
233234
return fmt.Sprintf("%s (in %s)", t.UTC().Format(time.RFC3339), strutil.FormatDuration(time.Until(t).Round(time.Second)))
234235
}
235236

237+
// buildCircuitBreakers creates per-backend circuit breakers from the configuration.
238+
func buildCircuitBreakers(cfg *config.Config) map[string]*circuitBreaker {
239+
cbs := make(map[string]*circuitBreaker)
240+
if cfg == nil {
241+
return cbs
242+
}
243+
for serverID, serverCfg := range cfg.Servers {
244+
threshold := serverCfg.RateLimitThreshold
245+
cooldown := time.Duration(serverCfg.RateLimitCooldown) * time.Second
246+
cbs[serverID] = newCircuitBreaker(serverID, threshold, cooldown)
247+
logCircuitBreaker.Printf("Created circuit breaker for server %s: threshold=%d, cooldown=%s",
248+
serverID, threshold, cooldown)
249+
}
250+
return cbs
251+
}
252+
236253
// extractRateLimitErrorText extracts the text content from a raw tool result
237254
// that has been identified as a rate-limit error. Returns the original backend
238255
// message so agents see the actual upstream error rather than a synthetic one.

‎internal/server/handlers.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ func handleClose(unifiedServer *UnifiedServer) http.Handler {
4242
// Only accept POST requests
4343
if r.Method != http.MethodPost {
4444
logHandlers.Printf("Close request rejected: invalid method=%s", r.Method)
45-
writeErrorResponse(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
45+
httputil.WriteErrorResponse(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
4646
return
4747
}
4848

‎internal/server/http_helpers.go‎

Lines changed: 1 addition & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package server
22

33
import (
44
"bytes"
5-
"context"
65
"encoding/json"
76
"io"
87
"log"
@@ -14,7 +13,6 @@ import (
1413
oteltrace "go.opentelemetry.io/otel/trace"
1514

1615
"github.com/github/gh-aw-mcpg/internal/auth"
17-
"github.com/github/gh-aw-mcpg/internal/guard"
1816
"github.com/github/gh-aw-mcpg/internal/httputil"
1917
"github.com/github/gh-aw-mcpg/internal/logger"
2018
"github.com/github/gh-aw-mcpg/internal/logger/sanitize"
@@ -45,13 +43,6 @@ func logRuntimeError(errorType, detail string, r *http.Request, serverName *stri
4543
timestamp, server, requestID, errorType, detail, r.URL.Path, r.Method)
4644
}
4745

48-
// writeErrorResponse writes a JSON error response with a consistent shape.
49-
// All HTTP error paths in the server package should use this helper to ensure
50-
// clients always receive application/json rather than text/plain.
51-
func writeErrorResponse(w http.ResponseWriter, statusCode int, code, message string) {
52-
httputil.WriteErrorResponse(w, statusCode, code, message)
53-
}
54-
5546
// rejectRequest logs a structured error, records a runtime error, and writes an
5647
// HTTP error response. This consolidates the 3-step rejection pattern that was
5748
// previously duplicated across auth and handler code paths.
@@ -64,7 +55,7 @@ func writeErrorResponse(w http.ResponseWriter, statusCode int, code, message str
6455
func rejectRequest(w http.ResponseWriter, r *http.Request, status int, code, msg, logCategory, runtimeErrType, runtimeDetail string) {
6556
logger.LogErrorMd(logCategory, "Request rejected: %s, remote=%s, path=%s", msg, r.RemoteAddr, r.URL.Path)
6657
logRuntimeError(runtimeErrType, runtimeDetail, r, nil)
67-
writeErrorResponse(w, status, code, msg)
58+
httputil.WriteErrorResponse(w, status, code, msg)
6859
}
6960

7061
// withResponseLogging wraps an http.Handler to log response bodies
@@ -79,24 +70,6 @@ func withResponseLogging(handler http.Handler) http.Handler {
7970
})
8071
}
8172

82-
// extractAndValidateSession extracts the session ID from the Authorization header
83-
// and logs connection details. Returns empty string if validation fails.
84-
func extractAndValidateSession(r *http.Request) string {
85-
logHelpers.Printf("Extracting session from request: remote=%s, path=%s", r.RemoteAddr, r.URL.Path)
86-
87-
authHeader := r.Header.Get("Authorization")
88-
sessionID := auth.ExtractSessionID(authHeader)
89-
90-
if sessionID == "" {
91-
logHelpers.Printf("Session extraction failed: no Authorization header, remote=%s", r.RemoteAddr)
92-
logger.LogError("client", "Rejected MCP client connection: no Authorization header, remote=%s, path=%s", r.RemoteAddr, r.URL.Path)
93-
return ""
94-
}
95-
96-
logHelpers.Printf("Session extracted successfully: sessionID=%s, remote=%s", auth.TruncateSessionID(sessionID), r.RemoteAddr)
97-
return sessionID
98-
}
99-
10073
// peekRequestBody reads all bytes from a POST request body and restores it
10174
// so downstream handlers can read it again.
10275
// Returns nil, nil for non-POST requests or requests with no body.
@@ -124,78 +97,6 @@ func peekRequestBody(r *http.Request) ([]byte, error) {
12497
return b, nil
12598
}
12699

127-
// logHTTPRequestBody logs the request body for debugging purposes.
128-
// It reads the body, logs it, and restores it so it can be read again.
129-
// The backendID parameter is optional and can be empty for unified mode.
130-
func logHTTPRequestBody(r *http.Request, sessionID, backendID string) {
131-
logHelpers.Printf("Checking request body: method=%s, hasBody=%v, sessionID=%s", r.Method, r.Body != nil, sessionID)
132-
133-
bodyBytes, err := peekRequestBody(r)
134-
if err != nil {
135-
logHelpers.Printf("Body read failed: err=%v", err)
136-
return
137-
}
138-
if len(bodyBytes) == 0 {
139-
logHelpers.Printf("Skipping body logging: not a POST request, no body present, or empty body")
140-
return
141-
}
142-
143-
logHelpers.Printf("Request body read: size=%d bytes, sessionID=%s, backendID=%s", len(bodyBytes), auth.TruncateSessionID(sessionID), backendID)
144-
145-
// Sanitize the body before logging
146-
sanitizedBody := sanitize.SanitizeString(string(bodyBytes))
147-
148-
// Log with backend context if provided (routed mode)
149-
if backendID != "" {
150-
logger.LogDebug("client", "MCP client request body, backend=%s, body=%s", backendID, sanitizedBody)
151-
} else {
152-
logger.LogDebug("client", "MCP request body, session=%s, body=%s", auth.TruncateSessionID(sessionID), sanitizedBody)
153-
}
154-
logHelpers.Print("Request body logged for debugging")
155-
}
156-
157-
// injectSessionContext stores the session ID and optional backend ID into the request context.
158-
// If backendID is empty, only session ID is injected (unified mode).
159-
// Returns the modified request with updated context.
160-
func injectSessionContext(r *http.Request, sessionID, backendID string) *http.Request {
161-
logHelpers.Printf("Injecting session context: sessionID=%s, backendID=%s", auth.TruncateSessionID(sessionID), backendID)
162-
163-
ctx := context.WithValue(r.Context(), SessionIDContextKey, sessionID)
164-
ctx = guard.SetAgentIDInContext(ctx, sessionID)
165-
166-
if backendID != "" {
167-
logHelpers.Printf("Adding backend ID to context: backendID=%s", backendID)
168-
ctx = context.WithValue(ctx, mcp.ContextKey("backend-id"), backendID)
169-
}
170-
171-
logHelpers.Print("Session context injected successfully")
172-
return r.WithContext(ctx)
173-
}
174-
175-
// setupSessionCallback extracts the session ID, logs the new connection, injects
176-
// the session into the request context, and returns the session ID.
177-
// Used by both routed and unified StreamableHTTP session establishment callbacks.
178-
func setupSessionCallback(r *http.Request, backendID string) (string, bool) {
179-
sessionID := extractAndValidateSession(r)
180-
if sessionID == "" {
181-
return "", false
182-
}
183-
184-
if backendID != "" {
185-
logger.LogInfo("client", "New MCP client connection, remote=%s, method=%s, path=%s, backend=%s, session=%s",
186-
r.RemoteAddr, r.Method, r.URL.Path, backendID, auth.TruncateSessionID(sessionID))
187-
} else {
188-
logger.LogInfo("client", "MCP connection established, remote=%s, method=%s, path=%s, session=%s",
189-
r.RemoteAddr, r.Method, r.URL.Path, auth.TruncateSessionID(sessionID))
190-
}
191-
192-
logHTTPRequestBody(r, sessionID, backendID)
193-
194-
*r = *injectSessionContext(r, sessionID, backendID)
195-
196-
return sessionID, true
197-
}
198-
199100
// WithOTELTracing wraps an http.Handler with an OpenTelemetry span for each request.
200101
// The span covers the full HTTP handler lifecycle and includes session ID, HTTP path,
201102
// and method as span attributes. The span context is propagated into the request context

‎internal/server/http_helpers_test.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111

1212
"github.com/github/gh-aw-mcpg/internal/config"
1313
"github.com/github/gh-aw-mcpg/internal/guard"
14+
"github.com/github/gh-aw-mcpg/internal/httputil"
1415
"github.com/github/gh-aw-mcpg/internal/mcp"
1516
"github.com/stretchr/testify/assert"
1617
"github.com/stretchr/testify/require"
@@ -732,7 +733,7 @@ func stringPtr(s string) *string {
732733
return &s
733734
}
734735

735-
// TestWriteErrorResponse verifies that writeErrorResponse writes a JSON error
736+
// TestWriteErrorResponse verifies that httputil.WriteErrorResponse writes a JSON error
736737
// body with "error" and "message" fields and the correct status code and Content-Type.
737738
func TestWriteErrorResponse(t *testing.T) {
738739
tests := []struct {
@@ -782,7 +783,7 @@ func TestWriteErrorResponse(t *testing.T) {
782783
for _, tt := range tests {
783784
t.Run(tt.name, func(t *testing.T) {
784785
w := httptest.NewRecorder()
785-
writeErrorResponse(w, tt.statusCode, tt.code, tt.message)
786+
httputil.WriteErrorResponse(w, tt.statusCode, tt.code, tt.message)
786787

787788
assert.Equal(t, tt.statusCode, w.Code, "status code should match")
788789
assert.Equal(t, "application/json", w.Header().Get("Content-Type"), "Content-Type should be application/json")

‎internal/server/session.go‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,16 @@ package server
33
import (
44
"context"
55
"fmt"
6+
"net/http"
67
"os"
78
"path/filepath"
89
"time"
910

1011
"github.com/github/gh-aw-mcpg/internal/auth"
12+
"github.com/github/gh-aw-mcpg/internal/guard"
1113
"github.com/github/gh-aw-mcpg/internal/logger"
14+
"github.com/github/gh-aw-mcpg/internal/logger/sanitize"
15+
"github.com/github/gh-aw-mcpg/internal/mcp"
1216
"github.com/github/gh-aw-mcpg/internal/syncutil"
1317
)
1418

@@ -110,3 +114,95 @@ func (us *UnifiedServer) getSessionKeys() []string {
110114
logSession.Printf("Active sessions: count=%d", len(keys))
111115
return keys
112116
}
117+
118+
// extractAndValidateSession extracts the session ID from the Authorization header
119+
// and logs connection details. Returns empty string if validation fails.
120+
func extractAndValidateSession(r *http.Request) string {
121+
logSession.Printf("Extracting session from request: remote=%s, path=%s", r.RemoteAddr, r.URL.Path)
122+
123+
authHeader := r.Header.Get("Authorization")
124+
sessionID := auth.ExtractSessionID(authHeader)
125+
126+
if sessionID == "" {
127+
logSession.Printf("Session extraction failed: missing or invalid Authorization header, remote=%s", r.RemoteAddr)
128+
logger.LogError("client", "Rejected MCP client connection: missing or invalid Authorization header, remote=%s, path=%s", r.RemoteAddr, r.URL.Path)
129+
return ""
130+
}
131+
132+
logSession.Printf("Session extracted successfully: sessionID=%s, remote=%s", auth.TruncateSessionID(sessionID), r.RemoteAddr)
133+
return sessionID
134+
}
135+
136+
// injectSessionContext stores the session ID and optional backend ID into the request context.
137+
// If backendID is empty, only session ID is injected (unified mode).
138+
// Returns the modified request with updated context.
139+
func injectSessionContext(r *http.Request, sessionID, backendID string) *http.Request {
140+
logSession.Printf("Injecting session context: sessionID=%s, backendID=%s", auth.TruncateSessionID(sessionID), backendID)
141+
142+
ctx := context.WithValue(r.Context(), SessionIDContextKey, sessionID)
143+
ctx = guard.SetAgentIDInContext(ctx, sessionID)
144+
145+
if backendID != "" {
146+
logSession.Printf("Adding backend ID to context: backendID=%s", backendID)
147+
ctx = context.WithValue(ctx, mcp.ContextKey("backend-id"), backendID)
148+
}
149+
150+
logSession.Print("Session context injected successfully")
151+
return r.WithContext(ctx)
152+
}
153+
154+
// setupSessionCallback extracts the session ID, logs the new connection, injects
155+
// the session into the request context, and returns the session ID.
156+
// Used by both routed and unified StreamableHTTP session establishment callbacks.
157+
func setupSessionCallback(r *http.Request, backendID string) (string, bool) {
158+
sessionID := extractAndValidateSession(r)
159+
if sessionID == "" {
160+
return "", false
161+
}
162+
163+
if backendID != "" {
164+
logger.LogInfo("client", "New MCP client connection, remote=%s, method=%s, path=%s, backend=%s, session=%s",
165+
r.RemoteAddr, r.Method, r.URL.Path, backendID, auth.TruncateSessionID(sessionID))
166+
} else {
167+
logger.LogInfo("client", "MCP connection established, remote=%s, method=%s, path=%s, session=%s",
168+
r.RemoteAddr, r.Method, r.URL.Path, auth.TruncateSessionID(sessionID))
169+
}
170+
171+
logHTTPRequestBody(r, sessionID, backendID)
172+
173+
*r = *injectSessionContext(r, sessionID, backendID)
174+
175+
return sessionID, true
176+
}
177+
178+
// logHTTPRequestBody logs the request body for debugging purposes.
179+
// It reads the body, logs it, and restores it so it can be read again.
180+
// The backendID parameter is optional and can be empty for unified mode.
181+
// It calls peekRequestBody (defined in http_helpers.go) which is a shared
182+
// HTTP utility also used by WithSDKLogging.
183+
func logHTTPRequestBody(r *http.Request, sessionID, backendID string) {
184+
logSession.Printf("Checking request body: method=%s, hasBody=%v, sessionID=%s", r.Method, r.Body != nil, auth.TruncateSessionID(sessionID))
185+
186+
bodyBytes, err := peekRequestBody(r)
187+
if err != nil {
188+
logSession.Printf("Body read failed: err=%v", err)
189+
return
190+
}
191+
if len(bodyBytes) == 0 {
192+
logSession.Printf("Skipping body logging: not a POST request, no body present, or empty body")
193+
return
194+
}
195+
196+
logSession.Printf("Request body read: size=%d bytes, sessionID=%s, backendID=%s", len(bodyBytes), auth.TruncateSessionID(sessionID), backendID)
197+
198+
// Sanitize the body before logging
199+
sanitizedBody := sanitize.SanitizeString(string(bodyBytes))
200+
201+
// Log with backend context if provided (routed mode)
202+
if backendID != "" {
203+
logger.LogDebug("client", "MCP client request body, backend=%s, body=%s", backendID, sanitizedBody)
204+
} else {
205+
logger.LogDebug("client", "MCP request body, session=%s, body=%s", auth.TruncateSessionID(sessionID), sanitizedBody)
206+
}
207+
logSession.Print("Request body logged for debugging")
208+
}

0 commit comments

Comments
 (0)