feat: Support completion handling for prompts and resources

- Added completionCompleteHandler to registeredPrompt and registeredResource for handling completion requests.
- Refactored handleCompletionComplete methods in promptManager and resourceManager to utilize new completion handlers.
- Added utility functions for extracting parameters from JSON-RPC requests.
- Updated lifecycleManager to check for completion availability in both prompt and resource managers.
- Implemented tests for completion handling in both promptManager and resourceManager, ensuring correct behavior for valid and invalid requests.

Author: FlySky-z

client.go

@@ -53,6 +53,8 @@ type Connector interface {
 	ListPrompts(ctx context.Context, req *ListPromptsRequest) (*ListPromptsResult, error)
 	// GetPrompt retrieves a specific prompt by name.
 	GetPrompt(ctx context.Context, req *GetPromptRequest) (*GetPromptResult, error)
+	// CompleteCompletion provides parameter auto-completion suggestions for prompts and resource URIs.
+	CompleteCompletion(ctx context.Context, req *CompleteCompletionRequest) (JSONRPCMessage, error)
 	// ListResources retrieves all available resources from the server.
 	ListResources(ctx context.Context, req *ListResourcesRequest) (*ListResourcesResult, error)
 	// ReadResource reads the content of a specific resource.
@@ -548,6 +550,42 @@ func (c *Client) GetPrompt(ctx context.Context, getPromptReq *GetPromptRequest)
 	return parseGetPromptResultFromJSON(rawResp)
 }
 
+// CompleteCompletion provides parameter auto-completion suggestions for prompts and resource URIs.
+func (c *Client) CompleteCompletion(ctx context.Context, completeReq *CompleteCompletionRequest) (*CompleteCompletionResult, error) {
+	// Check if initialized.
+	if !c.initialized {
+		return nil, errors.ErrNotInitialized
+	}
+
+	// Create request.
+	requestID := c.requestID.Add(1)
+	req := &JSONRPCRequest{
+		JSONRPC: JSONRPCVersion,
+		ID:      requestID,
+		Request: Request{
+			Method: MethodCompletionComplete,
+		},
+		Params: completeReq.Params,
+	}
+
+	rawResp, err := c.transport.sendRequest(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("complete completion request failed: %v", err)
+	}
+
+	// Check for error response
+	if isErrorResponse(rawResp) {
+		errResp, err := parseRawMessageToError(rawResp)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse error response: %w", err)
+		}
+		return nil, fmt.Errorf("complete completion error: %s (code: %d)",
+			errResp.Error.Message, errResp.Error.Code)
+	}
+
+	return parseCompleteCompletionResultFromJSON(rawResp)
+}
+
 // ListResources lists available resources.
 func (c *Client) ListResources(ctx context.Context, listResourcesReq *ListResourcesRequest) (*ListResourcesResult, error) {
 	// Check if initialized.

examples/resource_prompt_example/client/main.go

@@ -8,6 +8,7 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 
@@ -38,6 +39,10 @@ func initializeClient(ctx context.Context) (*mcp.Client, error) {
 	}
 	log.Printf("Connected to server: %s %s", initResult.ServerInfo.Name, initResult.ServerInfo.Version)
 
+	// Check server capabilities
+	capabilitiesJSON, _ := json.Marshal(initResult.Capabilities)
+	log.Printf("Server capabilities:%v", string(capabilitiesJSON))
+
 	return newClient, nil
 }
 
@@ -173,6 +178,77 @@ func handlePrompts(ctx context.Context, client *mcp.Client) error {
 	return nil
 }
 
+// handleCompletions demonstrates completion functionality for prompts and resources
+func handleCompletions(ctx context.Context, client *mcp.Client) error {
+	log.Printf("===== Prompt completion =====")
+
+	// Test completion for prompt arguments (code_review -> language)
+	promptCompletionReq := &mcp.CompleteCompletionRequest{}
+	promptCompletionReq.Params.Ref.Type = "ref/prompt"
+	promptCompletionReq.Params.Ref.Name = "code_review"
+	promptCompletionReq.Params.Argument.Name = "language"
+	promptCompletionReq.Params.Argument.Value = "p"
+
+	promptCompletionResult, err := client.CompleteCompletion(ctx, promptCompletionReq)
+	if err != nil {
+		log.Printf("Prompt completion error: %v", err)
+	} else {
+		log.Printf("Prompt completion for '%s' with prefix '%s': found %d suggestions",
+			promptCompletionReq.Params.Ref.Name,
+			promptCompletionReq.Params.Argument.Value,
+			len(promptCompletionResult.Completion.Values))
+		for i, value := range promptCompletionResult.Completion.Values {
+			log.Printf("  [%d] %s", i, value)
+		}
+	}
+
+	log.Printf("===== Resource completion =====")
+	// Test completion for resource (resource://example/completion)
+	resourceCompletionReq := &mcp.CompleteCompletionRequest{}
+	resourceCompletionReq.Params.Ref.Type = "ref/resource"
+	resourceCompletionReq.Params.Ref.URI = "resource://example/completion"
+	resourceCompletionReq.Params.Argument.Name = "query"
+	resourceCompletionReq.Params.Argument.Value = ""
+
+	resourceCompletionResult, err := client.CompleteCompletion(ctx, resourceCompletionReq)
+	if err != nil {
+		log.Printf("Resource completion error: %v", err)
+	} else {
+		log.Printf("Resource completion from '%s' for '%s': found %d suggestions",
+			resourceCompletionReq.Params.Ref.URI,
+			resourceCompletionReq.Params.Argument.Name,
+			len(resourceCompletionResult.Completion.Values))
+		for i, value := range resourceCompletionResult.Completion.Values {
+			log.Printf("  [%d] %s", i, value)
+		}
+	}
+
+	log.Printf("===== Resource template completion =====")
+
+	// Test completion for resource template (file-template -> path)
+	resourceTemplateCompletionReq := &mcp.CompleteCompletionRequest{}
+	resourceTemplateCompletionReq.Params.Ref.Type = "ref/resource"
+	resourceTemplateCompletionReq.Params.Ref.URI = "file://test"
+	resourceTemplateCompletionReq.Params.Argument.Name = "keyword"
+	resourceTemplateCompletionReq.Params.Argument.Value = "completion"
+
+	resourceTemplateCompletionResult, err := client.CompleteCompletion(ctx, resourceTemplateCompletionReq)
+	if err != nil {
+		log.Printf("Resource template completion error: %v", err)
+	} else {
+		log.Printf("Resource template completion from '%s' for '%s' with keyword '%s': found %d suggestions",
+			resourceTemplateCompletionReq.Params.Ref.URI,
+			resourceTemplateCompletionReq.Params.Argument.Name,
+			resourceTemplateCompletionReq.Params.Argument.Value,
+			len(resourceTemplateCompletionResult.Completion.Values))
+		for i, value := range resourceTemplateCompletionResult.Completion.Values {
+			log.Printf("  [%d] %s", i, value)
+		}
+	}
+
+	return nil
+}
+
 // handleTools manages tool-related operations including listing and calling tools
 func handleTools(ctx context.Context, client *mcp.Client) error {
 	log.Printf("===== List tools =====")
@@ -223,6 +299,8 @@ func main() {
 	}
 	defer client.Close()
 
+	//
+
 	// Handle resources
 	if err := handleResources(ctx, client); err != nil {
 		log.Printf("Error: %v", err)
@@ -233,6 +311,11 @@ func main() {
 		log.Printf("Error: %v", err)
 	}
 
+	// Handle completions
+	if err := handleCompletions(ctx, client); err != nil {
+		log.Printf("Error: %v", err)
+	}
+
 	// Handle tools
 	if err := handleTools(ctx, client); err != nil {
 		log.Printf("Error: %v", err)

examples/resource_prompt_example/server/main.go

@@ -83,6 +83,77 @@ func registerExampleResources(s *mcp.Server) {
 
 	s.RegisterResource(imageResource, imageHandler)
 	log.Printf("Registered image resource: %s", imageResource.Name)
+
+	// Register resource completion.
+	textResourceCompletion := &mcp.Resource{
+		URI:         "resource://example/completion",
+		Name:        "example-text-completion",
+		Description: "Example text resource with completion",
+		MimeType:    "text/plain",
+	}
+
+	// Define completion resource handler
+	resourceCompletionHandler := func(ctx context.Context, req *mcp.CompleteCompletionRequest) (*mcp.CompleteCompletionResult, error) {
+		result := &mcp.CompleteCompletionResult{}
+		if req.Params.Argument.Name == "query" {
+			if req.Params.Context.Arguments != nil {
+				// Use context argument if provided
+				context, ok := req.Params.Context.Arguments["context"]
+				if ok && context != "" {
+					result.Completion.Values = []string{fmt.Sprintf("Context: %s", context)}
+				} else {
+					result.Completion.Values = []string{"Context is empty"}
+				}
+			} else {
+				result.Completion.Values = []string{"First document text", "Second document text", "Third document text"}
+			}
+		} else {
+			result.Completion.Values = []string{"Unknown argument"}
+		}
+		result.Completion.Total = len(result.Completion.Values)
+		result.Completion.HasMore = false
+		return result, nil
+	}
+	s.RegisterResource(textResourceCompletion, nil,
+		mcp.WithResourceCompletion(resourceCompletionHandler),
+	)
+	log.Printf("Registered text resource completion: %s", textResourceCompletion.Name)
+
+	// Register resource template with completion
+	fileTemplate := mcp.NewResourceTemplate(
+		"file://{path}",
+		"example-file-template-completion",
+		mcp.WithTemplateDescription("Example file resource template with completion"),
+		mcp.WithTemplateMIMEType("text/plain"),
+	)
+
+	// Define completion handler for file paths
+	fileCompletionHandler := func(ctx context.Context, req *mcp.CompleteCompletionRequest, params map[string]string) (*mcp.CompleteCompletionResult, error) {
+		result := &mcp.CompleteCompletionResult{}
+		if req.Params.Argument.Name == "keyword" {
+			keyword := req.Params.Argument.Value
+			result.Completion.Values = []string{
+				fmt.Sprint("file://", params["path"], "_1?keyword=", keyword),
+				fmt.Sprint("file://", params["path"], "_2?keyword=", keyword),
+				fmt.Sprint("file://", params["path"], "_3?keyword=", keyword),
+			}
+		} else {
+			result.Completion.Values = []string{"unknown argument"}
+		}
+
+		total := len(result.Completion.Values)
+		result.Completion.Total = total
+		result.Completion.HasMore = false
+
+		return result, nil
+	}
+
+	s.RegisterResourceTemplate(
+		fileTemplate,
+		nil,
+		mcp.WithTemplateCompletion(fileCompletionHandler),
+	)
+	log.Printf("Registered file resource template completion: %s", fileTemplate.Name)
 }
 
 // Register example prompts.
@@ -162,6 +233,73 @@ func registerExamplePrompts(s *mcp.Server) {
 
 	s.RegisterPrompt(advancedPrompt, advancedPromptHandler)
 	log.Printf("Registered advanced prompt: %s", advancedPrompt.Name)
+
+	// Register prompt with completion support
+	codeReviewPrompt := &mcp.Prompt{
+		Name:        "code_review",
+		Description: "Code review prompt completion",
+		Arguments: []mcp.PromptArgument{
+			{
+				Name:        "language",
+				Description: "Programming language of the code",
+				Required:    true,
+			},
+		},
+	}
+
+	// Define completion prompt handler
+	codeReviewPromptHandler := func(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
+		return &mcp.GetPromptResult{
+			Messages: []mcp.PromptMessage{
+				{
+					Role: "user",
+					Content: mcp.TextContent{
+						Type: "text",
+						Text: fmt.Sprintf("Please review the following %s code:\n%s", req.Params.Arguments["language"], req.Params.Arguments["code"]),
+					},
+				},
+			},
+		}, nil
+	}
+
+	// Define completion handler for prompt arguments
+	codeReviewPromptCompletionHandler := func(ctx context.Context, req *mcp.CompleteCompletionRequest) (*mcp.CompleteCompletionResult, error) {
+		result := &mcp.CompleteCompletionResult{}
+		if req.Params.Argument.Name == "language" {
+			prefix := req.Params.Argument.Value
+			categories := []string{
+				"python", "pytorch", "javascript", "typescript",
+				"go", "java", "c++", "c#", "ruby", "php",
+			}
+
+			var matches []string
+			for _, cat := range categories {
+				if prefix == "" || len(prefix) == 0 {
+					matches = append(matches, cat)
+				} else if len(cat) >= len(prefix) && cat[:len(prefix)] == prefix {
+					matches = append(matches, cat)
+				}
+			}
+
+			// Limit to first 10 matches
+			if len(matches) > 10 {
+				matches = matches[:10]
+			}
+
+			result.Completion.Values = matches
+		} else {
+			result.Completion.Values = []string{"unknown argument"}
+		}
+
+		total := len(result.Completion.Values)
+		result.Completion.Total = total
+		result.Completion.HasMore = false
+
+		return result, nil
+	}
+
+	s.RegisterPrompt(codeReviewPrompt, codeReviewPromptHandler, mcp.WithPromptCompletion(codeReviewPromptCompletionHandler))
+	log.Printf("Registered prompt completion: %s", codeReviewPrompt.Name)
 }
 
 // Register example tools.

handler.go

@@ -8,6 +8,9 @@ package mcp
 
 import (
 	"context"
+
+	"trpc.group/trpc-go/trpc-mcp-go/internal/errors"
+	"trpc.group/trpc-go/trpc-mcp-go/internal/utils"
 )
 
 const (
@@ -199,7 +202,18 @@ func (h *mcpHandler) handlePromptsGet(ctx context.Context, req *JSONRPCRequest,
 }
 
 func (h *mcpHandler) handleCompletionComplete(ctx context.Context, req *JSONRPCRequest, session Session) (JSONRPCMessage, error) {
-	return h.promptManager.handleCompletionComplete(ctx, req)
+	ref := utils.ExtractMap(req.Params.(map[string]interface{}), "ref")
+	if ref == nil {
+		return newJSONRPCErrorResponse(req.ID, ErrCodeInvalidParams, errors.ErrMissingParams.Error(), nil), nil
+	}
+	switch utils.ExtractString(ref, "type") {
+	case "ref/prompt":
+		return h.promptManager.handleCompletionComplete(ctx, req)
+	case "ref/resource":
+		return h.resourceManager.handleCompletionComplete(ctx, req)
+	default:
+		return newJSONRPCErrorResponse(req.ID, ErrCodeInvalidParams, errors.ErrInvalidParams.Error(), nil), nil
+	}
 }
 
 // handleNotification implements the handler interface's handleNotification method

internal/errors/errors.go

@@ -16,6 +16,7 @@ var (
 	ErrInvalidToolFormat     = errors.New("invalid tool format")
 	ErrToolNotFound          = errors.New("tool not found")
 	ErrInvalidToolParams     = errors.New("invalid tool parameters")
+	ErrMethodNotFound        = errors.New("method not found")
 
 	// JSON-RPC related errors
 	ErrParseJSONRPC           = errors.New("failed to parse JSON-RPC message")

manager_lifecycle.go

@@ -127,13 +127,21 @@ func (m *lifecycleManager) updateCapabilities() {
 		capMap["resources"] = map[string]interface{}{
 			"listChanged": true,
 		}
+
+		if m.resourceManager.hasCompletionCompleteHandler() {
+			capMap["completions"] = map[string]interface{}{}
+		}
 	}
 
 	// If there is a prompt manager and prompts are registered, add prompt capabilities
 	if m.promptManager != nil && len(m.promptManager.getPrompts()) > 0 {
 		capMap["prompts"] = map[string]interface{}{
 			"listChanged": true,
 		}
+
+		if m.promptManager.hasCompletionCompleteHandler() {
+			capMap["completions"] = map[string]interface{}{}
+		}
 	}
 
 	// Preserve existing experimental features