@@ -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,19 @@ 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
106115}
107116
108117type SessionAgent interface {
109118 Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)
119+ BeginAccepted(sessionID string) *AcceptedRun
110120 SetModels(large Model, small Model)
111121 SetTools(tools []fantasy.AgentTool)
112122 SetSystemPrompt(systemPrompt string)
@@ -145,6 +155,32 @@ type sessionAgent struct {
145155
146156 messageQueue *csync.Map[string, []SessionAgentCall]
147157 activeRequests *csync.Map[string, context.CancelFunc]
158+
159+ // dispatchMu holds a per-session mutex that serializes the
160+ // accepted -> (cancel-on-entry | queued | active) transition in
161+ // Run against a concurrent Cancel. The lock is held only during
162+ // the brief handoff (no DB or LLM I/O under the lock).
163+ dispatchMu *csync.Map[string, *sync.Mutex]
164+ // acceptedRuns counts dispatched-but-not-yet-active runs per
165+ // session. A counter > 0 means a dispatched prompt is in flight
166+ // and has not yet completed the dispatch handoff in Run. Only
167+ // BeginAccepted increments it; only AcceptedRun.Close decrements
168+ // it.
169+ acceptedRuns *csync.Map[string, int]
170+ // pendingCancels records sessions whose dispatched-but-not-yet-
171+ // running call should observe a cancellation request. It is only
172+ // set by Cancel when acceptedRuns > 0, so an idle Escape never
173+ // poisons the next prompt.
174+ pendingCancels *csync.Map[string, struct{}]
175+ // dispatchMuCreate guards lazy creation of per-session entries in
176+ // dispatchMu so two goroutines can't race to lock different mutex
177+ // instances for the same session.
178+ dispatchMuCreate sync.Mutex
179+ // acceptedMu serializes increments/decrements of acceptedRuns. It
180+ // is separate from dispatchMu so AcceptedRun.Close (which may run
181+ // while Run holds dispatchMu for the same session) does not
182+ // deadlock by re-entering the dispatch lock.
183+ acceptedMu sync.Mutex
148184}
149185
150186type SessionAgentOptions struct {
@@ -180,7 +216,144 @@ func NewSessionAgent(
180216 runComplete: opts.RunComplete,
181217 messageQueue: csync.NewMap[string, []SessionAgentCall](),
182218 activeRequests: csync.NewMap[string, context.CancelFunc](),
219+ dispatchMu: csync.NewMap[string, *sync.Mutex](),
220+ acceptedRuns: csync.NewMap[string, int](),
221+ pendingCancels: csync.NewMap[string, struct{}](),
222+ }
223+}
224+
225+// AcceptedRun owns exactly one accept reservation taken by
226+// BeginAccepted. It is the only carrier of accept-state across the
227+// backend.runAgent / Coordinator.Run / sessionAgent.Run layers: a
228+// counter > 0 means a dispatched prompt is in flight and has not yet
229+// completed the dispatch handoff in Run. Close is the only way to
230+// release the reservation and is idempotent.
231+type AcceptedRun struct {
232+ agent *sessionAgent
233+ sessionID string
234+ done atomic.Bool
235+}
236+
237+// Close decrements the accept counter for this reservation. It is safe
238+// to call multiple times; only the first call has effect.
239+func (r *AcceptedRun) Close() {
240+ if r == nil {
241+ return
242+ }
243+ if !r.done.CompareAndSwap(false, true) {
244+ return
245+ }
246+ r.agent.endAccepted(r.sessionID)
247+}
248+
249+// SessionID exposes the session this reservation is for so the run path
250+// can use it without an extra parameter.
251+func (r *AcceptedRun) SessionID() string {
252+ if r == nil {
253+ return ""
254+ }
255+ return r.sessionID
256+}
257+
258+// BeginAccepted increments the accept counter for sessionID and returns
259+// a handle whose Close is the only way to decrement it. It is the only
260+// entry point that mutates acceptedRuns.
261+func (a *sessionAgent) BeginAccepted(sessionID string) *AcceptedRun {
262+ a.acceptedMu.Lock()
263+ defer a.acceptedMu.Unlock()
264+ count, _ := a.acceptedRuns.Get(sessionID)
265+ a.acceptedRuns.Set(sessionID, count+1)
266+ return &AcceptedRun{agent: a, sessionID: sessionID}
267+}
268+
269+// endAccepted decrements the accept counter for sessionID. It is only
270+// called via AcceptedRun.Close. It uses a dedicated lock (not the
271+// per-session dispatch mutex) so it can run while Run holds dispatchMu
272+// for the same session without deadlocking.
273+func (a *sessionAgent) endAccepted(sessionID string) {
274+ a.acceptedMu.Lock()
275+ defer a.acceptedMu.Unlock()
276+ count, ok := a.acceptedRuns.Get(sessionID)
277+ if !ok || count <= 1 {
278+ a.acceptedRuns.Del(sessionID)
279+ return
280+ }
281+ a.acceptedRuns.Set(sessionID, count-1)
282+}
283+
284+// sessionMu returns the per-session dispatch mutex, creating it on first
285+// use. Creation is guarded so concurrent callers always observe the same
286+// mutex instance for a given session.
287+func (a *sessionAgent) sessionMu(sessionID string) *sync.Mutex {
288+ if mu, ok := a.dispatchMu.Get(sessionID); ok {
289+ return mu
290+ }
291+ a.dispatchMuCreate.Lock()
292+ defer a.dispatchMuCreate.Unlock()
293+ if mu, ok := a.dispatchMu.Get(sessionID); ok {
294+ return mu
295+ }
296+ mu := &sync.Mutex{}
297+ a.dispatchMu.Set(sessionID, mu)
298+ return mu
299+}
300+
301+// enqueueCall appends call to the session's message queue. The
302+// OnComplete hook is stripped: the caller that supplied it (typically
303+// coordinator.Run) has its own retry/coalesce scope that ends when it
304+// returns, so by the time the queue drains nobody is left to consume the
305+// buffered terminal event. The recursive Run falls back to the default
306+// broker publish, which is what existing subscribers expect for queued
307+// turns.
308+func (a *sessionAgent) enqueueCall(call SessionAgentCall) {
309+ existing, ok := a.messageQueue.Get(call.SessionID)
310+ if !ok {
311+ existing = []SessionAgentCall{}
183312 }
313+ queued := call
314+ queued.OnComplete = nil
315+ queued.Accepted = nil
316+ existing = append(existing, queued)
317+ a.messageQueue.Set(call.SessionID, existing)
318+}
319+
320+// clearPendingCancel removes any pending-cancel record for sessionID. It
321+// takes the per-session dispatch lock so it is ordered against Cancel and
322+// the dispatch handoff.
323+func (a *sessionAgent) clearPendingCancel(sessionID string) {
324+ mu := a.sessionMu(sessionID)
325+ mu.Lock()
326+ defer mu.Unlock()
327+ a.pendingCancels.Del(sessionID)
328+}
329+
330+// persistCanceledTurn writes the user/assistant records for a turn that
331+// was canceled before (or just as) streaming would have produced them.
332+// It creates the user message only when it was not already created by an
333+// earlier createUserMessage call (userMsgCreated), then writes an
334+// assistant message with FinishReasonCanceled. Both writes use
335+// context.WithoutCancel(ctx) so workspace shutdown (which cancels the run
336+// context) can't drop them.
337+func (a *sessionAgent) persistCanceledTurn(ctx context.Context, call SessionAgentCall, userMsgCreated bool) error {
338+ writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
339+ defer cancel()
340+ if !userMsgCreated {
341+ if _, err := a.createUserMessage(writeCtx, call); err != nil {
342+ return err
343+ }
344+ }
345+ largeModel := a.largeModel.Get()
346+ assistant, err := a.messages.Create(writeCtx, call.SessionID, message.CreateMessageParams{
347+ Role: message.Assistant,
348+ Parts: []message.ContentPart{},
349+ Model: largeModel.ModelCfg.Model,
350+ Provider: largeModel.ModelCfg.Provider,
351+ })
352+ if err != nil {
353+ return err
354+ }
355+ assistant.AddFinish(message.FinishReasonCanceled, "User canceled request", "")
356+ return a.messages.Update(writeCtx, assistant)
184357}
185358
186359// ValidateCall performs the cheap structural validation that
@@ -204,22 +377,73 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
204377 return nil, err
205378 }
206379
207− // Queue the message if busy. Strip OnComplete: the caller that
208− // supplied the hook (typically coordinator.Run) has its own
209− // retry/coalesce scope that ends when it returns, so by the time
210− // the queue drains nobody is left to consume the buffered
211− // terminal event. The recursive Run will fall back to the
212− // default broker publish, which is what existing subscribers
213− // expect for queued turns.
214− if a.IsSessionBusy(call.SessionID) {
215− existing, ok := a.messageQueue.Get(call.SessionID)
216− if !ok {
217− existing = []SessionAgentCall{}
380+ // genCtx/cancel are the run context and its cancel func. For the
381+ // accepted (fire-and-forget) dispatch path they are created under
382+ // dispatchMu below so a concurrent Cancel can observe the
383+ // activeRequests entry before the assistant message exists. For
384+ // the in-process path they stay nil here and are created later,
385+ // preserving the original ordering.
386+ var (
387+ genCtx context.Context
388+ cancel context.CancelFunc
389+ activeRegistered bool
390+ userMsgCreated bool
391+ )
392+
393+ if call.Accepted != nil {
394+ // Serialize the accepted -> (cancel-on-entry | queued |
395+ // active) transition against a concurrent Cancel. Cancel takes
396+ // the same per-session lock, so every cancel observes at least
397+ // one of: pendingCancels, an activeRequests entry, or a
398+ // messageQueue entry it then clears.
399+ mu := a.sessionMu(call.SessionID)
400+ mu.Lock()
401+
402+ if _, pending := a.pendingCancels.Get(call.SessionID); pending {
403+ // Cancel-on-entry: a cancel arrived while this run was
404+ // dispatched but not yet active. Consume the pending
405+ // cancel, release the accept reservation, drop the lock,
406+ // and persist a canceled turn without entering Stream.
407+ a.pendingCancels.Del(call.SessionID)
408+ call.Accepted.Close()
409+ mu.Unlock()
410+ if err := a.persistCanceledTurn(ctx, call, false); err != nil {
411+ return nil, err
412+ }
413+ return nil, nil
414+ }
415+
416+ if a.IsSessionBusy(call.SessionID) {
417+ // Busy: an earlier prompt is active. Queue this call and
418+ // release the accept reservation. A Cancel arriving after
419+ // this point sees the active entry and clears the queue.
420+ a.enqueueCall(call)
421+ call.Accepted.Close()
422+ mu.Unlock()
423+ return nil, nil
218424 }
219− queued := call
220− queued.OnComplete = nil
221− existing = append(existing, queued)
222− a.messageQueue.Set(call.SessionID, existing)
425+
426+ // Idle: become the active run. Register the cancel func before
427+ // dropping the lock so a Cancel that arrives between here and
428+ // assistant creation is not lost.
429+ runCtx := context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
430+ genCtx, cancel = context.WithCancel(runCtx)
431+ a.activeRequests.Set(call.SessionID, cancel)
432+ activeRegistered = true
433+ call.Accepted.Close()
434+ mu.Unlock()
435+
436+ defer cancel()
437+ defer a.activeRequests.Del(call.SessionID)
438+ } else if a.IsSessionBusy(call.SessionID) {
439+ // Queue the message if busy. Strip OnComplete: the caller that
440+ // supplied the hook (typically coordinator.Run) has its own
441+ // retry/coalesce scope that ends when it returns, so by the time
442+ // the queue drains nobody is left to consume the buffered
443+ // terminal event. The recursive Run will fall back to the
444+ // default broker publish, which is what existing subscribers
445+ // expect for queued turns.
446+ a.enqueueCall(call)
223447 return nil, nil
224448 }
225449
@@ -282,15 +506,22 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
282506 if err != nil {
283507 return nil, err
284508 }
509+ userMsgCreated = true
285510
286511 // Add the session to the context.
287512 ctx = context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
288513
289− genCtx, cancel := context.WithCancel(ctx)
290− a.activeRequests.Set(call.SessionID, cancel)
514+ // For the accepted dispatch path the run context and cancel func
515+ // were already created and registered under dispatchMu above; reuse
516+ // them. For the in-process path create them here, preserving the
517+ // original ordering.
518+ if !activeRegistered {
519+ genCtx, cancel = context.WithCancel(ctx)
520+ a.activeRequests.Set(call.SessionID, cancel)
291521
292− defer cancel()
293− defer a.activeRequests.Del(call.SessionID)
522+ defer cancel()
523+ defer a.activeRequests.Del(call.SessionID)
524+ }
294525 // skipRunComplete is set just before the queued-recursion path so
295526 // the outer Run doesn't publish a RunComplete that would race
296527 // with — and be superseded by — the recursive call's own
@@ -390,14 +621,24 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
390621 // Use latest tools (updated by SetTools when MCP tools change).
391622 prepared.Tools = a.tools.Copy()
392623
624+ // Drain queued follow-up prompts, but skip them if a cancel
625+ // was recorded for the session while they sat in the queue:
626+ // a cancel that arrived after the queue insertion must not
627+ // let the queued prompt run as part of this step.
628+ dispatchLock := a.sessionMu(call.SessionID)
629+ dispatchLock.Lock()
630+ _, canceled := a.pendingCancels.Get(call.SessionID)
393631 queuedCalls, _ := a.messageQueue.Get(call.SessionID)
394632 a.messageQueue.Del(call.SessionID)
395− for _, queued := range queuedCalls {
396− userMessage, createErr := a.createUserMessage(callContext, queued)
397− if createErr != nil {
398− return callContext, prepared, createErr
633+ dispatchLock.Unlock()
634+ if !canceled {
635+ for _, queued := range queuedCalls {
636+ userMessage, createErr := a.createUserMessage(callContext, queued)
637+ if createErr != nil {
638+ return callContext, prepared, createErr
639+ }
640+ prepared.Messages = append(prepared.Messages, userMessage.ToAIMessage()...)
399641 }
400− prepared.Messages = append(prepared.Messages, userMessage.ToAIMessage()...)
401642 }
402643
403644 prepared.Messages = a.workaroundProviderMediaLimitations(prepared.Messages, largeModel)
@@ -594,6 +835,18 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
594835 isHyper := largeModel.ModelCfg.Provider == hyper.Name
595836 isCancelErr := errors.Is(err, context.Canceled)
596837 if currentAssistant == nil {
838+ // Cancel-before-assistant-creation window: the run was
839+ // canceled after activeRequests.Set but before PrepareStep
840+ // created the assistant message. Without this, the turn
841+ // would return with no FinishReasonCanceled marker and no
842+ // user-visible record. The user message was already created
843+ // above, so persistCanceledTurn only writes the assistant
844+ // record.
845+ if isCancelErr {
846+ if persistErr := a.persistCanceledTurn(ctx, call, userMsgCreated); persistErr != nil {
847+ return nil, persistErr
848+ }
849+ }
597850 return result, err
598851 }
599852 // Persist final state with a context detached from the run
@@ -741,8 +994,35 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
741994 })
742995 }
743996
997+ // Hand off to the next queued prompt (if any) under dispatchMu so
998+ // the transition from this finished run to the queued run is atomic
999+ // against a concurrent Cancel. activeRequests for this session was
1000+ // just deleted above, so without the lock there is a window in
1001+ // which the session looks idle and a cancel becomes a no-op that
1002+ // fails to stop the queued prompt. Holding the lock lets us observe
1003+ // a pending cancel recorded against the session and drop the queue
1004+ // instead of running it, and (for the recursion) hand a fresh
1005+ // accept reservation to the dequeued call so acceptedRuns stays > 0
1006+ // across the recursive Run's own dispatch handoff — keeping the
1007+ // session observable to Cancel for the entire transition and
1008+ // closing the dequeue -> re-register window.
1009+ mu := a.sessionMu(call.SessionID)
1010+ mu.Lock()
1011+ if _, pending := a.pendingCancels.Get(call.SessionID); pending {
1012+ // 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
1014+ // the queue instead of running it and consume the marker.
1015+ a.pendingCancels.Del(call.SessionID)
1016+ a.messageQueue.Del(call.SessionID)
1017+ mu.Unlock()
1018+ return result, err
1019+ }
7441020 queuedMessages, ok := a.messageQueue.Get(call.SessionID)
7451021 if !ok || len(queuedMessages) == 0 {
1022+ // No queued work. Clear any stale pending-cancel entry as a
1023+ // safety net so it can't catch a future run (no-op when unset).
1024+ a.pendingCancels.Del(call.SessionID)
1025+ mu.Unlock()
7461026 return result, err
7471027 }
7481028 // There are queued messages restart the loop. The recursive Run
@@ -753,6 +1033,14 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
7531033 skipRunComplete = true
7541034 firstQueuedMessage := queuedMessages[0]
7551035 a.messageQueue.Set(call.SessionID, queuedMessages[1:])
1036+ // Reserve a fresh accept for the dequeued prompt before dropping the
1037+ // lock so acceptedRuns > 0 across the handoff into the recursive
1038+ // Run. This closes the window between this dequeue and the recursive
1039+ // Run registering its activeRequests entry: a cancel arriving in
1040+ // that window now records a pending cancel (acceptedRuns > 0) that
1041+ // the recursive Run's accepted path observes as cancel-on-entry.
1042+ firstQueuedMessage.Accepted = a.BeginAccepted(call.SessionID)
1043+ mu.Unlock()
7561044 return a.Run(ctx, firstQueuedMessage)
7571045}
7581046
@@ -1305,6 +1593,16 @@ func summaryCompletionTokens(usage fantasy.Usage, summaryMessage message.Message
13051593}
13061594
13071595func (a *sessionAgent) Cancel(sessionID string) {
1596+ // Serialize against the dispatch handoff in Run so the accepted ->
1597+ // (cancel-on-entry | queued | active) transition is atomic against
1598+ // this cancel. Every cancel observes at least one of: an active
1599+ // request, an accepted run (recorded as a pending cancel), or a
1600+ // queue entry it then clears. If none of those hold, an idle Escape
1601+ // is a true no-op and must not poison the next prompt.
1602+ mu := a.sessionMu(sessionID)
1603+ mu.Lock()
1604+ defer mu.Unlock()
1605+
13081606 // Cancel regular requests. Don't use Take() here - we need the entry to
13091607 // remain in activeRequests so IsBusy() returns true until the goroutine
13101608 // fully completes (including error handling that may access the DB).
@@ -1320,6 +1618,20 @@ func (a *sessionAgent) Cancel(sessionID string) {
13201618 cancel()
13211619 }
13221620
1621+ // 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
1623+ // about to enter Run's busy-queue branch, while leaving an idle
1624+ // session untouched. Active and accepted are not mutually exclusive:
1625+ // when a run is active and a follow-up has been accepted, both the
1626+ // cancel above and this pending record fire.
1627+ a.acceptedMu.Lock()
1628+ count, ok := a.acceptedRuns.Get(sessionID)
1629+ a.acceptedMu.Unlock()
1630+ if ok && count > 0 {
1631+ slog.Debug("Recording pending cancel for accepted run", "session_id", sessionID)
1632+ a.pendingCancels.Set(sessionID, struct{}{})
1633+ }
1634+
13231635 if a.QueuedPrompts(sessionID) > 0 {
13241636 slog.Debug("Clearing queued prompts", "session_id", sessionID)
13251637 a.messageQueue.Del(sessionID)