Skip to content

Commit b02460d

Browse files
code-ashercode-asher
committed
Use common flags for agent client in MCP server
Before we were separately reading the environment, but we can reuse the global flags and the function for creating an agent client. Also, since the tests no longer need to set environment variables they can now be ran in parallel. This is not fully backwards compatible in that the old method would fall back to the non-agent client URL if the agent URL was not set, but that seems like undesired behavior since we are not doing that anywhere else.
1 parent 20b60a3 commit b02460d

File tree

2 files changed

+146
-147
lines changed

2 files changed

+146
-147
lines changed

‎cli/exp_mcp.go

Lines changed: 92 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"context"
66
"encoding/json"
77
"errors"
8-
"net/url"
98
"os"
109
"path/filepath"
1110
"slices"
@@ -24,6 +23,10 @@ import (
2423
"github.com/coder/serpent"
2524
)
2625

26+
const (
27+
envAppStatusSlug = "CODER_MCP_APP_STATUS_SLUG"
28+
)
29+
2730
func (r *RootCmd) mcpCommand() *serpent.Command {
2831
cmd := &serpent.Command{
2932
Use: "mcp",
@@ -110,7 +113,7 @@ func (*RootCmd) mcpConfigureClaudeDesktop() *serpent.Command {
110113
return cmd
111114
}
112115

113-
func (*RootCmd) mcpConfigureClaudeCode() *serpent.Command {
116+
func (r *RootCmd) mcpConfigureClaudeCode() *serpent.Command {
114117
var (
115118
claudeAPIKey string
116119
claudeConfigPath string
@@ -139,11 +142,12 @@ func (*RootCmd) mcpConfigureClaudeCode() *serpent.Command {
139142
binPath = testBinaryName
140143
}
141144
configureClaudeEnv := map[string]string{}
142-
agentToken, err := getAgentToken(fs)
145+
agentClient, err := r.createAgentClient()
143146
if err != nil {
144-
cliui.Warnf(inv.Stderr, "failed to get agent token: %s", err)
147+
cliui.Warnf(inv.Stderr, "failed to create agent client: %s", err)
145148
} else {
146-
configureClaudeEnv["CODER_AGENT_TOKEN"] = agentToken
149+
configureClaudeEnv[envAgentURL] = agentClient.SDK.URL.String()
150+
configureClaudeEnv[envAgentToken] = agentClient.SDK.SessionToken()
147151
}
148152
if claudeAPIKey == "" {
149153
if deprecatedCoderMCPClaudeAPIKey == "" {
@@ -154,7 +158,7 @@ func (*RootCmd) mcpConfigureClaudeCode() *serpent.Command {
154158
}
155159
}
156160
if appStatusSlug != "" {
157-
configureClaudeEnv["CODER_MCP_APP_STATUS_SLUG"] = appStatusSlug
161+
configureClaudeEnv[envAppStatusSlug] = appStatusSlug
158162
}
159163
if deprecatedSystemPromptEnv, ok := os.LookupEnv("SYSTEM_PROMPT"); ok {
160164
cliui.Warnf(inv.Stderr, "SYSTEM_PROMPT is deprecated, use CODER_MCP_CLAUDE_SYSTEM_PROMPT instead")
@@ -181,7 +185,7 @@ func (*RootCmd) mcpConfigureClaudeCode() *serpent.Command {
181185

182186
// Determine if we should include the reportTaskPrompt
183187
var reportTaskPrompt string
184-
if agentToken != "" && appStatusSlug != "" {
188+
if agentClient != nil && appStatusSlug != "" {
185189
// Only include the report task prompt if both agent token and app
186190
// status slug are defined. Otherwise, reporting a task will fail
187191
// and confuse the agent (and by extension, the user).
@@ -250,7 +254,7 @@ func (*RootCmd) mcpConfigureClaudeCode() *serpent.Command {
250254
{
251255
Name: "app-status-slug",
252256
Description: "The app status slug to use when running the Coder MCP server.",
253-
Env: "CODER_MCP_APP_STATUS_SLUG",
257+
Env: envAppStatusSlug,
254258
Flag: "claude-app-status-slug",
255259
Value: serpent.StringOf(&appStatusSlug),
256260
},
@@ -343,6 +347,12 @@ func (*RootCmd) mcpConfigureCursor() *serpent.Command {
343347
return cmd
344348
}
345349

350+
type mcpServer struct {
351+
agentClient *agentsdk.Client
352+
appStatusSlug string
353+
client *codersdk.Client
354+
}
355+
346356
func (r *RootCmd) mcpServer() *serpent.Command {
347357
var (
348358
client = new(codersdk.Client)
@@ -353,7 +363,53 @@ func (r *RootCmd) mcpServer() *serpent.Command {
353363
return &serpent.Command{
354364
Use: "server",
355365
Handler: func(inv *serpent.Invocation) error {
356-
return mcpServerHandler(inv, client, instructions, allowedTools, appStatusSlug)
366+
srv := &mcpServer{
367+
appStatusSlug: appStatusSlug,
368+
}
369+
370+
// Display client URL separately from authentication status.
371+
if client != nil && client.URL != nil {
372+
cliui.Infof(inv.Stderr, "URL : %s", client.URL.String())
373+
} else {
374+
cliui.Infof(inv.Stderr, "URL : Not configured")
375+
}
376+
377+
// Validate the client.
378+
if client != nil && client.URL != nil && client.SessionToken() != "" {
379+
me, err := client.User(inv.Context(), codersdk.Me)
380+
if err == nil {
381+
username := me.Username
382+
cliui.Infof(inv.Stderr, "Authentication : Successful")
383+
cliui.Infof(inv.Stderr, "User : %s", username)
384+
srv.client = client
385+
} else {
386+
cliui.Infof(inv.Stderr, "Authentication : Failed (%s)", err)
387+
cliui.Warnf(inv.Stderr, "Some tools that require authentication will not be available.")
388+
}
389+
} else {
390+
cliui.Infof(inv.Stderr, "Authentication : None")
391+
}
392+
393+
// Try to create an agent client for status reporting. Not validated.
394+
agentClient, err := r.createAgentClient()
395+
if err == nil {
396+
cliui.Infof(inv.Stderr, "Agent URL : %s", agentClient.SDK.URL.String())
397+
srv.agentClient = agentClient
398+
}
399+
if err != nil || appStatusSlug == "" {
400+
cliui.Infof(inv.Stderr, "Task reporter : Disabled")
401+
if err != nil {
402+
cliui.Warnf(inv.Stderr, "%s", err)
403+
}
404+
if appStatusSlug == "" {
405+
cliui.Warnf(inv.Stderr, "%s must be set", envAppStatusSlug)
406+
}
407+
} else {
408+
cliui.Infof(inv.Stderr, "Task reporter : Enabled")
409+
}
410+
411+
// Start the server.
412+
return srv.start(inv, instructions, allowedTools)
357413
},
358414
Short: "Start the Coder MCP server.",
359415
Middleware: serpent.Chain(
@@ -378,54 +434,24 @@ func (r *RootCmd) mcpServer() *serpent.Command {
378434
Name: "app-status-slug",
379435
Description: "When reporting a task, the coder_app slug under which to report the task.",
380436
Flag: "app-status-slug",
381-
Env: "CODER_MCP_APP_STATUS_SLUG",
437+
Env: envAppStatusSlug,
382438
Value: serpent.StringOf(&appStatusSlug),
383439
Default: "",
384440
},
385441
},
386442
}
387443
}
388444

389-
func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instructions string, allowedTools []string, appStatusSlug string) error {
445+
func (s *mcpServer) start(inv *serpent.Invocation, instructions string, allowedTools []string) error {
390446
ctx, cancel := context.WithCancel(inv.Context())
391447
defer cancel()
392448

393-
fs := afero.NewOsFs()
394-
395449
cliui.Infof(inv.Stderr, "Starting MCP server")
396450

397-
// Check authentication status
398-
var username string
399-
400-
// Check authentication status first
401-
if client != nil && client.URL != nil && client.SessionToken() != "" {
402-
// Try to validate the client
403-
me, err := client.User(ctx, codersdk.Me)
404-
if err == nil {
405-
username = me.Username
406-
cliui.Infof(inv.Stderr, "Authentication : Successful")
407-
cliui.Infof(inv.Stderr, "User : %s", username)
408-
} else {
409-
// Authentication failed but we have a client URL
410-
cliui.Warnf(inv.Stderr, "Authentication : Failed (%s)", err)
411-
cliui.Warnf(inv.Stderr, "Some tools that require authentication will not be available.")
412-
}
413-
} else {
414-
cliui.Infof(inv.Stderr, "Authentication : None")
415-
}
416-
417-
// Display URL separately from authentication status
418-
if client != nil && client.URL != nil {
419-
cliui.Infof(inv.Stderr, "URL : %s", client.URL.String())
420-
} else {
421-
cliui.Infof(inv.Stderr, "URL : Not configured")
422-
}
423-
424451
cliui.Infof(inv.Stderr, "Instructions : %q", instructions)
425452
if len(allowedTools) > 0 {
426453
cliui.Infof(inv.Stderr, "Allowed Tools : %v", allowedTools)
427454
}
428-
cliui.Infof(inv.Stderr, "Press Ctrl+C to stop the server")
429455

430456
// Capture the original stdin, stdout, and stderr.
431457
invStdin := inv.Stdin
@@ -443,68 +469,44 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
443469
server.WithInstructions(instructions),
444470
)
445471

446-
// Get the workspace agent token from the environment.
447-
toolOpts := make([]func(*toolsdk.Deps), 0)
448-
var hasAgentClient bool
449-
450-
var agentURL *url.URL
451-
if client != nil && client.URL != nil {
452-
agentURL = client.URL
453-
} else if agntURL, err := getAgentURL(); err == nil {
454-
agentURL = agntURL
455-
}
456-
457-
// First check if we have a valid client URL, which is required for agent client
458-
if agentURL == nil {
459-
cliui.Infof(inv.Stderr, "Agent URL : Not configured")
460-
} else {
461-
cliui.Infof(inv.Stderr, "Agent URL : %s", agentURL.String())
462-
agentToken, err := getAgentToken(fs)
463-
if err != nil || agentToken == "" {
464-
cliui.Warnf(inv.Stderr, "CODER_AGENT_TOKEN is not set, task reporting will not be available")
465-
} else {
466-
// Happy path: we have both URL and agent token
467-
agentClient := agentsdk.New(agentURL)
468-
agentClient.SetSessionToken(agentToken)
469-
toolOpts = append(toolOpts, toolsdk.WithAgentClient(agentClient))
470-
hasAgentClient = true
471-
}
472-
}
473-
474-
if (client == nil || client.URL == nil || client.SessionToken() == "") && !hasAgentClient {
472+
// If both clients are unauthorized, there are no tools we can enable.
473+
if s.client == nil && s.agentClient == nil {
475474
return xerrors.New(notLoggedInMessage)
476475
}
477476

478-
if appStatusSlug != "" {
479-
toolOpts = append(toolOpts, toolsdk.WithAppStatusSlug(appStatusSlug))
480-
} else {
481-
cliui.Warnf(inv.Stderr, "CODER_MCP_APP_STATUS_SLUG is not set, task reporting will not be available.")
477+
// Add tool dependencies.
478+
toolOpts := []func(*toolsdk.Deps){
479+
toolsdk.WithAgentClient(s.agentClient),
480+
toolsdk.WithAppStatusSlug(s.appStatusSlug),
482481
}
483482

484-
toolDeps, err := toolsdk.NewDeps(client, toolOpts...)
483+
toolDeps, err := toolsdk.NewDeps(s.client, toolOpts...)
485484
if err != nil {
486485
return xerrors.Errorf("failed to initialize tool dependencies: %w", err)
487486
}
488487

489-
// Register tools based on the allowlist (if specified)
488+
// Register tools based on the allowlist. Zero length means allow everything.
490489
for _, tool := range toolsdk.All {
491-
// Skip adding the coder_report_task tool if there is no agent client
492-
if !hasAgentClient && tool.Tool.Name == "coder_report_task" {
493-
cliui.Warnf(inv.Stderr, "Task reporting not available")
490+
// Skip if not allowed.
491+
if len(allowedTools) > 0 && !slices.ContainsFunc(allowedTools, func(t string) bool {
492+
return t == tool.Tool.Name
493+
}) {
494494
continue
495495
}
496496

497-
// Skip user-dependent tools if no authenticated user
498-
if !tool.UserClientOptional && username == "" {
497+
// Skip user-dependent tools if no authenticated user client.
498+
if !tool.UserClientOptional && s.client == nil {
499499
cliui.Warnf(inv.Stderr, "Tool %q requires authentication and will not be available", tool.Tool.Name)
500500
continue
501501
}
502502

503-
if len(allowedTools) == 0 || slices.ContainsFunc(allowedTools, func(t string) bool {
504-
return t == tool.Tool.Name
505-
}) {
506-
mcpSrv.AddTools(mcpFromSDK(tool, toolDeps))
503+
// Skip the coder_report_task tool if there is no agent client or slug.
504+
if tool.Tool.Name == "coder_report_task" && (s.agentClient == nil || s.appStatusSlug == "") {
505+
cliui.Warnf(inv.Stderr, "Tool %q requires the task reporter and will not be available", tool.Tool.Name)
506+
continue
507507
}
508+
509+
mcpSrv.AddTools(mcpFromSDK(tool, toolDeps))
508510
}
509511

510512
srv := server.NewStdioServer(mcpSrv)
@@ -515,11 +517,11 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
515517
done <- srvErr
516518
}()
517519

518-
if err := <-done; err != nil {
519-
if !errors.Is(err, context.Canceled) {
520-
cliui.Errorf(inv.Stderr, "Failed to start the MCP server: %s", err)
521-
return err
522-
}
520+
cliui.Infof(inv.Stderr, "Press Ctrl+C to stop the server")
521+
522+
if err := <-done; err != nil && !errors.Is(err, context.Canceled) {
523+
cliui.Errorf(inv.Stderr, "Failed to start the MCP server: %s", err)
524+
return err
523525
}
524526

525527
return nil
@@ -738,31 +740,6 @@ func indexOf(s, substr string) int {
738740
return -1
739741
}
740742

741-
func getAgentToken(fs afero.Fs) (string, error) {
742-
token, ok := os.LookupEnv("CODER_AGENT_TOKEN")
743-
if ok && token != "" {
744-
return token, nil
745-
}
746-
tokenFile, ok := os.LookupEnv("CODER_AGENT_TOKEN_FILE")
747-
if !ok {
748-
return "", xerrors.Errorf("CODER_AGENT_TOKEN or CODER_AGENT_TOKEN_FILE must be set for token auth")
749-
}
750-
bs, err := afero.ReadFile(fs, tokenFile)
751-
if err != nil {
752-
return "", xerrors.Errorf("failed to read agent token file: %w", err)
753-
}
754-
return string(bs), nil
755-
}
756-
757-
func getAgentURL() (*url.URL, error) {
758-
urlString, ok := os.LookupEnv("CODER_AGENT_URL")
759-
if !ok || urlString == "" {
760-
return nil, xerrors.New("CODEDR_AGENT_URL is empty")
761-
}
762-
763-
return url.Parse(urlString)
764-
}
765-
766743
// mcpFromSDK adapts a toolsdk.Tool to go-mcp's server.ServerTool.
767744
// It assumes that the tool responds with a valid JSON object.
768745
func mcpFromSDK(sdkTool toolsdk.GenericTool, tb toolsdk.Deps) server.ServerTool {

0 commit comments

Comments
 (0)