cbec491 · crush · git.iko.soy

git.iko.soy/crush

fix(server): complete queued prompts independently

Ensure queued prompts with their own submitted identity emit one terminal completion when canceled or dropped. Callers no longer wait on a completion that was folded into another prompt's turn. Co-Authored-By: Charm Crush <crush@charm.land>

commitcbec4916b3d9bcdaeaff75157071e5e87a7526ef authorChristian Rocha <christian@rocha.is> dateMon 1 Jun 2026, 20:35 (4 months ago) parentd7a814c

3 files changed+513−45

compare with main

unifiedsplit

internal/agent/agent.go+138−28 internal/agent/queued_runid_test.go+181 internal/agent/run_complete_test.go+194−17

internal/agent/agent.go+138−28

@@ -359,28 +359,96 @@ func (a *sessionAgent) enqueueCall(call SessionAgentCall) {

359359 a.messageQueue.Set(call.SessionID, existing)

360360}

361361

362−// drainUncanceledQueue copies and clears the session's message queue and

363−// returns the queued calls not covered by a pending cancel. The queue

364−// copy/delete and the cancel-mark check happen under the per-session

365−// dispatch mutex so the filtering is atomic against a concurrent Cancel:

366−// canceledBySeq requires the caller to hold that mutex, and evaluating it

367−// here (rather than after unlocking) prevents a cancel recorded between

368−// the drain and the check from being observed inconsistently. The

369−// returned survivors are processed by the caller without the lock held.

370−func (a *sessionAgent) drainUncanceledQueue(sessionID string) []SessionAgentCall {

362+// drainQueueForStep partitions the session's queued calls for the current

363+// streaming step under the per-session dispatch mutex so the filtering is

364+// atomic against a concurrent Cancel: canceledBySeq requires the caller to

365+// hold that mutex, and evaluating it here (rather than after unlocking)

366+// prevents a cancel recorded between the drain and the check from being

367+// observed inconsistently.

368+//

369+// Calls covered by a pending cancel are dropped; the dropped ones that

370+// carry a RunID are returned in canceledWithRunID so the caller can

371+// publish their terminal cancelled RunComplete (a caller waiting on that

372+// RunID, e.g. `crush run`, would otherwise hang). Uncanceled calls without

373+// a RunID are returned in fold to be folded into the active turn,

374+// preserving the existing follow-up behavior. Uncanceled calls that carry

375+// a RunID are left in the queue so each runs as its own turn via the

376+// recursive run path and publishes its own RunComplete, giving every

377+// RunID-bearing prompt an explicit lifecycle instead of being silently

378+// absorbed into another turn. fold is processed by the caller without the

379+// lock held.

380+func (a *sessionAgent) drainQueueForStep(sessionID string) (fold, canceledWithRunID []SessionAgentCall) {

371381 dispatchLock := a.sessionMu(sessionID)

372382 dispatchLock.Lock()

373383 defer dispatchLock.Unlock()

374384 queuedCalls, _ := a.messageQueue.Get(sessionID)

375− a.messageQueue.Del(sessionID)

376− survivors := queuedCalls[:0]

385+ var keep []SessionAgentCall

377386 for _, queued := range queuedCalls {

378387 if a.canceledBySeq(sessionID, queued.acceptSeq) {

388+ if queued.RunID != "" {

389+ canceledWithRunID = append(canceledWithRunID, queued)

390+ }

391+ continue

392+ }

393+ if queued.RunID != "" {

394+ keep = append(keep, queued)

395+ continue

396+ }

397+ fold = append(fold, queued)

398+ }

399+ if len(keep) == 0 {

400+ a.messageQueue.Del(sessionID)

401+ } else {

402+ a.messageQueue.Set(sessionID, keep)

403+ }

404+ return fold, canceledWithRunID

405+}

406+

407+// publishCanceledQueueDrops emits a terminal cancelled RunComplete for

408+// every dropped queued call that carries a RunID. A queued prompt removed

409+// from the queue without ever running — covered by a pending cancel, or

410+// cleared by Cancel/ClearQueue — would otherwise leave a caller blocked on

411+// that RunID: `crush run` ignores live message events and exits only on a

412+// RunComplete whose RunID matches. Calls without a RunID had no such waiter

413+// and are dropped silently as before. A detached, bounded context keeps the

414+// must-deliver publish alive even when the run context that triggered the

415+// drop is already canceled.

416+func (a *sessionAgent) publishCanceledQueueDrops(drops []SessionAgentCall) {

417+ var hasRunID bool

418+ for _, d := range drops {

419+ if d.RunID != "" {

420+ hasRunID = true

421+ break

422+ }

423+ }

424+ if !hasRunID {

425+ return

426+ }

427+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

428+ defer cancel()

429+ for _, d := range drops {

430+ if d.RunID == "" {

379431 continue

380432 }

381− survivors = append(survivors, queued)

433+ a.publishRunComplete(ctx, d, notify.RunComplete{

434+ SessionID: d.SessionID,

435+ RunID: d.RunID,

436+ Cancelled: true,

437+ })

382438 }

383− return survivors

439+}

440+

441+// clearQueueAndNotify removes all queued prompts for the session and

442+// publishes a terminal cancelled RunComplete for any that carried a RunID,

443+// so callers waiting on those RunIDs (e.g. `crush run`) are not left

444+// hanging when their queued prompt is discarded without running.

445+func (a *sessionAgent) clearQueueAndNotify(sessionID string) {

446+ queued, ok := a.messageQueue.Get(sessionID)

447+ a.messageQueue.Del(sessionID)

448+ if !ok {

449+ return

450+ }

451+ a.publishCanceledQueueDrops(queued)

384452}

385453

386454// clearPendingCancel removes any pending-cancel mark for sessionID. It

@@ -735,14 +803,20 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

735803 // Use latest tools (updated by SetTools when MCP tools change).

736804 prepared.Tools = a.tools.Copy()

737805

738− // Drain queued follow-up prompts, but skip any covered by a

739− // cancel recorded while they sat in the queue: a cancel that

740− // arrived after a prompt was queued must not let it run as

741− // part of this step. Coverage is per-call by accept sequence

742− // so a follow-up queued after the cancel (higher seq) is

743− // still folded in.

744− survivors := a.drainUncanceledQueue(call.SessionID)

745− for _, queued := range survivors {

806+ // Drain queued follow-up prompts for this step. Calls covered

807+ // by a cancel recorded while they sat in the queue are dropped:

808+ // a cancel that arrived after a prompt was queued must not let

809+ // it run as part of this step. Coverage is per-call by accept

810+ // sequence so a follow-up queued after the cancel (higher seq)

811+ // is not dropped. A dropped prompt carrying a RunID still gets

812+ // its terminal cancelled RunComplete so a caller waiting on it

813+ // does not hang. Uncanceled prompts without a RunID are folded

814+ // into this turn; uncanceled prompts with a RunID are left

815+ // queued so each runs as its own turn (with its own

816+ // RunComplete) via the recursive run path below.

817+ fold, canceledRunIDs := a.drainQueueForStep(call.SessionID)

818+ a.publishCanceledQueueDrops(canceledRunIDs)

819+ for _, queued := range fold {

746820 userMessage, createErr := a.createUserMessage(callContext, queued)

747821 if createErr != nil {

748822 return callContext, prepared, createErr

@@ -1125,14 +1199,22 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

11251199 // mark, or untracked); keep any queued after the cancel (higher

11261200 // sequence) so they still run.

11271201 var kept []SessionAgentCall

1202+ var canceledRunIDDrops []SessionAgentCall

11281203 for _, q := range queuedMessages {

11291204 if q.acceptSeq == 0 || q.acceptSeq <= mark {

1205+ if q.RunID != "" {

1206+ canceledRunIDDrops = append(canceledRunIDDrops, q)

1207+ }

11301208 continue

11311209 }

11321210 kept = append(kept, q)

11331211 }

11341212 queuedMessages = kept

11351213 a.messageQueue.Set(call.SessionID, kept)

1214+ // A dropped prompt carrying a RunID must still publish its

1215+ // terminal cancelled RunComplete so a caller waiting on that

1216+ // RunID does not hang.

1217+ a.publishCanceledQueueDrops(canceledRunIDDrops)

11361218 }

11371219 if len(queuedMessages) == 0 {

11381220 // No queued work. Clear the cancel mark only when no accepted

@@ -1151,12 +1233,29 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

11511233 mu.Unlock()

11521234 return result, err

11531235 }

1154− // There are queued messages restart the loop. The recursive Run

1155− // publishes its own RunComplete for the queued prompt, so suppress

1156− // the outer defer's emit to avoid a duplicate event whose Error

1157− // field would belong to the recursive turn but whose MessageID/Text

1158− // would belong to the outer turn.

1236+ // There are queued messages, restart the loop. Suppress the outer

1237+ // defer's emit: it would otherwise observe the recursive Run's retErr

1238+ // (named-return clobbering through the return below) against this

1239+ // turn's MessageID/Text and publish a mixed, racing event.

11591240 skipRunComplete = true

1241+ // Decide whether this turn still owes its own terminal RunComplete.

1242+ // Each submitted prompt with a RunID has its own lifecycle, so a turn

1243+ // that is finished and handing off to a *different* queued prompt must

1244+ // publish its own RunComplete here — leaving it to the recursive turn

1245+ // (which carries a different RunID) would hang a caller waiting on

1246+ // this turn's RunID. The exception is the summarize-continuation path,

1247+ // which re-queues this same call (same RunID) to resume after a

1248+ // summary; in that case the eventual terminal turn for this RunID

1249+ // publishes, so publishing now would double-emit.

1250+ outerOwesRunComplete := call.RunID != ""

1251+ if outerOwesRunComplete {

1252+ for _, q := range queuedMessages {

1253+ if q.RunID == call.RunID {

1254+ outerOwesRunComplete = false

1255+ break

1256+ }

1257+ }

1258+ }

11601259 firstQueuedMessage := queuedMessages[0]

11611260 a.messageQueue.Set(call.SessionID, queuedMessages[1:])

11621261 // Reserve a fresh accept for the dequeued prompt before dropping the

@@ -1167,6 +1266,17 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

11671266 // the recursive Run's accepted path observes as cancel-on-entry.

11681267 firstQueuedMessage.Accepted = a.BeginAccepted(call.SessionID)

11691268 mu.Unlock()

1269+ if outerOwesRunComplete {

1270+ complete := notify.RunComplete{SessionID: call.SessionID, RunID: call.RunID}

1271+ if currentAssistant != nil {

1272+ complete.MessageID = currentAssistant.ID

1273+ complete.Text = currentAssistant.Content().String()

1274+ }

1275+ if ctx.Err() != nil {

1276+ complete.Cancelled = true

1277+ }

1278+ a.publishRunComplete(ctx, call, complete)

1279+ }

11701280 return a.Run(ctx, firstQueuedMessage)

11711281}

11721282

@@ -1770,14 +1880,14 @@ func (a *sessionAgent) Cancel(sessionID string) {

17701880

17711881 if a.QueuedPrompts(sessionID) > 0 {

17721882 slog.Debug("Clearing queued prompts", "session_id", sessionID)

1773− a.messageQueue.Del(sessionID)

1883+ a.clearQueueAndNotify(sessionID)

17741884 }

17751885}

17761886

17771887func (a *sessionAgent) ClearQueue(sessionID string) {

17781888 if a.QueuedPrompts(sessionID) > 0 {

17791889 slog.Debug("Clearing queued prompts", "session_id", sessionID)

1780− a.messageQueue.Del(sessionID)

1890+ a.clearQueueAndNotify(sessionID)

17811891 }

17821892}

17831893

internal/agent/queued_runid_test.gonew file+181

@@ -0,0 +1,181 @@

1+package agent

3+import (

4+ "context"

5+ "errors"

6+ "sync/atomic"

7+ "testing"

8+ "time"

10+ "charm.land/catwalk/pkg/catwalk"

11+ "charm.land/fantasy"

12+ "github.com/charmbracelet/crush/internal/agent/notify"

13+ "github.com/charmbracelet/crush/internal/message"

14+ "github.com/charmbracelet/crush/internal/pubsub"

15+ "github.com/stretchr/testify/require"

16+)

17+

18+// gatedStreamModel streams a single text part followed by a clean finish,

19+// but blocks the very first Stream call until its gate is released. That

20+// lets a test hold a run "active" (past PrepareStep, inside Stream) just

21+// long enough to enqueue a follow-up prompt behind the busy session.

22+// Subsequent Stream calls (e.g. the recursive run draining the queue)

23+// proceed immediately.

24+type gatedStreamModel struct {

25+ text string

26+ gate chan struct{}

27+ entered chan struct{}

28+ calls atomic.Int64

29+}

30+

31+func (m *gatedStreamModel) Provider() string { return "fake" }

32+func (m *gatedStreamModel) Model() string { return "fake-model" }

33+

34+func (m *gatedStreamModel) Generate(ctx context.Context, call fantasy.Call) (*fantasy.Response, error) {

35+ return &fantasy.Response{

36+ Content: fantasy.ResponseContent{fantasy.TextContent{Text: m.text}},

37+ FinishReason: fantasy.FinishReasonStop,

38+ }, nil

39+}

40+

41+func (m *gatedStreamModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) {

42+ if m.calls.Add(1) == 1 {

43+ close(m.entered)

44+ select {

45+ case <-m.gate:

46+ case <-ctx.Done():

47+ }

48+ }

49+ text := m.text

50+ return func(yield func(fantasy.StreamPart) bool) {

51+ if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "1"}) {

52+ return

53+ }

54+ if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "1", Delta: text}) {

55+ return

56+ }

57+ if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "1"}) {

58+ return

59+ }

60+ yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop})

61+ }, nil

62+}

63+

64+func (m *gatedStreamModel) GenerateObject(ctx context.Context, call fantasy.ObjectCall) (*fantasy.ObjectResponse, error) {

65+ return nil, errors.New("not implemented")

66+}

67+

68+func (m *gatedStreamModel) StreamObject(ctx context.Context, call fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) {

69+ return nil, errors.New("not implemented")

70+}

71+

72+// TestRun_QueuedRunIDPromptRunsRecursivelyAndPublishesRunComplete is the

73+// end-to-end proof of fix 2: a prompt carrying a RunID that is queued

74+// behind a busy session must NOT be silently folded into the active turn.

75+// It runs as its own turn via the recursive run path and publishes its

76+// own terminal RunComplete, so a `crush run` caller blocking on that

77+// RunID does not hang. The active turn keeps its own RunComplete too.

78+func TestRun_QueuedRunIDPromptRunsRecursivelyAndPublishesRunComplete(t *testing.T) {

79+ t.Parallel()

80+

81+ env := testEnv(t)

82+ broker := pubsub.NewBroker[notify.RunComplete]()

83+ t.Cleanup(broker.Shutdown)

84+

85+ large := &gatedStreamModel{

86+ text: "done",

87+ gate: make(chan struct{}),

88+ entered: make(chan struct{}),

89+ }

90+ small := &finishStreamModel{text: "title"}

91+

92+ sa := NewSessionAgent(SessionAgentOptions{

93+ LargeModel: Model{Model: large, CatwalkCfg: catwalk.Model{ContextWindow: 200000, DefaultMaxTokens: 10000}},

94+ SmallModel: Model{Model: small, CatwalkCfg: catwalk.Model{ContextWindow: 200000, DefaultMaxTokens: 10000}},

95+ IsYolo: true,

96+ Sessions: env.sessions,

97+ Messages: env.messages,

98+ RunComplete: broker,

99+ }).(*sessionAgent)

100+

101+ sess, err := env.sessions.Create(t.Context(), "session")

102+ require.NoError(t, err)

103+

104+ subCtx, subCancel := context.WithCancel(t.Context())

105+ defer subCancel()

106+ ch := broker.Subscribe(subCtx)

107+

108+ // Start the main turn; it blocks inside Stream once active.

109+ mainDone := make(chan error, 1)

110+ go func() {

111+ _, runErr := sa.Run(t.Context(), SessionAgentCall{

112+ SessionID: sess.ID,

113+ RunID: "run-main",

114+ Prompt: "main",

115+ })

116+ mainDone <- runErr

117+ }()

118+

119+ // Wait until the main turn is active (inside Stream).

120+ select {

121+ case <-large.entered:

122+ case <-time.After(5 * time.Second):

123+ t.Fatal("main run never entered Stream")

124+ }

125+ require.True(t, sa.IsSessionBusy(sess.ID), "main run must be active before enqueueing the follow-up")

126+

127+ // Enqueue a RunID-bearing follow-up behind the busy session.

128+ res, err := sa.Run(t.Context(), SessionAgentCall{

129+ SessionID: sess.ID,

130+ RunID: "run-follow",

131+ Prompt: "follow",

132+ })

133+ require.NoError(t, err)

134+ require.Nil(t, res, "a busy-session follow-up must enqueue and return (nil, nil)")

135+ require.Equal(t, 1, sa.QueuedPrompts(sess.ID), "the follow-up must be queued, not folded")

136+

137+ // Release the main turn so it completes and hands off to the queue.

138+ close(large.gate)

139+ require.NoError(t, <-mainDone)

140+

141+ // Both turns must publish their own terminal RunComplete.

142+ got := map[string]notify.RunComplete{}

143+ deadline := time.After(5 * time.Second)

144+ for len(got) < 2 {

145+ select {

146+ case ev := <-ch:

147+ got[ev.Payload.RunID] = ev.Payload

148+ case <-deadline:

149+ t.Fatalf("timed out waiting for both RunCompletes; got %v", got)

150+ }

151+ }

152+

153+ main, ok := got["run-main"]

154+ require.True(t, ok, "the active turn must publish its own RunComplete")

155+ require.Empty(t, main.Error)

156+ require.False(t, main.Cancelled)

157+

158+ follow, ok := got["run-follow"]

159+ require.True(t, ok,

160+ "the queued RunID prompt must publish its own RunComplete instead of being folded silently")

161+ require.Empty(t, follow.Error)

162+ require.False(t, follow.Cancelled)

163+ require.Equal(t, "done", follow.Text, "the queued prompt ran as its own turn")

164+

165+ // Two distinct assistant turns prove the follow-up was not folded.

166+ msgs, err := env.messages.List(t.Context(), sess.ID)

167+ require.NoError(t, err)

168+ var assistants, follows int

169+ for _, m := range msgs {

170+ switch m.Role {

171+ case message.Assistant:

172+ assistants++

173+ case message.User:

174+ if m.Content().String() == "follow" {

175+ follows++

176+ }

177+ }

178+ }

179+ require.Equal(t, 2, assistants, "the active turn and the recursive turn each produce one assistant message")

180+ require.Equal(t, 1, follows, "the follow-up prompt is its own user turn")

181+}

internal/agent/run_complete_test.go+194−17

@@ -58,13 +58,15 @@ func TestSessionAgentRun_QueueStripsOnComplete(t *testing.T) {

5858 "RunComplete still correlates with the originating SendMessage")

5959}

6060

61−// TestDrainUncanceledQueue_FiltersUnderDispatchLock verifies that the

62−// queue drain evaluates the per-session cancel mark while holding the

63−// dispatch mutex (canceledBySeq's documented precondition). Queued calls

64−// at or below the cancel high-water mark are dropped, calls queued after

65−// the cancel (higher seq) survive, untracked enqueues (seq == 0) are

66−// dropped whenever any mark is present, and the queue is cleared.

67−func TestDrainUncanceledQueue_FiltersUnderDispatchLock(t *testing.T) {

61+// TestDrainQueueForStep_FiltersUnderDispatchLock verifies that the queue

62+// drain evaluates the per-session cancel mark while holding the dispatch

63+// mutex (canceledBySeq's documented precondition). Queued calls at or

64+// below the cancel high-water mark are dropped, calls queued after the

65+// cancel (higher seq) are folded, untracked enqueues (seq == 0) are

66+// dropped whenever any mark is present, and the queue is cleared. These

67+// calls carry no RunID, so all foldable survivors are returned for

68+// folding (the existing follow-up behavior).

69+func TestDrainQueueForStep_FiltersUnderDispatchLock(t *testing.T) {

6870 t.Parallel()

6971

7072 env := testEnv(t)

@@ -83,19 +85,21 @@ func TestDrainUncanceledQueue_FiltersUnderDispatchLock(t *testing.T) {

8385 // Cancel high-water mark at seq 2: seq <= 2 and seq == 0 are covered.

8486 a.cancelMark.Set(sessionID, 2)

8587

86− survivors := a.drainUncanceledQueue(sessionID)

88+ fold, canceledWithRunID := a.drainQueueForStep(sessionID)

8789

88− require.Len(t, survivors, 1,

89− "only the follow-up queued after the cancel (seq > mark) must survive")

90− require.Equal(t, "after", survivors[0].Prompt)

90+ require.Len(t, fold, 1,

91+ "only the follow-up queued after the cancel (seq > mark) must be folded")

92+ require.Equal(t, "after", fold[0].Prompt)

93+ require.Empty(t, canceledWithRunID,

94+ "no dropped call carried a RunID, so none need a terminal RunComplete")

9195

9296 _, ok := a.messageQueue.Get(sessionID)

93− require.False(t, ok, "drain must clear the session message queue")

97+ require.False(t, ok, "drain must clear the session message queue when nothing is kept")

9498}

9599

96−// TestDrainUncanceledQueue_NoMarkKeepsAll verifies that with no cancel

97−// mark recorded, every queued call survives the drain.

98−func TestDrainUncanceledQueue_NoMarkKeepsAll(t *testing.T) {

100+// TestDrainQueueForStep_NoMarkFoldsAllNonRunID verifies that with no

101+// cancel mark recorded, every queued call without a RunID is folded.

102+func TestDrainQueueForStep_NoMarkFoldsAllNonRunID(t *testing.T) {

99103 t.Parallel()

100104

101105 env := testEnv(t)

@@ -110,8 +114,80 @@ func TestDrainUncanceledQueue_NoMarkKeepsAll(t *testing.T) {

110114 {SessionID: sessionID, Prompt: "b", acceptSeq: 5},

111115 })

112116

113− survivors := a.drainUncanceledQueue(sessionID)

114− require.Len(t, survivors, 2, "no cancel mark means all queued calls survive")

117+ fold, canceledWithRunID := a.drainQueueForStep(sessionID)

118+ require.Len(t, fold, 2, "no cancel mark means all non-RunID queued calls are folded")

119+ require.Empty(t, canceledWithRunID)

120+}

121+

122+// TestDrainQueueForStep_KeepsRunIDPromptsQueued is the core of fix 2: a

123+// queued prompt that carries a RunID must NOT be folded into the active

124+// turn. Folding it would silently absorb it into another turn and never

125+// publish a RunComplete for its RunID, hanging a `crush run` caller that

126+// blocks on that event. Such prompts are left in the queue so the

127+// recursive run path gives each its own turn and its own RunComplete.

128+// Non-RunID prompts are still folded.

129+func TestDrainQueueForStep_KeepsRunIDPromptsQueued(t *testing.T) {

130+ t.Parallel()

131+

132+ env := testEnv(t)

133+ a := NewSessionAgent(SessionAgentOptions{

134+ Sessions: env.sessions,

135+ Messages: env.messages,

136+ }).(*sessionAgent)

137+

138+ const sessionID = "drain-runid"

139+ a.messageQueue.Set(sessionID, []SessionAgentCall{

140+ {SessionID: sessionID, Prompt: "fold-me", acceptSeq: 1},

141+ {SessionID: sessionID, RunID: "run-a", Prompt: "keep-me", acceptSeq: 2},

142+ {SessionID: sessionID, RunID: "run-b", Prompt: "keep-me-too", acceptSeq: 3},

143+ })

144+

145+ fold, canceledWithRunID := a.drainQueueForStep(sessionID)

146+

147+ require.Len(t, fold, 1, "only the non-RunID prompt is folded into the active turn")

148+ require.Equal(t, "fold-me", fold[0].Prompt)

149+ require.Empty(t, canceledWithRunID)

150+

151+ kept, ok := a.messageQueue.Get(sessionID)

152+ require.True(t, ok, "RunID-bearing prompts must remain queued for the recursive run path")

153+ require.Len(t, kept, 2)

154+ require.Equal(t, "run-a", kept[0].RunID)

155+ require.Equal(t, "run-b", kept[1].RunID)

156+}

157+

158+// TestDrainQueueForStep_ReportsCanceledRunIDDrops verifies that a queued

159+// prompt carrying a RunID that is dropped because a cancel covers it is

160+// reported in canceledWithRunID so the caller can publish its terminal

161+// cancelled RunComplete. A canceled prompt without a RunID is dropped

162+// silently as before.

163+func TestDrainQueueForStep_ReportsCanceledRunIDDrops(t *testing.T) {

164+ t.Parallel()

165+

166+ env := testEnv(t)

167+ a := NewSessionAgent(SessionAgentOptions{

168+ Sessions: env.sessions,

169+ Messages: env.messages,

170+ }).(*sessionAgent)

171+

172+ const sessionID = "drain-cancel-runid"

173+ a.messageQueue.Set(sessionID, []SessionAgentCall{

174+ {SessionID: sessionID, RunID: "run-canceled", Prompt: "canceled", acceptSeq: 1},

175+ {SessionID: sessionID, Prompt: "canceled-no-runid", acceptSeq: 1},

176+ {SessionID: sessionID, RunID: "run-survives", Prompt: "survives", acceptSeq: 5},

177+ })

178+ a.cancelMark.Set(sessionID, 2)

179+

180+ fold, canceledWithRunID := a.drainQueueForStep(sessionID)

181+

182+ require.Empty(t, fold, "no uncanceled non-RunID prompts to fold")

183+ require.Len(t, canceledWithRunID, 1,

184+ "only the dropped RunID-bearing prompt needs a terminal RunComplete")

185+ require.Equal(t, "run-canceled", canceledWithRunID[0].RunID)

186+

187+ kept, ok := a.messageQueue.Get(sessionID)

188+ require.True(t, ok)

189+ require.Len(t, kept, 1, "the uncanceled RunID prompt stays queued")

190+ require.Equal(t, "run-survives", kept[0].RunID)

115191}

116192

117193// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the

@@ -143,3 +219,104 @@ func TestRunCompletePublisher_MustDeliverOverTakesPublish(t *testing.T) {

143219 t.Fatal("PublishMustDeliver did not deliver event")

144220 }

145221}

222+

223+// requireSingleCancelledRunComplete reads exactly one RunComplete from ch,

224+// asserts it is the cancelled terminal event for runID, and verifies no

225+// second event arrives. This observes the published pubsub event rather

226+// than internal bookkeeping, which is the contract a `crush run` caller

227+// blocking on the broker actually relies on.

228+func requireSingleCancelledRunComplete(t *testing.T, ch <-chan pubsub.Event[notify.RunComplete], sessionID, runID string) {

229+ t.Helper()

230+ select {

231+ case ev := <-ch:

232+ require.Equal(t, runID, ev.Payload.RunID,

233+ "the published RunComplete must carry the dropped queued prompt's RunID")

234+ require.Equal(t, sessionID, ev.Payload.SessionID)

235+ require.True(t, ev.Payload.Cancelled,

236+ "a dropped queued prompt must publish a cancelled RunComplete")

237+ case <-time.After(5 * time.Second):

238+ t.Fatal("timed out waiting for the cancelled RunComplete")

239+ }

240+ select {

241+ case extra := <-ch:

242+ t.Fatalf("expected exactly one RunComplete, got a second: %+v", extra.Payload)

243+ case <-time.After(100 * time.Millisecond):

244+ }

245+}

246+

247+// TestCancel_QueuedRunIDPromptPublishesCancelledRunComplete proves the

248+// terminal-event behavior end-to-end: a RunID-bearing prompt sitting in

249+// the queue that is canceled while queued (via the public Cancel path,

250+// which routes through clearQueueAndNotify -> publishCanceledQueueDrops)

251+// must emit exactly one cancelled RunComplete on the broker for its

252+// RunID. A queued prompt without a RunID is dropped silently. This is the

253+// coverage the earlier drain test lacked: it asserted the returned

254+// bookkeeping slice, not the published event a `crush run` caller awaits.

255+func TestCancel_QueuedRunIDPromptPublishesCancelledRunComplete(t *testing.T) {

256+ t.Parallel()

257+

258+ env := testEnv(t)

259+ broker := pubsub.NewBroker[notify.RunComplete]()

260+ t.Cleanup(broker.Shutdown)

261+

262+ a := NewSessionAgent(SessionAgentOptions{

263+ Sessions: env.sessions,

264+ Messages: env.messages,

265+ RunComplete: broker,

266+ }).(*sessionAgent)

267+

268+ subCtx, subCancel := context.WithCancel(t.Context())

269+ defer subCancel()

270+ ch := broker.Subscribe(subCtx)

271+

272+ const sessionID = "cancel-queued-runid"

273+ a.messageQueue.Set(sessionID, []SessionAgentCall{

274+ {SessionID: sessionID, Prompt: "no-runid", acceptSeq: 1},

275+ {SessionID: sessionID, RunID: "run-queued", Prompt: "queued", acceptSeq: 2},

276+ })

277+

278+ a.Cancel(sessionID)

279+

280+ requireSingleCancelledRunComplete(t, ch, sessionID, "run-queued")

281+

282+ _, ok := a.messageQueue.Get(sessionID)

283+ require.False(t, ok, "Cancel must clear the queue")

284+}

285+

286+// TestDrainQueueForStep_DroppedRunIDPublishesCancelledRunComplete drives

287+// the production drain sequence (drainQueueForStep then

288+// publishCanceledQueueDrops, mirroring the PrepareStep handoff) and

289+// asserts the dropped RunID-bearing prompt actually publishes exactly one

290+// cancelled RunComplete on the broker. The companion bookkeeping test

291+// covers the returned slice; this one covers the observable terminal

292+// event.

293+func TestDrainQueueForStep_DroppedRunIDPublishesCancelledRunComplete(t *testing.T) {

294+ t.Parallel()

295+

296+ env := testEnv(t)

297+ broker := pubsub.NewBroker[notify.RunComplete]()

298+ t.Cleanup(broker.Shutdown)

299+

300+ a := NewSessionAgent(SessionAgentOptions{

301+ Sessions: env.sessions,

302+ Messages: env.messages,

303+ RunComplete: broker,

304+ }).(*sessionAgent)

305+

306+ subCtx, subCancel := context.WithCancel(t.Context())

307+ defer subCancel()

308+ ch := broker.Subscribe(subCtx)

309+

310+ const sessionID = "drain-drop-runid"

311+ a.messageQueue.Set(sessionID, []SessionAgentCall{

312+ {SessionID: sessionID, RunID: "run-dropped", Prompt: "dropped", acceptSeq: 1},

313+ {SessionID: sessionID, Prompt: "dropped-no-runid", acceptSeq: 1},

314+ })

315+ a.cancelMark.Set(sessionID, 2)

316+

317+ _, canceledWithRunID := a.drainQueueForStep(sessionID)

318+ require.Len(t, canceledWithRunID, 1)

319+ a.publishCanceledQueueDrops(canceledWithRunID)

320+

321+ requireSingleCancelledRunComplete(t, ch, sessionID, "run-dropped")

322+}