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 <[email protected]>
Reviewed-by: Jonathan Amsterdam <[email protected]>

Author: findleyr

internal/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
 	)

internal/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)
 	}

internal/mcp/examples/hello/main.go

@@ -6,33 +6,54 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"net/http"
 	"os"
 
 	"golang.org/x/tools/internal/mcp"
+	"golang.org/x/tools/internal/mcp/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()
 
 	server := mcp.NewServer("greeter", "v0.0.1", nil)
 	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 {

internal/mcp/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/mcp/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"
 
 	"golang.org/x/tools/internal/mcp/internal/jsonschema"
+	"golang.org/x/tools/internal/mcp/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
-}