diff --git a/examples/hello/main.go b/examples/hello/main.go index 6a42bbe..82e3d6c 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -6,26 +6,46 @@ package main import ( "context" + "encoding/json" "flag" "fmt" "net/http" "os" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/internal/protocol" ) var httpAddr = flag.String("http", "", "if set, use SSE HTTP at this address, instead of stdin/stdout") -type SayHiParams struct { +type HiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *SayHiParams) ([]mcp.Content, error) { +func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) ([]mcp.Content, error) { return []mcp.Content{ mcp.TextContent{Text: "Hi " + params.Name}, }, nil } +func PromptHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) (*protocol.GetPromptResult, error) { + // (see related TODOs about cleaning up content construction) + content, err := json.Marshal(protocol.TextContent{ + Type: "text", + Text: "Say hi to " + params.Name, + }) + if err != nil { + return nil, err + } + return &protocol.GetPromptResult{ + Description: "Code review prompt", + Messages: []protocol.PromptMessage{ + // TODO: move 'Content' to the protocol package. + {Role: "user", Content: json.RawMessage(content)}, + }, + }, nil +} + func main() { flag.Parse() @@ -33,6 +53,7 @@ func main() { server.AddTools(mcp.MakeTool("greet", "say hi", SayHi, mcp.Input( mcp.Property("name", mcp.Description("the name to say hi to")), ))) + server.AddPrompts(mcp.MakePrompt("greet", "", PromptHi)) if *httpAddr != "" { handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server { diff --git a/internal/jsonschema/infer.go b/internal/jsonschema/infer.go index c0a11e5..8d96de4 100644 --- a/internal/jsonschema/infer.go +++ b/internal/jsonschema/infer.go @@ -142,6 +142,7 @@ func parseField(f reflect.StructField) (name string, required, include bool) { } name = props[0] } + // TODO: support 'omitzero' as well. required = !slices.Contains(props[1:], "omitempty") } return name, required, true diff --git a/internal/protocol/generate.go b/internal/protocol/generate.go index 7c9163a..5172315 100644 --- a/internal/protocol/generate.go +++ b/internal/protocol/generate.go @@ -13,13 +13,11 @@ package main import ( "bytes" - "cmp" "encoding/json" "flag" "fmt" "go/format" "io" - "iter" "log" "net/http" "os" @@ -28,6 +26,7 @@ import ( "strings" "github.com/modelcontextprotocol/go-sdk/internal/jsonschema" + "github.com/modelcontextprotocol/go-sdk/internal/util" ) var schemaFile = flag.String("schema_file", "", "if set, use this file as the persistent schema file") @@ -54,31 +53,36 @@ var declarations = config{ "CallToolRequest": { Fields: config{"Params": {Name: "CallToolParams"}}, }, - "CallToolResult": { - Name: "CallToolResult", - }, + "CallToolResult": {Name: "CallToolResult"}, "CancelledNotification": { Fields: config{"Params": {Name: "CancelledParams"}}, }, "ClientCapabilities": {Name: "ClientCapabilities"}, - "Implementation": {Name: "Implementation"}, + "GetPromptRequest": { + Fields: config{"Params": {Name: "GetPromptParams"}}, + }, + "GetPromptResult": {Name: "GetPromptResult"}, + "Implementation": {Name: "Implementation"}, "InitializeRequest": { Fields: config{"Params": {Name: "InitializeParams"}}, }, - "InitializeResult": { - Name: "InitializeResult", - }, + "InitializeResult": {Name: "InitializeResult"}, "InitializedNotification": { Fields: config{"Params": {Name: "InitializedParams"}}, }, + "ListPromptsRequest": { + Fields: config{"Params": {Name: "ListPromptsParams"}}, + }, + "ListPromptsResult": {Name: "ListPromptsResult"}, "ListToolsRequest": { Fields: config{"Params": {Name: "ListToolsParams"}}, }, - "ListToolsResult": { - Name: "ListToolsResult", - }, - "RequestId": {Substitute: "any"}, // null|number|string - "Role": {Name: "Role"}, + "ListToolsResult": {Name: "ListToolsResult"}, + "Prompt": {Name: "Prompt"}, + "PromptMessage": {Name: "PromptMessage"}, + "PromptArgument": {Name: "PromptArgument"}, + "RequestId": {Substitute: "any"}, // null|number|string + "Role": {Name: "Role"}, "ServerCapabilities": { Name: "ServerCapabilities", Fields: config{ @@ -92,9 +96,7 @@ var declarations = config{ Name: "Tool", Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}}, }, - "ToolAnnotations": { - Name: "ToolAnnotations", - }, + "ToolAnnotations": {Name: "ToolAnnotations"}, } func main() { @@ -114,7 +116,7 @@ func main() { // writing types, we collect definitions and concatenate them later. This // also allows us to sort. named := make(map[string]*bytes.Buffer) - for name, def := range sorted(schema.Definitions) { + for name, def := range util.Sorted(schema.Definitions) { config := declarations[name] if config == nil { continue @@ -142,7 +144,7 @@ import ( `) // Write out types. - for _, b := range sorted(named) { + for _, b := range util.Sorted(named) { fmt.Fprintln(buf) fmt.Fprint(buf, b.String()) } @@ -242,8 +244,8 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma // unmarshal them into a map[string]any, or delay unmarshalling with // json.RawMessage. For now, use json.RawMessage as it defers the choice. if def.Type == "object" && canHaveAdditionalProperties(def) { - w.Write([]byte("json.RawMessage")) - return nil + w.Write([]byte("map[string]")) + return writeType(w, nil, def.AdditionalProperties, named) } if def.Type == "" { @@ -269,7 +271,7 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma case "object": fmt.Fprintf(w, "struct {\n") - for name, fieldDef := range sorted(def.Properties) { + for name, fieldDef := range util.Sorted(def.Properties) { if fieldDef.Description != "" { fmt.Fprintf(w, "%s\n", toComment(fieldDef.Description)) } @@ -385,28 +387,3 @@ func assert(cond bool, msg string) { panic(msg) } } - -// Helpers below are copied from gopls' moremaps package. - -// sorted returns an iterator over the entries of m in key order. -func sorted[M ~map[K]V, K cmp.Ordered, V any](m M) iter.Seq2[K, V] { - // TODO(adonovan): use maps.Sorted if proposal #68598 is accepted. - return func(yield func(K, V) bool) { - keys := keySlice(m) - slices.Sort(keys) - for _, k := range keys { - if !yield(k, m[k]) { - break - } - } - } -} - -// keySlice returns the keys of the map M, like slices.Collect(maps.Keys(m)). -func keySlice[M ~map[K]V, K comparable, V any](m M) []K { - r := make([]K, 0, len(m)) - for k := range m { - r = append(r, k) - } - return r -} diff --git a/internal/protocol/protocol.go b/internal/protocol/protocol.go index 5b5fdfe..1af7cff 100644 --- a/internal/protocol/protocol.go +++ b/internal/protocol/protocol.go @@ -29,8 +29,8 @@ type Annotations struct { } type CallToolParams struct { - Arguments json.RawMessage `json:"arguments,omitempty"` - Name string `json:"name"` + Arguments map[string]json.RawMessage `json:"arguments,omitempty"` + Name string `json:"name"` } // The server's response to a tool call. @@ -46,8 +46,8 @@ type CallToolParams struct { type CallToolResult struct { // This result property is reserved by the protocol to allow clients and servers // to attach additional metadata to their responses. - Meta json.RawMessage `json:"_meta,omitempty"` - Content []json.RawMessage `json:"content"` + Meta map[string]json.RawMessage `json:"_meta,omitempty"` + Content []json.RawMessage `json:"content"` // Whether the tool call ended in an error. // // If not set, this is assumed to be false (the call was successful). @@ -70,14 +70,31 @@ type CancelledParams struct { // additional capabilities. type ClientCapabilities struct { // Experimental, non-standard capabilities that the client supports. - Experimental json.RawMessage `json:"experimental,omitempty"` + Experimental map[string]map[string]json.RawMessage `json:"experimental,omitempty"` // Present if the client supports listing roots. Roots *struct { // Whether the client supports notifications for changes to the roots list. ListChanged bool `json:"listChanged,omitempty"` } `json:"roots,omitempty"` // Present if the client supports sampling from an LLM. - Sampling json.RawMessage `json:"sampling,omitempty"` + Sampling map[string]json.RawMessage `json:"sampling,omitempty"` +} + +type GetPromptParams struct { + // 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"` +} + +// 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"` + // An optional description for the prompt. + Description string `json:"description,omitempty"` + Messages []PromptMessage `json:"messages"` } // Describes the name and version of an MCP implementation. @@ -99,8 +116,8 @@ type InitializeParams struct { type InitializeResult struct { // This result property is reserved by the protocol to allow clients and servers // to attach additional metadata to their responses. - Meta json.RawMessage `json:"_meta,omitempty"` - Capabilities ServerCapabilities `json:"capabilities"` + Meta map[string]json.RawMessage `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 @@ -114,7 +131,24 @@ type InitializeResult struct { ServerInfo Implementation `json:"serverInfo"` } -type InitializedParams json.RawMessage +type InitializedParams map[string]json.RawMessage + +type ListPromptsParams struct { + // An opaque token representing the current pagination position. If provided, + // the server should return results starting after this cursor. + Cursor string `json:"cursor,omitempty"` +} + +// 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"` + // 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"` +} type ListToolsParams struct { // An opaque token representing the current pagination position. If provided, @@ -126,19 +160,48 @@ type ListToolsParams struct { type ListToolsResult struct { // This result property is reserved by the protocol to allow clients and servers // to attach additional metadata to their responses. - Meta json.RawMessage `json:"_meta,omitempty"` + Meta map[string]json.RawMessage `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"` } +// A prompt or prompt template that the server offers. +type Prompt struct { + // A list of arguments to use for templating the prompt. + Arguments []PromptArgument `json:"arguments,omitempty"` + // An optional description of what this prompt provides + Description string `json:"description,omitempty"` + // The name of the prompt or prompt template. + Name string `json:"name"` +} + +// Describes an argument that a prompt can accept. +type PromptArgument struct { + // A human-readable description of the argument. + Description string `json:"description,omitempty"` + // The name of the argument. + Name string `json:"name"` + // Whether this argument must be provided. + Required bool `json:"required,omitempty"` +} + // Present if the server offers any prompt templates. type PromptCapabilities struct { // Whether this server supports notifications for changes to the prompt list. ListChanged bool `json:"listChanged,omitempty"` } +// Describes a message returned as part of a prompt. +// +// This is similar to `SamplingMessage`, but also supports the embedding of +// resources from the MCP server. +type PromptMessage struct { + Content json.RawMessage `json:"content"` + Role Role `json:"role"` +} + // Present if the server offers any resources to read. type ResourceCapabilities struct { // Whether this server supports notifications for changes to the resource list. @@ -155,11 +218,11 @@ type Role string // additional capabilities. type ServerCapabilities struct { // Present if the server supports argument autocompletion suggestions. - Completions json.RawMessage `json:"completions,omitempty"` + Completions map[string]json.RawMessage `json:"completions,omitempty"` // Experimental, non-standard capabilities that the server supports. - Experimental json.RawMessage `json:"experimental,omitempty"` + Experimental map[string]map[string]json.RawMessage `json:"experimental,omitempty"` // Present if the server supports sending log messages to the client. - Logging json.RawMessage `json:"logging,omitempty"` + Logging map[string]json.RawMessage `json:"logging,omitempty"` // Present if the server offers any prompt templates. Prompts *PromptCapabilities `json:"prompts,omitempty"` // Present if the server offers any resources to read. diff --git a/internal/util/util.go b/internal/util/util.go new file mode 100644 index 0000000..03a5dbb --- /dev/null +++ b/internal/util/util.go @@ -0,0 +1,36 @@ +// 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 util + +import ( + "cmp" + "iter" + "slices" +) + +// Helpers below are copied from gopls' moremaps package. + +// sorted returns an iterator over the entries of m in key order. +func Sorted[M ~map[K]V, K cmp.Ordered, V any](m M) iter.Seq2[K, V] { + // TODO(adonovan): use maps.Sorted if proposal #68598 is accepted. + return func(yield func(K, V) bool) { + keys := KeySlice(m) + slices.Sort(keys) + for _, k := range keys { + if !yield(k, m[k]) { + break + } + } + } +} + +// keySlice returns the keys of the map M, like slices.Collect(maps.Keys(m)). +func KeySlice[M ~map[K]V, K comparable, V any](m M) []K { + r := make([]K, 0, len(m)) + for k := range m { + r = append(r, k) + } + return r +} diff --git a/mcp/client.go b/mcp/client.go index 1700561..540715e 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -150,6 +150,33 @@ func (sc *ServerConnection) Ping(ctx context.Context) error { return call(ctx, sc.conn, "ping", nil, nil) } +// ListPrompts lists prompts that are currently available on the server. +func (sc *ServerConnection) ListPrompts(ctx context.Context) ([]protocol.Prompt, error) { + var ( + params = &protocol.ListPromptsParams{} + result protocol.ListPromptsResult + ) + if err := call(ctx, sc.conn, "prompts/list", params, &result); err != nil { + return nil, err + } + return result.Prompts, nil +} + +// GetPrompt gets a prompt from the server. +func (sc *ServerConnection) GetPrompt(ctx context.Context, name string, args map[string]string) (*protocol.GetPromptResult, error) { + var ( + params = &protocol.GetPromptParams{ + Name: name, + Arguments: args, + } + result = &protocol.GetPromptResult{} + ) + if err := call(ctx, sc.conn, "prompts/get", params, result); err != nil { + return nil, err + } + return result, nil +} + // ListTools lists tools that are currently available on the server. func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, error) { var ( @@ -164,23 +191,27 @@ func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, err // CallTool calls the tool with the given name and arguments. // -// TODO: make the following true: +// TODO(jba): make the following true: // If the provided arguments do not conform to the schema for the given tool, // the call fails. -func (sc *ServerConnection) CallTool(ctx context.Context, name string, args any) (_ []Content, err error) { +func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[string]any) (_ []Content, err error) { defer func() { if err != nil { err = fmt.Errorf("calling tool %q: %w", name, err) } }() - argJSON, err := json.Marshal(args) - if err != nil { - return nil, fmt.Errorf("marshaling args: %v", err) + argsJSON := make(map[string]json.RawMessage) + for name, arg := range args { + argJSON, err := json.Marshal(arg) + if err != nil { + return nil, fmt.Errorf("marshaling argument %s: %v", name, err) + } + argsJSON[name] = argJSON } var ( params = &protocol.CallToolParams{ Name: name, - Arguments: argJSON, + Arguments: argsJSON, } result protocol.CallToolResult ) diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index d989eb1..7d6215d 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -53,7 +53,7 @@ func TestCmdTransport(t *testing.T) { if err != nil { log.Fatal(err) } - got, err := serverConn.CallTool(ctx, "greet", SayHiParams{Name: "user"}) + got, err := serverConn.CallTool(ctx, "greet", map[string]any{"name": "user"}) if err != nil { log.Fatal(err) } diff --git a/mcp/mcp.go b/mcp/mcp.go index c278234..86d05ee 100644 --- a/mcp/mcp.go +++ b/mcp/mcp.go @@ -63,8 +63,12 @@ // // # TODO // +// - Support all content types. // - Support pagination. +// - Support completion. +// - Support oauth. // - Support all client/server operations. +// - Pass the client connection in the context. // - Support streamable HTTP transport. // - Support multiple versions of the spec. // - Implement full JSON schema support, with both client-side and diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index c5d8f88..418ba09 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -6,6 +6,7 @@ package mcp import ( "context" + "encoding/json" "errors" "fmt" "slices" @@ -41,9 +42,34 @@ func TestEndToEnd(t *testing.T) { // The 'fail' tool returns this error. failure := errors.New("mcp failure") - s.AddTools(MakeTool("fail", "just fail", func(context.Context, *ClientConnection, struct{}) ([]Content, error) { - return nil, failure - })) + s.AddTools( + MakeTool("fail", "just fail", func(context.Context, *ClientConnection, struct{}) ([]Content, error) { + return nil, failure + }), + ) + + s.AddPrompts( + MakePrompt("code_review", "do a code review", func(_ context.Context, _ *ClientConnection, params struct{ Code string }) (*protocol.GetPromptResult, error) { + // TODO(rfindley): clean up this handling of content. + content, err := json.Marshal(protocol.TextContent{ + Type: "text", + Text: "Please review the following code: " + params.Code, + }) + if err != nil { + return nil, err + } + return &protocol.GetPromptResult{ + Description: "Code review prompt", + Messages: []protocol.PromptMessage{ + // TODO: move 'Content' to the protocol package. + {Role: "user", Content: json.RawMessage(content)}, + }, + }, nil + }), + MakePrompt("fail", "", func(_ context.Context, _ *ClientConnection, params struct{}) (*protocol.GetPromptResult, error) { + return nil, failure + }), + ) // Connect the server. cc, err := s.Connect(ctx, st, nil) @@ -80,6 +106,34 @@ func TestEndToEnd(t *testing.T) { t.Fatalf("ping failed: %v", err) } + gotPrompts, err := sc.ListPrompts(ctx) + if err != nil { + t.Errorf("prompts/list failed: %v", err) + } + wantPrompts := []protocol.Prompt{ + { + Name: "code_review", + Description: "do a code review", + Arguments: []protocol.PromptArgument{{Name: "Code", Required: true}}, + }, + {Name: "fail"}, + } + if diff := cmp.Diff(wantPrompts, gotPrompts); diff != "" { + t.Fatalf("prompts/list mismatch (-want +got):\n%s", diff) + } + + gotReview, err := sc.GetPrompt(ctx, "code_review", map[string]string{"Code": "1+1"}) + if err != nil { + t.Fatal(err) + } + // TODO: assert on the full review, once content is easier to create. + if got, want := gotReview.Description, "Code review prompt"; got != want { + t.Errorf("prompts/get 'code_review': got description %q, want %q", got, want) + } + if _, err := sc.GetPrompt(ctx, "fail", map[string]string{}); err == nil || !strings.Contains(err.Error(), failure.Error()) { + t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) + } + gotTools, err := sc.ListTools(ctx) if err != nil { t.Errorf("tools/list failed: %v", err) @@ -107,7 +161,7 @@ func TestEndToEnd(t *testing.T) { t.Fatalf("tools/list mismatch (-want +got):\n%s", diff) } - gotHi, err := sc.CallTool(ctx, "greet", hiParams{"user"}) + gotHi, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"}) if err != nil { t.Fatal(err) } @@ -116,7 +170,7 @@ func TestEndToEnd(t *testing.T) { t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) } - if _, err := sc.CallTool(ctx, "fail", struct{}{}); err == nil || !strings.Contains(err.Error(), failure.Error()) { + if _, err := sc.CallTool(ctx, "fail", map[string]any{}); err == nil || !strings.Contains(err.Error(), failure.Error()) { t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) } @@ -175,12 +229,12 @@ func TestServerClosing(t *testing.T) { } wg.Done() }() - if _, err := sc.CallTool(ctx, "greet", hiParams{"user"}); err != nil { + if _, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"}); err != nil { t.Fatalf("after connecting: %v", err) } cc.Close() wg.Wait() - if _, err := sc.CallTool(ctx, "greet", hiParams{"user"}); !errors.Is(err, ErrConnectionClosed) { + if _, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"}); !errors.Is(err, ErrConnectionClosed) { t.Errorf("after disconnection, got error %v, want EOF", err) } } @@ -244,7 +298,7 @@ func TestCancellation(t *testing.T) { defer sc.Close() ctx, cancel := context.WithCancel(context.Background()) - go sc.CallTool(ctx, "slow", struct{}{}) + go sc.CallTool(ctx, "slow", map[string]any{}) <-start cancel() select { diff --git a/mcp/prompt.go b/mcp/prompt.go new file mode 100644 index 0000000..7ccf646 --- /dev/null +++ b/mcp/prompt.go @@ -0,0 +1,122 @@ +// 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 ( + "context" + "encoding/json" + "fmt" + "reflect" + "slices" + + "github.com/modelcontextprotocol/go-sdk/internal/jsonschema" + "github.com/modelcontextprotocol/go-sdk/internal/protocol" + "github.com/modelcontextprotocol/go-sdk/internal/util" +) + +// A PromptHandler handles a call to prompts/get. +type PromptHandler func(context.Context, *ClientConnection, map[string]string) (*protocol.GetPromptResult, error) + +// A Prompt is a prompt definition bound to a prompt handler. +type Prompt struct { + Definition protocol.Prompt + Handler PromptHandler +} + +// MakePrompt is a helper to use reflection to create a prompt for the given +// handler. +// +// The arguments for the prompt are extracted from the request type for the +// handler. The handler request type must be a struct consisting only of fields +// of type string or *string. The argument names for the resulting prompt +// definition correspond to the JSON names of the request fields, and any +// fields that are not marked "omitempty" are considered required. +func MakePrompt[TReq any](name, description string, handler func(context.Context, *ClientConnection, TReq) (*protocol.GetPromptResult, error), opts ...PromptOption) *Prompt { + schema, err := jsonschema.For[TReq]() + if err != nil { + panic(err) + } + if schema.Type != "object" || !reflect.DeepEqual(schema.AdditionalProperties, &jsonschema.Schema{Not: &jsonschema.Schema{}}) { + panic(fmt.Sprintf("handler request type must be a struct")) + } + prompt := &Prompt{ + Definition: protocol.Prompt{ + Name: name, + Description: description, + }, + } + required := make(map[string]bool) + for _, p := range schema.Required { + required[p] = true + } + for name, prop := range util.Sorted(schema.Properties) { + if prop.Type != "string" { + panic(fmt.Sprintf("handler type must consist only of string fields")) + } + prompt.Definition.Arguments = append(prompt.Definition.Arguments, protocol.PromptArgument{ + Name: name, + Description: prop.Description, + Required: required[name], + }) + } + prompt.Handler = func(ctx context.Context, cc *ClientConnection, args map[string]string) (*protocol.GetPromptResult, error) { + // For simplicity, just marshal and unmarshal the arguments. + // This could be avoided in the future. + rawArgs, err := json.Marshal(args) + if err != nil { + return nil, err + } + var v TReq + if err := unmarshalSchema(rawArgs, schema, &v); err != nil { + return nil, err + } + return handler(ctx, cc, v) + } + for _, opt := range opts { + opt.set(prompt) + } + return prompt +} + +// A PromptOption configures the behavior of a Prompt. +type PromptOption interface { + set(*Prompt) +} + +type promptSetter func(*Prompt) + +func (s promptSetter) set(p *Prompt) { s(p) } + +// Argument configures the 'schema' of a prompt argument. +// If the argument does not exist, it is added. +// +// Since prompt arguments are not a full JSON schema, Argument only accepts +// Required and Description, and panics when encountering any other option. +func Argument(name string, opts ...SchemaOption) PromptOption { + return promptSetter(func(p *Prompt) { + i := slices.IndexFunc(p.Definition.Arguments, func(arg protocol.PromptArgument) bool { + return arg.Name == name + }) + var arg protocol.PromptArgument + if i < 0 { + i = len(p.Definition.Arguments) + arg = protocol.PromptArgument{Name: name} + p.Definition.Arguments = append(p.Definition.Arguments, arg) + } else { + arg = p.Definition.Arguments[i] + } + for _, opt := range opts { + switch v := opt.(type) { + case required: + arg.Required = bool(v) + case description: + arg.Description = string(v) + default: + panic(fmt.Sprintf("unsupported prompt argument schema option %T", opt)) + } + } + p.Definition.Arguments[i] = arg + }) +} diff --git a/mcp/prompt_test.go b/mcp/prompt_test.go new file mode 100644 index 0000000..ca7f48d --- /dev/null +++ b/mcp/prompt_test.go @@ -0,0 +1,54 @@ +// 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_test + +import ( + "context" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/internal/protocol" +) + +// testPromptHandler is used for type inference in TestMakePrompt. +func testPromptHandler[T any](context.Context, *mcp.ClientConnection, T) (*protocol.GetPromptResult, error) { + panic("not implemented") +} + +func TestMakePrompt(t *testing.T) { + tests := []struct { + prompt *mcp.Prompt + want []protocol.PromptArgument + }{ + { + mcp.MakePrompt("empty", "", testPromptHandler[struct{}]), + nil, + }, + { + mcp.MakePrompt("add_arg", "", testPromptHandler[struct{}], mcp.Argument("x")), + []protocol.PromptArgument{{Name: "x"}}, + }, + { + mcp.MakePrompt("combo", "", testPromptHandler[struct { + Name string `json:"name"` + Country string `json:"country,omitempty"` + State string + }], + mcp.Argument("name", mcp.Description("the person's name")), + mcp.Argument("State", mcp.Required(false))), + []protocol.PromptArgument{ + {Name: "State"}, + {Name: "country"}, + {Name: "name", Required: true, Description: "the person's name"}, + }, + }, + } + for _, test := range tests { + if diff := cmp.Diff(test.want, test.prompt.Definition.Arguments); diff != "" { + t.Errorf("MakePrompt(%v) mismatch (-want +got):\n%s", test.prompt.Definition.Name, diff) + } + } +} diff --git a/mcp/server.go b/mcp/server.go index e2a4c27..6c590ab 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -26,6 +26,7 @@ type Server struct { opts ServerOptions mu sync.Mutex + prompts []*Prompt tools []*Tool clients []*ClientConnection } @@ -53,6 +54,15 @@ func NewServer(name, version string, opts *ServerOptions) *Server { } } +// AddPrompts adds the given prompts to the server. +// +// TODO(rfindley): notify connected clients of any changes. +func (s *Server) AddPrompts(prompts ...*Prompt) { + s.mu.Lock() + defer s.mu.Unlock() + s.prompts = append(s.prompts, prompts...) +} + // AddTools adds the given tools to the server. // // TODO(rfindley): notify connected clients of any changes. @@ -71,6 +81,33 @@ func (s *Server) Clients() iter.Seq[*ClientConnection] { return slices.Values(clients) } +func (s *Server) listPrompts(_ context.Context, _ *ClientConnection, params *protocol.ListPromptsParams) (*protocol.ListPromptsResult, error) { + s.mu.Lock() + defer s.mu.Unlock() + + res := new(protocol.ListPromptsResult) + for _, p := range s.prompts { + res.Prompts = append(res.Prompts, p.Definition) + } + return res, nil +} + +func (s *Server) getPrompt(ctx context.Context, cc *ClientConnection, params *protocol.GetPromptParams) (*protocol.GetPromptResult, error) { + s.mu.Lock() + var prompt *Prompt + if i := slices.IndexFunc(s.prompts, func(t *Prompt) bool { + return t.Definition.Name == params.Name + }); i >= 0 { + prompt = s.prompts[i] + } + s.mu.Unlock() + + if prompt == nil { + return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name) + } + return prompt.Handler(ctx, cc, params.Arguments) +} + func (s *Server) listTools(_ context.Context, _ *ClientConnection, params *protocol.ListToolsParams) (*protocol.ListToolsResult, error) { s.mu.Lock() defer s.mu.Unlock() @@ -185,6 +222,12 @@ func (cc *ClientConnection) handle(ctx context.Context, req *jsonrpc2.Request) ( // The spec says that 'ping' expects an empty object result. return struct{}{}, nil + case "prompts/list": + return dispatch(ctx, cc, req, cc.server.listPrompts) + + case "prompts/get": + return dispatch(ctx, cc, req, cc.server.getPrompt) + case "tools/list": return dispatch(ctx, cc, req, cc.server.listTools) @@ -216,8 +259,11 @@ func (cc *ClientConnection) initialize(ctx context.Context, _ *ClientConnection, // TODO(rfindley): support multiple protocol versions. ProtocolVersion: "2024-11-05", Capabilities: protocol.ServerCapabilities{ + Prompts: &protocol.PromptCapabilities{ + ListChanged: false, // not yet supported + }, Tools: &protocol.ToolCapabilities{ - ListChanged: true, + ListChanged: false, // not yet supported }, }, Instructions: cc.server.opts.Instructions, @@ -249,5 +295,11 @@ func dispatch[TConn, TParams, TResult any](ctx context.Context, conn TConn, req if err := json.Unmarshal(req.Params, ¶ms); err != nil { return nil, err } - return f(ctx, conn, params) + // Important: avoid returning a typed nil, as it can't be handled by the + // jsonrpc2 package. + res, err := f(ctx, conn, params) + if err != nil { + return nil, err + } + return res, nil } diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 98e87c0..63826b5 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -40,7 +40,7 @@ func ExampleServer() { log.Fatal(err) } - content, err := serverConnection.CallTool(ctx, "greet", SayHiParams{Name: "user"}) + content, err := serverConnection.CallTool(ctx, "greet", map[string]any{"name": "user"}) if err != nil { log.Fatal(err) } diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index c781f84..baf6dc8 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -40,7 +40,7 @@ func ExampleSSEHandler() { } defer serverConn.Close() - content, err := serverConn.CallTool(ctx, "add", AddParams{1, 2}) + content, err := serverConn.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2}) if err != nil { log.Fatal(err) } diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 7f0da7d..52135c8 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -44,7 +44,7 @@ func TestSSEServer(t *testing.T) { t.Fatal(err) } cc := <-clients - gotHi, err := sc.CallTool(ctx, "greet", hiParams{"user"}) + gotHi, err := sc.CallTool(ctx, "greet", map[string]any{"name": "user"}) if err != nil { t.Fatal(err) } diff --git a/mcp/tool.go b/mcp/tool.go index ed37cae..3de32d5 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -14,7 +14,7 @@ import ( ) // A ToolHandler handles a call to tools/call. -type ToolHandler func(context.Context, *ClientConnection, json.RawMessage) (*protocol.CallToolResult, error) +type ToolHandler func(context.Context, *ClientConnection, map[string]json.RawMessage) (*protocol.CallToolResult, error) // A Tool is a tool definition that is bound to a tool handler. type Tool struct { @@ -30,20 +30,30 @@ type Tool struct { // handler, and used to unmmarshal and validate requests to the handler. This // schema may be customized using the [Input] option. // -// It is the caller's responsibility that the handler request type can produce -// a valid schema, as documented by [jsonschema.ForType]; otherwise, MakeTool -// panics. +// The handler request type must translate to a valid schema, as documented by +// [jsonschema.ForType]; otherwise, MakeTool panics. +// +// TODO: just have the handler return a CallToolResult: returning []Content is +// going to be inconsistent with other server features. func MakeTool[TReq any](name, description string, handler func(context.Context, *ClientConnection, TReq) ([]Content, error), opts ...ToolOption) *Tool { schema, err := jsonschema.For[TReq]() if err != nil { panic(err) } - wrapped := func(ctx context.Context, cc *ClientConnection, args json.RawMessage) (*protocol.CallToolResult, error) { + wrapped := func(ctx context.Context, cc *ClientConnection, args map[string]json.RawMessage) (*protocol.CallToolResult, error) { + // For simplicity, just marshal and unmarshal the arguments. + // This could be avoided in the future. + rawArgs, err := json.Marshal(args) + if err != nil { + return nil, err + } var v TReq - if err := unmarshalSchema(args, schema, &v); err != nil { + if err := unmarshalSchema(rawArgs, schema, &v); err != nil { return nil, err } content, err := handler(ctx, cc, v) + // TODO: investigate why server errors are embedded in this strange way, + // rather than returned as jsonrpc2 server errors. if err != nil { return &protocol.CallToolResult{ Content: marshalContent([]Content{TextContent{Text: err.Error()}}), @@ -140,6 +150,8 @@ func Required(v bool) SchemaOption { return required(v) } +// required must be a distinguished type as it needs special handling to mutate +// the parent schema, and to mutate prompt arguments. type required bool func (required) set(s *jsonschema.Schema) { @@ -154,10 +166,16 @@ func Enum(values ...any) SchemaOption { } // Description sets the provided schema description. -func Description(description string) SchemaOption { - return schemaSetter(func(schema *jsonschema.Schema) { - schema.Description = description - }) +func Description(desc string) SchemaOption { + return description(desc) +} + +// description must be a distinguished type so that it can be handled by prompt +// options. +type description string + +func (d description) set(s *jsonschema.Schema) { + s.Description = string(d) } // Schema overrides the inferred schema with a shallow copy of the given diff --git a/mcp/tool_test.go b/mcp/tool_test.go index 36dd8f4..a657fc0 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -13,8 +13,8 @@ import ( "github.com/modelcontextprotocol/go-sdk/internal/jsonschema" ) -// testHandler is used for type inference in TestMakeTool. -func testHandler[T any](context.Context, *mcp.ClientConnection, T) ([]mcp.Content, error) { +// testToolHandler is used for type inference in TestMakeTool. +func testToolHandler[T any](context.Context, *mcp.ClientConnection, T) ([]mcp.Content, error) { panic("not implemented") } @@ -24,7 +24,7 @@ func TestMakeTool(t *testing.T) { want *jsonschema.Schema }{ { - mcp.MakeTool("basic", "", testHandler[struct { + mcp.MakeTool("basic", "", testToolHandler[struct { Name string `json:"name"` }]), &jsonschema.Schema{ @@ -37,7 +37,7 @@ func TestMakeTool(t *testing.T) { }, }, { - mcp.MakeTool("enum", "", testHandler[struct{ Name string }], mcp.Input( + mcp.MakeTool("enum", "", testToolHandler[struct{ Name string }], mcp.Input( mcp.Property("Name", mcp.Enum("x", "y", "z")), )), &jsonschema.Schema{ @@ -50,7 +50,7 @@ func TestMakeTool(t *testing.T) { }, }, { - mcp.MakeTool("required", "", testHandler[struct { + mcp.MakeTool("required", "", testToolHandler[struct { Name string `json:"name"` Language string `json:"language"` X int `json:"x,omitempty"` @@ -70,7 +70,7 @@ func TestMakeTool(t *testing.T) { }, }, { - mcp.MakeTool("set_schema", "", testHandler[struct { + mcp.MakeTool("set_schema", "", testToolHandler[struct { X int `json:"x,omitempty"` Y int `json:"y,omitempty"` }], mcp.Input(