mcp: keep request log levels scoped to each request (#1168)
Fixes a request-scoping bug in the new-protocol logging path. `ServerSession.handle` decoded `_meta[io.modelcontextprotocol/logLevel]`, but stored it in `ServerSessionState.LogLevel`. Concurrent new-protocol requests on the same session could therefore clear or inherit another request's log threshold before the original handler emitted logs. This keeps the new-protocol log level on the handler context instead. `ServerSession.Log` and `LoggingHandler.Enabled` read that context value first, while legacy `logging/setLevel` behavior still uses session state. Tests: - `go test ./mcp -run TestServerLogLevelDoesNotLeakBetweenNewProtocolRequests -count=1` - `go test -race ./mcp -run TestServerLogLevelDoesNotLeakBetweenNewProtocolRequests -count=1` - `go test ./mcp -run TestEndToEnd -count=1` - `go test ./mcp -run TestValidateRequestMeta -count=1` - `go test ./mcp -run TestStreamableStateless_NewProtocolSession_NoFakeInit -count=1` - `go test ./... -skip TestSSELocalhostProtection|TestStreamableLocalhostProtection` A full `go test ./...` run in my local environment still fails the two localhost protection cases for `0.0.0.0 via localhost rejects evil.com`; those same two tests also fail when run directly and are unrelated to this logging change. Co-authored-by: Guglielmo Colombo <guglielmoc@google.com>
This commit is contained in:
parent
0eb1600eef
commit
ba1496a9aa
3 changed files with 93 additions and 5 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1739,9 +1739,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.
|
||||
|
|
@ -1929,9 +1932,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 {
|
||||
|
|
|
|||
Loading…
Reference in a new issue