git.iko.soy/crush
on main help sign in
perf: batch streaming message updates
6938dedChristian Rocha
internal/agent/agent.go+14
compare with main
unifiedsplit
@@ -245,6 +245,15 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
245
245
246 defer cancel()
246 defer cancel()
247 defer a.activeRequests.Del(call.SessionID)
247 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+ }()
248
257
249 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)
258 history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)
250
259
@@ -653,6 +662,11 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
653 a.activeRequests.Set(sessionID, cancel)
662 a.activeRequests.Set(sessionID, cancel)
654 defer a.activeRequests.Del(sessionID)
663 defer a.activeRequests.Del(sessionID)
655 defer cancel()
664 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+ }()
656
670
657 agent := fantasy.NewAgent(largeModel.Model,
671 agent := fantasy.NewAgent(largeModel.Model,
658 fantasy.WithSystemPrompt(string(summaryPrompt)),
672 fantasy.WithSystemPrompt(string(summaryPrompt)),
internal/app/app.go+13−3
@@ -582,13 +582,23 @@ func (app *App) Shutdown() {
582 app.AgentCoordinator.CancelAll()
582 app.AgentCoordinator.CancelAll()
583 }
583 }
584
584
585− // Now run remaining cleanup tasks in parallel.
586− var wg sync.WaitGroup
587−
588 // Shared shutdown context for all timeout-bounded cleanup.
585 // Shared shutdown context for all timeout-bounded cleanup.
589 shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
586 shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
590 defer cancel()
587 defer cancel()
591
588
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+
592 // Send exit event
602 // Send exit event
593 wg.Go(func() {
603 wg.Go(func() {
594 event.AppExited()
604 event.AppExited()
internal/backend/session.go+6
@@ -72,6 +72,12 @@ func (b *Backend) ListSessionMessages(ctx context.Context, workspaceID, sessionI
72 return nil, err
72 return nil, err
73 }
73 }
74
74
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+ }
75 return ws.Messages.List(ctx, sessionID)
81 return ws.Messages.List(ctx, sessionID)
76}
82}
77
83
internal/message/message.go+325−16
@@ -5,6 +5,7 @@ import (
5 "database/sql"
5 "database/sql"
6 "encoding/json"
6 "encoding/json"
7 "fmt"
7 "fmt"
8+ "sync"
8 "time"
9 "time"
9
10
10 "github.com/charmbracelet/crush/internal/db"
11 "github.com/charmbracelet/crush/internal/db"
@@ -12,6 +13,13 @@ import (
12 "github.com/google/uuid"
13 "github.com/google/uuid"
13)
14)
14
15
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+
15type CreateMessageParams struct {
23type CreateMessageParams struct {
16 Role MessageRole
24 Role MessageRole
17 Parts []ContentPart
25 Parts []ContentPart
@@ -20,6 +28,21 @@ type CreateMessageParams struct {
20 IsSummaryMessage bool
28 IsSummaryMessage bool
21}
29}
22
30
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.
23type Service interface {
46type Service interface {
24 pubsub.Subscriber[Message]
47 pubsub.Subscriber[Message]
25 Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error)
48 Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error)
@@ -30,20 +53,87 @@ type Service interface {
30 ListAllUserMessages(ctx context.Context) ([]Message, error)
53 ListAllUserMessages(ctx context.Context) ([]Message, error)
31 Delete(ctx context.Context, id string) error
54 Delete(ctx context.Context, id string) error
32 DeleteSessionMessages(ctx context.Context, sessionID string) error
55 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
33}
101}
34
102
35type service struct {
103type service struct {
36 *pubsub.Broker[Message]
104 *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
38}
110}
39
111
40−func NewService(q db.Querier) Service {
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 {
41− return &service{
119+ return func(s *service) {
42− Broker: pubsub.NewBroker[Message](),
120+ s.debounce = d
43− q: q,
44 }
121 }
45}
122}
46
123
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+
47func (s *service) Delete(ctx context.Context, id string) error {
137func (s *service) Delete(ctx context.Context, id string) error {
48 message, err := s.Get(ctx, id)
138 message, err := s.Get(ctx, id)
49 if err != nil {
139 if err != nil {
@@ -53,6 +143,16 @@ func (s *service) Delete(ctx context.Context, id string) error {
53 if err != nil {
143 if err != nil {
54 return err
144 return err
55 }
145 }
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()
56 // Clone the message before publishing to avoid race conditions with
156 // Clone the message before publishing to avoid race conditions with
57 // concurrent modifications to the Parts slice.
157 // concurrent modifications to the Parts slice.
58 s.Publish(pubsub.DeletedEvent, message.Clone())
158 s.Publish(pubsub.DeletedEvent, message.Clone())
@@ -111,31 +211,240 @@ func (s *service) DeleteSessionMessages(ctx context.Context, sessionID string) e
111 return nil
211 return nil
112}
212}
113
213
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.
114−func (s *service) Update(ctx context.Context, message Message) error {
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 {
115− parts, err := marshalParts(message.Parts)
394+ parts, err := marshalParts(msg.Parts)
116 if err != nil {
395 if err != nil {
117 return err
396 return err
118 }
397 }
119 finishedAt := sql.NullInt64{}
398 finishedAt := sql.NullInt64{}
120− if f := message.FinishPart(); f != nil {
399+ if f := msg.FinishPart(); f != nil {
121 finishedAt.Int64 = f.Time
400 finishedAt.Int64 = f.Time
122 finishedAt.Valid = true
401 finishedAt.Valid = true
123 }
402 }
124− err = s.q.UpdateMessage(ctx, db.UpdateMessageParams{
403+ if err := s.q.UpdateMessage(ctx, db.UpdateMessageParams{
125− ID: message.ID,
404+ ID: msg.ID,
126 Parts: string(parts),
405 Parts: string(parts),
127 FinishedAt: finishedAt,
406 FinishedAt: finishedAt,
128− })
129− if err != nil {
407+ }); err != nil {
130 return err
408 return err
131 }
409 }
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())
136 return nil
410 return nil
137}
411}
138
412
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+
139func (s *service) Get(ctx context.Context, id string) (Message, error) {
448func (s *service) Get(ctx context.Context, id string) (Message, error) {
140 dbMessage, err := s.q.GetMessage(ctx, id)
449 dbMessage, err := s.q.GetMessage(ctx, id)
141 if err != nil {
450 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.
1package pubsub
23package pubsub
2
24
3import (
25import (
4 "context"
26 "context"
27+ "log/slog"
5 "sync"
28 "sync"
29+ "sync/atomic"
30+ "time"
6)
31)
7
32
33+const (
8−const bufferSize = 64
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+)
9
41
10type Broker[T any] struct {
42type Broker[T any] struct {
11− subs map[chan Event[T]]struct{}
43+ subs map[chan Event[T]]struct{}
12− mu sync.RWMutex
44+ mu sync.RWMutex
13− done chan struct{}
45+ done chan struct{}
14− subCount int
46+ subCount int
15− maxEvents int
47+ maxEvents int
16− channelBufferSize int
48+ channelBufferSize int
49+ mustDeliverTimeout time.Duration
50+ dropCount atomic.Uint64
51+ mustDeliverDropCount atomic.Uint64
17}
52}
18
53
19func NewBroker[T any]() *Broker[T] {
54func NewBroker[T any]() *Broker[T] {
@@ -22,13 +57,27 @@ func NewBroker[T any]() *Broker[T] {
22
57
23func NewBrokerWithOptions[T any](channelBufferSize, maxEvents int) *Broker[T] {
58func NewBrokerWithOptions[T any](channelBufferSize, maxEvents int) *Broker[T] {
24 return &Broker[T]{
59 return &Broker[T]{
25− subs: make(map[chan Event[T]]struct{}),
60+ subs: make(map[chan Event[T]]struct{}),
26− done: make(chan struct{}),
61+ done: make(chan struct{}),
27− maxEvents: maxEvents,
62+ maxEvents: maxEvents,
28− channelBufferSize: channelBufferSize,
63+ channelBufferSize: channelBufferSize,
64+ mustDeliverTimeout: defaultMustDeliverTimeout,
29 }
65 }
30}
66}
31
67
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+
32func (b *Broker[T]) Shutdown() {
81func (b *Broker[T]) Shutdown() {
33 select {
82 select {
34 case <-b.done: // Already closed
83 case <-b.done: // Already closed
@@ -90,6 +139,25 @@ func (b *Broker[T]) GetSubscriberCount() int {
90 return b.subCount
139 return b.subCount
91}
140}
92
141
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.
93func (b *Broker[T]) Publish(t EventType, payload T) {
161func (b *Broker[T]) Publish(t EventType, payload T) {
94 b.mu.RLock()
162 b.mu.RLock()
95 defer b.mu.RUnlock()
163 defer b.mu.RUnlock()
@@ -106,8 +174,57 @@ func (b *Broker[T]) Publish(t EventType, payload T) {
106 select {
174 select {
107 case sub <- event:
175 case sub <- event:
108 default:
176 default:
109− // Channel is full, subscriber is slow - skip this event
177+ // Channel is full, subscriber is slow - skip this event.
110− // This prevents blocking the publisher
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
111 }
228 }
112 }
229 }
113}
230}
internal/workspace/app_workspace.go+6
@@ -70,6 +70,12 @@ func (w *AppWorkspace) ParseAgentToolSessionID(sessionID string) (string, string
70// -- Messages --
70// -- Messages --
71
71
72func (w *AppWorkspace) ListMessages(ctx context.Context, sessionID string) ([]message.Message, error) {
72func (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+ }
73 return w.app.Messages.List(ctx, sessionID)
79 return w.app.Messages.List(ctx, sessionID)
74}
80}
75
81