Skip to content

Commit ce247f8

Browse files
samthanawallasamthanawalla
committed
mcp: implement completion
This CL implements a completion handler for the server and a complete method for the client.
1 parent 09181c2 commit ce247f8

File tree

7 files changed

+286
-5
lines changed

7 files changed

+286
-5
lines changed

‎examples/completion/main.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package main
6+
7+
import (
8+
"context"
9+
"fmt"
10+
"log"
11+
12+
"github.com/modelcontextprotocol/go-sdk/mcp"
13+
)
14+
15+
// This example demonstrates the minimal code to declare and assign
16+
// a CompletionHandler to an MCP Server's options.
17+
func main() {
18+
// Define your custom CompletionHandler logic
19+
myCompletionHandler := func(_ context.Context, _ *mcp.ServerSession, params *mcp.CompleteParams) (*mcp.CompleteResult, error) {
20+
// In a real application, you'd implement actual completion logic here
21+
// For this example, we return a fixed set of suggestions.
22+
var suggestions []string
23+
switch params.Ref.Type {
24+
case "ref/prompt":
25+
suggestions = []string{"suggestion1", "suggestion2", "suggestion3"}
26+
case "ref/resource":
27+
suggestions = []string{"suggestion4", "suggestion5", "suggestion6"}
28+
default:
29+
return nil, fmt.Errorf("unrecognized content type %s", params.Ref.Type)
30+
}
31+
32+
return &mcp.CompleteResult{
33+
Completion: mcp.CompletionResultDetails{
34+
HasMore: false,
35+
Total: int64(len(suggestions)),
36+
Values: suggestions,
37+
},
38+
}, nil
39+
}
40+
41+
// Create the MCP Server instance and assign the handler
42+
// No server running, just showing the configuration.
43+
_ = mcp.NewServer("myServer", "v1.0.0", &mcp.ServerOptions{
44+
CompletionHandler: myCompletionHandler,
45+
})
46+
47+
log.Println("MCP Server instance created with a CompletionHandler assigned (but not running).")
48+
log.Println("This example demonstrates configuration, not live interaction.")
49+
}

‎mcp/client.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,7 @@ func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]
233233

234234
// clientMethodInfos maps from the RPC method name to serverMethodInfos.
235235
var clientMethodInfos = map[string]methodInfo{
236+
methodComplete: newMethodInfo(sessionMethod((*ClientSession).Complete)),
236237
methodPing: newMethodInfo(sessionMethod((*ClientSession).ping)),
237238
methodListRoots: newMethodInfo(clientMethod((*Client).listRoots)),
238239
methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage)),
@@ -328,6 +329,10 @@ func (cs *ClientSession) ReadResource(ctx context.Context, params *ReadResourceP
328329
return handleSend[*ReadResourceResult](ctx, cs, methodReadResource, orZero[Params](params))
329330
}
330331

332+
func (cs *ClientSession) Complete(ctx context.Context, params *CompleteParams) (*CompleteResult, error) {
333+
return handleSend[*CompleteResult](ctx, cs, methodComplete, orZero[Params](params))
334+
}
335+
331336
func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (Result, error) {
332337
return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params)
333338
}

‎mcp/completion.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package mcp
6+
7+
import (
8+
"context"
9+
"encoding/json"
10+
"fmt"
11+
)
12+
13+
type Reference struct {
14+
Type string `json:"type"`
15+
Name string `json:"name,omitempty"`
16+
URI string ` json:"uri,omitempty"`
17+
}
18+
19+
func (r *Reference) UnmarshalJSON(data []byte) error {
20+
type wireReference Reference // for naive unmarshaling
21+
var r2 wireReference
22+
if err := json.Unmarshal(data, &r2); err != nil {
23+
return err
24+
}
25+
switch r2.Type {
26+
case "ref/prompt", "ref/resource":
27+
default:
28+
return fmt.Errorf("unrecognized content type %s", r2.Type)
29+
}
30+
*r = Reference(r2)
31+
return nil
32+
}
33+
34+
// A CompletionHandler handles completion.
35+
type CompletionHandler func(context.Context, *ServerSession, *CompleteParams) (*CompleteResult, error)

‎mcp/completion_test.go‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package mcp_test
6+
7+
import (
8+
"encoding/json"
9+
"testing"
10+
11+
"github.com/google/go-cmp/cmp"
12+
"github.com/modelcontextprotocol/go-sdk/mcp"
13+
)
14+
15+
func TestReference(t *testing.T) {
16+
tests := []struct {
17+
name string
18+
in mcp.Reference
19+
want string
20+
wantErr bool
21+
}{
22+
{
23+
name: "PromptReference",
24+
in: mcp.Reference{Type: "ref/prompt", Name: "my_prompt"},
25+
want: `{"type":"ref/prompt","name":"my_prompt"}`,
26+
},
27+
{
28+
name: "ResourceReference",
29+
in: mcp.Reference{Type: "ref/resource", URI: "file:///path/to/resource.txt"},
30+
want: `{"type":"ref/resource","uri":"file:///path/to/resource.txt"}`,
31+
},
32+
{
33+
name: "PromptReference with empty name",
34+
in: mcp.Reference{Type: "ref/prompt", Name: ""},
35+
want: `{"type":"ref/prompt"}`,
36+
},
37+
{
38+
name: "ResourceReference with empty URI",
39+
in: mcp.Reference{Type: "ref/resource", URI: ""},
40+
want: `{"type":"ref/resource"}`,
41+
},
42+
{
43+
name: "Unrecognized Type",
44+
in: mcp.Reference{Type: "ref/unknown", Name: "something"},
45+
want: `{"type":"ref/unknown","name":"something"}`,
46+
wantErr: true,
47+
},
48+
{
49+
name: "Missing Type Field",
50+
in: mcp.Reference{Name: "missing"},
51+
want: `{"type":"","name":"missing"}`,
52+
wantErr: true,
53+
},
54+
{
55+
name: "Invalid JSON Format",
56+
in: mcp.Reference{},
57+
want: `{"type":""}`,
58+
wantErr: true,
59+
},
60+
}
61+
for _, test := range tests {
62+
t.Run(test.name, func(t *testing.T) {
63+
// Test Marshal
64+
got, err := json.Marshal(test.in)
65+
if err != nil {
66+
t.Fatalf("json.Marshal(%v) failed: %v", test.in, err)
67+
}
68+
if diff := cmp.Diff(test.want, string(got)); diff != "" {
69+
t.Errorf("json.Marshal(%v) mismatch (-want +got):\n%s", test.in, diff)
70+
}
71+
72+
// Test Unmarshal
73+
var unmarshaled mcp.Reference
74+
err = json.Unmarshal([]byte(test.want), &unmarshaled)
75+
76+
if test.wantErr {
77+
if err == nil {
78+
t.Fatalf("json.Unmarshal(%q) should have failed", test.want)
79+
}
80+
return
81+
}
82+
if err != nil {
83+
t.Fatalf("json.Unmarshal(%q) failed: %v", test.want, err)
84+
}
85+
if diff := cmp.Diff(test.in, unmarshaled); diff != "" {
86+
t.Errorf("json.Unmarshal(%q) mismatch (-want +got):\n%s", test.want, diff)
87+
}
88+
})
89+
}
90+
}
91+
92+
func TestCompleteParams(t *testing.T) {
93+
// Test CompleteParams
94+
params := mcp.CompleteParams{
95+
Ref: &mcp.Reference{
96+
Type: "ref/prompt",
97+
Name: "my_prompt",
98+
},
99+
Argument: mcp.CompleteParamsArgument{
100+
Name: "language",
101+
Value: "go",
102+
},
103+
}
104+
wantParamsJSON := `{"argument":{"name":"language","value":"go"},"ref":{"type":"ref/prompt","name":"my_prompt"}}`
105+
106+
gotParamsJSON, err := json.Marshal(params)
107+
if err != nil {
108+
t.Fatalf("json.Marshal(CompleteParams) failed: %v", err)
109+
}
110+
if diff := cmp.Diff(wantParamsJSON, string(gotParamsJSON)); diff != "" {
111+
t.Errorf("CompleteParams marshal mismatch (-want +got):\n%s", diff)
112+
}
113+
114+
var unmarshaledParams mcp.CompleteParams
115+
if err := json.Unmarshal([]byte(wantParamsJSON), &unmarshaledParams); err != nil {
116+
t.Fatalf("json.Unmarshal(CompleteParams) failed: %v", err)
117+
}
118+
if diff := cmp.Diff(params, unmarshaledParams); diff != "" {
119+
t.Errorf("CompleteParams unmarshal mismatch (-want +got):\n%s", diff)
120+
}
121+
}
122+
123+
func TestCompleteResult(t *testing.T) {
124+
res := mcp.CompleteResult{
125+
Completion: mcp.CompletionResultDetails{
126+
Values: []string{"golang", "google", "goroutine"},
127+
Total: 10,
128+
HasMore: true,
129+
},
130+
}
131+
wantResJSON := `{"completion":{"hasMore":true,"total":10,"values":["golang","google","goroutine"]}}`
132+
gotResJSON, err := json.Marshal(res)
133+
if err != nil {
134+
t.Fatalf("json.Marshal(CompleteResult) failed: %v", err)
135+
}
136+
if diff := cmp.Diff(wantResJSON, string(gotResJSON)); diff != "" {
137+
t.Errorf("CompleteResult marshal mismatch (-want +got):\n%s", diff)
138+
}
139+
140+
var unmarshaledRes mcp.CompleteResult
141+
if err := json.Unmarshal([]byte(wantResJSON), &unmarshaledRes); err != nil {
142+
t.Fatalf("json.Unmarshal(CompleteResult) failed: %v", err)
143+
}
144+
if diff := cmp.Diff(res, unmarshaledRes); diff != "" {
145+
t.Errorf("CompleteResult unmarshal mismatch (-want +got):\n%s", diff)
146+
}
147+
}

‎mcp/protocol.go‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,36 @@ type ClientCapabilities struct {
163163
Elicitation *ElicitationCapabilities `json:"elicitation,omitempty"`
164164
}
165165

166+
type CompleteParamsArgument struct {
167+
// The name of the argument
168+
Name string `json:"name"`
169+
// The value of the argument to use for completion matching.
170+
Value string `json:"value"`
171+
}
172+
173+
type CompleteParams struct {
174+
// This property is reserved by the protocol to allow clients and servers to
175+
// attach additional metadata to their responses.
176+
Meta `json:"_meta,omitempty"`
177+
// The argument's information
178+
Argument CompleteParamsArgument `json:"argument"`
179+
Ref *Reference `json:"ref"`
180+
}
181+
182+
type CompletionResultDetails struct {
183+
HasMore bool `json:"hasMore,omitempty"`
184+
Total int64 `json:"total,omitempty"`
185+
Values []string `json:"values"`
186+
}
187+
188+
// The server's response to a completion/complete request
189+
type CompleteResult struct {
190+
// This property is reserved by the protocol to allow clients and servers to
191+
// attach additional metadata to their responses.
192+
Meta `json:"_meta,omitempty"`
193+
Completion CompletionResultDetails `json:"completion"`
194+
}
195+
166196
type CreateMessageParams struct {
167197
// This property is reserved by the protocol to allow clients and servers to
168198
// attach additional metadata to their responses.
@@ -817,6 +847,9 @@ type implementation struct {
817847
Version string `json:"version"`
818848
}
819849

850+
// Present if the server supports argument autocompletion suggestions.
851+
type completionCapabilities struct{}
852+
820853
// Present if the server supports sending log messages to the client.
821854
type loggingCapabilities struct{}
822855

@@ -839,7 +872,7 @@ type resourceCapabilities struct {
839872
// additional capabilities.
840873
type serverCapabilities struct {
841874
// Present if the server supports argument autocompletion suggestions.
842-
Completions struct{} `json:"completions,omitempty"`
875+
Completions completionCapabilities `json:"completions,omitempty"`
843876
// Experimental, non-standard capabilities that the server supports.
844877
Experimental map[string]struct{} `json:"experimental,omitempty"`
845878
// Present if the server supports sending log messages to the client.

‎mcp/server.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ type ServerOptions struct {
5757
RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams)
5858
// If non-nil, called when "notifications/progress" is received.
5959
ProgressNotificationHandler func(context.Context, *ServerSession, *ProgressNotificationParams)
60+
// If non-nil, called when "completion/complete" is received.
61+
CompletionHandler CompletionHandler
6062
}
6163

6264
// NewServer creates a new MCP server. The resulting server has no features:
@@ -226,6 +228,13 @@ func (s *Server) RemoveResourceTemplates(uriTemplates ...string) {
226228
func() bool { return s.resourceTemplates.remove(uriTemplates...) })
227229
}
228230

231+
func (s *Server) complete(ctx context.Context, ss *ServerSession, params *CompleteParams) (Result, error) {
232+
if s.opts.CompletionHandler == nil {
233+
return nil, jsonrpc2.ErrMethodNotFound
234+
}
235+
return s.opts.CompletionHandler(ctx, ss, params)
236+
}
237+
229238
// changeAndNotify is called when a feature is added or removed.
230239
// It calls change, which should do the work and report whether a change actually occurred.
231240
// If there was a change, it notifies a snapshot of the sessions.
@@ -563,6 +572,7 @@ func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]
563572

564573
// serverMethodInfos maps from the RPC method name to serverMethodInfos.
565574
var serverMethodInfos = map[string]methodInfo{
575+
methodComplete: newMethodInfo(serverMethod((*Server).complete)),
566576
methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize)),
567577
methodPing: newMethodInfo(sessionMethod((*ServerSession).ping)),
568578
methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts)),
@@ -646,6 +656,7 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
646656
// reject unsupported features.
647657
ProtocolVersion: version,
648658
Capabilities: &serverCapabilities{
659+
Completions: completionCapabilities{},
649660
Prompts: &promptCapabilities{
650661
ListChanged: true,
651662
},

‎mcp/streamable_test.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -143,10 +143,11 @@ func TestStreamableServerTransport(t *testing.T) {
143143
initReq := req(1, "initialize", &InitializeParams{})
144144
initResp := resp(1, &InitializeResult{
145145
Capabilities: &serverCapabilities{
146-
Logging: &loggingCapabilities{},
147-
Prompts: &promptCapabilities{ListChanged: true},
148-
Resources: &resourceCapabilities{ListChanged: true},
149-
Tools: &toolCapabilities{ListChanged: true},
146+
Completions: completionCapabilities{},
147+
Logging: &loggingCapabilities{},
148+
Prompts: &promptCapabilities{ListChanged: true},
149+
Resources: &resourceCapabilities{ListChanged: true},
150+
Tools: &toolCapabilities{ListChanged: true},
150151
},
151152
ProtocolVersion: "2025-03-26",
152153
ServerInfo: &implementation{Name: "testServer", Version: "v1.0.0"},

0 commit comments

Comments
 (0)