cbec491 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(server): complete queued prompts independently

cbec491Christian Rocha5 months ago

internal/agent/agent.go+138−28

compare with main

unifiedsplit

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

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

360}

361

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

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

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:

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

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

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

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

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

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

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

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

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

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

379+// lock held.

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

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

371 dispatchLock := a.sessionMu(sessionID)

381 dispatchLock := a.sessionMu(sessionID)

372 dispatchLock.Lock()

382 dispatchLock.Lock()

373 defer dispatchLock.Unlock()

383 defer dispatchLock.Unlock()

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

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

375− a.messageQueue.Del(sessionID)

385+ var keep []SessionAgentCall

376− survivors := queuedCalls[:0]

377 for _, queued := range queuedCalls {

386 for _, queued := range queuedCalls {

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

387 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 == "" {

379 continue

431 continue

380 }

432 }

381− survivors = append(survivors, queued)

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

434+ SessionID: d.SessionID,

435+ RunID: d.RunID,

436+ Cancelled: true,

437+ })

382 }

438 }

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 {

383− return survivors

449+ return

450+ }

451+ a.publishCanceledQueueDrops(queued)

384}

452}

385

453

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

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

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

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

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

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

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

737

805

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

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

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

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

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

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

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

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

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

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

743− // still folded in.

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

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

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)

745− for _, queued := range survivors {

819+ for _, queued := range fold {

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

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

747 if createErr != nil {

821 if createErr != nil {

748 return callContext, prepared, createErr

822 return callContext, prepared, createErr

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

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

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

1126 // sequence) so they still run.

1200 // sequence) so they still run.

1127 var kept []SessionAgentCall

1201 var kept []SessionAgentCall

1202+ var canceledRunIDDrops []SessionAgentCall

1128 for _, q := range queuedMessages {

1203 for _, q := range queuedMessages {

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

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

1205+ if q.RunID != "" {

1206+ canceledRunIDDrops = append(canceledRunIDDrops, q)

1207+ }

1130 continue

1208 continue

1131 }

1209 }

1132 kept = append(kept, q)

1210 kept = append(kept, q)

1133 }

1211 }

1134 queuedMessages = kept

1212 queuedMessages = kept

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

1213 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)

1136 }

1218 }

1137 if len(queuedMessages) == 0 {

1219 if len(queuedMessages) == 0 {

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

1220 // 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 *

1151 mu.Unlock()

1233 mu.Unlock()

1152 return result, err

1234 return result, err

1153 }

1235 }

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

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

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

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

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

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

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

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

1158− // would belong to the outer turn.

1159 skipRunComplete = true

1240 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+ }

1160 firstQueuedMessage := queuedMessages[0]

1259 firstQueuedMessage := queuedMessages[0]

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

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

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

1261 // 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 *

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

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

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

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

1169 mu.Unlock()

1268 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+ }

1170 return a.Run(ctx, firstQueuedMessage)

1280 return a.Run(ctx, firstQueuedMessage)

1171}

1281}

1172

1282

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

1770

1880

1771 if a.QueuedPrompts(sessionID) > 0 {

1881 if a.QueuedPrompts(sessionID) > 0 {

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

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

1773− a.messageQueue.Del(sessionID)

1883+ a.clearQueueAndNotify(sessionID)

1774 }

1884 }

1775}

1885}

1776

1886

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

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

1778 if a.QueuedPrompts(sessionID) > 0 {

1888 if a.QueuedPrompts(sessionID) > 0 {

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

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

1780− a.messageQueue.Del(sessionID)

1890+ a.clearQueueAndNotify(sessionID)

1781 }

1891 }

1782}

1892}

1783

1893

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) {

58 "RunComplete still correlates with the originating SendMessage")

59}

61−// TestDrainUncanceledQueue_FiltersUnderDispatchLock verifies that the

61+// TestDrainQueueForStep_FiltersUnderDispatchLock verifies that the queue

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

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

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

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

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

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

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

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

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

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

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

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) {

68 t.Parallel()

70 t.Parallel()

70 env := testEnv(t)

72 env := testEnv(t)

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

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

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

84 a.cancelMark.Set(sessionID, 2)

86 a.cancelMark.Set(sessionID, 2)

86− survivors := a.drainUncanceledQueue(sessionID)

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

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

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

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

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

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

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")

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

96 _, 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")

94}

98}

96−// TestDrainUncanceledQueue_NoMarkKeepsAll verifies that with no cancel

100+// TestDrainQueueForStep_NoMarkFoldsAllNonRunID verifies that with no

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

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

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

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

99 t.Parallel()

103 t.Parallel()

100

104

101 env := testEnv(t)

105 env := testEnv(t)

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

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

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

111 })

115 })

112

116

113− survivors := a.drainUncanceledQueue(sessionID)

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

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

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)

115}

191}

116

192

117// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the

193// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the

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

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

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

144 }

220 }

145}

221}

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+}