diff --git a/docs/server.md b/docs/server.md index 966da750..d7bc314a 100644 --- a/docs/server.md +++ b/docs/server.md @@ -975,27 +975,33 @@ results: clients use `ttlMs` as a freshness hint to reduce polling, and `cacheScope` (`"public"` or `"private"`) controls whether shared intermediaries may cache the response. -**Server-side defaults.** After paginating, the SDK sets `CacheScope = "public"` -and leaves `TTLMs = 0` (which the client cache treats as "immediately -stale"). A handler that wants its responses cached must set `TTLMs` -explicitly on the returned result. +**Server-side defaults.** The SDK fills in `CacheScope = "public"` whenever a +result leaves it empty, since the field is required on the wire. Everything +else is left as produced: a `ResourceHandler` that sets `TTLMs` or `CacheScope` +on its `ReadResourceResult` keeps those values. The list results and +`server/discover` are built by the SDK itself, so they start out with +`TTLMs = 0`, which a client cache treats as "immediately stale". + +**Setting a policy.** `ServerOptions.SetCacheable` decides the fields for every +result that carries them. It runs after the corresponding handler, receives +what that handler produced, and may keep, adjust, or replace it. The request is +passed in so that policy can vary by method, session, or authorization context. ```go -mcp.AddTool(server, &mcp.Tool{Name: "expensive"}, func(...) (*mcp.CallToolResult, any, error) { - // ... -}) -// Override the default for tools/list: -server.AddSendingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler { - return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { - res, err := next(ctx, method, req) - if lr, ok := res.(*mcp.ListToolsResult); ok { - lr.TTLMs = 30_000 // 30s cache freshness hint +server := mcp.NewServer(impl, &mcp.ServerOptions{ + SetCacheable: func(_ context.Context, req mcp.Request, c *mcp.Cacheable) { + // A 30s freshness hint, except where the handler chose its own. + if c.TTLMs == 0 { + c.TTLMs = 30_000 } - return res, err - } + }, }) ``` +`SetCacheable` may run with the server's lock held, so it must not call back +into the `Server`: adding or removing a feature, or ranging over +`Server.Sessions`, deadlocks. + ## Utilities ### Completion diff --git a/internal/docs/server.src.md b/internal/docs/server.src.md index 8ad704a7..82b956e9 100644 --- a/internal/docs/server.src.md +++ b/internal/docs/server.src.md @@ -374,27 +374,33 @@ results: clients use `ttlMs` as a freshness hint to reduce polling, and `cacheScope` (`"public"` or `"private"`) controls whether shared intermediaries may cache the response. -**Server-side defaults.** After paginating, the SDK sets `CacheScope = "public"` -and leaves `TTLMs = 0` (which the client cache treats as "immediately -stale"). A handler that wants its responses cached must set `TTLMs` -explicitly on the returned result. +**Server-side defaults.** The SDK fills in `CacheScope = "public"` whenever a +result leaves it empty, since the field is required on the wire. Everything +else is left as produced: a `ResourceHandler` that sets `TTLMs` or `CacheScope` +on its `ReadResourceResult` keeps those values. The list results and +`server/discover` are built by the SDK itself, so they start out with +`TTLMs = 0`, which a client cache treats as "immediately stale". + +**Setting a policy.** `ServerOptions.SetCacheable` decides the fields for every +result that carries them. It runs after the corresponding handler, receives +what that handler produced, and may keep, adjust, or replace it. The request is +passed in so that policy can vary by method, session, or authorization context. ```go -mcp.AddTool(server, &mcp.Tool{Name: "expensive"}, func(...) (*mcp.CallToolResult, any, error) { - // ... -}) -// Override the default for tools/list: -server.AddSendingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler { - return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { - res, err := next(ctx, method, req) - if lr, ok := res.(*mcp.ListToolsResult); ok { - lr.TTLMs = 30_000 // 30s cache freshness hint +server := mcp.NewServer(impl, &mcp.ServerOptions{ + SetCacheable: func(_ context.Context, req mcp.Request, c *mcp.Cacheable) { + // A 30s freshness hint, except where the handler chose its own. + if c.TTLMs == 0 { + c.TTLMs = 30_000 } - return res, err - } + }, }) ``` +`SetCacheable` may run with the server's lock held, so it must not call back +into the `Server`: adding or removing a feature, or ranging over +`Server.Sessions`, deadlocks. + ## Utilities ### Completion diff --git a/mcp/protocol.go b/mcp/protocol.go index 7a881d48..4867e544 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -1191,9 +1191,16 @@ func (c Cacheable) GetTTLMs() int { return c.TTLMs } // GetCacheScope returns the cache scope. func (c Cacheable) GetCacheScope() string { return c.CacheScope } -// setDefaultCacheableValues sets the default values for the cacheable fields. -func (c *Cacheable) setDefaultCacheableValues() { - c.CacheScope = "public" +// normalize fills in the protocol default for any cache field left +// unset. An absent cacheScope means "public", but the field is required on the +// wire, so the default is materialized here rather than sent empty. +// +// Values already present are preserved: this must not undo a decision made by +// a resource handler or by [ServerOptions.SetCacheable]. +func (c *Cacheable) normalize() { + if c.CacheScope == "" { + c.CacheScope = "public" + } } // The server's response to a prompts/list request from the client. diff --git a/mcp/server.go b/mcp/server.go index b935385f..8e2f4bbc 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -174,6 +174,14 @@ type ServerOptions struct { // true, since stateless servers do not maintain sessions. GetSessionID func() string + // SetCacheable, if non-nil, decides the cache-control fields (ttlMs and + // cacheScope) of every result that carries them: server/discover, the + // four list methods, and resources/read. + // + // It is called once per such result, after the corresponding handler has + // returned, with c holding the values that handler produced. + SetCacheable func(ctx context.Context, req Request, c *Cacheable) + // SupportedProtocolVersions, if non-empty, restricts the protocol versions // this server advertises. If empty, every version returned by // [SupportedProtocolVersions] is used. @@ -870,7 +878,7 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] { return slices.Values(clients) } -func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) { +func (s *Server) listPrompts(ctx context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { @@ -885,7 +893,7 @@ func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListP if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.resolveCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -915,7 +923,7 @@ func (s *Server) getPrompt(ctx context.Context, req *GetPromptRequest) (*GetProm // the server's capabilities, the server's identity, and the server's // instructions, allowing clients to negotiate without performing the legacy // initialize handshake. -func (s *Server) discover(_ context.Context, req *ServerRequest[*DiscoverParams]) (*DiscoverResult, error) { +func (s *Server) discover(ctx context.Context, req *ServerRequest[*DiscoverParams]) (*DiscoverResult, error) { req.Session.mu.Lock() versions := req.Session.supportedVersions req.Session.mu.Unlock() @@ -947,7 +955,7 @@ func (s *Server) discover(_ context.Context, req *ServerRequest[*DiscoverParams] Capabilities: s.capabilities(), Instructions: s.opts.Instructions, } - res.setDefaultCacheableValues() + s.resolveCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -968,7 +976,7 @@ func filterSupportedVersions(t Transport, versions []string) []string { return out } -func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListToolsResult, error) { +func (s *Server) listTools(ctx context.Context, req *ListToolsRequest) (*ListToolsResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { @@ -983,7 +991,7 @@ func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListTools if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.resolveCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -1016,7 +1024,7 @@ func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolR return res, err } -func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) { +func (s *Server) listResources(ctx context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { @@ -1031,11 +1039,11 @@ func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*L if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.resolveCacheable(ctx, req, &res.Cacheable) return res, nil } -func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) { +func (s *Server) listResourceTemplates(ctx context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) { s.mu.Lock() defer s.mu.Unlock() if req.Params == nil { @@ -1051,7 +1059,7 @@ func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTempl if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.resolveCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -1075,7 +1083,7 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R if err := handleMultiRoundTripResult(req.Session, s.opts.Logger, res); err != nil { return nil, err } - res.setDefaultCacheableValues() + s.resolveCacheable(ctx, req, &res.Cacheable) if res.resultType == resultTypeInputRequired { return res, nil } @@ -1094,6 +1102,22 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R return res, nil } +// resolveCacheable settles the cache-control fields of an outgoing result. +// +// On entry c holds whatever the feature handler produced, which is the zero +// value for results the SDK builds itself. [ServerOptions.SetCacheable] may +// then rewrite it, and any cacheScope still unset is filled with the protocol +// default. +// +// The list methods call this with s.mu held, so SetCacheable must not call +// back into the server; see its documentation. +func (s *Server) resolveCacheable(ctx context.Context, req Request, c *Cacheable) { + if s.opts.SetCacheable != nil { + s.opts.SetCacheable(ctx, req, c) + } + c.normalize() +} + // lookupResourceHandler returns the resource handler and MIME type for the resource or // resource template matching uri. If none, the last return value is false. func (s *Server) lookupResourceHandler(uri string) (ResourceHandler, string, bool) { diff --git a/mcp/server_test.go b/mcp/server_test.go index 41f49515..11e75e23 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -1692,6 +1692,86 @@ func TestServerSession_RejectsServerInitiated(t *testing.T) { } } +func TestServerSetCacheable(t *testing.T) { + ctx := context.Background() + // The values an opinionated resource handler sets on its own result. + handlerValue := Cacheable{TTLMs: 9_000, CacheScope: "private"} + + for _, tc := range []struct { + name string + callback func(context.Context, Request, *Cacheable) + wantSDK Cacheable // for a result whose fields nobody set + wantHandler Cacheable // for a result the handler set fields on + }{ + { + name: "no callback", + wantSDK: Cacheable{TTLMs: 0, CacheScope: "public"}, + wantHandler: handlerValue, + }, + { + name: "callback overrules the handler", + callback: func(_ context.Context, _ Request, c *Cacheable) { + c.TTLMs, c.CacheScope = 60_000, "private" + }, + wantSDK: Cacheable{TTLMs: 60_000, CacheScope: "private"}, + wantHandler: Cacheable{TTLMs: 60_000, CacheScope: "private"}, + }, + { + // The callback is handed what the handler produced, so it can + // single out the results that expressed no opinion. The scope it + // leaves empty is still filled with the protocol default. + name: "callback defers to the handler", + callback: func(_ context.Context, _ Request, c *Cacheable) { + if c.CacheScope == "" { + c.TTLMs = 1_000 + } + }, + wantSDK: Cacheable{TTLMs: 1_000, CacheScope: "public"}, + wantHandler: handlerValue, + }, + } { + t.Run(tc.name, func(t *testing.T) { + s := NewServer(testImpl, &ServerOptions{SetCacheable: tc.callback}) + s.AddResource(&Resource{URI: "test://plain", Name: "plain"}, + func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{Contents: []*ResourceContents{{URI: req.Params.URI, Text: "x"}}}, nil + }) + s.AddResource(&Resource{URI: "test://opinionated", Name: "opinionated"}, + func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{ + Cacheable: handlerValue, + Contents: []*ResourceContents{{URI: req.Params.URI, Text: "x"}}, + }, nil + }) + cs := mustConnect(t, s, nil) + + check := func(label string, got, want Cacheable) { + t.Helper() + if got != want { + t.Errorf("%s Cacheable = %+v, want %+v", label, got, want) + } + } + + tools, err := cs.ListTools(ctx, nil) // built by the SDK, never by a handler + if err != nil { + t.Fatal(err) + } + check("tools/list", tools.Cacheable, tc.wantSDK) + + for uri, want := range map[string]Cacheable{ + "test://plain": tc.wantSDK, + "test://opinionated": tc.wantHandler, + } { + res, err := cs.ReadResource(ctx, &ReadResourceParams{URI: uri}) + if err != nil { + t.Fatal(err) + } + check(uri, res.Cacheable, want) + } + }) + } +} + // TestServerSupportedProtocolVersions verifies that // [ServerOptions.SupportedProtocolVersions] narrows both the versions the // server advertises in server/discover and the versions it negotiates,