From 5a2ba3ef52602d1e8c9f8c48d3de117ea86211bf Mon Sep 17 00:00:00 2001 From: King Star Date: Fri, 14 Aug 2026 04:20:18 +0800 Subject: [PATCH] mcp: keep request log levels scoped to each request --- mcp/logging.go | 10 +++++++ mcp/mcp_test.go | 76 +++++++++++++++++++++++++++++++++++++++++++++++++ mcp/server.go | 12 ++++---- 3 files changed, 93 insertions(+), 5 deletions(-) diff --git a/mcp/logging.go b/mcp/logging.go index 402e2432..e8d0cd73 100644 --- a/mcp/logging.go +++ b/mcp/logging.go @@ -68,6 +68,13 @@ func compareLevels(l1, l2 LoggingLevel) int { return cmp.Compare(mcpLevelToSlog(l1), mcpLevelToSlog(l2)) } +type logLevelContextKey struct{} + +func logLevelFromContext(ctx context.Context) (LoggingLevel, bool) { + v, ok := ctx.Value(logLevelContextKey{}).(LoggingLevel) + return v, ok +} + // LoggingHandlerOptions are options for a LoggingHandler. // // Deprecated: the logging feature is deprecated as of protocol version @@ -146,6 +153,9 @@ func NewLoggingHandler(ss *ServerSession, opts *LoggingHandlerOptions) *LoggingH func (h *LoggingHandler) Enabled(ctx context.Context, level slog.Level) bool { // This is also checked in ServerSession.LoggingMessage, so checking it here // is just an optimization that skips building the JSON. + if mcpLevel, ok := logLevelFromContext(ctx); ok { + return mcpLevel != "" && level >= mcpLevelToSlog(mcpLevel) + } h.ss.mu.Lock() mcpLevel := h.ss.state.LogLevel h.ss.mu.Unlock() diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index d9d9b3af..ade30de4 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -3439,3 +3439,79 @@ func TestCallCustomMethodTypedNilParams(t *testing.T) { t.Fatalf("CallCustomMethod with typed-nil params: %v", err) } } + +func TestServerLogLevelDoesNotLeakBetweenNewProtocolRequests(t *testing.T) { + ctx := context.Background() + s := NewServer(testImpl, nil) + _, st := NewInMemoryTransports() + ss, err := s.Connect(ctx, st, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.Close() }) + + logged := make(chan LoggingLevel, 1) + s.AddSendingMiddleware(func(next MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { + if method == notificationLoggingMessage { + logged <- req.GetParams().(*LoggingMessageParams).Level + return nil, nil + } + return next(ctx, method, req) + } + }) + + started := make(chan struct{}) + release := make(chan struct{}) + AddTool(s, &Tool{Name: "blocked-log"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + close(started) + <-release + if err := req.Session.Log(ctx, &LoggingMessageParams{Level: "warning", Data: "request log"}); err != nil { + return nil, nil, err + } + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + AddTool(s, &Tool{Name: "noop"}, func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + + withLogLevel := &CallToolParams{Name: "blocked-log"} + withLogLevel.SetMeta(newProtocolMeta("warning")) + errc := make(chan error, 1) + go func() { + _, err := ss.handle(ctx, req(1, methodCallTool, withLogLevel)) + errc <- err + }() + + <-started + withoutLogLevel := &CallToolParams{Name: "noop"} + withoutLogLevel.SetMeta(newProtocolMeta("")) + if _, err := ss.handle(ctx, req(2, methodCallTool, withoutLogLevel)); err != nil { + t.Fatal(err) + } + close(release) + if err := <-errc; err != nil { + t.Fatal(err) + } + + select { + case got := <-logged: + if got != "warning" { + t.Fatalf("logged level = %q, want warning", got) + } + default: + t.Fatal("request-scoped warning log was suppressed after another request cleared session log level") + } +} + +func newProtocolMeta(logLevel LoggingLevel) Meta { + m := Meta{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: testImpl, + MetaKeyClientCapabilities: (&ClientCapabilities{}).toV2(), + } + if logLevel != "" { + m[MetaKeyLogLevel] = logLevel + } + return m +} diff --git a/mcp/server.go b/mcp/server.go index c189a8ee..cad764e2 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1736,9 +1736,12 @@ func (ss *ServerSession) Elicit(ctx context.Context, params *ElicitParams) (*Eli // (at least twelve months). See // https://modelcontextprotocol.io/seps/2577-deprecate-roots-sampling-and-logging. func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams) error { - ss.mu.Lock() - logLevel := ss.state.LogLevel - ss.mu.Unlock() + logLevel, ok := logLevelFromContext(ctx) + if !ok { + ss.mu.Lock() + logLevel = ss.state.LogLevel + ss.mu.Unlock() + } if logLevel == "" { // The spec is unclear, but seems to imply that no log messages are sent until the client // sets the level. @@ -1926,9 +1929,8 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any, // server->client calls and notifications to the incoming request from which // they originated. See [idContextKey] for details. ctx = context.WithValue(ctx, idContextKey{}, req.ID) - // For new-protocol requests, propagate the per-request log level. if validatedMeta.usesNewProtocol { - ss.setLevel(ctx, &SetLoggingLevelParams{Level: validatedMeta.logLevel}) + ctx = context.WithValue(ctx, logLevelContextKey{}, validatedMeta.logLevel) } res, err := handleReceive(ctx, ss, req) if err != nil {