From 6ff541aa65c9c7f4d81e49ad15ff7e6ce11acd80 Mon Sep 17 00:00:00 2001 From: Luke Marsden Date: Fri, 3 Jul 2026 12:59:18 +0100 Subject: [PATCH] fix(api): add trailing-edge DB flush for streamed interaction writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming DB write is throttled to dbWriteInterval (5s) leading-edge with no trailing flush, so when an agent pauses mid-turn (e.g. before a tool call) the persisted interaction sits up to 5s behind the live stream. The frontend publish path already had a trailing flushTimer; the DB path did not. Any consumer reading the DB during that window — the useLiveInteraction 3s-poll fallback, a page-reload snapshot, or any other reader — saw stale/truncated text. Mirror the publish flushTimer with a dbFlushTimer (500ms): a throttled streaming update now schedules a trailing DB flush that persists the latest content shortly after the last chunk. Continuous streaming keeps resetting the timer, so writes still happen at the 5s leading cadence (bounding TOAST churn); only pauses and burst tails trigger a catch-up write. Extract the write into flushStreamingFieldsToDB (column-scoped, shared by both paths) and stop the timer at the existing teardown points. Verified live in the inner Helix: DB catch-up on a mid-turn pause drops from up to 5s to ~500ms; the live view was already current and is unchanged. Adds regression test TestMessageAdded_TrailingDBFlush. Co-Authored-By: Claude Opus 4.8 Spec-Ref: helix-specs@e9bbf4f02:002212_read-design2026-07-03 --- .../server/websocket_external_agent_sync.go | 107 ++++++++++++++---- .../websocket_external_agent_sync_test.go | 72 ++++++++++++ 2 files changed, 160 insertions(+), 19 deletions(-) diff --git a/api/pkg/server/websocket_external_agent_sync.go b/api/pkg/server/websocket_external_agent_sync.go index c3756fa188..3a43349da0 100644 --- a/api/pkg/server/websocket_external_agent_sync.go +++ b/api/pkg/server/websocket_external_agent_sync.go @@ -85,6 +85,11 @@ type streamingContext struct { // Trailing-edge flush timer: fires after publishInterval to drain any // patches that were skipped by the throttle when no new event arrived. flushTimer *time.Timer + // Trailing-edge DB flush timer: mirrors flushTimer for the throttled DB + // write. Fires after dbTrailingFlushInterval so the persisted interaction + // (read by the 3s poll fallback and page-reload snapshots) never sits + // more than ~dbTrailingFlushInterval behind the live stream during a pause. + dbFlushTimer *time.Timer // Per-entry delta tracking: tracks entries sent to frontend so we can compute per-entry diffs previousEntries []wsprotocol.ResponseEntry // Message accumulator: persists across handleMessageAdded calls so that @@ -116,6 +121,13 @@ const ( // publishInterval is the minimum time between frontend pubsub events during streaming. // Frontend batches to requestAnimationFrame (~16ms), so faster is wasted work. publishInterval = 50 * time.Millisecond + + // dbTrailingFlushInterval is the delay after the last streamed chunk before + // the trailing-edge DB flush fires. Kept small enough that a fallback/reload + // read is never badly stale, but large enough that continuous streaming + // (which keeps resetting the timer) still writes at the dbWriteInterval + // cadence rather than on every chunk — bounding TOAST churn. + dbTrailingFlushInterval = 500 * time.Millisecond ) // External agent WebSocket connections @@ -1333,28 +1345,41 @@ func (apiServer *HelixAPIServer) handleMessageAdded(sessionID string, syncMsg *t // serializing multi-MB JSON on every message. now := time.Now() if now.Sub(sctx.lastDBWrite) >= dbWriteInterval { - acc.Rebuild() - targetInteraction.ResponseMessage = acc.Content - targetInteraction.LastZedMessageOffset = acc.Offset - if entriesJSON, entErr := json.Marshal(acc.Entries()); entErr == nil { - _ = json.Unmarshal(entriesJSON, &targetInteraction.ResponseEntries) + // Leading-edge DB write; cancel any pending trailing flush since + // we're persisting the latest state right now. + if sctx.dbFlushTimer != nil { + sctx.dbFlushTimer.Stop() + sctx.dbFlushTimer = nil } - // Column-scoped write: never touch state/completed/error here, - // so that a concurrent handleTurnCancelled / handleMessageCompleted - // transition can't be clobbered by this in-flight streaming flush. - if err := apiServer.Controller.Options.Store.UpdateInteractionStreamingFields( - context.Background(), - targetInteraction.ID, - targetInteraction.GenerationID, - targetInteraction.ResponseMessage, - targetInteraction.ResponseEntries, - targetInteraction.LastZedMessageOffset, - targetInteraction.LastZedMessageID, - ); err != nil { + if err := apiServer.flushStreamingFieldsToDB(sctx); err != nil { return fmt.Errorf("failed to update interaction %s: %w", targetInteraction.ID, err) } - sctx.lastDBWrite = now - sctx.dirty = false + } else { + // Trailing-edge DB flush: mirror the frontend publish flushTimer + // so the persisted interaction (used by the 3s poll fallback and + // page-reload snapshots) is never more than dbTrailingFlushInterval + // behind the live stream during a pause. Without this the DB sits + // up to dbWriteInterval (5s) stale whenever the agent pauses + // mid-turn (e.g. before a tool call). Each new chunk resets the + // timer, so continuous streaming still writes at the dbWriteInterval + // cadence via the leading-edge branch above. + if sctx.dbFlushTimer != nil { + sctx.dbFlushTimer.Stop() + } + trailingInteractionID := targetInteraction.ID + sctx.dbFlushTimer = time.AfterFunc(dbTrailingFlushInterval, func() { + sctx.mu.Lock() + defer sctx.mu.Unlock() + sctx.dbFlushTimer = nil + if !sctx.dirty { + return + } + if err := apiServer.flushStreamingFieldsToDB(sctx); err != nil { + log.Error().Err(err). + Str("interaction_id", trailingInteractionID). + Msg("Failed to write interaction in trailing DB flush") + } + }) } log.Debug(). @@ -1624,6 +1649,10 @@ func (apiServer *HelixAPIServer) getOrCreateStreamingContext(ctx context.Context sctx.flushTimer.Stop() sctx.flushTimer = nil } + if sctx.dbFlushTimer != nil { + sctx.dbFlushTimer.Stop() + sctx.dbFlushTimer = nil + } } sctx.mu.Unlock() @@ -1857,6 +1886,10 @@ func (apiServer *HelixAPIServer) flushAndClearStreamingContext(ctx context.Conte sctx.flushTimer.Stop() sctx.flushTimer = nil } + if sctx.dbFlushTimer != nil { + sctx.dbFlushTimer.Stop() + sctx.dbFlushTimer = nil + } if sctx.interaction != nil { if sctx.dirty { @@ -4035,6 +4068,42 @@ func computePatch(previousContent, newContent string) (patchOffset int, patch st return utf16Off, newContent[byteOff:], totalLength } +// flushStreamingFieldsToDB rebuilds the accumulator content for the current +// interaction and persists the streaming columns (response_message, +// response_entries, offset, last message id). It is a column-scoped write: it +// never touches state/completed/error, so a concurrent +// handleTurnCancelled / handleMessageCompleted transition can't be clobbered by +// an in-flight streaming flush. The caller must hold sctx.mu. It is a no-op if +// there is no interaction or accumulator to flush. On success it updates +// lastDBWrite and clears the dirty flag. +func (apiServer *HelixAPIServer) flushStreamingFieldsToDB(sctx *streamingContext) error { + if sctx.accumulator == nil || sctx.interaction == nil { + return nil + } + acc := sctx.accumulator + it := sctx.interaction + acc.Rebuild() + it.ResponseMessage = acc.Content + it.LastZedMessageOffset = acc.Offset + if entriesJSON, err := json.Marshal(acc.Entries()); err == nil { + _ = json.Unmarshal(entriesJSON, &it.ResponseEntries) + } + if err := apiServer.Controller.Options.Store.UpdateInteractionStreamingFields( + context.Background(), + it.ID, + it.GenerationID, + it.ResponseMessage, + it.ResponseEntries, + it.LastZedMessageOffset, + it.LastZedMessageID, + ); err != nil { + return err + } + sctx.lastDBWrite = time.Now() + sctx.dirty = false + return nil +} + // publishEntryPatchesToFrontend sends per-entry delta patches for structured streaming. // Each entry gets its own string patch (offset/patch/length) so unchanged entries cost // zero bytes on the wire. The frontend maintains a ResponseEntry[] and applies patches diff --git a/api/pkg/server/websocket_external_agent_sync_test.go b/api/pkg/server/websocket_external_agent_sync_test.go index c3cac12a2f..f39affd9ea 100644 --- a/api/pkg/server/websocket_external_agent_sync_test.go +++ b/api/pkg/server/websocket_external_agent_sync_test.go @@ -453,6 +453,78 @@ func (s *WebSocketSyncSuite) TestMessageAdded_AssistantNewMessageID_MultiEntry() s.NoError(err) } +// TestMessageAdded_TrailingDBFlush verifies the trailing-edge DB flush: a +// streaming update that arrives within dbWriteInterval of the last DB write is +// throttled (no immediate write), but is still persisted shortly afterwards by +// the dbFlushTimer — rather than sitting up to dbWriteInterval stale in the DB. +// This closes the cause-#1 gap where a mid-turn pause left the persisted +// interaction (and thus the poll-fallback / reload snapshot) badly stale. +func (s *WebSocketSyncSuite) TestMessageAdded_TrailingDBFlush() { + s.server.contextMappings["thread-tf"] = "ses_tf" + + session := &types.Session{ID: "ses_tf", Owner: "user-1"} + existingInteraction := &types.Interaction{ + ID: "int-tf", + SessionID: "ses_tf", + State: types.InteractionStateWaiting, + ResponseMessage: "Hello", + LastZedMessageID: "msg-A", + } + + // Pre-seed a streaming context whose last DB write was just now, so the next + // message_added takes the throttled (else) branch and schedules a trailing + // flush instead of writing immediately. + s.server.streamingContexts["ses_tf"] = &streamingContext{ + session: session, + interaction: existingInteraction, + interactionID: "int-tf", + lastDBWrite: time.Now(), + lastPublish: time.Now(), + } + + // The trailing flush is the ONLY expected DB write; capture its content and + // signal when it lands. + written := make(chan string, 1) + s.store.EXPECT().UpdateInteractionStreamingFields(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, _ string, _ int, responseMessage string, _ datatypes.JSON, _ int, _ string) error { + written <- responseMessage + return nil + }, + ).Times(1) + + s.store.EXPECT().GetCommentByInteractionID(gomock.Any(), "int-tf"). + Return(nil, store.ErrNotFound).AnyTimes() + s.store.EXPECT().GetPendingCommentByPlanningSessionID(gomock.Any(), "ses_tf"). + Return(nil, nil).AnyTimes() + + syncMsg := &types.SyncMessage{ + EventType: "message_added", + Data: map[string]interface{}{ + "acp_thread_id": "thread-tf", + "message_id": "msg-A", + "content": "Hello, world!", + "role": "assistant", + }, + } + + s.NoError(s.server.handleMessageAdded("agent-1", syncMsg)) + + // No immediate write (throttled). + select { + case <-written: + s.Fail("DB write happened immediately; expected it to be throttled to the trailing flush") + case <-time.After(dbTrailingFlushInterval / 2): + } + + // Trailing flush fires shortly after and persists the latest content. + select { + case msg := <-written: + s.Equal("Hello, world!", msg) + case <-time.After(2 * time.Second): + s.Fail("trailing DB flush did not fire") + } +} + // TestMessageAdded_PriorInteractionMessageIDsAreFiltered reproduces the // cross-interaction response_entries leak surfaced by the e2e RESPONSE // ENTRIES ISOLATION VALIDATION step in Drone build #1024 (tag 2.11.0).