@@ -21,6 +21,7 @@ import (
26 "charm.land/catwalk/pkg/catwalk"
27 "charm.land/catwalk/pkg/catwalk"
@@ -103,10 +104,28 @@ type SessionAgentCall struct {
103 // recursion drains, so falling back to the default broker
104 // recursion drains, so falling back to the default broker
104 // publish keeps the event visible to subscribers.
105 // publish keeps the event visible to subscribers.
105 OnComplete func(notify.RunComplete)
106 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.
108type SessionAgent interface {
126type SessionAgent interface {
109 Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)
127 Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)
128+ BeginAccepted(sessionID string) *AcceptedRun
110 SetModels(large Model, small Model)
129 SetModels(large Model, small Model)
111 SetTools(tools []fantasy.AgentTool)
130 SetTools(tools []fantasy.AgentTool)
112 SetSystemPrompt(systemPrompt string)
131 SetSystemPrompt(systemPrompt string)
@@ -145,6 +164,43 @@ type sessionAgent struct {
146 messageQueue *csync.Map[string, []SessionAgentCall]
165 messageQueue *csync.Map[string, []SessionAgentCall]
147 activeRequests *csync.Map[string, context.CancelFunc]
166 activeRequests *csync.Map[string, context.CancelFunc]
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
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
150type SessionAgentOptions struct {
206type SessionAgentOptions struct {
@@ -180,33 +236,403 @@ func NewSessionAgent(
180 runComplete: opts.RunComplete,
236 runComplete: opts.RunComplete,
181 messageQueue: csync.NewMap[string, []SessionAgentCall](),
237 messageQueue: csync.NewMap[string, []SessionAgentCall](),
182 activeRequests: csync.NewMap[string, context.CancelFunc](),
238 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](),
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 {
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.
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() {
268+ if !r.done.CompareAndSwap(false, true) {
271+ r.agent.endAccepted(r.sessionID)
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 {
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 {
288+ defer a.acceptedMu.Unlock()
289+ count, _ := a.acceptedRuns.Get(sessionID)
290+ a.acceptedRuns.Set(sessionID, count+1)
292+ return &AcceptedRun{agent: a, sessionID: sessionID, seq: a.acceptSeqGen}
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.
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) {
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)
317+ a.acceptedRuns.Set(sessionID, count-1)
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 {
327+ a.dispatchMuCreate.Lock()
328+ defer a.dispatchMuCreate.Unlock()
329+ if mu, ok := a.dispatchMu.Get(sessionID); ok {
333+ a.dispatchMu.Set(sessionID, mu)
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
344+func (a *sessionAgent) enqueueCall(call SessionAgentCall) {
345+ existing, ok := a.messageQueue.Get(call.SessionID)
347+ existing = []SessionAgentCall{}
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
356+ queued.OnComplete = nil
357+ queued.Accepted = nil
358+ existing = append(existing, queued)
359+ a.messageQueue.Set(call.SessionID, existing)
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.
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
380+func (a *sessionAgent) drainQueueForStep(sessionID string) (fold, canceledWithRunID []SessionAgentCall) {
381+ dispatchLock := a.sessionMu(sessionID)
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)
393+ if queued.RunID != "" {
394+ keep = append(keep, queued)
397+ fold = append(fold, queued)
400+ a.messageQueue.Del(sessionID)
402+ a.messageQueue.Set(sessionID, keep)
404+ return fold, canceledWithRunID
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) {
418+ for _, d := range drops {
427+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
429+ for _, d := range drops {
433+ a.publishRunComplete(ctx, d, notify.RunComplete{
434+ SessionID: d.SessionID,
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)
451+ a.publishCanceledQueueDrops(queued)
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)
461+ a.cancelMark.Del(sessionID)
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 {
479+ return seq == 0 || seq <= mark
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)
492+ if !userMsgCreated {
493+ if _, err := a.createUserMessage(writeCtx, call); err != nil {
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,
507+ assistant.AddFinish(message.FinishReasonCanceled, "User canceled request", "")
508+ return a.messages.Update(writeCtx, assistant)
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
520+func (a *sessionAgent) publishRunComplete(ctx context.Context, call SessionAgentCall, complete notify.RunComplete) {
521+ if call.OnComplete != nil {
522+ call.OnComplete(complete)
525+ if a.runComplete == nil {
528+ a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete)
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 {
187 if call.Prompt == "" && !message.ContainsTextAttachment(call.Attachments) {
538 if call.Prompt == "" && !message.ContainsTextAttachment(call.Attachments) {
188− return nil, ErrEmptyPrompt
539+ return ErrEmptyPrompt
190 if call.SessionID == "" {
541 if call.SessionID == "" {
191− return nil, ErrSessionMissing
542+ return ErrSessionMissing
194− // Queue the message if busy. Strip OnComplete: the caller that
545+}
195− // supplied the hook (typically coordinator.Run) has its own
546+
196− // retry/coalesce scope that ends when it returns, so by the time
547+func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *fantasy.AgentResult, retErr error) {
197− // the queue drains nobody is left to consume the buffered
548+ if err := ValidateCall(call); err != nil {
198− // terminal event. The recursive Run will fall back to the
549+ return nil, err
199− // default broker publish, which is what existing subscribers
550+ }
200− // expect for queued turns.
551+
201− if a.IsSessionBusy(call.SessionID) {
552+ // genCtx/cancel are the run context and its cancel func. For the
202− existing, ok := a.messageQueue.Get(call.SessionID)
553+ // accepted (fire-and-forget) dispatch path they are created under
203− if !ok {
554+ // dispatchMu below so a concurrent Cancel can observe the
204− existing = []SessionAgentCall{}
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.
559+ genCtx context.Context
560+ cancel context.CancelFunc
561+ activeRegistered bool
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)
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
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()
591+ complete := notify.RunComplete{
592+ SessionID: call.SessionID,
596+ if err := a.persistCanceledTurn(ctx, call, false); err != nil {
597+ complete.Error = err.Error()
598+ a.publishRunComplete(ctx, call, complete)
601+ a.publishRunComplete(ctx, call, complete)
207− queued.OnComplete = nil
605+ if a.IsSessionBusy(call.SessionID) {
208− existing = append(existing, queued)
606+ // Busy: an earlier prompt is active. Queue this call and
209− a.messageQueue.Set(call.SessionID, existing)
607+ // release the accept reservation. A Cancel arriving after
608+ // this point sees the active entry and clears the queue.
610+ call.Accepted.Close()
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()
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.
210 return nil, nil
636 return nil, nil
@@ -269,15 +695,22 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
269 if err != nil {
695 if err != nil {
270 return nil, err
696 return nil, err
698+ userMsgCreated = true
273 // Add the session to the context.
700 // Add the session to the context.
274 ctx = context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
701 ctx = context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
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 {
276− genCtx, cancel := context.WithCancel(ctx)
708+ genCtx, cancel = context.WithCancel(ctx)
277− a.activeRequests.Set(call.SessionID, cancel)
709+ a.activeRequests.Set(call.SessionID, cancel)
279− defer cancel()
711+ defer cancel()
280− defer a.activeRequests.Del(call.SessionID)
712+ defer a.activeRequests.Del(call.SessionID)
281 // skipRunComplete is set just before the queued-recursion path so
714 // skipRunComplete is set just before the queued-recursion path so
282 // the outer Run doesn't publish a RunComplete that would race
715 // the outer Run doesn't publish a RunComplete that would race
283 // with — and be superseded by — the recursive call's own
716 // with — and be superseded by — the recursive call's own
@@ -305,7 +738,13 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
305 // (the pubsub broker fan-in does not serialize publishes from
738 // (the pubsub broker fan-in does not serialize publishes from
306 // different upstream brokers).
739 // different upstream brokers).
307 defer func() {
740 defer func() {
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)
308− if flushErr := a.messages.FlushAll(ctx); flushErr != nil {
747+ if flushErr := a.messages.FlushAll(flushCtx); flushErr != nil {
309 slog.Error("Failed to flush pending message updates after run", "error", flushErr)
748 slog.Error("Failed to flush pending message updates after run", "error", flushErr)
311 if skipRunComplete {
750 if skipRunComplete {
@@ -329,14 +768,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
329 // the authoritative terminal event so a momentarily-full
768 // the authoritative terminal event so a momentarily-full
330 // subscriber channel can't silently drop it and hang
769 // subscriber channel can't silently drop it and hang
331 // non-interactive clients waiting on RunComplete.
770 // non-interactive clients waiting on RunComplete.
332− if call.OnComplete != nil {
771+ a.publishRunComplete(ctx, call, complete)
333− call.OnComplete(complete)
336− if a.runComplete == nil {
339− a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete)
342 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)
774 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)
@@ -371,9 +803,20 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
371 // Use latest tools (updated by SetTools when MCP tools change).
803 // Use latest tools (updated by SetTools when MCP tools change).
372 prepared.Tools = a.tools.Copy()
804 prepared.Tools = a.tools.Copy()
374− queuedCalls, _ := a.messageQueue.Get(call.SessionID)
806+ // Drain queued follow-up prompts for this step. Calls covered
375− a.messageQueue.Del(call.SessionID)
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)
376− for _, queued := range queuedCalls {
819+ for _, queued := range fold {
377 userMessage, createErr := a.createUserMessage(callContext, queued)
820 userMessage, createErr := a.createUserMessage(callContext, queued)
378 if createErr != nil {
821 if createErr != nil {
379 return callContext, prepared, createErr
822 return callContext, prepared, createErr
@@ -575,13 +1018,34 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
575 isHyper := largeModel.ModelCfg.Provider == hyper.Name
1018 isHyper := largeModel.ModelCfg.Provider == hyper.Name
576 isCancelErr := errors.Is(err, context.Canceled)
1019 isCancelErr := errors.Is(err, context.Canceled)
577 if currentAssistant == nil {
1020 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
1029+ if persistErr := a.persistCanceledTurn(ctx, call, userMsgCreated); persistErr != nil {
1030+ return nil, persistErr
578 return result, err
1033 return result, err
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()
580 // Ensure we finish thinking on error to close the reasoning state.
1044 // Ensure we finish thinking on error to close the reasoning state.
581 currentAssistant.FinishThinking()
1045 currentAssistant.FinishThinking()
582 toolCalls := currentAssistant.ToolCalls()
1046 toolCalls := currentAssistant.ToolCalls()
583− // INFO: we use the parent context here because the genCtx has been cancelled.
1047+ // INFO: we use the cleanup context here because the genCtx has been cancelled.
584− msgs, createErr := a.messages.List(ctx, currentAssistant.SessionID)
1048+ msgs, createErr := a.messages.List(cleanupCtx, currentAssistant.SessionID)
585 if createErr != nil {
1049 if createErr != nil {
586 return nil, createErr
1050 return nil, createErr
@@ -590,7 +1054,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
590 tc.Finished = true
1054 tc.Finished = true
591 tc.Input = "{}"
1055 tc.Input = "{}"
592 currentAssistant.AddToolCall(tc)
1056 currentAssistant.AddToolCall(tc)
593− updateErr := a.messages.Update(ctx, *currentAssistant)
1057+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)
594 if updateErr != nil {
1058 if updateErr != nil {
595 return nil, updateErr
1059 return nil, updateErr
@@ -623,7 +1087,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
623 Content: content,
1087 Content: content,
624 IsError: true,
1088 IsError: true,
626− _, createErr = a.messages.Create(ctx, currentAssistant.SessionID, message.CreateMessageParams{
1090+ _, createErr = a.messages.Create(cleanupCtx, currentAssistant.SessionID, message.CreateMessageParams{
627 Role: message.Tool,
1091 Role: message.Tool,
628 Parts: []message.ContentPart{
1092 Parts: []message.ContentPart{
629 toolResult,
1093 toolResult,
@@ -662,9 +1126,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
662 } else {
1126 } else {
663 currentAssistant.AddFinish(message.FinishReasonError, defaultTitle, err.Error())
1127 currentAssistant.AddFinish(message.FinishReasonError, defaultTitle, err.Error())
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
666 // cancelled.
1130 // cancelled.
667− updateErr := a.messages.Update(ctx, *currentAssistant)
1131+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)
668 if updateErr != nil {
1132 if updateErr != nil {
669 return nil, updateErr
1133 return nil, updateErr
@@ -705,18 +1169,106 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
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)
708− queuedMessages, ok := a.messageQueue.Get(call.SessionID)
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 {
1198+ canceledRunIDDrops = append(canceledRunIDDrops, q)
1202+ kept = append(kept, q)
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)
709− if !ok || len(queuedMessages) == 0 {
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
1218+ a.messageQueue.Del(call.SessionID)
1219+ a.acceptedMu.Lock()
1220+ inFlight, _ := a.acceptedRuns.Get(call.SessionID)
1221+ a.acceptedMu.Unlock()
1223+ a.cancelMark.Del(call.SessionID)
710 return result, err
1226 return result, err
712− // There are queued messages restart the loop. The recursive Run
1228+ // There are queued messages, restart the loop. Suppress the outer
713− // publishes its own RunComplete for the queued prompt, so suppress
1229+ // defer's emit: it would otherwise observe the recursive Run's retErr
714− // the outer defer's emit to avoid a duplicate event whose Error
1230+ // (named-return clobbering through the return below) against this
715− // field would belong to the recursive turn but whose MessageID/Text
1231+ // turn's MessageID/Text and publish a mixed, racing event.
716− // would belong to the outer turn.
717 skipRunComplete = true
1232 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
718 firstQueuedMessage := queuedMessages[0]
1251 firstQueuedMessage := queuedMessages[0]
719 a.messageQueue.Set(call.SessionID, queuedMessages[1:])
1252 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)
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()
1267+ if ctx.Err() != nil {
1268+ complete.Cancelled = true
1270+ a.publishRunComplete(ctx, call, complete)
720 return a.Run(ctx, firstQueuedMessage)
1272 return a.Run(ctx, firstQueuedMessage)
@@ -1269,6 +1821,16 @@ func summaryCompletionTokens(usage fantasy.Usage, summaryMessage message.Message
1271func (a *sessionAgent) Cancel(sessionID string) {
1823func (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)
1272 // Cancel regular requests. Don't use Take() here - we need the entry to
1834 // Cancel regular requests. Don't use Take() here - we need the entry to
1273 // remain in activeRequests so IsBusy() returns true until the goroutine
1835 // remain in activeRequests so IsBusy() returns true until the goroutine
1274 // fully completes (including error handling that may access the DB).
1836 // fully completes (including error handling that may access the DB).
@@ -1284,16 +1846,40 @@ func (a *sessionAgent) Cancel(sessionID string) {
1284 cancel()
1846 cancel()
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.
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))
1287 if a.QueuedPrompts(sessionID) > 0 {
1873 if a.QueuedPrompts(sessionID) > 0 {
1288 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1874 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1289− a.messageQueue.Del(sessionID)
1875+ a.clearQueueAndNotify(sessionID)
1293func (a *sessionAgent) ClearQueue(sessionID string) {
1879func (a *sessionAgent) ClearQueue(sessionID string) {
1294 if a.QueuedPrompts(sessionID) > 0 {
1880 if a.QueuedPrompts(sessionID) > 0 {
1295 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1881 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1296− a.messageQueue.Del(sessionID)
1882+ a.clearQueueAndNotify(sessionID)