internal/mcp: support prompts

Add support for 'prompts/list' and 'prompts/get', and registering
prompts on the server using AddPrompts. Add a 'MakePrompt' helper to
construct prompts from a prompt handler using reflection.

Change-Id: I479cb9c9b99313cca99640c1f5ea1939363759e6
Reviewed-on: https://go-review.googlesource.com/c/tools/+/669015
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

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

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
-}

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.

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
+}