fix: handle new user messages correctly

code-asher · code-asher · commit 7454e4963fe9 · 2025-06-26T13:02:54.000-08:00
Two things:

1. Message IDs can be zero-based.
2. We were not actually updating the last user message ID.

I refactored a little to keep track of the last message in lastReport.
diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go
@@ -362,11 +362,19 @@ func (*RootCmd) mcpConfigureCursor() *serpent.Command {
 }
 
 type taskReport struct {
-	link         string
-	messageID    int64
+	// link is optional.
+	link string
+	// messageID must be set if this update is from a *user* message. A user
+	// message only happens when interacting via the AI AgentAPI (as opposed to
+	// interacting with the terminal directly).
+	messageID *int64
+	// selfReported must be set if the update is directly from the AI agent
+	// (as opposed to the screen watcher).
 	selfReported bool
-	state        codersdk.WorkspaceAppStatusState
-	summary      string
+	// state must always be set.
+	state codersdk.WorkspaceAppStatusState
+	// summary is optional.
+	summary string
 }
 
 type mcpServer struct {
@@ -388,17 +396,17 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 	return &serpent.Command{
 		Use: "server",
 		Handler: func(inv *serpent.Invocation) error {
-			// 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
 			var lastReport taskReport
 			// Create a queue that skips duplicates and preserves summaries.
 			queue := cliutil.NewQueue[taskReport](512).WithPredicate(func(report taskReport) (taskReport, bool) {
-				// Use "working" status if this is a new user message.  If this is not a
-				// new user message, and the status is "working" and not self-reported
-				// (meaning it came from the screen watcher), then it means one of two
-				// things:
+				// If this is a user message, discard if it is not new.
+				if report.messageID != nil &&
+					lastReport.messageID != nil && *lastReport.messageID >= *report.messageID {
+					return report, false
+				}
+				// If this is not a user message, and the status is "working" and not
+				// self-reported (meaning it came from the screen watcher), then it
+				// means one of two things:
 				//
 				// 1. The AI agent is not working; the user is interacting with the
 				//    terminal directly.
@@ -415,18 +423,27 @@ 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 {
-					report.state = codersdk.WorkspaceAppStatusStateWorking
-				} else if report.state == codersdk.WorkspaceAppStatusStateWorking && !report.selfReported && lastReport.state != "" {
+				if report.messageID == nil &&
+					report.state == codersdk.WorkspaceAppStatusStateWorking &&
+					!report.selfReported && lastReport.state != "" {
 					return report, false
 				}
+				// Keep track of the last message ID.
+				if report.messageID == nil {
+					report.messageID = lastReport.messageID
+				}
 				// Preserve previous message and URI if there was no message.
 				if report.summary == "" {
 					report.summary = lastReport.summary
 					if report.link == "" {
 						report.link = lastReport.link
 					}
 				}
+				// Avoid queuing empty statuses (this would probably indicate a
+				// developer error)
+				if report.state == "" {
+					return report, false
+				}
 				// Avoid queueing duplicate updates.
 				if report.state == lastReport.state &&
 					report.link == lastReport.link &&
@@ -607,7 +624,8 @@ 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,
+							state:     codersdk.WorkspaceAppStatusStateWorking,
 						})
 						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,33 @@ func TestExpMcpReporter(t *testing.T) {
 						URI:     "",
 					},
 				},
+				// Zero ID should be accepted.
+				{
+					event: makeMessageEvent(0, agentapi.RoleUser),
+					expected: &codersdk.WorkspaceAppStatus{
+						State:   codersdk.WorkspaceAppStatusStateWorking,
+						Message: "",
+						URI:     "",
+					},
+				},
+				// Stable again.
+				{
+					event: makeStatusEvent(agentapi.StatusStable),
+					expected: &codersdk.WorkspaceAppStatus{
+						State:   codersdk.WorkspaceAppStatusStateIdle,
+						Message: "",
+						URI:     "",
+					},
+				},
+				// Next ID.
+				{
+					event: makeMessageEvent(1, agentapi.RoleUser),
+					expected: &codersdk.WorkspaceAppStatus{
+						State:   codersdk.WorkspaceAppStatusStateWorking,
+						Message: "",
+						URI:     "",
+					},
+				},
 			},
 		},
 	}
@@ -954,6 +981,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
 						}
