Move update predicates to push phase

code-asher · code-asher · commit 866b721a14d7 · 2025-06-12T14:37:26.000-08:00
Instead of the pop phase.  This ensures we do not queue up updates that
will just end up being discarded once they are popped (which could take
some time due to latency to coderd).

It also has the side effect of preserving summaries even when the queue
gets too big, because now we preserve them as part of pushing, before
they might get lost due to getting dropped while we wait on coderd.
diff --git a/cli/cliutil/queue.go b/cli/cliutil/queue.go
@@ -4,6 +4,8 @@ import (
 	"sync"
 
 	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/codersdk"
 )
 
 // Queue is a FIFO queue with a fixed size.  If the size is exceeded, the first
@@ -14,6 +16,7 @@ type Queue[T any] struct {
 	mu     sync.Mutex
 	size   int
 	closed bool
+	pred   func(x T) (T, bool)
 }
 
 // NewQueue creates a queue with the given size.
@@ -26,6 +29,13 @@ func NewQueue[T any](size int) *Queue[T] {
 	return q
 }
 
+// WithPredicate adds the given predicate function, which can control what is
+// pushed to the queue.
+func (q *Queue[T]) WithPredicate(pred func(x T) (T, bool)) *Queue[T] {
+	q.pred = pred
+	return q
+}
+
 // Close aborts any pending pops and makes future pushes error.
 func (q *Queue[T]) Close() {
 	q.mu.Lock()
@@ -41,6 +51,15 @@ func (q *Queue[T]) Push(x T) error {
 	if q.closed {
 		return xerrors.New("queue has been closed")
 	}
+	// Potentially mutate or skip the push using the predicate.
+	if q.pred != nil {
+		var ok bool
+		x, ok = q.pred(x)
+		if !ok {
+			return nil
+		}
+	}
+	// Remove the first item from the queue if it has gotten too big.
 	if len(q.items) >= q.size {
 		q.items = q.items[1:]
 	}
@@ -70,3 +89,72 @@ func (q *Queue[T]) Len() int {
 	defer q.mu.Unlock()
 	return len(q.items)
 }
+
+type reportTask struct {
+	link         string
+	messageID    int64
+	selfReported bool
+	state        codersdk.WorkspaceAppStatusState
+	summary      string
+}
+
+// statusQueue is a Queue that:
+// 1. Only pushes items that are not duplicates.
+// 2. Preserves the existing message and URI when one a message is not provided.
+// 3. Ignores "working" updates from the status watcher.
+type StatusQueue struct {
+	Queue[reportTask]
+	// lastMessageID is the ID of the last *user* message that we saw.  A user
+	// message only happens when interacting via the API (as opposed to
+	// interacting with the terminal directly).
+	lastMessageID int64
+}
+
+func (q *StatusQueue) Push(report reportTask) error {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	if q.closed {
+		return xerrors.New("queue has been closed")
+	}
+	var lastReport reportTask
+	if len(q.items) > 0 {
+		lastReport = q.items[len(q.items)-1]
+	}
+	// 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:
+	// 1. The LLM is still working, in which case our last status will already
+	//    have been "working", so there is nothing to do.
+	// 2. The user has interacted with the terminal directly.  For now, we are
+	//    ignoring these updates.  This risks missing cases where the user
+	//    manually submits a new prompt and the LLM 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.  In the future, if we can
+	//    reliably distinguish between user and LLM activity, we can change this.
+	if report.messageID > q.lastMessageID {
+		report.state = codersdk.WorkspaceAppStatusStateWorking
+	} else if report.state == codersdk.WorkspaceAppStatusStateWorking && !report.selfReported {
+		q.mu.Unlock()
+		return nil
+	}
+	// 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 queueing duplicate updates.
+	if report.state == lastReport.state &&
+		report.link == lastReport.link &&
+		report.summary == lastReport.summary {
+		return nil
+	}
+	// Drop the first item if the queue has gotten too big.
+	if len(q.items) >= q.size {
+		q.items = q.items[1:]
+	}
+	q.items = append(q.items, report)
+	q.cond.Broadcast()
+	return nil
+}
diff --git a/cli/cliutil/queue_test.go b/cli/cliutil/queue_test.go
@@ -82,4 +82,29 @@ func TestQueue(t *testing.T) {
 		err := q.Push(10)
 		require.Error(t, err)
 	})
+
+	t.Run("WithPredicate", func(t *testing.T) {
+		t.Parallel()
+
+		q := cliutil.NewQueue[int](10)
+		q.WithPredicate(func(n int) (int, bool) {
+			if n == 2 {
+				return n, false
+			}
+			return n + 1, true
+		})
+
+		for i := 0; i < 5; i++ {
+			err := q.Push(i)
+			require.NoError(t, err)
+		}
+
+		got := []int{}
+		for i := 0; i < 4; i++ {
+			val, ok := q.Pop()
+			require.True(t, ok)
+			got = append(got, val)
+		}
+		require.Equal(t, []int{1, 2, 4, 5}, got)
+	})
 }
diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go
@@ -361,7 +361,7 @@ func (*RootCmd) mcpConfigureCursor() *serpent.Command {
 	return cmd
 }
 
-type reportTask struct {
+type taskReport struct {
 	link         string
 	messageID    int64
 	selfReported bool
@@ -374,7 +374,7 @@ type mcpServer struct {
 	appStatusSlug string
 	client        *codersdk.Client
 	llmClient     *agentapi.Client
-	queue         *cliutil.Queue[reportTask]
+	queue         *cliutil.Queue[taskReport]
 }
 
 func (r *RootCmd) mcpServer() *serpent.Command {
@@ -388,9 +388,50 @@ 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 LLM agent API (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:
+				// 1. The LLM is still working, so there is nothing to update.
+				// 2. The LLM stopped working, then the user has interacted with the
+				//    terminal directly.  For now, we are ignoring these updates.  This
+				//    risks missing cases where the user manually submits a new prompt
+				//    and the LLM 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.  In the future, if we can reliably
+				//    distinguish between user and LLM activity, we can change this.
+				if report.messageID > lastUserMessageID {
+					report.state = codersdk.WorkspaceAppStatusStateWorking
+				} else if report.state == codersdk.WorkspaceAppStatusStateWorking && !report.selfReported {
+					return report, false
+				}
+				// 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 queueing duplicate updates.
+				if report.state == lastReport.state &&
+					report.link == lastReport.link &&
+					report.summary == lastReport.summary {
+					return report, false
+				}
+				lastReport = report
+				return report, true
+			})
+
 			srv := &mcpServer{
 				appStatusSlug: appStatusSlug,
-				queue:         cliutil.NewQueue[reportTask](100),
+				queue:         queue,
 			}
 
 			// Display client URL separately from authentication status.
@@ -505,35 +546,6 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 }
 
 func (s *mcpServer) startReporter(ctx context.Context, inv *serpent.Invocation) {
-	// lastMessageID is the ID of the last *user* message that we saw.  A user
-	// message only happens when interacting via the API (as opposed to
-	// interacting with the terminal directly).
-	var lastMessageID int64
-	shouldUpdate := func(item reportTask) codersdk.WorkspaceAppStatusState {
-		// Always send self-reported updates.
-		if item.selfReported {
-			return item.state
-		}
-		// Always send completed states.
-		switch item.state {
-		case codersdk.WorkspaceAppStatusStateComplete,
-			codersdk.WorkspaceAppStatusStateFailure:
-			return item.state
-		}
-		// Always send "working" when there is a new user message, since we know the
-		// LLM will begin work soon if it has not already.
-		if item.messageID > lastMessageID {
-			return codersdk.WorkspaceAppStatusStateWorking
-		}
-		// Otherwise, if the state is "working" and there have been no new user
-		// messages, it means either that the LLM is still working or it means the
-		// user has interacted with the terminal directly.  For now, we are ignoring
-		// these updates.  This risks missing cases where the user manually submits
-		// a new prompt and the LLM becomes active and does not update itself, but
-		// it avoids spamming useless status updates.
-		return ""
-	}
-	var lastPayload agentsdk.PatchAppStatus
 	go func() {
 		for {
 			// TODO: Even with the queue, there is still the potential that a message
@@ -545,45 +557,15 @@ func (s *mcpServer) startReporter(ctx context.Context, inv *serpent.Invocation)
 				return
 			}
 
-			state := shouldUpdate(item)
-			if state == "" {
-				continue
-			}
-
-			if item.messageID != 0 {
-				lastMessageID = item.messageID
-			}
-
-			payload := agentsdk.PatchAppStatus{
+			err := s.agentClient.PatchAppStatus(ctx, agentsdk.PatchAppStatus{
 				AppSlug: s.appStatusSlug,
 				Message: item.summary,
 				URI:     item.link,
-				State:   state,
-			}
-
-			// Preserve previous message and URI if there was no message.
-			if payload.Message == "" {
-				payload.Message = lastPayload.Message
-				if payload.URI == "" {
-					payload.URI = lastPayload.URI
-				}
-			}
-
-			// Avoid sending duplicate updates.
-			if lastPayload.State == payload.State &&
-				lastPayload.URI == payload.URI &&
-				lastPayload.Message == payload.Message {
-				continue
-			}
-
-			err := s.agentClient.PatchAppStatus(ctx, payload)
+				State:   item.state,
+			})
 			if err != nil && !errors.Is(err, context.Canceled) {
 				cliui.Warnf(inv.Stderr, "Failed to report task status: %s", err)
 			}
-
-			if err == nil {
-				lastPayload = payload
-			}
 		}
 	}()
 }
@@ -607,7 +589,7 @@ func (s *mcpServer) startWatcher(ctx context.Context, inv *serpent.Invocation) {
 					if ev.Status == agentapi.StatusStable {
 						state = codersdk.WorkspaceAppStatusStateComplete
 					}
-					err := s.queue.Push(reportTask{
+					err := s.queue.Push(taskReport{
 						state: state,
 					})
 					if err != nil {
@@ -616,7 +598,7 @@ func (s *mcpServer) startWatcher(ctx context.Context, inv *serpent.Invocation) {
 					}
 				case agentapi.EventMessageUpdate:
 					if ev.Role == agentapi.RoleUser {
-						err := s.queue.Push(reportTask{
+						err := s.queue.Push(taskReport{
 							messageID: ev.Id,
 						})
 						if err != nil {
@@ -667,7 +649,7 @@ func (s *mcpServer) startServer(ctx context.Context, inv *serpent.Invocation, in
 	// Add tool dependencies.
 	toolOpts := []func(*toolsdk.Deps){
 		toolsdk.WithTaskReporter(func(args toolsdk.ReportTaskArgs) error {
-			return s.queue.Push(reportTask{
+			return s.queue.Push(taskReport{
 				link:         args.Link,
 				selfReported: true,
 				state:        codersdk.WorkspaceAppStatusState(args.State),
