34995e9 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(server): prevent cancels from affecting future prompts

34995e9Christian Rocha5 months ago

internal/agent/accepted_run_test.go+7−2

compare with main

unifiedsplit

@@ -28,8 +28,13 @@ func (a *sessionAgent) acceptedCount(sessionID string) int {

28}

30func (a *sessionAgent) hasPendingCancel(sessionID string) bool {

31+ mark, ok := a.cancelMark.Get(sessionID)

32+ return ok && mark > 0

33+}

34+

35+func (a *sessionAgent) pendingCancelMark(sessionID string) uint64 {

31− _, ok := a.pendingCancels.Get(sessionID)

36+ mark, _ := a.cancelMark.Get(sessionID)

32− return ok

37+ return mark

33}

38}

35func TestAcceptedRun_CloseIsIdempotent(t *testing.T) {

40func TestAcceptedRun_CloseIsIdempotent(t *testing.T) {

internal/agent/agent.go+173−54

@@ -112,6 +112,15 @@ type SessionAgentCall struct {

112 // (in-process / local callers like AppWorkspace), behavior is

113 // unchanged and no accept tracking applies.

114 Accepted *AcceptedRun

115+ // acceptSeq carries the accept sequence of the handle that produced

116+ // this call after it has been enqueued and its Accepted handle

117+ // stripped. The queue-drain paths compare it against a session's

118+ // cancel mark so a follow-up queued before a cancel is dropped while

119+ // one queued after the cancel survives. 0 means untracked (an

120+ // in-process enqueue with no accept reservation), which the drain

121+ // paths treat as covered by any present mark, preserving the

122+ // pre-sequence behavior.

123+ acceptSeq uint64

115}

124}

116

125

117type SessionAgent interface {

126type SessionAgent interface {

@@ -167,20 +176,31 @@ type sessionAgent struct {

167 // BeginAccepted increments it; only AcceptedRun.Close decrements

176 // BeginAccepted increments it; only AcceptedRun.Close decrements

168 // it.

177 // it.

169 acceptedRuns *csync.Map[string, int]

178 acceptedRuns *csync.Map[string, int]

170− // pendingCancels records sessions whose dispatched-but-not-yet-

179+ // cancelMark records, per session, a high-water accept sequence: an

171− // running call should observe a cancellation request. It is only

180+ // accepted handle is canceled by it iff the handle's sequence is at

181+ // or below the mark. Cancel raises the mark to the latest sequence

182+ // assigned at cancel time, so a single Cancel covers every prompt

183+ // accepted-but-not-yet-active then, while a prompt accepted later

184+ // (higher sequence) is never poisoned. Absent or 0 means no pending

172− // set by Cancel when acceptedRuns > 0, so an idle Escape never

185+ // cancel. It is only raised by Cancel when acceptedRuns > 0, so an

173− // poisons the next prompt.

186+ // idle Escape never records a mark.

174− pendingCancels *csync.Map[string, struct{}]

187+ cancelMark *csync.Map[string, uint64]

175 // dispatchMuCreate guards lazy creation of per-session entries in

188 // dispatchMuCreate guards lazy creation of per-session entries in

176 // dispatchMu so two goroutines can't race to lock different mutex

189 // dispatchMu so two goroutines can't race to lock different mutex

177 // instances for the same session.

190 // instances for the same session.

178 dispatchMuCreate sync.Mutex

191 dispatchMuCreate sync.Mutex

179− // acceptedMu serializes increments/decrements of acceptedRuns. It

192+ // acceptedMu serializes increments/decrements of acceptedRuns and

193+ // the assignment of accept sequence numbers from acceptSeqGen. It

180 // is separate from dispatchMu so AcceptedRun.Close (which may run

194 // is separate from dispatchMu so AcceptedRun.Close (which may run

181 // while Run holds dispatchMu for the same session) does not

195 // while Run holds dispatchMu for the same session) does not

182 // deadlock by re-entering the dispatch lock.

196 // deadlock by re-entering the dispatch lock.

183 acceptedMu sync.Mutex

197 acceptedMu sync.Mutex

198+ // acceptSeqGen is the monotonic source of accept sequence numbers.

199+ // Each BeginAccepted increments it under acceptedMu and stamps the

200+ // returned handle, so sequences strictly increase in accept order

201+ // across the agent. Cancel uses its current value as the per-session

202+ // high-water mark.

203+ acceptSeqGen uint64

184}

204}

185

205

186type SessionAgentOptions struct {

206type SessionAgentOptions struct {

@@ -218,7 +238,7 @@ func NewSessionAgent(

218 activeRequests: csync.NewMap[string, context.CancelFunc](),

238 activeRequests: csync.NewMap[string, context.CancelFunc](),

219 dispatchMu: csync.NewMap[string, *sync.Mutex](),

239 dispatchMu: csync.NewMap[string, *sync.Mutex](),

220 acceptedRuns: csync.NewMap[string, int](),

240 acceptedRuns: csync.NewMap[string, int](),

221− pendingCancels: csync.NewMap[string, struct{}](),

241+ cancelMark: csync.NewMap[string, uint64](),

222 }

242 }

223}

243}

224

244

@@ -231,7 +251,12 @@ func NewSessionAgent(

231type AcceptedRun struct {

251type AcceptedRun struct {

232 agent *sessionAgent

252 agent *sessionAgent

233 sessionID string

253 sessionID string

254+ // seq is the monotonic accept sequence stamped by BeginAccepted. A

255+ // cancel covers this handle iff seq is at or below the session's

256+ // cancel mark, so a handle accepted after a cancel (higher seq) is

257+ // never poisoned by it.

258+ seq uint64

234− done atomic.Bool

259+ done atomic.Bool

235}

260}

236

261

237// Close decrements the accept counter for this reservation. It is safe

262// Close decrements the accept counter for this reservation. It is safe

@@ -263,19 +288,30 @@ func (a *sessionAgent) BeginAccepted(sessionID string) *AcceptedRun {

263 defer a.acceptedMu.Unlock()

288 defer a.acceptedMu.Unlock()

264 count, _ := a.acceptedRuns.Get(sessionID)

289 count, _ := a.acceptedRuns.Get(sessionID)

265 a.acceptedRuns.Set(sessionID, count+1)

290 a.acceptedRuns.Set(sessionID, count+1)

291+ a.acceptSeqGen++

266− return &AcceptedRun{agent: a, sessionID: sessionID}

292+ return &AcceptedRun{agent: a, sessionID: sessionID, seq: a.acceptSeqGen}

267}

293}

268

294

269// endAccepted decrements the accept counter for sessionID. It is only

295// endAccepted decrements the accept counter for sessionID. It is only

270// called via AcceptedRun.Close. It uses a dedicated lock (not the

296// called via AcceptedRun.Close. It uses a dedicated lock (not the

271// per-session dispatch mutex) so it can run while Run holds dispatchMu

297// per-session dispatch mutex) so it can run while Run holds dispatchMu

272// for the same session without deadlocking.

298// for the same session without deadlocking.

299+//

300+// When the count reaches zero the session's cancel mark is dropped: no

301+// accepted handle remains for it to cover, and any handle accepted later

302+// gets a strictly higher sequence that the mark would not match anyway.

303+// Handles canceled on entry never reach RunComplete, so this is the only

304+// place that clears the mark for an all-canceled batch. Sibling handles

305+// covered by the same mark are serialized on the per-session dispatch

306+// mutex and read the mark before they Close, so this never clears it out

307+// from under a covered handle still waiting to enter Run.

273func (a *sessionAgent) endAccepted(sessionID string) {

308func (a *sessionAgent) endAccepted(sessionID string) {

274 a.acceptedMu.Lock()

309 a.acceptedMu.Lock()

275 defer a.acceptedMu.Unlock()

310 defer a.acceptedMu.Unlock()

276 count, ok := a.acceptedRuns.Get(sessionID)

311 count, ok := a.acceptedRuns.Get(sessionID)

277 if !ok || count <= 1 {

312 if !ok || count <= 1 {

278 a.acceptedRuns.Del(sessionID)

313 a.acceptedRuns.Del(sessionID)

314+ a.cancelMark.Del(sessionID)

279 return

315 return

280 }

316 }

281 a.acceptedRuns.Set(sessionID, count-1)

317 a.acceptedRuns.Set(sessionID, count-1)

@@ -311,20 +347,44 @@ func (a *sessionAgent) enqueueCall(call SessionAgentCall) {

311 existing = []SessionAgentCall{}

347 existing = []SessionAgentCall{}

312 }

348 }

313 queued := call

349 queued := call

350+ if call.Accepted != nil {

351+ // Preserve the accept sequence after the handle is stripped so

352+ // the queue-drain paths can tell a follow-up queued before a

353+ // cancel (covered by the mark) from one queued after it.

354+ queued.acceptSeq = call.Accepted.seq

355+ }

314 queued.OnComplete = nil

356 queued.OnComplete = nil

315 queued.Accepted = nil

357 queued.Accepted = nil

316 existing = append(existing, queued)

358 existing = append(existing, queued)

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

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

318}

360}

319

361

320−// clearPendingCancel removes any pending-cancel record for sessionID. It

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

321−// takes the per-session dispatch lock so it is ordered against Cancel and

363+// takes the per-session dispatch lock so it is ordered against Cancel

322−// the dispatch handoff.

364+// and the dispatch handoff.

323func (a *sessionAgent) clearPendingCancel(sessionID string) {

365func (a *sessionAgent) clearPendingCancel(sessionID string) {

324 mu := a.sessionMu(sessionID)

366 mu := a.sessionMu(sessionID)

325 mu.Lock()

367 mu.Lock()

326 defer mu.Unlock()

368 defer mu.Unlock()

327− a.pendingCancels.Del(sessionID)

369+ a.cancelMark.Del(sessionID)

370+}

371+

372+// canceledBySeq reports whether an accepted handle or queued call with

373+// the given accept sequence is covered by a pending cancel for the

374+// session. Callers must hold the session's dispatch mutex. A tracked

375+// sequence (seq > 0) is covered only when it is at or below the cancel

376+// high-water mark, so a prompt accepted after the cancel (higher seq) is

377+// never poisoned. An untracked sequence (seq == 0, an in-process enqueue

378+// with no accept reservation) is covered whenever any mark is present,

379+// preserving the pre-sequence behavior. The mark is not consumed: it

380+// stays so every sibling handle it covers observes the same cancel, and

381+// a later handle (higher seq) ignores it regardless.

382+func (a *sessionAgent) canceledBySeq(sessionID string, seq uint64) bool {

383+ mark, ok := a.cancelMark.Get(sessionID)

384+ if !ok || mark == 0 {

385+ return false

386+ }

387+ return seq == 0 || seq <= mark

328}

388}

329

389

330// persistCanceledTurn writes the user/assistant records for a turn that

390// persistCanceledTurn writes the user/assistant records for a turn that

@@ -356,6 +416,26 @@ func (a *sessionAgent) persistCanceledTurn(ctx context.Context, call SessionAgen

356 return a.messages.Update(writeCtx, assistant)

416 return a.messages.Update(writeCtx, assistant)

357}

417}

358

418

419+// publishRunComplete emits the authoritative terminal event for a turn.

420+// It honors the per-call OnComplete hook when set (so the coordinator can

421+// coalesce retries) and otherwise falls back to the RunComplete broker.

422+// ctx is used only for the bounded-blocking must-deliver publish; the

423+// terminal payload is supplied by the caller. This is the single emit path

424+// shared by the streaming defer and the cancel-on-entry early return so a

425+// caller waiting on RunComplete (e.g. `crush run` with a RunID) always

426+// observes exactly one terminal event regardless of which Run branch ends

427+// the turn.

428+func (a *sessionAgent) publishRunComplete(ctx context.Context, call SessionAgentCall, complete notify.RunComplete) {

429+ if call.OnComplete != nil {

430+ call.OnComplete(complete)

431+ return

432+ }

433+ if a.runComplete == nil {

434+ return

435+ }

436+ a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete)

437+}

438+

359// ValidateCall performs the cheap structural validation that

439// ValidateCall performs the cheap structural validation that

360// sessionAgent.Run requires before a call can be dispatched: a call must

440// sessionAgent.Run requires before a call can be dispatched: a call must

361// carry either a non-empty prompt or a text attachment, and it must name a

441// carry either a non-empty prompt or a text attachment, and it must name a

@@ -394,22 +474,39 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

394 // Serialize the accepted -> (cancel-on-entry | queued |

474 // Serialize the accepted -> (cancel-on-entry | queued |

395 // active) transition against a concurrent Cancel. Cancel takes

475 // active) transition against a concurrent Cancel. Cancel takes

396 // the same per-session lock, so every cancel observes at least

476 // the same per-session lock, so every cancel observes at least

397− // one of: pendingCancels, an activeRequests entry, or a

477+ // one of: a cancel mark, an activeRequests entry, or a

398 // messageQueue entry it then clears.

478 // messageQueue entry it then clears.

399 mu := a.sessionMu(call.SessionID)

479 mu := a.sessionMu(call.SessionID)

400 mu.Lock()

480 mu.Lock()

401

481

402− if _, pending := a.pendingCancels.Get(call.SessionID); pending {

482+ if a.canceledBySeq(call.SessionID, call.Accepted.seq) {

403 // Cancel-on-entry: a cancel arrived while this run was

483 // Cancel-on-entry: a cancel arrived while this run was

404− // dispatched but not yet active. Consume the pending

484+ // dispatched but not yet active, and this handle's accept

485+ // sequence is at or below the session's cancel mark. The

486+ // mark is left in place so sibling handles it also covers

405− // cancel, release the accept reservation, drop the lock,

487+ // observe the same cancel; release the accept reservation,

406− // and persist a canceled turn without entering Stream.

488+ // drop the lock, and persist a canceled turn without

407− a.pendingCancels.Del(call.SessionID)

489+ // entering Stream.

490+ //

491+ // This path returns before the streaming defer that

492+ // publishes RunComplete is installed, so emit the terminal

493+ // event explicitly. Without it, a caller waiting on

494+ // RunComplete for this RunID (e.g. `crush run`, which

495+ // ignores message events and blocks on RunComplete) would

496+ // hang on an immediately-canceled accepted run.

408 call.Accepted.Close()

497 call.Accepted.Close()

409 mu.Unlock()

498 mu.Unlock()

499+ complete := notify.RunComplete{

500+ SessionID: call.SessionID,

501+ RunID: call.RunID,

502+ Cancelled: true,

503+ }

410 if err := a.persistCanceledTurn(ctx, call, false); err != nil {

504 if err := a.persistCanceledTurn(ctx, call, false); err != nil {

505+ complete.Error = err.Error()

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

411 return nil, err

507 return nil, err

412 }

508 }

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

413 return nil, nil

510 return nil, nil

414 }

511 }

415

512

@@ -579,14 +676,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

579 // the authoritative terminal event so a momentarily-full

676 // the authoritative terminal event so a momentarily-full

580 // subscriber channel can't silently drop it and hang

677 // subscriber channel can't silently drop it and hang

581 // non-interactive clients waiting on RunComplete.

678 // non-interactive clients waiting on RunComplete.

582− if call.OnComplete != nil {

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

583− call.OnComplete(complete)

584− return

585− }

586− if a.runComplete == nil {

587− return

588− }

589− a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete)

590 }()

680 }()

591

681

592 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)

682 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)

@@ -621,24 +711,26 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

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

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

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

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

623

713

624− // Drain queued follow-up prompts, but skip them if a cancel

714+ // Drain queued follow-up prompts, but skip any covered by a

625− // was recorded for the session while they sat in the queue:

715+ // cancel recorded while they sat in the queue: a cancel that

626− // a cancel that arrived after the queue insertion must not

716+ // arrived after a prompt was queued must not let it run as

627− // let the queued prompt run as part of this step.

717+ // part of this step. Coverage is per-call by accept sequence

718+ // so a follow-up queued after the cancel (higher seq) is

719+ // still folded in.

628 dispatchLock := a.sessionMu(call.SessionID)

720 dispatchLock := a.sessionMu(call.SessionID)

629 dispatchLock.Lock()

721 dispatchLock.Lock()

630− _, canceled := a.pendingCancels.Get(call.SessionID)

631 queuedCalls, _ := a.messageQueue.Get(call.SessionID)

722 queuedCalls, _ := a.messageQueue.Get(call.SessionID)

632 a.messageQueue.Del(call.SessionID)

723 a.messageQueue.Del(call.SessionID)

633 dispatchLock.Unlock()

724 dispatchLock.Unlock()

634− if !canceled {

635− for _, queued := range queuedCalls {

725+ for _, queued := range queuedCalls {

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

726+ if a.canceledBySeq(call.SessionID, queued.acceptSeq) {

637− if createErr != nil {

727+ continue

638− return callContext, prepared, createErr

639− }

640− prepared.Messages = append(prepared.Messages, userMessage.ToAIMessage()...)

641 }

728 }

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

730+ if createErr != nil {

731+ return callContext, prepared, createErr

732+ }

733+ prepared.Messages = append(prepared.Messages, userMessage.ToAIMessage()...)

642 }

734 }

643

735

644 prepared.Messages = a.workaroundProviderMediaLimitations(prepared.Messages, largeModel)

736 prepared.Messages = a.workaroundProviderMediaLimitations(prepared.Messages, largeModel)

@@ -1008,20 +1100,37 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

1008 // closing the dequeue -> re-register window.

1100 // closing the dequeue -> re-register window.

1009 mu := a.sessionMu(call.SessionID)

1101 mu := a.sessionMu(call.SessionID)

1010 mu.Lock()

1102 mu.Lock()

1011− if _, pending := a.pendingCancels.Get(call.SessionID); pending {

1103+ queuedMessages, _ := a.messageQueue.Get(call.SessionID)

1104+ if mark, ok := a.cancelMark.Get(call.SessionID); ok && mark > 0 && len(queuedMessages) > 0 {

1012 // A cancel was recorded for this session (e.g. it arrived while

1105 // A cancel was recorded for this session (e.g. it arrived while

1013− // this run was active and a follow-up had been accepted). Drop

1106+ // this run was active and follow-ups had been queued). Drop the

1014− // the queue instead of running it and consume the marker.

1107+ // queued prompts it covers (accept sequence at or below the

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

1109+ // sequence) so they still run.

1110+ var kept []SessionAgentCall

1111+ for _, q := range queuedMessages {

1112+ if q.acceptSeq == 0 || q.acceptSeq <= mark {

1113+ continue

1114+ }

1115+ kept = append(kept, q)

1116+ }

1117+ queuedMessages = kept

1015− a.pendingCancels.Del(call.SessionID)

1118+ a.messageQueue.Set(call.SessionID, kept)

1119+ }

1120+ if len(queuedMessages) == 0 {

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

1122+ // run remains in flight that it might still cover; otherwise a

1123+ // sibling prompt (sequence at or below the mark) waiting to

1124+ // enter Run would lose its cancellation. When accepted runs are

1125+ // gone, this also clears a stale mark so it can't catch a

1126+ // future run.

1016 a.messageQueue.Del(call.SessionID)

1127 a.messageQueue.Del(call.SessionID)

1017− mu.Unlock()

1128+ a.acceptedMu.Lock()

1018− return result, err

1019− }

1020− queuedMessages, ok := a.messageQueue.Get(call.SessionID)

1129+ inFlight, _ := a.acceptedRuns.Get(call.SessionID)

1021− if !ok || len(queuedMessages) == 0 {

1130+ a.acceptedMu.Unlock()

1022− // No queued work. Clear any stale pending-cancel entry as a

1131+ if inFlight == 0 {

1023− // safety net so it can't catch a future run (no-op when unset).

1024− a.pendingCancels.Del(call.SessionID)

1132+ a.cancelMark.Del(call.SessionID)

1133+ }

1025 mu.Unlock()

1134 mu.Unlock()

1026 return result, err

1135 return result, err

1027 }

1136 }

@@ -1619,17 +1728,27 @@ func (a *sessionAgent) Cancel(sessionID string) {

1619 }

1728 }

1620

1729

1621 // Record a pending cancel only when a dispatched-but-not-yet-active

1730 // Record a pending cancel only when a dispatched-but-not-yet-active

1622− // run exists. This catches a run still in the goroutine scheduler or

1731+ // run exists. This catches runs still in the goroutine scheduler or

1623 // about to enter Run's busy-queue branch, while leaving an idle

1732 // about to enter Run's busy-queue branch, while leaving an idle

1624 // session untouched. Active and accepted are not mutually exclusive:

1733 // session untouched. Active and accepted are not mutually exclusive:

1625 // when a run is active and a follow-up has been accepted, both the

1734 // when a run is active and a follow-up has been accepted, both the

1626 // cancel above and this pending record fire.

1735 // cancel above and this pending record fire.

1736+ //

1737+ // Raise the session's cancel mark to the latest accept sequence

1738+ // assigned so far. Every prompt currently accepted-but-not-yet-

1739+ // active has a sequence at or below that value, so one cancel covers

1740+ // all of them; a prompt accepted after this cancel gets a strictly

1741+ // higher sequence and is never poisoned. Using max keeps repeated

1742+ // cancels idempotent while the same prompts are in flight and lets a

1743+ // later cancel extend coverage to prompts accepted since.

1627 a.acceptedMu.Lock()

1744 a.acceptedMu.Lock()

1628 count, ok := a.acceptedRuns.Get(sessionID)

1745 count, ok := a.acceptedRuns.Get(sessionID)

1746+ mark := a.acceptSeqGen

1629 a.acceptedMu.Unlock()

1747 a.acceptedMu.Unlock()

1630 if ok && count > 0 {

1748 if ok && count > 0 {

1631− slog.Debug("Recording pending cancel for accepted run", "session_id", sessionID)

1749+ slog.Debug("Recording cancel mark for accepted runs", "session_id", sessionID, "count", count, "mark", mark)

1632− a.pendingCancels.Set(sessionID, struct{}{})

1750+ existing, _ := a.cancelMark.Get(sessionID)

1751+ a.cancelMark.Set(sessionID, max(existing, mark))

1633 }

1752 }

1634

1753

1635 if a.QueuedPrompts(sessionID) > 0 {

1754 if a.QueuedPrompts(sessionID) > 0 {

internal/agent/dispatch_cancel_test.go+252−3

@@ -5,9 +5,12 @@ import (

5 "errors"

6 "sync/atomic"

7 "testing"

8+ "time"

9 "charm.land/fantasy"

10 "charm.land/fantasy"

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

10 "github.com/charmbracelet/crush/internal/message"

12 "github.com/charmbracelet/crush/internal/message"

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

11 "github.com/stretchr/testify/assert"

14 "github.com/stretchr/testify/assert"

12 "github.com/stretchr/testify/require"

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

13)

16)

@@ -140,7 +143,7 @@ func TestRun_PrepareStepDrainSkipsQueuedOnPendingCancel(t *testing.T) {

140 // A follow-up prompt sits queued for the session.

143 // A follow-up prompt sits queued for the session.

141 sa.enqueueCall(SessionAgentCall{SessionID: sess.ID, Prompt: "queued-followup"})

144 sa.enqueueCall(SessionAgentCall{SessionID: sess.ID, Prompt: "queued-followup"})

142 // A cancel was recorded for the session while it sat in the queue.

145 // A cancel was recorded for the session while it sat in the queue.

143− sa.pendingCancels.Set(sess.ID, struct{}{})

146+ sa.cancelMark.Set(sess.ID, 1)

144

147

145 result, err := sa.Run(t.Context(), SessionAgentCall{

148 result, err := sa.Run(t.Context(), SessionAgentCall{

146 SessionID: sess.ID,

149 SessionID: sess.ID,

@@ -177,8 +180,8 @@ func TestRun_NormalCompletionClearsStalePendingCancel(t *testing.T) {

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

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

178 require.NoError(t, err)

181 require.NoError(t, err)

179

182

180− // A stale pending cancel lingers (no queued work, no accepted run).

183+ // A stale cancel mark lingers (no queued work, no accepted run).

181− sa.pendingCancels.Set(sess.ID, struct{}{})

184+ sa.cancelMark.Set(sess.ID, 1)

182

185

183 result, err := sa.Run(t.Context(), SessionAgentCall{

186 result, err := sa.Run(t.Context(), SessionAgentCall{

184 SessionID: sess.ID,

187 SessionID: sess.ID,

@@ -196,3 +199,249 @@ func TestRun_NormalCompletionClearsStalePendingCancel(t *testing.T) {

196 assert.Equal(t, message.Assistant, msgs[1].Role)

199 assert.Equal(t, message.Assistant, msgs[1].Role)

197 assert.Equal(t, message.FinishReasonEndTurn, msgs[1].FinishReason())

200 assert.Equal(t, message.FinishReasonEndTurn, msgs[1].FinishReason())

198}

201}

202+

203+// newCancelTestAgentWithRunComplete builds a DB-backed sessionAgent wired

204+// to a RunComplete broker so tests can observe the terminal event a

205+// RunID-bearing caller (e.g. `crush run`) blocks on.

206+func newCancelTestAgentWithRunComplete(t *testing.T) (*sessionAgent, fakeEnv, *pubsub.Broker[notify.RunComplete]) {

207+ t.Helper()

208+ env := testEnv(t)

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

210+ t.Cleanup(broker.Shutdown)

211+ sa := NewSessionAgent(SessionAgentOptions{

212+ Sessions: env.sessions,

213+ Messages: env.messages,

214+ RunComplete: broker,

215+ }).(*sessionAgent)

216+ return sa, env, broker

217+}

218+

219+// TestRun_CancelOnEntryPublishesRunComplete covers the first review

220+// finding: the cancel-on-entry path returned before the streaming defer

221+// that publishes RunComplete was installed. A caller that dispatches a

222+// run with a RunID and blocks on RunComplete (ignoring message events,

223+// like `crush run`) would hang on an immediately-canceled accepted run.

224+// The cancel-on-entry path must publish a terminal RunComplete carrying

225+// the originating RunID.

226+func TestRun_CancelOnEntryPublishesRunComplete(t *testing.T) {

227+ t.Parallel()

228+ sa, env, broker := newCancelTestAgentWithRunComplete(t)

229+

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

231+ require.NoError(t, err)

232+

233+ ctx, cancel := context.WithCancel(t.Context())

234+ defer cancel()

235+ ch := broker.Subscribe(ctx)

236+

237+ accept := sa.BeginAccepted(sess.ID)

238+ // A cancel arrives in the accepted-but-not-yet-active window.

239+ sa.Cancel(sess.ID)

240+ require.True(t, sa.hasPendingCancel(sess.ID))

241+

242+ result, err := sa.Run(t.Context(), SessionAgentCall{

243+ SessionID: sess.ID,

244+ RunID: "run-cancel-on-entry",

245+ Prompt: "hello",

246+ Accepted: accept,

247+ })

248+ require.NoError(t, err)

249+ require.Nil(t, result)

250+

251+ select {

252+ case got := <-ch:

253+ assert.Equal(t, "run-cancel-on-entry", got.Payload.RunID,

254+ "RunComplete must echo the originating RunID")

255+ assert.Equal(t, sess.ID, got.Payload.SessionID)

256+ assert.True(t, got.Payload.Cancelled,

257+ "cancel-on-entry RunComplete must be marked cancelled")

258+ case <-time.After(2 * time.Second):

259+ t.Fatal("cancel-on-entry must publish RunComplete; a RunID caller would hang otherwise")

260+ }

261+}

262+

263+// TestCancel_TwoAcceptedBothObserveCancellation covers the second review

264+// finding: a single cancel with two accepted-not-yet-active prompts must

265+// cancel both. The cancel raises the session's high-water mark to the

266+// latest accept sequence, so every prompt accepted-but-not-yet-active at

267+// cancel time is covered and both take the cancel-on-entry path.

268+func TestCancel_TwoAcceptedBothObserveCancellation(t *testing.T) {

269+ t.Parallel()

270+ sa, env := newCancelTestAgent(t)

271+

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

273+ require.NoError(t, err)

274+

275+ // Two prompts are accepted-but-not-yet-active for the same session.

276+ accept1 := sa.BeginAccepted(sess.ID)

277+ accept2 := sa.BeginAccepted(sess.ID)

278+ require.Equal(t, 2, sa.acceptedCount(sess.ID))

279+

280+ // A single cancel arrives before either becomes active.

281+ sa.Cancel(sess.ID)

282+ require.Equal(t, accept2.seq, sa.pendingCancelMark(sess.ID),

283+ "one cancel must mark every currently-accepted prompt as canceled")

284+ require.GreaterOrEqual(t, sa.pendingCancelMark(sess.ID), accept1.seq,

285+ "the mark must cover the earlier accepted prompt too")

286+

287+ // Both prompts enter Run; each must take cancel-on-entry, not run.

288+ r1, err := sa.Run(t.Context(), SessionAgentCall{

289+ SessionID: sess.ID,

290+ Prompt: "first",

291+ Accepted: accept1,

292+ })

293+ require.NoError(t, err)

294+ require.Nil(t, r1)

295+

296+ r2, err := sa.Run(t.Context(), SessionAgentCall{

297+ SessionID: sess.ID,

298+ Prompt: "second",

299+ Accepted: accept2,

300+ })

301+ require.NoError(t, err)

302+ require.Nil(t, r2)

303+

304+ require.False(t, sa.hasPendingCancel(sess.ID),

305+ "both reserved units must be consumed")

306+ require.Equal(t, 0, sa.acceptedCount(sess.ID),

307+ "both accept reservations must be released")

308+

309+ // Each canceled-on-entry turn writes a user + canceled assistant

310+ // message, and neither prompt was enqueued to run normally.

311+ require.Equal(t, 0, sa.QueuedPrompts(sess.ID),

312+ "neither accepted prompt may be enqueued to run normally")

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

314+ require.NoError(t, err)

315+ require.Len(t, msgs, 4, "two canceled turns produce two user + two assistant messages")

316+ var canceled int

317+ for _, m := range msgs {

318+ if m.Role == message.Assistant {

319+ assert.Equal(t, message.FinishReasonCanceled, m.FinishReason())

320+ canceled++

321+ }

322+ }

323+ require.Equal(t, 2, canceled, "both turns must finish canceled")

324+}

325+

326+// TestRun_IdleCancelDoesNotPoisonNextPrompt covers the idle-cancel

327+// no-op guarantee end-to-end: an Escape on an idle session must not

328+// record a pending cancel that leaks into the next accepted prompt, which

329+// must run normally to completion.

330+func TestRun_IdleCancelDoesNotPoisonNextPrompt(t *testing.T) {

331+ t.Parallel()

332+ sa, env := newStreamTestAgent(t)

333+

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

335+ require.NoError(t, err)

336+

337+ // Idle Escape: no accepted run, no active request.

338+ sa.Cancel(sess.ID)

339+ require.False(t, sa.hasPendingCancel(sess.ID),

340+ "idle cancel must not record a pending cancel")

341+

342+ // The next accepted prompt must run normally, not cancel on entry.

343+ accept := sa.BeginAccepted(sess.ID)

344+ result, err := sa.Run(t.Context(), SessionAgentCall{

345+ SessionID: sess.ID,

346+ Prompt: "next",

347+ Accepted: accept,

348+ })

349+ require.NoError(t, err)

350+ require.NotNil(t, result, "next prompt must run normally after an idle cancel")

351+

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

353+ require.NoError(t, err)

354+ require.Len(t, msgs, 2)

355+ assert.Equal(t, message.User, msgs[0].Role)

356+ assert.Equal(t, message.Assistant, msgs[1].Role)

357+ assert.Equal(t, message.FinishReasonEndTurn, msgs[1].FinishReason(),

358+ "the prompt must finish normally, not canceled")

359+}

360+

361+// TestCancel_AcceptedAfterCancelIsNotPoisoned is the regression test for

362+// the reviewer's finding: a counted session-level pending cancel let a

363+// prompt accepted after the cancel enter Run first and consume a unit

364+// reserved for the earlier prompts. With a sequence high-water mark, a

365+// single cancel covers exactly the prompts accepted-but-not-yet-active at

366+// cancel time (A and B); a prompt accepted afterwards (C) gets a higher

367+// sequence and must run normally without consuming A or B's cancellation.

368+// C is run first to prove it neither cancels nor drains the mark, then A

369+// and B are run and must both cancel on entry.

370+func TestCancel_AcceptedAfterCancelIsNotPoisoned(t *testing.T) {

371+ t.Parallel()

372+ sa, env := newStreamTestAgent(t)

373+

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

375+ require.NoError(t, err)

376+

377+ // A and B are accepted-but-not-yet-active.

378+ acceptA := sa.BeginAccepted(sess.ID)

379+ acceptB := sa.BeginAccepted(sess.ID)

380+

381+ // One cancel arrives covering both A and B.

382+ sa.Cancel(sess.ID)

383+ require.True(t, sa.hasPendingCancel(sess.ID))

384+ require.Equal(t, acceptB.seq, sa.pendingCancelMark(sess.ID),

385+ "the mark must cover every prompt accepted before the cancel")

386+

387+ // C is accepted AFTER the cancel; its sequence is above the mark.

388+ acceptC := sa.BeginAccepted(sess.ID)

389+ require.Greater(t, acceptC.seq, sa.pendingCancelMark(sess.ID),

390+ "a prompt accepted after the cancel must not be covered by the mark")

391+

392+ // Run C first. It must run normally to completion and must not

393+ // consume or clear the cancellation reserved for A and B.

394+ rc, err := sa.Run(t.Context(), SessionAgentCall{

395+ SessionID: sess.ID,

396+ Prompt: "C",

397+ Accepted: acceptC,

398+ })

399+ require.NoError(t, err)

400+ require.NotNil(t, rc, "C was accepted after the cancel and must run normally")

401+ require.True(t, sa.hasPendingCancel(sess.ID),

402+ "running C must not drain the cancellation reserved for A and B")

403+

404+ // Now A and B run. Both were accepted before the cancel and must

405+ // take the cancel-on-entry path.

406+ ra, err := sa.Run(t.Context(), SessionAgentCall{

407+ SessionID: sess.ID,

408+ Prompt: "A",

409+ Accepted: acceptA,

410+ })

411+ require.NoError(t, err)

412+ require.Nil(t, ra, "A must cancel on entry, not run")

413+

414+ rb, err := sa.Run(t.Context(), SessionAgentCall{

415+ SessionID: sess.ID,

416+ Prompt: "B",

417+ Accepted: acceptB,

418+ })

419+ require.NoError(t, err)

420+ require.Nil(t, rb, "B must cancel on entry, not run")

421+

422+ require.False(t, sa.hasPendingCancel(sess.ID),

423+ "the mark clears once all covered handles are resolved")

424+ require.Equal(t, 0, sa.acceptedCount(sess.ID))

425+ require.Equal(t, 0, sa.QueuedPrompts(sess.ID),

426+ "neither A nor B may be enqueued to run normally")

427+

428+ // C produced a normal turn; A and B each produced a canceled turn.

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

430+ require.NoError(t, err)

431+ require.Len(t, msgs, 6, "C normal + A canceled + B canceled = 3 user + 3 assistant")

432+

433+ var normal, canceled int

434+ for _, m := range msgs {

435+ if m.Role != message.Assistant {

436+ continue

437+ }

438+ switch m.FinishReason() {

439+ case message.FinishReasonEndTurn:

440+ normal++

441+ case message.FinishReasonCanceled:

442+ canceled++

443+ }

444+ }

445+ require.Equal(t, 1, normal, "only C finished normally")

446+ require.Equal(t, 2, canceled, "both A and B finished canceled")

447+}