Skip to content

Commit 478a0bd

Browse files
AjmeraParth132Yuan325
AjmeraParth132
and
Yuan325
authored
feat: update/add detailed telemetry for stdio and http mcp transports (#1987)
## Description This PR adds consistent and actionable telemetry for MCP sessions across HTTP and STDIO transports, enabling quick visibility into toolset discovery and tool invocation activity with minimal setup. ## PR Checklist - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [x] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change --------- Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
1 parent 2d341ac commit 478a0bd

6 files changed

Lines changed: 251 additions & 27 deletions

File tree

‎internal/server/mcp.go‎

Lines changed: 164 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,11 @@ import (
3636
v20241105 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20241105"
3737
v20250326 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250326"
3838
"github.com/googleapis/genai-toolbox/internal/util"
39+
"go.opentelemetry.io/otel"
3940
"go.opentelemetry.io/otel/attribute"
4041
"go.opentelemetry.io/otel/codes"
4142
"go.opentelemetry.io/otel/metric"
43+
"go.opentelemetry.io/otel/trace"
4244
)
4345

4446
type sseSession struct {
@@ -116,6 +118,55 @@ type stdioSession struct {
116118
writer io.Writer
117119
}
118120

121+
// traceContextCarrier implements propagation.TextMapCarrier for extracting trace context from _meta
122+
type traceContextCarrier map[string]string
123+
124+
func (c traceContextCarrier) Get(key string) string {
125+
return c[key]
126+
}
127+
128+
func (c traceContextCarrier) Set(key, value string) {
129+
c[key] = value
130+
}
131+
132+
func (c traceContextCarrier) Keys() []string {
133+
keys := make([]string, 0, len(c))
134+
for k := range c {
135+
keys = append(keys, k)
136+
}
137+
return keys
138+
}
139+
140+
// extractTraceContext extracts W3C Trace Context from params._meta
141+
func extractTraceContext(ctx context.Context, body []byte) context.Context {
142+
// Try to parse the request to extract _meta
143+
var req struct {
144+
Params struct {
145+
Meta struct {
146+
Traceparent string `json:"traceparent,omitempty"`
147+
Tracestate string `json:"tracestate,omitempty"`
148+
} `json:"_meta,omitempty"`
149+
} `json:"params,omitempty"`
150+
}
151+
152+
if err := json.Unmarshal(body, &req); err != nil {
153+
return ctx
154+
}
155+
156+
// If traceparent is present, extract the context
157+
if req.Params.Meta.Traceparent != "" {
158+
carrier := traceContextCarrier{
159+
"traceparent": req.Params.Meta.Traceparent,
160+
}
161+
if req.Params.Meta.Tracestate != "" {
162+
carrier["tracestate"] = req.Params.Meta.Tracestate
163+
}
164+
return otel.GetTextMapPropagator().Extract(ctx, carrier)
165+
}
166+
167+
return ctx
168+
}
169+
119170
func NewStdioSession(s *Server, stdin io.Reader, stdout io.Writer) *stdioSession {
120171
stdioSession := &stdioSession{
121172
server: s,
@@ -142,18 +193,29 @@ func (s *stdioSession) readInputStream(ctx context.Context) error {
142193
}
143194
return err
144195
}
145-
v, res, err := processMcpMessage(ctx, []byte(line), s.server, s.protocol, "", "", nil)
196+
// This ensures the transport span becomes a child of the client span
197+
msgCtx := extractTraceContext(ctx, []byte(line))
198+
199+
// Create span for STDIO transport
200+
msgCtx, span := s.server.instrumentation.Tracer.Start(msgCtx, "toolbox/server/mcp/stdio",
201+
trace.WithSpanKind(trace.SpanKindServer),
202+
)
203+
defer span.End()
204+
205+
v, res, err := processMcpMessage(msgCtx, []byte(line), s.server, s.protocol, "", "", nil, "")
146206
if err != nil {
147207
// errors during the processing of message will generate a valid MCP Error response.
148208
// server can continue to run.
149-
s.server.logger.ErrorContext(ctx, err.Error())
209+
s.server.logger.ErrorContext(msgCtx, err.Error())
210+
span.SetStatus(codes.Error, err.Error())
150211
}
212+
151213
if v != "" {
152214
s.protocol = v
153215
}
154216
// no responses for notifications
155217
if res != nil {
156-
if err = s.write(ctx, res); err != nil {
218+
if err = s.write(msgCtx, res); err != nil {
157219
return err
158220
}
159221
}
@@ -239,7 +301,9 @@ func mcpRouter(s *Server) (chi.Router, error) {
239301

240302
// sseHandler handles sse initialization and message.
241303
func sseHandler(s *Server, w http.ResponseWriter, r *http.Request) {
242-
ctx, span := s.instrumentation.Tracer.Start(r.Context(), "toolbox/server/mcp/sse")
304+
ctx, span := s.instrumentation.Tracer.Start(r.Context(), "toolbox/server/mcp/sse",
305+
trace.WithSpanKind(trace.SpanKindServer),
306+
)
243307
r = r.WithContext(ctx)
244308

245309
sessionId := uuid.New().String()
@@ -335,9 +399,27 @@ func methodNotAllowed(s *Server, w http.ResponseWriter, r *http.Request) {
335399
func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
336400
w.Header().Set("Content-Type", "application/json")
337401

338-
ctx, span := s.instrumentation.Tracer.Start(r.Context(), "toolbox/server/mcp")
402+
ctx := r.Context()
403+
ctx = util.WithLogger(ctx, s.logger)
404+
405+
// Read body first so we can extract trace context
406+
body, err := io.ReadAll(r.Body)
407+
if err != nil {
408+
// Generate a new uuid if unable to decode
409+
id := uuid.New().String()
410+
s.logger.DebugContext(ctx, err.Error())
411+
render.JSON(w, r, jsonrpc.NewError(id, jsonrpc.PARSE_ERROR, err.Error(), nil))
412+
return
413+
}
414+
415+
// This ensures the transport span becomes a child of the client span
416+
ctx = extractTraceContext(ctx, body)
417+
418+
// Create span for HTTP transport
419+
ctx, span := s.instrumentation.Tracer.Start(ctx, "toolbox/server/mcp/http",
420+
trace.WithSpanKind(trace.SpanKindServer),
421+
)
339422
r = r.WithContext(ctx)
340-
ctx = util.WithLogger(r.Context(), s.logger)
341423

342424
var sessionId, protocolVersion string
343425
var session *sseSession
@@ -379,7 +461,6 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
379461
s.logger.DebugContext(ctx, fmt.Sprintf("toolset name: %s", toolsetName))
380462
span.SetAttributes(attribute.String("toolset_name", toolsetName))
381463

382-
var err error
383464
defer func() {
384465
if err != nil {
385466
span.SetStatus(codes.Error, err.Error())
@@ -398,17 +479,9 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
398479
)
399480
}()
400481

401-
// Read and returns a body from io.Reader
402-
body, err := io.ReadAll(r.Body)
403-
if err != nil {
404-
// Generate a new uuid if unable to decode
405-
id := uuid.New().String()
406-
s.logger.DebugContext(ctx, err.Error())
407-
render.JSON(w, r, jsonrpc.NewError(id, jsonrpc.PARSE_ERROR, err.Error(), nil))
408-
return
409-
}
482+
networkProtocolVersion := fmt.Sprintf("%d.%d", r.ProtoMajor, r.ProtoMinor)
410483

411-
v, res, err := processMcpMessage(ctx, body, s, protocolVersion, toolsetName, promptsetName, r.Header)
484+
v, res, err := processMcpMessage(ctx, body, s, protocolVersion, toolsetName, promptsetName, r.Header, networkProtocolVersion)
412485
if err != nil {
413486
s.logger.DebugContext(ctx, fmt.Errorf("error processing message: %w", err).Error())
414487
}
@@ -458,7 +531,7 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
458531
}
459532

460533
// processMcpMessage process the messages received from clients
461-
func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVersion string, toolsetName string, promptsetName string, header http.Header) (string, any, error) {
534+
func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVersion string, toolsetName string, promptsetName string, header http.Header, networkProtocolVersion string) (string, any, error) {
462535
logger, err := util.LoggerFromContext(ctx)
463536
if err != nil {
464537
return "", jsonrpc.NewError("", jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
@@ -494,31 +567,95 @@ func processMcpMessage(ctx context.Context, body []byte, s *Server, protocolVers
494567
return "", jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
495568
}
496569

570+
// Create method-specific span with semantic conventions
571+
// Note: Trace context is already extracted and set in ctx by the caller
572+
ctx, span := s.instrumentation.Tracer.Start(ctx, baseMessage.Method,
573+
trace.WithSpanKind(trace.SpanKindServer),
574+
)
575+
defer span.End()
576+
577+
// Determine network transport and protocol based on header presence
578+
networkTransport := "pipe" // default for stdio
579+
networkProtocolName := "stdio"
580+
if header != nil {
581+
networkTransport = "tcp" // HTTP/SSE transport
582+
networkProtocolName = "http"
583+
}
584+
585+
// Set required semantic attributes for span according to OTEL MCP semcov
586+
// ref: https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/#server
587+
span.SetAttributes(
588+
attribute.String("mcp.method.name", baseMessage.Method),
589+
attribute.String("network.transport", networkTransport),
590+
attribute.String("network.protocol.name", networkProtocolName),
591+
)
592+
593+
// Set network protocol version if available
594+
if networkProtocolVersion != "" {
595+
span.SetAttributes(attribute.String("network.protocol.version", networkProtocolVersion))
596+
}
597+
598+
// Set MCP protocol version if available
599+
if protocolVersion != "" {
600+
span.SetAttributes(attribute.String("mcp.protocol.version", protocolVersion))
601+
}
602+
603+
// Set request ID
604+
if baseMessage.Id != nil {
605+
span.SetAttributes(attribute.String("jsonrpc.request.id", fmt.Sprintf("%v", baseMessage.Id)))
606+
}
607+
608+
// Set toolset name
609+
span.SetAttributes(attribute.String("toolset.name", toolsetName))
610+
497611
// Check if message is a notification
498612
if baseMessage.Id == nil {
499613
err := mcp.NotificationHandler(ctx, body)
614+
if err != nil {
615+
span.SetStatus(codes.Error, err.Error())
616+
}
500617
return "", nil, err
501618
}
502619

620+
// Process the method
503621
switch baseMessage.Method {
504622
case mcputil.INITIALIZE:
505-
res, v, err := mcp.InitializeResponse(ctx, baseMessage.Id, body, s.version)
623+
result, version, err := mcp.InitializeResponse(ctx, baseMessage.Id, body, s.version)
506624
if err != nil {
507-
return "", res, err
625+
span.SetStatus(codes.Error, err.Error())
626+
if rpcErr, ok := result.(jsonrpc.JSONRPCError); ok {
627+
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
628+
}
629+
return "", result, err
508630
}
509-
return v, res, err
631+
span.SetAttributes(attribute.String("mcp.protocol.version", version))
632+
return version, result, err
510633
default:
511634
toolset, ok := s.ResourceMgr.GetToolset(toolsetName)
512635
if !ok {
513-
err = fmt.Errorf("toolset does not exist")
514-
return "", jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
636+
err := fmt.Errorf("toolset does not exist")
637+
rpcErr := jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil)
638+
span.SetStatus(codes.Error, err.Error())
639+
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
640+
return "", rpcErr, err
515641
}
516642
promptset, ok := s.ResourceMgr.GetPromptset(promptsetName)
517643
if !ok {
518-
err = fmt.Errorf("promptset does not exist")
519-
return "", jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
644+
err := fmt.Errorf("promptset does not exist")
645+
rpcErr := jsonrpc.NewError(baseMessage.Id, jsonrpc.INVALID_REQUEST, err.Error(), nil)
646+
span.SetStatus(codes.Error, err.Error())
647+
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
648+
return "", rpcErr, err
649+
}
650+
result, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, toolset, promptset, s.ResourceMgr, body, header)
651+
if err != nil {
652+
span.SetStatus(codes.Error, err.Error())
653+
// Set error.type based on JSON-RPC error code
654+
if rpcErr, ok := result.(jsonrpc.JSONRPCError); ok {
655+
span.SetAttributes(attribute.Int("jsonrpc.error.code", rpcErr.Error.Code))
656+
span.SetAttributes(attribute.String("error.type", rpcErr.Error.String()))
657+
}
520658
}
521-
res, err := mcp.ProcessMethod(ctx, protocolVersion, baseMessage.Id, baseMessage.Method, toolset, promptset, s.ResourceMgr, body, header)
522-
return "", res, err
659+
return "", result, err
523660
}
524661
}

‎internal/server/mcp/jsonrpc/jsonrpc.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,9 @@ type Request struct {
4545
// notifications. The receiver is not obligated to provide these
4646
// notifications.
4747
ProgressToken ProgressToken `json:"progressToken,omitempty"`
48+
// W3C Trace Context fields for distributed tracing
49+
Traceparent string `json:"traceparent,omitempty"`
50+
Tracestate string `json:"tracestate,omitempty"`
4851
} `json:"_meta,omitempty"`
4952
} `json:"params,omitempty"`
5053
}
@@ -97,6 +100,24 @@ type Error struct {
97100
Data interface{} `json:"data,omitempty"`
98101
}
99102

103+
// String returns the error type as a string based on the error code.
104+
func (e Error) String() string {
105+
switch e.Code {
106+
case METHOD_NOT_FOUND:
107+
return "method_not_found"
108+
case INVALID_PARAMS:
109+
return "invalid_params"
110+
case INTERNAL_ERROR:
111+
return "internal_error"
112+
case PARSE_ERROR:
113+
return "parse_error"
114+
case INVALID_REQUEST:
115+
return "invalid_request"
116+
default:
117+
return "jsonrpc_error"
118+
}
119+
}
120+
100121
// JSONRPCError represents a non-successful (error) response to a request.
101122
type JSONRPCError struct {
102123
Jsonrpc string `json:"jsonrpc"`

‎internal/server/mcp/v20241105/method.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ import (
2828
"github.com/googleapis/genai-toolbox/internal/tools"
2929
"github.com/googleapis/genai-toolbox/internal/util"
3030
"github.com/googleapis/genai-toolbox/internal/util/parameters"
31+
"go.opentelemetry.io/otel/attribute"
32+
"go.opentelemetry.io/otel/trace"
3133
)
3234

3335
// ProcessMethod returns a response for the request.
@@ -101,6 +103,14 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
101103
toolName := req.Params.Name
102104
toolArgument := req.Params.Arguments
103105
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
106+
107+
// Update span name and set gen_ai attributes
108+
span := trace.SpanFromContext(ctx)
109+
span.SetName(fmt.Sprintf("%s %s", TOOLS_CALL, toolName))
110+
span.SetAttributes(
111+
attribute.String("gen_ai.tool.name", toolName),
112+
attribute.String("gen_ai.operation.name", "execute_tool"),
113+
)
104114
tool, ok := resourceMgr.GetTool(toolName)
105115
if !ok {
106116
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
@@ -310,6 +320,11 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *r
310320

311321
promptName := req.Params.Name
312322
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
323+
324+
// Update span name and set gen_ai attributes
325+
span := trace.SpanFromContext(ctx)
326+
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
327+
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
313328
prompt, ok := resourceMgr.GetPrompt(promptName)
314329
if !ok {
315330
err := fmt.Errorf("prompt with name %q does not exist", promptName)

‎internal/server/mcp/v20250326/method.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ import (
2828
"github.com/googleapis/genai-toolbox/internal/tools"
2929
"github.com/googleapis/genai-toolbox/internal/util"
3030
"github.com/googleapis/genai-toolbox/internal/util/parameters"
31+
"go.opentelemetry.io/otel/attribute"
32+
"go.opentelemetry.io/otel/trace"
3133
)
3234

3335
// ProcessMethod returns a response for the request.
@@ -101,6 +103,15 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
101103
toolName := req.Params.Name
102104
toolArgument := req.Params.Arguments
103105
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
106+
107+
// Update span name and set gen_ai attributes
108+
span := trace.SpanFromContext(ctx)
109+
span.SetName(fmt.Sprintf("%s %s", TOOLS_CALL, toolName))
110+
span.SetAttributes(
111+
attribute.String("gen_ai.tool.name", toolName),
112+
attribute.String("gen_ai.operation.name", "execute_tool"),
113+
)
114+
104115
tool, ok := resourceMgr.GetTool(toolName)
105116
if !ok {
106117
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
@@ -309,6 +320,12 @@ func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *r
309320

310321
promptName := req.Params.Name
311322
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
323+
324+
// Update span name and set gen_ai attributes
325+
span := trace.SpanFromContext(ctx)
326+
span.SetName(fmt.Sprintf("%s %s", PROMPTS_GET, promptName))
327+
span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName))
328+
312329
prompt, ok := resourceMgr.GetPrompt(promptName)
313330
if !ok {
314331
err := fmt.Errorf("prompt with name %q does not exist", promptName)

0 commit comments

Comments
 (0)