From 16add939087ed02c199e1b93b065381ba376e15b Mon Sep 17 00:00:00 2001
From: Cian Johnston <cian@coder.com>
Date: Thu, 2 Apr 2026 12:29:47 +0100
Subject: [PATCH] fix(coderd/x/chatd): stabilize subagent pubsub completion
 test (#23944)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- stabilize `TestAwaitSubagentCompletion/CompletesViaPubsub` by waiting
for durable completion state before sending the synthetic pubsub wake
- add coverage for successful subagent completion with an empty report

> 🤖 Written by a Coder Agent. Reviewed by a human.
---
 coderd/x/chatd/subagent_internal_test.go | 67 +++++++++++++++++++++---
 1 file changed, 59 insertions(+), 8 deletions(-)

diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go
index 655a606b7c..7c0645457d 100644
--- a/coderd/x/chatd/subagent_internal_test.go
+++ b/coderd/x/chatd/subagent_internal_test.go
@@ -1012,8 +1012,10 @@ func TestAwaitSubagentCompletion(t *testing.T) {
 		setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
 
 		// Trap the fallback poll ticker to know when the
-		// function has subscribed to pubsub and entered
-		// its select loop.
+		// function has entered the wait setup path. We still
+		// need an explicit subscription handshake below because
+		// the ticker can be created before SubscribeWithErr has
+		// finished registering the listener.
 		tickTrap := mClock.Trap().NewTicker("chatd", "subagent_poll")
 
 		type awaitResult struct {
@@ -1029,19 +1031,47 @@ func TestAwaitSubagentCompletion(t *testing.T) {
 			resultCh <- awaitResult{chat, report, err}
 		}()
 
-		// Wait for the ticker to be created (confirms pubsub
-		// subscription is set up and select loop entered).
+		// Wait for the ticker to be created so the waiter has
+		// entered its setup path, then subscribe our own probe on
+		// the same channel. Because MemoryPubsub publishes only to
+		// listeners already present at Publish time, waiting for
+		// our probe to receive a message proves the waiter's
+		// subscription is also registered before we assert on the
+		// wake-up behavior.
 		tickTrap.MustWait(ctx).MustRelease(ctx)
 		tickTrap.Close()
 
-		// Transition child and publish. The pubsub notification
-		// wakes the function without needing a clock advance.
+		probeCh := make(chan struct{}, 1)
+		cancelProbe, err := ps.SubscribeWithErr(
+			coderdpubsub.ChatStreamNotifyChannel(child.ID),
+			func(_ context.Context, _ []byte, _ error) {
+				select {
+				case probeCh <- struct{}{}:
+				default:
+				}
+			},
+		)
+		require.NoError(t, err)
+		defer cancelProbe()
+
+		// Transition the child first, then publish once the
+		// durable completion state is observable. Pubsub only
+		// wakes the waiter; it does not guarantee the report is
+		// visible in the same instant as the notification.
 		setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
 		insertAssistantMessage(ctx, t, db, child.ID, model.ID, "pubsub result")
-		_ = ps.Publish(
+		require.EventuallyWithT(t, func(c *assert.CollectT) {
+			chat, report, done, err := server.checkSubagentCompletion(ctx, child.ID)
+			require.NoError(c, err)
+			assert.True(c, done)
+			assert.Equal(c, child.ID, chat.ID)
+			assert.Equal(c, "pubsub result", report)
+		}, testutil.WaitMedium, testutil.IntervalFast)
+		require.NoError(t, ps.Publish(
 			coderdpubsub.ChatStreamNotifyChannel(child.ID),
 			[]byte("done"),
-		)
+		))
+		testutil.RequireReceive(ctx, t, probeCh)
 
 		result := testutil.RequireReceive(ctx, t, resultCh)
 		require.NoError(t, result.err)
@@ -1049,6 +1079,27 @@ func TestAwaitSubagentCompletion(t *testing.T) {
 		assert.Equal(t, "pubsub result", result.report)
 	})
 
+	t.Run("AlreadyWaitingNoReport", func(t *testing.T) {
+		t.Parallel()
+		ctx := chatdTestContext(t)
+
+		parent, child := createParentChildChats(ctx, t, server, user, model)
+
+		// signalWake from CreateChat may trigger immediate processing.
+		// Wait for it to settle, then set the terminal state we need.
+		// This case should return immediately, so use the shared
+		// real-clock server instead of a mock clock.
+		server.inflight.Wait()
+		setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
+
+		gotChat, report, err := server.awaitSubagentCompletion(
+			ctx, parent.ID, child.ID, 5*time.Second,
+		)
+		require.NoError(t, err)
+		assert.Equal(t, child.ID, gotChat.ID)
+		assert.Empty(t, report)
+	})
+
 	t.Run("Timeout", func(t *testing.T) {
 		t.Parallel()