e14d47a · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(noninteractive): crush run reliability in client/server mode

e14d47aChristian Rocha5 months ago

internal/agent/agent.go+95−7

compare with main

unifiedsplit

@@ -70,7 +70,16 @@ var (

70)

72type SessionAgentCall struct {

73− SessionID string

73+ SessionID string

74+ // RunID, when non-empty, is the caller-supplied correlator that

75+ // gets echoed back on the notify.RunComplete event emitted for

76+ // this turn. It is preserved when the call is enqueued behind a

77+ // busy session so the queued turn's terminal event is still

78+ // recognisable to the original caller. Callers that need a

79+ // reliable completion contract (e.g. `crush run` against a

80+ // session that may be busy) MUST set it; SessionID alone is

81+ // ambiguous when concurrent turns share the same session.

82+ RunID string

74 Prompt string

83 Prompt string

75 ProviderOptions fantasy.ProviderOptions

84 ProviderOptions fantasy.ProviderOptions

76 Attachments []message.Attachment

85 Attachments []message.Attachment

@@ -81,6 +90,19 @@ type SessionAgentCall struct {

81 FrequencyPenalty *float64

90 FrequencyPenalty *float64

82 PresencePenalty *float64

91 PresencePenalty *float64

83 NonInteractive bool

92 NonInteractive bool

93+ // OnComplete, when non-nil, replaces the default RunComplete

94+ // publish path: the inner Run hands the terminal payload to this

95+ // callback instead of emitting it on the RunComplete broker. The

96+ // coordinator uses this hook to coalesce the unauthorized →

97+ // re-auth → retry chain into a single user-visible terminal

98+ // event, so non-interactive clients (e.g. `crush run`) don't

99+ // exit on a stale failed-attempt RunComplete before the

100+ // successful retry. It is intentionally stripped when queueing

101+ // a busy-session call (see Run): the originating

102+ // coordinator.Run has long returned by the time the queued

103+ // recursion drains, so falling back to the default broker

104+ // publish keeps the event visible to subscribers.

105+ OnComplete func(notify.RunComplete)

84}

106}

107

86type SessionAgent interface {

108type SessionAgent interface {

@@ -119,6 +141,7 @@ type sessionAgent struct {

119 disableAutoSummarize bool

141 disableAutoSummarize bool

120 isYolo bool

142 isYolo bool

121 notify pubsub.Publisher[notify.Notification]

143 notify pubsub.Publisher[notify.Notification]

144+ runComplete pubsub.Publisher[notify.RunComplete]

122

145

123 messageQueue *csync.Map[string, []SessionAgentCall]

146 messageQueue *csync.Map[string, []SessionAgentCall]

124 activeRequests *csync.Map[string, context.CancelFunc]

147 activeRequests *csync.Map[string, context.CancelFunc]

@@ -136,6 +159,7 @@ type SessionAgentOptions struct {

136 Messages message.Service

159 Messages message.Service

137 Tools []fantasy.AgentTool

160 Tools []fantasy.AgentTool

138 Notify pubsub.Publisher[notify.Notification]

161 Notify pubsub.Publisher[notify.Notification]

162+ RunComplete pubsub.Publisher[notify.RunComplete]

139}

163}

140

164

141func NewSessionAgent(

165func NewSessionAgent(

@@ -153,12 +177,13 @@ func NewSessionAgent(

153 tools: csync.NewSliceFrom(opts.Tools),

177 tools: csync.NewSliceFrom(opts.Tools),

154 isYolo: opts.IsYolo,

178 isYolo: opts.IsYolo,

155 notify: opts.Notify,

179 notify: opts.Notify,

180+ runComplete: opts.RunComplete,

156 messageQueue: csync.NewMap[string, []SessionAgentCall](),

181 messageQueue: csync.NewMap[string, []SessionAgentCall](),

157 activeRequests: csync.NewMap[string, context.CancelFunc](),

182 activeRequests: csync.NewMap[string, context.CancelFunc](),

158 }

183 }

159}

184}

160

185

161−func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) {

186+func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *fantasy.AgentResult, retErr error) {

162 if call.Prompt == "" && !message.ContainsTextAttachment(call.Attachments) {

187 if call.Prompt == "" && !message.ContainsTextAttachment(call.Attachments) {

163 return nil, ErrEmptyPrompt

188 return nil, ErrEmptyPrompt

164 }

189 }

@@ -166,13 +191,21 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

166 return nil, ErrSessionMissing

191 return nil, ErrSessionMissing

167 }

192 }

168

193

169− // Queue the message if busy

194+ // Queue the message if busy. Strip OnComplete: the caller that

195+ // supplied the hook (typically coordinator.Run) has its own

196+ // retry/coalesce scope that ends when it returns, so by the time

197+ // the queue drains nobody is left to consume the buffered

198+ // terminal event. The recursive Run will fall back to the

199+ // default broker publish, which is what existing subscribers

200+ // expect for queued turns.

170 if a.IsSessionBusy(call.SessionID) {

201 if a.IsSessionBusy(call.SessionID) {

171 existing, ok := a.messageQueue.Get(call.SessionID)

202 existing, ok := a.messageQueue.Get(call.SessionID)

172 if !ok {

203 if !ok {

173 existing = []SessionAgentCall{}

204 existing = []SessionAgentCall{}

174 }

205 }

206+ queued := call

207+ queued.OnComplete = nil

175− existing = append(existing, call)

208+ existing = append(existing, queued)

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

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

177 return nil, nil

210 return nil, nil

178 }

211 }

@@ -245,14 +278,65 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

245

278

246 defer cancel()

279 defer cancel()

247 defer a.activeRequests.Del(call.SessionID)

280 defer a.activeRequests.Del(call.SessionID)

281+ // skipRunComplete is set just before the queued-recursion path so

282+ // the outer Run doesn't publish a RunComplete that would race

283+ // with — and be superseded by — the recursive call's own

284+ // RunComplete (each queued user prompt is its own turn and

285+ // publishes exactly one terminal event).

286+ var skipRunComplete bool

287+ // currentAssistant is declared here so the deferred RunComplete

288+ // publish below can capture the pointer that PrepareStep will

289+ // later (re)assign for each streaming step. The final assistant

290+ // message of the turn is the value reachable through this

291+ // pointer when the defer runs.

292+ var currentAssistant *message.Message

248 // Drain any debounced message updates before returning. message.Service

293 // Drain any debounced message updates before returning. message.Service

249 // already flushes synchronously on terminal updates, but a defer here

294 // already flushes synchronously on terminal updates, but a defer here

250 // guarantees the contract at every Run exit (success, error, panic

295 // guarantees the contract at every Run exit (success, error, panic

251 // recovery upstream) without callers needing to know.

296 // recovery upstream) without callers needing to know.

297+ //

298+ // After the flush completes — meaning all per-message

299+ // Publish(UpdatedEvent) calls have fired and been buffered into

300+ // every subscriber's channel — publish the authoritative

301+ // RunComplete event for this turn. The flush-then-publish order

302+ // gives well-behaved clients the best chance of seeing the final

303+ // message event before RunComplete; the embedded Text field

304+ // reconciles for clients that observe the events out of order

305+ // (the pubsub broker fan-in does not serialize publishes from

306+ // different upstream brokers).

252 defer func() {

307 defer func() {

253 if flushErr := a.messages.FlushAll(ctx); flushErr != nil {

308 if flushErr := a.messages.FlushAll(ctx); flushErr != nil {

254 slog.Error("Failed to flush pending message updates after run", "error", flushErr)

309 slog.Error("Failed to flush pending message updates after run", "error", flushErr)

255 }

310 }

311+ if skipRunComplete {

312+ return

313+ }

314+ complete := notify.RunComplete{SessionID: call.SessionID, RunID: call.RunID}

315+ if currentAssistant != nil {

316+ complete.MessageID = currentAssistant.ID

317+ complete.Text = currentAssistant.Content().String()

318+ }

319+ if retErr != nil {

320+ complete.Error = retErr.Error()

321+ complete.Cancelled = errors.Is(retErr, context.Canceled)

322+ } else if ctx.Err() != nil {

323+ complete.Cancelled = true

324+ }

325+ // Prefer the per-call hook when supplied so the coordinator

326+ // can coalesce retries (e.g. unauthorized → re-auth → retry)

327+ // into a single user-visible terminal event. The fallback

328+ // must-deliver publish applies bounded-blocking semantics to

329+ // the authoritative terminal event so a momentarily-full

330+ // subscriber channel can't silently drop it and hang

331+ // non-interactive clients waiting on RunComplete.

332+ if call.OnComplete != nil {

333+ call.OnComplete(complete)

334+ return

335+ }

336+ if a.runComplete == nil {

337+ return

338+ }

339+ a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete)

256 }()

340 }()

257

341

258 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)

342 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)

@@ -260,7 +344,6 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

260 startTime := time.Now()

344 startTime := time.Now()

261 a.eventPromptSent(call.SessionID)

345 a.eventPromptSent(call.SessionID)

262

346

263− var currentAssistant *message.Message

264 var stepMessages []fantasy.Message

347 var stepMessages []fantasy.Message

265 var shouldSummarize bool

348 var shouldSummarize bool

266 // Don't send MaxOutputTokens if 0 — some providers (e.g. LM Studio) reject it

349 // Don't send MaxOutputTokens if 0 — some providers (e.g. LM Studio) reject it

@@ -268,7 +351,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

268 if call.MaxOutputTokens > 0 {

351 if call.MaxOutputTokens > 0 {

269 maxOutputTokens = &call.MaxOutputTokens

352 maxOutputTokens = &call.MaxOutputTokens

270 }

353 }

271− result, err := agent.Stream(genCtx, fantasy.AgentStreamCall{

354+ result, err = agent.Stream(genCtx, fantasy.AgentStreamCall{

272 Prompt: message.PromptWithTextAttachments(call.Prompt, call.Attachments),

355 Prompt: message.PromptWithTextAttachments(call.Prompt, call.Attachments),

273 Files: files,

356 Files: files,

274 Messages: history,

357 Messages: history,

@@ -634,7 +717,12 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

634 if !ok || len(queuedMessages) == 0 {

717 if !ok || len(queuedMessages) == 0 {

635 return result, err

718 return result, err

636 }

719 }

637− // There are queued messages restart the loop.

720+ // There are queued messages restart the loop. The recursive Run

721+ // publishes its own RunComplete for the queued prompt, so suppress

722+ // the outer defer's emit to avoid a duplicate event whose Error

723+ // field would belong to the recursive turn but whose MessageID/Text

724+ // would belong to the outer turn.

725+ skipRunComplete = true

638 firstQueuedMessage := queuedMessages[0]

726 firstQueuedMessage := queuedMessages[0]

639 a.messageQueue.Set(call.SessionID, queuedMessages[1:])

727 a.messageQueue.Set(call.SessionID, queuedMessages[1:])

640 return a.Run(ctx, firstQueuedMessage)

728 return a.Run(ctx, firstQueuedMessage)

internal/agent/coordinator.go+34−1

@@ -97,6 +97,7 @@ type coordinator struct {

97 filetracker filetracker.Service

98 lspManager *lsp.Manager

99 notify pubsub.Publisher[notify.Notification]

100+ runComplete pubsub.Publisher[notify.RunComplete]

100

101

101 currentAgent SessionAgent

102 currentAgent SessionAgent

102 agents map[string]SessionAgent

103 agents map[string]SessionAgent

@@ -119,6 +120,7 @@ func NewCoordinator(

119 filetracker filetracker.Service,

120 filetracker filetracker.Service,

120 lspManager *lsp.Manager,

121 lspManager *lsp.Manager,

121 notify pubsub.Publisher[notify.Notification],

122 notify pubsub.Publisher[notify.Notification],

123+ runComplete pubsub.Publisher[notify.RunComplete],

122 skillsMgr *skills.Manager,

124 skillsMgr *skills.Manager,

123) (Coordinator, error) {

125) (Coordinator, error) {

124 // Skills are pre-discovered by the caller (see app.New /

126 // Skills are pre-discovered by the caller (see app.New /

@@ -143,6 +145,7 @@ func NewCoordinator(

143 filetracker: filetracker,

145 filetracker: filetracker,

144 lspManager: lspManager,

146 lspManager: lspManager,

145 notify: notify,

147 notify: notify,

148+ runComplete: runComplete,

146 agents: make(map[string]SessionAgent),

149 agents: make(map[string]SessionAgent),

147 allSkills: allSkills,

150 allSkills: allSkills,

148 activeSkills: activeSkills,

151 activeSkills: activeSkills,

@@ -210,9 +213,34 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,

210 slog.Error("Failed to refresh OAuth2 token. Proceeding with existing token.", "error", err)

213 slog.Error("Failed to refresh OAuth2 token. Proceeding with existing token.", "error", err)

211 }

214 }

212

215

216+ // Coalesce per-attempt RunComplete payloads so only the final

217+ // outcome reaches subscribers. Without this, the first attempt's

218+ // failed RunComplete (unauthorized) would race ahead of the

219+ // retry's success, and `crush run` would exit on the stale error

220+ // before ever seeing the retry result. Each attempt's

221+ // SessionAgentCall.OnComplete hook overwrites latest; we publish

222+ // exactly once after retries resolve, via PublishMustDeliver, so

223+ // a momentarily-full subscriber buffer can't silently drop the

224+ // terminal event.

225+ var (

226+ latest notify.RunComplete

227+ hasLatest bool

228+ )

229+ onComplete := func(rc notify.RunComplete) {

230+ latest = rc

231+ hasLatest = true

232+ }

233+ // Propagate the caller-supplied RunID (set via agent.WithRunID

234+ // at the HTTP boundary in backend.SendMessage) onto the

235+ // SessionAgentCall so the terminal RunComplete event echoes it

236+ // back. Both attempts in the retry chain reuse the same RunID;

237+ // the coalesce closure publishes the final outcome under that

238+ // same correlator.

239+ runID := RunIDFromContext(ctx)

213 run := func() (*fantasy.AgentResult, error) {

240 run := func() (*fantasy.AgentResult, error) {

214 return c.currentAgent.Run(ctx, SessionAgentCall{

241 return c.currentAgent.Run(ctx, SessionAgentCall{

215 SessionID: sessionID,

242 SessionID: sessionID,

243+ RunID: runID,

216 Prompt: prompt,

244 Prompt: prompt,

217 Attachments: attachments,

245 Attachments: attachments,

218 MaxOutputTokens: maxTokens,

246 MaxOutputTokens: maxTokens,

@@ -222,6 +250,7 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,

222 TopK: topK,

250 TopK: topK,

223 FrequencyPenalty: freqPenalty,

251 FrequencyPenalty: freqPenalty,

224 PresencePenalty: presPenalty,

252 PresencePenalty: presPenalty,

253+ OnComplete: onComplete,

225 })

254 })

226 }

255 }

227 beforeLoaded := c.skillTracker.LoadedNames()

256 beforeLoaded := c.skillTracker.LoadedNames()

@@ -230,10 +259,13 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,

230

259

231 if c.isUnauthorized(originalErr) {

260 if c.isUnauthorized(originalErr) {

232 if err := c.retryAfterUnauthorized(ctx, providerCfg); err == nil {

261 if err := c.retryAfterUnauthorized(ctx, providerCfg); err == nil {

233− return run()

262+ result, originalErr = run()

234 }

263 }

235 }

264 }

236

265

266+ if hasLatest && c.runComplete != nil {

267+ c.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, latest)

268+ }

237 return result, originalErr

269 return result, originalErr

238}

270}

239

271

@@ -452,6 +484,7 @@ func (c *coordinator) buildAgent(ctx context.Context, prompt *prompt.Prompt, age

452 Messages: c.messages,

484 Messages: c.messages,

453 Tools: nil,

485 Tools: nil,

454 Notify: c.notify,

486 Notify: c.notify,

487+ RunComplete: c.runComplete,

455 })

488 })

456

489

457 c.readyWg.Go(func() error {

490 c.readyWg.Go(func() error {

internal/agent/notify/notify.go+28

@@ -21,3 +21,31 @@ type Notification struct {

21 Type Type

22 ProviderID string

23}

24+

25+// RunComplete is the authoritative end-of-run signal for a session.

26+// It is published exactly once per top-level agent run (per

27+// [sessionAgent.Run] invocation that actually executed) after all

28+// message updates for the turn have been flushed via

29+// message.Service.FlushAll. Carries the final assistant text and

30+// message ID so non-interactive clients can reconcile stdout even if

31+// SSE events arrive out of order or are dropped by the broker. Error

32+// is non-empty when the run terminated with an error; Cancelled is

33+// true when the run terminated due to context cancellation. The two

34+// are mutually exclusive in the success case but may overlap when a

35+// cancel triggers a downstream error.

36+//

37+// RunID identifies the specific request that produced this event.

38+// It is the value the caller set on `proto.AgentMessage.RunID` (or

39+// equivalently propagated via agent.WithRunID on the context that

40+// reaches the coordinator); empty when no caller set one. Filtering

41+// by RunID lets a client correlate a SendMessage call with its

42+// terminal event even when the session is busy and other turns are

43+// finishing on the same session.

44+type RunComplete struct {

45+ SessionID string

46+ RunID string

47+ MessageID string

48+ Text string

49+ Error string

50+ Cancelled bool

51+}

internal/agent/run_complete_test.gonew file+89

@@ -0,0 +1,89 @@

1+package agent

3+import (

4+ "context"

5+ "testing"

6+ "time"

8+ "github.com/charmbracelet/crush/internal/agent/notify"

9+ "github.com/charmbracelet/crush/internal/pubsub"

10+ "github.com/stretchr/testify/require"

11+)

12+

13+// TestSessionAgentRun_QueueStripsOnComplete verifies that when a Run

14+// call is enqueued (because the session is already busy), the

15+// OnComplete hook is NOT propagated onto the queued copy. The hook

16+// belongs to the caller's retry/coalesce scope (typically

17+// coordinator.Run) which has already returned by the time the queue

18+// drains; carrying it forward would silently funnel the terminal

19+// event into a closure nobody reads, and subscribers (`crush run`)

20+// would hang waiting for a RunComplete that never publishes.

21+func TestSessionAgentRun_QueueStripsOnComplete(t *testing.T) {

22+ t.Parallel()

23+

24+ env := testEnv(t)

25+ a := NewSessionAgent(SessionAgentOptions{

26+ Sessions: env.sessions,

27+ Messages: env.messages,

28+ }).(*sessionAgent)

29+

30+ const sessionID = "queued-session"

31+ // Mark the session as busy so Run takes the queue branch

32+ // without needing a real model.

33+ a.activeRequests.Set(sessionID, func() {})

34+

35+ var called bool

36+ hook := func(notify.RunComplete) { called = true }

37+

38+ res, err := a.Run(t.Context(), SessionAgentCall{

39+ SessionID: sessionID,

40+ RunID: "run-xyz",

41+ Prompt: "queued prompt",

42+ OnComplete: hook,

43+ })

44+ require.NoError(t, err)

45+ require.Nil(t, res, "queued Run must return (nil, nil)")

46+ require.False(t, called,

47+ "OnComplete must not fire on the enqueue path; the caller's scope is still live")

48+

49+ queued, ok := a.messageQueue.Get(sessionID)

50+ require.True(t, ok)

51+ require.Len(t, queued, 1)

52+ require.Nil(t, queued[0].OnComplete,

53+ "queued SessionAgentCall must have OnComplete stripped so the drain falls back to the default broker publish")

54+ require.Equal(t, "queued prompt", queued[0].Prompt,

55+ "all other fields must be preserved on the queued copy")

56+ require.Equal(t, "run-xyz", queued[0].RunID,

57+ "RunID must be preserved on the queued copy so the drained turn's "+

58+ "RunComplete still correlates with the originating SendMessage")

59+}

60+

61+// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the

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

63+// concrete Publisher implementation and must satisfy both Publish and

64+// PublishMustDeliver. The coordinator's final RunComplete emit relies

65+// on PublishMustDeliver to apply bounded-blocking semantics so a

66+// momentarily-full subscriber buffer can't silently drop the

67+// authoritative end-of-run event.

68+func TestRunCompletePublisher_MustDeliverOverTakesPublish(t *testing.T) {

69+ t.Parallel()

70+

71+ broker := pubsub.NewBroker[notify.RunComplete]()

72+ t.Cleanup(broker.Shutdown)

73+

74+ // Subscribe before publishing so the event is delivered.

75+ ctx, cancel := context.WithCancel(t.Context())

76+ defer cancel()

77+ ch := broker.Subscribe(ctx)

78+

79+ rc := notify.RunComplete{SessionID: "S", MessageID: "m", Text: "ok"}

80+ var pub pubsub.Publisher[notify.RunComplete] = broker

81+ pub.PublishMustDeliver(t.Context(), pubsub.UpdatedEvent, rc)

82+

83+ select {

84+ case got := <-ch:

85+ require.Equal(t, rc, got.Payload)

86+ case <-time.After(time.Second):

87+ t.Fatal("PublishMustDeliver did not deliver event")

88+ }

89+}

internal/agent/runid.gonew file+33

@@ -0,0 +1,33 @@

1+package agent

3+import "context"

5+// runIDContextKey is the unexported context key used to carry a

6+// caller-supplied RunID from the workspace HTTP boundary

7+// (backend.SendMessage) down into coordinator.Run without forcing a

8+// breaking change to the Coordinator.Run signature. The value is

9+// then copied onto SessionAgentCall.RunID by the coordinator so the

10+// agent's terminal RunComplete event can echo it back to the

11+// originating caller.

12+type runIDContextKey struct{}

13+

14+// WithRunID returns ctx tagged with a per-request RunID. It is the

15+// boundary helper for callers that need their SendMessage→Run

16+// terminal event to be uniquely correlatable (e.g. `crush run`

17+// against a session that may be busy). Empty runIDs are stored

18+// as-is; downstream code treats an empty RunID as "caller did not

19+// supply one" and falls back to SessionID-only correlation.

20+func WithRunID(ctx context.Context, runID string) context.Context {

21+ return context.WithValue(ctx, runIDContextKey{}, runID)

22+}

23+

24+// RunIDFromContext returns the RunID set by [WithRunID], or "" if

25+// none was set or the value is not a string. Exported because the

26+// coordinator and tests in other packages need to read it; safe to

27+// call on any context.

28+func RunIDFromContext(ctx context.Context) string {

29+ if v, ok := ctx.Value(runIDContextKey{}).(string); ok {

30+ return v

31+ }

32+ return ""

33+}

internal/app/app.go+42

@@ -75,6 +75,13 @@ type App struct {

75 globalCtx context.Context

76 cleanupFuncs []func(context.Context) error

77 agentNotifications *pubsub.Broker[notify.Notification]

78+ // runCompletions is the authoritative per-run completion signal,

79+ // emitted once per top-level agent turn after all message

80+ // updates have been flushed. Bridged into app.events so SSE

81+ // subscribers (notably `crush run` in client/server mode) can

82+ // drive their exit on a deterministic, payload-bearing event

83+ // instead of guessing from message finish parts.

84+ runCompletions *pubsub.Broker[notify.RunComplete]

78}

85}

80// New initializes a new application instance. skillsMgr carries the

87// New initializes a new application instance. skillsMgr carries the

@@ -110,6 +117,7 @@ func New(ctx context.Context, conn *sql.DB, store *config.ConfigStore, skillsMgr

110 serviceEventsWG: &sync.WaitGroup{},

117 serviceEventsWG: &sync.WaitGroup{},

111 tuiWG: &sync.WaitGroup{},

118 tuiWG: &sync.WaitGroup{},

112 agentNotifications: pubsub.NewBroker[notify.Notification](),

119 agentNotifications: pubsub.NewBroker[notify.Notification](),

120+ runCompletions: pubsub.NewBroker[notify.RunComplete](),

113 }

121 }

114

122

115 app.setupEvents()

123 app.setupEvents()

@@ -485,6 +493,7 @@ func (app *App) setupEvents() {

485 setupSubscriber(ctx, app.serviceEventsWG, "permissions-notifications", app.Permissions.SubscribeNotifications, app.events)

493 setupSubscriber(ctx, app.serviceEventsWG, "permissions-notifications", app.Permissions.SubscribeNotifications, app.events)

486 setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)

494 setupSubscriber(ctx, app.serviceEventsWG, "history", app.History.Subscribe, app.events)

487 setupSubscriber(ctx, app.serviceEventsWG, "agent-notifications", app.agentNotifications.Subscribe, app.events)

495 setupSubscriber(ctx, app.serviceEventsWG, "agent-notifications", app.agentNotifications.Subscribe, app.events)

496+ setupSubscriberMustDeliver(ctx, app.serviceEventsWG, "run-completions", app.runCompletions.Subscribe, app.events)

488 setupSubscriber(ctx, app.serviceEventsWG, "mcp", mcp.SubscribeEvents, app.events)

497 setupSubscriber(ctx, app.serviceEventsWG, "mcp", mcp.SubscribeEvents, app.events)

489 setupSubscriber(ctx, app.serviceEventsWG, "lsp", SubscribeLSPEvents, app.events)

498 setupSubscriber(ctx, app.serviceEventsWG, "lsp", SubscribeLSPEvents, app.events)

490 if app.Skills != nil {

499 if app.Skills != nil {

@@ -524,6 +533,38 @@ func setupSubscriber[T any](

524 })

533 })

525}

534}

526

535

536+// setupSubscriberMustDeliver is the bounded-blocking fan-in variant of

537+// setupSubscriber: it re-publishes upstream events onto the shared

538+// app.events broker using PublishMustDeliver instead of Publish. Use

539+// this for terminal events that subscribers cannot tolerate losing —

540+// notably RunComplete, which is the authoritative end-of-run signal

541+// for `crush run`. A lossy fan-in here can drop the only terminal

542+// event and hang non-interactive clients waiting on it.

543+func setupSubscriberMustDeliver[T any](

544+ ctx context.Context,

545+ wg *sync.WaitGroup,

546+ name string,

547+ subscriber func(context.Context) <-chan pubsub.Event[T],

548+ broker *pubsub.Broker[tea.Msg],

549+) {

550+ wg.Go(func() {

551+ subCh := subscriber(ctx)

552+ for {

553+ select {

554+ case event, ok := <-subCh:

555+ if !ok {

556+ slog.Debug("Subscription channel closed", "name", name)

557+ return

558+ }

559+ broker.PublishMustDeliver(ctx, pubsub.UpdatedEvent, tea.Msg(event))

560+ case <-ctx.Done():

561+ slog.Debug("Subscription cancelled", "name", name)

562+ return

563+ }

564+ }

565+ })

566+}

567+

527func (app *App) InitCoderAgent(ctx context.Context) error {

568func (app *App) InitCoderAgent(ctx context.Context) error {

528 coderAgentCfg := app.config.Config().Agents[config.AgentCoder]

569 coderAgentCfg := app.config.Config().Agents[config.AgentCoder]

529 if coderAgentCfg.ID == "" {

570 if coderAgentCfg.ID == "" {

@@ -540,6 +581,7 @@ func (app *App) InitCoderAgent(ctx context.Context) error {

540 app.FileTracker,

581 app.FileTracker,

541 app.LSPManager,

582 app.LSPManager,

542 app.agentNotifications,

583 app.agentNotifications,

584+ app.runCompletions,

543 app.Skills,

585 app.Skills,

544 )

586 )

545 if err != nil {

587 if err != nil {

internal/app/testing.go+3

@@ -34,6 +34,7 @@ func NewForTest(ctx context.Context) *App {

34 serviceEventsWG: &sync.WaitGroup{},

35 tuiWG: &sync.WaitGroup{},

36 agentNotifications: pubsub.NewBroker[notify.Notification](),

37+ runCompletions: pubsub.NewBroker[notify.RunComplete](),

37 }

38 }

39 eventsCtx, cancel := context.WithCancel(ctx)

40 eventsCtx, cancel := context.WithCancel(ctx)

@@ -44,6 +45,8 @@ func NewForTest(ctx context.Context) *App {

44 app.Permissions.SubscribeNotifications, app.events)

45 app.Permissions.SubscribeNotifications, app.events)

45 setupSubscriber(eventsCtx, app.serviceEventsWG, "agent-notifications",

46 setupSubscriber(eventsCtx, app.serviceEventsWG, "agent-notifications",

46 app.agentNotifications.Subscribe, app.events)

47 app.agentNotifications.Subscribe, app.events)

48+ setupSubscriber(eventsCtx, app.serviceEventsWG, "run-completions",

49+ app.runCompletions.Subscribe, app.events)

47 app.cleanupFuncs = append(app.cleanupFuncs, func(context.Context) error {

50 app.cleanupFuncs = append(app.cleanupFuncs, func(context.Context) error {

48 cancel()

51 cancel()

49 app.serviceEventsWG.Wait()

52 app.serviceEventsWG.Wait()

internal/backend/agent.go+11

@@ -3,12 +3,20 @@ package backend

3import (

4 "context"

6+ "github.com/charmbracelet/crush/internal/agent"

6 "github.com/charmbracelet/crush/internal/config"

7 "github.com/charmbracelet/crush/internal/config"

7 "github.com/charmbracelet/crush/internal/proto"

8 "github.com/charmbracelet/crush/internal/proto"

10// SendMessage sends a prompt to the agent coordinator for the given

11// SendMessage sends a prompt to the agent coordinator for the given

11// workspace and session.

12// workspace and session.

13+//

14+// When msg.RunID is non-empty it is attached to the context via

15+// agent.WithRunID so the coordinator can stamp the resulting

16+// SessionAgentCall (and therefore the terminal notify.RunComplete

17+// event) with that correlator. This is the only way for the

18+// originating client to distinguish its own turn's RunComplete from

19+// any concurrent turn that finishes on the same session.

12func (b *Backend) SendMessage(ctx context.Context, workspaceID string, msg proto.AgentMessage) error {

20func (b *Backend) SendMessage(ctx context.Context, workspaceID string, msg proto.AgentMessage) error {

13 ws, err := b.GetWorkspace(workspaceID)

21 ws, err := b.GetWorkspace(workspaceID)

14 if err != nil {

22 if err != nil {

@@ -19,6 +27,9 @@ func (b *Backend) SendMessage(ctx context.Context, workspaceID string, msg proto

19 return ErrAgentNotInitialized

27 return ErrAgentNotInitialized

20 }

28 }

30+ if msg.RunID != "" {

31+ ctx = agent.WithRunID(ctx, msg.RunID)

32+ }

22 _, err = ws.AgentCoordinator.Run(ctx, msg.SessionID, msg.Prompt, proto.AttachmentsToMessage(msg.Attachments)...)

33 _, err = ws.AgentCoordinator.Run(ctx, msg.SessionID, msg.Prompt, proto.AttachmentsToMessage(msg.Attachments)...)

23 return err

34 return err

24}

35}

internal/client/proto.go+14−1

@@ -231,6 +231,12 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er

231 if !sendEvent(ctx, events, e) {

232 return

233 }

234+ case pubsub.PayloadTypeRunComplete:

235+ var e pubsub.Event[proto.RunComplete]

236+ _ = json.Unmarshal(p.Payload, &e)

237+ if !sendEvent(ctx, events, e) {

238+ return

239+ }

234 default:

240 default:

235 slog.Warn("Unknown event type", "type", p.Type)

241 slog.Warn("Unknown event type", "type", p.Type)

236 continue

242 continue

@@ -400,9 +406,16 @@ func (c *Client) UpdateAgent(ctx context.Context, id string) error {

400}

406}

401

407

402// SendMessage sends a message to the agent for a workspace.

408// SendMessage sends a message to the agent for a workspace.

409+//

410+// When runID is non-empty it is echoed back on the resulting

411+// proto.RunComplete event, giving the caller a unique correlator

412+// for completion detection. Pass "" when the caller does not need

413+// to distinguish its own turn's terminal event from any concurrent

414+// turn on the same session (e.g. interactive TUI usage).

403−func (c *Client) SendMessage(ctx context.Context, id string, sessionID, prompt string, attachments ...message.Attachment) error {

415+func (c *Client) SendMessage(ctx context.Context, id string, sessionID, runID, prompt string, attachments ...message.Attachment) error {

404 rsp, err := c.post(ctx, fmt.Sprintf("/workspaces/%s/agent", id), nil, jsonBody(proto.AgentMessage{

416 rsp, err := c.post(ctx, fmt.Sprintf("/workspaces/%s/agent", id), nil, jsonBody(proto.AgentMessage{

405 SessionID: sessionID,

417 SessionID: sessionID,

418+ RunID: runID,

406 Prompt: prompt,

419 Prompt: prompt,

407 Attachments: proto.AttachmentsFromMessage(attachments),

420 Attachments: proto.AttachmentsFromMessage(attachments),

408 }), http.Header{"Content-Type": []string{"application/json"}})

421 }), http.Header{"Content-Type": []string{"application/json"}})

internal/cmd/run.go+140−36

@@ -3,6 +3,7 @@ package cmd

3import (

4 "context"

5 "fmt"

6+ "io"

6 "log/slog"

7 "log/slog"

7 "os"

8 "os"

8 "os/signal"

9 "os/signal"

@@ -24,6 +25,7 @@ import (

24 "github.com/charmbracelet/x/ansi"

25 "github.com/charmbracelet/x/ansi"

25 "github.com/charmbracelet/x/exp/charmtone"

26 "github.com/charmbracelet/x/exp/charmtone"

26 "github.com/charmbracelet/x/term"

27 "github.com/charmbracelet/x/term"

28+ "github.com/google/uuid"

27 "github.com/spf13/cobra"

29 "github.com/spf13/cobra"

28)

30)

@@ -243,12 +245,22 @@ func runNonInteractive(

243 return fmt.Errorf("failed to subscribe to events: %w", err)

245 return fmt.Errorf("failed to subscribe to events: %w", err)

244 }

246 }

245

247

248+ // Mint a per-call RunID so we can correlate the terminal

249+ // RunComplete with *this* SendMessage even if the session was

250+ // busy and another turn finished first. Without it the stream

251+ // loop would exit on whichever RunComplete arrived first for

252+ // the same session and drop the queued prompt's output.

253+ runID := uuid.New().String()

246− if err := c.SendMessage(ctx, ws.ID, sess.ID, prompt); err != nil {

254+ if err := c.SendMessage(ctx, ws.ID, sess.ID, runID, prompt); err != nil {

247 return fmt.Errorf("failed to send message: %w", err)

255 return fmt.Errorf("failed to send message: %w", err)

248 }

256 }

249

257

258+ stream := &runStream{

259+ sessionID: sess.ID,

260+ runID: runID,

261+ out: os.Stdout,

250− messageReadBytes := make(map[string]int)

262+ read: make(map[string]int),

251− var printed bool

263+ }

252

264

253 defer func() {

265 defer func() {

254 if progress && stderrTTY {

266 if progress && stderrTTY {

@@ -269,49 +281,141 @@ func runNonInteractive(

269 return nil

281 return nil

270 }

282 }

271

283

272− switch e := ev.(type) {

273− case pubsub.Event[proto.Message]:

274− msg := e.Payload

275− if msg.SessionID != sess.ID || msg.Role != proto.Assistant || len(msg.Parts) == 0 {

276− continue

277− }

278− stopSpinner()

284+ done, err := stream.handle(ev, stopSpinner)

285+ if err != nil {

286+ return err

287+ }

288+ if done {

289+ return nil

290+ }

279

291

280− content := msg.Content().String()

292+ case <-ctx.Done():

281− readBytes := messageReadBytes[msg.ID]

293+ stopSpinner()

294+ return ctx.Err()

295+ }

296+ }

297+}

282

298

283− if len(content) < readBytes {

299+// runStream tracks the per-message stdout cursor and the

284− slog.Error("Non-interactive: message content shorter than read bytes",

300+// reconciliation state used by [runNonInteractive] to translate

285− "message_length", len(content), "read_bytes", readBytes)

301+// streaming SSE events into a final, complete stdout for `crush run`.

286− return fmt.Errorf("message content is shorter than read bytes: %d < %d", len(content), readBytes)

302+// It is split out so the state machine can be exercised in unit tests

303+// without spinning up the full server/client harness.

304+//

305+// runID, when non-empty, is the authoritative correlator for the

306+// terminal RunComplete event: the stream suppresses live message

307+// events and only exits on a RunComplete whose RunID matches, so a

308+// turn that finishes first on the same session (e.g. when our prompt

309+// was queued behind a busy session) cannot contaminate stdout or

310+// terminate us prematurely. When empty (older servers, tests that

311+// don't supply one) the stream falls back to SessionID-only matching

312+// and live message streaming, which is still correct for the

313+// single-turn case.

314+type runStream struct {

315+ sessionID string

316+ runID string

317+ out io.Writer

318+ read map[string]int

319+ printed bool

287− }

320+}

288

321

289− part := content[readBytes:]

322+// handle processes one SSE event. Returns done=true when the run

290− if readBytes == 0 {

323+// loop should exit (RunComplete observed); returns an error only

291− part = strings.TrimLeft(part, " \t")

324+// when the agent run failed (not on context cancel — that path is

325+// handled by the caller's select). stopSpinner is called on the

326+// first observable assistant output and on completion; passing nil

327+// is safe for tests.

328+func (s *runStream) handle(ev any, stopSpinner func()) (done bool, err error) {

329+ stop := func() {

330+ if stopSpinner != nil {

331+ stopSpinner()

332+ }

333+ }

334+ switch e := ev.(type) {

335+ case pubsub.Event[proto.Message]:

336+ msg := e.Payload

337+ if msg.SessionID != s.sessionID || msg.Role != proto.Assistant || len(msg.Parts) == 0 {

338+ return false, nil

339+ }

340+ if s.runID != "" {

341+ return false, nil

292− }

342+ }

293− if printed || strings.TrimSpace(part) != "" {

343+ stop()

294− printed = true

344+

295− fmt.Fprint(os.Stdout, part)

345+ content := msg.Content().String()

296− }

346+ readBytes := s.read[msg.ID]

297− messageReadBytes[msg.ID] = len(content)

347+ if len(content) < readBytes {

348+ slog.Error("Non-interactive: message content shorter than read bytes",

349+ "message_length", len(content), "read_bytes", readBytes)

350+ return false, fmt.Errorf("message content is shorter than read bytes: %d < %d", len(content), readBytes)

351+ }

298

352

299− if msg.IsFinished() {

353+ part := content[readBytes:]

300− return nil

354+ if readBytes == 0 {

355+ part = strings.TrimLeft(part, " \t")

356+ }

357+ if s.printed || strings.TrimSpace(part) != "" {

358+ s.printed = true

359+ fmt.Fprint(s.out, part)

360+ }

361+ s.read[msg.ID] = len(content)

362+ return false, nil

363+

364+ case pubsub.Event[proto.RunComplete]:

365+ // RunComplete is the authoritative end-of-run signal. We

366+ // exit on it instead of guessing from message finish parts,

367+ // which fire on every tool-call step too and were the

368+ // source of the regression where `crush run` exited

369+ // mid-turn on finish.reason == tool_use.

370+ //

371+ // Correlation:

372+ // - if we minted a RunID for this SendMessage, only the

373+ // event whose RunID matches is ours; any other turn

374+ // finishing first on the same session (busy-session

375+ // queue path) must be ignored.

376+ // - if we have no RunID (older server, tests), fall back

377+ // to SessionID matching.

378+ if s.runID != "" {

379+ if e.Payload.RunID != s.runID {

380+ return false, nil

381+ }

382+ } else if e.Payload.SessionID != s.sessionID {

383+ return false, nil

384+ }

385+ stop()

386+ if e.Payload.Error != "" && !e.Payload.Cancelled {

387+ return true, fmt.Errorf("agent run failed: %s", e.Payload.Error)

388+ }

389+ // Reconcile stdout against the authoritative final

390+ // assistant text carried in the event. The pubsub fan-in

391+ // does not serialize publishes across upstream brokers, so

392+ // the final message event may not have reached this loop

393+ // yet; the embedded Text field is the backstop that

394+ // guarantees the full final text always appears on stdout.

395+ if e.Payload.MessageID != "" {

396+ full := e.Payload.Text

397+ readBytes := s.read[e.Payload.MessageID]

398+ if readBytes < len(full) {

399+ tail := full[readBytes:]

400+ if readBytes == 0 {

401+ tail = strings.TrimLeft(tail, " \t")

301 }

402 }

302−

403+ if s.printed || strings.TrimSpace(tail) != "" {

303− case pubsub.Event[proto.AgentEvent]:

404+ s.printed = true

304− if e.Payload.Error != nil {

405+ fmt.Fprint(s.out, tail)

305− stopSpinner()

306− return fmt.Errorf("agent error: %w", e.Payload.Error)

307 }

406 }

308 }

407 }

408+ }

409+ return true, nil

309

410

310− case <-ctx.Done():

411+ case pubsub.Event[proto.AgentEvent]:

311− stopSpinner()

412+ if e.Payload.Error != nil {

312− return ctx.Err()

413+ stop()

414+ return true, fmt.Errorf("agent error: %w", e.Payload.Error)

313 }

415 }

416+ return false, nil

314 }

417 }

418+ return false, nil

315}

419}

316

420

317// waitForAgent polls GetAgentInfo until the agent is ready, with a

421// waitForAgent polls GetAgentInfo until the agent is ready, with a

internal/cmd/run_stream_test.gonew file+329

@@ -0,0 +1,329 @@

1+package cmd

3+import (

4+ "bytes"

5+ "testing"

6+ "time"

8+ "github.com/charmbracelet/crush/internal/proto"

9+ "github.com/charmbracelet/crush/internal/pubsub"

10+ "github.com/stretchr/testify/require"

11+)

12+

13+// TestRunStream_ToolUseDoesNotTerminate is the regression test for

14+// the original bug: a tool-call assistant message has a Finish part

15+// with reason=tool_use and used to terminate `crush run` early via

16+// the discarded `msg.IsFinished()` exit condition. With the new

17+// RunComplete-driven loop, tool_use finishes must keep the stream

18+// alive so the post-tool final text still reaches stdout.

19+func TestRunStream_ToolUseDoesNotTerminate(t *testing.T) {

20+ t.Parallel()

21+

22+ s := &runStream{sessionID: "S", out: &bytes.Buffer{}, read: map[string]int{}}

23+

24+ toolUse := proto.Message{

25+ ID: "m1",

26+ SessionID: "S",

27+ Role: proto.Assistant,

28+ Parts: []proto.ContentPart{

29+ proto.TextContent{Text: ""},

30+ proto.Finish{Reason: proto.FinishReasonToolUse, Time: time.Now().Unix()},

31+ },

32+ }

33+ done, err := s.handle(pubsub.Event[proto.Message]{Payload: toolUse}, nil)

34+ require.NoError(t, err)

35+ require.False(t, done, "tool_use finish must NOT terminate the run loop")

36+}

37+

38+// TestRunStream_RunCompleteExits verifies the happy path: streaming

39+// assistant text then RunComplete terminates with the full final

40+// text on stdout. Together with the tool_use test above this

41+// nails down the "tool use + final text" sequence that the original

42+// bug truncated.

43+func TestRunStream_RunCompleteExits(t *testing.T) {

44+ t.Parallel()

45+

46+ buf := &bytes.Buffer{}

47+ s := &runStream{sessionID: "S", out: buf, read: map[string]int{}}

48+

49+ // Tool-use step.

50+ done, err := s.handle(pubsub.Event[proto.Message]{Payload: proto.Message{

51+ ID: "m1", SessionID: "S", Role: proto.Assistant,

52+ Parts: []proto.ContentPart{

53+ proto.TextContent{Text: ""},

54+ proto.Finish{Reason: proto.FinishReasonToolUse},

55+ },

56+ }}, nil)

57+ require.NoError(t, err)

58+ require.False(t, done)

59+

60+ // Final assistant message stream.

61+ done, err = s.handle(pubsub.Event[proto.Message]{Payload: proto.Message{

62+ ID: "m2", SessionID: "S", Role: proto.Assistant,

63+ Parts: []proto.ContentPart{

64+ proto.TextContent{Text: "VERDICT: APPROVED"},

65+ proto.Finish{Reason: proto.FinishReasonEndTurn},

66+ },

67+ }}, nil)

68+ require.NoError(t, err)

69+ require.False(t, done, "message finish (even end_turn) must not exit; RunComplete is the only terminal signal")

70+

71+ // RunComplete.

72+ done, err = s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

73+ SessionID: "S",

74+ MessageID: "m2",

75+ Text: "VERDICT: APPROVED",

76+ }}, nil)

77+ require.NoError(t, err)

78+ require.True(t, done)

79+ require.Equal(t, "VERDICT: APPROVED", buf.String())

80+}

81+

82+// TestRunStream_ReconcilesOnOutOfOrderRunComplete is the worst-case

83+// ordering scenario: RunComplete reaches the client BEFORE any of

84+// the streaming assistant message events for the turn (the pubsub

85+// fan-in across upstream brokers does not preserve cross-broker

86+// ordering). The embedded Text field must rescue stdout so the

87+// caller still sees the complete final text.

88+func TestRunStream_ReconcilesOnOutOfOrderRunComplete(t *testing.T) {

89+ t.Parallel()

90+

91+ buf := &bytes.Buffer{}

92+ s := &runStream{sessionID: "S", out: buf, read: map[string]int{}}

93+

94+ done, err := s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

95+ SessionID: "S",

96+ MessageID: "m2",

97+ Text: "VERDICT: APPROVED",

98+ }}, nil)

99+ require.NoError(t, err)

100+ require.True(t, done)

101+ require.Equal(t, "VERDICT: APPROVED", buf.String(),

102+ "RunComplete must reconcile stdout when message events did not arrive in time")

103+}

104+

105+// TestRunStream_ReconcilesPartialStream covers the realistic case

106+// where some streaming output reached stdout before RunComplete

107+// arrived: the reconciliation pass must append only the unread tail,

108+// never duplicate the prefix.

109+func TestRunStream_ReconcilesPartialStream(t *testing.T) {

110+ t.Parallel()

111+

112+ buf := &bytes.Buffer{}

113+ s := &runStream{sessionID: "S", out: buf, read: map[string]int{}}

114+

115+ _, err := s.handle(pubsub.Event[proto.Message]{Payload: proto.Message{

116+ ID: "m2", SessionID: "S", Role: proto.Assistant,

117+ Parts: []proto.ContentPart{proto.TextContent{Text: "VERDICT: "}},

118+ }}, nil)

119+ require.NoError(t, err)

120+

121+ _, err = s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

122+ SessionID: "S",

123+ MessageID: "m2",

124+ Text: "VERDICT: APPROVED",

125+ }}, nil)

126+ require.NoError(t, err)

127+ require.Equal(t, "VERDICT: APPROVED", buf.String())

128+}

129+

130+// TestRunStream_IgnoresOtherSessions ensures multi-session

131+// subscribers (e.g. a TUI watching workspace events while `crush

132+// run` is in flight against the same workspace) do not cause

133+// premature exit on RunComplete for a different session.

134+func TestRunStream_IgnoresOtherSessions(t *testing.T) {

135+ t.Parallel()

136+

137+ s := &runStream{sessionID: "S", out: &bytes.Buffer{}, read: map[string]int{}}

138+ done, err := s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

139+ SessionID: "OTHER",

140+ MessageID: "x",

141+ Text: "noise",

142+ }}, nil)

143+ require.NoError(t, err)

144+ require.False(t, done)

145+}

146+

147+// TestRunStream_ErrorRunComplete surfaces a failing run as a

148+// non-nil error from `crush run` so shells and CI catch it via

149+// exit status.

150+func TestRunStream_ErrorRunComplete(t *testing.T) {

151+ t.Parallel()

152+

153+ s := &runStream{sessionID: "S", out: &bytes.Buffer{}, read: map[string]int{}}

154+ done, err := s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

155+ SessionID: "S",

156+ Error: "model temporarily unavailable",

157+ }}, nil)

158+ require.True(t, done)

159+ require.Error(t, err)

160+ require.Contains(t, err.Error(), "model temporarily unavailable")

161+}

162+

163+// TestRunStream_CancelledRunCompleteIsClean ensures a cancelled

164+// run (e.g. Ctrl+C while `crush run` waits) exits cleanly rather

165+// than reporting the cancellation as a failure.

166+func TestRunStream_CancelledRunCompleteIsClean(t *testing.T) {

167+ t.Parallel()

168+

169+ s := &runStream{sessionID: "S", out: &bytes.Buffer{}, read: map[string]int{}}

170+ done, err := s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

171+ SessionID: "S",

172+ Error: "context canceled",

173+ Cancelled: true,

174+ }}, nil)

175+ require.True(t, done)

176+ require.NoError(t, err)

177+}

178+

179+// TestRunStream_LeadingWhitespaceTrimmedOnce mirrors the

180+// pre-existing trim of leading whitespace on the first byte of

181+// stdout: the trim must happen exactly once even when stdout is

182+// first produced by the RunComplete reconciliation path rather

183+// than the live stream.

184+func TestRunStream_LeadingWhitespaceTrimmedOnce(t *testing.T) {

185+ t.Parallel()

186+

187+ buf := &bytes.Buffer{}

188+ s := &runStream{sessionID: "S", out: buf, read: map[string]int{}}

189+

190+ done, err := s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

191+ SessionID: "S",

192+ MessageID: "m2",

193+ Text: " \tactual output",

194+ }}, nil)

195+ require.NoError(t, err)

196+ require.True(t, done)

197+ require.Equal(t, "actual output", buf.String())

198+}

199+

200+// TestRunStream_StopSpinnerInvokedOnFirstOutput verifies the

201+// spinner is stopped exactly when meaningful output starts (either

202+// a streamed assistant message or the reconciliation tail). This

203+// matches the prior behaviour and prevents the spinner from

204+// painting over the final response on TTYs.

205+func TestRunStream_StopSpinnerInvokedOnFirstOutput(t *testing.T) {

206+ t.Parallel()

207+

208+ calls := 0

209+ stop := func() { calls++ }

210+ s := &runStream{sessionID: "S", out: &bytes.Buffer{}, read: map[string]int{}}

211+ _, _ = s.handle(pubsub.Event[proto.Message]{Payload: proto.Message{

212+ ID: "m1", SessionID: "S", Role: proto.Assistant,

213+ Parts: []proto.ContentPart{proto.TextContent{Text: "hi"}},

214+ }}, stop)

215+ require.GreaterOrEqual(t, calls, 1, "spinner must stop once stdout has content")

216+}

217+

218+// TestRunStream_RunIDFiltersForeignTurns covers the busy-session

219+// queue scenario: `crush run --continue` attaches to a session

220+// whose currently running turn finishes first, publishing its

221+// RunComplete on the same session ID. Without per-run correlation

222+// the stream would exit on that foreign event and drop our own

223+// queued turn's output. With RunID filtering the foreign event is

224+// ignored and only the matching RunComplete terminates the stream.

225+func TestRunStream_RunIDFiltersForeignTurns(t *testing.T) {

226+ t.Parallel()

227+

228+ const sessionID = "S"

229+ const myRun = "run-mine"

230+ const otherRun = "run-other"

231+

232+ buf := &bytes.Buffer{}

233+ s := &runStream{

234+ sessionID: sessionID,

235+ runID: myRun,

236+ out: buf,

237+ read: map[string]int{},

238+ }

239+

240+ // The busy session's existing turn emits more text before it finishes.

241+ done, err := s.handle(pubsub.Event[proto.Message]{Payload: proto.Message{

242+ ID: "other-msg",

243+ SessionID: sessionID,

244+ Role: proto.Assistant,

245+ Parts: []proto.ContentPart{proto.TextContent{Text: "noise from another turn"}},

246+ }}, nil)

247+ require.NoError(t, err)

248+ require.False(t, done,

249+ "foreign message on same session must not terminate our run")

250+ require.Empty(t, buf.String(),

251+ "foreign message on same session must not write to our stdout")

252+

253+ // The busy session's existing turn finishes first.

254+ done, err = s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

255+ SessionID: sessionID,

256+ RunID: otherRun,

257+ MessageID: "other-msg",

258+ Text: "noise from another turn",

259+ }}, nil)

260+ require.NoError(t, err)

261+ require.False(t, done,

262+ "foreign RunComplete on same session must not terminate our run")

263+ require.Empty(t, buf.String(),

264+ "foreign RunComplete must not write to our stdout")

265+

266+ // Our own queued turn eventually finishes.

267+ done, err = s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

268+ SessionID: sessionID,

269+ RunID: myRun,

270+ MessageID: "my-msg",

271+ Text: "OK",

272+ }}, nil)

273+ require.NoError(t, err)

274+ require.True(t, done, "matching RunID must terminate the stream")

275+ require.Equal(t, "OK", buf.String())

276+}

277+

278+func TestRunStream_RunIDSuppressesLiveMessagesAndPrintsRunComplete(t *testing.T) {

279+ t.Parallel()

280+

281+ buf := &bytes.Buffer{}

282+ s := &runStream{

283+ sessionID: "S",

284+ runID: "run-mine",

285+ out: buf,

286+ read: map[string]int{},

287+ }

288+

289+ done, err := s.handle(pubsub.Event[proto.Message]{Payload: proto.Message{

290+ ID: "my-msg",

291+ SessionID: "S",

292+ Role: proto.Assistant,

293+ Parts: []proto.ContentPart{proto.TextContent{Text: "streamed prefix"}},

294+ }}, nil)

295+ require.NoError(t, err)

296+ require.False(t, done)

297+ require.Empty(t, buf.String())

298+

299+ done, err = s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

300+ SessionID: "S",

301+ RunID: "run-mine",

302+ MessageID: "my-msg",

303+ Text: "streamed prefix final",

304+ }}, nil)

305+ require.NoError(t, err)

306+ require.True(t, done)

307+ require.Equal(t, "streamed prefix final", buf.String())

308+}

309+

310+// TestRunStream_NoRunIDFallsBackToSessionID preserves the older

311+// behaviour for callers (and tests) that don't supply a RunID:

312+// SessionID-only matching still terminates the stream on the

313+// session's RunComplete. This keeps the contract backwards

314+// compatible with servers that don't echo RunID and with the

315+// pre-existing TestRunStream_* assertions.

316+func TestRunStream_NoRunIDFallsBackToSessionID(t *testing.T) {

317+ t.Parallel()

318+

319+ buf := &bytes.Buffer{}

320+ s := &runStream{sessionID: "S", out: buf, read: map[string]int{}}

321+ done, err := s.handle(pubsub.Event[proto.RunComplete]{Payload: proto.RunComplete{

322+ SessionID: "S",

323+ MessageID: "m2",

324+ Text: "DONE",

325+ }}, nil)

326+ require.NoError(t, err)

327+ require.True(t, done)

328+ require.Equal(t, "DONE", buf.String())

329+}

internal/proto/proto.go+39

@@ -46,6 +46,31 @@ type CurrentSession struct {

46 SessionID string `json:"session_id"`

47}

49+// RunComplete is the authoritative end-of-run signal for a session,

50+// emitted exactly once per top-level agent turn after all message

51+// updates for the turn have flushed. Clients that need a reliable

52+// completion contract (notably `crush run` in client/server mode)

53+// should listen for this event filtered by RunID (preferred) — or

54+// by SessionID when no RunID was supplied — and use Text and

55+// MessageID to reconcile any output they have already streamed from

56+// earlier message events. Error is non-empty when the run terminated

57+// with an error; Cancelled is true when terminated due to context

58+// cancellation.

59+//

60+// RunID echoes the value the caller set on AgentMessage.RunID. It is

61+// the only safe correlator when the caller's prompt was queued

62+// behind a busy session: another turn's RunComplete for the same

63+// SessionID may arrive first, and filtering by SessionID alone

64+// would terminate the caller before its own turn ran.

65+type RunComplete struct {

66+ SessionID string `json:"session_id"`

67+ RunID string `json:"run_id,omitempty"`

68+ MessageID string `json:"message_id"`

69+ Text string `json:"text,omitempty"`

70+ Error string `json:"error,omitempty"`

71+ Cancelled bool `json:"cancelled,omitempty"`

72+}

73+

49// SkillInfo describes a visible skill exposed to a frontend.

74// SkillInfo describes a visible skill exposed to a frontend.

50type SkillInfo struct {

75type SkillInfo struct {

51 ID string `json:"id"`

76 ID string `json:"id"`

@@ -89,8 +114,22 @@ func (a AgentInfo) IsZero() bool {

89}

114}

115

91// AgentMessage represents a message sent to the agent.

116// AgentMessage represents a message sent to the agent.

117+//

118+// RunID, when non-empty, is echoed back on the [RunComplete] event

119+// emitted for the resulting turn. Callers that need to correlate a

120+// specific SendMessage with its terminal event (notably

121+// `crush run`, which may attach to a busy session whose currently

122+// running turn finishes first) should set it to a fresh unique

123+// value before the request. Server-side propagation flows through

124+// agent.WithRunID on the request context into the

125+// SessionAgentCall; it is preserved across the busy-session queue.

126+// When empty the resulting RunComplete carries an empty RunID and

127+// callers must fall back to SessionID-only filtering, which

128+// remains correct only when no other turns are in flight for the

129+// same session.

92type AgentMessage struct {

130type AgentMessage struct {

93 SessionID string `json:"session_id"`

131 SessionID string `json:"session_id"`

132+ RunID string `json:"run_id,omitempty"`

94 Prompt string `json:"prompt"`

133 Prompt string `json:"prompt"`

95 Attachments []Attachment `json:"attachments,omitempty"`

134 Attachments []Attachment `json:"attachments,omitempty"`

96}

135}

internal/pubsub/events.go+8

@@ -26,6 +26,7 @@ const (

26 PayloadTypeAgentEvent PayloadType = "agent_event"

27 PayloadTypeConfigChanged PayloadType = "config_changed"

28 PayloadTypeSkillsEvent PayloadType = "skills_event"

29+ PayloadTypeRunComplete PayloadType = "run_complete"

29)

30)

31// Payload wraps a discriminated JSON payload with a type tag.

32// Payload wraps a discriminated JSON payload with a type tag.

@@ -50,7 +51,14 @@ type (

50 }

51 }

52 // Publisher can publish events of type T.

53 // Publisher can publish events of type T.

54+ //

55+ // Publish is best-effort and lossy under back-pressure;

56+ // PublishMustDeliver applies the bounded-blocking semantics used

57+ // for terminal events that must reach subscribers (finish, tool

58+ // result, error, cancel, RunComplete). See [Broker.Publish] and

59+ // [Broker.PublishMustDeliver].

53 Publisher[T any] interface {

60 Publisher[T any] interface {

54 Publish(EventType, T)

61 Publish(EventType, T)

62+ PublishMustDeliver(context.Context, EventType, T)

55 }

63 }

56)

64)

internal/server/events.go+12

@@ -93,6 +93,18 @@ func wrapEvent(ev any) *pubsub.Payload {

93 Type: proto.AgentEventType(e.Payload.Type),

94 },

95 })

96+ case pubsub.Event[notify.RunComplete]:

97+ return envelope(pubsub.PayloadTypeRunComplete, pubsub.Event[proto.RunComplete]{

98+ Type: e.Type,

99+ Payload: proto.RunComplete{

100+ SessionID: e.Payload.SessionID,

101+ RunID: e.Payload.RunID,

102+ MessageID: e.Payload.MessageID,

103+ Text: e.Payload.Text,

104+ Error: e.Payload.Error,

105+ Cancelled: e.Payload.Cancelled,

106+ },

107+ })

96 case pubsub.Event[proto.ConfigChanged]:

108 case pubsub.Event[proto.ConfigChanged]:

97 return envelope(pubsub.PayloadTypeConfigChanged, e)

109 return envelope(pubsub.PayloadTypeConfigChanged, e)

98 case pubsub.Event[skills.Event]:

110 case pubsub.Event[skills.Event]:

internal/server/events_test.go+65

@@ -5,6 +5,7 @@ import (

5 "errors"

6 "testing"

8+ "github.com/charmbracelet/crush/internal/agent/notify"

8 "github.com/charmbracelet/crush/internal/message"

9 "github.com/charmbracelet/crush/internal/message"

9 "github.com/charmbracelet/crush/internal/proto"

10 "github.com/charmbracelet/crush/internal/proto"

10 "github.com/charmbracelet/crush/internal/pubsub"

11 "github.com/charmbracelet/crush/internal/pubsub"

@@ -83,3 +84,67 @@ func TestSkillsEventToProto_RoundTrip(t *testing.T) {

83 require.Equal(t, proto.SkillStateError, decoded.Payload.States[1].State)

84 require.Equal(t, proto.SkillStateError, decoded.Payload.States[1].State)

84 require.Equal(t, "bad frontmatter", decoded.Payload.States[1].Error)

85 require.Equal(t, "bad frontmatter", decoded.Payload.States[1].Error)

85}

86}

87+

88+// TestRunCompleteToProto_RoundTrip verifies that the authoritative

89+// per-run completion event survives the SSE envelope conversion with

90+// all reconciliation fields intact. SessionID, MessageID, and Text

91+// are what non-interactive clients (e.g. `crush run`) rely on to

92+// terminate the run loop and guarantee final text on stdout when

93+// message events arrive out of order.

94+func TestRunCompleteToProto_RoundTrip(t *testing.T) {

95+ t.Parallel()

96+

97+ src := pubsub.Event[notify.RunComplete]{

98+ Type: pubsub.UpdatedEvent,

99+ Payload: notify.RunComplete{

100+ SessionID: "S",

101+ RunID: "run-42",

102+ MessageID: "M",

103+ Text: "VERDICT: APPROVED",

104+ Error: "",

105+ Cancelled: false,

106+ },

107+ }

108+

109+ env := wrapEvent(src)

110+ require.NotNil(t, env)

111+ require.Equal(t, pubsub.PayloadTypeRunComplete, env.Type)

112+

113+ var decoded pubsub.Event[proto.RunComplete]

114+ require.NoError(t, json.Unmarshal(env.Payload, &decoded))

115+ require.Equal(t, pubsub.UpdatedEvent, decoded.Type)

116+ require.Equal(t, "S", decoded.Payload.SessionID)

117+ require.Equal(t, "run-42", decoded.Payload.RunID,

118+ "RunID must survive the SSE envelope so clients can correlate "+

119+ "this event with the SendMessage call that produced it")

120+ require.Equal(t, "M", decoded.Payload.MessageID)

121+ require.Equal(t, "VERDICT: APPROVED", decoded.Payload.Text)

122+ require.Empty(t, decoded.Payload.Error)

123+ require.False(t, decoded.Payload.Cancelled)

124+}

125+

126+// TestRunCompleteToProto_Error verifies that error- and cancel-shaped

127+// RunComplete events round-trip cleanly so clients can distinguish

128+// "agent failed" (returns non-zero from `crush run`) from "agent

129+// cancelled by user" (clean exit).

130+func TestRunCompleteToProto_Error(t *testing.T) {

131+ t.Parallel()

132+

133+ src := pubsub.Event[notify.RunComplete]{

134+ Type: pubsub.UpdatedEvent,

135+ Payload: notify.RunComplete{

136+ SessionID: "S",

137+ MessageID: "M",

138+ Text: "partial",

139+ Error: "context canceled",

140+ Cancelled: true,

141+ },

142+ }

143+

144+ env := wrapEvent(src)

145+ require.NotNil(t, env)

146+ var decoded pubsub.Event[proto.RunComplete]

147+ require.NoError(t, json.Unmarshal(env.Payload, &decoded))

148+ require.Equal(t, "context canceled", decoded.Payload.Error)

149+ require.True(t, decoded.Payload.Cancelled)

150+}

internal/workspace/client_workspace.go+23−1

@@ -179,7 +179,11 @@ func (w *ClientWorkspace) ListAllUserMessages(ctx context.Context) ([]message.Me

179// -- Agent --

180

181func (w *ClientWorkspace) AgentRun(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) error {

182+ // The interactive TUI does not consume notify.RunComplete for

183+ // completion detection (it observes message events directly),

184+ // so passing an empty RunID is correct here: it skips the

185+ // correlator stamping path without functional consequences.

182− return w.client.SendMessage(ctx, w.workspaceID(), sessionID, prompt, attachments...)

186+ return w.client.SendMessage(ctx, w.workspaceID(), sessionID, "", prompt, attachments...)

183}

187}

184

188

185func (w *ClientWorkspace) AgentCancel(sessionID string) {

189func (w *ClientWorkspace) AgentCancel(sessionID string) {

@@ -707,6 +711,24 @@ func (w *ClientWorkspace) translateEvent(ev any) tea.Msg {

707 Type: notify.Type(e.Payload.Type),

711 Type: notify.Type(e.Payload.Type),

708 },

712 },

709 }

713 }

714+ case pubsub.Event[proto.RunComplete]:

715+ // Translate the wire-level proto.RunComplete back into the

716+ // agent's domain notify.RunComplete. Without this case the

717+ // default branch below warns on every run completion in the

718+ // server-mode TUI, even though the TUI itself doesn't act

719+ // on RunComplete — converting silently keeps the workspace

720+ // event bridge symmetric with the server-side wrapEvent.

721+ return pubsub.Event[notify.RunComplete]{

722+ Type: e.Type,

723+ Payload: notify.RunComplete{

724+ SessionID: e.Payload.SessionID,

725+ RunID: e.Payload.RunID,

726+ MessageID: e.Payload.MessageID,

727+ Text: e.Payload.Text,

728+ Error: e.Payload.Error,

729+ Cancelled: e.Payload.Cancelled,

730+ },

731+ }

710 case pubsub.Event[proto.SkillsEvent]:

732 case pubsub.Event[proto.SkillsEvent]:

711 states := protoToSkillStates(e.Payload.States)

733 states := protoToSkillStates(e.Payload.States)

712 if w.skills != nil {

734 if w.skills != nil {