internal/mcp: future-proof prompt handling

Pass GetPromptParams as needed, in case that type ever gets additional
properties.

Change-Id: I83fa37bb2af5633a8c2bcdc47a85c3030dfe70fb
Reviewed-on: https://go-review.googlesource.com/c/tools/+/674895
Auto-Submit: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

examples/hello/main.go

@@ -16,17 +16,17 @@ import (
 
 var httpAddr = flag.String("http", "", "if set, use SSE HTTP at this address, instead of stdin/stdout")
 
-type HiParams struct {
+type HiArgs struct {
 	Name string `json:"name"`
 }
 
-func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) ([]*mcp.Content, error) {
+func SayHi(ctx context.Context, cc *mcp.ServerSession, params *HiArgs) ([]*mcp.Content, error) {
 	return []*mcp.Content{
 		mcp.NewTextContent("Hi " + params.Name),
 	}, nil
 }
 
-func PromptHi(ctx context.Context, cc *mcp.ServerSession, params *HiParams) (*mcp.GetPromptResult, error) {
+func PromptHi(ctx context.Context, cc *mcp.ServerSession, args *HiArgs, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
 	return &mcp.GetPromptResult{
 		Description: "Code review prompt",
 		Messages: []*mcp.PromptMessage{

mcp/mcp_test.go

@@ -73,15 +73,16 @@ func TestEndToEnd(t *testing.T) {
 	)
 
 	s.AddPrompts(
-		NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerSession, params struct{ Code string }) (*GetPromptResult, error) {
-			return &GetPromptResult{
-				Description: "Code review prompt",
-				Messages: []*PromptMessage{
-					{Role: "user", Content: NewTextContent("Please review the following code: " + params.Code)},
-				},
-			}, nil
-		}),
-		NewPrompt("fail", "", func(_ context.Context, _ *ServerSession, params struct{}) (*GetPromptResult, error) {
+		NewPrompt("code_review", "do a code review",
+			func(_ context.Context, _ *ServerSession, params struct{ Code string }, _ *GetPromptParams) (*GetPromptResult, error) {
+				return &GetPromptResult{
+					Description: "Code review prompt",
+					Messages: []*PromptMessage{
+						{Role: "user", Content: NewTextContent("Please review the following code: " + params.Code)},
+					},
+				}, nil
+			}),
+		NewPrompt("fail", "", func(_ context.Context, _ *ServerSession, args struct{}, _ *GetPromptParams) (*GetPromptResult, error) {
 			return nil, failure
 		}),
 	)

mcp/prompt.go

@@ -16,7 +16,7 @@ import (
 )
 
 // A PromptHandler handles a call to prompts/get.
-type PromptHandler func(context.Context, *ServerSession, map[string]string) (*GetPromptResult, error)
+type PromptHandler func(context.Context, *ServerSession, *GetPromptParams) (*GetPromptResult, error)
 
 // A Prompt is a prompt definition bound to a prompt handler.
 type ServerPrompt struct {
@@ -24,15 +24,17 @@ type ServerPrompt struct {
 	Handler PromptHandler
 }
 
-// NewPrompt is a helper to use reflection to create a prompt for the given
-// handler.
+// NewPrompt is a helper that uses 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 NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
+//
+// The handler is passed [GetPromptParams] so it can have access to prompt parameters other than name and arguments.
+// At present, there are no such parameters.
+func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq, *GetPromptParams) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
 	schema, err := jsonschema.For[TReq]()
 	if err != nil {
 		panic(err)
@@ -60,10 +62,10 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
 			Required:    required[name],
 		})
 	}
-	prompt.Handler = func(ctx context.Context, cc *ServerSession, args map[string]string) (*GetPromptResult, error) {
+	prompt.Handler = func(ctx context.Context, ss *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
 		// For simplicity, just marshal and unmarshal the arguments.
 		// This could be avoided in the future.
-		rawArgs, err := json.Marshal(args)
+		rawArgs, err := json.Marshal(params.Arguments)
 		if err != nil {
 			return nil, err
 		}
@@ -71,7 +73,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
 		if err := unmarshalSchema(rawArgs, schema, &v); err != nil {
 			return nil, err
 		}
-		return handler(ctx, cc, v)
+		return handler(ctx, ss, v, params)
 	}
 	for _, opt := range opts {
 		opt.set(prompt)

mcp/prompt_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 // testPromptHandler is used for type inference in TestNewPrompt.
-func testPromptHandler[T any](context.Context, *mcp.ServerSession, T) (*mcp.GetPromptResult, error) {
+func testPromptHandler[T any](context.Context, *mcp.ServerSession, T, *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
 	panic("not implemented")
 }
 

mcp/server.go

@@ -180,7 +180,7 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPr
 		// TODO: surface the error code over the wire, instead of flattening it into the string.
 		return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name)
 	}
-	return prompt.Handler(ctx, cc, params.Arguments)
+	return prompt.Handler(ctx, cc, params)
 }
 
 func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) {