diff --git a/mcp/protocol.go b/mcp/protocol.go index 47cb9cbd..97a9475d 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -1192,7 +1192,19 @@ func (c Cacheable) GetTTLMs() int { return c.TTLMs } func (c Cacheable) GetCacheScope() string { return c.CacheScope } // setDefaultCacheableValues sets the default values for the cacheable fields. -func (c *Cacheable) setDefaultCacheableValues() { +// When defaults is non-nil, TTLMs is copied and CacheScope is copied only when +// non-empty. An empty CacheScope on defaults (or a nil defaults pointer) falls +// back to "public", matching the protocol default for an absent cacheScope. +// Used by [Server.applyDefaultCacheable] after optionally invoking +// ServerOptions.DefaultCacheable. +func (c *Cacheable) setDefaultCacheableValues(defaults *Cacheable) { + if defaults != nil { + c.TTLMs = defaults.TTLMs + if defaults.CacheScope != "" { + c.CacheScope = defaults.CacheScope + return + } + } c.CacheScope = "public" } diff --git a/mcp/server.go b/mcp/server.go index c189a8ee..a32bac37 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -169,6 +169,21 @@ type ServerOptions struct { // GetSessionID is not consulted when [StreamableHTTPOptions.Stateless] is // true, since stateless servers do not maintain sessions. GetSessionID func() string + + // DefaultCacheable, if non-nil, is called for each SDK-generated + // cacheable result (server/discover, *_list, resources/read) to + // determine the Cacheable values stamped on the response. The + // request is provided so policy can vary by method, session, + // tenant, or authentication state. + // + // If nil, SDK-generated results use CacheScope "public" and + // TTLMs 0, matching the historical defaults. + // + // If the callback returns an empty CacheScope, the SDK still stamps + // "public" (the protocol default for an absent cacheScope). + // Receiving middleware can still overwrite Cacheable on a + // per-result basis after the SDK stamps these values. + DefaultCacheable func(ctx context.Context, req Request) Cacheable } // NewServer creates a new MCP server. The resulting server has no features: @@ -836,7 +851,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 { @@ -851,7 +866,7 @@ func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListP if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.applyDefaultCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -881,7 +896,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() @@ -913,7 +928,7 @@ func (s *Server) discover(_ context.Context, req *ServerRequest[*DiscoverParams] Capabilities: s.capabilities(), Instructions: s.opts.Instructions, } - res.setDefaultCacheableValues() + s.applyDefaultCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -934,7 +949,7 @@ func filterSupportedVersions(t Transport) []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 { @@ -949,7 +964,7 @@ func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListTools if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.applyDefaultCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -982,7 +997,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 { @@ -997,11 +1012,11 @@ func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*L if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.applyDefaultCacheable(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 { @@ -1017,7 +1032,7 @@ func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTempl if err != nil { return nil, err } - res.setDefaultCacheableValues() + s.applyDefaultCacheable(ctx, req, &res.Cacheable) return res, nil } @@ -1038,7 +1053,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.applyDefaultCacheable(ctx, req, &res.Cacheable) if res.resultType == resultTypeInputRequired { return res, nil } @@ -1057,6 +1072,18 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R return res, nil } +// applyDefaultCacheable stamps Cacheable defaults from ServerOptions.DefaultCacheable +// onto c. When the option is nil, historical public/0 defaults apply. An empty +// CacheScope from the callback still falls back to "public". +func (s *Server) applyDefaultCacheable(ctx context.Context, req Request, c *Cacheable) { + var defaults *Cacheable + if s.opts.DefaultCacheable != nil { + d := s.opts.DefaultCacheable(ctx, req) + defaults = &d + } + c.setDefaultCacheableValues(defaults) +} + // 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 75d7455c..d90ae875 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -1691,3 +1691,98 @@ func TestServerSession_RejectsServerInitiated(t *testing.T) { } } } + +func TestServerDefaultCacheable(t *testing.T) { + ctx := context.Background() + + for _, tc := range []struct { + name string + callback func(context.Context, Request) Cacheable + want Cacheable + }{ + { + name: "historical defaults", + want: Cacheable{TTLMs: 0, CacheScope: "public"}, + }, + { + name: "private with TTL", + callback: func(context.Context, Request) Cacheable { + return Cacheable{TTLMs: 60_000, CacheScope: "private"} + }, + want: Cacheable{TTLMs: 60_000, CacheScope: "private"}, + }, + { + // Callback may set only TTLMs; empty CacheScope must still become "public". + name: "empty CacheScope falls back to public", + callback: func(context.Context, Request) Cacheable { + return Cacheable{TTLMs: 60_000} + }, + want: Cacheable{TTLMs: 60_000, CacheScope: "public"}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + server := NewServer(testImpl, &ServerOptions{DefaultCacheable: tc.callback}) + AddTool(server, &Tool{Name: "t", Description: "d"}, + func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + server.AddPrompt(&Prompt{Name: "p"}, func(context.Context, *GetPromptRequest) (*GetPromptResult, error) { + return &GetPromptResult{}, nil + }) + server.AddResource(&Resource{URI: "test://r", Name: "r"}, func(context.Context, *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{Contents: []*ResourceContents{{URI: "test://r", Text: "x"}}}, nil + }) + server.AddResourceTemplate(&ResourceTemplate{URITemplate: "test://{id}", Name: "rt"}, + func(context.Context, *ReadResourceRequest) (*ReadResourceResult, error) { + return &ReadResourceResult{Contents: []*ResourceContents{{Text: "x"}}}, nil + }) + + ct, st := NewInMemoryTransports() + if _, err := server.Connect(ctx, st, nil); err != nil { + t.Fatal(err) + } + cs, err := NewClient(testImpl, nil).Connect(ctx, ct, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = cs.Close() }) + + check := func(t *testing.T, label string, got Cacheable) { + t.Helper() + if got != tc.want { + t.Errorf("%s Cacheable = %+v, want %+v", label, got, tc.want) + } + } + + tools, err := cs.ListTools(ctx, nil) + if err != nil { + t.Fatal(err) + } + check(t, "ListTools", tools.Cacheable) + + prompts, err := cs.ListPrompts(ctx, nil) + if err != nil { + t.Fatal(err) + } + check(t, "ListPrompts", prompts.Cacheable) + + resources, err := cs.ListResources(ctx, nil) + if err != nil { + t.Fatal(err) + } + check(t, "ListResources", resources.Cacheable) + + templates, err := cs.ListResourceTemplates(ctx, nil) + if err != nil { + t.Fatal(err) + } + check(t, "ListResourceTemplates", templates.Cacheable) + + read, err := cs.ReadResource(ctx, &ReadResourceParams{URI: "test://r"}) + if err != nil { + t.Fatal(err) + } + check(t, "ReadResource", read.Cacheable) + }) + } +}