d7a814c · crush · git.iko.soy

git.iko.soy/crush

fix(server): close a queued-prompt cancel race

Check queued prompt cancellation while holding the same session handoff lock used for accepted and active prompts. Queue draining now follows the same ordering guarantees as prompt dispatch. Co-Authored-By: Charm Crush <crush@charm.land>

commitd7a814c540c258f06cb798c29beff676ed4b130b authorChristian Rocha <christian@rocha.is> dateMon 1 Jun 2026, 20:11 (4 months ago) parentc6d1590

2 files changed+82−9

compare with main

unifiedsplit

internal/agent/agent.go+26−9 internal/agent/run_complete_test.go+56

internal/agent/agent.go+26−9

@@ -359,6 +359,30 @@ func (a *sessionAgent) enqueueCall(call SessionAgentCall) {

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

360360}

361361

362+// drainUncanceledQueue copies and clears the session's message queue and

363+// returns the queued calls not covered by a pending cancel. The queue

364+// copy/delete and the cancel-mark check happen under the per-session

365+// dispatch mutex so the filtering is atomic against a concurrent Cancel:

366+// canceledBySeq requires the caller to hold that mutex, and evaluating it

367+// here (rather than after unlocking) prevents a cancel recorded between

368+// the drain and the check from being observed inconsistently. The

369+// returned survivors are processed by the caller without the lock held.

370+func (a *sessionAgent) drainUncanceledQueue(sessionID string) []SessionAgentCall {

371+ dispatchLock := a.sessionMu(sessionID)

372+ dispatchLock.Lock()

373+ defer dispatchLock.Unlock()

374+ queuedCalls, _ := a.messageQueue.Get(sessionID)

375+ a.messageQueue.Del(sessionID)

376+ survivors := queuedCalls[:0]

377+ for _, queued := range queuedCalls {

378+ if a.canceledBySeq(sessionID, queued.acceptSeq) {

379+ continue

380+ }

381+ survivors = append(survivors, queued)

382+ }

383+ return survivors

384+}

385+

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

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

364388// and the dispatch handoff.

@@ -717,15 +741,8 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

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

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

719743 // still folded in.

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

721− dispatchLock.Lock()

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

723− a.messageQueue.Del(call.SessionID)

724− dispatchLock.Unlock()

725− for _, queued := range queuedCalls {

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

727− continue

728− }

744+ survivors := a.drainUncanceledQueue(call.SessionID)

745+ for _, queued := range survivors {

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

730747 if createErr != nil {

731748 return callContext, prepared, createErr

internal/agent/run_complete_test.go+56

@@ -58,6 +58,62 @@ func TestSessionAgentRun_QueueStripsOnComplete(t *testing.T) {

5858 "RunComplete still correlates with the originating SendMessage")

5959}

6060

61+// TestDrainUncanceledQueue_FiltersUnderDispatchLock verifies that the

62+// queue drain evaluates the per-session cancel mark while holding the

63+// dispatch mutex (canceledBySeq's documented precondition). Queued calls

64+// at or below the cancel high-water mark are dropped, calls queued after

65+// the cancel (higher seq) survive, untracked enqueues (seq == 0) are

66+// dropped whenever any mark is present, and the queue is cleared.

67+func TestDrainUncanceledQueue_FiltersUnderDispatchLock(t *testing.T) {

68+ t.Parallel()

69+

70+ env := testEnv(t)

71+ a := NewSessionAgent(SessionAgentOptions{

72+ Sessions: env.sessions,

73+ Messages: env.messages,

74+ }).(*sessionAgent)

75+

76+ const sessionID = "drain-session"

77+ a.messageQueue.Set(sessionID, []SessionAgentCall{

78+ {SessionID: sessionID, Prompt: "below", acceptSeq: 1},

79+ {SessionID: sessionID, Prompt: "at-mark", acceptSeq: 2},

80+ {SessionID: sessionID, Prompt: "after", acceptSeq: 3},

81+ {SessionID: sessionID, Prompt: "untracked", acceptSeq: 0},

82+ })

83+ // Cancel high-water mark at seq 2: seq <= 2 and seq == 0 are covered.

84+ a.cancelMark.Set(sessionID, 2)

85+

86+ survivors := a.drainUncanceledQueue(sessionID)

87+

88+ require.Len(t, survivors, 1,

89+ "only the follow-up queued after the cancel (seq > mark) must survive")

90+ require.Equal(t, "after", survivors[0].Prompt)

91+

92+ _, ok := a.messageQueue.Get(sessionID)

93+ require.False(t, ok, "drain must clear the session message queue")

94+}

95+

96+// TestDrainUncanceledQueue_NoMarkKeepsAll verifies that with no cancel

97+// mark recorded, every queued call survives the drain.

98+func TestDrainUncanceledQueue_NoMarkKeepsAll(t *testing.T) {

99+ t.Parallel()

100+

101+ env := testEnv(t)

102+ a := NewSessionAgent(SessionAgentOptions{

103+ Sessions: env.sessions,

104+ Messages: env.messages,

105+ }).(*sessionAgent)

106+

107+ const sessionID = "drain-nomark"

108+ a.messageQueue.Set(sessionID, []SessionAgentCall{

109+ {SessionID: sessionID, Prompt: "a", acceptSeq: 0},

110+ {SessionID: sessionID, Prompt: "b", acceptSeq: 5},

111+ })

112+

113+ survivors := a.drainUncanceledQueue(sessionID)

114+ require.Len(t, survivors, 2, "no cancel mark means all queued calls survive")

115+}

116+

61117// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the

62118// pubsub.Publisher interface change end-to-end: a Broker is the only

63119// concrete Publisher implementation and must satisfy both Publish and