fix: handle new user messages correctly

code-asher · code-asher · commit 62b94b783118 · 2025-06-26T12:28:51.000-08:00
Two things:

1. Message IDs can be zero-based.
2. We were not actually updating the last user message ID.
diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go
@@ -363,7 +363,7 @@ func (*RootCmd) mcpConfigureCursor() *serpent.Command {
 
 type taskReport struct {
 	link         string
-	messageID    int64
+	messageID    *int64
 	selfReported bool
 	state        codersdk.WorkspaceAppStatusState
 	summary      string
@@ -391,7 +391,7 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 			// lastUserMessageID is the ID of the last *user* message that we saw.  A
 			// user message only happens when interacting via the AI AgentAPI (as
 			// opposed to interacting with the terminal directly).
-			var lastUserMessageID int64
+			lastUserMessageID := int64(-1)
 			var lastReport taskReport
 			// Create a queue that skips duplicates and preserves summaries.
 			queue := cliutil.NewQueue[taskReport](512).WithPredicate(func(report taskReport) (taskReport, bool) {
@@ -415,8 +415,9 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 				// user manually submits a new prompt and the AI agent becomes active
 				// (and does not update itself), but it avoids spamming useless status
 				// updates as the user is typing, so the tradeoff is worth it.
-				if report.messageID > lastUserMessageID {
+				if report.messageID != nil && *report.messageID > lastUserMessageID {
 					report.state = codersdk.WorkspaceAppStatusStateWorking
+					lastUserMessageID = *report.messageID
 				} else if report.state == codersdk.WorkspaceAppStatusStateWorking && !report.selfReported && lastReport.state != "" {
 					return report, false
 				}
@@ -607,7 +608,7 @@ func (s *mcpServer) startWatcher(ctx context.Context, inv *serpent.Invocation) {
 				case agentapi.EventMessageUpdate:
 					if ev.Role == agentapi.RoleUser {
 						err := s.queue.Push(taskReport{
-							messageID: ev.Id,
+							messageID: &ev.Id,
 						})
 						if err != nil {
 							cliui.Warnf(inv.Stderr, "Failed to queue update: %s", err)
diff --git a/cli/exp_mcp_test.go b/cli/exp_mcp_test.go
@@ -835,7 +835,7 @@ func TestExpMcpReporter(t *testing.T) {
 				},
 				// Agent messages are ignored.
 				{
-					event: makeMessageEvent(1, agentapi.RoleAgent),
+					event: makeMessageEvent(0, agentapi.RoleAgent),
 				},
 				// AI agent reports that it failed and URI is blank.
 				{
@@ -854,7 +854,7 @@ func TestExpMcpReporter(t *testing.T) {
 				// ... but this time we have a new user message so we know there is AI
 				// agent activity.  This time the "working" update will not be skipped.
 				{
-					event: makeMessageEvent(2, agentapi.RoleUser),
+					event: makeMessageEvent(1, agentapi.RoleUser),
 					expected: &codersdk.WorkspaceAppStatus{
 						State:   codersdk.WorkspaceAppStatusStateWorking,
 						Message: "oops",
@@ -895,6 +895,15 @@ func TestExpMcpReporter(t *testing.T) {
 						URI:     "",
 					},
 				},
+				// Zero ID should be accepted.
+				{
+					event: makeMessageEvent(0, agentapi.RoleUser),
+					expected: &codersdk.WorkspaceAppStatus{
+						State:   codersdk.WorkspaceAppStatusStateWorking,
+						Message: "",
+						URI:     "",
+					},
+				},
 			},
 		},
 	}
@@ -954,6 +963,7 @@ func TestExpMcpReporter(t *testing.T) {
 					case w, ok := <-watcher:
 						require.True(t, ok, "watch channel closed")
 						if w.LatestAppStatus != nil && w.LatestAppStatus.ID != lastAppStatus.ID {
+							t.Logf("Got status update: %s > %s", lastAppStatus.State, w.LatestAppStatus.State)
 							lastAppStatus = *w.LatestAppStatus
 							return lastAppStatus
 						}
