@@ -21,6 +21,7 @@ import (
2121 "strconv"
2222 "strings"
2323 "sync"
24+ "sync/atomic"
2425 "time"
2526
2627 "charm.land/catwalk/pkg/catwalk"
@@ -103,10 +104,28 @@ type SessionAgentCall struct {
103104 // recursion drains, so falling back to the default broker
104105 // publish keeps the event visible to subscribers.
105106 OnComplete func(notify.RunComplete)
107+ // Accepted, when non-nil, is the accept reservation taken by
108+ // BeginAccepted before the call was dispatched onto a goroutine
109+ // (the client/server fire-and-forget path). Run consumes it under
110+ // dispatchMu[SessionID] once the accepted -> (cancel-on-entry |
111+ // queued | active) transition has been chosen. When nil
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
106124}
107125
108126type SessionAgent interface {
109127 Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)
128+ BeginAccepted(sessionID string) *AcceptedRun
110129 SetModels(large Model, small Model)
111130 SetTools(tools []fantasy.AgentTool)
112131 SetSystemPrompt(systemPrompt string)
@@ -145,6 +164,43 @@ type sessionAgent struct {
145164
146165 messageQueue *csync.Map[string, []SessionAgentCall]
147166 activeRequests *csync.Map[string, context.CancelFunc]
167+
168+ // dispatchMu holds a per-session mutex that serializes the
169+ // accepted -> (cancel-on-entry | queued | active) transition in
170+ // Run against a concurrent Cancel. The lock is held only during
171+ // the brief handoff (no DB or LLM I/O under the lock).
172+ dispatchMu *csync.Map[string, *sync.Mutex]
173+ // acceptedRuns counts dispatched-but-not-yet-active runs per
174+ // session. A counter > 0 means a dispatched prompt is in flight
175+ // and has not yet completed the dispatch handoff in Run. Only
176+ // BeginAccepted increments it; only AcceptedRun.Close decrements
177+ // it.
178+ acceptedRuns *csync.Map[string, int]
179+ // cancelMark records, per session, a high-water accept sequence: an
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
185+ // cancel. It is only raised by Cancel when acceptedRuns > 0, so an
186+ // idle Escape never records a mark.
187+ cancelMark *csync.Map[string, uint64]
188+ // dispatchMuCreate guards lazy creation of per-session entries in
189+ // dispatchMu so two goroutines can't race to lock different mutex
190+ // instances for the same session.
191+ dispatchMuCreate sync.Mutex
192+ // acceptedMu serializes increments/decrements of acceptedRuns and
193+ // the assignment of accept sequence numbers from acceptSeqGen. It
194+ // is separate from dispatchMu so AcceptedRun.Close (which may run
195+ // while Run holds dispatchMu for the same session) does not
196+ // deadlock by re-entering the dispatch lock.
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
148204}
149205
150206type SessionAgentOptions struct {
@@ -180,33 +236,403 @@ func NewSessionAgent(
180236 runComplete: opts.RunComplete,
181237 messageQueue: csync.NewMap[string, []SessionAgentCall](),
182238 activeRequests: csync.NewMap[string, context.CancelFunc](),
239+ dispatchMu: csync.NewMap[string, *sync.Mutex](),
240+ acceptedRuns: csync.NewMap[string, int](),
241+ cancelMark: csync.NewMap[string, uint64](),
183242 }
184243}
185244
186−func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *fantasy.AgentResult, retErr error) {
245+// AcceptedRun owns exactly one accept reservation taken by
246+// BeginAccepted. It is the only carrier of accept-state across the
247+// backend.runAgent / Coordinator.Run / sessionAgent.Run layers: a
248+// counter > 0 means a dispatched prompt is in flight and has not yet
249+// completed the dispatch handoff in Run. Close is the only way to
250+// release the reservation and is idempotent.
251+type AcceptedRun struct {
252+ agent *sessionAgent
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
259+ done atomic.Bool
260+}
261+
262+// Close decrements the accept counter for this reservation. It is safe
263+// to call multiple times; only the first call has effect.
264+func (r *AcceptedRun) Close() {
265+ if r == nil {
266+ return
267+ }
268+ if !r.done.CompareAndSwap(false, true) {
269+ return
270+ }
271+ r.agent.endAccepted(r.sessionID)
272+}
273+
274+// SessionID exposes the session this reservation is for so the run path
275+// can use it without an extra parameter.
276+func (r *AcceptedRun) SessionID() string {
277+ if r == nil {
278+ return ""
279+ }
280+ return r.sessionID
281+}
282+
283+// BeginAccepted increments the accept counter for sessionID and returns
284+// a handle whose Close is the only way to decrement it. It is the only
285+// entry point that mutates acceptedRuns.
286+func (a *sessionAgent) BeginAccepted(sessionID string) *AcceptedRun {
287+ a.acceptedMu.Lock()
288+ defer a.acceptedMu.Unlock()
289+ count, _ := a.acceptedRuns.Get(sessionID)
290+ a.acceptedRuns.Set(sessionID, count+1)
291+ a.acceptSeqGen++
292+ return &AcceptedRun{agent: a, sessionID: sessionID, seq: a.acceptSeqGen}
293+}
294+
295+// endAccepted decrements the accept counter for sessionID. It is only
296+// called via AcceptedRun.Close. It uses a dedicated lock (not the
297+// per-session dispatch mutex) so it can run while Run holds dispatchMu
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.
308+func (a *sessionAgent) endAccepted(sessionID string) {
309+ a.acceptedMu.Lock()
310+ defer a.acceptedMu.Unlock()
311+ count, ok := a.acceptedRuns.Get(sessionID)
312+ if !ok || count <= 1 {
313+ a.acceptedRuns.Del(sessionID)
314+ a.cancelMark.Del(sessionID)
315+ return
316+ }
317+ a.acceptedRuns.Set(sessionID, count-1)
318+}
319+
320+// sessionMu returns the per-session dispatch mutex, creating it on first
321+// use. Creation is guarded so concurrent callers always observe the same
322+// mutex instance for a given session.
323+func (a *sessionAgent) sessionMu(sessionID string) *sync.Mutex {
324+ if mu, ok := a.dispatchMu.Get(sessionID); ok {
325+ return mu
326+ }
327+ a.dispatchMuCreate.Lock()
328+ defer a.dispatchMuCreate.Unlock()
329+ if mu, ok := a.dispatchMu.Get(sessionID); ok {
330+ return mu
331+ }
332+ mu := &sync.Mutex{}
333+ a.dispatchMu.Set(sessionID, mu)
334+ return mu
335+}
336+
337+// enqueueCall appends call to the session's message queue. The
338+// OnComplete hook is stripped: the caller that supplied it (typically
339+// coordinator.Run) has its own retry/coalesce scope that ends when it
340+// returns, so by the time the queue drains nobody is left to consume the
341+// buffered terminal event. The recursive Run falls back to the default
342+// broker publish, which is what existing subscribers expect for queued
343+// turns.
344+func (a *sessionAgent) enqueueCall(call SessionAgentCall) {
345+ existing, ok := a.messageQueue.Get(call.SessionID)
346+ if !ok {
347+ existing = []SessionAgentCall{}
348+ }
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+ }
356+ queued.OnComplete = nil
357+ queued.Accepted = nil
358+ existing = append(existing, queued)
359+ a.messageQueue.Set(call.SessionID, existing)
360+}
361+
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) {
381+ dispatchLock := a.sessionMu(sessionID)
382+ dispatchLock.Lock()
383+ defer dispatchLock.Unlock()
384+ queuedCalls, _ := a.messageQueue.Get(sessionID)
385+ var keep []SessionAgentCall
386+ for _, queued := range queuedCalls {
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 == "" {
431+ continue
432+ }
433+ a.publishRunComplete(ctx, d, notify.RunComplete{
434+ SessionID: d.SessionID,
435+ RunID: d.RunID,
436+ Cancelled: true,
437+ })
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 {
449+ return
450+ }
451+ a.publishCanceledQueueDrops(queued)
452+}
453+
454+// clearPendingCancel removes any pending-cancel mark for sessionID. It
455+// takes the per-session dispatch lock so it is ordered against Cancel
456+// and the dispatch handoff.
457+func (a *sessionAgent) clearPendingCancel(sessionID string) {
458+ mu := a.sessionMu(sessionID)
459+ mu.Lock()
460+ defer mu.Unlock()
461+ a.cancelMark.Del(sessionID)
462+}
463+
464+// canceledBySeq reports whether an accepted handle or queued call with
465+// the given accept sequence is covered by a pending cancel for the
466+// session. Callers must hold the session's dispatch mutex. A tracked
467+// sequence (seq > 0) is covered only when it is at or below the cancel
468+// high-water mark, so a prompt accepted after the cancel (higher seq) is
469+// never poisoned. An untracked sequence (seq == 0, an in-process enqueue
470+// with no accept reservation) is covered whenever any mark is present,
471+// preserving the pre-sequence behavior. The mark is not consumed: it
472+// stays so every sibling handle it covers observes the same cancel, and
473+// a later handle (higher seq) ignores it regardless.
474+func (a *sessionAgent) canceledBySeq(sessionID string, seq uint64) bool {
475+ mark, ok := a.cancelMark.Get(sessionID)
476+ if !ok || mark == 0 {
477+ return false
478+ }
479+ return seq == 0 || seq <= mark
480+}
481+
482+// persistCanceledTurn writes the user/assistant records for a turn that
483+// was canceled before (or just as) streaming would have produced them.
484+// It creates the user message only when it was not already created by an
485+// earlier createUserMessage call (userMsgCreated), then writes an
486+// assistant message with FinishReasonCanceled. Both writes use
487+// context.WithoutCancel(ctx) so workspace shutdown (which cancels the run
488+// context) can't drop them.
489+func (a *sessionAgent) persistCanceledTurn(ctx context.Context, call SessionAgentCall, userMsgCreated bool) error {
490+ writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
491+ defer cancel()
492+ if !userMsgCreated {
493+ if _, err := a.createUserMessage(writeCtx, call); err != nil {
494+ return err
495+ }
496+ }
497+ largeModel := a.largeModel.Get()
498+ assistant, err := a.messages.Create(writeCtx, call.SessionID, message.CreateMessageParams{
499+ Role: message.Assistant,
500+ Parts: []message.ContentPart{},
501+ Model: largeModel.ModelCfg.Model,
502+ Provider: largeModel.ModelCfg.Provider,
503+ })
504+ if err != nil {
505+ return err
506+ }
507+ assistant.AddFinish(message.FinishReasonCanceled, "User canceled request", "")
508+ return a.messages.Update(writeCtx, assistant)
509+}
510+
511+// publishRunComplete emits the authoritative terminal event for a turn.
512+// It honors the per-call OnComplete hook when set (so the coordinator can
513+// coalesce retries) and otherwise falls back to the RunComplete broker.
514+// ctx is used only for the bounded-blocking must-deliver publish; the
515+// terminal payload is supplied by the caller. This is the single emit path
516+// shared by the streaming defer and the cancel-on-entry early return so a
517+// caller waiting on RunComplete (e.g. `crush run` with a RunID) always
518+// observes exactly one terminal event regardless of which Run branch ends
519+// the turn.
520+func (a *sessionAgent) publishRunComplete(ctx context.Context, call SessionAgentCall, complete notify.RunComplete) {
521+ if call.OnComplete != nil {
522+ call.OnComplete(complete)
523+ return
524+ }
525+ if a.runComplete == nil {
526+ return
527+ }
528+ a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete)
529+}
530+
531+// ValidateCall performs the cheap structural validation that
532+// sessionAgent.Run requires before a call can be dispatched: a call must
533+// carry either a non-empty prompt or a text attachment, and it must name a
534+// session. It is exported so callers that accept a run before dispatching it
535+// (e.g. backend.SendMessage) can apply the same checks and keep the error
536+// contract consistent.
537+func ValidateCall(call SessionAgentCall) error {
187538 if call.Prompt == "" && !message.ContainsTextAttachment(call.Attachments) {
188− return nil, ErrEmptyPrompt
539+ return ErrEmptyPrompt
189540 }
190541 if call.SessionID == "" {
191− return nil, ErrSessionMissing
192− }
193−
194− // Queue the message if busy. Strip OnComplete: the caller that
195− // supplied the hook (typically coordinator.Run) has its own
196− // retry/coalesce scope that ends when it returns, so by the time
197− // the queue drains nobody is left to consume the buffered
198− // terminal event. The recursive Run will fall back to the
199− // default broker publish, which is what existing subscribers
200− // expect for queued turns.
201− if a.IsSessionBusy(call.SessionID) {
202− existing, ok := a.messageQueue.Get(call.SessionID)
203− if !ok {
204− existing = []SessionAgentCall{}
542+ return ErrSessionMissing
543+ }
544+ return nil
545+}
546+
547+func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *fantasy.AgentResult, retErr error) {
548+ if err := ValidateCall(call); err != nil {
549+ return nil, err
550+ }
551+
552+ // genCtx/cancel are the run context and its cancel func. For the
553+ // accepted (fire-and-forget) dispatch path they are created under
554+ // dispatchMu below so a concurrent Cancel can observe the
555+ // activeRequests entry before the assistant message exists. For
556+ // the in-process path they stay nil here and are created later,
557+ // preserving the original ordering.
558+ var (
559+ genCtx context.Context
560+ cancel context.CancelFunc
561+ activeRegistered bool
562+ userMsgCreated bool
563+ )
564+
565+ if call.Accepted != nil {
566+ // Serialize the accepted -> (cancel-on-entry | queued |
567+ // active) transition against a concurrent Cancel. Cancel takes
568+ // the same per-session lock, so every cancel observes at least
569+ // one of: a cancel mark, an activeRequests entry, or a
570+ // messageQueue entry it then clears.
571+ mu := a.sessionMu(call.SessionID)
572+ mu.Lock()
573+
574+ if a.canceledBySeq(call.SessionID, call.Accepted.seq) {
575+ // Cancel-on-entry: a cancel arrived while this run was
576+ // dispatched but not yet active, and this handle's accept
577+ // sequence is at or below the session's cancel mark. The
578+ // mark is left in place so sibling handles it also covers
579+ // observe the same cancel; release the accept reservation,
580+ // drop the lock, and persist a canceled turn without
581+ // entering Stream.
582+ //
583+ // This path returns before the streaming defer that
584+ // publishes RunComplete is installed, so emit the terminal
585+ // event explicitly. Without it, a caller waiting on
586+ // RunComplete for this RunID (e.g. `crush run`, which
587+ // ignores message events and blocks on RunComplete) would
588+ // hang on an immediately-canceled accepted run.
589+ call.Accepted.Close()
590+ mu.Unlock()
591+ complete := notify.RunComplete{
592+ SessionID: call.SessionID,
593+ RunID: call.RunID,
594+ Cancelled: true,
595+ }
596+ if err := a.persistCanceledTurn(ctx, call, false); err != nil {
597+ complete.Error = err.Error()
598+ a.publishRunComplete(ctx, call, complete)
599+ return nil, err
600+ }
601+ a.publishRunComplete(ctx, call, complete)
602+ return nil, nil
205603 }
206− queued := call
207− queued.OnComplete = nil
208− existing = append(existing, queued)
209− a.messageQueue.Set(call.SessionID, existing)
604+
605+ if a.IsSessionBusy(call.SessionID) {
606+ // Busy: an earlier prompt is active. Queue this call and
607+ // release the accept reservation. A Cancel arriving after
608+ // this point sees the active entry and clears the queue.
609+ a.enqueueCall(call)
610+ call.Accepted.Close()
611+ mu.Unlock()
612+ return nil, nil
613+ }
614+
615+ // Idle: become the active run. Register the cancel func before
616+ // dropping the lock so a Cancel that arrives between here and
617+ // assistant creation is not lost.
618+ runCtx := context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
619+ genCtx, cancel = context.WithCancel(runCtx)
620+ a.activeRequests.Set(call.SessionID, cancel)
621+ activeRegistered = true
622+ call.Accepted.Close()
623+ mu.Unlock()
624+
625+ defer cancel()
626+ defer a.activeRequests.Del(call.SessionID)
627+ } else if a.IsSessionBusy(call.SessionID) {
628+ // Queue the message if busy. Strip OnComplete: the caller that
629+ // supplied the hook (typically coordinator.Run) has its own
630+ // retry/coalesce scope that ends when it returns, so by the time
631+ // the queue drains nobody is left to consume the buffered
632+ // terminal event. The recursive Run will fall back to the
633+ // default broker publish, which is what existing subscribers
634+ // expect for queued turns.
635+ a.enqueueCall(call)
210636 return nil, nil
211637 }
212638
@@ -269,15 +695,22 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
269695 if err != nil {
270696 return nil, err
271697 }
698+ userMsgCreated = true
272699
273700 // Add the session to the context.
274701 ctx = context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
275702
276− genCtx, cancel := context.WithCancel(ctx)
277− a.activeRequests.Set(call.SessionID, cancel)
703+ // For the accepted dispatch path the run context and cancel func
704+ // were already created and registered under dispatchMu above; reuse
705+ // them. For the in-process path create them here, preserving the
706+ // original ordering.
707+ if !activeRegistered {
708+ genCtx, cancel = context.WithCancel(ctx)
709+ a.activeRequests.Set(call.SessionID, cancel)
278710
279− defer cancel()
280− defer a.activeRequests.Del(call.SessionID)
711+ defer cancel()
712+ defer a.activeRequests.Del(call.SessionID)
713+ }
281714 // skipRunComplete is set just before the queued-recursion path so
282715 // the outer Run doesn't publish a RunComplete that would race
283716 // with — and be superseded by — the recursive call's own
@@ -305,7 +738,13 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
305738 // (the pubsub broker fan-in does not serialize publishes from
306739 // different upstream brokers).
307740 defer func() {
308− if flushErr := a.messages.FlushAll(ctx); flushErr != nil {
741+ // Use a context detached from the run context: workspace
742+ // shutdown cancels ctx before this goroutine returns, but the
743+ // buffered streaming deltas must still land before the DB is
744+ // closed. A short timeout bounds the flush.
745+ flushCtx, flushCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
746+ defer flushCancel()
747+ if flushErr := a.messages.FlushAll(flushCtx); flushErr != nil {
309748 slog.Error("Failed to flush pending message updates after run", "error", flushErr)
310749 }
311750 if skipRunComplete {
@@ -329,14 +768,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
329768 // the authoritative terminal event so a momentarily-full
330769 // subscriber channel can't silently drop it and hang
331770 // non-interactive clients waiting on RunComplete.
332− if call.OnComplete != nil {
333− call.OnComplete(complete)
334− return
335− }
336− if a.runComplete == nil {
337− return
338− }
339− a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete)
771+ a.publishRunComplete(ctx, call, complete)
340772 }()
341773
342774 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)
@@ -371,9 +803,20 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
371803 // Use latest tools (updated by SetTools when MCP tools change).
372804 prepared.Tools = a.tools.Copy()
373805
374− queuedCalls, _ := a.messageQueue.Get(call.SessionID)
375− a.messageQueue.Del(call.SessionID)
376− for _, queued := range queuedCalls {
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 {
377820 userMessage, createErr := a.createUserMessage(callContext, queued)
378821 if createErr != nil {
379822 return callContext, prepared, createErr
@@ -575,13 +1018,34 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
5751018 isHyper := largeModel.ModelCfg.Provider == hyper.Name
5761019 isCancelErr := errors.Is(err, context.Canceled)
5771020 if currentAssistant == nil {
1021+ // Cancel-before-assistant-creation window: the run was
1022+ // canceled after activeRequests.Set but before PrepareStep
1023+ // created the assistant message. Without this, the turn
1024+ // would return with no FinishReasonCanceled marker and no
1025+ // user-visible record. The user message was already created
1026+ // above, so persistCanceledTurn only writes the assistant
1027+ // record.
1028+ if isCancelErr {
1029+ if persistErr := a.persistCanceledTurn(ctx, call, userMsgCreated); persistErr != nil {
1030+ return nil, persistErr
1031+ }
1032+ }
5781033 return result, err
5791034 }
1035+ // Persist final state with a context detached from the run
1036+ // context. The run context (ctx) is derived from the
1037+ // workspace context, which workspace shutdown cancels before
1038+ // agent goroutines finish; using ctx here would drop the
1039+ // final assistant state. WithoutCancel keeps the values
1040+ // (e.g. session ID) while ignoring cancellation, and a short
1041+ // timeout bounds the cleanup writes.
1042+ cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
1043+ defer cleanupCancel()
5801044 // Ensure we finish thinking on error to close the reasoning state.
5811045 currentAssistant.FinishThinking()
5821046 toolCalls := currentAssistant.ToolCalls()
583− // INFO: we use the parent context here because the genCtx has been cancelled.
584− msgs, createErr := a.messages.List(ctx, currentAssistant.SessionID)
1047+ // INFO: we use the cleanup context here because the genCtx has been cancelled.
1048+ msgs, createErr := a.messages.List(cleanupCtx, currentAssistant.SessionID)
5851049 if createErr != nil {
5861050 return nil, createErr
5871051 }
@@ -590,7 +1054,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
5901054 tc.Finished = true
5911055 tc.Input = "{}"
5921056 currentAssistant.AddToolCall(tc)
593− updateErr := a.messages.Update(ctx, *currentAssistant)
1057+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)
5941058 if updateErr != nil {
5951059 return nil, updateErr
5961060 }
@@ -623,7 +1087,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
6231087 Content: content,
6241088 IsError: true,
6251089 }
626− _, createErr = a.messages.Create(ctx, currentAssistant.SessionID, message.CreateMessageParams{
1090+ _, createErr = a.messages.Create(cleanupCtx, currentAssistant.SessionID, message.CreateMessageParams{
6271091 Role: message.Tool,
6281092 Parts: []message.ContentPart{
6291093 toolResult,
@@ -662,9 +1126,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
6621126 } else {
6631127 currentAssistant.AddFinish(message.FinishReasonError, defaultTitle, err.Error())
6641128 }
665− // Note: we use the parent context here because the genCtx has been
1129+ // Note: we use the cleanup context here because the genCtx has been
6661130 // cancelled.
667− updateErr := a.messages.Update(ctx, *currentAssistant)
1131+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)
6681132 if updateErr != nil {
6691133 return nil, updateErr
6701134 }
@@ -705,18 +1169,106 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
7051169 })
7061170 }
7071171
708− queuedMessages, ok := a.messageQueue.Get(call.SessionID)
709− if !ok || len(queuedMessages) == 0 {
1172+ // Hand off to the next queued prompt (if any) under dispatchMu so
1173+ // the transition from this finished run to the queued run is atomic
1174+ // against a concurrent Cancel. activeRequests for this session was
1175+ // just deleted above, so without the lock there is a window in
1176+ // which the session looks idle and a cancel becomes a no-op that
1177+ // fails to stop the queued prompt. Holding the lock lets us observe
1178+ // a pending cancel recorded against the session and drop the queue
1179+ // instead of running it, and (for the recursion) hand a fresh
1180+ // accept reservation to the dequeued call so acceptedRuns stays > 0
1181+ // across the recursive Run's own dispatch handoff — keeping the
1182+ // session observable to Cancel for the entire transition and
1183+ // closing the dequeue -> re-register window.
1184+ mu := a.sessionMu(call.SessionID)
1185+ mu.Lock()
1186+ queuedMessages, _ := a.messageQueue.Get(call.SessionID)
1187+ if mark, ok := a.cancelMark.Get(call.SessionID); ok && mark > 0 && len(queuedMessages) > 0 {
1188+ // A cancel was recorded for this session (e.g. it arrived while
1189+ // this run was active and follow-ups had been queued). Drop the
1190+ // queued prompts it covers (accept sequence at or below the
1191+ // mark, or untracked); keep any queued after the cancel (higher
1192+ // sequence) so they still run.
1193+ var kept []SessionAgentCall
1194+ var canceledRunIDDrops []SessionAgentCall
1195+ for _, q := range queuedMessages {
1196+ if q.acceptSeq == 0 || q.acceptSeq <= mark {
1197+ if q.RunID != "" {
1198+ canceledRunIDDrops = append(canceledRunIDDrops, q)
1199+ }
1200+ continue
1201+ }
1202+ kept = append(kept, q)
1203+ }
1204+ queuedMessages = kept
1205+ a.messageQueue.Set(call.SessionID, kept)
1206+ // A dropped prompt carrying a RunID must still publish its
1207+ // terminal cancelled RunComplete so a caller waiting on that
1208+ // RunID does not hang.
1209+ a.publishCanceledQueueDrops(canceledRunIDDrops)
1210+ }
1211+ if len(queuedMessages) == 0 {
1212+ // No queued work. Clear the cancel mark only when no accepted
1213+ // run remains in flight that it might still cover; otherwise a
1214+ // sibling prompt (sequence at or below the mark) waiting to
1215+ // enter Run would lose its cancellation. When accepted runs are
1216+ // gone, this also clears a stale mark so it can't catch a
1217+ // future run.
1218+ a.messageQueue.Del(call.SessionID)
1219+ a.acceptedMu.Lock()
1220+ inFlight, _ := a.acceptedRuns.Get(call.SessionID)
1221+ a.acceptedMu.Unlock()
1222+ if inFlight == 0 {
1223+ a.cancelMark.Del(call.SessionID)
1224+ }
1225+ mu.Unlock()
7101226 return result, err
7111227 }
712− // There are queued messages restart the loop. The recursive Run
713− // publishes its own RunComplete for the queued prompt, so suppress
714− // the outer defer's emit to avoid a duplicate event whose Error
715− // field would belong to the recursive turn but whose MessageID/Text
716− // would belong to the outer turn.
1228+ // There are queued messages, restart the loop. Suppress the outer
1229+ // defer's emit: it would otherwise observe the recursive Run's retErr
1230+ // (named-return clobbering through the return below) against this
1231+ // turn's MessageID/Text and publish a mixed, racing event.
7171232 skipRunComplete = true
1233+ // Decide whether this turn still owes its own terminal RunComplete.
1234+ // Each submitted prompt with a RunID has its own lifecycle, so a turn
1235+ // that is finished and handing off to a *different* queued prompt must
1236+ // publish its own RunComplete here — leaving it to the recursive turn
1237+ // (which carries a different RunID) would hang a caller waiting on
1238+ // this turn's RunID. The exception is the summarize-continuation path,
1239+ // which re-queues this same call (same RunID) to resume after a
1240+ // summary; in that case the eventual terminal turn for this RunID
1241+ // publishes, so publishing now would double-emit.
1242+ outerOwesRunComplete := call.RunID != ""
1243+ if outerOwesRunComplete {
1244+ for _, q := range queuedMessages {
1245+ if q.RunID == call.RunID {
1246+ outerOwesRunComplete = false
1247+ break
1248+ }
1249+ }
1250+ }
7181251 firstQueuedMessage := queuedMessages[0]
7191252 a.messageQueue.Set(call.SessionID, queuedMessages[1:])
1253+ // Reserve a fresh accept for the dequeued prompt before dropping the
1254+ // lock so acceptedRuns > 0 across the handoff into the recursive
1255+ // Run. This closes the window between this dequeue and the recursive
1256+ // Run registering its activeRequests entry: a cancel arriving in
1257+ // that window now records a pending cancel (acceptedRuns > 0) that
1258+ // the recursive Run's accepted path observes as cancel-on-entry.
1259+ firstQueuedMessage.Accepted = a.BeginAccepted(call.SessionID)
1260+ mu.Unlock()
1261+ if outerOwesRunComplete {
1262+ complete := notify.RunComplete{SessionID: call.SessionID, RunID: call.RunID}
1263+ if currentAssistant != nil {
1264+ complete.MessageID = currentAssistant.ID
1265+ complete.Text = currentAssistant.Content().String()
1266+ }
1267+ if ctx.Err() != nil {
1268+ complete.Cancelled = true
1269+ }
1270+ a.publishRunComplete(ctx, call, complete)
1271+ }
7201272 return a.Run(ctx, firstQueuedMessage)
7211273}
7221274
@@ -1269,6 +1821,16 @@ func summaryCompletionTokens(usage fantasy.Usage, summaryMessage message.Message
12691821}
12701822
12711823func (a *sessionAgent) Cancel(sessionID string) {
1824+ // Serialize against the dispatch handoff in Run so the accepted ->
1825+ // (cancel-on-entry | queued | active) transition is atomic against
1826+ // this cancel. Every cancel observes at least one of: an active
1827+ // request, an accepted run (recorded as a pending cancel), or a
1828+ // queue entry it then clears. If none of those hold, an idle Escape
1829+ // is a true no-op and must not poison the next prompt.
1830+ mu := a.sessionMu(sessionID)
1831+ mu.Lock()
1832+ defer mu.Unlock()
1833+
12721834 // Cancel regular requests. Don't use Take() here - we need the entry to
12731835 // remain in activeRequests so IsBusy() returns true until the goroutine
12741836 // fully completes (including error handling that may access the DB).
@@ -1284,16 +1846,40 @@ func (a *sessionAgent) Cancel(sessionID string) {
12841846 cancel()
12851847 }
12861848
1849+ // Record a pending cancel only when a dispatched-but-not-yet-active
1850+ // run exists. This catches runs still in the goroutine scheduler or
1851+ // about to enter Run's busy-queue branch, while leaving an idle
1852+ // session untouched. Active and accepted are not mutually exclusive:
1853+ // when a run is active and a follow-up has been accepted, both the
1854+ // cancel above and this pending record fire.
1855+ //
1856+ // Raise the session's cancel mark to the latest accept sequence
1857+ // assigned so far. Every prompt currently accepted-but-not-yet-
1858+ // active has a sequence at or below that value, so one cancel covers
1859+ // all of them; a prompt accepted after this cancel gets a strictly
1860+ // higher sequence and is never poisoned. Using max keeps repeated
1861+ // cancels idempotent while the same prompts are in flight and lets a
1862+ // later cancel extend coverage to prompts accepted since.
1863+ a.acceptedMu.Lock()
1864+ count, ok := a.acceptedRuns.Get(sessionID)
1865+ mark := a.acceptSeqGen
1866+ a.acceptedMu.Unlock()
1867+ if ok && count > 0 {
1868+ slog.Debug("Recording cancel mark for accepted runs", "session_id", sessionID, "count", count, "mark", mark)
1869+ existing, _ := a.cancelMark.Get(sessionID)
1870+ a.cancelMark.Set(sessionID, max(existing, mark))
1871+ }
1872+
12871873 if a.QueuedPrompts(sessionID) > 0 {
12881874 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1289− a.messageQueue.Del(sessionID)
1875+ a.clearQueueAndNotify(sessionID)
12901876 }
12911877}
12921878
12931879func (a *sessionAgent) ClearQueue(sessionID string) {
12941880 if a.QueuedPrompts(sessionID) > 0 {
12951881 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1296− a.messageQueue.Del(sessionID)
1882+ a.clearQueueAndNotify(sessionID)
12971883 }
12981884}
12991885