Skip to content

Commit df1ca76

Browse files
johnstcnjohnstcn
committed
add more tools!
1 parent ba898b9 commit df1ca76

File tree

3 files changed

+274
-4
lines changed

3 files changed

+274
-4
lines changed

‎mcp/mcp.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,12 @@ func New(ctx context.Context, client *codersdk.Client, opts ...Option) io.Closer
7676

7777
mcptools.RegisterCoderReportTask(mcpSrv, client, logger)
7878
mcptools.RegisterCoderWhoami(mcpSrv, client)
79+
mcptools.RegisterCoderListTemplates(mcpSrv, client)
7980
mcptools.RegisterCoderListWorkspaces(mcpSrv, client)
81+
mcptools.RegisterCoderGetWorkspace(mcpSrv, client)
8082
mcptools.RegisterCoderWorkspaceExec(mcpSrv, client)
83+
mcptools.RegisterCoderStartWorkspace(mcpSrv, client)
84+
mcptools.RegisterCoderStopWorkspace(mcpSrv, client)
8185

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

‎mcp/tools/tools_coder.go

Lines changed: 184 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,26 @@ func RegisterCoderListWorkspaces(ta ToolAdder, client *codersdk.Client) {
3535
ta.AddTool(toolCoderListWorkspaces, handleCoderListWorkspaces(client))
3636
}
3737

38+
func RegisterCoderGetWorkspace(ta ToolAdder, client *codersdk.Client) {
39+
ta.AddTool(toolCoderGetWorkspace, handleCoderGetWorkspace(client))
40+
}
41+
3842
func RegisterCoderWorkspaceExec(ta ToolAdder, client *codersdk.Client) {
3943
ta.AddTool(toolCoderWorkspaceExec, handleCoderWorkspaceExec(client))
4044
}
4145

46+
func RegisterCoderListTemplates(ta ToolAdder, client *codersdk.Client) {
47+
ta.AddTool(toolCoderListTemplates, handleCoderListTemplates(client))
48+
}
49+
50+
func RegisterCoderStartWorkspace(ta ToolAdder, client *codersdk.Client) {
51+
ta.AddTool(toolCoderStartWorkspace, handleCoderStartWorkspace(client))
52+
}
53+
54+
func RegisterCoderStopWorkspace(ta ToolAdder, client *codersdk.Client) {
55+
ta.AddTool(toolCoderStopWorkspace, handleCoderStopWorkspace(client))
56+
}
57+
4258
var (
4359
toolCoderReportTask = mcp.NewTool("coder_report_task",
4460
mcp.WithDescription(`Report progress on a task.`),
@@ -61,16 +77,31 @@ Good Summaries:
6177
mcp.WithNumber(`offset`, mcp.Description(`The offset to start listing workspaces from. Defaults to 0.`), mcp.DefaultNumber(0)),
6278
mcp.WithNumber(`limit`, mcp.Description(`The maximum number of workspaces to list. Defaults to 10.`), mcp.DefaultNumber(10)),
6379
)
80+
toolCoderGetWorkspace = mcp.NewTool("coder_get_workspace",
81+
mcp.WithDescription(`Get information about a workspace on a given Coder deployment.`),
82+
mcp.WithString("workspace", mcp.Description(`The workspace ID or name to get.`), mcp.Required()),
83+
)
6484
toolCoderWorkspaceExec = mcp.NewTool("coder_workspace_exec",
6585
mcp.WithDescription(`Execute a command in a remote workspace on a given Coder deployment.`),
6686
// Required parameters.
6787
mcp.WithString("workspace", mcp.Description(`The workspace ID or name in which to execute the command in. The workspace must be running.`), mcp.Required()),
6888
mcp.WithString("command", mcp.Description(`The command to execute. Changing the working directory is not currently supported, so you may need to preface the command with 'cd /some/path && <my-command>'.`), mcp.Required()),
6989
)
90+
toolCoderListTemplates = mcp.NewTool("coder_list_templates",
91+
mcp.WithDescription(`List all templates on a given Coder deployment.`),
92+
)
93+
toolCoderStartWorkspace = mcp.NewTool("coder_start_workspace",
94+
mcp.WithDescription(`Start a workspace on a given Coder deployment.`),
95+
mcp.WithString("workspace", mcp.Description(`The workspace ID or name to start.`), mcp.Required()),
96+
)
97+
toolCoderStopWorkspace = mcp.NewTool("coder_stop_workspace",
98+
mcp.WithDescription(`Stop a workspace on a given Coder deployment.`),
99+
mcp.WithString("workspace", mcp.Description(`The workspace ID or name to stop.`), mcp.Required()),
100+
)
70101
)
71102

72103
// Example payload:
73-
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_report_task", "arguments": {"summary": "I'm working on the login page.", "link": "https://github.com/coder/coder/pull/1234", "emoji": "🔍", "done": false, "coder_url": "http://localhost:3000", "coder_session_token": "REDACTED"}}}
104+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_report_task", "arguments": {"summary": "I'm working on the login page.", "link": "https://github.com/coder/coder/pull/1234", "emoji": "🔍", "done": false}}}
74105
func handleCoderReportTask(log slog.Logger, client *codersdk.Client) mcpserver.ToolHandlerFunc {
75106
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
76107
if client == nil {
@@ -123,7 +154,7 @@ func handleCoderReportTask(log slog.Logger, client *codersdk.Client) mcpserver.T
123154
}
124155

125156
// Example payload:
126-
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_whoami", "arguments": {"coder_url": "http://localhost:3000", "coder_session_token": "REDACTED"}}}
157+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_whoami", "arguments": {}}}
127158
func handleCoderWhoami(client *codersdk.Client) mcpserver.ToolHandlerFunc {
128159
return func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
129160
if client == nil {
@@ -148,7 +179,7 @@ func handleCoderWhoami(client *codersdk.Client) mcpserver.ToolHandlerFunc {
148179
}
149180

150181
// Example payload:
151-
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_list_workspaces", "arguments": {"owner": "me", "offset": 0, "limit": 10, "coder_url": "http://localhost:3000", "coder_session_token": "REDACTED"}}}
182+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_list_workspaces", "arguments": {"owner": "me", "offset": 0, "limit": 10}}}
152183
func handleCoderListWorkspaces(client *codersdk.Client) mcpserver.ToolHandlerFunc {
153184
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
154185
if client == nil {
@@ -194,7 +225,39 @@ func handleCoderListWorkspaces(client *codersdk.Client) mcpserver.ToolHandlerFun
194225
}
195226

196227
// Example payload:
197-
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_workspace_exec", "arguments": {"workspace": "dev", "command": "ps -ef", "coder_url": "http://localhost:3000", "coder_session_token": "REDACTED"}}}
228+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_get_workspace", "arguments": {"workspace": "dev"}}}
229+
func handleCoderGetWorkspace(client *codersdk.Client) mcpserver.ToolHandlerFunc {
230+
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
231+
if client == nil {
232+
return nil, xerrors.New("developer error: client is required")
233+
}
234+
args := request.Params.Arguments
235+
236+
wsArg, ok := args["workspace"].(string)
237+
if !ok {
238+
return nil, xerrors.New("workspace is required")
239+
}
240+
241+
workspace, err := getWorkspaceByIDOrOwnerName(ctx, client, wsArg)
242+
if err != nil {
243+
return nil, xerrors.Errorf("failed to fetch workspace: %w", err)
244+
}
245+
246+
workspaceJSON, err := json.Marshal(workspace)
247+
if err != nil {
248+
return nil, xerrors.Errorf("failed to encode workspace: %w", err)
249+
}
250+
251+
return &mcp.CallToolResult{
252+
Content: []mcp.Content{
253+
mcp.NewTextContent(string(workspaceJSON)),
254+
},
255+
}, nil
256+
}
257+
}
258+
259+
// Example payload:
260+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_workspace_exec", "arguments": {"workspace": "dev", "command": "ps -ef"}}}
198261
func handleCoderWorkspaceExec(client *codersdk.Client) mcpserver.ToolHandlerFunc {
199262
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
200263
if client == nil {
@@ -265,6 +328,123 @@ func handleCoderWorkspaceExec(client *codersdk.Client) mcpserver.ToolHandlerFunc
265328
}
266329
}
267330

331+
// Example payload:
332+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_list_templates", "arguments": {}}}
333+
func handleCoderListTemplates(client *codersdk.Client) mcpserver.ToolHandlerFunc {
334+
return func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
335+
if client == nil {
336+
return nil, xerrors.New("developer error: client is required")
337+
}
338+
templates, err := client.Templates(ctx, codersdk.TemplateFilter{})
339+
if err != nil {
340+
return nil, xerrors.Errorf("failed to fetch templates: %w", err)
341+
}
342+
343+
templateJSON, err := json.Marshal(templates)
344+
if err != nil {
345+
return nil, xerrors.Errorf("failed to encode templates: %w", err)
346+
}
347+
348+
return &mcp.CallToolResult{
349+
Content: []mcp.Content{
350+
mcp.NewTextContent(string(templateJSON)),
351+
},
352+
}, nil
353+
}
354+
}
355+
356+
// Example payload:
357+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_start_workspace", "arguments": {"workspace": "dev"}}}
358+
func handleCoderStartWorkspace(client *codersdk.Client) mcpserver.ToolHandlerFunc {
359+
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
360+
if client == nil {
361+
return nil, xerrors.New("developer error: client is required")
362+
}
363+
364+
args := request.Params.Arguments
365+
366+
wsArg, ok := args["workspace"].(string)
367+
if !ok {
368+
return nil, xerrors.New("workspace is required")
369+
}
370+
371+
workspace, err := getWorkspaceByIDOrOwnerName(ctx, client, wsArg)
372+
if err != nil {
373+
return nil, xerrors.Errorf("failed to fetch workspace: %w", err)
374+
}
375+
376+
switch workspace.LatestBuild.Status {
377+
case codersdk.WorkspaceStatusPending, codersdk.WorkspaceStatusStarting, codersdk.WorkspaceStatusRunning, codersdk.WorkspaceStatusCanceling:
378+
return nil, xerrors.Errorf("workspace is %s", workspace.LatestBuild.Status)
379+
}
380+
381+
wb, err := client.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{
382+
Transition: codersdk.WorkspaceTransitionStart,
383+
})
384+
if err != nil {
385+
return nil, xerrors.Errorf("failed to start workspace: %w", err)
386+
}
387+
388+
resp := map[string]any{"status": wb.Status, "transition": wb.Transition}
389+
respJSON, err := json.Marshal(resp)
390+
if err != nil {
391+
return nil, xerrors.Errorf("failed to encode workspace build: %w", err)
392+
}
393+
394+
return &mcp.CallToolResult{
395+
Content: []mcp.Content{
396+
mcp.NewTextContent(string(respJSON)),
397+
},
398+
}, nil
399+
}
400+
}
401+
402+
// Example payload:
403+
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_stop_workspace", "arguments": {"workspace": "dev"}}}
404+
func handleCoderStopWorkspace(client *codersdk.Client) mcpserver.ToolHandlerFunc {
405+
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
406+
if client == nil {
407+
return nil, xerrors.New("developer error: client is required")
408+
}
409+
410+
args := request.Params.Arguments
411+
412+
wsArg, ok := args["workspace"].(string)
413+
if !ok {
414+
return nil, xerrors.New("workspace is required")
415+
}
416+
417+
workspace, err := getWorkspaceByIDOrOwnerName(ctx, client, wsArg)
418+
if err != nil {
419+
return nil, xerrors.Errorf("failed to fetch workspace: %w", err)
420+
}
421+
422+
switch workspace.LatestBuild.Status {
423+
case codersdk.WorkspaceStatusPending, codersdk.WorkspaceStatusStopping, codersdk.WorkspaceStatusStopped, codersdk.WorkspaceStatusCanceling:
424+
return nil, xerrors.Errorf("workspace is %s", workspace.LatestBuild.Status)
425+
}
426+
427+
wb, err := client.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{
428+
Transition: codersdk.WorkspaceTransitionStop,
429+
})
430+
if err != nil {
431+
return nil, xerrors.Errorf("failed to stop workspace: %w", err)
432+
}
433+
434+
resp := map[string]any{"status": wb.Status, "transition": wb.Transition}
435+
respJSON, err := json.Marshal(resp)
436+
if err != nil {
437+
return nil, xerrors.Errorf("failed to encode workspace build: %w", err)
438+
}
439+
440+
return &mcp.CallToolResult{
441+
Content: []mcp.Content{
442+
mcp.NewTextContent(string(respJSON)),
443+
},
444+
}, nil
445+
}
446+
}
447+
268448
func getWorkspaceByIDOrOwnerName(ctx context.Context, client *codersdk.Client, identifier string) (codersdk.Workspace, error) {
269449
if wsid, err := uuid.Parse(identifier); err == nil {
270450
return client.Workspace(ctx, wsid)

‎mcp/tools/tools_coder_test.go

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,29 @@ func TestCoderTools(t *testing.T) {
5757
mcptools.RegisterCoderReportTask(mcpSrv, memberClient, slogtest.Make(t, nil))
5858
mcptools.RegisterCoderWhoami(mcpSrv, memberClient)
5959
mcptools.RegisterCoderListWorkspaces(mcpSrv, memberClient)
60+
mcptools.RegisterCoderGetWorkspace(mcpSrv, memberClient)
6061
mcptools.RegisterCoderWorkspaceExec(mcpSrv, memberClient)
62+
mcptools.RegisterCoderListTemplates(mcpSrv, memberClient)
63+
mcptools.RegisterCoderStartWorkspace(mcpSrv, memberClient)
64+
mcptools.RegisterCoderStopWorkspace(mcpSrv, memberClient)
65+
66+
t.Run("coder_list_templates", func(t *testing.T) {
67+
// When: the coder_list_templates tool is called
68+
ctr := makeJSONRPCRequest(t, "tools/call", "coder_list_templates", map[string]any{})
69+
70+
pty.WriteLine(ctr)
71+
_ = pty.ReadLine(ctx) // skip the echo
72+
73+
templates, err := memberClient.Templates(ctx, codersdk.TemplateFilter{})
74+
require.NoError(t, err)
75+
templatesJSON, err := json.Marshal(templates)
76+
require.NoError(t, err)
77+
78+
// Then: the response is a list of templates visible to the user.
79+
expected := makeJSONRPCTextResponse(t, string(templatesJSON))
80+
actual := pty.ReadLine(ctx)
81+
testutil.RequireJSONEq(t, expected, actual)
82+
})
6183

6284
t.Run("coder_report_task", func(t *testing.T) {
6385
// When: the coder_report_task tool is called
@@ -119,6 +141,26 @@ func TestCoderTools(t *testing.T) {
119141
testutil.RequireJSONEq(t, expected, actual)
120142
})
121143

144+
t.Run("coder_get_workspace", func(t *testing.T) {
145+
// When: the coder_get_workspace tool is called
146+
ctr := makeJSONRPCRequest(t, "tools/call", "coder_get_workspace", map[string]any{
147+
"workspace": r.Workspace.ID.String(),
148+
})
149+
150+
pty.WriteLine(ctr)
151+
_ = pty.ReadLine(ctx) // skip the echo
152+
153+
ws, err := memberClient.Workspace(ctx, r.Workspace.ID)
154+
require.NoError(t, err)
155+
wsJSON, err := json.Marshal(ws)
156+
require.NoError(t, err)
157+
158+
// Then: the response is a valid JSON respresentation of the workspace.
159+
expected := makeJSONRPCTextResponse(t, string(wsJSON))
160+
actual := pty.ReadLine(ctx)
161+
testutil.RequireJSONEq(t, expected, actual)
162+
})
163+
122164
t.Run("coder_workspace_exec", func(t *testing.T) {
123165
// When: the coder_workspace_exec tools is called with a command
124166
randString := testutil.GetRandomName(t)
@@ -137,6 +179,50 @@ func TestCoderTools(t *testing.T) {
137179
actual := pty.ReadLine(ctx)
138180
testutil.RequireJSONEq(t, expected, actual)
139181
})
182+
183+
t.Run("coder_stop_workspace", func(t *testing.T) {
184+
// Given: a separate workspace in the running state
185+
stopWs := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
186+
OrganizationID: owner.OrganizationID,
187+
OwnerID: member.ID,
188+
}).WithAgent().Do()
189+
190+
// When: the coder_stop_workspace tool is called
191+
ctr := makeJSONRPCRequest(t, "tools/call", "coder_stop_workspace", map[string]any{
192+
"workspace": stopWs.Workspace.ID.String(),
193+
})
194+
195+
pty.WriteLine(ctr)
196+
_ = pty.ReadLine(ctx) // skip the echo
197+
198+
// Then: the response is as expected.
199+
expected := makeJSONRPCTextResponse(t, `{"status":"pending","transition":"stop"}`) // no provisionerd yet
200+
actual := pty.ReadLine(ctx)
201+
testutil.RequireJSONEq(t, expected, actual)
202+
})
203+
204+
t.Run("coder_start_workspace", func(t *testing.T) {
205+
// Given: a separate workspace in the stopped state
206+
stopWs := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
207+
OrganizationID: owner.OrganizationID,
208+
OwnerID: member.ID,
209+
}).Seed(database.WorkspaceBuild{
210+
Transition: database.WorkspaceTransitionStop,
211+
}).Do()
212+
213+
// When: the coder_start_workspace tool is called
214+
ctr := makeJSONRPCRequest(t, "tools/call", "coder_start_workspace", map[string]any{
215+
"workspace": stopWs.Workspace.ID.String(),
216+
})
217+
218+
pty.WriteLine(ctr)
219+
_ = pty.ReadLine(ctx) // skip the echo
220+
221+
// Then: the response is as expected
222+
expected := makeJSONRPCTextResponse(t, `{"status":"pending","transition":"start"}`) // no provisionerd yet
223+
actual := pty.ReadLine(ctx)
224+
testutil.RequireJSONEq(t, expected, actual)
225+
})
140226
}
141227

142228
// makeJSONRPCRequest is a helper function that makes a JSON RPC request.

0 commit comments

Comments
 (0)