From 7130f0f081cfbe90857653b41ed97c1b2c306060 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Wed, 21 May 2025 07:11:50 -0400 Subject: [PATCH] internal/mcp: meta and progress token Meta is a struct that contains a map[string]any, and a ProgressToken field. Each Params and Result type has a Meta as a value field. Meta marshals and unmarshals specially, so the ProgressToken field becomes part of the map. The marshalling code is generalized in anticipation of every Params and Result growing a map of additional properties. Change-Id: I6d2f34270d33501c96437f24f91ef3e51d20d955 Reviewed-on: https://go-review.googlesource.com/c/tools/+/674956 Reviewed-by: Sam Thanawalla LUCI-TryBot-Result: Go LUCI Auto-Submit: Jonathan Amsterdam Reviewed-by: Robert Findley --- design/design.md | 17 ++- jsonschema/validate.go | 11 +- jsonschema/validate_test.go | 21 ++++ mcp/client.go | 22 ++-- mcp/generate.go | 35 +++++-- mcp/mcp_test.go | 2 +- mcp/protocol.go | 201 +++++++++++++++++++++++++----------- mcp/server.go | 14 +-- mcp/shared.go | 85 ++++++++++++--- mcp/shared_test.go | 70 +++++++++++++ mcp/util.go | 137 ++++++++++++++++++++++++ mcp/util_test.go | 48 +++++++++ 12 files changed, 548 insertions(+), 115 deletions(-) create mode 100644 mcp/shared_test.go create mode 100644 mcp/util_test.go diff --git a/design/design.md b/design/design.md index 3b2be85..1f9dd63 100644 --- a/design/design.md +++ b/design/design.md @@ -243,7 +243,7 @@ Types needed for the protocol are generated from the [JSON schema of the MCP spe These types will be included in the `mcp` package, but will be unexported unless they are needed for the user-facing API. Notably, JSON-RPC request types are elided, since they are handled by the `jsonrpc2` package and should not be observed by the user. -For user-provided data, we use `json.RawMessage`, so that marshalling/unmarshalling can be delegated to the business logic of the client or server. +For user-provided data, we use `json.RawMessage` or `map[string]any`, depending on the use case. For union types, which can't be represented in Go (specifically `Content` and `ResourceContents`), we prefer distinguished unions: struct types with fields corresponding to the union of all properties for union elements. @@ -255,9 +255,9 @@ type ReadResourceParams struct { } type CallToolResult struct { - Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Content []Content `json:"content"` - IsError bool `json:"isError,omitempty"` + Meta Meta `json:"_meta,omitempty"` + Content []Content `json:"content"` + IsError bool `json:"isError,omitempty"` } // Content is the wire format for content. @@ -276,6 +276,7 @@ type Content struct { func NewTextContent(text string) *Content // etc. ``` +The `Meta` type includes a `map[string]any` for arbitrary data, and a `ProgressToken` field. **Differences from mcp-go**: these types are largely similar, but our type generator flattens types rather than using struct embedding. @@ -480,15 +481,21 @@ The server observes a client cancellation as a cancelled context. ### Progress handling -A caller can request progress notifications by setting the `ProgressToken` field on any request. +A caller can request progress notifications by setting the `Meta.ProgressToken` field on any request. ```go type XXXParams struct { // where XXX is each type of call + Meta Meta ... +} + +type Meta struct { + Data map[string]any ProgressToken any // string or int } ``` + Handlers can notify their peer about progress by calling the `NotifyProgress` method. The notification is only sent if the peer requested it by providing a progress token. ```go diff --git a/jsonschema/validate.go b/jsonschema/validate.go index 75c4911..e146281 100644 --- a/jsonschema/validate.go +++ b/jsonschema/validate.go @@ -616,14 +616,19 @@ func structPropertiesOf(t reflect.Type) propertyMap { // jsonName returns the name for f as would be used by [json.Marshal]. // That is the name in the json struct tag, or the field name if there is no tag. -// If f is not exported or the tag name is "-", jsonName returns "", false. +// If f is not exported or the tag is "-", jsonName returns "", false. func jsonName(f reflect.StructField) (string, bool) { if !f.IsExported() { return "", false } if tag, ok := f.Tag.Lookup("json"); ok { - if name, _, _ := strings.Cut(tag, ","); name != "" { - return name, name != "-" + name, _, found := strings.Cut(tag, ",") + // "-" means omit, but "-," means the name is "-" + if name == "-" && !found { + return "", false + } + if name != "" { + return name, true } } return f.Name, true diff --git a/jsonschema/validate_test.go b/jsonschema/validate_test.go index fca4258..8e1a568 100644 --- a/jsonschema/validate_test.go +++ b/jsonschema/validate_test.go @@ -10,6 +10,7 @@ import ( "net/url" "os" "path/filepath" + "reflect" "strings" "testing" ) @@ -171,6 +172,26 @@ func TestStructInstance(t *testing.T) { } } +func TestJSONName(t *testing.T) { + type S struct { + A int + B int `json:","` + C int `json:"-"` + D int `json:"-,"` + E int `json:"echo"` + F int `json:"foxtrot,omitempty"` + g int `json:"golf"` + } + want := []string{"A", "B", "", "-", "echo", "foxtrot", ""} + tt := reflect.TypeFor[S]() + for i := range tt.NumField() { + got, _ := jsonName(tt.Field(i)) + if got != want[i] { + t.Errorf("got %q, want %q", got, want[i]) + } + } +} + // loadRemote loads a remote reference used in the test suite. func loadRemote(uri *url.URL) (*Schema, error) { // Anything with localhost:1234 refers to the remotes directory in the test suite repo. diff --git a/mcp/client.go b/mcp/client.go index d9917b6..2f8c654 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -162,7 +162,7 @@ func (c *Client) RemoveRoots(uris ...string) { // changeAndNotify is called when a feature is added or removed. // It calls change, which should do the work and report whether a change actually occurred. // If there was a change, it notifies a snapshot of the sessions. -func (c *Client) changeAndNotify(notification string, params any, change func() bool) { +func (c *Client) changeAndNotify(notification string, params Params, change func() bool) { var sessions []*ClientSession // Lock for the change, but not for the notification. c.mu.Lock() @@ -231,8 +231,8 @@ func (cs *ClientSession) methodHandler() MethodHandler[ClientSession] { // getConn implements [session.getConn]. func (cs *ClientSession) getConn() *jsonrpc2.Connection { return cs.conn } -func (c *ClientSession) ping(ct context.Context, params *PingParams) (struct{}, error) { - return struct{}{}, nil +func (c *ClientSession) ping(ct context.Context, params *PingParams) (Result, error) { + return emptyResult{}, nil } // Ping makes an MCP "ping" request to the server. @@ -296,29 +296,21 @@ func (c *ClientSession) ReadResource(ctx context.Context, params *ReadResourcePa return standardCall[ReadResourceResult](ctx, c.conn, methodReadResource, params) } -func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (any, error) { +func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (Result, error) { return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params) } -func (c *Client) callPromptChangedHandler(ctx context.Context, s *ClientSession, params *PromptListChangedParams) (any, error) { +func (c *Client) callPromptChangedHandler(ctx context.Context, s *ClientSession, params *PromptListChangedParams) (Result, error) { return callNotificationHandler(ctx, c.opts.PromptListChangedHandler, s, params) } -func (c *Client) callResourceChangedHandler(ctx context.Context, s *ClientSession, params *ResourceListChangedParams) (any, error) { +func (c *Client) callResourceChangedHandler(ctx context.Context, s *ClientSession, params *ResourceListChangedParams) (Result, error) { return callNotificationHandler(ctx, c.opts.ResourceListChangedHandler, s, params) } -func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, params *LoggingMessageParams) (any, error) { +func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, params *LoggingMessageParams) (Result, error) { if h := c.opts.LoggingMessageHandler; h != nil { h(ctx, cs, params) } return nil, nil } - -func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) { - var result TRes - if err := call(ctx, conn, method, params, &result); err != nil { - return nil, err - } - return &result, nil -} diff --git a/mcp/generate.go b/mcp/generate.go index 719520d..67fa2f6 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -299,13 +299,21 @@ func loadSchema(schemaFile string) (data []byte, err error) { func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, named map[string]*bytes.Buffer) error { var w io.Writer = io.Discard - if typeName := config.Name; typeName != "-" { + var typeName string + if typeName = config.Name; typeName != "-" { if typeName == "" { typeName = configName } if _, ok := named[typeName]; ok { return nil } + // The JSON schema does not accurately represent the source of truth, which is typescript. + // Every Params and Result type should have a _meta property. + // Also, those with a progress token will turn into a struct; we want the progress token to + // be a map item. So replace all metas. + if strings.HasSuffix(typeName, "Params") || strings.HasSuffix(typeName, "Result") { + def.Properties["_meta"] = metaSchema + } buf := new(bytes.Buffer) w = buf named[typeName] = buf @@ -318,6 +326,12 @@ func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, nam return err // Better error here? } fmt.Fprintf(w, "\n") + + // Any decl with a _meta field gets a GetMeta method. + if _, ok := def.Properties["_meta"]; ok { + fmt.Fprintf(w, "\nfunc (x *%s) GetMeta() *Meta { return &x.Meta }", typeName) + } + return nil } @@ -354,11 +368,7 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma // For types that explicitly allow additional properties, we can either // unmarshal them into a map[string]any, or delay unmarshalling with - // json.RawMessage. For now, use json.RawMessage as it defers the choice. - // - // TODO(jba): further refine this classification of object schemas. - // For example, the typescript "object" type, which should map to a Go "any", - // is represented in schema.json by `{type: object, properties: {}, additionalProperties: true}`. + // json.RawMessage. We use any. if def.Type == "object" && canHaveAdditionalProperties(def) && def.Properties == nil { w.Write([]byte("map[string]")) return writeType(w, nil, def.AdditionalProperties, named) @@ -372,7 +382,7 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma fmt.Fprintf(w, "*Content") } else { // E.g. union types. - fmt.Fprintf(w, "json.RawMessage") + fmt.Fprintf(w, "any") } } else { switch def.Type { @@ -398,6 +408,11 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma if fieldDef.Description != "" { fmt.Fprintf(w, "%s\n", toComment(fieldDef.Description)) } + if name == "_meta" { + fmt.Fprintln(w, "\tMeta Meta `json:\"_meta,omitempty\"`") + continue + } + export := exportName(name) fmt.Fprintf(w, "\t%s ", export) @@ -551,6 +566,12 @@ func isStruct(s *jsonschema.Schema) bool { return s.Type == "object" && s.Properties != nil && !canHaveAdditionalProperties(s) } +// The schema for "_meta". +// We only need the description: the rest is a special case. +var metaSchema = &jsonschema.Schema{ + Description: "This property is reserved by the protocol to allow clients and servers to attach additional metadata to their responses.", +} + // schemaJSON returns the JSON for s. // For debugging. func schemaJSON(s *jsonschema.Schema) string { diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index b3b7c5f..5549d42 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -620,7 +620,7 @@ func TestMiddleware(t *testing.T) { // with the given prefix. func traceCalls[S ClientSession | ServerSession](w io.Writer, prefix string) Middleware[S] { return func(h MethodHandler[S]) MethodHandler[S] { - return func(ctx context.Context, sess *S, method string, params any) (any, error) { + return func(ctx context.Context, sess *S, method string, params Params) (Result, error) { fmt.Fprintf(w, "%s >%s\n", prefix, method) defer fmt.Fprintf(w, "%s <%s\n", prefix, method) return h(ctx, sess, method, params) diff --git a/mcp/protocol.go b/mcp/protocol.go index ca424c4..4405631 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -29,10 +29,15 @@ type Annotations struct { } type CallToolParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` Arguments json.RawMessage `json:"arguments,omitempty"` Name string `json:"name"` } +func (x *CallToolParams) GetMeta() *Meta { return &x.Meta } + // The server's response to a tool call. // // Any errors that originate from the tool SHOULD be reported inside the result @@ -44,17 +49,22 @@ type CallToolParams struct { // server does not support tool calls, or any other exceptional conditions, // should be reported as an MCP error response. type CallToolResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Content []*Content `json:"content"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` + Content []*Content `json:"content"` // Whether the tool call ended in an error. // // If not set, this is assumed to be false (the call was successful). IsError bool `json:"isError,omitempty"` } +func (x *CallToolResult) GetMeta() *Meta { return &x.Meta } + type CancelledParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An optional string describing the reason for the cancellation. This MAY be // logged or presented to the user. Reason string `json:"reason,omitempty"` @@ -65,6 +75,8 @@ type CancelledParams struct { RequestID any `json:"requestId"` } +func (x *CancelledParams) GetMeta() *Meta { return &x.Meta } + // Capabilities a client may support. Known capabilities are defined here, in // this schema, but this is not a closed set: any client can define its own, // additional capabilities. @@ -82,6 +94,9 @@ type ClientCapabilities struct { } type CreateMessageParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // A request to include context from one or more MCP servers (including the // caller), to be attached to the prompt. The client MAY ignore this request. IncludeContext string `json:"includeContext,omitempty"` @@ -103,15 +118,17 @@ type CreateMessageParams struct { Temperature float64 `json:"temperature,omitempty"` } +func (x *CreateMessageParams) GetMeta() *Meta { return &x.Meta } + // The client's response to a sampling/create_message request from the server. // The client should inform the user before returning the sampled message, to // allow them to inspect the response (human in the loop) and decide whether to // allow the server to see it. type CreateMessageResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Content *Content `json:"content"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` + Content *Content `json:"content"` // The name of the model that generated the message. Model string `json:"model"` Role Role `json:"role"` @@ -119,24 +136,36 @@ type CreateMessageResult struct { StopReason string `json:"stopReason,omitempty"` } +func (x *CreateMessageResult) GetMeta() *Meta { return &x.Meta } + type GetPromptParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // Arguments to use for templating the prompt. Arguments map[string]string `json:"arguments,omitempty"` // The name of the prompt or prompt template. Name string `json:"name"` } +func (x *GetPromptParams) GetMeta() *Meta { return &x.Meta } + // The server's response to a prompts/get request from the client. type GetPromptResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An optional description for the prompt. Description string `json:"description,omitempty"` Messages []*PromptMessage `json:"messages"` } +func (x *GetPromptResult) GetMeta() *Meta { return &x.Meta } + type InitializeParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` Capabilities *ClientCapabilities `json:"capabilities"` ClientInfo *implementation `json:"clientInfo"` // The latest version of the Model Context Protocol that the client supports. @@ -144,13 +173,15 @@ type InitializeParams struct { ProtocolVersion string `json:"protocolVersion"` } +func (x *InitializeParams) GetMeta() *Meta { return &x.Meta } + // After receiving an initialize request from the client, the server sends this // response. type InitializeResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Capabilities *serverCapabilities `json:"capabilities"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` + Capabilities *serverCapabilities `json:"capabilities"` // Instructions describing how to use the server and its features. // // This can be used by clients to improve the LLM's understanding of available @@ -164,83 +195,108 @@ type InitializeResult struct { ServerInfo *implementation `json:"serverInfo"` } +func (x *InitializeResult) GetMeta() *Meta { return &x.Meta } + type InitializedParams struct { - // This parameter name is reserved by MCP to allow clients and servers to attach - // additional metadata to their notifications. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` } +func (x *InitializedParams) GetMeta() *Meta { return &x.Meta } + type ListPromptsParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An opaque token representing the current pagination position. If provided, // the server should return results starting after this cursor. Cursor string `json:"cursor,omitempty"` } +func (x *ListPromptsParams) GetMeta() *Meta { return &x.Meta } + // The server's response to a prompts/list request from the client. type ListPromptsResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An opaque token representing the pagination position after the last returned // result. If present, there may be more results available. NextCursor string `json:"nextCursor,omitempty"` Prompts []*Prompt `json:"prompts"` } +func (x *ListPromptsResult) GetMeta() *Meta { return &x.Meta } + type ListResourcesParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An opaque token representing the current pagination position. If provided, // the server should return results starting after this cursor. Cursor string `json:"cursor,omitempty"` } +func (x *ListResourcesParams) GetMeta() *Meta { return &x.Meta } + // The server's response to a resources/list request from the client. type ListResourcesResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An opaque token representing the pagination position after the last returned // result. If present, there may be more results available. NextCursor string `json:"nextCursor,omitempty"` Resources []*Resource `json:"resources"` } +func (x *ListResourcesResult) GetMeta() *Meta { return &x.Meta } + type ListRootsParams struct { - Meta struct { - // If specified, the caller is requesting out-of-band progress notifications for - // this request (as represented by notifications/progress). The value of this - // parameter is an opaque token that will be attached to any subsequent - // notifications. The receiver is not obligated to provide these notifications. - ProgressToken any `json:"progressToken,omitempty"` - } `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` } +func (x *ListRootsParams) GetMeta() *Meta { return &x.Meta } + // The client's response to a roots/list request from the server. This result // contains an array of Root objects, each representing a root directory or file // that the server can operate on. type ListRootsResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Roots []*Root `json:"roots"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` + Roots []*Root `json:"roots"` } +func (x *ListRootsResult) GetMeta() *Meta { return &x.Meta } + type ListToolsParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An opaque token representing the current pagination position. If provided, // the server should return results starting after this cursor. Cursor string `json:"cursor,omitempty"` } +func (x *ListToolsParams) GetMeta() *Meta { return &x.Meta } + // The server's response to a tools/list request from the client. type ListToolsResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // An opaque token representing the pagination position after the last returned // result. If present, there may be more results available. NextCursor string `json:"nextCursor,omitempty"` Tools []*Tool `json:"tools"` } +func (x *ListToolsResult) GetMeta() *Meta { return &x.Meta } + // The severity of a log message. // // These map to syslog message severities, as specified in RFC-5424: @@ -248,6 +304,9 @@ type ListToolsResult struct { type LoggingLevel string type LoggingMessageParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // The data to be logged, such as a string message or an object. Any JSON // serializable type is allowed here. Data any `json:"data"` @@ -257,6 +316,8 @@ type LoggingMessageParams struct { Logger string `json:"logger,omitempty"` } +func (x *LoggingMessageParams) GetMeta() *Meta { return &x.Meta } + // Hints to use for model selection. // // Keys not declared here are currently left unspecified by the spec and are up @@ -310,15 +371,13 @@ type ModelPreferences struct { } type PingParams struct { - Meta struct { - // If specified, the caller is requesting out-of-band progress notifications for - // this request (as represented by notifications/progress). The value of this - // parameter is an opaque token that will be attached to any subsequent - // notifications. The receiver is not obligated to provide these notifications. - ProgressToken any `json:"progressToken,omitempty"` - } `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` } +func (x *PingParams) GetMeta() *Meta { return &x.Meta } + // A prompt or prompt template that the server offers. type Prompt struct { // A list of arguments to use for templating the prompt. @@ -340,11 +399,13 @@ type PromptArgument struct { } type PromptListChangedParams struct { - // This parameter name is reserved by MCP to allow clients and servers to attach - // additional metadata to their notifications. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` } +func (x *PromptListChangedParams) GetMeta() *Meta { return &x.Meta } + // Describes a message returned as part of a prompt. // // This is similar to `SamplingMessage`, but also supports the embedding of @@ -355,19 +416,26 @@ type PromptMessage struct { } type ReadResourceParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // The URI of the resource to read. The URI can use any protocol; it is up to // the server how to interpret it. URI string `json:"uri"` } +func (x *ReadResourceParams) GetMeta() *Meta { return &x.Meta } + // The server's response to a resources/read request from the client. type ReadResourceResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Contents *ResourceContents `json:"contents"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` + Contents *ResourceContents `json:"contents"` } +func (x *ReadResourceResult) GetMeta() *Meta { return &x.Meta } + // A known resource that the server is capable of reading. type Resource struct { // Optional annotations for the client. @@ -394,11 +462,13 @@ type Resource struct { } type ResourceListChangedParams struct { - // This parameter name is reserved by MCP to allow clients and servers to attach - // additional metadata to their notifications. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` } +func (x *ResourceListChangedParams) GetMeta() *Meta { return &x.Meta } + // The sender or recipient of messages and data in a conversation. type Role string @@ -415,11 +485,13 @@ type Root struct { } type RootsListChangedParams struct { - // This parameter name is reserved by MCP to allow clients and servers to attach - // additional metadata to their notifications. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` } +func (x *RootsListChangedParams) GetMeta() *Meta { return &x.Meta } + // Present if the client supports sampling from an LLM. type SamplingCapabilities struct { } @@ -431,12 +503,17 @@ type SamplingMessage struct { } type SetLevelParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` // The level of logging that the client wants to receive from the server. The // server should send all logs at this level and higher (i.e., more severe) to // the client as notifications/message. Level LoggingLevel `json:"level"` } +func (x *SetLevelParams) GetMeta() *Meta { return &x.Meta } + // Definition for a tool the client can call. type Tool struct { // Optional additional tool information. @@ -490,11 +567,13 @@ type ToolAnnotations struct { } type ToolListChangedParams struct { - // This parameter name is reserved by MCP to allow clients and servers to attach - // additional metadata to their notifications. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` } +func (x *ToolListChangedParams) GetMeta() *Meta { return &x.Meta } + // Describes the name and version of an MCP implementation. type implementation struct { Name string `json:"name"` diff --git a/mcp/server.go b/mcp/server.go index 7e1a089..3930a62 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -157,7 +157,7 @@ func (s *Server) RemoveResources(uris ...string) { // changeAndNotify is called when a feature is added or removed. // It calls change, which should do the work and report whether a change actually occurred. // If there was a change, it notifies a snapshot of the sessions. -func (s *Server) changeAndNotify(notification string, params any, change func() bool) { +func (s *Server) changeAndNotify(notification string, params Params, change func() bool) { var sessions []*ServerSession // Lock for the change, but not for the notification. s.mu.Lock() @@ -358,11 +358,11 @@ func (s *Server) Connect(ctx context.Context, t Transport) (*ServerSession, erro return connect(ctx, t, s) } -func (s *Server) callInitializedHandler(ctx context.Context, ss *ServerSession, params *InitializedParams) (any, error) { +func (s *Server) callInitializedHandler(ctx context.Context, ss *ServerSession, params *InitializedParams) (Result, error) { return callNotificationHandler(ctx, s.opts.InitializedHandler, ss, params) } -func (s *Server) callRootsListChangedHandler(ctx context.Context, ss *ServerSession, params *RootsListChangedParams) (any, error) { +func (s *Server) callRootsListChangedHandler(ctx context.Context, ss *ServerSession, params *RootsListChangedParams) (Result, error) { return callNotificationHandler(ctx, s.opts.RootsListChangedHandler, ss, params) } @@ -519,15 +519,15 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam }, nil } -func (ss *ServerSession) ping(context.Context, *PingParams) (struct{}, error) { - return struct{}{}, nil +func (ss *ServerSession) ping(context.Context, *PingParams) (Result, error) { + return emptyResult{}, nil } -func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (struct{}, error) { +func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (Result, error) { ss.mu.Lock() defer ss.mu.Unlock() ss.logLevel = params.Level - return struct{}{}, nil + return emptyResult{}, nil } // Close performs a graceful shutdown of the connection, preventing new diff --git a/mcp/shared.go b/mcp/shared.go index 4262620..fc7ab96 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -14,6 +14,7 @@ import ( "encoding/json" "fmt" "log" + "reflect" "slices" "time" @@ -21,12 +22,10 @@ import ( ) // A MethodHandler handles MCP messages. -// The params argument is an XXXParams struct pointer, such as *GetPromptParams. -// For methods, a MethodHandler must return either an XXResult struct pointer and a nil error, or -// nil with a non-nil error. -// For notifications, a MethodHandler must return nil, nil. +// For methods, exactly one of the return values must be nil. +// For notifications, both must be nil. type MethodHandler[S ClientSession | ServerSession] func( - ctx context.Context, _ *S, method string, params any) (result any, err error) + ctx context.Context, _ *S, method string, params Params) (result Result, err error) // Middleware is a function from MethodHandlers to MethodHandlers. type Middleware[S ClientSession | ServerSession] func(MethodHandler[S]) MethodHandler[S] @@ -56,7 +55,7 @@ func toSession[S ClientSession | ServerSession](sess *S) session[S] { } // defaultMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware. -func defaultMethodHandler[S ClientSession | ServerSession](ctx context.Context, sess *S, method string, params any) (any, error) { +func defaultMethodHandler[S ClientSession | ServerSession](ctx context.Context, sess *S, method string, params Params) (Result, error) { info, ok := toSession(sess).methodInfos()[method] if !ok { // This can be called from user code, with an arbitrary value for method. @@ -87,7 +86,7 @@ func handleRequest[S ClientSession | ServerSession](ctx context.Context, req *js // methodInfo is information about invoking a method. type methodInfo[TSession ClientSession | ServerSession] struct { // unmarshal params from the wire into an XXXParams struct - unmarshalParams func(json.RawMessage) (any, error) + unmarshalParams func(json.RawMessage) (Params, error) // run the code for the method handleMethod MethodHandler[TSession] } @@ -99,40 +98,40 @@ type methodInfo[TSession ClientSession | ServerSession] struct { // - R: results // A typedMethodHandler is like a MethodHandler, but with type information. -type typedMethodHandler[S, P, R any] func(context.Context, *S, P) (R, error) +type typedMethodHandler[S any, P Params, R Result] func(context.Context, *S, P) (R, error) // newMethodInfo creates a methodInfo from a typedMethodHandler. -func newMethodInfo[S ClientSession | ServerSession, P, R any](d typedMethodHandler[S, P, R]) methodInfo[S] { +func newMethodInfo[S ClientSession | ServerSession, P Params, R Result](d typedMethodHandler[S, P, R]) methodInfo[S] { return methodInfo[S]{ - unmarshalParams: func(m json.RawMessage) (any, error) { + unmarshalParams: func(m json.RawMessage) (Params, error) { var p P if err := json.Unmarshal(m, &p); err != nil { return nil, fmt.Errorf("unmarshaling %q into a %T: %w", m, p, err) } return p, nil }, - handleMethod: func(ctx context.Context, ss *S, _ string, params any) (any, error) { + handleMethod: func(ctx context.Context, ss *S, _ string, params Params) (Result, error) { return d(ctx, ss, params.(P)) }, } } // serverMethod is glue for creating a typedMethodHandler from a method on Server. -func serverMethod[P, R any](f func(*Server, context.Context, *ServerSession, P) (R, error)) typedMethodHandler[ServerSession, P, R] { +func serverMethod[P Params, R Result](f func(*Server, context.Context, *ServerSession, P) (R, error)) typedMethodHandler[ServerSession, P, R] { return func(ctx context.Context, ss *ServerSession, p P) (R, error) { return f(ss.server, ctx, ss, p) } } // clientMethod is glue for creating a typedMethodHandler from a method on Server. -func clientMethod[P, R any](f func(*Client, context.Context, *ClientSession, P) (R, error)) typedMethodHandler[ClientSession, P, R] { +func clientMethod[P Params, R Result](f func(*Client, context.Context, *ClientSession, P) (R, error)) typedMethodHandler[ClientSession, P, R] { return func(ctx context.Context, cs *ClientSession, p P) (R, error) { return f(cs.client, ctx, cs, p) } } // sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession. -func sessionMethod[S ClientSession | ServerSession, P, R any](f func(*S, context.Context, P) (R, error)) typedMethodHandler[S, P, R] { +func sessionMethod[S ClientSession | ServerSession, P Params, R Result](f func(*S, context.Context, P) (R, error)) typedMethodHandler[S, P, R] { return func(ctx context.Context, sess *S, p P) (R, error) { return f(sess, ctx, p) } @@ -151,7 +150,7 @@ const ( CodeUnsupportedMethod = -31001 ) -func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context.Context, h func(context.Context, *S, *P), sess *S, params *P) (any, error) { +func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context.Context, h func(context.Context, *S, *P), sess *S, params *P) (Result, error) { if h != nil { h(ctx, sess, params) } @@ -160,7 +159,7 @@ func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context // notifySessions calls Notify on all the sessions. // Should be called on a copy of the peer sessions. -func notifySessions[S ClientSession | ServerSession](sessions []*S, method string, params any) { +func notifySessions[S ClientSession | ServerSession](sessions []*S, method string, params Params) { if sessions == nil { return } @@ -174,3 +173,57 @@ func notifySessions[S ClientSession | ServerSession](sessions []*S, method strin } } } + +func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) { + var result TRes + if err := call(ctx, conn, method, params, &result); err != nil { + return nil, err + } + return &result, nil +} + +type Meta struct { + Data map[string]any `json:",omitempty"` + // For params, the progress token can be nil, a string or an integer. + // It should be nil for results. + ProgressToken any `json:"progressToken,omitempty"` +} + +type metaSansMethods Meta // avoid infinite recursion during marshaling + +func (m Meta) MarshalJSON() ([]byte, error) { + if p := m.ProgressToken; p != nil { + if k := reflect.ValueOf(p).Kind(); k != reflect.Int && k != reflect.String { + return nil, fmt.Errorf("bad type %T for Meta.ProgressToken: must be int or string", p) + } + } + // If ProgressToken is nil, accept Data["progressToken"]. We can't call marshalStructWithMap + // in that case because it will complain about duplicate fields. (We'd have to + // make it much smarter to avoid that; not worth it.) + if m.ProgressToken == nil { + return json.Marshal(m.Data) + } + return marshalStructWithMap((*metaSansMethods)(&m), "Data") +} + +func (m *Meta) UnmarshalJSON(data []byte) error { + return unmarshalStructWithMap(data, (*metaSansMethods)(m), "Data") +} + +// Params is a parameter (input) type for an MCP call or notification. +type Params interface { + // Returns a pointer to the params's Meta field. + GetMeta() *Meta +} + +// Result is a result of an MCP call. +type Result interface { + // Returns a pointer to the result's Meta field. + GetMeta() *Meta +} + +// emptyResult is returned by methods that have no result, like ping. +// Those methods cannot return nil, because jsonrpc2 cannot handle nils. +type emptyResult struct{} + +func (emptyResult) GetMeta() *Meta { panic("should never be called") } diff --git a/mcp/shared_test.go b/mcp/shared_test.go new file mode 100644 index 0000000..82aeae7 --- /dev/null +++ b/mcp/shared_test.go @@ -0,0 +1,70 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestMetaMarshal(t *testing.T) { + // Verify that Meta values round-trip. + for _, meta := range []Meta{ + {Data: nil, ProgressToken: nil}, + {Data: nil, ProgressToken: "p"}, + {Data: map[string]any{"d": true}, ProgressToken: nil}, + {Data: map[string]any{"d": true}, ProgressToken: "p"}, + } { + got := roundTrip(t, meta) + if !cmp.Equal(got, meta) { + t.Errorf("\ngot %#v\nwant %#v", got, meta) + } + } + + // Check errors. + for _, tt := range []struct { + meta Meta + want string + }{ + { + Meta{Data: map[string]any{"progressToken": "p"}, ProgressToken: 1}, + "duplicate", + }, + { + Meta{ProgressToken: true}, + "bad type", + }, + } { + _, err := json.Marshal(tt.meta) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Errorf("%+v: got %v, want error containing %q", tt.meta, err, tt.want) + } + } + + // Accept progressToken in map if the field is nil. + // It will unmarshal by populating ProgressToken. + meta := Meta{Data: map[string]any{"progressToken": "p"}} + got := roundTrip(t, meta) + want := Meta{ProgressToken: "p"} + if !cmp.Equal(got, want) { + t.Errorf("got %+v, want %+v", got, want) + } +} + +func roundTrip[T any](t *testing.T, v T) T { + t.Helper() + bytes, err := json.Marshal(v) + if err != nil { + t.Fatal(err) + } + var res T + if err := json.Unmarshal(bytes, &res); err != nil { + t.Fatal(err) + } + return res +} diff --git a/mcp/util.go b/mcp/util.go index 102c088..cc037b1 100644 --- a/mcp/util.go +++ b/mcp/util.go @@ -6,6 +6,11 @@ package mcp import ( "crypto/rand" + "encoding/json" + "fmt" + "reflect" + "strings" + "sync" ) func assert(cond bool, msg string) { @@ -27,3 +32,135 @@ func randText() string { } return string(src) } + +// marshalStructWithMap marshals its first argument to JSON, treating the field named +// mapField as an embedded map. The first argument must be a pointer to +// a struct. The underlying type of mapField must be a map[string]any, and it must have +// an "omitempty" json tag. +// +// For example, given this struct: +// +// type S struct { +// A int +// Extra map[string] any `json:,omitempty` +// } +// +// and this value: +// +// s := S{A: 1, Extra: map[string]any{"B": 2}} +// +// the call marshalJSONWithMap(s, "Extra") would return +// +// {"A": 1, "B": 2} +// +// It is an error if the map contains the same key as another struct field's +// JSON name. +// +// marshalStructWithMap calls json.Marshal on a value of type T, so T must not +// have a MarshalJSON method that calls this function, on pain of infinite regress. +// +// TODO: avoid this restriction on T by forcing it to marshal in a default way. +// See https://go.dev/play/p/EgXKJHxEx_R. +func marshalStructWithMap[T any](s *T, mapField string) ([]byte, error) { + // Marshal the struct and the map separately, and concatenate the bytes. + // This strategy is dramatically less complicated than + // constructing a synthetic struct or map with the combined keys. + if s == nil { + return []byte("null"), nil + } + s2 := *s + vMapField := reflect.ValueOf(&s2).Elem().FieldByName(mapField) + mapVal := vMapField.Interface().(map[string]any) + + // Check for duplicates. + names := jsonNames(reflect.TypeFor[T]()) + for key := range mapVal { + if names[key] { + return nil, fmt.Errorf("map key %q duplicates struct field", key) + } + } + + // Clear the map field, relying on the omitempty tag to omit it. + vMapField.Set(reflect.Zero(vMapField.Type())) + structBytes, err := json.Marshal(s2) + if err != nil { + return nil, fmt.Errorf("marshalStructWithMap(%+v): %w", s, err) + } + if len(mapVal) == 0 { + return structBytes, nil + } + mapBytes, err := json.Marshal(mapVal) + if err != nil { + return nil, err + } + if len(structBytes) == 2 { // must be "{}" + return mapBytes, nil + } + // "{X}" + "{Y}" => "{X,Y}" + res := append(structBytes[:len(structBytes)-1], ',') + res = append(res, mapBytes[1:]...) + return res, nil +} + +// unmarshalStructWithMap is the inverse of marshalStructWithMap. +// T has the same restrictions as in that function. +func unmarshalStructWithMap[T any](data []byte, v *T, mapField string) error { + // Unmarshal into the struct, ignoring unknown fields. + if err := json.Unmarshal(data, v); err != nil { + return err + } + // Unmarshal into the map. + m := map[string]any{} + if err := json.Unmarshal(data, &m); err != nil { + return err + } + // Delete from the map the fields of the struct. + for n := range jsonNames(reflect.TypeFor[T]()) { + delete(m, n) + } + if len(m) != 0 { + reflect.ValueOf(v).Elem().FieldByName(mapField).Set(reflect.ValueOf(m)) + } + return nil +} + +var jsonNamesMap sync.Map // from reflect.Type to map[string]bool + +// jsonNames returns the set of JSON object keys that t will marshal into. +// t must be a struct type. +func jsonNames(t reflect.Type) map[string]bool { + // Lock not necessary: at worst we'll duplicate work. + if val, ok := jsonNamesMap.Load(t); ok { + return val.(map[string]bool) + } + m := map[string]bool{} + for i := range t.NumField() { + if n, ok := jsonName(t.Field(i)); ok { + m[n] = true + } + } + jsonNamesMap.Store(t, m) + return m +} + +// jsonName returns the name for f as would be used by [json.Marshal]. +// That is the name in the json struct tag, or the field name if there is no tag. +// If f is not exported or the tag is "-", jsonName returns "", false. +// +// Copied from jsonschema/validate.go. +func jsonName(f reflect.StructField) (string, bool) { + if !f.IsExported() { + return "", false + } + if tag, ok := f.Tag.Lookup("json"); ok { + name, _, found := strings.Cut(tag, ",") + // "-" means omit, but "-," means the name is "-" + if name == "-" && !found { + return "", false + } + if name != "" { + return name, true + } + } + return f.Name, true +} diff --git a/mcp/util_test.go b/mcp/util_test.go new file mode 100644 index 0000000..f2cb0f5 --- /dev/null +++ b/mcp/util_test.go @@ -0,0 +1,48 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "strings" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" +) + +func TestMarshalStructWithMap(t *testing.T) { + type S struct { + A int + B string `json:"b,omitempty"` + u bool + M map[string]any `json:",omitempty"` + } + t.Run("basic", func(t *testing.T) { + s := S{A: 1, B: "two", M: map[string]any{"!@#": true}} + got, err := marshalStructWithMap(&s, "M") + if err != nil { + t.Fatal(err) + } + want := `{"A":1,"b":"two","!@#":true}` + if g := string(got); g != want { + t.Errorf("\ngot %s\nwant %s", g, want) + } + + var un S + if err := unmarshalStructWithMap(got, &un, "M"); err != nil { + t.Fatal(err) + } + if diff := cmp.Diff(s, un, cmpopts.IgnoreUnexported(S{})); diff != "" { + t.Errorf("mismatch (-want, +got):\n%s", diff) + } + }) + t.Run("duplicate", func(t *testing.T) { + s := S{A: 1, B: "two", M: map[string]any{"b": "dup"}} + _, err := marshalStructWithMap(&s, "M") + if err == nil || !strings.Contains(err.Error(), "duplicate") { + t.Errorf("got %v, want error with 'duplicate'", err) + } + }) +}