Only override agent status if agent api is configured

code-asher · code-asher · commit 45a56850fd79 · 2025-06-30T11:00:26.000-08:00
diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go
@@ -674,16 +674,22 @@ func (s *mcpServer) startServer(ctx context.Context, inv *serpent.Invocation, in
 	}
 
 	// Add tool dependencies.
+	hasAgentAPI := s.aiAgentAPIClient != nil
 	toolOpts := []func(*toolsdk.Deps){
 		toolsdk.WithTaskReporter(func(args toolsdk.ReportTaskArgs) error {
+			// The agent does not reliably report its status correctly.  If AgentAPI
+			// is enabled, we will always set the status to "working" when we get an
+			// MCP message, and rely on the screen watcher to eventually catch the
+			// idle state.
+			state := codersdk.WorkspaceAppStatusStateWorking
+			if !hasAgentAPI {
+				state = codersdk.WorkspaceAppStatusState(args.State)
+			}
 			return s.queue.Push(taskReport{
 				link:         args.Link,
 				selfReported: true,
-				// The agent does not reliably report its status correctly.  We will
-				// always set the status to "working" when we get an MCP message, and
-				// rely on the screen watcher to eventually catch the idle state.
-				state:   codersdk.WorkspaceAppStatusStateWorking,
-				summary: args.Summary,
+				state:        state,
+				summary:      args.Summary,
 			})
 		}),
 	}
diff --git a/cli/exp_mcp_test.go b/cli/exp_mcp_test.go
@@ -793,8 +793,9 @@ func TestExpMcpReporter(t *testing.T) {
 	}
 
 	runs := []struct {
-		name  string
-		tests []test
+		name            string
+		tests           []test
+		disableAgentAPI bool
 	}{
 		// In this run the AI agent starts with a state change but forgets to update
 		// that it finished.
@@ -954,6 +955,29 @@ func TestExpMcpReporter(t *testing.T) {
 				},
 			},
 		},
+		// When AgentAPI is not being used, we accept agent state updates as-is.
+		{
+			name: "KeepAgentState",
+			tests: []test{
+				{
+					state:   codersdk.WorkspaceAppStatusStateWorking,
+					summary: "doing work",
+					expected: &codersdk.WorkspaceAppStatus{
+						State:   codersdk.WorkspaceAppStatusStateWorking,
+						Message: "doing work",
+					},
+				},
+				{
+					state:   codersdk.WorkspaceAppStatusStateIdle,
+					summary: "finished",
+					expected: &codersdk.WorkspaceAppStatus{
+						State:   codersdk.WorkspaceAppStatusStateIdle,
+						Message: "finished",
+					},
+				},
+			},
+			disableAgentAPI: true,
+		},
 	}
 
 	for _, run := range runs {
@@ -1019,15 +1043,19 @@ func TestExpMcpReporter(t *testing.T) {
 				}
 			}
 
-			inv, _ := clitest.New(t,
+			args := []string{
 				"exp", "mcp", "server",
-				// We need the agent credentials, AI AgentAPI url, and a slug for reporting.
+				// We need the agent credentials, AI AgentAPI url (if not
+				// disabled), and a slug for reporting.
 				"--agent-url", client.URL.String(),
 				"--agent-token", r.AgentToken,
 				"--app-status-slug", "vscode",
-				"--ai-agentapi-url", aiAgentAPIURL,
 				"--allowed-tools=coder_report_task",
-			)
+			}
+			if !run.disableAgentAPI {
+				args = append(args, "--ai-agentapi-url", aiAgentAPIURL)
+			}
+			inv, _ := clitest.New(t, args...)
 			inv = inv.WithContext(ctx)
 
 			pty := ptytest.New(t)
