Skip to content

Commit 92daffb

Browse files
garrettladleygarrettladley
authored
fix(observability): split mcp transaction names (#16)
1 parent d16c620 commit 92daffb

4 files changed

Lines changed: 277 additions & 6 deletions

File tree

‎internal/httpserver/server.go‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"net/http"
1010
"os"
1111
"os/signal"
12+
"strings"
1213
"syscall"
1314
"time"
1415

@@ -81,13 +82,9 @@ func Run(ctx context.Context, cfg Config, logger *slog.Logger) error {
8182
middleware.RateLimit(store, cfg.RateLimit, logger),
8283
)(mux)
8384
handler = otelhttp.NewHandler(handler, "http.server",
84-
otelhttp.WithSpanNameFormatter(func(_ string, r *http.Request) string {
85-
if r.Pattern != "" {
86-
return r.Method + " " + r.Pattern
87-
}
88-
return r.Method + " " + r.URL.Path
89-
}),
85+
otelhttp.WithSpanNameFormatter(httpSpanName),
9086
)
87+
handler = middleware.MCPRequestMetadata(handler)
9188

9289
baseCtx, cancelBase := context.WithCancel(ctx)
9390
defer cancelBase()
@@ -150,6 +147,24 @@ func observabilityOptions(cfg config.Observability) observability.Options {
150147
}
151148
}
152149

150+
func httpSpanName(_ string, r *http.Request) string {
151+
name := r.Method + " " + r.URL.Path
152+
if r.Pattern != "" {
153+
if strings.HasPrefix(r.Pattern, r.Method+" ") {
154+
name = r.Pattern
155+
} else {
156+
name = r.Method + " " + r.Pattern
157+
}
158+
}
159+
if method := r.Header.Get(middleware.HeaderInternalMCPMethod); method != "" {
160+
name += " " + method
161+
if mcpName := r.Header.Get(middleware.HeaderInternalMCPName); mcpName != "" {
162+
name += " " + mcpName
163+
}
164+
}
165+
return name
166+
}
167+
153168
func health(w http.ResponseWriter, _ *http.Request) {
154169
w.Header().Set("Content-Type", "application/json")
155170
_, _ = fmt.Fprintln(w, `{"status":"ok"}`)

‎internal/httpserver/server_test.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import (
44
"net/http"
55
"net/http/httptest"
66
"testing"
7+
8+
"github.com/garrettladley/pkgsite-mcp/internal/middleware"
79
)
810

911
func TestHealthDoesNotExposeBuildMetadata(t *testing.T) {
@@ -24,3 +26,37 @@ func TestHealthDoesNotExposeBuildMetadata(t *testing.T) {
2426
t.Fatalf("content type = %q, want %q", got, want)
2527
}
2628
}
29+
30+
func TestHTTPSpanNameIncludesMCPMethodAndName(t *testing.T) {
31+
t.Parallel()
32+
33+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/mcp", nil)
34+
req.Header.Set(middleware.HeaderInternalMCPMethod, "tools/call")
35+
req.Header.Set(middleware.HeaderInternalMCPName, "pkgsite_search")
36+
37+
if got, want := httpSpanName("", req), "POST /mcp tools/call pkgsite_search"; got != want {
38+
t.Fatalf("httpSpanName() = %q, want %q", got, want)
39+
}
40+
}
41+
42+
func TestHTTPSpanNameIncludesMCPMethodWithoutName(t *testing.T) {
43+
t.Parallel()
44+
45+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/mcp", nil)
46+
req.Header.Set(middleware.HeaderInternalMCPMethod, "tools/list")
47+
48+
if got, want := httpSpanName("", req), "POST /mcp tools/list"; got != want {
49+
t.Fatalf("httpSpanName() = %q, want %q", got, want)
50+
}
51+
}
52+
53+
func TestHTTPSpanNameFallsBackToHTTPRoute(t *testing.T) {
54+
t.Parallel()
55+
56+
req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/health", nil)
57+
req.Pattern = "GET /health"
58+
59+
if got, want := httpSpanName("", req), "GET /health"; got != want {
60+
t.Fatalf("httpSpanName() = %q, want %q", got, want)
61+
}
62+
}

‎internal/middleware/mcp_metadata.go‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package middleware
2+
3+
import (
4+
"bytes"
5+
"encoding/json"
6+
"io"
7+
"net/http"
8+
"strings"
9+
)
10+
11+
const (
12+
HeaderInternalMCPMethod = "X-Pkgsite-Mcp-Method"
13+
HeaderInternalMCPName = "X-Pkgsite-Mcp-Name"
14+
15+
maxMCPMetadataBodyBytes int64 = 64 << 10
16+
)
17+
18+
type mcpJSONRPCRequest struct {
19+
Method string `json:"method"`
20+
Params json.RawMessage `json:"params"`
21+
}
22+
23+
// MCPRequestMetadata extracts bounded MCP routing metadata before HTTP
24+
// instrumentation names the root server span.
25+
func MCPRequestMetadata(next http.Handler) http.Handler {
26+
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
27+
r.Header.Del(HeaderInternalMCPMethod)
28+
r.Header.Del(HeaderInternalMCPName)
29+
30+
if r.Method == http.MethodPost && r.URL.Path == "/mcp" && r.Body != nil && r.ContentLength >= 0 && r.ContentLength <= maxMCPMetadataBodyBytes {
31+
method, name := readMCPRequestMetadata(r)
32+
if method != "" {
33+
r.Header.Set(HeaderInternalMCPMethod, method)
34+
}
35+
if name != "" {
36+
r.Header.Set(HeaderInternalMCPName, name)
37+
}
38+
}
39+
next.ServeHTTP(w, r)
40+
})
41+
}
42+
43+
func readMCPRequestMetadata(r *http.Request) (string, string) {
44+
body, err := io.ReadAll(io.LimitReader(r.Body, maxMCPMetadataBodyBytes+1))
45+
r.Body = io.NopCloser(bytes.NewReader(body))
46+
if err != nil {
47+
return "", ""
48+
}
49+
if int64(len(body)) > maxMCPMetadataBodyBytes {
50+
return "", ""
51+
}
52+
53+
var req mcpJSONRPCRequest
54+
if err := json.Unmarshal(body, &req); err != nil {
55+
return "", ""
56+
}
57+
method := strings.TrimSpace(req.Method)
58+
if method == "" {
59+
return "", ""
60+
}
61+
return method, safeMCPName(method, req.Params)
62+
}
63+
64+
func safeMCPName(method string, params json.RawMessage) string {
65+
switch method {
66+
case "tools/call", "prompts/get":
67+
var named struct {
68+
Name string `json:"name"`
69+
}
70+
if err := json.Unmarshal(params, &named); err != nil {
71+
return ""
72+
}
73+
return strings.TrimSpace(named.Name)
74+
default:
75+
return ""
76+
}
77+
}

‎internal/middleware/mcp_metadata_test.go‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
package middleware
2+
3+
import (
4+
"errors"
5+
"io"
6+
"net/http"
7+
"net/http/httptest"
8+
"strings"
9+
"testing"
10+
)
11+
12+
func TestMCPRequestMetadataExtractsToolCall(t *testing.T) {
13+
t.Parallel()
14+
15+
const body = `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"pkgsite_search","arguments":{"query":"slices"}}}`
16+
handler := MCPRequestMetadata(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
17+
if got, want := r.Header.Get(HeaderInternalMCPMethod), "tools/call"; got != want {
18+
t.Fatalf("method = %q, want %q", got, want)
19+
}
20+
if got, want := r.Header.Get(HeaderInternalMCPName), "pkgsite_search"; got != want {
21+
t.Fatalf("name = %q, want %q", got, want)
22+
}
23+
gotBody, err := io.ReadAll(r.Body)
24+
if err != nil {
25+
t.Fatal(err)
26+
}
27+
if string(gotBody) != body {
28+
t.Fatalf("body = %q, want %q", string(gotBody), body)
29+
}
30+
}))
31+
32+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/mcp", strings.NewReader(body))
33+
handler.ServeHTTP(httptest.NewRecorder(), req)
34+
}
35+
36+
func TestMCPRequestMetadataExtractsListMethodWithoutName(t *testing.T) {
37+
t.Parallel()
38+
39+
handler := MCPRequestMetadata(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
40+
if got, want := r.Header.Get(HeaderInternalMCPMethod), "tools/list"; got != want {
41+
t.Fatalf("method = %q, want %q", got, want)
42+
}
43+
if got := r.Header.Get(HeaderInternalMCPName); got != "" {
44+
t.Fatalf("name = %q, want empty", got)
45+
}
46+
}))
47+
48+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))
49+
handler.ServeHTTP(httptest.NewRecorder(), req)
50+
}
51+
52+
func TestMCPRequestMetadataSkipsNonMCPRequest(t *testing.T) {
53+
t.Parallel()
54+
55+
handler := MCPRequestMetadata(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
56+
if got := r.Header.Get(HeaderInternalMCPMethod); got != "" {
57+
t.Fatalf("method = %q, want empty", got)
58+
}
59+
}))
60+
61+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/other", strings.NewReader(`{"method":"tools/list"}`))
62+
handler.ServeHTTP(httptest.NewRecorder(), req)
63+
}
64+
65+
func TestMCPRequestMetadataClearsCallerSuppliedInternalHeaders(t *testing.T) {
66+
t.Parallel()
67+
68+
handler := MCPRequestMetadata(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
69+
if got := r.Header.Get(HeaderInternalMCPMethod); got != "" {
70+
t.Fatalf("method = %q, want empty", got)
71+
}
72+
if got := r.Header.Get(HeaderInternalMCPName); got != "" {
73+
t.Fatalf("name = %q, want empty", got)
74+
}
75+
}))
76+
77+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/other", strings.NewReader(`{"method":"tools/list"}`))
78+
req.Header.Set(HeaderInternalMCPMethod, "tools/call")
79+
req.Header.Set(HeaderInternalMCPName, "pkgsite_search")
80+
handler.ServeHTTP(httptest.NewRecorder(), req)
81+
}
82+
83+
func TestMCPRequestMetadataSkipsUnknownLengthBodyWithoutTruncating(t *testing.T) {
84+
t.Parallel()
85+
86+
body := strings.Repeat("x", int(maxMCPMetadataBodyBytes)+1)
87+
handler := MCPRequestMetadata(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
88+
if got := r.Header.Get(HeaderInternalMCPMethod); got != "" {
89+
t.Fatalf("method = %q, want empty", got)
90+
}
91+
gotBody, err := io.ReadAll(r.Body)
92+
if err != nil {
93+
t.Fatal(err)
94+
}
95+
if string(gotBody) != body {
96+
t.Fatalf("body length = %d, want %d", len(gotBody), len(body))
97+
}
98+
}))
99+
100+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/mcp", strings.NewReader(body))
101+
req.ContentLength = -1
102+
req.Header.Set(HeaderInternalMCPMethod, "tools/list")
103+
handler.ServeHTTP(httptest.NewRecorder(), req)
104+
}
105+
106+
func TestMCPRequestMetadataRestoresPartialBodyAfterReadError(t *testing.T) {
107+
t.Parallel()
108+
109+
const body = `{"jsonrpc":"2.0"`
110+
handler := MCPRequestMetadata(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
111+
if got := r.Header.Get(HeaderInternalMCPMethod); got != "" {
112+
t.Fatalf("method = %q, want empty", got)
113+
}
114+
gotBody, err := io.ReadAll(r.Body)
115+
if err != nil {
116+
t.Fatal(err)
117+
}
118+
if string(gotBody) != body {
119+
t.Fatalf("body = %q, want %q", string(gotBody), body)
120+
}
121+
}))
122+
123+
req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/mcp", &errorReader{data: body})
124+
req.ContentLength = int64(len(body))
125+
handler.ServeHTTP(httptest.NewRecorder(), req)
126+
}
127+
128+
type errorReader struct {
129+
data string
130+
done bool
131+
}
132+
133+
func (r *errorReader) Close() error {
134+
return nil
135+
}
136+
137+
func (r *errorReader) Read(p []byte) (int, error) {
138+
if r.done {
139+
return 0, io.EOF
140+
}
141+
r.done = true
142+
return copy(p, r.data), errors.New("read failed")
143+
}

0 commit comments

Comments
 (0)