feat: autocompletion for prompt args and resource template URI args

mcp/prompts.go

@@ -68,6 +68,8 @@ type PromptArgument struct {
 	// Whether this argument must be provided.
 	// If true, clients must include this argument when calling prompts/get.
 	Required bool `json:"required,omitempty"`
+	// Optional CompletionHandlerFunc for autocompleting the argument value.
+	CompletionHandler *CompletionHandlerFunc `json:"-"`
 }
 
 // Role represents the sender or recipient of messages and data in a
@@ -168,3 +170,10 @@ func RequiredArgument() ArgumentOption {
 		arg.Required = true
 	}
 }
+
+// ArgumentCompletion configures an autocomplete handler for the argument.
+func ArgumentCompletion(handler CompletionHandlerFunc) ArgumentOption {
+	return func(arg *PromptArgument) {
+		arg.CompletionHandler = &handler
+	}
+}

mcp/resources.go

@@ -97,3 +97,14 @@ func WithTemplateAnnotations(audience []Role, priority float64) ResourceTemplate
 		t.Annotations.Priority = priority
 	}
 }
+
+// WithTemplateArgumentCompletion adds an autocomplete handler for the specified argument of ResourceTemplate.
+// The argument should be one of the variables referenced by the URI template.
+func WithTemplateArgumentCompletion(argument string, handler CompletionHandlerFunc) ResourceTemplateOption {
+	return func(t *ResourceTemplate) {
+		if t.URITemplate.ArgumentCompletionHandlers == nil {
+			t.URITemplate.ArgumentCompletionHandlers = make(map[string]CompletionHandlerFunc)
+		}
+		t.URITemplate.ArgumentCompletionHandlers[argument] = handler
+	}
+}

mcp/types.go

@@ -3,6 +3,7 @@
 package mcp
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"maps"
@@ -50,6 +51,10 @@ const (
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/tools/
 	MethodToolsCall MCPMethod = "tools/call"
 
+	// MethodCompletion provides autocompletion suggestions for URI arguments.
+	// https://modelcontextprotocol.io/specification/draft/server/utilities/completion
+	MethodCompletion MCPMethod = "completion/complete"
+
 	// MethodSetLogLevel configures the minimum log level for client
 	// https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging
 	MethodSetLogLevel MCPMethod = "logging/setLevel"
@@ -69,8 +74,14 @@ const (
 	MethodNotificationToolsListChanged = "notifications/tools/list_changed"
 )
 
+// CompletionHandlerFunc handles completion requests.
+type CompletionHandlerFunc func(ctx context.Context, request CompleteRequest) (*CompleteResult, error)
+
 type URITemplate struct {
 	*uritemplate.Template
+
+	// Optional mapping of URI template arguments to CompletionHandlerFunc for autocompleting each argument's value.
+	ArgumentCompletionHandlers map[string]CompletionHandlerFunc `json:"-"`
 }
 
 func (t *URITemplate) MarshalJSON() ([]byte, error) {
@@ -456,6 +467,8 @@ type ServerCapabilities struct {
 	Experimental map[string]any `json:"experimental,omitempty"`
 	// Present if the server supports sending log messages to the client.
 	Logging *struct{} `json:"logging,omitempty"`
+	// Present if the server supports autocompletion for prompt and resource template arguments.
+	Completion *struct{} `json:"completion,omitempty"`
 	// Present if the server offers any prompt templates.
 	Prompts *struct {
 		// Whether this server supports notifications for changes to the prompt list.
@@ -965,6 +978,10 @@ type CompleteParams struct {
 		// The value of the argument to use for completion matching.
 		Value string `json:"value"`
 	} `json:"argument"`
+	Context struct {
+		// Previously completed arguments for this reference.
+		Arguments map[string]string `json:"arguments,omitempty"`
+	} `json:"context,omitempty"`
 }
 
 // CompleteResult is the server's response to a completion/complete request
@@ -984,18 +1001,25 @@ type CompleteResult struct {
 
 // ResourceReference is a reference to a resource or resource template definition.
 type ResourceReference struct {
-	Type string `json:"type"`
+	Type RefType `json:"type"`
 	// The URI or URI template of the resource.
 	URI string `json:"uri"`
 }
 
 // PromptReference identifies a prompt.
 type PromptReference struct {
-	Type string `json:"type"`
+	Type RefType `json:"type"`
 	// The name of the prompt or prompt template
 	Name string `json:"name"`
 }
 
+type RefType string
+
+const (
+	RefTypeResource RefType = "ref/resource"
+	RefTypePrompt   RefType = "ref/prompt"
+)
+
 /* Roots */
 
 // ListRootsRequest is sent from the server to request a list of root URIs from the client. Roots allow

mcp/utils.go

@@ -514,6 +514,35 @@ func ParseContent(contentMap map[string]any) (Content, error) {
 	return nil, fmt.Errorf("unsupported content type: %s", contentType)
 }
 
+func ParseCompletionReference(req CompleteRequest) (*PromptReference, *ResourceReference, error) {
+	ref, ok := req.Params.Ref.(map[string]interface{})
+	if !ok {
+		return nil, nil, fmt.Errorf("params.ref must be a mapping")
+	}
+
+	refType, ok := ref["type"].(string)
+	if !ok {
+		return nil, nil, fmt.Errorf("params.ref.type must be a string")
+	}
+
+	switch RefType(refType) {
+	case RefTypeResource:
+		if uri, ok := ref["uri"].(string); !ok {
+			return nil, nil, fmt.Errorf("params.ref.uri must be a string")
+		} else {
+			return nil, &ResourceReference{Type: RefType(refType), URI: uri}, nil
+		}
+	case RefTypePrompt:
+		if name, ok := ref["name"].(string); !ok {
+			return nil, nil, fmt.Errorf("params.ref.name must be a string")
+		} else {
+			return &PromptReference{Type: RefType(refType), Name: name}, nil, nil
+		}
+	default:
+		return nil, nil, fmt.Errorf("unexpected value for params.ref.type: %s", refType)
+	}
+}
+
 func ParseGetPromptResult(rawMessage *json.RawMessage) (*GetPromptResult, error) {
 	if rawMessage == nil {
 		return nil, fmt.Errorf("response is nil")

server/internal/gen/data.go

@@ -107,5 +107,15 @@ var MCPRequestTypes = []MCPRequestType{
 		HookName:       "CallTool",
 		UnmarshalError: "invalid call tool request",
 		HandlerFunc:    "handleToolCall",
+	}, {
+		MethodName:     "MethodCompletion",
+		ParamType:      "CompleteRequest",
+		ResultType:     "CompleteResult",
+		Group:          "completions",
+		GroupName:      "Completions",
+		GroupHookName:  "Completion",
+		HookName:       "Complete",
+		UnmarshalError: "invalid completion request",
+		HandlerFunc:    "handleCompletion",
 	},
 }

server/server.go

@@ -171,10 +171,11 @@ func WithPaginationLimit(limit int) ServerOption {
 
 // serverCapabilities defines the supported features of the MCP server
 type serverCapabilities struct {
-	tools     *toolCapabilities
-	resources *resourceCapabilities
-	prompts   *promptCapabilities
-	logging   *bool
+	tools       *toolCapabilities
+	resources   *resourceCapabilities
+	prompts     *promptCapabilities
+	completions *bool
+	logging     *bool
 }
 
 // resourceCapabilities defines the supported resource-related features
@@ -374,6 +375,10 @@ func (s *MCPServer) AddResourceTemplate(
 	}
 	s.resourcesMu.Unlock()
 
+	if len(template.URITemplate.ArgumentCompletionHandlers) > 0 {
+		s.implicitlyRegisterCompletionCapabilities()
+	}
+
 	// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
 	if s.capabilities.resources.listChanged {
 		// Send notification to all initialized sessions
@@ -392,6 +397,15 @@ func (s *MCPServer) AddPrompts(prompts ...ServerPrompt) {
 	}
 	s.promptsMu.Unlock()
 
+	for _, entry := range prompts {
+		for _, arg := range entry.Prompt.Arguments {
+			if arg.CompletionHandler != nil {
+				s.implicitlyRegisterCompletionCapabilities()
+				break
+			}
+		}
+	}
+
 	// When the list of available prompts changes, servers that declared the listChanged capability SHOULD send a notification.
 	if s.capabilities.prompts.listChanged {
 		// Send notification to all initialized sessions
@@ -402,6 +416,12 @@ func (s *MCPServer) AddPrompts(prompts ...ServerPrompt) {
 // AddPrompt registers a new prompt handler with the given name
 func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
 	s.AddPrompts(ServerPrompt{Prompt: prompt, Handler: handler})
+	for _, arg := range prompt.Arguments {
+		if arg.CompletionHandler != nil {
+			s.implicitlyRegisterCompletionCapabilities()
+			break
+		}
+	}
 }
 
 // DeletePrompts removes prompts from the server
@@ -446,6 +466,13 @@ func (s *MCPServer) implicitlyRegisterResourceCapabilities() {
 	)
 }
 
+func (s *MCPServer) implicitlyRegisterCompletionCapabilities() {
+	s.implicitlyRegisterCapabilities(
+		func() bool { return s.capabilities.completions != nil },
+		func() { s.capabilities.completions = mcp.ToBoolPtr(true) },
+	)
+}
+
 func (s *MCPServer) implicitlyRegisterPromptCapabilities() {
 	s.implicitlyRegisterCapabilities(
 		func() bool { return s.capabilities.prompts != nil },
@@ -562,6 +589,10 @@ func (s *MCPServer) handleInitialize(
 		capabilities.Logging = &struct{}{}
 	}
 
+	if s.capabilities.completions != nil && *s.capabilities.completions {
+		capabilities.Completion = &struct{}{}
+	}
+
 	result := mcp.InitializeResult{
 		ProtocolVersion: s.protocolVersion(request.Params.ProtocolVersion),
 		ServerInfo: mcp.Implementation{
@@ -1037,6 +1068,64 @@ func (s *MCPServer) handleToolCall(
 	return result, nil
 }
 
+func (s *MCPServer) handleCompletion(
+	ctx context.Context,
+	id any,
+	request mcp.CompleteRequest,
+) (*mcp.CompleteResult, *requestError) {
+
+	promptRef, resourceRef, err := mcp.ParseCompletionReference(request)
+	if err != nil {
+		return nil, &requestError{
+			id:   id,
+			code: mcp.INVALID_PARAMS,
+			err:  err,
+		}
+	}
+
+	var handler *mcp.CompletionHandlerFunc
+	if promptRef != nil {
+		s.promptsMu.RLock()
+		if prompt, ok := s.prompts[promptRef.Name]; ok {
+			for _, arg := range prompt.Arguments {
+				if arg.Name == request.Params.Argument.Name {
+					handler = arg.CompletionHandler
+				}
+			}
+		}
+		s.promptsMu.RUnlock()
+	} else if resourceRef != nil {
+		s.resourcesMu.RLock()
+		if resourceTemplate, ok := s.resourceTemplates[resourceRef.URI]; ok {
+			if h, ok := resourceTemplate.template.URITemplate.ArgumentCompletionHandlers[request.Params.Argument.Name]; ok {
+				handler = &h
+			}
+		}
+		s.resourcesMu.RUnlock()
+	}
+
+	if handler == nil {
+		return nil, &requestError{
+			id:   id,
+			code: mcp.INVALID_PARAMS,
+			err: fmt.Errorf(
+				"no completion handler for argument '%s' of reference %s",
+				request.Params.Argument.Name, request.Params.Ref,
+			),
+		}
+	}
+
+	result, err := (*handler)(ctx, request)
+	if err != nil {
+		return nil, &requestError{
+			id:   id,
+			code: mcp.INTERNAL_ERROR,
+			err:  err,
+		}
+	}
+	return result, nil
+}
+
 func (s *MCPServer) handleNotification(
 	ctx context.Context,
 	notification mcp.JSONRPCNotification,