Skip to content

Commit 9127d62

Browse files
johnstcnjohnstcn
committed
add cli exp mcp command
1 parent 8445d28 commit 9127d62

File tree

6 files changed

+188
-17
lines changed

6 files changed

+188
-17
lines changed

‎cli/exp.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ func (r *RootCmd) expCmd() *serpent.Command {
1313
Children: []*serpent.Command{
1414
r.scaletestCmd(),
1515
r.errorExample(),
16+
r.mcpCommand(),
1617
r.promptExample(),
1718
r.rptyCommand(),
1819
},

‎cli/exp_mcp.go

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package cli
2+
3+
import (
4+
"context"
5+
"errors"
6+
7+
"cdr.dev/slog"
8+
"cdr.dev/slog/sloggers/sloghuman"
9+
"github.com/coder/coder/v2/cli/cliui"
10+
"github.com/coder/coder/v2/codersdk"
11+
codermcp "github.com/coder/coder/v2/mcp"
12+
"github.com/coder/serpent"
13+
)
14+
15+
func (r *RootCmd) mcpCommand() *serpent.Command {
16+
var (
17+
client = new(codersdk.Client)
18+
instructions string
19+
)
20+
return &serpent.Command{
21+
Use: "mcp",
22+
Handler: func(inv *serpent.Invocation) error {
23+
return mcpHandler(inv, client, instructions)
24+
},
25+
Short: "Start an MCP server that can be used to interact with a Coder depoyment.",
26+
Middleware: serpent.Chain(
27+
r.InitClient(client),
28+
),
29+
Options: []serpent.Option{
30+
{
31+
Name: "instructions",
32+
Description: "The instructions to pass to the MCP server.",
33+
Flag: "instructions",
34+
Value: serpent.StringOf(&instructions),
35+
},
36+
},
37+
}
38+
}
39+
40+
func mcpHandler(inv *serpent.Invocation, client *codersdk.Client, instructions string) error {
41+
ctx, cancel := context.WithCancel(inv.Context())
42+
defer cancel()
43+
44+
logger := slog.Make(sloghuman.Sink(inv.Stdout))
45+
46+
me, err := client.User(ctx, codersdk.Me)
47+
if err != nil {
48+
cliui.Errorf(inv.Stderr, "Failed to log in to the Coder deployment.")
49+
cliui.Errorf(inv.Stderr, "Please check your URL and credentials.")
50+
cliui.Errorf(inv.Stderr, "Tip: Run `coder whoami` to check your credentials.")
51+
return err
52+
}
53+
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)
57+
cliui.Infof(inv.Stderr, "Press Ctrl+C to stop the server")
58+
59+
// Capture the original stdin, stdout, and stderr.
60+
invStdin := inv.Stdin
61+
invStdout := inv.Stdout
62+
invStderr := inv.Stderr
63+
defer func() {
64+
inv.Stdin = invStdin
65+
inv.Stdout = invStdout
66+
inv.Stderr = invStderr
67+
}()
68+
69+
closer := codermcp.New(ctx, client,
70+
codermcp.WithInstructions(instructions),
71+
codermcp.WithLogger(&logger),
72+
codermcp.WithStdin(invStdin),
73+
codermcp.WithStdout(invStdout),
74+
)
75+
76+
<-ctx.Done()
77+
if err := closer.Close(); err != nil {
78+
if !errors.Is(err, context.Canceled) {
79+
cliui.Errorf(inv.Stderr, "Failed to stop the MCP server: %s", err)
80+
return err
81+
}
82+
}
83+
return nil
84+
}

‎cli/exp_mcp_test.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package cli_test
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"testing"
7+
8+
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/require"
10+
11+
"github.com/coder/coder/v2/cli/clitest"
12+
"github.com/coder/coder/v2/coderd/coderdtest"
13+
"github.com/coder/coder/v2/pty/ptytest"
14+
"github.com/coder/coder/v2/testutil"
15+
)
16+
17+
func TestExpMcp(t *testing.T) {
18+
t.Parallel()
19+
20+
t.Run("OK", func(t *testing.T) {
21+
t.Parallel()
22+
23+
ctx := testutil.Context(t, testutil.WaitShort)
24+
cancelCtx, cancel := context.WithCancel(ctx)
25+
t.Cleanup(cancel)
26+
27+
client := coderdtest.New(t, nil)
28+
_ = coderdtest.CreateFirstUser(t, client)
29+
inv, root := clitest.New(t, "exp", "mcp")
30+
inv = inv.WithContext(cancelCtx)
31+
32+
pty := ptytest.New(t)
33+
inv.Stdin = pty.Input()
34+
inv.Stdout = pty.Output()
35+
clitest.SetupConfig(t, client, root)
36+
37+
cmdDone := tGo(t, func() {
38+
err := inv.Run()
39+
assert.NoError(t, err)
40+
})
41+
42+
payload := `{"jsonrpc":"2.0","id":1,"method":"initialize"}`
43+
pty.WriteLine(payload)
44+
_ = pty.ReadLine(ctx) // ignore echoed output
45+
output := pty.ReadLine(ctx)
46+
cancel()
47+
<-cmdDone
48+
49+
// Ensure the initialize output is valid JSON
50+
t.Logf("/initialize output: %s", output)
51+
var initializeResponse map[string]interface{}
52+
err := json.Unmarshal([]byte(output), &initializeResponse)
53+
require.NoError(t, err)
54+
require.Equal(t, "2.0", initializeResponse["jsonrpc"])
55+
require.Equal(t, 1.0, initializeResponse["id"])
56+
require.NotNil(t, initializeResponse["result"])
57+
})
58+
59+
t.Run("NoCredentials", func(t *testing.T) {
60+
t.Parallel()
61+
62+
ctx := testutil.Context(t, testutil.WaitShort)
63+
cancelCtx, cancel := context.WithCancel(ctx)
64+
t.Cleanup(cancel)
65+
66+
client := coderdtest.New(t, nil)
67+
inv, root := clitest.New(t, "exp", "mcp")
68+
inv = inv.WithContext(cancelCtx)
69+
70+
pty := ptytest.New(t)
71+
inv.Stdin = pty.Input()
72+
inv.Stdout = pty.Output()
73+
clitest.SetupConfig(t, client, root)
74+
75+
err := inv.Run()
76+
assert.ErrorContains(t, err, "your session has expired")
77+
})
78+
}

‎mcp/mcp.go

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,41 +21,40 @@ type mcpOptions struct {
2121
out io.Writer
2222
instructions string
2323
logger *slog.Logger
24-
client *codersdk.Client
2524
}
2625

26+
// Option is a function that configures the MCP server.
2727
type Option func(*mcpOptions)
2828

29+
// WithInstructions sets the instructions for the MCP server.
2930
func WithInstructions(instructions string) Option {
3031
return func(o *mcpOptions) {
3132
o.instructions = instructions
3233
}
3334
}
3435

36+
// WithLogger sets the logger for the MCP server.
3537
func WithLogger(logger *slog.Logger) Option {
3638
return func(o *mcpOptions) {
3739
o.logger = logger
3840
}
3941
}
4042

43+
// WithStdin sets the input reader for the MCP server.
4144
func WithStdin(in io.Reader) Option {
4245
return func(o *mcpOptions) {
4346
o.in = in
4447
}
4548
}
4649

50+
// WithStdout sets the output writer for the MCP server.
4751
func WithStdout(out io.Writer) Option {
4852
return func(o *mcpOptions) {
4953
o.out = out
5054
}
5155
}
5256

53-
func WithClient(client *codersdk.Client) Option {
54-
return func(o *mcpOptions) {
55-
o.client = client
56-
}
57-
}
58-
57+
// New creates a new MCP server with the given client and options.
5958
func New(ctx context.Context, client *codersdk.Client, opts ...Option) io.Closer {
6059
options := &mcpOptions{
6160
in: os.Stdin,
@@ -75,10 +74,10 @@ func New(ctx context.Context, client *codersdk.Client, opts ...Option) io.Closer
7574

7675
logger := slog.Make(sloghuman.Sink(os.Stdout))
7776

78-
mcptools.RegisterCoderReportTask(mcpSrv, options.client, logger)
79-
mcptools.RegisterCoderWhoami(mcpSrv, options.client)
80-
mcptools.RegisterCoderListWorkspaces(mcpSrv, options.client)
81-
mcptools.RegisterCoderWorkspaceExec(mcpSrv, options.client)
77+
mcptools.RegisterCoderReportTask(mcpSrv, client, logger)
78+
mcptools.RegisterCoderWhoami(mcpSrv, client)
79+
mcptools.RegisterCoderListWorkspaces(mcpSrv, client)
80+
mcptools.RegisterCoderWorkspaceExec(mcpSrv, client)
8281

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

‎mcp/tools/tools_coder.go

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,12 @@ Good Summaries:
7171

7272
// Example payload:
7373
// {"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"}}}
74-
func handleCoderReportTask(log slog.Logger, _ *codersdk.Client) mcpserver.ToolHandlerFunc {
74+
func handleCoderReportTask(log slog.Logger, client *codersdk.Client) mcpserver.ToolHandlerFunc {
7575
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
76+
if client == nil {
77+
return nil, xerrors.New("developer error: client is required")
78+
}
79+
7680
args := request.Params.Arguments
7781

7882
summary, ok := args["summary"].(string)
@@ -122,6 +126,9 @@ func handleCoderReportTask(log slog.Logger, _ *codersdk.Client) mcpserver.ToolHa
122126
// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_whoami", "arguments": {"coder_url": "http://localhost:3000", "coder_session_token": "REDACTED"}}}
123127
func handleCoderWhoami(client *codersdk.Client) mcpserver.ToolHandlerFunc {
124128
return func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
129+
if client == nil {
130+
return nil, xerrors.New("developer error: client is required")
131+
}
125132
me, err := client.User(ctx, codersdk.Me)
126133
if err != nil {
127134
return nil, xerrors.Errorf("Failed to fetch the current user: %s", err.Error())
@@ -144,6 +151,9 @@ func handleCoderWhoami(client *codersdk.Client) mcpserver.ToolHandlerFunc {
144151
// {"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"}}}
145152
func handleCoderListWorkspaces(client *codersdk.Client) mcpserver.ToolHandlerFunc {
146153
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
154+
if client == nil {
155+
return nil, xerrors.New("developer error: client is required")
156+
}
147157
args := request.Params.Arguments
148158

149159
owner, ok := args["owner"].(string)
@@ -187,6 +197,9 @@ func handleCoderListWorkspaces(client *codersdk.Client) mcpserver.ToolHandlerFun
187197
// {"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"}}}
188198
func handleCoderWorkspaceExec(client *codersdk.Client) mcpserver.ToolHandlerFunc {
189199
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
200+
if client == nil {
201+
return nil, xerrors.New("developer error: client is required")
202+
}
190203
args := request.Params.Arguments
191204

192205
wsArg, ok := args["workspace"].(string)

‎mcp/tools/tools_coder_test.go

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import (
2323

2424
// These tests are dependent on the state of the coder server.
2525
// Running them in parallel is prone to racy behavior.
26-
// nolint:tparallel
26+
// nolint:tparallel,paralleltest
2727
func TestCoderTools(t *testing.T) {
2828
t.Parallel()
2929

@@ -60,7 +60,6 @@ func TestCoderTools(t *testing.T) {
6060
mcptools.RegisterCoderWorkspaceExec(mcpSrv, memberClient)
6161

6262
t.Run("coder_report_task", func(t *testing.T) {
63-
// nolint:tparallel
6463
// When: the coder_report_task tool is called
6564
ctr := makeJSONRPCRequest(t, "tools/call", "coder_report_task", map[string]any{
6665
"summary": "Test summary",
@@ -82,7 +81,6 @@ func TestCoderTools(t *testing.T) {
8281
})
8382

8483
t.Run("coder_whoami", func(t *testing.T) {
85-
// nolint:tparallel
8684
// When: the coder_whoami tool is called
8785
me, err := memberClient.User(ctx, codersdk.Me)
8886
require.NoError(t, err)
@@ -101,7 +99,6 @@ func TestCoderTools(t *testing.T) {
10199
})
102100

103101
t.Run("coder_list_workspaces", func(t *testing.T) {
104-
// nolint:tparallel
105102
// When: the coder_list_workspaces tool is called
106103
ctr := makeJSONRPCRequest(t, "tools/call", "coder_list_workspaces", map[string]any{
107104
"coder_url": client.URL.String(),
@@ -123,7 +120,6 @@ func TestCoderTools(t *testing.T) {
123120
})
124121

125122
t.Run("coder_workspace_exec", func(t *testing.T) {
126-
// nolint:tparallel
127123
// When: the coder_workspace_exec tools is called with a command
128124
randString := testutil.GetRandomName(t)
129125
ctr := makeJSONRPCRequest(t, "tools/call", "coder_workspace_exec", map[string]any{

0 commit comments

Comments
 (0)