Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitc54620f

Browse files
committed
allow specifying only allowed tools
1 parentdf1ca76 commitc54620f

File tree

6 files changed

+274
-135
lines changed

6 files changed

+274
-135
lines changed

‎cli/exp_mcp.go

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,12 @@ func (r *RootCmd) mcpCommand() *serpent.Command {
1616
var (
1717
client=new(codersdk.Client)
1818
instructionsstring
19+
allowedTools []string
1920
)
2021
return&serpent.Command{
2122
Use:"mcp",
2223
Handler:func(inv*serpent.Invocation)error {
23-
returnmcpHandler(inv,client,instructions)
24+
returnmcpHandler(inv,client,instructions,allowedTools)
2425
},
2526
Short:"Start an MCP server that can be used to interact with a Coder depoyment.",
2627
Middleware:serpent.Chain(
@@ -33,11 +34,17 @@ func (r *RootCmd) mcpCommand() *serpent.Command {
3334
Flag:"instructions",
3435
Value:serpent.StringOf(&instructions),
3536
},
37+
{
38+
Name:"allowed-tools",
39+
Description:"Comma-separated list of allowed tools. If not specified, all tools are allowed.",
40+
Flag:"allowed-tools",
41+
Value:serpent.StringArrayOf(&allowedTools),
42+
},
3643
},
3744
}
3845
}
3946

40-
funcmcpHandler(inv*serpent.Invocation,client*codersdk.Client,instructionsstring)error {
47+
funcmcpHandler(inv*serpent.Invocation,client*codersdk.Client,instructionsstring,allowedTools []string)error {
4148
ctx,cancel:=context.WithCancel(inv.Context())
4249
defercancel()
4350

@@ -51,9 +58,12 @@ func mcpHandler(inv *serpent.Invocation, client *codersdk.Client, instructions s
5158
returnerr
5259
}
5360
cliui.Infof(inv.Stderr,"Starting MCP server")
54-
cliui.Infof(inv.Stderr,"User : %s",me.Username)
55-
cliui.Infof(inv.Stderr,"URL : %s",client.URL)
56-
cliui.Infof(inv.Stderr,"Instructions : %q",instructions)
61+
cliui.Infof(inv.Stderr,"User : %s",me.Username)
62+
cliui.Infof(inv.Stderr,"URL : %s",client.URL)
63+
cliui.Infof(inv.Stderr,"Instructions : %q",instructions)
64+
iflen(allowedTools)>0 {
65+
cliui.Infof(inv.Stderr,"Allowed Tools : %v",allowedTools)
66+
}
5767
cliui.Infof(inv.Stderr,"Press Ctrl+C to stop the server")
5868

5969
// Capture the original stdin, stdout, and stderr.
@@ -66,12 +76,19 @@ func mcpHandler(inv *serpent.Invocation, client *codersdk.Client, instructions s
6676
inv.Stderr=invStderr
6777
}()
6878

69-
closer:=codermcp.New(ctx,client,
79+
options:=[]codermcp.Option{
7080
codermcp.WithInstructions(instructions),
7181
codermcp.WithLogger(&logger),
7282
codermcp.WithStdin(invStdin),
7383
codermcp.WithStdout(invStdout),
74-
)
84+
}
85+
86+
// Add allowed tools option if specified
87+
iflen(allowedTools)>0 {
88+
options=append(options,codermcp.WithAllowedTools(allowedTools))
89+
}
90+
91+
closer:=codermcp.New(ctx,client,options...)
7592

7693
<-ctx.Done()
7794
iferr:=closer.Close();err!=nil {

‎cli/exp_mcp_test.go

Lines changed: 61 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package cli_test
33
import (
44
"context"
55
"encoding/json"
6+
"slices"
67
"testing"
78

89
"github.com/stretchr/testify/assert"
@@ -17,6 +18,61 @@ import (
1718
funcTestExpMcp(t*testing.T) {
1819
t.Parallel()
1920

21+
t.Run("AllowedTools",func(t*testing.T) {
22+
t.Parallel()
23+
24+
ctx:=testutil.Context(t,testutil.WaitShort)
25+
cancelCtx,cancel:=context.WithCancel(ctx)
26+
t.Cleanup(cancel)
27+
28+
// Given: a running coder deployment
29+
client:=coderdtest.New(t,nil)
30+
_=coderdtest.CreateFirstUser(t,client)
31+
32+
// Given: we run the exp mcp command with allowed tools set
33+
inv,root:=clitest.New(t,"exp","mcp","--allowed-tools=coder_whoami,coder_list_templates")
34+
inv=inv.WithContext(cancelCtx)
35+
36+
pty:=ptytest.New(t)
37+
inv.Stdin=pty.Input()
38+
inv.Stdout=pty.Output()
39+
clitest.SetupConfig(t,client,root)
40+
41+
cmdDone:=make(chanstruct{})
42+
gofunc() {
43+
deferclose(cmdDone)
44+
err:=inv.Run()
45+
assert.NoError(t,err)
46+
}()
47+
48+
// When: we send a tools/list request
49+
toolsPayload:=`{"jsonrpc":"2.0","id":2,"method":"tools/list"}`
50+
pty.WriteLine(toolsPayload)
51+
_=pty.ReadLine(ctx)// ignore echoed output
52+
output:=pty.ReadLine(ctx)
53+
54+
cancel()
55+
<-cmdDone
56+
57+
// Then: we should only see the allowed tools in the response
58+
vartoolsResponsestruct {
59+
Resultstruct {
60+
Tools []struct {
61+
Namestring`json:"name"`
62+
}`json:"tools"`
63+
}`json:"result"`
64+
}
65+
err:=json.Unmarshal([]byte(output),&toolsResponse)
66+
require.NoError(t,err)
67+
require.Len(t,toolsResponse.Result.Tools,2,"should have exactly 2 tools")
68+
foundTools:=make([]string,0,2)
69+
for_,tool:=rangetoolsResponse.Result.Tools {
70+
foundTools=append(foundTools,tool.Name)
71+
}
72+
slices.Sort(foundTools)
73+
require.Equal(t, []string{"coder_list_templates","coder_whoami"},foundTools)
74+
})
75+
2076
t.Run("OK",func(t*testing.T) {
2177
t.Parallel()
2278

@@ -34,10 +90,12 @@ func TestExpMcp(t *testing.T) {
3490
inv.Stdout=pty.Output()
3591
clitest.SetupConfig(t,client,root)
3692

37-
cmdDone:=tGo(t,func() {
93+
cmdDone:=make(chanstruct{})
94+
gofunc() {
95+
deferclose(cmdDone)
3896
err:=inv.Run()
3997
assert.NoError(t,err)
40-
})
98+
}()
4199

42100
payload:=`{"jsonrpc":"2.0","id":1,"method":"initialize"}`
43101
pty.WriteLine(payload)
@@ -75,4 +133,4 @@ func TestExpMcp(t *testing.T) {
75133
err:=inv.Run()
76134
assert.ErrorContains(t,err,"your session has expired")
77135
})
78-
}
136+
}

‎mcp/mcp.go

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ type mcpOptions struct {
2121
out io.Writer
2222
instructionsstring
2323
logger*slog.Logger
24+
allowedTools []string
2425
}
2526

2627
// Option is a function that configures the MCP server.
@@ -54,6 +55,13 @@ func WithStdout(out io.Writer) Option {
5455
}
5556
}
5657

58+
// WithAllowedTools sets the allowed tools for the MCP server.
59+
funcWithAllowedTools(tools []string)Option {
60+
returnfunc(o*mcpOptions) {
61+
o.allowedTools=tools
62+
}
63+
}
64+
5765
// New creates a new MCP server with the given client and options.
5866
funcNew(ctx context.Context,client*codersdk.Client,opts...Option) io.Closer {
5967
options:=&mcpOptions{
@@ -74,14 +82,15 @@ func New(ctx context.Context, client *codersdk.Client, opts ...Option) io.Closer
7482

7583
logger:=slog.Make(sloghuman.Sink(os.Stdout))
7684

77-
mcptools.RegisterCoderReportTask(mcpSrv,client,logger)
78-
mcptools.RegisterCoderWhoami(mcpSrv,client)
79-
mcptools.RegisterCoderListTemplates(mcpSrv,client)
80-
mcptools.RegisterCoderListWorkspaces(mcpSrv,client)
81-
mcptools.RegisterCoderGetWorkspace(mcpSrv,client)
82-
mcptools.RegisterCoderWorkspaceExec(mcpSrv,client)
83-
mcptools.RegisterCoderStartWorkspace(mcpSrv,client)
84-
mcptools.RegisterCoderStopWorkspace(mcpSrv,client)
85+
// Register tools based on the allowed list (if specified)
86+
reg:=mcptools.AllTools()
87+
iflen(options.allowedTools)>0 {
88+
reg=reg.WithOnlyAllowed(options.allowedTools...)
89+
}
90+
reg.Register(mcpSrv, mcptools.ToolDeps{
91+
Client:client,
92+
Logger:&logger,
93+
})
8594

8695
srv:=server.NewStdioServer(mcpSrv)
8796
srv.SetErrorLogger(log.New(options.out,"",log.LstdFlags))

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp