Repository: googleapis/genai-toolbox
File: internal/server/mcp.go (Go)

Code before:

func (s *stdioSession) Start(ctx context.Context) error {
	return s.readInputStream(ctx)
}

// readInputStream reads requests/notifications from MCP clients through stdin
func (s *stdioSession) readInputStream(ctx context.Context) error {
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		line, err := s.readLine(ctx)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}

		// Parse the message to determine the trace name
		var baseMessage jsonrpc.BaseMessage
		traceName := "toolbox/server/mcp"
		if err := json.Unmarshal([]byte(line), &baseMessage); err == nil && baseMessage.Method != "" {
			switch baseMessage.Method {
			case "tools/list":
				traceName = "toolbox/server/toolset/get"
			case "tools/call":
				traceName = "toolbox/server/tool/invoke"
			}
		}

		// Create a span for this message
		spanCtx, span := s.server.instrumentation.Tracer.Start(ctx, traceName)
		span.SetAttributes(attribute.String("transport", "stdio"))

		v, toolName, res, processErr := processMcpMessage(spanCtx, []byte(line), s.server, s.protocol, "", "", nil)
		if toolName != "" {
			span.SetAttributes(attribute.String("tool_name", toolName))
		}
		if processErr != nil {
			// errors during the processing of message will generate a valid MCP Error response.
			// server can continue to run.
			s.server.logger.ErrorContext(spanCtx, processErr.Error())
			span.SetStatus(codes.Error, processErr.Error())
		}
		span.End()

		if v != "" {
			s.protocol = v
		}
		// no responses for notifications

Review comment (suggestion):
please refer to the comment (https://github.com/googleapis/genai-toolbox/issues/1633#issuecomment-3439457943) Will it be better to add it in the processMcpMessage() function? so that all transport does get the same child span? and since there's where the method function is processed (ideally we want to prevent checking for baseMessage.Method here).

Code after:
func (c traceContextCarrier) Set(key, value string) {
	c[key] = value
}

func (c traceContextCarrier) Keys() []string {
	keys := make([]string, 0, len(c))
	for k := range c {
		keys = append(keys, k)
	}
	return keys
}

// extractTraceContext extracts W3C Trace Context from params._meta
func extractTraceContext(ctx context.Context, body []byte) context.Context {
	// Try to parse the request to extract _meta
	var req struct {
		Params struct {
			Meta struct {
				Traceparent string `json:"traceparent,omitempty"`
				Tracestate  string `json:"tracestate,omitempty"`
			} `json:"_meta,omitempty"`
		} `json:"params,omitempty"`
	}

	if err := json.Unmarshal(body, &req); err != nil {
		return ctx
	}

	// If traceparent is present, extract the context
	if req.Params.Meta.Traceparent != "" {
		carrier := traceContextCarrier{
			"traceparent": req.Params.Meta.Traceparent,
		}
		if req.Params.Meta.Tracestate != "" {
			carrier["tracestate"] = req.Params.Meta.Tracestate
		}
		return otel.GetTextMapPropagator().Extract(ctx, carrier)
	}

	return ctx
}

func NewStdioSession(s *Server, stdin io.Reader, stdout io.Writer) *stdioSession {
	stdioSession := &stdioSession{
		server: s,
		reader: bufio.NewReader(stdin),
		writer: stdout,
	}
	return stdioSession
}