git.iko.soy/crush
on main help sign in
perf: batch streaming message updates

Group rapid streaming updates into one save and one notification per short window instead of one per token. Important updates like finishes, tool calls, and errors still go through immediately. Cuts database writes and UI redraws by orders of magnitude during long responses. Co-Authored-By: Charm Crush <crush@charm.land>

commit6938dedd6cde0378f428b2a90d1591f69dbce6eb authorChristian Rocha <christian@rocha.is> dateTue 12 May 2026, 10:29 (5 months ago) parent771117c
7 files changed+1189−32
compare with main
unifiedsplit
internal/agent/agent.go+14 internal/app/app.go+13−3 internal/backend/session.go+6 internal/message/message.go+325−16 internal/message/message_test.go+695 internal/pubsub/broker.go+130−13 internal/workspace/app_workspace.go+6
internal/agent/agent.go+14
@@ -245,6 +245,15 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
245245
246246 defer cancel()
247247 defer a.activeRequests.Del(call.SessionID)
248+ // Drain any debounced message updates before returning. message.Service
249+ // already flushes synchronously on terminal updates, but a defer here
250+ // guarantees the contract at every Run exit (success, error, panic
251+ // recovery upstream) without callers needing to know.
252+ defer func() {
253+ if flushErr := a.messages.FlushAll(ctx); flushErr != nil {
254+ slog.Error("Failed to flush pending message updates after run", "error", flushErr)
255+ }
256+ }()
248257
249258 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)
250259
@@ -653,6 +662,11 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
653662 a.activeRequests.Set(sessionID, cancel)
654663 defer a.activeRequests.Del(sessionID)
655664 defer cancel()
665+ defer func() {
666+ if flushErr := a.messages.FlushAll(ctx); flushErr != nil {
667+ slog.Error("Failed to flush pending message updates after summarize", "error", flushErr)
668+ }
669+ }()
656670
657671 agent := fantasy.NewAgent(largeModel.Model,
658672 fantasy.WithSystemPrompt(string(summaryPrompt)),
internal/app/app.go+13−3
@@ -582,13 +582,23 @@ func (app *App) Shutdown() {
582582 app.AgentCoordinator.CancelAll()
583583 }
584584
585− // Now run remaining cleanup tasks in parallel.
586− var wg sync.WaitGroup
587−
588585 // Shared shutdown context for all timeout-bounded cleanup.
589586 shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
590587 defer cancel()
591588
589+ // Drain any debounced message updates before the DB-close cleanup
590+ // runs in the parallel block below. message.Service buffers
591+ // streaming deltas (see internal/message/message.go) and we must
592+ // land them while the connection is still open.
593+ if app.Messages != nil {
594+ if err := app.Messages.FlushAll(shutdownCtx); err != nil {
595+ slog.Error("Failed to flush pending message updates on shutdown", "error", err)
596+ }
597+ }
598+
599+ // Now run remaining cleanup tasks in parallel.
600+ var wg sync.WaitGroup
601+
592602 // Send exit event
593603 wg.Go(func() {
594604 event.AppExited()
internal/backend/session.go+6
@@ -72,6 +72,12 @@ func (b *Backend) ListSessionMessages(ctx context.Context, workspaceID, sessionI
7272 return nil, err
7373 }
7474
75+ // Drain debounced updates so HTTP clients (and the TUI on session
76+ // switch) observe the latest in-memory state rather than racing the
77+ // debounce timer in message.Service.
78+ if err := ws.Messages.FlushAll(ctx); err != nil {
79+ return nil, err
80+ }
7581 return ws.Messages.List(ctx, sessionID)
7682}
7783
internal/message/message.go+325−16
@@ -5,6 +5,7 @@ import (
55 "database/sql"
66 "encoding/json"
77 "fmt"
8+ "sync"
89 "time"
910
1011 "github.com/charmbracelet/crush/internal/db"
@@ -12,6 +13,13 @@ import (
1213 "github.com/google/uuid"
1314)
1415
16+// defaultUpdateDebounce is the default debounce window for [Service.Update].
17+// Streaming deltas that arrive within the window are coalesced into a
18+// single SQL write and a single pubsub event. Terminal updates
19+// (finish/error/cancel/tool-call structural changes) bypass the
20+// debounce and flush synchronously.
21+const defaultUpdateDebounce = 33 * time.Millisecond
22+
1523type CreateMessageParams struct {
1624 Role MessageRole
1725 Parts []ContentPart
@@ -20,6 +28,21 @@ type CreateMessageParams struct {
2028 IsSummaryMessage bool
2129}
2230
31+// Service is the public interface to the message store.
32+//
33+// [Service.Update] is eventually consistent: it accepts new state into
34+// an in-memory buffer and writes it to SQLite plus publishes a
35+// [pubsub.UpdatedEvent] on the next debounce tick (default
36+// [defaultUpdateDebounce]) or on the next terminal-state update,
37+// whichever comes first. Terminal-state updates — those that finish
38+// the message, add or finish a tool call, or end a reasoning section —
39+// flush synchronously before [Service.Update] returns.
40+//
41+// Callers that need stronger ordering (e.g. tests, shutdown,
42+// session-switch reads) must use [Service.Flush] or [Service.FlushAll]
43+// before reading via [Service.Get] / [Service.List]. Without an
44+// explicit flush, a read can race the debounce timer and miss the
45+// most recent in-memory state.
2346type Service interface {
2447 pubsub.Subscriber[Message]
2548 Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error)
@@ -30,20 +53,87 @@ type Service interface {
3053 ListAllUserMessages(ctx context.Context) ([]Message, error)
3154 Delete(ctx context.Context, id string) error
3255 DeleteSessionMessages(ctx context.Context, sessionID string) error
56+
57+ // Flush synchronously drains any pending debounced state for the
58+ // given message ID, performs the SQL write, and publishes the
59+ // resulting [pubsub.UpdatedEvent]. Idempotent; cheap no-op if no
60+ // updates are pending. Use this before any read that must observe
61+ // the latest [Service.Update].
62+ Flush(ctx context.Context, id string) error
63+
64+ // FlushAll synchronously drains pending debounced state for every
65+ // message known to the service. Intended for shutdown and
66+ // session-switch paths.
67+ FlushAll(ctx context.Context) error
68+}
69+
70+// pendingState holds the in-memory coalescing buffer for a single
71+// message ID. All fields except where noted are guarded by
72+// service.mu. The flushing flag serializes concurrent flushers for
73+// the same ID so SQL writes never reorder.
74+type pendingState struct {
75+ // latest is the most recent [Message] passed to [Service.Update]
76+ // that has not yet been flushed.
77+ latest Message
78+
79+ // dirty is true when latest contains state that has not been
80+ // written to SQL since the last successful flush.
81+ dirty bool
82+
83+ // flushing is true while a goroutine is performing the SQL write
84+ // for this ID. New updates are still accepted (and re-mark dirty)
85+ // but other flushers must back off.
86+ flushing bool
87+
88+ // timer is the active debounce timer, or nil if no flush is
89+ // scheduled. Stopped and reset when a terminal update preempts
90+ // the debounce window.
91+ timer *time.Timer
92+
93+ // lastFlushed is the snapshot most recently written to SQL. Used
94+ // as the baseline for terminal-state detection.
95+ lastFlushed Message
96+
97+ // hasFlushed is false until the first successful write for this
98+ // ID; until then lastFlushed is the zero value and must not be
99+ // treated as a real prior state.
100+ hasFlushed bool
33101}
34102
35103type service struct {
36104 *pubsub.Broker[Message]
37− q db.Querier
105+ q db.Querier
106+ debounce time.Duration
107+
108+ mu sync.Mutex
109+ pending map[string]*pendingState
38110}
39111
40−func NewService(q db.Querier) Service {
41− return &service{
42− Broker: pubsub.NewBroker[Message](),
43− q: q,
112+// ServiceOption configures a [Service] at construction.
113+type ServiceOption func(*service)
114+
115+// WithDebounce overrides the debounce window for [Service.Update]. A
116+// zero or negative value disables debouncing entirely (every update
117+// flushes synchronously). Intended primarily for tests.
118+func WithDebounce(d time.Duration) ServiceOption {
119+ return func(s *service) {
120+ s.debounce = d
44121 }
45122}
46123
124+func NewService(q db.Querier, opts ...ServiceOption) Service {
125+ s := &service{
126+ Broker: pubsub.NewBroker[Message](),
127+ q: q,
128+ debounce: defaultUpdateDebounce,
129+ pending: make(map[string]*pendingState),
130+ }
131+ for _, opt := range opts {
132+ opt(s)
133+ }
134+ return s
135+}
136+
47137func (s *service) Delete(ctx context.Context, id string) error {
48138 message, err := s.Get(ctx, id)
49139 if err != nil {
@@ -53,6 +143,16 @@ func (s *service) Delete(ctx context.Context, id string) error {
53143 if err != nil {
54144 return err
55145 }
146+ // Drop any pending coalesced state for this ID. We never want to
147+ // flush back over a deleted row.
148+ s.mu.Lock()
149+ if p, ok := s.pending[id]; ok {
150+ if p.timer != nil {
151+ p.timer.Stop()
152+ }
153+ delete(s.pending, id)
154+ }
155+ s.mu.Unlock()
56156 // Clone the message before publishing to avoid race conditions with
57157 // concurrent modifications to the Parts slice.
58158 s.Publish(pubsub.DeletedEvent, message.Clone())
@@ -111,31 +211,240 @@ func (s *service) DeleteSessionMessages(ctx context.Context, sessionID string) e
111211 return nil
112212}
113213
114−func (s *service) Update(ctx context.Context, message Message) error {
115− parts, err := marshalParts(message.Parts)
214+// Update accepts a new state for a message and either flushes
215+// synchronously (terminal updates, debounce <= 0) or buffers it until
216+// the next debounce tick. See [Service] for the contract.
217+func (s *service) Update(ctx context.Context, msg Message) error {
218+ cloned := msg.Clone()
219+
220+ // Zero or negative debounce: flush every update synchronously. This
221+ // preserves the pre-coalescing behaviour for tests and any caller
222+ // that explicitly opted out via [WithDebounce].
223+ if s.debounce <= 0 {
224+ s.mu.Lock()
225+ p, ok := s.pending[msg.ID]
226+ if !ok {
227+ p = &pendingState{}
228+ s.pending[msg.ID] = p
229+ }
230+ p.latest = cloned
231+ p.dirty = true
232+ s.mu.Unlock()
233+ return s.flushOne(ctx, msg.ID, true)
234+ }
235+
236+ s.mu.Lock()
237+ p, ok := s.pending[msg.ID]
238+ if !ok {
239+ p = &pendingState{}
240+ s.pending[msg.ID] = p
241+ }
242+ p.latest = cloned
243+ p.dirty = true
244+
245+ var prev *Message
246+ if p.hasFlushed {
247+ prev = &p.lastFlushed
248+ }
249+ terminal := shouldFlushNow(prev, &cloned)
250+
251+ if terminal {
252+ if p.timer != nil {
253+ p.timer.Stop()
254+ p.timer = nil
255+ }
256+ s.mu.Unlock()
257+ return s.flushOne(ctx, msg.ID, true)
258+ }
259+
260+ // Debounce: schedule a single flush per pending state. If a flush
261+ // is already running we let it finish; the trailing dirty bit will
262+ // be picked up by the next Update or by Flush.
263+ if p.timer == nil && !p.flushing {
264+ id := msg.ID
265+ p.timer = time.AfterFunc(s.debounce, func() {
266+ // Detached from caller ctx so a cancelled stream context
267+ // does not strand the buffered write.
268+ _ = s.flushOne(context.Background(), id, false)
269+ })
270+ }
271+ s.mu.Unlock()
272+ return nil
273+}
274+
275+// Flush implements [Service.Flush].
276+func (s *service) Flush(ctx context.Context, id string) error {
277+ return s.flushOne(ctx, id, true)
278+}
279+
280+// FlushAll implements [Service.FlushAll]. It snapshots every ID with
281+// outstanding work — either dirty buffered state or a flush already in
282+// flight — then drains each one. Picking up in-flight IDs ensures
283+// FlushAll cannot return while a timer-fired write is still mid-SQL,
284+// which is what shutdown and session-switch callers rely on.
285+func (s *service) FlushAll(ctx context.Context) error {
286+ s.mu.Lock()
287+ ids := make([]string, 0, len(s.pending))
288+ for id, p := range s.pending {
289+ if p.dirty || p.flushing {
290+ ids = append(ids, id)
291+ }
292+ }
293+ s.mu.Unlock()
294+ var firstErr error
295+ for _, id := range ids {
296+ if err := s.flushOne(ctx, id, true); err != nil && firstErr == nil {
297+ firstErr = err
298+ }
299+ }
300+ return firstErr
301+}
302+
303+// flushOne drains a single message ID. When syncCaller is true the
304+// caller is willing to wait through a concurrent in-flight flush so
305+// that, on return, lastFlushed equals latest at the moment of return.
306+// When false (timer-fired path) we bail if another flusher is already
307+// running; that flusher will pick up the trailing dirty bit.
308+//
309+// Order matters: a sync caller must wait for any in-flight flush to
310+// drain even when the buffer is currently clean — that in-flight
311+// write has not yet updated the SQL row, so returning early would
312+// violate the contract that on success lastFlushed reflects the most
313+// recent state.
314+func (s *service) flushOne(ctx context.Context, id string, syncCaller bool) error {
315+ for {
316+ s.mu.Lock()
317+ p, ok := s.pending[id]
318+ if !ok {
319+ s.mu.Unlock()
320+ return nil
321+ }
322+ if p.flushing {
323+ if !syncCaller {
324+ s.mu.Unlock()
325+ return nil
326+ }
327+ s.mu.Unlock()
328+ // Brief yield; in-flight write should land in <1ms typical.
329+ time.Sleep(time.Millisecond)
330+ continue
331+ }
332+ if !p.dirty {
333+ s.mu.Unlock()
334+ return nil
335+ }
336+
337+ if p.timer != nil {
338+ p.timer.Stop()
339+ p.timer = nil
340+ }
341+ snap := p.latest
342+ // Decide whether this snapshot represents a terminal event
343+ // against the prior baseline. We must do this before resetting
344+ // dirty/flushing because shouldFlushNow looks at p.lastFlushed
345+ // (which is what was on disk before this write).
346+ var prev *Message
347+ if p.hasFlushed {
348+ prev = &p.lastFlushed
349+ }
350+ isTerminal := shouldFlushNow(prev, &snap)
351+ p.flushing = true
352+ p.dirty = false
353+ s.mu.Unlock()
354+
355+ err := s.write(ctx, snap)
356+
357+ s.mu.Lock()
358+ p.flushing = false
359+ if err == nil {
360+ p.lastFlushed = snap
361+ p.hasFlushed = true
362+ } else {
363+ // Restore dirty so the next caller retries.
364+ p.dirty = true
365+ }
366+ // If a delta arrived during the SQL write and we are a sync
367+ // caller, the user expects that delta to land too.
368+ wasDirty := p.dirty
369+ s.mu.Unlock()
370+
371+ if err != nil {
372+ return err
373+ }
374+
375+ // Terminal events — message finished, tool call added or
376+ // finished, reasoning ended — use the bounded must-deliver
377+ // path so they never get dropped under channel contention.
378+ if isTerminal {
379+ s.PublishMustDeliver(ctx, pubsub.UpdatedEvent, snap)
380+ } else {
381+ s.Publish(pubsub.UpdatedEvent, snap)
382+ }
383+
384+ if wasDirty && syncCaller {
385+ continue
386+ }
387+ return nil
388+ }
389+}
390+
391+// write performs the unguarded SQL write + UpdatedAt stamp. Caller
392+// owns publishing.
393+func (s *service) write(ctx context.Context, msg Message) error {
394+ parts, err := marshalParts(msg.Parts)
116395 if err != nil {
117396 return err
118397 }
119398 finishedAt := sql.NullInt64{}
120− if f := message.FinishPart(); f != nil {
399+ if f := msg.FinishPart(); f != nil {
121400 finishedAt.Int64 = f.Time
122401 finishedAt.Valid = true
123402 }
124− err = s.q.UpdateMessage(ctx, db.UpdateMessageParams{
125− ID: message.ID,
403+ if err := s.q.UpdateMessage(ctx, db.UpdateMessageParams{
404+ ID: msg.ID,
126405 Parts: string(parts),
127406 FinishedAt: finishedAt,
128− })
129− if err != nil {
407+ }); err != nil {
130408 return err
131409 }
132− message.UpdatedAt = time.Now().Unix()
133− // Clone the message before publishing to avoid race conditions with
134− // concurrent modifications to the Parts slice.
135− s.Publish(pubsub.UpdatedEvent, message.Clone())
136410 return nil
137411}
138412
413+// shouldFlushNow returns true when next represents a structural
414+// change that must not be silently coalesced: the message just
415+// finished, the tool-call set grew, a tool call transitioned to
416+// finished, or reasoning just finished. prev is the last-flushed
417+// snapshot (or nil if no write has landed yet).
418+func shouldFlushNow(prev, next *Message) bool {
419+ if next.IsFinished() {
420+ return true
421+ }
422+
423+ var prevCalls []ToolCall
424+ var prevReasoningFinishedAt int64
425+ if prev != nil {
426+ prevCalls = prev.ToolCalls()
427+ prevReasoningFinishedAt = prev.ReasoningContent().FinishedAt
428+ }
429+ nextCalls := next.ToolCalls()
430+ if len(nextCalls) != len(prevCalls) {
431+ return true
432+ }
433+ for i := range nextCalls {
434+ // Bounds-safe: lengths are equal here.
435+ if nextCalls[i].Finished != prevCalls[i].Finished {
436+ return true
437+ }
438+ // A tool call's input only matters once it has landed (Finished
439+ // flips true). Earlier deltas to Input are debounced with the
440+ // rest of the streaming state.
441+ }
442+ if next.ReasoningContent().FinishedAt > 0 && prevReasoningFinishedAt == 0 {
443+ return true
444+ }
445+ return false
446+}
447+
139448func (s *service) Get(ctx context.Context, id string) (Message, error) {
140449 dbMessage, err := s.q.GetMessage(ctx, id)
141450 if err != nil {
internal/message/message_test.gonew file+695
@@ -0,0 +1,695 @@
1+package message
2+
3+import (
4+ "context"
5+ "sync"
6+ "sync/atomic"
7+ "testing"
8+ "time"
9+
10+ "github.com/charmbracelet/crush/internal/db"
11+ "github.com/charmbracelet/crush/internal/pubsub"
12+ "github.com/charmbracelet/crush/internal/session"
13+ "github.com/stretchr/testify/require"
14+)
15+
16+// slowUpdateQuerier wraps a [db.Querier] and forces UpdateMessage to
17+// hang on a release channel. Used to simulate an in-flight SQL write.
18+type slowUpdateQuerier struct {
19+ db.Querier
20+ release chan struct{}
21+ started chan struct{}
22+ startOnce sync.Once
23+}
24+
25+func (s *slowUpdateQuerier) UpdateMessage(ctx context.Context, arg db.UpdateMessageParams) error {
26+ s.startOnce.Do(func() { close(s.started) })
27+ select {
28+ case <-s.release:
29+ case <-ctx.Done():
30+ return ctx.Err()
31+ }
32+ return s.Querier.UpdateMessage(ctx, arg)
33+}
34+
35+// newTestService spins up a fresh in-memory message.Service backed by a
36+// temporary on-disk SQLite database. Returns the service plus a session
37+// ID to attach messages to.
38+func newTestService(t *testing.T, opts ...ServiceOption) (Service, string) {
39+ t.Helper()
40+ conn, err := db.Connect(t.Context(), t.TempDir())
41+ require.NoError(t, err)
42+ t.Cleanup(func() { _ = conn.Close() })
43+
44+ q := db.New(conn)
45+ sessions := session.NewService(q, conn)
46+ sess, err := sessions.Create(t.Context(), "test")
47+ require.NoError(t, err)
48+
49+ svc := NewService(q, opts...)
50+ return svc, sess.ID
51+}
52+
53+// eventCollector consumes broker events into a slice in a goroutine
54+// and exposes thread-safe Snapshot / Reset helpers for assertions.
55+type eventCollector struct {
56+ mu sync.Mutex
57+ events []pubsub.Event[Message]
58+}
59+
60+func collect(ctx context.Context, sub <-chan pubsub.Event[Message]) *eventCollector {
61+ c := &eventCollector{}
62+ go func() {
63+ for {
64+ select {
65+ case <-ctx.Done():
66+ return
67+ case ev, ok := <-sub:
68+ if !ok {
69+ return
70+ }
71+ c.mu.Lock()
72+ c.events = append(c.events, ev)
73+ c.mu.Unlock()
74+ }
75+ }
76+ }()
77+ return c
78+}
79+
80+func (c *eventCollector) snapshot() []pubsub.Event[Message] {
81+ c.mu.Lock()
82+ defer c.mu.Unlock()
83+ out := make([]pubsub.Event[Message], len(c.events))
84+ copy(out, c.events)
85+ return out
86+}
87+
88+func (c *eventCollector) reset() {
89+ c.mu.Lock()
90+ defer c.mu.Unlock()
91+ c.events = nil
92+}
93+
94+func TestUpdate_DebouncesTextDeltas(t *testing.T) {
95+ t.Parallel()
96+
97+ // Long-enough debounce that we can verify nothing flushes prematurely.
98+ svc, sessionID := newTestService(t, WithDebounce(50*time.Millisecond))
99+
100+ subCtx, cancelSub := context.WithCancel(t.Context())
101+ defer cancelSub()
102+ sub := svc.Subscribe(subCtx)
103+ collector := collect(subCtx, sub)
104+
105+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{
106+ Role: Assistant,
107+ })
108+ require.NoError(t, err)
109+ // Drop the CreatedEvent emitted by Create.
110+ time.Sleep(5 * time.Millisecond)
111+ collector.reset()
112+
113+ // Push 5 deltas inside a single debounce window.
114+ for i := 0; i < 5; i++ {
115+ msg.AppendContent("a")
116+ require.NoError(t, svc.Update(t.Context(), msg))
117+ }
118+
119+ // Before the debounce expires no UpdatedEvent should have landed.
120+ time.Sleep(10 * time.Millisecond)
121+ require.Empty(t, collector.snapshot(), "no events should land before debounce window expires")
122+
123+ // Wait for the debounce timer to fire.
124+ require.Eventually(t, func() bool {
125+ return len(collector.snapshot()) >= 1
126+ }, time.Second, 5*time.Millisecond)
127+ events := collector.snapshot()
128+ require.Len(t, events, 1, "5 deltas should coalesce into 1 UpdatedEvent")
129+ require.Equal(t, pubsub.UpdatedEvent, events[0].Type)
130+ require.Equal(t, "aaaaa", events[0].Payload.Content().Text)
131+
132+ // Final state must be persisted.
133+ got, err := svc.Get(t.Context(), msg.ID)
134+ require.NoError(t, err)
135+ require.Equal(t, "aaaaa", got.Content().Text)
136+}
137+
138+func TestUpdate_TerminalUpdatesFlushSynchronously(t *testing.T) {
139+ t.Parallel()
140+
141+ svc, sessionID := newTestService(t, WithDebounce(time.Hour))
142+
143+ subCtx, cancelSub := context.WithCancel(t.Context())
144+ defer cancelSub()
145+ sub := svc.Subscribe(subCtx)
146+ collector := collect(subCtx, sub)
147+
148+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
149+ require.NoError(t, err)
150+ time.Sleep(5 * time.Millisecond)
151+ collector.reset()
152+
153+ // AddFinish makes the message terminal; Update must flush
154+ // synchronously even with a 1-hour debounce.
155+ msg.AppendContent("done")
156+ msg.AddFinish(FinishReasonEndTurn, "", "")
157+ require.NoError(t, svc.Update(t.Context(), msg))
158+
159+ require.Eventually(t, func() bool {
160+ return len(collector.snapshot()) >= 1
161+ }, time.Second, 5*time.Millisecond,
162+ "terminal update must publish without waiting for debounce")
163+ events := collector.snapshot()
164+ require.Len(t, events, 1)
165+ require.True(t, events[0].Payload.IsFinished())
166+
167+ got, err := svc.Get(t.Context(), msg.ID)
168+ require.NoError(t, err)
169+ require.True(t, got.IsFinished())
170+}
171+
172+func TestUpdate_ToolCallStructuralChangeFlushes(t *testing.T) {
173+ t.Parallel()
174+
175+ svc, sessionID := newTestService(t, WithDebounce(time.Hour))
176+
177+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
178+ require.NoError(t, err)
179+
180+ // Adding a new tool call is a structural change → sync flush.
181+ msg.AddToolCall(ToolCall{ID: "tc1", Name: "view", Finished: false})
182+ require.NoError(t, svc.Update(t.Context(), msg))
183+
184+ got, err := svc.Get(t.Context(), msg.ID)
185+ require.NoError(t, err)
186+ require.Len(t, got.ToolCalls(), 1)
187+ require.Equal(t, "tc1", got.ToolCalls()[0].ID)
188+
189+ // Marking the tool call finished is also a structural change.
190+ msg.AddToolCall(ToolCall{ID: "tc1", Name: "view", Input: "{}", Finished: true})
191+ require.NoError(t, svc.Update(t.Context(), msg))
192+
193+ got, err = svc.Get(t.Context(), msg.ID)
194+ require.NoError(t, err)
195+ require.True(t, got.ToolCalls()[0].Finished)
196+}
197+
198+func TestUpdate_ReasoningEndFlushes(t *testing.T) {
199+ t.Parallel()
200+
201+ svc, sessionID := newTestService(t, WithDebounce(time.Hour))
202+
203+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
204+ require.NoError(t, err)
205+
206+ // Reasoning deltas alone debounce.
207+ msg.AppendReasoningContent("hmm")
208+ require.NoError(t, svc.Update(t.Context(), msg))
209+ got, err := svc.Get(t.Context(), msg.ID)
210+ require.NoError(t, err)
211+ require.Empty(t, got.ReasoningContent().Thinking, "reasoning delta should still be in the debounce buffer")
212+
213+ // FinishThinking sets FinishedAt → terminal flush.
214+ msg.FinishThinking()
215+ require.NoError(t, svc.Update(t.Context(), msg))
216+
217+ got, err = svc.Get(t.Context(), msg.ID)
218+ require.NoError(t, err)
219+ require.Equal(t, "hmm", got.ReasoningContent().Thinking)
220+ require.NotZero(t, got.ReasoningContent().FinishedAt)
221+}
222+
223+func TestFlush_DrainsPendingDebouncedUpdates(t *testing.T) {
224+ t.Parallel()
225+
226+ svc, sessionID := newTestService(t, WithDebounce(time.Hour))
227+
228+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
229+ require.NoError(t, err)
230+ msg.AppendContent("buffered")
231+ require.NoError(t, svc.Update(t.Context(), msg))
232+
233+ // Without a flush the SQL row is unchanged from Create.
234+ got, err := svc.Get(t.Context(), msg.ID)
235+ require.NoError(t, err)
236+ require.Empty(t, got.Content().Text)
237+
238+ require.NoError(t, svc.Flush(t.Context(), msg.ID))
239+
240+ got, err = svc.Get(t.Context(), msg.ID)
241+ require.NoError(t, err)
242+ require.Equal(t, "buffered", got.Content().Text)
243+
244+ // Subsequent Flush is a no-op.
245+ require.NoError(t, svc.Flush(t.Context(), msg.ID))
246+}
247+
248+func TestFlushAll_DrainsAllPending(t *testing.T) {
249+ t.Parallel()
250+
251+ svc, sessionID := newTestService(t, WithDebounce(time.Hour))
252+
253+ const n = 5
254+ msgs := make([]Message, n)
255+ for i := range msgs {
256+ m, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
257+ require.NoError(t, err)
258+ m.AppendContent("hi")
259+ require.NoError(t, svc.Update(t.Context(), m))
260+ msgs[i] = m
261+ }
262+
263+ require.NoError(t, svc.FlushAll(t.Context()))
264+
265+ for _, m := range msgs {
266+ got, err := svc.Get(t.Context(), m.ID)
267+ require.NoError(t, err)
268+ require.Equal(t, "hi", got.Content().Text, "FlushAll should drain every pending message")
269+ }
270+}
271+
272+func TestUpdate_OrderingMatchesNonCoalesced(t *testing.T) {
273+ t.Parallel()
274+
275+ // Compare the final state after coalesced vs zero-debounce updates.
276+ // A sequence of interleaved text/reasoning/tool-call updates must
277+ // converge to the same final DB row either way.
278+ build := func(svc Service, sessionID string) Message {
279+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
280+ require.NoError(t, err)
281+ msg.AppendReasoningContent("thinking 1 ")
282+ require.NoError(t, svc.Update(t.Context(), msg))
283+ msg.AppendReasoningContent("thinking 2")
284+ require.NoError(t, svc.Update(t.Context(), msg))
285+ msg.FinishThinking()
286+ require.NoError(t, svc.Update(t.Context(), msg))
287+ msg.AppendContent("hello ")
288+ require.NoError(t, svc.Update(t.Context(), msg))
289+ msg.AppendContent("world")
290+ require.NoError(t, svc.Update(t.Context(), msg))
291+ msg.AddToolCall(ToolCall{ID: "tc", Name: "x", Finished: false})
292+ require.NoError(t, svc.Update(t.Context(), msg))
293+ msg.AddToolCall(ToolCall{ID: "tc", Name: "x", Input: "{}", Finished: true})
294+ require.NoError(t, svc.Update(t.Context(), msg))
295+ msg.AddFinish(FinishReasonEndTurn, "", "")
296+ require.NoError(t, svc.Update(t.Context(), msg))
297+ return msg
298+ }
299+
300+ coalesced, sid1 := newTestService(t, WithDebounce(20*time.Millisecond))
301+ a := build(coalesced, sid1)
302+ require.NoError(t, coalesced.FlushAll(t.Context()))
303+ gotA, err := coalesced.Get(t.Context(), a.ID)
304+ require.NoError(t, err)
305+
306+ immediate, sid2 := newTestService(t, WithDebounce(0))
307+ b := build(immediate, sid2)
308+ gotB, err := immediate.Get(t.Context(), b.ID)
309+ require.NoError(t, err)
310+
311+ require.Equal(t, gotA.Content().Text, gotB.Content().Text)
312+ require.Equal(t, gotA.ReasoningContent().Thinking, gotB.ReasoningContent().Thinking)
313+ require.Equal(t, len(gotA.ToolCalls()), len(gotB.ToolCalls()))
314+ require.Equal(t, gotA.IsFinished(), gotB.IsFinished())
315+}
316+
317+func TestDelete_DropsPendingState(t *testing.T) {
318+ t.Parallel()
319+
320+ svc, sessionID := newTestService(t, WithDebounce(time.Hour))
321+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
322+ require.NoError(t, err)
323+ msg.AppendContent("dropped")
324+ require.NoError(t, svc.Update(t.Context(), msg))
325+
326+ require.NoError(t, svc.Delete(t.Context(), msg.ID))
327+
328+ // FlushAll after Delete must not write to the deleted row.
329+ require.NoError(t, svc.FlushAll(t.Context()))
330+
331+ _, err = svc.Get(t.Context(), msg.ID)
332+ require.Error(t, err, "deleted message must remain deleted")
333+}
334+
335+func TestBroker_PublishLossyDropCounter(t *testing.T) {
336+ t.Parallel()
337+
338+ // Tiny channel buffer so we can saturate from a single sender.
339+ b := pubsub.NewBrokerWithOptions[int](1, 1000)
340+ defer b.Shutdown()
341+
342+ subCtx, cancel := context.WithCancel(t.Context())
343+ defer cancel()
344+ sub := b.Subscribe(subCtx)
345+ require.NotNil(t, sub)
346+
347+ // Don't read from sub. Saturate the buffer.
348+ for range 100 {
349+ b.Publish(pubsub.UpdatedEvent, 1)
350+ }
351+
352+ require.GreaterOrEqual(t, b.DropCount(), uint64(1),
353+ "lossy Publish must increment the drop counter under contention")
354+}
355+
356+func TestBroker_PublishMustDeliverHonorsTimeout(t *testing.T) {
357+ t.Parallel()
358+
359+ b := pubsub.NewBrokerWithOptions[int](1, 1000)
360+ b.SetMustDeliverTimeout(20 * time.Millisecond)
361+ defer b.Shutdown()
362+
363+ subCtx, cancel := context.WithCancel(t.Context())
364+ defer cancel()
365+ sub := b.Subscribe(subCtx)
366+ require.NotNil(t, sub)
367+
368+ // Saturate: one event sits in the buffer, the second must wait.
369+ b.Publish(pubsub.UpdatedEvent, 1)
370+
371+ // PublishMustDeliver should block up to 20ms then drop.
372+ start := time.Now()
373+ b.PublishMustDeliver(t.Context(), pubsub.UpdatedEvent, 2)
374+ elapsed := time.Since(start)
375+
376+ require.GreaterOrEqual(t, elapsed, 20*time.Millisecond,
377+ "PublishMustDeliver should block at least the timeout under contention")
378+ require.Less(t, elapsed, 200*time.Millisecond,
379+ "PublishMustDeliver must not block indefinitely")
380+ require.GreaterOrEqual(t, b.MustDeliverDropCount(), uint64(1),
381+ "timeout must increment the must-deliver drop counter")
382+}
383+
384+func TestBroker_PublishMustDeliverWithReader(t *testing.T) {
385+ t.Parallel()
386+
387+ b := pubsub.NewBrokerWithOptions[int](1, 1000)
388+ b.SetMustDeliverTimeout(50 * time.Millisecond)
389+ defer b.Shutdown()
390+
391+ subCtx, cancel := context.WithCancel(t.Context())
392+ defer cancel()
393+ sub := b.Subscribe(subCtx)
394+
395+ var received atomic.Uint64
396+ done := make(chan struct{})
397+ go func() {
398+ defer close(done)
399+ for {
400+ select {
401+ case <-subCtx.Done():
402+ return
403+ case _, ok := <-sub:
404+ if !ok {
405+ return
406+ }
407+ received.Add(1)
408+ }
409+ }
410+ }()
411+
412+ for i := range 10 {
413+ b.PublishMustDeliver(t.Context(), pubsub.UpdatedEvent, i)
414+ }
415+
416+ // All 10 should land within the must-deliver timeout window.
417+ require.Eventually(t, func() bool { return received.Load() == 10 },
418+ time.Second, 5*time.Millisecond,
419+ "all must-deliver events should reach an active subscriber")
420+ require.Zero(t, b.MustDeliverDropCount(),
421+ "no drops expected when subscriber drains promptly")
422+}
423+
424+func TestUpdate_TerminalEventUsesMustDeliver(t *testing.T) {
425+ t.Parallel()
426+
427+ svc, sessionID := newTestService(t, WithDebounce(time.Hour))
428+
429+ subCtx, cancel := context.WithCancel(t.Context())
430+ defer cancel()
431+ sub := svc.Subscribe(subCtx)
432+
433+ var seenFinish atomic.Bool
434+ done := make(chan struct{})
435+ go func() {
436+ defer close(done)
437+ for {
438+ select {
439+ case <-subCtx.Done():
440+ return
441+ case ev, ok := <-sub:
442+ if !ok {
443+ return
444+ }
445+ if ev.Type == pubsub.UpdatedEvent && ev.Payload.IsFinished() {
446+ seenFinish.Store(true)
447+ }
448+ }
449+ }
450+ }()
451+
452+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
453+ require.NoError(t, err)
454+ msg.AppendContent("final")
455+ msg.AddFinish(FinishReasonEndTurn, "", "")
456+ require.NoError(t, svc.Update(t.Context(), msg))
457+
458+ require.Eventually(t, func() bool { return seenFinish.Load() },
459+ time.Second, 10*time.Millisecond,
460+ "terminal update must reach subscribers via the must-deliver path")
461+}
462+
463+func TestUpdate_ZeroDebounceFlushesEveryUpdate(t *testing.T) {
464+ t.Parallel()
465+
466+ svc, sessionID := newTestService(t, WithDebounce(0))
467+
468+ msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant})
469+ require.NoError(t, err)
470+
471+ for i := 0; i < 3; i++ {
472+ msg.AppendContent("x")
473+ require.NoError(t, svc.Update(t.Context(), msg))
474+ got, err := svc.Get(t.Context(), msg.ID)
475+ require.NoError(t, err)
476+ require.Len(t, got.Content().Text, i+1, "every update must land synchronously when debounce is 0")
477+ }
478+}
479+
480+// TestFlush_WaitsForInFlightWrite reproduces the failure where Flush
481+// or FlushAll could return before a concurrent in-flight SQL write
482+// completed. We block UpdateMessage on a release channel, fire the
483+// debounce timer, then call Flush and assert it does not return until
484+// the in-flight write actually lands.
485+func TestFlush_WaitsForInFlightWrite(t *testing.T) {
486+ t.Parallel()
487+
488+ conn, err := db.Connect(t.Context(), t.TempDir())
489+ require.NoError(t, err)
490+ t.Cleanup(func() { _ = conn.Close() })
491+
492+ q := db.New(conn)
493+ sessions := session.NewService(q, conn)
494+ sess, err := sessions.Create(t.Context(), "test")
495+ require.NoError(t, err)
496+
497+ slow := &slowUpdateQuerier{
498+ Querier: q,
499+ release: make(chan struct{}),
500+ started: make(chan struct{}),
501+ }
502+ // Short debounce so the timer fires quickly.
503+ svc := NewService(slow, WithDebounce(10*time.Millisecond))
504+
505+ msg, err := svc.Create(t.Context(), sess.ID, CreateMessageParams{Role: Assistant})
506+ require.NoError(t, err)
507+ msg.AppendContent("payload")
508+ require.NoError(t, svc.Update(t.Context(), msg))
509+
510+ // Wait for the timer-fired flush to enter UpdateMessage.
511+ select {
512+ case <-slow.started:
513+ case <-time.After(time.Second):
514+ t.Fatal("timer-fired flush never reached UpdateMessage")
515+ }
516+
517+ // At this point the buffer is dirty=false but flushing=true. A
518+ // naive Flush would early-return on !dirty. Spawn Flush in a
519+ // goroutine and assert it has not returned while the write is
520+ // still blocked.
521+ flushDone := make(chan error, 1)
522+ go func() { flushDone <- svc.Flush(t.Context(), msg.ID) }()
523+
524+ select {
525+ case err := <-flushDone:
526+ t.Fatalf("Flush returned %v while in-flight write was still blocked", err)
527+ case <-time.After(50 * time.Millisecond):
528+ // Expected: Flush is correctly waiting.
529+ }
530+
531+ // Release the slow write; Flush must now return cleanly.
532+ close(slow.release)
533+ select {
534+ case err := <-flushDone:
535+ require.NoError(t, err)
536+ case <-time.After(time.Second):
537+ t.Fatal("Flush did not return after in-flight write completed")
538+ }
539+
540+ // The SQL row should now reflect the buffered payload.
541+ got, err := svc.Get(t.Context(), msg.ID)
542+ require.NoError(t, err)
543+ require.Equal(t, "payload", got.Content().Text)
544+}
545+
546+// TestFlushAll_WaitsForInFlightWrite asserts FlushAll picks up IDs
547+// whose buffer is currently flushing (dirty=false) so shutdown and
548+// session-switch callers don't return while an SQL write is mid-flight.
549+func TestFlushAll_WaitsForInFlightWrite(t *testing.T) {
550+ t.Parallel()
551+
552+ conn, err := db.Connect(t.Context(), t.TempDir())
553+ require.NoError(t, err)
554+ t.Cleanup(func() { _ = conn.Close() })
555+
556+ q := db.New(conn)
557+ sessions := session.NewService(q, conn)
558+ sess, err := sessions.Create(t.Context(), "test")
559+ require.NoError(t, err)
560+
561+ slow := &slowUpdateQuerier{
562+ Querier: q,
563+ release: make(chan struct{}),
564+ started: make(chan struct{}),
565+ }
566+ svc := NewService(slow, WithDebounce(10*time.Millisecond))
567+
568+ msg, err := svc.Create(t.Context(), sess.ID, CreateMessageParams{Role: Assistant})
569+ require.NoError(t, err)
570+ msg.AppendContent("payload")
571+ require.NoError(t, svc.Update(t.Context(), msg))
572+
573+ select {
574+ case <-slow.started:
575+ case <-time.After(time.Second):
576+ t.Fatal("timer-fired flush never reached UpdateMessage")
577+ }
578+
579+ flushDone := make(chan error, 1)
580+ go func() { flushDone <- svc.FlushAll(t.Context()) }()
581+
582+ select {
583+ case err := <-flushDone:
584+ t.Fatalf("FlushAll returned %v while in-flight write was still blocked", err)
585+ case <-time.After(50 * time.Millisecond):
586+ }
587+
588+ close(slow.release)
589+ select {
590+ case err := <-flushDone:
591+ require.NoError(t, err)
592+ case <-time.After(time.Second):
593+ t.Fatal("FlushAll did not return after in-flight write completed")
594+ }
595+
596+ got, err := svc.Get(t.Context(), msg.ID)
597+ require.NoError(t, err)
598+ require.Equal(t, "payload", got.Content().Text)
599+}
600+
601+// TestUpdate_StructuralFlushUsesMustDeliver covers the second review
602+// finding: structural terminal events (tool-call add, tool-call
603+// finish, reasoning end) must publish via the must-deliver path even
604+// when the message itself is not yet IsFinished.
605+//
606+// We detect which path was taken by saturating a subscriber's channel
607+// buffer with no reader. With a short must-deliver timeout, the
608+// must-deliver path increments [pubsub.Broker.MustDeliverDropCount]
609+// after the timeout expires; the lossy path increments
610+// [pubsub.Broker.DropCount] immediately. The two counters are
611+// disjoint, so they precisely identify which call site published the
612+// event.
613+func TestUpdate_StructuralFlushUsesMustDeliver(t *testing.T) {
614+ t.Parallel()
615+
616+ cases := []struct {
617+ name string
618+ mut func(*Message)
619+ }{
620+ {
621+ name: "tool call add",
622+ mut: func(m *Message) {
623+ m.AddToolCall(ToolCall{ID: "tc1", Name: "view"})
624+ },
625+ },
626+ {
627+ name: "tool call finish",
628+ mut: func(m *Message) {
629+ m.AddToolCall(ToolCall{ID: "tc1", Name: "view", Input: "{}", Finished: true})
630+ },
631+ },
632+ {
633+ name: "reasoning end",
634+ mut: func(m *Message) {
635+ m.AppendReasoningContent("hmm")
636+ m.FinishThinking()
637+ },
638+ },
639+ }
640+
641+ for _, tc := range cases {
642+ t.Run(tc.name, func(t *testing.T) {
643+ t.Parallel()
644+
645+ conn, err := db.Connect(t.Context(), t.TempDir())
646+ require.NoError(t, err)
647+ t.Cleanup(func() { _ = conn.Close() })
648+
649+ q := db.New(conn)
650+ sessions := session.NewService(q, conn)
651+ sess, err := sessions.Create(t.Context(), "test")
652+ require.NoError(t, err)
653+
654+ // Replace the default broker with a tiny buffer + short
655+ // must-deliver timeout so we can fully saturate from a
656+ // single sender and observe drops without long waits.
657+ svc := NewService(q, WithDebounce(time.Hour))
658+ impl := svc.(*service)
659+ impl.Shutdown()
660+ impl.Broker = pubsub.NewBrokerWithOptions[Message](1, 1000)
661+ impl.SetMustDeliverTimeout(40 * time.Millisecond)
662+
663+ subCtx, cancel := context.WithCancel(t.Context())
664+ defer cancel()
665+ sub := svc.Subscribe(subCtx)
666+
667+ msg, err := svc.Create(t.Context(), sess.ID, CreateMessageParams{Role: Assistant})
668+ require.NoError(t, err)
669+
670+ // Saturate the subscriber's buffer (capacity 1). The
671+ // CreatedEvent from Create above already left one event
672+ // in the buffer; we never read sub, so the next publish
673+ // has nowhere to go.
674+ _ = sub // intentionally not drained.
675+
676+ // Drive the structural change. With debounce=1h, Update
677+ // flushes synchronously and routes through whichever
678+ // publish path the service chose for structural events.
679+ tc.mut(&msg)
680+ require.NoError(t, svc.Update(t.Context(), msg))
681+
682+ // Must-deliver timeout (40ms) should have expired with
683+ // no drain. If structural events are routed through
684+ // must-deliver: MustDeliverDropCount > 0, DropCount
685+ // unchanged. If routed through lossy Publish:
686+ // DropCount > 0, MustDeliverDropCount == 0.
687+ require.Eventually(t, func() bool {
688+ return impl.MustDeliverDropCount() >= 1
689+ }, time.Second, 5*time.Millisecond,
690+ "structural terminal event should publish via must-deliver, not lossy Publish")
691+ require.Zero(t, impl.DropCount(),
692+ "structural terminal event must not be silently dropped via lossy Publish")
693+ })
694+ }
695+}
internal/pubsub/broker.go+130−13
@@ -1,19 +1,54 @@
1+// Package pubsub provides a lightweight in-process broker for fan-out
2+// event delivery between services and the UI.
3+//
4+// Delivery semantics:
5+//
6+// - [Broker.Publish] is best-effort and lossy under contention. If a
7+// subscriber's channel is full, the event is dropped for that
8+// subscriber and a counter is incremented. This is the right choice
9+// for high-frequency intermediate updates (e.g. streaming token
10+// deltas) where only the latest state matters.
11+//
12+// - [Broker.PublishMustDeliver] is bounded-blocking. For each
13+// subscriber it first tries a non-blocking send, then falls back to
14+// a per-subscriber blocking send with a hard timeout. On timeout the
15+// event is dropped for that subscriber, an error is logged, and the
16+// must-deliver drop counter is incremented. The publisher never
17+// blocks indefinitely. This is the right choice for terminal events
18+// (finish, tool result, error, cancel) that must not be silently
19+// coalesced away.
20+//
21+// Drop counters ([Broker.DropCount], [Broker.MustDeliverDropCount]) are
22+// exposed so callers can surface saturation in telemetry.
123package pubsub
224
325import (
426 "context"
27+ "log/slog"
528 "sync"
29+ "sync/atomic"
30+ "time"
631)
732
8−const bufferSize = 64
33+const (
34+ bufferSize = 64
35+
36+ // defaultMustDeliverTimeout is the per-subscriber upper bound on how
37+ // long [Broker.PublishMustDeliver] will block waiting for buffer
38+ // space before giving up on that subscriber.
39+ defaultMustDeliverTimeout = 50 * time.Millisecond
40+)
941
1042type Broker[T any] struct {
11− subs map[chan Event[T]]struct{}
12− mu sync.RWMutex
13− done chan struct{}
14− subCount int
15− maxEvents int
16− channelBufferSize int
43+ subs map[chan Event[T]]struct{}
44+ mu sync.RWMutex
45+ done chan struct{}
46+ subCount int
47+ maxEvents int
48+ channelBufferSize int
49+ mustDeliverTimeout time.Duration
50+ dropCount atomic.Uint64
51+ mustDeliverDropCount atomic.Uint64
1752}
1853
1954func NewBroker[T any]() *Broker[T] {
@@ -22,13 +57,27 @@ func NewBroker[T any]() *Broker[T] {
2257
2358func NewBrokerWithOptions[T any](channelBufferSize, maxEvents int) *Broker[T] {
2459 return &Broker[T]{
25− subs: make(map[chan Event[T]]struct{}),
26− done: make(chan struct{}),
27− maxEvents: maxEvents,
28− channelBufferSize: channelBufferSize,
60+ subs: make(map[chan Event[T]]struct{}),
61+ done: make(chan struct{}),
62+ maxEvents: maxEvents,
63+ channelBufferSize: channelBufferSize,
64+ mustDeliverTimeout: defaultMustDeliverTimeout,
2965 }
3066}
3167
68+// SetMustDeliverTimeout overrides the per-subscriber timeout used by
69+// [Broker.PublishMustDeliver]. A zero or negative value resets to the
70+// default. Intended primarily for tests.
71+func (b *Broker[T]) SetMustDeliverTimeout(d time.Duration) {
72+ b.mu.Lock()
73+ defer b.mu.Unlock()
74+ if d <= 0 {
75+ b.mustDeliverTimeout = defaultMustDeliverTimeout
76+ return
77+ }
78+ b.mustDeliverTimeout = d
79+}
80+
3281func (b *Broker[T]) Shutdown() {
3382 select {
3483 case <-b.done: // Already closed
@@ -90,6 +139,25 @@ func (b *Broker[T]) GetSubscriberCount() int {
90139 return b.subCount
91140}
92141
142+// DropCount returns the cumulative number of events dropped by
143+// [Broker.Publish] because a subscriber's channel was full.
144+func (b *Broker[T]) DropCount() uint64 {
145+ return b.dropCount.Load()
146+}
147+
148+// MustDeliverDropCount returns the cumulative number of events dropped
149+// by [Broker.PublishMustDeliver] after the per-subscriber timeout
150+// expired.
151+func (b *Broker[T]) MustDeliverDropCount() uint64 {
152+ return b.mustDeliverDropCount.Load()
153+}
154+
155+// Publish delivers an event to every active subscriber.
156+//
157+// Delivery is non-blocking and lossy: if a subscriber's channel is full
158+// the event is dropped for that subscriber and [Broker.DropCount] is
159+// incremented. Use [Broker.PublishMustDeliver] for events that must not
160+// be silently dropped.
93161func (b *Broker[T]) Publish(t EventType, payload T) {
94162 b.mu.RLock()
95163 defer b.mu.RUnlock()
@@ -106,8 +174,57 @@ func (b *Broker[T]) Publish(t EventType, payload T) {
106174 select {
107175 case sub <- event:
108176 default:
109− // Channel is full, subscriber is slow - skip this event
110− // This prevents blocking the publisher
177+ // Channel is full, subscriber is slow - skip this event.
178+ // Lossy by design; counted so saturation is observable.
179+ b.dropCount.Add(1)
180+ }
181+ }
182+}
183+
184+// PublishMustDeliver delivers an event with bounded-blocking semantics.
185+// For each subscriber it first attempts a non-blocking send, then falls
186+// back to a blocking send bounded by a per-subscriber timeout (default
187+// [defaultMustDeliverTimeout]). On timeout the event is dropped for
188+// that subscriber, [Broker.MustDeliverDropCount] is incremented, and an
189+// error is logged. The publisher never blocks indefinitely.
190+//
191+// Use this for terminal events that must reach subscribers (finish,
192+// tool result, error, cancel). Callers must still tolerate rare drops
193+// after timeout — recovery is the subscriber's responsibility (e.g. a
194+// re-fetch on the next session-visible event).
195+func (b *Broker[T]) PublishMustDeliver(ctx context.Context, t EventType, payload T) {
196+ b.mu.RLock()
197+ defer b.mu.RUnlock()
198+
199+ select {
200+ case <-b.done:
201+ return
202+ default:
203+ }
204+
205+ event := Event[T]{Type: t, Payload: payload}
206+ timeout := b.mustDeliverTimeout
207+
208+ for sub := range b.subs {
209+ // Fast path: non-blocking send.
210+ select {
211+ case sub <- event:
212+ continue
213+ default:
214+ }
215+
216+ // Slow path: bounded blocking send.
217+ timer := time.NewTimer(timeout)
218+ select {
219+ case sub <- event:
220+ timer.Stop()
221+ case <-timer.C:
222+ b.mustDeliverDropCount.Add(1)
223+ slog.Error("PublishMustDeliver timed out delivering event",
224+ "type", t, "timeout", timeout)
225+ case <-ctx.Done():
226+ timer.Stop()
227+ return
111228 }
112229 }
113230}
internal/workspace/app_workspace.go+6
@@ -70,6 +70,12 @@ func (w *AppWorkspace) ParseAgentToolSessionID(sessionID string) (string, string
7070// -- Messages --
7171
7272func (w *AppWorkspace) ListMessages(ctx context.Context, sessionID string) ([]message.Message, error) {
73+ // Drain any debounced updates so the caller observes the latest
74+ // in-memory state. message.Service buffers streaming deltas and a
75+ // cold List would otherwise miss them at session-switch time.
76+ if err := w.app.Messages.FlushAll(ctx); err != nil {
77+ return nil, err
78+ }
7379 return w.app.Messages.List(ctx, sessionID)
7480}
7581