git.iko.soy/crush
on main help sign in
chore(server): honor cancels immediately after prompt acceptance
4347015Christian Rocha
internal/agent/accepted_run_test.gonew file+226
compare with main
unifiedsplit
@@ -0,0 +1,226 @@
1+package agent
2+
3+import (
4+ "context"
5+ "testing"
6+
7+ "github.com/charmbracelet/crush/internal/message"
8+ "github.com/stretchr/testify/assert"
9+ "github.com/stretchr/testify/require"
10+)
11+
12+// newCancelTestAgent builds a DB-backed sessionAgent with no model. The
13+// tests here exercise the dispatch/cancel/persist paths, none of which
14+// reach agent.Stream, so a model is unnecessary.
15+func newCancelTestAgent(t *testing.T) (*sessionAgent, fakeEnv) {
16+ t.Helper()
17+ env := testEnv(t)
18+ sa := NewSessionAgent(SessionAgentOptions{
19+ Sessions: env.sessions,
20+ Messages: env.messages,
21+ }).(*sessionAgent)
22+ return sa, env
23+}
24+
25+func (a *sessionAgent) acceptedCount(sessionID string) int {
26+ c, _ := a.acceptedRuns.Get(sessionID)
27+ return c
28+}
29+
30+func (a *sessionAgent) hasPendingCancel(sessionID string) bool {
31+ _, ok := a.pendingCancels.Get(sessionID)
32+ return ok
33+}
34+
35+func TestAcceptedRun_CloseIsIdempotent(t *testing.T) {
36+ t.Parallel()
37+ sa, _ := newCancelTestAgent(t)
38+
39+ accept := sa.BeginAccepted("sid")
40+ require.Equal(t, "sid", accept.SessionID())
41+ require.Equal(t, 1, sa.acceptedCount("sid"))
42+
43+ accept.Close()
44+ require.Equal(t, 0, sa.acceptedCount("sid"))
45+
46+ // Repeated Close must not underflow the counter.
47+ accept.Close()
48+ accept.Close()
49+ require.Equal(t, 0, sa.acceptedCount("sid"))
50+}
51+
52+func TestAcceptedRun_MultipleReservations(t *testing.T) {
53+ t.Parallel()
54+ sa, _ := newCancelTestAgent(t)
55+
56+ a1 := sa.BeginAccepted("sid")
57+ a2 := sa.BeginAccepted("sid")
58+ require.Equal(t, 2, sa.acceptedCount("sid"))
59+
60+ a1.Close()
61+ require.Equal(t, 1, sa.acceptedCount("sid"))
62+
63+ a2.Close()
64+ require.Equal(t, 0, sa.acceptedCount("sid"))
65+}
66+
67+func TestAcceptedRun_NilSafe(t *testing.T) {
68+ t.Parallel()
69+ var accept *AcceptedRun
70+ require.Equal(t, "", accept.SessionID())
71+ // Must not panic.
72+ accept.Close()
73+}
74+
75+func TestCancel_IdleDoesNotRecordPendingCancel(t *testing.T) {
76+ t.Parallel()
77+ sa, _ := newCancelTestAgent(t)
78+
79+ // No accepted run, no active request: a true no-op.
80+ sa.Cancel("sid")
81+ require.False(t, sa.hasPendingCancel("sid"))
82+}
83+
84+func TestCancel_AcceptedRecordsPendingCancel(t *testing.T) {
85+ t.Parallel()
86+ sa, _ := newCancelTestAgent(t)
87+
88+ accept := sa.BeginAccepted("sid")
89+ defer accept.Close()
90+
91+ sa.Cancel("sid")
92+ require.True(t, sa.hasPendingCancel("sid"))
93+}
94+
95+func TestCancel_SecondCancelWhilePendingIsNoOp(t *testing.T) {
96+ t.Parallel()
97+ sa, _ := newCancelTestAgent(t)
98+
99+ accept := sa.BeginAccepted("sid")
100+ defer accept.Close()
101+
102+ sa.Cancel("sid")
103+ require.True(t, sa.hasPendingCancel("sid"))
104+
105+ // A second cancel while a pending cancel is already recorded must
106+ // remain a single pending cancel; one Run consumes exactly one.
107+ sa.Cancel("sid")
108+ require.True(t, sa.hasPendingCancel("sid"))
109+}
110+
111+func TestRun_CancelOnEntryPersistsCanceledTurn(t *testing.T) {
112+ t.Parallel()
113+ sa, env := newCancelTestAgent(t)
114+
115+ sess, err := env.sessions.Create(t.Context(), "session")
116+ require.NoError(t, err)
117+
118+ accept := sa.BeginAccepted(sess.ID)
119+ // A cancel arrives in the accepted-but-not-yet-active window.
120+ sa.Cancel(sess.ID)
121+ require.True(t, sa.hasPendingCancel(sess.ID))
122+
123+ result, err := sa.Run(t.Context(), SessionAgentCall{
124+ SessionID: sess.ID,
125+ Prompt: "hello",
126+ Accepted: accept,
127+ })
128+ require.NoError(t, err)
129+ require.Nil(t, result)
130+
131+ // The pending cancel was consumed and the accept released.
132+ require.False(t, sa.hasPendingCancel(sess.ID))
133+ require.Equal(t, 0, sa.acceptedCount(sess.ID))
134+
135+ msgs, err := env.messages.List(t.Context(), sess.ID)
136+ require.NoError(t, err)
137+ require.Len(t, msgs, 2)
138+ assert.Equal(t, message.User, msgs[0].Role)
139+ assert.Equal(t, message.Assistant, msgs[1].Role)
140+ assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason())
141+}
142+
143+func TestPersistCanceledTurn_WritesBothWhenUserMissing(t *testing.T) {
144+ t.Parallel()
145+ sa, env := newCancelTestAgent(t)
146+
147+ sess, err := env.sessions.Create(t.Context(), "session")
148+ require.NoError(t, err)
149+
150+ err = sa.persistCanceledTurn(t.Context(), SessionAgentCall{
151+ SessionID: sess.ID,
152+ Prompt: "hello",
153+ }, false)
154+ require.NoError(t, err)
155+
156+ msgs, err := env.messages.List(t.Context(), sess.ID)
157+ require.NoError(t, err)
158+ require.Len(t, msgs, 2)
159+ assert.Equal(t, message.User, msgs[0].Role)
160+ assert.Equal(t, message.Assistant, msgs[1].Role)
161+ assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason())
162+}
163+
164+func TestPersistCanceledTurn_WritesAssistantOnlyWhenUserCreated(t *testing.T) {
165+ t.Parallel()
166+ sa, env := newCancelTestAgent(t)
167+
168+ sess, err := env.sessions.Create(t.Context(), "session")
169+ require.NoError(t, err)
170+
171+ // Simulate PrepareStep having already created the user message.
172+ _, err = sa.createUserMessage(t.Context(), SessionAgentCall{
173+ SessionID: sess.ID,
174+ Prompt: "hello",
175+ })
176+ require.NoError(t, err)
177+
178+ err = sa.persistCanceledTurn(t.Context(), SessionAgentCall{
179+ SessionID: sess.ID,
180+ Prompt: "hello",
181+ }, true)
182+ require.NoError(t, err)
183+
184+ msgs, err := env.messages.List(t.Context(), sess.ID)
185+ require.NoError(t, err)
186+ require.Len(t, msgs, 2)
187+ assert.Equal(t, message.User, msgs[0].Role)
188+ assert.Equal(t, message.Assistant, msgs[1].Role)
189+ assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason())
190+}
191+
192+func TestPersistCanceledTurn_SucceedsWithCanceledContext(t *testing.T) {
193+ t.Parallel()
194+ sa, env := newCancelTestAgent(t)
195+
196+ sess, err := env.sessions.Create(t.Context(), "session")
197+ require.NoError(t, err)
198+
199+ // Simulate workspace shutdown having already canceled the run
200+ // context. WithoutCancel must let the writes through.
201+ ctx, cancel := context.WithCancel(t.Context())
202+ cancel()
203+
204+ err = sa.persistCanceledTurn(ctx, SessionAgentCall{
205+ SessionID: sess.ID,
206+ Prompt: "hello",
207+ }, false)
208+ require.NoError(t, err)
209+
210+ msgs, err := env.messages.List(t.Context(), sess.ID)
211+ require.NoError(t, err)
212+ require.Len(t, msgs, 2)
213+}
214+
215+func TestClearPendingCancel(t *testing.T) {
216+ t.Parallel()
217+ sa, _ := newCancelTestAgent(t)
218+
219+ accept := sa.BeginAccepted("sid")
220+ defer accept.Close()
221+ sa.Cancel("sid")
222+ require.True(t, sa.hasPendingCancel("sid"))
223+
224+ sa.clearPendingCancel("sid")
225+ require.False(t, sa.hasPendingCancel("sid"))
226+}
internal/agent/agent.go+336−24
@@ -21,6 +21,7 @@ import (
21 "strconv"
21 "strconv"
22 "strings"
22 "strings"
23 "sync"
23 "sync"
24+ "sync/atomic"
24 "time"
25 "time"
25
26
26 "charm.land/catwalk/pkg/catwalk"
27 "charm.land/catwalk/pkg/catwalk"
@@ -103,10 +104,19 @@ 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
106}
115}
107
116
108type SessionAgent interface {
117type SessionAgent interface {
109 Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)
118 Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)
119+ BeginAccepted(sessionID string) *AcceptedRun
110 SetModels(large Model, small Model)
120 SetModels(large Model, small Model)
111 SetTools(tools []fantasy.AgentTool)
121 SetTools(tools []fantasy.AgentTool)
112 SetSystemPrompt(systemPrompt string)
122 SetSystemPrompt(systemPrompt string)
@@ -145,6 +155,32 @@ type sessionAgent struct {
145
155
146 messageQueue *csync.Map[string, []SessionAgentCall]
156 messageQueue *csync.Map[string, []SessionAgentCall]
147 activeRequests *csync.Map[string, context.CancelFunc]
157 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
148}
184}
149
185
150type SessionAgentOptions struct {
186type SessionAgentOptions struct {
@@ -180,7 +216,144 @@ func NewSessionAgent(
180 runComplete: opts.RunComplete,
216 runComplete: opts.RunComplete,
181 messageQueue: csync.NewMap[string, []SessionAgentCall](),
217 messageQueue: csync.NewMap[string, []SessionAgentCall](),
182 activeRequests: csync.NewMap[string, context.CancelFunc](),
218 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{}
183 }
312 }
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)
184}
357}
185
358
186// ValidateCall performs the cheap structural validation that
359// ValidateCall performs the cheap structural validation that
@@ -204,22 +377,73 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
204 return nil, err
377 return nil, err
205 }
378 }
206
379
207− // Queue the message if busy. Strip OnComplete: the caller that
380+ // genCtx/cancel are the run context and its cancel func. For the
208− // supplied the hook (typically coordinator.Run) has its own
381+ // accepted (fire-and-forget) dispatch path they are created under
209− // retry/coalesce scope that ends when it returns, so by the time
382+ // dispatchMu below so a concurrent Cancel can observe the
210− // the queue drains nobody is left to consume the buffered
383+ // activeRequests entry before the assistant message exists. For
211− // terminal event. The recursive Run will fall back to the
384+ // the in-process path they stay nil here and are created later,
212− // default broker publish, which is what existing subscribers
385+ // preserving the original ordering.
213− // expect for queued turns.
386+ var (
214− if a.IsSessionBusy(call.SessionID) {
387+ genCtx context.Context
215− existing, ok := a.messageQueue.Get(call.SessionID)
388+ cancel context.CancelFunc
216− if !ok {
389+ activeRegistered bool
217− existing = []SessionAgentCall{}
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
218 }
424 }
219− queued := call
425+
220− queued.OnComplete = nil
426+ // Idle: become the active run. Register the cancel func before
221− existing = append(existing, queued)
427+ // dropping the lock so a Cancel that arrives between here and
222− a.messageQueue.Set(call.SessionID, existing)
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)
223 return nil, nil
447 return nil, nil
224 }
448 }
225
449
@@ -282,15 +506,22 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
282 if err != nil {
506 if err != nil {
283 return nil, err
507 return nil, err
284 }
508 }
509+ userMsgCreated = true
285
510
286 // Add the session to the context.
511 // Add the session to the context.
287 ctx = context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
512 ctx = context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID)
288
513
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 {
289− genCtx, cancel := context.WithCancel(ctx)
519+ genCtx, cancel = context.WithCancel(ctx)
290− a.activeRequests.Set(call.SessionID, cancel)
520+ a.activeRequests.Set(call.SessionID, cancel)
291
521
292− defer cancel()
522+ defer cancel()
293− defer a.activeRequests.Del(call.SessionID)
523+ defer a.activeRequests.Del(call.SessionID)
524+ }
294 // skipRunComplete is set just before the queued-recursion path so
525 // skipRunComplete is set just before the queued-recursion path so
295 // the outer Run doesn't publish a RunComplete that would race
526 // the outer Run doesn't publish a RunComplete that would race
296 // with — and be superseded by — the recursive call's own
527 // with — and be superseded by — the recursive call's own
@@ -390,14 +621,24 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
390 // Use latest tools (updated by SetTools when MCP tools change).
621 // Use latest tools (updated by SetTools when MCP tools change).
391 prepared.Tools = a.tools.Copy()
622 prepared.Tools = a.tools.Copy()
392
623
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)
393 queuedCalls, _ := a.messageQueue.Get(call.SessionID)
631 queuedCalls, _ := a.messageQueue.Get(call.SessionID)
394 a.messageQueue.Del(call.SessionID)
632 a.messageQueue.Del(call.SessionID)
633+ dispatchLock.Unlock()
634+ if !canceled {
395− for _, queued := range queuedCalls {
635+ for _, queued := range queuedCalls {
396− userMessage, createErr := a.createUserMessage(callContext, queued)
636+ userMessage, createErr := a.createUserMessage(callContext, queued)
397− if createErr != nil {
637+ if createErr != nil {
398− return callContext, prepared, createErr
638+ return callContext, prepared, createErr
639+ }
640+ prepared.Messages = append(prepared.Messages, userMessage.ToAIMessage()...)
399 }
641 }
400− prepared.Messages = append(prepared.Messages, userMessage.ToAIMessage()...)
401 }
642 }
402
643
403 prepared.Messages = a.workaroundProviderMediaLimitations(prepared.Messages, largeModel)
644 prepared.Messages = a.workaroundProviderMediaLimitations(prepared.Messages, largeModel)
@@ -594,6 +835,18 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
594 isHyper := largeModel.ModelCfg.Provider == hyper.Name
835 isHyper := largeModel.ModelCfg.Provider == hyper.Name
595 isCancelErr := errors.Is(err, context.Canceled)
836 isCancelErr := errors.Is(err, context.Canceled)
596 if currentAssistant == nil {
837 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+ }
597 return result, err
850 return result, err
598 }
851 }
599 // Persist final state with a context detached from the run
852 // Persist final state with a context detached from the run
@@ -741,8 +994,35 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
741 })
994 })
742 }
995 }
743
996
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+ }
744 queuedMessages, ok := a.messageQueue.Get(call.SessionID)
1020 queuedMessages, ok := a.messageQueue.Get(call.SessionID)
745 if !ok || len(queuedMessages) == 0 {
1021 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()
746 return result, err
1026 return result, err
747 }
1027 }
748 // There are queued messages restart the loop. The recursive Run
1028 // There are queued messages restart the loop. The recursive Run
@@ -753,6 +1033,14 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
753 skipRunComplete = true
1033 skipRunComplete = true
754 firstQueuedMessage := queuedMessages[0]
1034 firstQueuedMessage := queuedMessages[0]
755 a.messageQueue.Set(call.SessionID, queuedMessages[1:])
1035 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()
756 return a.Run(ctx, firstQueuedMessage)
1044 return a.Run(ctx, firstQueuedMessage)
757}
1045}
758
1046
@@ -1305,6 +1593,16 @@ func summaryCompletionTokens(usage fantasy.Usage, summaryMessage message.Message
1305}
1593}
1306
1594
1307func (a *sessionAgent) Cancel(sessionID string) {
1595func (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+
1308 // Cancel regular requests. Don't use Take() here - we need the entry to
1606 // Cancel regular requests. Don't use Take() here - we need the entry to
1309 // remain in activeRequests so IsBusy() returns true until the goroutine
1607 // remain in activeRequests so IsBusy() returns true until the goroutine
1310 // fully completes (including error handling that may access the DB).
1608 // fully completes (including error handling that may access the DB).
@@ -1320,6 +1618,20 @@ func (a *sessionAgent) Cancel(sessionID string) {
1320 cancel()
1618 cancel()
1321 }
1619 }
1322
1620
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+
1323 if a.QueuedPrompts(sessionID) > 0 {
1635 if a.QueuedPrompts(sessionID) > 0 {
1324 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1636 slog.Debug("Clearing queued prompts", "session_id", sessionID)
1325 a.messageQueue.Del(sessionID)
1637 a.messageQueue.Del(sessionID)
internal/agent/coordinator.go+33
@@ -81,6 +81,15 @@ type Coordinator interface {
81 // INFO: (kujtim) this is not used yet we will use this when we have multiple agents
81 // INFO: (kujtim) this is not used yet we will use this when we have multiple agents
82 // SetMainAgent(string)
82 // SetMainAgent(string)
83 Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error)
83 Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error)
84+ // RunAccepted runs a call that was already accepted via
85+ // BeginAccepted on the fire-and-forget dispatch path. The handle is
86+ // the only carrier of accept-state across the backend.runAgent /
87+ // Coordinator / sessionAgent.Run layers: it reaches
88+ // sessionAgent.Run as SessionAgentCall.Accepted, where it is
89+ // consumed under dispatchMu once the accepted -> (cancel-on-entry |
90+ // queued | active) transition is chosen.
91+ RunAccepted(ctx context.Context, accept *AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error)
92+ BeginAccepted(sessionID string) *AcceptedRun
84 Cancel(sessionID string)
93 Cancel(sessionID string)
85 CancelAll()
94 CancelAll()
86 IsSessionBusy(sessionID string) bool
95 IsSessionBusy(sessionID string) bool
@@ -179,6 +188,20 @@ func NewCoordinator(
179
188
180// Run implements Coordinator.
189// Run implements Coordinator.
181func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
190func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
191+ return c.run(ctx, nil, sessionID, prompt, attachments...)
192+}
193+
194+// RunAccepted implements Coordinator.
195+func (c *coordinator) RunAccepted(ctx context.Context, accept *AcceptedRun, sessionID string, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
196+ return c.run(ctx, accept, sessionID, prompt, attachments...)
197+}
198+
199+// run is the shared implementation behind Run and RunAccepted. When
200+// accept is non-nil it is threaded onto the SessionAgentCall as
201+// Accepted so sessionAgent.Run can consume the accept reservation under
202+// dispatchMu; when nil (the in-process/local path) no accept tracking
203+// applies.
204+func (c *coordinator) run(ctx context.Context, accept *AcceptedRun, sessionID string, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
182 if err := c.readyWg.Wait(); err != nil {
205 if err := c.readyWg.Wait(); err != nil {
183 return nil, err
206 return nil, err
184 }
207 }
@@ -256,6 +279,7 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,
256 FrequencyPenalty: freqPenalty,
279 FrequencyPenalty: freqPenalty,
257 PresencePenalty: presPenalty,
280 PresencePenalty: presPenalty,
258 OnComplete: onComplete,
281 OnComplete: onComplete,
282+ Accepted: accept,
259 })
283 })
260 }
284 }
261 beforeLoaded := c.skillTracker.LoadedNames()
285 beforeLoaded := c.skillTracker.LoadedNames()
@@ -989,6 +1013,15 @@ func isExactoSupported(modelID string) bool {
989 return slices.Contains(supportedModels, modelID)
1013 return slices.Contains(supportedModels, modelID)
990}
1014}
991
1015
1016+// BeginAccepted reserves an accept slot for sessionID on the active
1017+// agent and returns the ownership handle. It is the fire-and-forget
1018+// dispatch path's only way to mark a run as accepted-but-not-yet-active
1019+// so a cancel arriving before the run registers in activeRequests is not
1020+// lost.
1021+func (c *coordinator) BeginAccepted(sessionID string) *AcceptedRun {
1022+ return c.currentAgent.BeginAccepted(sessionID)
1023+}
1024+
992func (c *coordinator) Cancel(sessionID string) {
1025func (c *coordinator) Cancel(sessionID string) {
993 c.currentAgent.Cancel(sessionID)
1026 c.currentAgent.Cancel(sessionID)
994}
1027}
internal/agent/coordinator_test.go+4
@@ -25,6 +25,10 @@ func (m *mockSessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fan
25 return m.runFunc(ctx, call)
25 return m.runFunc(ctx, call)
26}
26}
27
27
28+func (m *mockSessionAgent) BeginAccepted(sessionID string) *AcceptedRun {
29+ return &AcceptedRun{sessionID: sessionID}
30+}
31+
28func (m *mockSessionAgent) Model() Model { return m.model }
32func (m *mockSessionAgent) Model() Model { return m.model }
29func (m *mockSessionAgent) SetModels(large, small Model) {}
33func (m *mockSessionAgent) SetModels(large, small Model) {}
30func (m *mockSessionAgent) SetTools(tools []fantasy.AgentTool) {}
34func (m *mockSessionAgent) SetTools(tools []fantasy.AgentTool) {}
internal/agent/dispatch_cancel_test.gonew file+198
@@ -0,0 +1,198 @@
1+package agent
2+
3+import (
4+ "context"
5+ "errors"
6+ "sync/atomic"
7+ "testing"
8+
9+ "charm.land/fantasy"
10+ "github.com/charmbracelet/crush/internal/message"
11+ "github.com/stretchr/testify/assert"
12+ "github.com/stretchr/testify/require"
13+)
14+
15+// finishStreamModel is a minimal fantasy.LanguageModel that streams a
16+// single text part followed by a normal (FinishReasonStop) finish. It
17+// is enough to drive sessionAgent.Run through PrepareStep and a clean
18+// completion without a recorded provider cassette.
19+type finishStreamModel struct {
20+ text string
21+}
22+
23+func (m *finishStreamModel) Provider() string { return "fake" }
24+func (m *finishStreamModel) Model() string { return "fake-model" }
25+
26+func (m *finishStreamModel) Generate(ctx context.Context, call fantasy.Call) (*fantasy.Response, error) {
27+ return &fantasy.Response{
28+ Content: fantasy.ResponseContent{fantasy.TextContent{Text: m.text}},
29+ FinishReason: fantasy.FinishReasonStop,
30+ }, nil
31+}
32+
33+func (m *finishStreamModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) {
34+ text := m.text
35+ return func(yield func(fantasy.StreamPart) bool) {
36+ if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "1"}) {
37+ return
38+ }
39+ if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "1", Delta: text}) {
40+ return
41+ }
42+ if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "1"}) {
43+ return
44+ }
45+ yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop})
46+ }, nil
47+}
48+
49+func (m *finishStreamModel) GenerateObject(ctx context.Context, call fantasy.ObjectCall) (*fantasy.ObjectResponse, error) {
50+ return nil, errors.New("not implemented")
51+}
52+
53+func (m *finishStreamModel) StreamObject(ctx context.Context, call fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) {
54+ return nil, errors.New("not implemented")
55+}
56+
57+func newStreamTestAgent(t *testing.T) (*sessionAgent, fakeEnv) {
58+ t.Helper()
59+ env := testEnv(t)
60+ model := &finishStreamModel{text: "done"}
61+ sa := testSessionAgent(env, model, model, "system").(*sessionAgent)
62+ return sa, env
63+}
64+
65+// TestCancel_ActiveAndAcceptedFiresBothBranches covers the case where a
66+// session is actively running (activeRequests set) AND a follow-up has
67+// been accepted (acceptedRuns > 0). A single Cancel must fire both: it
68+// invokes the active cancel func and records a pending cancel for the
69+// accepted follow-up.
70+func TestCancel_ActiveAndAcceptedFiresBothBranches(t *testing.T) {
71+ t.Parallel()
72+ sa, _ := newCancelTestAgent(t)
73+
74+ const sid = "sid"
75+ var activeCanceled atomic.Bool
76+ sa.activeRequests.Set(sid, func() { activeCanceled.Store(true) })
77+
78+ accept := sa.BeginAccepted(sid)
79+ defer accept.Close()
80+
81+ sa.Cancel(sid)
82+
83+ require.True(t, activeCanceled.Load(), "active cancel func must fire")
84+ require.True(t, sa.hasPendingCancel(sid), "accepted follow-up must record a pending cancel")
85+}
86+
87+// TestRun_BusyWithPendingCancelTakesCancelOnEntry covers the busy-queue
88+// branch consulting pendingCancels: when the session is busy AND a
89+// cancel has been recorded for an accepted follow-up, Run must take the
90+// cancel-on-entry path (persist a canceled turn) instead of enqueueing
91+// the call behind the active run.
92+func TestRun_BusyWithPendingCancelTakesCancelOnEntry(t *testing.T) {
93+ t.Parallel()
94+ sa, env := newCancelTestAgent(t)
95+
96+ sess, err := env.sessions.Create(t.Context(), "session")
97+ require.NoError(t, err)
98+
99+ // Make the session look busy: an earlier prompt is active.
100+ sa.activeRequests.Set(sess.ID, func() {})
101+
102+ accept := sa.BeginAccepted(sess.ID)
103+ // A cancel arrives while this follow-up is accepted-but-not-active.
104+ sa.Cancel(sess.ID)
105+ require.True(t, sa.hasPendingCancel(sess.ID))
106+
107+ result, err := sa.Run(t.Context(), SessionAgentCall{
108+ SessionID: sess.ID,
109+ Prompt: "follow-up",
110+ Accepted: accept,
111+ })
112+ require.NoError(t, err)
113+ require.Nil(t, result)
114+
115+ // The follow-up was canceled on entry, not enqueued.
116+ require.Equal(t, 0, sa.QueuedPrompts(sess.ID),
117+ "cancel-on-entry must not enqueue the follow-up behind the active run")
118+ require.False(t, sa.hasPendingCancel(sess.ID), "pending cancel must be consumed")
119+ require.Equal(t, 0, sa.acceptedCount(sess.ID), "accept reservation must be released")
120+
121+ msgs, err := env.messages.List(t.Context(), sess.ID)
122+ require.NoError(t, err)
123+ require.Len(t, msgs, 2)
124+ assert.Equal(t, message.User, msgs[0].Role)
125+ assert.Equal(t, message.Assistant, msgs[1].Role)
126+ assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason())
127+}
128+
129+// TestRun_PrepareStepDrainSkipsQueuedOnPendingCancel verifies that the
130+// queue drain inside PrepareStep skips queued follow-up prompts when a
131+// cancel has been recorded for the session: the queued prompt must not
132+// be folded into the active turn as an extra user message.
133+func TestRun_PrepareStepDrainSkipsQueuedOnPendingCancel(t *testing.T) {
134+ t.Parallel()
135+ sa, env := newStreamTestAgent(t)
136+
137+ sess, err := env.sessions.Create(t.Context(), "session")
138+ require.NoError(t, err)
139+
140+ // A follow-up prompt sits queued for the session.
141+ sa.enqueueCall(SessionAgentCall{SessionID: sess.ID, Prompt: "queued-followup"})
142+ // A cancel was recorded for the session while it sat in the queue.
143+ sa.pendingCancels.Set(sess.ID, struct{}{})
144+
145+ result, err := sa.Run(t.Context(), SessionAgentCall{
146+ SessionID: sess.ID,
147+ Prompt: "main",
148+ })
149+ require.NoError(t, err)
150+ require.NotNil(t, result)
151+
152+ // Only the main prompt produced a user message; the queued
153+ // follow-up was skipped, not folded into the turn.
154+ msgs, err := env.messages.List(t.Context(), sess.ID)
155+ require.NoError(t, err)
156+ var userMsgs []message.Message
157+ for _, m := range msgs {
158+ if m.Role == message.User {
159+ userMsgs = append(userMsgs, m)
160+ }
161+ }
162+ require.Len(t, userMsgs, 1, "queued follow-up must not create a user message")
163+ assert.Equal(t, "main", userMsgs[0].Content().String())
164+
165+ // The queue was drained and the pending cancel consumed.
166+ require.Equal(t, 0, sa.QueuedPrompts(sess.ID))
167+ require.False(t, sa.hasPendingCancel(sess.ID))
168+}
169+
170+// TestRun_NormalCompletionClearsStalePendingCancel verifies that a Run
171+// which completes normally clears any stale pending-cancel entry for the
172+// session, so it cannot catch a future run.
173+func TestRun_NormalCompletionClearsStalePendingCancel(t *testing.T) {
174+ t.Parallel()
175+ sa, env := newStreamTestAgent(t)
176+
177+ sess, err := env.sessions.Create(t.Context(), "session")
178+ require.NoError(t, err)
179+
180+ // A stale pending cancel lingers (no queued work, no accepted run).
181+ sa.pendingCancels.Set(sess.ID, struct{}{})
182+
183+ result, err := sa.Run(t.Context(), SessionAgentCall{
184+ SessionID: sess.ID,
185+ Prompt: "main",
186+ })
187+ require.NoError(t, err)
188+ require.NotNil(t, result)
189+
190+ require.False(t, sa.hasPendingCancel(sess.ID),
191+ "normal completion must clear the stale pending cancel")
192+
193+ msgs, err := env.messages.List(t.Context(), sess.ID)
194+ require.NoError(t, err)
195+ require.Len(t, msgs, 2)
196+ assert.Equal(t, message.Assistant, msgs[1].Role)
197+ assert.Equal(t, message.FinishReasonEndTurn, msgs[1].FinishReason())
198+}
internal/server/agent_cancel_test.go+7
@@ -60,6 +60,13 @@ func (s *runCoordinator) Run(ctx context.Context, sessionID, prompt string, atta
60 return nil, s.returnFn(ctx)
60 return nil, s.returnFn(ctx)
61}
61}
62
62
63+func (s *runCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
64+ return s.Run(ctx, sessionID, prompt, attachments...)
65+}
66+
67+func (s *runCoordinator) BeginAccepted(sessionID string) *agent.AcceptedRun {
68+ return nil
69+}
63func (s *runCoordinator) Cancel(string) {}
70func (s *runCoordinator) Cancel(string) {}
64func (s *runCoordinator) CancelAll() {}
71func (s *runCoordinator) CancelAll() {}
65func (s *runCoordinator) IsBusy() bool { return false }
72func (s *runCoordinator) IsBusy() bool { return false }
internal/server/sessions_isbusy_test.go+8
@@ -30,6 +30,14 @@ type stubCoordinator struct {
30func (s *stubCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
30func (s *stubCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
31 return nil, nil
31 return nil, nil
32}
32}
33+
34+func (s *stubCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) {
35+ return nil, nil
36+}
37+
38+func (s *stubCoordinator) BeginAccepted(sessionID string) *agent.AcceptedRun {
39+ return nil
40+}
33func (s *stubCoordinator) Cancel(string) {}
41func (s *stubCoordinator) Cancel(string) {}
34func (s *stubCoordinator) CancelAll() {}
42func (s *stubCoordinator) CancelAll() {}
35func (s *stubCoordinator) IsBusy() bool { return false }
43func (s *stubCoordinator) IsBusy() bool { return false }