From e8f124de85b47e19152cd538d03765325087f2c2 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 1 Apr 2025 14:33:03 +0100 Subject: [PATCH] fix(mcp): actually report task status --- cli/exp_mcp.go | 72 ++++++++++++++++++++------ mcp/mcp.go | 135 +++++++++++++++++------------------------------- mcp/mcp_test.go | 50 ++++++++++++++---- 3 files changed, 144 insertions(+), 113 deletions(-) diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go index a5af41d9103a6..b46a8b4d7f03a 100644 --- a/cli/exp_mcp.go +++ b/cli/exp_mcp.go @@ -4,14 +4,18 @@ import ( "context" "encoding/json" "errors" - "log" "os" "path/filepath" + "github.com/mark3labs/mcp-go/server" + "golang.org/x/xerrors" + "cdr.dev/slog" "cdr.dev/slog/sloggers/sloghuman" + "github.com/coder/coder/v2/buildinfo" "github.com/coder/coder/v2/cli/cliui" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agentsdk" codermcp "github.com/coder/coder/v2/mcp" "github.com/coder/serpent" ) @@ -191,14 +195,16 @@ func (*RootCmd) mcpConfigureCursor() *serpent.Command { func (r *RootCmd) mcpServer() *serpent.Command { var ( - client = new(codersdk.Client) - instructions string - allowedTools []string + client = new(codersdk.Client) + instructions string + allowedTools []string + appStatusSlug string + mcpServerAgent bool ) return &serpent.Command{ Use: "server", Handler: func(inv *serpent.Invocation) error { - return mcpServerHandler(inv, client, instructions, allowedTools) + return mcpServerHandler(inv, client, instructions, allowedTools, appStatusSlug, mcpServerAgent) }, Short: "Start the Coder MCP server.", Middleware: serpent.Chain( @@ -209,24 +215,39 @@ func (r *RootCmd) mcpServer() *serpent.Command { Name: "instructions", Description: "The instructions to pass to the MCP server.", Flag: "instructions", + Env: "CODER_MCP_INSTRUCTIONS", Value: serpent.StringOf(&instructions), }, { Name: "allowed-tools", Description: "Comma-separated list of allowed tools. If not specified, all tools are allowed.", Flag: "allowed-tools", + Env: "CODER_MCP_ALLOWED_TOOLS", Value: serpent.StringArrayOf(&allowedTools), }, + { + Name: "app-status-slug", + Description: "When reporting a task, the coder_app slug under which to report the task.", + Flag: "app-status-slug", + Env: "CODER_MCP_APP_STATUS_SLUG", + Value: serpent.StringOf(&appStatusSlug), + Default: "", + }, + { + Flag: "agent", + Env: "CODER_MCP_SERVER_AGENT", + Description: "Start the MCP server in agent mode, with a different set of tools.", + Value: serpent.BoolOf(&mcpServerAgent), + }, }, } } -func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instructions string, allowedTools []string) error { +//nolint:revive // control coupling +func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instructions string, allowedTools []string, appStatusSlug string, mcpServerAgent bool) error { ctx, cancel := context.WithCancel(inv.Context()) defer cancel() - logger := slog.Make(sloghuman.Sink(inv.Stdout)) - me, err := client.User(ctx, codersdk.Me) if err != nil { cliui.Errorf(inv.Stderr, "Failed to log in to the Coder deployment.") @@ -253,19 +274,40 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct inv.Stderr = invStderr }() - options := []codermcp.Option{ - codermcp.WithInstructions(instructions), - codermcp.WithLogger(&logger), + mcpSrv := server.NewMCPServer( + "Coder Agent", + buildinfo.Version(), + server.WithInstructions(instructions), + ) + + // Create a separate logger for the tools. + toolLogger := slog.Make(sloghuman.Sink(invStderr)) + + toolDeps := codermcp.ToolDeps{ + Client: client, + Logger: &toolLogger, + AppStatusSlug: appStatusSlug, + AgentClient: agentsdk.New(client.URL), + } + + if mcpServerAgent { + // Get the workspace agent token from the environment. + agentToken, ok := os.LookupEnv("CODER_AGENT_TOKEN") + if !ok || agentToken == "" { + return xerrors.New("CODER_AGENT_TOKEN is not set") + } + toolDeps.AgentClient.SetSessionToken(agentToken) } - // Add allowed tools option if specified + // Register tools based on the allowlist (if specified) + reg := codermcp.AllTools() if len(allowedTools) > 0 { - options = append(options, codermcp.WithAllowedTools(allowedTools)) + reg = reg.WithOnlyAllowed(allowedTools...) } - srv := codermcp.NewStdio(client, options...) - srv.SetErrorLogger(log.New(invStderr, "", log.LstdFlags)) + reg.Register(mcpSrv, toolDeps) + srv := server.NewStdioServer(mcpSrv) done := make(chan error) go func() { defer close(done) diff --git a/mcp/mcp.go b/mcp/mcp.go index 80e0f341e16e6..0dd01ccdc5fdd 100644 --- a/mcp/mcp.go +++ b/mcp/mcp.go @@ -6,7 +6,6 @@ import ( "encoding/json" "errors" "io" - "os" "slices" "strings" "time" @@ -17,76 +16,12 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog" - "cdr.dev/slog/sloggers/sloghuman" - "github.com/coder/coder/v2/buildinfo" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/codersdk/workspacesdk" ) -type mcpOptions struct { - instructions string - logger *slog.Logger - allowedTools []string -} - -// Option is a function that configures the MCP server. -type Option func(*mcpOptions) - -// WithInstructions sets the instructions for the MCP server. -func WithInstructions(instructions string) Option { - return func(o *mcpOptions) { - o.instructions = instructions - } -} - -// WithLogger sets the logger for the MCP server. -func WithLogger(logger *slog.Logger) Option { - return func(o *mcpOptions) { - o.logger = logger - } -} - -// WithAllowedTools sets the allowed tools for the MCP server. -func WithAllowedTools(tools []string) Option { - return func(o *mcpOptions) { - o.allowedTools = tools - } -} - -// NewStdio creates a new MCP stdio server with the given client and options. -// It is the responsibility of the caller to start and stop the server. -func NewStdio(client *codersdk.Client, opts ...Option) *server.StdioServer { - options := &mcpOptions{ - instructions: ``, - logger: ptr.Ref(slog.Make(sloghuman.Sink(os.Stdout))), - } - for _, opt := range opts { - opt(options) - } - - mcpSrv := server.NewMCPServer( - "Coder Agent", - buildinfo.Version(), - server.WithInstructions(options.instructions), - ) - - logger := slog.Make(sloghuman.Sink(os.Stdout)) - - // Register tools based on the allowed list (if specified) - reg := AllTools() - if len(options.allowedTools) > 0 { - reg = reg.WithOnlyAllowed(options.allowedTools...) - } - reg.Register(mcpSrv, ToolDeps{ - Client: client, - Logger: &logger, - }) - - srv := server.NewStdioServer(mcpSrv) - return srv -} - // allTools is the list of all available tools. When adding a new tool, // make sure to update this list. var allTools = ToolRegistry{ @@ -120,6 +55,8 @@ Choose an emoji that helps the user understand the current phase at a glance.`), mcp.WithBoolean("done", mcp.Description(`Whether the overall task the user requested is complete. Set to true only when the entire requested operation is finished successfully. For multi-step processes, use false until all steps are complete.`), mcp.Required()), + mcp.WithBoolean("need_user_attention", mcp.Description(`Whether the user needs to take action on the task. +Set to true if the task is in a failed state or if the user needs to take action to continue.`), mcp.Required()), ), MakeHandler: handleCoderReportTask, }, @@ -265,8 +202,10 @@ Can be either "start" or "stop".`)), // ToolDeps contains all dependencies needed by tool handlers type ToolDeps struct { - Client *codersdk.Client - Logger *slog.Logger + Client *codersdk.Client + AgentClient *agentsdk.Client + Logger *slog.Logger + AppStatusSlug string } // ToolHandler associates a tool with its handler creation function @@ -313,18 +252,23 @@ func AllTools() ToolRegistry { } type handleCoderReportTaskArgs struct { - Summary string `json:"summary"` - Link string `json:"link"` - Emoji string `json:"emoji"` - Done bool `json:"done"` + Summary string `json:"summary"` + Link string `json:"link"` + Emoji string `json:"emoji"` + Done bool `json:"done"` + NeedUserAttention bool `json:"need_user_attention"` } // Example payload: -// {"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}}} +// {"jsonrpc":"2.0","id":1,"method":"tools/call", "params": {"name": "coder_report_task", "arguments": {"summary": "I need help with the login page.", "link": "https://github.com/coder/coder/pull/1234", "emoji": "🔍", "done": false, "need_user_attention": true}}} func handleCoderReportTask(deps ToolDeps) server.ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { - if deps.Client == nil { - return nil, xerrors.New("developer error: client is required") + if deps.AgentClient == nil { + return nil, xerrors.New("developer error: agent client is required") + } + + if deps.AppStatusSlug == "" { + return nil, xerrors.New("No app status slug provided, set CODER_MCP_APP_STATUS_SLUG when running the MCP server to report tasks.") } // Convert the request parameters to a json.RawMessage so we can unmarshal @@ -334,20 +278,33 @@ func handleCoderReportTask(deps ToolDeps) server.ToolHandlerFunc { return nil, xerrors.Errorf("failed to unmarshal arguments: %w", err) } - // TODO: Waiting on support for tasks. - deps.Logger.Info(ctx, "report task tool called", slog.F("summary", args.Summary), slog.F("link", args.Link), slog.F("done", args.Done), slog.F("emoji", args.Emoji)) - /* - err := sdk.PostTask(ctx, agentsdk.PostTaskRequest{ - Reporter: "claude", - Summary: summary, - URL: link, - Completion: done, - Icon: emoji, - }) - if err != nil { - return nil, err - } - */ + deps.Logger.Info(ctx, "report task tool called", + slog.F("summary", args.Summary), + slog.F("link", args.Link), + slog.F("emoji", args.Emoji), + slog.F("done", args.Done), + slog.F("need_user_attention", args.NeedUserAttention), + ) + + newStatus := agentsdk.PatchAppStatus{ + AppSlug: deps.AppStatusSlug, + Message: args.Summary, + URI: args.Link, + Icon: args.Emoji, + NeedsUserAttention: args.NeedUserAttention, + State: codersdk.WorkspaceAppStatusStateWorking, + } + + if args.Done { + newStatus.State = codersdk.WorkspaceAppStatusStateComplete + } + if args.NeedUserAttention { + newStatus.State = codersdk.WorkspaceAppStatusStateFailure + } + + if err := deps.AgentClient.PatchAppStatus(ctx, newStatus); err != nil { + return nil, xerrors.Errorf("failed to patch app status: %w", err) + } return &mcp.CallToolResult{ Content: []mcp.Content{ diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 1144d9265aa15..c5cf000efcfa3 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -17,7 +17,9 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agentsdk" codermcp "github.com/coder/coder/v2/mcp" + "github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/pty/ptytest" "github.com/coder/coder/v2/testutil" ) @@ -39,7 +41,14 @@ func TestCoderTools(t *testing.T) { r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ OrganizationID: owner.OrganizationID, OwnerID: member.ID, - }).WithAgent().Do() + }).WithAgent(func(agents []*proto.Agent) []*proto.Agent { + agents[0].Apps = []*proto.App{ + { + Slug: "some-agent-app", + }, + } + return agents + }).Do() // Note: we want to test the list_workspaces tool before starting the // workspace agent. Starting the workspace agent will modify the workspace @@ -65,9 +74,12 @@ func TestCoderTools(t *testing.T) { // Register tools using our registry logger := slogtest.Make(t, nil) + agentClient := agentsdk.New(memberClient.URL) codermcp.AllTools().Register(mcpSrv, codermcp.ToolDeps{ - Client: memberClient, - Logger: &logger, + Client: memberClient, + Logger: &logger, + AppStatusSlug: "some-agent-app", + AgentClient: agentClient, }) t.Run("coder_list_templates", func(t *testing.T) { @@ -86,24 +98,44 @@ func TestCoderTools(t *testing.T) { }) t.Run("coder_report_task", func(t *testing.T) { + // Given: the MCP server has an agent token. + oldAgentToken := agentClient.SDK.SessionToken() + agentClient.SetSessionToken(r.AgentToken) + t.Cleanup(func() { + agentClient.SDK.SetSessionToken(oldAgentToken) + }) // When: the coder_report_task tool is called ctr := makeJSONRPCRequest(t, "tools/call", "coder_report_task", map[string]any{ "summary": "Test summary", "link": "https://example.com", "emoji": "🔍", "done": false, - "coder_url": client.URL.String(), - "coder_session_token": client.SessionToken(), + "need_user_attention": true, }) pty.WriteLine(ctr) _ = pty.ReadLine(ctx) // skip the echo - // Then: the response is a success message. - // TODO: check the task was created. This functionality is not yet implemented. - expected := makeJSONRPCTextResponse(t, "Thanks for reporting!") + // Then: positive feedback is given to the reporting agent. actual := pty.ReadLine(ctx) - testutil.RequireJSONEq(t, expected, actual) + require.Contains(t, actual, "Thanks for reporting!") + + // Then: the response is a success message. + ws, err := memberClient.Workspace(ctx, r.Workspace.ID) + require.NoError(t, err, "failed to get workspace") + agt, err := memberClient.WorkspaceAgent(ctx, ws.LatestBuild.Resources[0].Agents[0].ID) + require.NoError(t, err, "failed to get workspace agent") + require.NotEmpty(t, agt.Apps, "workspace agent should have an app") + require.NotEmpty(t, agt.Apps[0].Statuses, "workspace agent app should have a status") + st := agt.Apps[0].Statuses[0] + // require.Equal(t, ws.ID, st.WorkspaceID, "workspace app status should have the correct workspace id") + require.Equal(t, agt.ID, st.AgentID, "workspace app status should have the correct agent id") + require.Equal(t, agt.Apps[0].ID, st.AppID, "workspace app status should have the correct app id") + require.Equal(t, codersdk.WorkspaceAppStatusStateFailure, st.State, "workspace app status should be in the failure state") + require.Equal(t, "Test summary", st.Message, "workspace app status should have the correct message") + require.Equal(t, "https://example.com", st.URI, "workspace app status should have the correct uri") + require.Equal(t, "🔍", st.Icon, "workspace app status should have the correct icon") + require.True(t, st.NeedsUserAttention, "workspace app status should need user attention") }) t.Run("coder_whoami", func(t *testing.T) {