Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
12 changes: 12 additions & 0 deletions internal/testing/testotel/collector.go
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,18 @@ func (o *OTLPCollector) TakeSpan() *tracev1.Span {
}
}

// DrainSpans drains all pending spans from the channel.
func (o *OTLPCollector) DrainSpans() {
for {
select {
case <-o.spanCh:
// Discard the span
case <-time.After(otlpTimeout):
return // No more spans available
}
}
}

// DrainMetrics returns metrics or nil if none were recorded.
func (o *OTLPCollector) DrainMetrics() *metricsv1.ResourceMetrics {
select {
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