From b1780319e47aaf6a74edb373ffbc5244ac15391c Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Mon, 19 May 2025 10:16:11 -0400 Subject: [PATCH] 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 Reviewed-by: Sam Thanawalla LUCI-TryBot-Result: Go LUCI --- examples/hello/main.go | 6 +++--- mcp/mcp_test.go | 19 ++++++++++--------- mcp/prompt.go | 16 +++++++++------- mcp/prompt_test.go | 2 +- mcp/server.go | 2 +- 5 files changed, 24 insertions(+), 21 deletions(-) diff --git a/examples/hello/main.go b/examples/hello/main.go index 40b4b22..aa76da7 100644 --- a/examples/hello/main.go +++ b/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{ diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 6de25a4..b3b7c5f 100644 --- a/mcp/mcp_test.go +++ b/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 }), ) diff --git a/mcp/prompt.go b/mcp/prompt.go index b17dbfc..28265d9 100644 --- a/mcp/prompt.go +++ b/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) diff --git a/mcp/prompt_test.go b/mcp/prompt_test.go index b706ced..55a2b0c 100644 --- a/mcp/prompt_test.go +++ b/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") } diff --git a/mcp/server.go b/mcp/server.go index 40fa913..127b8b8 100644 --- a/mcp/server.go +++ b/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) {