Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions internal/mcpproxy/handlers.go
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,7 @@ func (m *MCPProxy) servePOST(w http.ResponseWriter, r *http.Request) {
}
err = m.handleSetLoggingLevel(ctx, s, w, msg, p, span)
case "ping":
// Ping is intentionally not traced as it's a lightweight health check.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @nacx

err = m.handlePing(ctx, w, msg)
case "prompts/list":
p := &mcp.ListPromptsParams{}
Expand Down
38 changes: 37 additions & 1 deletion internal/tracing/mcp.go
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,9 @@ func (m mcpTracer) StartSpanAndInjectMeta(ctx context.Context, req *jsonrpc.Requ
parentCtx := m.propagator.Extract(ctx, mc)

// Start the span with options appropriate for the semantic convention.
newCtx, span := m.tracer.Start(parentCtx, "mcp.request", trace.WithSpanKind(trace.SpanKindClient))
// Convert method name to span name following mcp-go SDK patterns
spanName := getSpanName(req.Method)
newCtx, span := m.tracer.Start(parentCtx, spanName, trace.WithSpanKind(trace.SpanKindClient))

// Always inject trace context into the header mutation if provided.
// This ensures trace propagation works even for unsampled spans.
Expand Down Expand Up @@ -174,3 +176,37 @@ func (c metaMapCarrier) Keys() []string {

return keys
}

// getSpanName converts MCP method names to span names following mcp-go SDK patterns.
func getSpanName(method string) string {
switch method {
case "initialize":
return "Initialize"
case "tools/list":
return "ListTools"
case "tools/call":
return "CallTool"
case "prompts/list":
return "ListPrompts"
case "prompts/get":
return "GetPrompt"
case "resources/list":
return "ListResources"
case "resources/read":
return "ReadResource"
case "resources/subscribe":
return "Subscribe"
case "resources/unsubscribe":
return "Unsubscribe"
case "resources/templates/list":
return "ListResourceTemplates"
case "logging/setLevel":
return "SetLoggingLevel"
case "completion/complete":
return "Complete"
case "ping":
return "Ping"
default:
return method
}
}
105 changes: 105 additions & 0 deletions internal/tracing/mcp_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
package tracing

import (
"context"
"testing"

"github.com/modelcontextprotocol/go-sdk/jsonrpc"
Expand All @@ -15,6 +16,7 @@ import (
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/sdk/trace"
"go.opentelemetry.io/otel/sdk/trace/tracetest"
oteltrace "go.opentelemetry.io/otel/trace"
)

func TestTracer_StartSpanAndInjectMeta(t *testing.T) {
Expand Down Expand Up @@ -134,3 +136,106 @@ func Test_getMCPAttributes(t *testing.T) {
})
}
}

func Test_getSpanName(t *testing.T) {
tests := []struct {
method string
expected string
}{
{method: "initialize", expected: "Initialize"},
{method: "tools/list", expected: "ListTools"},
{method: "tools/call", expected: "CallTool"},
{method: "prompts/list", expected: "ListPrompts"},
{method: "prompts/get", expected: "GetPrompt"},
{method: "resources/list", expected: "ListResources"},
{method: "resources/read", expected: "ReadResource"},
{method: "resources/subscribe", expected: "Subscribe"},
{method: "resources/unsubscribe", expected: "Unsubscribe"},
{method: "resources/templates/list", expected: "ListResourceTemplates"},
{method: "logging/setLevel", expected: "SetLoggingLevel"},
{method: "completion/complete", expected: "Complete"},
{method: "ping", expected: "Ping"},
}

for _, tt := range tests {
t.Run(tt.method, func(t *testing.T) {
actual := getSpanName(tt.method)
require.Equal(t, tt.expected, actual)
})
}
}

func TestMCPTracer_SpanName(t *testing.T) {
tests := []struct {
name string
method string
params mcp.Params
expectedSpanName string
}{
{
name: "tools/list",
method: "tools/list",
params: &mcp.ListToolsParams{},
expectedSpanName: "ListTools",
},
{
name: "tools/call",
method: "tools/call",
params: &mcp.CallToolParams{Name: "test-tool"},
expectedSpanName: "CallTool",
},
{
name: "prompts/list",
method: "prompts/list",
params: &mcp.ListPromptsParams{},
expectedSpanName: "ListPrompts",
},
{
name: "prompts/get",
method: "prompts/get",
params: &mcp.GetPromptParams{Name: "test-prompt"},
expectedSpanName: "GetPrompt",
},
{
name: "resources/list",
method: "resources/list",
params: &mcp.ListResourcesParams{},
expectedSpanName: "ListResources",
},
{
name: "resources/read",
method: "resources/read",
params: &mcp.ReadResourceParams{URI: "test://uri"},
expectedSpanName: "ReadResource",
},
{
name: "initialize",
method: "initialize",
params: &mcp.InitializeParams{},
expectedSpanName: "Initialize",
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
exporter := tracetest.NewInMemoryExporter()
tp := trace.NewTracerProvider(trace.WithSyncer(exporter))

tracer := newMCPTracer(tp.Tracer("test"), autoprop.NewTextMapPropagator())

reqID, _ := jsonrpc.MakeID("test-id")
req := &jsonrpc.Request{ID: reqID, Method: tt.method}

span := tracer.StartSpanAndInjectMeta(context.Background(), req, tt.params)
require.NotNil(t, span)
span.EndSpan()

spans := exporter.GetSpans()
require.Len(t, spans, 1)
actualSpan := spans[0]

require.Equal(t, tt.expectedSpanName, actualSpan.Name)
require.Equal(t, oteltrace.SpanKindClient, actualSpan.SpanKind)
})
}
}
13 changes: 7 additions & 6 deletions tests/extproc/mcp/env.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ import (

"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/stretchr/testify/require"
v1 "go.opentelemetry.io/proto/otlp/trace/v1"

"github.com/envoyproxy/ai-gateway/internal/filterapi"
"github.com/envoyproxy/ai-gateway/internal/testing/testotel"
Expand Down Expand Up @@ -241,16 +240,18 @@ func (m *mcpEnv) newSession(t *testing.T) *mcpSession {
ret.session, err = m.client.Connect(t.Context(), &mcp.StreamableClientTransport{Endpoint: m.baseURL}, nil)
require.NoError(t, err)
span := m.collector.TakeSpan()
require.Equal(t, "mcp.request", span.Name)
require.Equal(t, v1.Span_SPAN_KIND_CLIENT, span.Kind)

t.Log("created new MCP session with ID ", ret.session.ID(), ", first span: ", span.String())
require.NotNil(t, span)
requireMCPSpan(t, span, "Initialize", map[string]string{
"mcp.method.name": "initialize",
"mcp.client.name": "demo-http-client",
"mcp.client.title": "",
"mcp.client.version": "0.1.0",
})

// **NOTE*** Do not add any direct access to the in-memory server session. Otherwise, the tests will result in
// not being able to run in end-to-end tests. The test code must solely operate through the client side sessions.
t.Cleanup(func() {
require.NoError(t, ret.session.Close())
_ = ret.session.Close()
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should be a linter for this somewhere. @anuraaga come across any that notice when accidental use of require.RequreXX n a cleanup? t.LogXX is fine, but require is bad

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's a close one but no cigar

https://github.com/Antonboom/testifylint?tab=readme-ov-file#go-require

Probably not too hard to adapt that to apply to both cases though

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cc @nacx

m.mux.Lock()
defer m.mux.Unlock()
delete(m.sessions, ret.session.ID())
Expand Down
Loading
Loading