From f246d3fd8aba7485baff3debc4bece01768dfd54 Mon Sep 17 00:00:00 2001 From: "vellum-apollo-bot[bot]" <221061047+vellum-apollo-bot[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 19:25:34 +0000 Subject: [PATCH 1/5] feat(agent-loop): partial-persist assistant content mid-turn (B6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the durability gap left by B3: a row reserved at `llm_call_started` stayed empty for the full duration of a turn, so a refresh mid-turn rendered an empty assistant bubble where the in-progress reply should have been. B6 mirrors streamed text + tool_use blocks into a per-row accumulator on `EventHandlerState` and flushes via `updateContent` on a first-fire-wins debounce: • Time gate: 250ms since last flush (PARTIAL_PERSIST_DEBOUNCE_MS) • Size gate: 1024 new bytes (PARTIAL_PERSIST_SIZE_THRESHOLD) `handleTextDelta` appends + schedules; `handleToolUse` flushes eagerly (tool_use blocks are atomic); `handleMessageComplete` clears the pending timer before the final authoritative flush, so a late debounce can't race the indexer/projector. Indexer + attention projector still fire ONLY at message_complete. Partial rows are never indexed — memory recall would otherwise surface mid-turn fragments. Secrets are redacted from text blocks on every flush, matching the discipline at message_complete. Failure modes: • Provider error: B3's orphan cleanup already deletes the row; its partial content goes with the orphan (new test pins this). • Retry path: B3's `assistantRowAwaitingFinalization` flag triggers cleanup before the new reservation; `handleLlmCallStarted` resets the accumulator for the fresh row. • Stranded partial flush failure: logged and swallowed — the next debounce tick or final flush at message_complete recovers. Thinking blocks are intentionally NOT mirrored mid-turn. The final `event.message.content` still carries them, so finalized rows remain authoritative. Spec: scratch/b6-partial-persist-spec.md Tests: 7 new (debounce time gate, size gate, eager flush, timer clear, no mid-turn indexing, redaction, provider-error cleanup). All 71 tests in conversation-agent-loop.test.ts pass. --- .../__tests__/conversation-agent-loop.test.ts | 466 ++++++++++++++++++ .../conversation-agent-loop-handlers.ts | 233 +++++++++ 2 files changed, 699 insertions(+) diff --git a/assistant/src/__tests__/conversation-agent-loop.test.ts b/assistant/src/__tests__/conversation-agent-loop.test.ts index b5cff4feae6..080130f70eb 100644 --- a/assistant/src/__tests__/conversation-agent-loop.test.ts +++ b/assistant/src/__tests__/conversation-agent-loop.test.ts @@ -3444,6 +3444,472 @@ describe("session-agent-loop", () => { }); }); + describe("B6 partial persistence", () => { + // The pre-B6 flow reserves an empty assistant row at `llm_call_started` + // (`content: "[]"`) and never touches that row again until + // `handleMessageComplete` fires the single authoritative + // `updateContent`. Between those events the row is empty for the full + // duration of a turn — a browser refresh mid-turn sees nothing where + // the in-progress assistant reply should be. + // + // B6 closes that durability gap with a debounced partial flush from + // `handleTextDelta` (250ms time gate / 1024-char size gate, first-fire + // wins) and an eager flush from `handleToolUse` (tool_use blocks are + // atomic, no debounce needed). The indexer + projector still fire + // ONLY at `message_complete` — partial rows are never indexed. + // + // These tests pin down the wire-level contract by counting + // `updateMessageContent` calls and inspecting the JSON payload of the + // partial-flush writes. The indexing / sync-invalidation paths are + // covered by the B3 block above. + + test("debounced time gate flushes one partial write after PARTIAL_PERSIST_DEBOUNCE_MS", async () => { + mockMessageById = { + id: "msg-reserve", + conversationId: "test-conv", + createdAt: 1234567, + role: "assistant", + content: "[]", + metadata: null, + }; + + const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { + await onEvent({ type: "llm_call_started" }); + // Two small deltas — well under the 1024-char size gate — should + // schedule a single debounced flush. + onEvent({ type: "text_delta", text: "Hello, " }); + onEvent({ type: "text_delta", text: "world." }); + // Wait long enough for the 250ms debounce to fire. + await new Promise((resolve) => setTimeout(resolve, 350)); + await onEvent({ + type: "message_complete", + message: { + role: "assistant", + content: [{ type: "text", text: "Hello, world." }], + }, + }); + onEvent({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + model: "test", + providerDurationMs: 50, + }); + return [ + ...messages, + { + role: "assistant" as const, + content: [ + { type: "text", text: "Hello, world." }, + ] as ContentBlock[], + }, + ]; + }; + + const ctx = makeCtx({ agentLoopRun }); + await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); + + // Exactly two `updateContent` calls land: + // 1. the debounced partial flush after both deltas accumulated, and + // 2. the final authoritative flush in `handleMessageComplete`. + // Without the debounce gate this would be one-per-delta + one final + // (3). Without the partial flush at all it would be just 1. + expect(updateMessageContentMock).toHaveBeenCalledTimes(2); + const calls = updateMessageContentMock.mock.calls as unknown as Array< + [string, string] + >; + const partialFlush = calls[0]; + expect(partialFlush?.[0]).toBe("msg-reserve"); + const partialBlocks = JSON.parse(partialFlush?.[1] ?? "[]") as Array<{ + type: string; + text?: string; + }>; + expect(partialBlocks).toEqual([{ type: "text", text: "Hello, world." }]); + }); + + test("size gate forces an immediate flush past PARTIAL_PERSIST_SIZE_THRESHOLD without waiting", async () => { + mockMessageById = { + id: "msg-reserve", + conversationId: "test-conv", + createdAt: 1234567, + role: "assistant", + content: "[]", + metadata: null, + }; + + // A single delta over the 1024-char size gate should bypass the + // debounce timer entirely and flush immediately. + const bigChunk = "x".repeat(1500); + + const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { + await onEvent({ type: "llm_call_started" }); + onEvent({ type: "text_delta", text: bigChunk }); + // Yield a microtask so the size-gated `void flushAccumulatedContent` + // can resolve. NO 250ms wait — proves the size gate is the + // mechanism that fired the flush. + await new Promise((resolve) => setImmediate(resolve)); + await onEvent({ + type: "message_complete", + message: { + role: "assistant", + content: [{ type: "text", text: bigChunk }], + }, + }); + onEvent({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + model: "test", + providerDurationMs: 50, + }); + return [ + ...messages, + { + role: "assistant" as const, + content: [{ type: "text", text: bigChunk }] as ContentBlock[], + }, + ]; + }; + + const ctx = makeCtx({ agentLoopRun }); + await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); + + // Partial flush + final flush. + expect(updateMessageContentMock).toHaveBeenCalledTimes(2); + const calls = updateMessageContentMock.mock.calls as unknown as Array< + [string, string] + >; + const partialBlocks = JSON.parse(calls[0]?.[1] ?? "[]") as Array<{ + type: string; + text?: string; + }>; + expect(partialBlocks[0]?.text?.length).toBe(1500); + }); + + test("handleToolUse flushes eagerly without waiting for the debounce", async () => { + mockMessageById = { + id: "msg-reserve", + conversationId: "test-conv", + createdAt: 1234567, + role: "assistant", + content: "[]", + metadata: null, + }; + + const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { + await onEvent({ type: "llm_call_started" }); + // A short text delta + a tool_use. The tool_use must flush + // immediately — the text delta alone wouldn't trip the size gate + // and is well under the 250ms time gate. + onEvent({ type: "text_delta", text: "Let me check that." }); + onEvent({ + type: "tool_use", + id: "tu-eager-flush", + name: "file_read", + input: { path: "/foo" }, + }); + await new Promise((resolve) => setImmediate(resolve)); + // Tool result + message_complete so the loop exits cleanly. + onEvent({ + type: "tool_result", + tool_use_id: "tu-eager-flush", + content: "ok", + is_error: false, + }); + await onEvent({ + type: "message_complete", + message: { + role: "assistant", + content: [ + { type: "text", text: "Let me check that." }, + { + type: "tool_use", + id: "tu-eager-flush", + name: "file_read", + input: { path: "/foo" }, + }, + ], + }, + }); + onEvent({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + model: "test", + providerDurationMs: 50, + }); + return [ + ...messages, + { + role: "assistant" as const, + content: [ + { type: "text", text: "Let me check that." }, + { + type: "tool_use", + id: "tu-eager-flush", + name: "file_read", + input: { path: "/foo" }, + }, + ] as ContentBlock[], + }, + ]; + }; + + const ctx = makeCtx({ agentLoopRun }); + await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); + + // Eager flush from handleToolUse + final flush. + expect(updateMessageContentMock).toHaveBeenCalledTimes(2); + const calls = updateMessageContentMock.mock.calls as unknown as Array< + [string, string] + >; + const partialBlocks = JSON.parse(calls[0]?.[1] ?? "[]") as Array<{ + type: string; + text?: string; + name?: string; + id?: string; + }>; + // The eager flush captures everything to date: the accumulated text + // block AND the freshly-pushed tool_use block. + expect(partialBlocks).toEqual([ + { type: "text", text: "Let me check that." }, + { + type: "tool_use", + id: "tu-eager-flush", + name: "file_read", + input: { path: "/foo" }, + }, + ]); + }); + + test("handleMessageComplete clears any pending debounce timer before the final flush", async () => { + mockMessageById = { + id: "msg-reserve", + conversationId: "test-conv", + createdAt: 1234567, + role: "assistant", + content: "[]", + metadata: null, + }; + + const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { + await onEvent({ type: "llm_call_started" }); + // Short delta — schedules a debounce timer but does NOT trip the + // size gate. message_complete then arrives immediately after, + // before the 250ms timer can fire. + onEvent({ type: "text_delta", text: "Quick reply." }); + await onEvent({ + type: "message_complete", + message: { + role: "assistant", + content: [{ type: "text", text: "Quick reply." }], + }, + }); + onEvent({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + model: "test", + providerDurationMs: 50, + }); + // Wait past the original debounce window to prove a late timer + // does NOT fire a stray partial flush. + await new Promise((resolve) => setTimeout(resolve, 350)); + return [ + ...messages, + { + role: "assistant" as const, + content: [{ type: "text", text: "Quick reply." }] as ContentBlock[], + }, + ]; + }; + + const ctx = makeCtx({ agentLoopRun }); + await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); + + // Only the final flush from `handleMessageComplete` lands. The + // debounced partial would have fired around T+250ms; the timer-clear + // at the top of `handleMessageComplete` cancels it. + expect(updateMessageContentMock).toHaveBeenCalledTimes(1); + }); + + test("partial flushes never trigger the indexer or attention projector", async () => { + mockMessageById = { + id: "msg-reserve", + conversationId: "test-conv", + createdAt: 1234567, + role: "assistant", + content: "[]", + metadata: null, + }; + // Use the size gate so the partial flush definitely fires inside + // the agent loop run. + const bigChunk = "y".repeat(1500); + + const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { + await onEvent({ type: "llm_call_started" }); + onEvent({ type: "text_delta", text: bigChunk }); + await new Promise((resolve) => setImmediate(resolve)); + // Snapshot the indexer/projector call counts AFTER the partial + // flush has run but BEFORE message_complete. They must be zero. + const indexerCallsBeforeComplete = + indexMessageNowMock.mock.calls.length; + const projectorCallsBeforeComplete = + projectAssistantMessageMock.mock.calls.length; + // Stash on a side channel the assertion phase can read. + (ctx as unknown as { __b6Snapshot?: [number, number] }).__b6Snapshot = [ + indexerCallsBeforeComplete, + projectorCallsBeforeComplete, + ]; + await onEvent({ + type: "message_complete", + message: { + role: "assistant", + content: [{ type: "text", text: bigChunk }], + }, + }); + onEvent({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + model: "test", + providerDurationMs: 50, + }); + return [ + ...messages, + { + role: "assistant" as const, + content: [{ type: "text", text: bigChunk }] as ContentBlock[], + }, + ]; + }; + + const ctx = makeCtx({ agentLoopRun }); + await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); + + const snapshot = (ctx as unknown as { __b6Snapshot?: [number, number] }) + .__b6Snapshot; + expect(snapshot).toBeDefined(); + // Indexer + projector were both ZERO during the mid-turn partial + // flush — they only fire from `handleMessageComplete` after the + // authoritative `updateContent`. + expect(snapshot![0]).toBe(0); + expect(snapshot![1]).toBe(0); + // After the loop completes the indexer + projector each ran exactly + // once (the B3 finalize path). + expect(indexMessageNowMock).toHaveBeenCalledTimes(1); + expect(projectAssistantMessageMock).toHaveBeenCalledTimes(1); + }); + + test("partial flushes redact secrets from text blocks before writing", async () => { + mockMessageById = { + id: "msg-reserve", + conversationId: "test-conv", + createdAt: 1234567, + role: "assistant", + content: "[]", + metadata: null, + }; + // A GitHub PAT-shaped token mid-stream — the redaction discipline + // mirrors `handleMessageComplete`'s final flush so a refresh mid-turn + // never sees plaintext credentials in the persisted row. + const ghToken = "ghp_" + "a".repeat(36); + const payload = "Here's the key: " + ghToken + " enjoy."; + + const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { + await onEvent({ type: "llm_call_started" }); + // Pad with size-gate content so the partial flush definitely lands. + onEvent({ type: "text_delta", text: payload }); + onEvent({ type: "text_delta", text: "z".repeat(1500) }); + await new Promise((resolve) => setImmediate(resolve)); + await onEvent({ + type: "message_complete", + message: { + role: "assistant", + content: [{ type: "text", text: payload + "z".repeat(1500) }], + }, + }); + onEvent({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + model: "test", + providerDurationMs: 50, + }); + return [ + ...messages, + { + role: "assistant" as const, + content: [ + { type: "text", text: payload + "z".repeat(1500) }, + ] as ContentBlock[], + }, + ]; + }; + + const ctx = makeCtx({ agentLoopRun }); + await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); + + expect(updateMessageContentMock).toHaveBeenCalledTimes(2); + const partialPayload = ( + updateMessageContentMock.mock.calls[0] as unknown as [string, string] + )[1]; + // The raw PAT must never appear in the persisted snapshot. The + // redaction substitute is implementation-defined; the contract here + // is "the literal token string is gone". + expect(partialPayload).not.toContain(ghToken); + }); + + test("provider-error cleanup deletes a row that has accumulated partial content", async () => { + // Regression check: B3's orphan-cleanup branch already deletes the + // reserved row when the LLM call exits via `provider_error`. B6 + // writes partial content to that row mid-turn; the cleanup must + // still fire and the row (along with its partial content) must + // still be deleted before the synthetic error message lands. + reserveMessageMock.mockImplementationOnce(async () => ({ + id: "msg-orphan-with-partial", + })); + + const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { + await onEvent({ type: "llm_call_started" }); + // A size-gated delta lands a partial flush BEFORE the provider + // error fires. + onEvent({ type: "text_delta", text: "z".repeat(1500) }); + await new Promise((resolve) => setImmediate(resolve)); + onEvent({ + type: "provider_error", + error: new Error("upstream 500"), + rawRequest: { model: "gpt-4.1", messages: [] }, + actualProvider: "openai", + }); + onEvent({ + type: "error", + error: new Error("upstream 500"), + }); + return messages; + }; + + const ctx = makeCtx({ agentLoopRun }); + await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); + + // Partial flush fired exactly once (before the provider error). + // The orphan row was then deleted; the synthetic error message is + // inserted separately via `addMessage` (`mock-msg-id`) and never + // touched by `updateContent`. + const partialFlushes = ( + updateMessageContentMock.mock.calls as unknown as Array< + [string, string] + > + ).filter(([id]) => id === "msg-orphan-with-partial"); + expect(partialFlushes).toHaveLength(1); + expect(deleteMessageByIdMock).toHaveBeenCalledTimes(1); + const deleteCall = deleteMessageByIdMock.mock.calls[0] as unknown as [ + string, + ]; + expect(deleteCall[0]).toBe("msg-orphan-with-partial"); + }); + }); + describe("pkbSystemReminderBlock metadata persistence", () => { test("persists pkbSystemReminderBlock in full mode with PKB active", async () => { const reminder = "\npkb content\n"; diff --git a/assistant/src/daemon/conversation-agent-loop-handlers.ts b/assistant/src/daemon/conversation-agent-loop-handlers.ts index c4003906fc7..40dc1b70344 100644 --- a/assistant/src/daemon/conversation-agent-loop-handlers.ts +++ b/assistant/src/daemon/conversation-agent-loop-handlers.ts @@ -90,6 +90,29 @@ import type { const log = getLogger("agent-loop-handlers"); +// ── B6: partial-persistence tunables ───────────────────────────────── +// +// `handleTextDelta` accumulates streamed text into +// `state.accumulatedContentBlocks` and flushes the row via +// `updateMessageContent` on a first-fire-wins debounce: +// +// • Time gate: at most one flush per PARTIAL_PERSIST_DEBOUNCE_MS +// (~4 writes/sec on a steady stream). Cheap enough that a refresh +// mid-turn never lands more than ~250ms behind the wire. +// • Size gate: PARTIAL_PERSIST_SIZE_THRESHOLD bytes of new text since +// the last flush forces an immediate flush regardless of timer. +// Catches bursts that wouldn't trip the time gate. +// +// `handleToolUse` flushes eagerly — tool_use blocks are atomic and small +// relative to text streams, so debounce gives no win there. +// +// Indexer + projector still fire ONLY at `handleMessageComplete`; the +// row's `content` may temporarily reflect a partial assistant turn but +// is never indexed mid-stream. See B6 spec at +// `scratch/b6-partial-persist-spec.md`. +const PARTIAL_PERSIST_DEBOUNCE_MS = 250; +const PARTIAL_PERSIST_SIZE_THRESHOLD = 1024; + /** * Build a {@link TurnContext} from the handler's deps for pipeline logging * and plugin attribution. @@ -239,6 +262,40 @@ export interface EventHandlerState { readonly serverToolStartedAt: Map; /** Original input from server_tool_start, keyed by tool_use_id, so the complete handler can read the query. */ readonly serverToolInputs: Map>; + /** + * B6 partial-persistence accumulator. Mirrors the in-progress + * assistant content (text + tool_use blocks) emitted to the wire + * since the current row's `handleLlmCallStarted`, so a flush can + * write the partial snapshot to the row's `content` column without + * waiting for `message_complete`. Reset at `handleLlmCallStarted` + * for each new assistant row reservation (a tool-flush + second LLM + * call within the same run reserves a fresh row and starts a fresh + * accumulator). Final flush in `handleMessageComplete` uses the + * authoritative `event.message.content` instead — this accumulator + * is for *mid-turn* snapshots only. + * + * Thinking blocks are intentionally NOT mirrored here. The provider + * emits them via `thinking_delta`; including them mid-turn would + * make the persisted snapshot more chatty than the wire emit (which + * already drops thinking text from the on-screen display). The + * final `event.message.content` still carries thinking blocks, so + * a finalized row reflects them — only the mid-turn snapshot omits. + */ + accumulatedContentBlocks: ContentBlock[]; + /** + * Bytes of new text appended to {@link accumulatedContentBlocks} + * since the last flush. Tripping `PARTIAL_PERSIST_SIZE_THRESHOLD` + * forces an immediate flush regardless of the debounce timer; reset + * to 0 after every flush. + */ + accumulatedCharsSinceLastFlush: number; + /** + * Active debounce timer for partial persistence. `undefined` when + * idle (no text since last flush, or a size-gate / eager flush just + * fired). Cleared at the start of `handleMessageComplete` so the + * final authoritative flush never races a debounced partial write. + */ + pendingPartialFlushTimer: ReturnType | undefined; } /** Immutable context shared across event handlers within a single agent loop run. */ @@ -299,9 +356,146 @@ export function createEventHandlerState(): EventHandlerState { turnStartedAt: Date.now(), serverToolStartedAt: new Map(), serverToolInputs: new Map(), + accumulatedContentBlocks: [], + accumulatedCharsSinceLastFlush: 0, + pendingPartialFlushTimer: undefined, }; } +// ── B6: partial-persistence helpers ────────────────────────────────── + +/** + * Append a chunk of streamed text to the partial-persist accumulator. + * + * If the tail block is already a `text` block, the chunk is appended in + * place (the producer of `event.message.content` at finalize fuses + * consecutive deltas into a single text block, so we mirror that + * shape). Otherwise a new text block is pushed. + * + * Returns the new total of unflushed bytes so the caller can drive the + * size gate. + */ +function appendTextToAccumulator( + state: EventHandlerState, + text: string, +): number { + if (text.length === 0) return state.accumulatedCharsSinceLastFlush; + const tail = state.accumulatedContentBlocks.at(-1); + if (tail && tail.type === "text") { + tail.text = tail.text + text; + } else { + state.accumulatedContentBlocks.push({ type: "text", text }); + } + state.accumulatedCharsSinceLastFlush += text.length; + return state.accumulatedCharsSinceLastFlush; +} + +/** + * Reset partial-persist accumulator state. Called from + * `handleLlmCallStarted` after reserving a fresh assistant row (each + * row has its own accumulator) and after the final flush in + * `handleMessageComplete`. Idempotent and safe to call repeatedly. + */ +function resetPartialPersistAccumulator(state: EventHandlerState): void { + if (state.pendingPartialFlushTimer !== undefined) { + clearTimeout(state.pendingPartialFlushTimer); + state.pendingPartialFlushTimer = undefined; + } + state.accumulatedContentBlocks = []; + state.accumulatedCharsSinceLastFlush = 0; +} + +/** + * Flush the partial-persist accumulator to the assistant row via + * `updateContent`. Applies the same redaction discipline as the final + * flush in `handleMessageComplete` (text blocks get + * `redactSecrets`-cleaned; non-text blocks pass through). UI surfaces + * and directive cleaning are NOT applied — surfaces aren't known until + * finalize, and directives in mid-turn text render harmlessly as raw + * markup if a user happens to refresh before they're stripped. + * + * No-ops if: + * • there's no reserved assistant row id (pre-`handleLlmCallStarted` + * or post-cleanup), or + * • the accumulator is empty (nothing meaningful has streamed yet). + * + * After a successful pipeline write the timer is cleared and the size + * counter resets to 0. The accumulator itself is *retained* so a + * subsequent partial flush rewrites the full content (the accumulated + * blocks are the in-memory mirror of what should be in the row). + * + * Failures are logged but never thrown — a missed partial flush is + * always recoverable by either the next debounce tick or the final + * authoritative flush at `handleMessageComplete`. Throwing here would + * tear down the agent loop turn for a degradation that the eventual + * finalize will resolve. + */ +async function flushAccumulatedContent( + state: EventHandlerState, + deps: EventHandlerDeps, +): Promise { + if (state.pendingPartialFlushTimer !== undefined) { + clearTimeout(state.pendingPartialFlushTimer); + state.pendingPartialFlushTimer = undefined; + } + const messageId = state.lastAssistantMessageId; + if (messageId === undefined) return; + if (state.accumulatedContentBlocks.length === 0) return; + + const redacted = state.accumulatedContentBlocks.map((block) => { + if (block.type === "text") { + return { ...block, text: redactSecrets(block.text) }; + } + return block; + }); + const contentJson = JSON.stringify(redacted); + + try { + await runPipeline( + "persistence", + getMiddlewaresFor("persistence"), + defaultPersistenceTerminal, + { + op: "updateContent", + messageId, + content: contentJson, + }, + buildHandlerTurnContext(deps), + DEFAULT_TIMEOUTS.persistence, + ); + state.accumulatedCharsSinceLastFlush = 0; + } catch (err) { + deps.rlog.warn( + { err, messageId }, + "B6 partial flush of accumulated assistant content failed; finalize at message_complete will recover", + ); + } +} + +/** + * Either flush immediately (size gate tripped) or schedule a debounced + * flush (time gate). First-fire wins — if a timer is already pending + * and the size gate has not tripped, this is a no-op. + * + * Called from `handleTextDelta` after appending the drained delta. + * `handleToolUse` flushes eagerly via `flushAccumulatedContent` instead + * of going through this gate. + */ +function scheduleOrFireFlush( + state: EventHandlerState, + deps: EventHandlerDeps, +): void { + if (state.accumulatedCharsSinceLastFlush >= PARTIAL_PERSIST_SIZE_THRESHOLD) { + void flushAccumulatedContent(state, deps); + return; + } + if (state.pendingPartialFlushTimer !== undefined) return; + state.pendingPartialFlushTimer = setTimeout(() => { + state.pendingPartialFlushTimer = undefined; + void flushAccumulatedContent(state, deps); + }, PARTIAL_PERSIST_DEBOUNCE_MS); +} + // ── Shared Helper ──────────────────────────────────────────────────── // providerNameOverride should be supplied when the caller already knows the @@ -545,6 +739,12 @@ export async function handleLlmCallStarted( )) as PersistReserveResult; state.lastAssistantMessageId = reserveResult.message.id; state.assistantRowAwaitingFinalization = true; + // B6: fresh row → fresh accumulator. If an earlier (failed) LLM call + // within the same run left partial state behind, the + // `assistantRowAwaitingFinalization` cleanup above already deleted + // the orphan row, so the accumulator content would point at a + // non-existent id. Reset here so the new row starts from zero. + resetPartialPersistAccumulator(state); deps.onEvent({ type: "assistant_turn_start", messageId: reserveResult.message.id, @@ -583,6 +783,12 @@ function handleTextDelta( messageId: state.lastAssistantMessageId, }); if (deps.shouldGenerateTitle) state.firstAssistantText += drained.emitText; + // B6: mirror the drained delta into the partial-persist accumulator + // (the wire's drained text, not the raw event.text, so a refresh + // mid-turn sees the same content the user was watching live — raw + // directive markup never reaches the persisted snapshot). + appendTextToAccumulator(state, drained.emitText); + scheduleOrFireFlush(state, deps); } } @@ -651,6 +857,19 @@ export function handleToolUse( toolUseId: event.id, messageId: state.lastAssistantMessageId, }); + // B6: tool_use blocks are atomic (the AgentEvent only emits once the + // provider has finished streaming the input), so push the full block + // onto the accumulator and flush eagerly. No debounce — tool blocks + // are infrequent compared to text deltas and a refresh between + // tool_use_start and the corresponding tool_result is the + // most-visible mid-turn refresh case. + state.accumulatedContentBlocks.push({ + type: "tool_use", + id: event.id, + name: event.name, + input: event.input, + }); + void flushAccumulatedContent(state, deps); } export function handleToolUsePreviewStart( @@ -1145,6 +1364,20 @@ export async function handleMessageComplete( // Reset per-turn tool tracking for the new turn. state.currentTurnToolUseIds = []; + // B6: cancel any pending debounced partial flush before the + // authoritative `updateContent` below. Without this, a timer that + // fires between this line and the final pipeline call could + // double-write the row (idempotent in content but wastes a write) + // or worse, race ahead of the indexer/projector read and serve a + // stale snapshot. The accumulator itself is left untouched until + // after finalize — the orchestrator's downstream paths don't read + // it, but a defensive reset on the next turn's + // `handleLlmCallStarted` keeps state hygienic regardless. + if (state.pendingPartialFlushTimer !== undefined) { + clearTimeout(state.pendingPartialFlushTimer); + state.pendingPartialFlushTimer = undefined; + } + // Flush any remaining directive display buffer if (state.pendingDirectiveDisplayBuffer.length > 0) { deps.onEvent({ From 196a2d86af1949af234b1cf9460615ea631b4e20 Mon Sep 17 00:00:00 2001 From: "vellum-apollo-bot[bot]" <242025090+vellum-apollo-bot[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 19:43:13 +0000 Subject: [PATCH 2/5] test(agent-loop): fix camelCase field names + widen partialBlocks type Two TS errors in conversation-agent-loop.test.ts: 1. tool_result event used snake_case (`tool_use_id`, `is_error`) instead of the AgentEvent shape's camelCase (`toolUseId`, `isError`). The snake_case shape is the wire/block form, not the in-process event. 2. `partialBlocks` type didn't allow `input` on tool_use blocks, so `expect(partialBlocks).toEqual([... { input: { path: '/foo' } }])` failed type-check. Widened with `input?: Record`. --- assistant/src/__tests__/conversation-agent-loop.test.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/assistant/src/__tests__/conversation-agent-loop.test.ts b/assistant/src/__tests__/conversation-agent-loop.test.ts index 080130f70eb..30b79a3f5df 100644 --- a/assistant/src/__tests__/conversation-agent-loop.test.ts +++ b/assistant/src/__tests__/conversation-agent-loop.test.ts @@ -3612,9 +3612,9 @@ describe("session-agent-loop", () => { // Tool result + message_complete so the loop exits cleanly. onEvent({ type: "tool_result", - tool_use_id: "tu-eager-flush", + toolUseId: "tu-eager-flush", content: "ok", - is_error: false, + isError: false, }); await onEvent({ type: "message_complete", @@ -3668,6 +3668,7 @@ describe("session-agent-loop", () => { text?: string; name?: string; id?: string; + input?: Record; }>; // The eager flush captures everything to date: the accumulated text // block AND the freshly-pushed tool_use block. From 26b1cc6a130471837d5cc75cc5f2ffe3a67dc33b Mon Sep 17 00:00:00 2001 From: "vellum-apollo-bot[bot]" <242025090+vellum-apollo-bot[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 21:02:43 +0000 Subject: [PATCH 3/5] refactor(agent-loop): timer-only partial persist + shared content build pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review feedback on #32602: - Drop size-gate; debounce timer only (Vargas) - Drop eager flush in handleToolUse — AgentLoop emits tool_use AFTER message_complete, so any flush there would overwrite the finalized row (Codex P1 / Vargas line 872) - Extract buildPersistedAssistantContent shared between the partial flush and handleMessageComplete so both writes use the same cleanAssistantContent + surfaces + redactSecrets pipeline (Vargas line 378) - Track in-flight partial flush promise on EventHandlerState; await it at message_complete before the authoritative final write so a partial pipeline call dispatched moments before finalize cannot settle after it (Codex P2) - Scrub B6 / PR-number references from comments (Vargas line 266) - Tests: drop size-gate test, replace eager-flush test with 'handleToolUse does NOT flush' assertion, rewrite indexer/redact/ provider-error tests to use the 250ms debounce instead of size-gate --- .../__tests__/conversation-agent-loop.test.ts | 187 +++------- .../conversation-agent-loop-handlers.ts | 333 ++++++++++-------- 2 files changed, 240 insertions(+), 280 deletions(-) diff --git a/assistant/src/__tests__/conversation-agent-loop.test.ts b/assistant/src/__tests__/conversation-agent-loop.test.ts index 30b79a3f5df..0c33bd1173d 100644 --- a/assistant/src/__tests__/conversation-agent-loop.test.ts +++ b/assistant/src/__tests__/conversation-agent-loop.test.ts @@ -3444,24 +3444,26 @@ describe("session-agent-loop", () => { }); }); - describe("B6 partial persistence", () => { - // The pre-B6 flow reserves an empty assistant row at `llm_call_started` - // (`content: "[]"`) and never touches that row again until + describe("partial persistence", () => { + // The legacy flow reserves an empty assistant row at `llm_call_started` + // (`content: "[]"`) and never touches it again until // `handleMessageComplete` fires the single authoritative // `updateContent`. Between those events the row is empty for the full // duration of a turn — a browser refresh mid-turn sees nothing where // the in-progress assistant reply should be. // - // B6 closes that durability gap with a debounced partial flush from - // `handleTextDelta` (250ms time gate / 1024-char size gate, first-fire - // wins) and an eager flush from `handleToolUse` (tool_use blocks are - // atomic, no debounce needed). The indexer + projector still fire - // ONLY at `message_complete` — partial rows are never indexed. + // Partial persistence closes that durability gap with a debounced + // flush from `handleTextDelta` (250ms timer). `handleToolUse` + // intentionally does NOT flush — `AgentLoop.run` emits `tool_use` + // strictly AFTER `message_complete`, so any flush from that handler + // would land after the authoritative finalize and overwrite the + // finalized row. The indexer + projector still fire ONLY at + // `message_complete` — partial rows are never indexed. // // These tests pin down the wire-level contract by counting // `updateMessageContent` calls and inspecting the JSON payload of the // partial-flush writes. The indexing / sync-invalidation paths are - // covered by the B3 block above. + // covered by the pre-allocation block above. test("debounced time gate flushes one partial write after PARTIAL_PERSIST_DEBOUNCE_MS", async () => { mockMessageById = { @@ -3527,66 +3529,12 @@ describe("session-agent-loop", () => { expect(partialBlocks).toEqual([{ type: "text", text: "Hello, world." }]); }); - test("size gate forces an immediate flush past PARTIAL_PERSIST_SIZE_THRESHOLD without waiting", async () => { - mockMessageById = { - id: "msg-reserve", - conversationId: "test-conv", - createdAt: 1234567, - role: "assistant", - content: "[]", - metadata: null, - }; - - // A single delta over the 1024-char size gate should bypass the - // debounce timer entirely and flush immediately. - const bigChunk = "x".repeat(1500); - - const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { - await onEvent({ type: "llm_call_started" }); - onEvent({ type: "text_delta", text: bigChunk }); - // Yield a microtask so the size-gated `void flushAccumulatedContent` - // can resolve. NO 250ms wait — proves the size gate is the - // mechanism that fired the flush. - await new Promise((resolve) => setImmediate(resolve)); - await onEvent({ - type: "message_complete", - message: { - role: "assistant", - content: [{ type: "text", text: bigChunk }], - }, - }); - onEvent({ - type: "usage", - inputTokens: 10, - outputTokens: 5, - model: "test", - providerDurationMs: 50, - }); - return [ - ...messages, - { - role: "assistant" as const, - content: [{ type: "text", text: bigChunk }] as ContentBlock[], - }, - ]; - }; - - const ctx = makeCtx({ agentLoopRun }); - await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); - - // Partial flush + final flush. - expect(updateMessageContentMock).toHaveBeenCalledTimes(2); - const calls = updateMessageContentMock.mock.calls as unknown as Array< - [string, string] - >; - const partialBlocks = JSON.parse(calls[0]?.[1] ?? "[]") as Array<{ - type: string; - text?: string; - }>; - expect(partialBlocks[0]?.text?.length).toBe(1500); - }); - - test("handleToolUse flushes eagerly without waiting for the debounce", async () => { + test("handleToolUse does NOT trigger a partial flush of its own", async () => { + // `AgentLoop.run` emits `tool_use` strictly AFTER `message_complete`, + // so a flush from the tool_use handler would land after the + // authoritative final `updateContent` and overwrite the finalized + // row (Codex P1 / Vargas review feedback). The handler must be a + // no-op for the partial-persist accumulator. mockMessageById = { id: "msg-reserve", conversationId: "test-conv", @@ -3598,21 +3546,21 @@ describe("session-agent-loop", () => { const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { await onEvent({ type: "llm_call_started" }); - // A short text delta + a tool_use. The tool_use must flush - // immediately — the text delta alone wouldn't trip the size gate - // and is well under the 250ms time gate. - onEvent({ type: "text_delta", text: "Let me check that." }); + // No text delta — only a tool_use. If `handleToolUse` were + // flushing, this would land a partial write before + // `message_complete`. onEvent({ type: "tool_use", - id: "tu-eager-flush", + id: "tu-no-flush", name: "file_read", input: { path: "/foo" }, }); + // Yield a microtask so any (incorrectly) fire-and-forget + // pipeline call has a chance to land before message_complete. await new Promise((resolve) => setImmediate(resolve)); - // Tool result + message_complete so the loop exits cleanly. onEvent({ type: "tool_result", - toolUseId: "tu-eager-flush", + toolUseId: "tu-no-flush", content: "ok", isError: false, }); @@ -3621,10 +3569,9 @@ describe("session-agent-loop", () => { message: { role: "assistant", content: [ - { type: "text", text: "Let me check that." }, { type: "tool_use", - id: "tu-eager-flush", + id: "tu-no-flush", name: "file_read", input: { path: "/foo" }, }, @@ -3643,10 +3590,9 @@ describe("session-agent-loop", () => { { role: "assistant" as const, content: [ - { type: "text", text: "Let me check that." }, { type: "tool_use", - id: "tu-eager-flush", + id: "tu-no-flush", name: "file_read", input: { path: "/foo" }, }, @@ -3658,29 +3604,10 @@ describe("session-agent-loop", () => { const ctx = makeCtx({ agentLoopRun }); await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); - // Eager flush from handleToolUse + final flush. - expect(updateMessageContentMock).toHaveBeenCalledTimes(2); - const calls = updateMessageContentMock.mock.calls as unknown as Array< - [string, string] - >; - const partialBlocks = JSON.parse(calls[0]?.[1] ?? "[]") as Array<{ - type: string; - text?: string; - name?: string; - id?: string; - input?: Record; - }>; - // The eager flush captures everything to date: the accumulated text - // block AND the freshly-pushed tool_use block. - expect(partialBlocks).toEqual([ - { type: "text", text: "Let me check that." }, - { - type: "tool_use", - id: "tu-eager-flush", - name: "file_read", - input: { path: "/foo" }, - }, - ]); + // Only the authoritative final flush from `handleMessageComplete` + // lands. A partial flush from `handleToolUse` would have made this + // 2; that's the regression this test guards against. + expect(updateMessageContentMock).toHaveBeenCalledTimes(1); }); test("handleMessageComplete clears any pending debounce timer before the final flush", async () => { @@ -3743,14 +3670,13 @@ describe("session-agent-loop", () => { content: "[]", metadata: null, }; - // Use the size gate so the partial flush definitely fires inside - // the agent loop run. - const bigChunk = "y".repeat(1500); const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { await onEvent({ type: "llm_call_started" }); - onEvent({ type: "text_delta", text: bigChunk }); - await new Promise((resolve) => setImmediate(resolve)); + onEvent({ type: "text_delta", text: "hello world" }); + // Wait past the 250ms debounce so the partial flush definitely + // lands BEFORE message_complete fires. + await new Promise((resolve) => setTimeout(resolve, 350)); // Snapshot the indexer/projector call counts AFTER the partial // flush has run but BEFORE message_complete. They must be zero. const indexerCallsBeforeComplete = @@ -3758,7 +3684,9 @@ describe("session-agent-loop", () => { const projectorCallsBeforeComplete = projectAssistantMessageMock.mock.calls.length; // Stash on a side channel the assertion phase can read. - (ctx as unknown as { __b6Snapshot?: [number, number] }).__b6Snapshot = [ + ( + ctx as unknown as { __partialSnapshot?: [number, number] } + ).__partialSnapshot = [ indexerCallsBeforeComplete, projectorCallsBeforeComplete, ]; @@ -3766,7 +3694,7 @@ describe("session-agent-loop", () => { type: "message_complete", message: { role: "assistant", - content: [{ type: "text", text: bigChunk }], + content: [{ type: "text", text: "hello world" }], }, }); onEvent({ @@ -3780,7 +3708,9 @@ describe("session-agent-loop", () => { ...messages, { role: "assistant" as const, - content: [{ type: "text", text: bigChunk }] as ContentBlock[], + content: [ + { type: "text", text: "hello world" }, + ] as ContentBlock[], }, ]; }; @@ -3788,8 +3718,9 @@ describe("session-agent-loop", () => { const ctx = makeCtx({ agentLoopRun }); await runAgentLoopImpl(ctx, "hi", "msg-1", () => {}); - const snapshot = (ctx as unknown as { __b6Snapshot?: [number, number] }) - .__b6Snapshot; + const snapshot = ( + ctx as unknown as { __partialSnapshot?: [number, number] } + ).__partialSnapshot; expect(snapshot).toBeDefined(); // Indexer + projector were both ZERO during the mid-turn partial // flush — they only fire from `handleMessageComplete` after the @@ -3797,7 +3728,7 @@ describe("session-agent-loop", () => { expect(snapshot![0]).toBe(0); expect(snapshot![1]).toBe(0); // After the loop completes the indexer + projector each ran exactly - // once (the B3 finalize path). + // once (the pre-allocation finalize path). expect(indexMessageNowMock).toHaveBeenCalledTimes(1); expect(projectAssistantMessageMock).toHaveBeenCalledTimes(1); }); @@ -3819,15 +3750,14 @@ describe("session-agent-loop", () => { const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { await onEvent({ type: "llm_call_started" }); - // Pad with size-gate content so the partial flush definitely lands. onEvent({ type: "text_delta", text: payload }); - onEvent({ type: "text_delta", text: "z".repeat(1500) }); - await new Promise((resolve) => setImmediate(resolve)); + // Wait past the 250ms debounce so the partial flush lands. + await new Promise((resolve) => setTimeout(resolve, 350)); await onEvent({ type: "message_complete", message: { role: "assistant", - content: [{ type: "text", text: payload + "z".repeat(1500) }], + content: [{ type: "text", text: payload }], }, }); onEvent({ @@ -3841,9 +3771,7 @@ describe("session-agent-loop", () => { ...messages, { role: "assistant" as const, - content: [ - { type: "text", text: payload + "z".repeat(1500) }, - ] as ContentBlock[], + content: [{ type: "text", text: payload }] as ContentBlock[], }, ]; }; @@ -3862,21 +3790,22 @@ describe("session-agent-loop", () => { }); test("provider-error cleanup deletes a row that has accumulated partial content", async () => { - // Regression check: B3's orphan-cleanup branch already deletes the - // reserved row when the LLM call exits via `provider_error`. B6 - // writes partial content to that row mid-turn; the cleanup must - // still fire and the row (along with its partial content) must - // still be deleted before the synthetic error message lands. + // Regression check: the pre-allocation orphan-cleanup branch + // already deletes the reserved row when the LLM call exits via + // `provider_error`. Partial-persist writes content to that row + // mid-turn; the cleanup must still fire and the row (along with + // its partial content) must still be deleted before the synthetic + // error message lands. reserveMessageMock.mockImplementationOnce(async () => ({ id: "msg-orphan-with-partial", })); const agentLoopRun: AgentLoopRun = async (messages, onEvent) => { await onEvent({ type: "llm_call_started" }); - // A size-gated delta lands a partial flush BEFORE the provider + // A debounced delta lands a partial flush BEFORE the provider // error fires. - onEvent({ type: "text_delta", text: "z".repeat(1500) }); - await new Promise((resolve) => setImmediate(resolve)); + onEvent({ type: "text_delta", text: "hello world" }); + await new Promise((resolve) => setTimeout(resolve, 350)); onEvent({ type: "provider_error", error: new Error("upstream 500"), diff --git a/assistant/src/daemon/conversation-agent-loop-handlers.ts b/assistant/src/daemon/conversation-agent-loop-handlers.ts index 40dc1b70344..09fbb1b084a 100644 --- a/assistant/src/daemon/conversation-agent-loop-handlers.ts +++ b/assistant/src/daemon/conversation-agent-loop-handlers.ts @@ -90,28 +90,19 @@ import type { const log = getLogger("agent-loop-handlers"); -// ── B6: partial-persistence tunables ───────────────────────────────── +// ── Partial-persistence tunables ───────────────────────────────────── // // `handleTextDelta` accumulates streamed text into // `state.accumulatedContentBlocks` and flushes the row via -// `updateMessageContent` on a first-fire-wins debounce: -// -// • Time gate: at most one flush per PARTIAL_PERSIST_DEBOUNCE_MS -// (~4 writes/sec on a steady stream). Cheap enough that a refresh -// mid-turn never lands more than ~250ms behind the wire. -// • Size gate: PARTIAL_PERSIST_SIZE_THRESHOLD bytes of new text since -// the last flush forces an immediate flush regardless of timer. -// Catches bursts that wouldn't trip the time gate. -// -// `handleToolUse` flushes eagerly — tool_use blocks are atomic and small -// relative to text streams, so debounce gives no win there. +// `updateMessageContent` on a debounce timer: at most one flush per +// `PARTIAL_PERSIST_DEBOUNCE_MS` (~4 writes/sec on a steady stream). +// Cheap enough that a refresh mid-turn never lands more than ~250ms +// behind the wire. // // Indexer + projector still fire ONLY at `handleMessageComplete`; the // row's `content` may temporarily reflect a partial assistant turn but -// is never indexed mid-stream. See B6 spec at -// `scratch/b6-partial-persist-spec.md`. +// is never indexed mid-stream. const PARTIAL_PERSIST_DEBOUNCE_MS = 250; -const PARTIAL_PERSIST_SIZE_THRESHOLD = 1024; /** * Build a {@link TurnContext} from the handler's deps for pipeline logging @@ -263,39 +254,39 @@ export interface EventHandlerState { /** Original input from server_tool_start, keyed by tool_use_id, so the complete handler can read the query. */ readonly serverToolInputs: Map>; /** - * B6 partial-persistence accumulator. Mirrors the in-progress - * assistant content (text + tool_use blocks) emitted to the wire - * since the current row's `handleLlmCallStarted`, so a flush can - * write the partial snapshot to the row's `content` column without - * waiting for `message_complete`. Reset at `handleLlmCallStarted` - * for each new assistant row reservation (a tool-flush + second LLM - * call within the same run reserves a fresh row and starts a fresh - * accumulator). Final flush in `handleMessageComplete` uses the - * authoritative `event.message.content` instead — this accumulator - * is for *mid-turn* snapshots only. + * Partial-persistence accumulator. Mirrors the in-progress assistant + * text emitted to the wire since the current row's + * `handleLlmCallStarted`, so a debounced flush can write the partial + * snapshot to the row's `content` column without waiting for + * `message_complete`. Reset at `handleLlmCallStarted` for each new + * assistant row reservation (a tool-flush + second LLM call within + * the same run reserves a fresh row and starts a fresh accumulator). + * Final flush in `handleMessageComplete` uses the authoritative + * `event.message.content` instead — this accumulator is for + * *mid-turn* snapshots only. * - * Thinking blocks are intentionally NOT mirrored here. The provider - * emits them via `thinking_delta`; including them mid-turn would - * make the persisted snapshot more chatty than the wire emit (which - * already drops thinking text from the on-screen display). The - * final `event.message.content` still carries thinking blocks, so - * a finalized row reflects them — only the mid-turn snapshot omits. + * Only text blocks are accumulated. `tool_use` events arrive AFTER + * `message_complete` finalizes the row, so mirroring them mid-turn + * would have no effect; `thinking` blocks are intentionally omitted + * to keep the persisted snapshot in lockstep with what `cleanAssistantContent` + * produces at finalize (which the partial flush also runs through). */ accumulatedContentBlocks: ContentBlock[]; - /** - * Bytes of new text appended to {@link accumulatedContentBlocks} - * since the last flush. Tripping `PARTIAL_PERSIST_SIZE_THRESHOLD` - * forces an immediate flush regardless of the debounce timer; reset - * to 0 after every flush. - */ - accumulatedCharsSinceLastFlush: number; /** * Active debounce timer for partial persistence. `undefined` when - * idle (no text since last flush, or a size-gate / eager flush just - * fired). Cleared at the start of `handleMessageComplete` so the - * final authoritative flush never races a debounced partial write. + * idle (no text since last flush). Cleared at the start of + * `handleMessageComplete` so the final authoritative flush never + * races a debounced partial write. */ pendingPartialFlushTimer: ReturnType | undefined; + /** + * In-flight partial flush write (the most recently dispatched + * `flushAccumulatedContent` pipeline call). `handleMessageComplete` + * awaits this before its final `updateContent` so a partial write + * that started just before finalize can never overwrite the + * authoritative content. + */ + pendingPartialFlushPromise: Promise | undefined; } /** Immutable context shared across event handlers within a single agent loop run. */ @@ -357,12 +348,74 @@ export function createEventHandlerState(): EventHandlerState { serverToolStartedAt: new Map(), serverToolInputs: new Map(), accumulatedContentBlocks: [], - accumulatedCharsSinceLastFlush: 0, pendingPartialFlushTimer: undefined, + pendingPartialFlushPromise: undefined, }; } -// ── B6: partial-persistence helpers ────────────────────────────────── +// ── Partial-persistence helpers ────────────────────────────────────── + +/** + * Type of a single UI surface that may be merged into persisted assistant + * content. Mirrors the inline structural type on + * `AgentLoopConversationContext.currentTurnSurfaces` — sharing it via + * this alias keeps {@link buildPersistedAssistantContent} usable in + * tests without routing them through the full context. + */ +type AssistantSurface = + AgentLoopConversationContext["currentTurnSurfaces"] extends ReadonlyArray< + infer S + > + ? S + : never; + +/** + * Build the canonical persisted-content array for an assistant row, + * applying the same pipeline both `handleMessageComplete` (the + * authoritative final write) and partial flushes use: + * + * 1. {@link cleanAssistantContent} — drops Anthropic placeholder + * sentinel text blocks and strips attachment directives from + * text. Directive results (`directives`, `warnings`) are NOT + * returned here — handleMessageComplete owns the canonical + * accumulation path for those. + * 2. UI surface blocks (if any) are appended. + * 3. Text blocks are run through {@link redactSecrets}. + * + * Used in two places: + * • final flush in `handleMessageComplete` with `event.message.content` + * and `deps.ctx.currentTurnSurfaces`. + * • partial flush in `flushAccumulatedContent` with the in-memory + * accumulator and an empty surfaces array (surfaces are produced + * by tool execution, which runs after `message_complete`). + */ +function buildPersistedAssistantContent( + rawBlocks: readonly ContentBlock[], + surfaces: readonly AssistantSurface[], +): ContentBlock[] { + const { cleanedContent } = cleanAssistantContent(rawBlocks); + const cleaned = cleanedContent as ContentBlock[]; + const withSurfaces: ContentBlock[] = [...cleaned]; + for (const surface of surfaces) { + withSurfaces.push({ + type: "ui_surface", + surfaceId: surface.surfaceId, + surfaceType: surface.surfaceType, + title: surface.title, + data: surface.data, + actions: surface.actions, + display: surface.display, + ...(surface.persistent ? { persistent: true } : {}), + } as unknown as ContentBlock); + } + return withSurfaces.map((block) => { + if (block.type === "text") { + const tb = block as Extract; + return { ...tb, text: redactSecrets(tb.text) }; + } + return block; + }); +} /** * Append a chunk of streamed text to the partial-persist accumulator. @@ -371,23 +424,18 @@ export function createEventHandlerState(): EventHandlerState { * place (the producer of `event.message.content` at finalize fuses * consecutive deltas into a single text block, so we mirror that * shape). Otherwise a new text block is pushed. - * - * Returns the new total of unflushed bytes so the caller can drive the - * size gate. */ function appendTextToAccumulator( state: EventHandlerState, text: string, -): number { - if (text.length === 0) return state.accumulatedCharsSinceLastFlush; +): void { + if (text.length === 0) return; const tail = state.accumulatedContentBlocks.at(-1); if (tail && tail.type === "text") { tail.text = tail.text + text; } else { state.accumulatedContentBlocks.push({ type: "text", text }); } - state.accumulatedCharsSinceLastFlush += text.length; - return state.accumulatedCharsSinceLastFlush; } /** @@ -402,53 +450,43 @@ function resetPartialPersistAccumulator(state: EventHandlerState): void { state.pendingPartialFlushTimer = undefined; } state.accumulatedContentBlocks = []; - state.accumulatedCharsSinceLastFlush = 0; + state.pendingPartialFlushPromise = undefined; } /** * Flush the partial-persist accumulator to the assistant row via - * `updateContent`. Applies the same redaction discipline as the final - * flush in `handleMessageComplete` (text blocks get - * `redactSecrets`-cleaned; non-text blocks pass through). UI surfaces - * and directive cleaning are NOT applied — surfaces aren't known until - * finalize, and directives in mid-turn text render harmlessly as raw - * markup if a user happens to refresh before they're stripped. + * `updateContent`. Reuses {@link buildPersistedAssistantContent} so the + * partial snapshot lands in the same shape `handleMessageComplete` + * would produce (with `surfaces=[]` — surfaces emerge from tool + * execution which runs after `message_complete`). * * No-ops if: * • there's no reserved assistant row id (pre-`handleLlmCallStarted` * or post-cleanup), or * • the accumulator is empty (nothing meaningful has streamed yet). * - * After a successful pipeline write the timer is cleared and the size - * counter resets to 0. The accumulator itself is *retained* so a - * subsequent partial flush rewrites the full content (the accumulated - * blocks are the in-memory mirror of what should be in the row). + * The pending-promise handle is published on `state` so + * `handleMessageComplete` can await it before the authoritative final + * write, serializing the two `updateContent` calls and preventing a + * partial write from overwriting the finalized row. * * Failures are logged but never thrown — a missed partial flush is * always recoverable by either the next debounce tick or the final - * authoritative flush at `handleMessageComplete`. Throwing here would - * tear down the agent loop turn for a degradation that the eventual - * finalize will resolve. + * authoritative flush at `handleMessageComplete`. */ async function flushAccumulatedContent( state: EventHandlerState, deps: EventHandlerDeps, ): Promise { - if (state.pendingPartialFlushTimer !== undefined) { - clearTimeout(state.pendingPartialFlushTimer); - state.pendingPartialFlushTimer = undefined; - } const messageId = state.lastAssistantMessageId; if (messageId === undefined) return; if (state.accumulatedContentBlocks.length === 0) return; - const redacted = state.accumulatedContentBlocks.map((block) => { - if (block.type === "text") { - return { ...block, text: redactSecrets(block.text) }; - } - return block; - }); - const contentJson = JSON.stringify(redacted); + const built = buildPersistedAssistantContent( + state.accumulatedContentBlocks, + [], + ); + const contentJson = JSON.stringify(built); try { await runPipeline( @@ -463,36 +501,41 @@ async function flushAccumulatedContent( buildHandlerTurnContext(deps), DEFAULT_TIMEOUTS.persistence, ); - state.accumulatedCharsSinceLastFlush = 0; } catch (err) { deps.rlog.warn( { err, messageId }, - "B6 partial flush of accumulated assistant content failed; finalize at message_complete will recover", + "partial flush of accumulated assistant content failed; finalize at message_complete will recover", ); } } /** - * Either flush immediately (size gate tripped) or schedule a debounced - * flush (time gate). First-fire wins — if a timer is already pending - * and the size gate has not tripped, this is a no-op. + * Schedule a debounced partial flush. First-scheduled wins — if a timer + * is already pending, this is a no-op (the next flush will pick up any + * text appended in the meantime). Called from `handleTextDelta` after + * appending the drained delta. * - * Called from `handleTextDelta` after appending the drained delta. - * `handleToolUse` flushes eagerly via `flushAccumulatedContent` instead - * of going through this gate. + * `tool_use` events are intentionally not flush triggers: the + * `AgentLoop.run` emit order is `message_complete` THEN `tool_use`, so + * any flush from a `tool_use` handler would land after finalize and + * overwrite the authoritative row. */ -function scheduleOrFireFlush( +function schedulePartialFlush( state: EventHandlerState, deps: EventHandlerDeps, ): void { - if (state.accumulatedCharsSinceLastFlush >= PARTIAL_PERSIST_SIZE_THRESHOLD) { - void flushAccumulatedContent(state, deps); - return; - } if (state.pendingPartialFlushTimer !== undefined) return; state.pendingPartialFlushTimer = setTimeout(() => { state.pendingPartialFlushTimer = undefined; - void flushAccumulatedContent(state, deps); + const flushPromise = flushAccumulatedContent(state, deps); + state.pendingPartialFlushPromise = flushPromise; + // Clear the promise handle once the flush settles so the next + // scheduling cycle starts fresh. + void flushPromise.finally(() => { + if (state.pendingPartialFlushPromise === flushPromise) { + state.pendingPartialFlushPromise = undefined; + } + }); }, PARTIAL_PERSIST_DEBOUNCE_MS); } @@ -739,7 +782,7 @@ export async function handleLlmCallStarted( )) as PersistReserveResult; state.lastAssistantMessageId = reserveResult.message.id; state.assistantRowAwaitingFinalization = true; - // B6: fresh row → fresh accumulator. If an earlier (failed) LLM call + // Fresh row → fresh accumulator. If an earlier (failed) LLM call // within the same run left partial state behind, the // `assistantRowAwaitingFinalization` cleanup above already deleted // the orphan row, so the accumulator content would point at a @@ -783,12 +826,12 @@ function handleTextDelta( messageId: state.lastAssistantMessageId, }); if (deps.shouldGenerateTitle) state.firstAssistantText += drained.emitText; - // B6: mirror the drained delta into the partial-persist accumulator - // (the wire's drained text, not the raw event.text, so a refresh - // mid-turn sees the same content the user was watching live — raw - // directive markup never reaches the persisted snapshot). + // Mirror the drained delta (not raw event.text) into the + // partial-persist accumulator so a refresh mid-turn sees the same + // content the user was watching live — raw directive markup never + // reaches the persisted snapshot. appendTextToAccumulator(state, drained.emitText); - scheduleOrFireFlush(state, deps); + schedulePartialFlush(state, deps); } } @@ -857,19 +900,12 @@ export function handleToolUse( toolUseId: event.id, messageId: state.lastAssistantMessageId, }); - // B6: tool_use blocks are atomic (the AgentEvent only emits once the - // provider has finished streaming the input), so push the full block - // onto the accumulator and flush eagerly. No debounce — tool blocks - // are infrequent compared to text deltas and a refresh between - // tool_use_start and the corresponding tool_result is the - // most-visible mid-turn refresh case. - state.accumulatedContentBlocks.push({ - type: "tool_use", - id: event.id, - name: event.name, - input: event.input, - }); - void flushAccumulatedContent(state, deps); + // No partial-persist flush from here: `AgentLoop.run` emits + // `tool_use` strictly AFTER `message_complete`, so any flush from + // this handler would land after the authoritative final + // `updateContent` and overwrite the finalized row. The tool_use + // block lands in the finalized content via `event.message.content` + // at `handleMessageComplete`. } export function handleToolUsePreviewStart( @@ -1364,19 +1400,28 @@ export async function handleMessageComplete( // Reset per-turn tool tracking for the new turn. state.currentTurnToolUseIds = []; - // B6: cancel any pending debounced partial flush before the - // authoritative `updateContent` below. Without this, a timer that - // fires between this line and the final pipeline call could - // double-write the row (idempotent in content but wastes a write) - // or worse, race ahead of the indexer/projector read and serve a - // stale snapshot. The accumulator itself is left untouched until - // after finalize — the orchestrator's downstream paths don't read - // it, but a defensive reset on the next turn's - // `handleLlmCallStarted` keeps state hygienic regardless. + // Cancel any pending debounced partial flush and await an already + // in-flight one before the authoritative `updateContent` below. + // Without the timer-clear, a timer that fires during this handler + // could double-write (idempotent in content but wastes a write) or + // race ahead of the indexer/projector and serve a stale snapshot. + // Without the await, a partial pipeline call that was dispatched a + // moment before this handler can settle AFTER the final write and + // overwrite the authoritative row. if (state.pendingPartialFlushTimer !== undefined) { clearTimeout(state.pendingPartialFlushTimer); state.pendingPartialFlushTimer = undefined; } + if (state.pendingPartialFlushPromise !== undefined) { + try { + await state.pendingPartialFlushPromise; + } catch { + // The partial flush swallows its own pipeline errors via + // `rlog.warn`; the `try`/`catch` here is defensive against + // future changes that might surface them. + } + state.pendingPartialFlushPromise = undefined; + } // Flush any remaining directive display buffer if (state.pendingDirectiveDisplayBuffer.length > 0) { @@ -1448,13 +1493,14 @@ export async function handleMessageComplete( state.pendingToolResults.clear(); } - // Clean assistant content and accumulate directives - const { - cleanedContent, - directives: msgDirectives, - warnings: msgWarnings, - } = cleanAssistantContent(event.message.content); - const cleanedBlocks = cleanedContent as ContentBlock[]; + // Accumulate directives + warnings from the assistant content for + // downstream attachment processing. `cleanAssistantContent` is also + // called inside {@link buildPersistedAssistantContent} below; running + // it here separately is the cheapest way to keep the directive + // side-effects local to this handler while letting the shared helper + // own the persisted-content shape. + const { directives: msgDirectives, warnings: msgWarnings } = + cleanAssistantContent(event.message.content); state.accumulatedDirectives.push(...msgDirectives); state.directiveWarnings.push(...msgWarnings); if (msgDirectives.length > 0) { @@ -1476,31 +1522,14 @@ export async function handleMessageComplete( // are applied in handleToolResult after all tools for the turn complete, // then the persisted message is updated via updateMessageContent. - // Build content with UI surfaces - const contentWithSurfaces: ContentBlock[] = [...cleanedBlocks]; - for (const surface of deps.ctx.currentTurnSurfaces) { - contentWithSurfaces.push({ - type: "ui_surface", - surfaceId: surface.surfaceId, - surfaceType: surface.surfaceType, - title: surface.title, - data: surface.data, - actions: surface.actions, - display: surface.display, - ...(surface.persistent ? { persistent: true } : {}), - } as unknown as ContentBlock); - } - - // Redact known-pattern secrets from assistant text blocks before they are - // written to durable storage. Non-text blocks (images, UI surfaces) pass - // through unchanged. The live model history retains the original values. - const contentForPersistence = contentWithSurfaces.map((block) => { - if (block.type === "text") { - const tb = block as Extract; - return { ...tb, text: redactSecrets(tb.text) }; - } - return block; - }); + // Build the canonical persisted content (cleaned + surfaces + + // redacted) via the shared helper. The partial-persist flush uses + // the same helper with `surfaces=[]` so a mid-turn snapshot lands in + // the same shape as the finalize. + const contentForPersistence = buildPersistedAssistantContent( + event.message.content as ContentBlock[], + deps.ctx.currentTurnSurfaces, + ); // The row was reserved at `llm_call_started` (with channel metadata // stamped at that point) and `state.lastAssistantMessageId` carries its @@ -1652,8 +1681,10 @@ export async function handleMessageComplete( deps.ctx.currentTurnSurfaces = []; - // Emit trace event - const charCount = cleanedBlocks + // Emit trace event. Char count is computed from the cleaned + + // redacted text blocks (UI surface blocks filtered out via the + // type guard) — same shape as what was just persisted. + const charCount = contentForPersistence .filter( (b): b is Extract => b.type === "text", ) From 6b970043eebcc7f1052918c52569a52fddafb042 Mon Sep 17 00:00:00 2001 From: "vellum-apollo-bot[bot]" <242025090+vellum-apollo-bot[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 22:10:10 +0000 Subject: [PATCH 4/5] refactor(agent-loop): host running content on AgentLoopConversationContext Address review feedback round 2 on #32602: - Move the partial-persist accumulator from EventHandlerState to AgentLoopConversationContext.currentMessageContent. The running view now lives at the turn level alongside currentTurnSurfaces, making it the single source of truth for in-flight message content. Handler becomes a writer; flushAccumulatedContent reads from ctx. - Export AssistantSurface as a named interface from conversation-agent-loop.ts so callers can import it directly instead of infer-extracting from the context type. currentTurnSurfaces is now typed as AssistantSurface[]. - Bump partial-persist debounce from 250ms to 1000ms. - Compress doc-strings to one-liners on the tunable block, EventHandlerState fields, helpers, and the build pipeline. - Delete the explanatory comment in handleToolUse (deleted handler body is self-explanatory after the eager flush removal). Test mocks updated; debounce-wait increased to 1100ms. 70/70 partial-persist tests pass; typecheck clean. All 20 tests that touch currentTurnSurfaces pass individually. --- .../__tests__/conversation-agent-loop.test.ts | 11 +- .../conversation-agent-loop-handlers.ts | 180 ++++-------------- .../src/daemon/conversation-agent-loop.ts | 42 ++-- assistant/src/daemon/conversation.ts | 24 +-- 4 files changed, 75 insertions(+), 182 deletions(-) diff --git a/assistant/src/__tests__/conversation-agent-loop.test.ts b/assistant/src/__tests__/conversation-agent-loop.test.ts index 0c33bd1173d..aadd161a27a 100644 --- a/assistant/src/__tests__/conversation-agent-loop.test.ts +++ b/assistant/src/__tests__/conversation-agent-loop.test.ts @@ -613,6 +613,7 @@ function makeCtx( pendingSurfaceActions: new Map(), surfaceActionRequestIds: new Set(), currentTurnSurfaces: [], + currentMessageContent: [], workingDir: "/tmp", workspaceTopLevelContext: null, @@ -3482,7 +3483,7 @@ describe("session-agent-loop", () => { onEvent({ type: "text_delta", text: "Hello, " }); onEvent({ type: "text_delta", text: "world." }); // Wait long enough for the 250ms debounce to fire. - await new Promise((resolve) => setTimeout(resolve, 350)); + await new Promise((resolve) => setTimeout(resolve, 1100)); await onEvent({ type: "message_complete", message: { @@ -3642,7 +3643,7 @@ describe("session-agent-loop", () => { }); // Wait past the original debounce window to prove a late timer // does NOT fire a stray partial flush. - await new Promise((resolve) => setTimeout(resolve, 350)); + await new Promise((resolve) => setTimeout(resolve, 1100)); return [ ...messages, { @@ -3676,7 +3677,7 @@ describe("session-agent-loop", () => { onEvent({ type: "text_delta", text: "hello world" }); // Wait past the 250ms debounce so the partial flush definitely // lands BEFORE message_complete fires. - await new Promise((resolve) => setTimeout(resolve, 350)); + await new Promise((resolve) => setTimeout(resolve, 1100)); // Snapshot the indexer/projector call counts AFTER the partial // flush has run but BEFORE message_complete. They must be zero. const indexerCallsBeforeComplete = @@ -3752,7 +3753,7 @@ describe("session-agent-loop", () => { await onEvent({ type: "llm_call_started" }); onEvent({ type: "text_delta", text: payload }); // Wait past the 250ms debounce so the partial flush lands. - await new Promise((resolve) => setTimeout(resolve, 350)); + await new Promise((resolve) => setTimeout(resolve, 1100)); await onEvent({ type: "message_complete", message: { @@ -3805,7 +3806,7 @@ describe("session-agent-loop", () => { // A debounced delta lands a partial flush BEFORE the provider // error fires. onEvent({ type: "text_delta", text: "hello world" }); - await new Promise((resolve) => setTimeout(resolve, 350)); + await new Promise((resolve) => setTimeout(resolve, 1100)); onEvent({ type: "provider_error", error: new Error("upstream 500"), diff --git a/assistant/src/daemon/conversation-agent-loop-handlers.ts b/assistant/src/daemon/conversation-agent-loop-handlers.ts index 09fbb1b084a..97c6f70958e 100644 --- a/assistant/src/daemon/conversation-agent-loop-handlers.ts +++ b/assistant/src/daemon/conversation-agent-loop-handlers.ts @@ -67,7 +67,10 @@ import { cleanAssistantContent, drainDirectiveDisplayBuffer, } from "./assistant-attachments.js"; -import type { AgentLoopConversationContext } from "./conversation-agent-loop.js"; +import type { + AgentLoopConversationContext, + AssistantSurface, +} from "./conversation-agent-loop.js"; import { buildConversationErrorMessage, classifyConversationError, @@ -91,18 +94,9 @@ import type { const log = getLogger("agent-loop-handlers"); // ── Partial-persistence tunables ───────────────────────────────────── -// -// `handleTextDelta` accumulates streamed text into -// `state.accumulatedContentBlocks` and flushes the row via -// `updateMessageContent` on a debounce timer: at most one flush per -// `PARTIAL_PERSIST_DEBOUNCE_MS` (~4 writes/sec on a steady stream). -// Cheap enough that a refresh mid-turn never lands more than ~250ms -// behind the wire. -// -// Indexer + projector still fire ONLY at `handleMessageComplete`; the -// row's `content` may temporarily reflect a partial assistant turn but -// is never indexed mid-stream. -const PARTIAL_PERSIST_DEBOUNCE_MS = 250; +// Debounce for mid-turn `updateContent` writes from text deltas. +// Indexer + projector still fire ONLY at `handleMessageComplete`. +const PARTIAL_PERSIST_DEBOUNCE_MS = 1000; /** * Build a {@link TurnContext} from the handler's deps for pipeline logging @@ -253,39 +247,9 @@ export interface EventHandlerState { readonly serverToolStartedAt: Map; /** Original input from server_tool_start, keyed by tool_use_id, so the complete handler can read the query. */ readonly serverToolInputs: Map>; - /** - * Partial-persistence accumulator. Mirrors the in-progress assistant - * text emitted to the wire since the current row's - * `handleLlmCallStarted`, so a debounced flush can write the partial - * snapshot to the row's `content` column without waiting for - * `message_complete`. Reset at `handleLlmCallStarted` for each new - * assistant row reservation (a tool-flush + second LLM call within - * the same run reserves a fresh row and starts a fresh accumulator). - * Final flush in `handleMessageComplete` uses the authoritative - * `event.message.content` instead — this accumulator is for - * *mid-turn* snapshots only. - * - * Only text blocks are accumulated. `tool_use` events arrive AFTER - * `message_complete` finalizes the row, so mirroring them mid-turn - * would have no effect; `thinking` blocks are intentionally omitted - * to keep the persisted snapshot in lockstep with what `cleanAssistantContent` - * produces at finalize (which the partial flush also runs through). - */ - accumulatedContentBlocks: ContentBlock[]; - /** - * Active debounce timer for partial persistence. `undefined` when - * idle (no text since last flush). Cleared at the start of - * `handleMessageComplete` so the final authoritative flush never - * races a debounced partial write. - */ + /** Active debounce timer for partial persistence; `undefined` when idle. */ pendingPartialFlushTimer: ReturnType | undefined; - /** - * In-flight partial flush write (the most recently dispatched - * `flushAccumulatedContent` pipeline call). `handleMessageComplete` - * awaits this before its final `updateContent` so a partial write - * that started just before finalize can never overwrite the - * authoritative content. - */ + /** In-flight partial flush write awaited at finalize to avoid overwrite races. */ pendingPartialFlushPromise: Promise | undefined; } @@ -347,7 +311,6 @@ export function createEventHandlerState(): EventHandlerState { turnStartedAt: Date.now(), serverToolStartedAt: new Map(), serverToolInputs: new Map(), - accumulatedContentBlocks: [], pendingPartialFlushTimer: undefined, pendingPartialFlushPromise: undefined, }; @@ -355,40 +318,7 @@ export function createEventHandlerState(): EventHandlerState { // ── Partial-persistence helpers ────────────────────────────────────── -/** - * Type of a single UI surface that may be merged into persisted assistant - * content. Mirrors the inline structural type on - * `AgentLoopConversationContext.currentTurnSurfaces` — sharing it via - * this alias keeps {@link buildPersistedAssistantContent} usable in - * tests without routing them through the full context. - */ -type AssistantSurface = - AgentLoopConversationContext["currentTurnSurfaces"] extends ReadonlyArray< - infer S - > - ? S - : never; - -/** - * Build the canonical persisted-content array for an assistant row, - * applying the same pipeline both `handleMessageComplete` (the - * authoritative final write) and partial flushes use: - * - * 1. {@link cleanAssistantContent} — drops Anthropic placeholder - * sentinel text blocks and strips attachment directives from - * text. Directive results (`directives`, `warnings`) are NOT - * returned here — handleMessageComplete owns the canonical - * accumulation path for those. - * 2. UI surface blocks (if any) are appended. - * 3. Text blocks are run through {@link redactSecrets}. - * - * Used in two places: - * • final flush in `handleMessageComplete` with `event.message.content` - * and `deps.ctx.currentTurnSurfaces`. - * • partial flush in `flushAccumulatedContent` with the in-memory - * accumulator and an empty surfaces array (surfaces are produced - * by tool execution, which runs after `message_complete`). - */ +/** Canonical persisted-content build: clean → append surfaces → redact. */ function buildPersistedAssistantContent( rawBlocks: readonly ContentBlock[], surfaces: readonly AssistantSurface[], @@ -417,73 +347,44 @@ function buildPersistedAssistantContent( }); } -/** - * Append a chunk of streamed text to the partial-persist accumulator. - * - * If the tail block is already a `text` block, the chunk is appended in - * place (the producer of `event.message.content` at finalize fuses - * consecutive deltas into a single text block, so we mirror that - * shape). Otherwise a new text block is pushed. - */ -function appendTextToAccumulator( - state: EventHandlerState, +/** Append a streamed text chunk to `ctx.currentMessageContent`, fusing into tail text block. */ +function appendTextToCurrentMessage( + ctx: AgentLoopConversationContext, text: string, ): void { if (text.length === 0) return; - const tail = state.accumulatedContentBlocks.at(-1); + const tail = ctx.currentMessageContent.at(-1); if (tail && tail.type === "text") { tail.text = tail.text + text; } else { - state.accumulatedContentBlocks.push({ type: "text", text }); + ctx.currentMessageContent.push({ type: "text", text }); } } -/** - * Reset partial-persist accumulator state. Called from - * `handleLlmCallStarted` after reserving a fresh assistant row (each - * row has its own accumulator) and after the final flush in - * `handleMessageComplete`. Idempotent and safe to call repeatedly. - */ -function resetPartialPersistAccumulator(state: EventHandlerState): void { +/** Reset partial-persist accumulator and any pending flush state. Idempotent. */ +function resetPartialPersistAccumulator( + state: EventHandlerState, + ctx: AgentLoopConversationContext, +): void { if (state.pendingPartialFlushTimer !== undefined) { clearTimeout(state.pendingPartialFlushTimer); state.pendingPartialFlushTimer = undefined; } - state.accumulatedContentBlocks = []; + ctx.currentMessageContent = []; state.pendingPartialFlushPromise = undefined; } -/** - * Flush the partial-persist accumulator to the assistant row via - * `updateContent`. Reuses {@link buildPersistedAssistantContent} so the - * partial snapshot lands in the same shape `handleMessageComplete` - * would produce (with `surfaces=[]` — surfaces emerge from tool - * execution which runs after `message_complete`). - * - * No-ops if: - * • there's no reserved assistant row id (pre-`handleLlmCallStarted` - * or post-cleanup), or - * • the accumulator is empty (nothing meaningful has streamed yet). - * - * The pending-promise handle is published on `state` so - * `handleMessageComplete` can await it before the authoritative final - * write, serializing the two `updateContent` calls and preventing a - * partial write from overwriting the finalized row. - * - * Failures are logged but never thrown — a missed partial flush is - * always recoverable by either the next debounce tick or the final - * authoritative flush at `handleMessageComplete`. - */ +/** Flush `ctx.currentMessageContent` to the row via the persistence pipeline. */ async function flushAccumulatedContent( state: EventHandlerState, deps: EventHandlerDeps, ): Promise { const messageId = state.lastAssistantMessageId; if (messageId === undefined) return; - if (state.accumulatedContentBlocks.length === 0) return; + if (deps.ctx.currentMessageContent.length === 0) return; const built = buildPersistedAssistantContent( - state.accumulatedContentBlocks, + deps.ctx.currentMessageContent, [], ); const contentJson = JSON.stringify(built); @@ -509,17 +410,7 @@ async function flushAccumulatedContent( } } -/** - * Schedule a debounced partial flush. First-scheduled wins — if a timer - * is already pending, this is a no-op (the next flush will pick up any - * text appended in the meantime). Called from `handleTextDelta` after - * appending the drained delta. - * - * `tool_use` events are intentionally not flush triggers: the - * `AgentLoop.run` emit order is `message_complete` THEN `tool_use`, so - * any flush from a `tool_use` handler would land after finalize and - * overwrite the authoritative row. - */ +/** Schedule a debounced partial flush. First-scheduled wins; no-op when timer pending. */ function schedulePartialFlush( state: EventHandlerState, deps: EventHandlerDeps, @@ -529,8 +420,6 @@ function schedulePartialFlush( state.pendingPartialFlushTimer = undefined; const flushPromise = flushAccumulatedContent(state, deps); state.pendingPartialFlushPromise = flushPromise; - // Clear the promise handle once the flush settles so the next - // scheduling cycle starts fresh. void flushPromise.finally(() => { if (state.pendingPartialFlushPromise === flushPromise) { state.pendingPartialFlushPromise = undefined; @@ -787,7 +676,7 @@ export async function handleLlmCallStarted( // `assistantRowAwaitingFinalization` cleanup above already deleted // the orphan row, so the accumulator content would point at a // non-existent id. Reset here so the new row starts from zero. - resetPartialPersistAccumulator(state); + resetPartialPersistAccumulator(state, deps.ctx); deps.onEvent({ type: "assistant_turn_start", messageId: reserveResult.message.id, @@ -826,11 +715,9 @@ function handleTextDelta( messageId: state.lastAssistantMessageId, }); if (deps.shouldGenerateTitle) state.firstAssistantText += drained.emitText; - // Mirror the drained delta (not raw event.text) into the - // partial-persist accumulator so a refresh mid-turn sees the same - // content the user was watching live — raw directive markup never - // reaches the persisted snapshot. - appendTextToAccumulator(state, drained.emitText); + // Mirror the drained delta into ctx.currentMessageContent so partial + // flushes mid-turn see the same content the user is watching live. + appendTextToCurrentMessage(deps.ctx, drained.emitText); schedulePartialFlush(state, deps); } } @@ -900,12 +787,6 @@ export function handleToolUse( toolUseId: event.id, messageId: state.lastAssistantMessageId, }); - // No partial-persist flush from here: `AgentLoop.run` emits - // `tool_use` strictly AFTER `message_complete`, so any flush from - // this handler would land after the authoritative final - // `updateContent` and overwrite the finalized row. The tool_use - // block lands in the finalized content via `event.message.content` - // at `handleMessageComplete`. } export function handleToolUsePreviewStart( @@ -1558,6 +1439,9 @@ export async function handleMessageComplete( DEFAULT_TIMEOUTS.persistence, ); state.assistantRowAwaitingFinalization = false; + // Reset the partial-persist mirror so subsequent calls in this turn + // start with an empty running view. + deps.ctx.currentMessageContent = []; // ── Indexing + attention projection (restored from the pre-B3 `add` path) ── // `reserveMessage` + `updateMessageContent` are CRUD-only: they don't run diff --git a/assistant/src/daemon/conversation-agent-loop.ts b/assistant/src/daemon/conversation-agent-loop.ts index ad7a5be065f..216a2969611 100644 --- a/assistant/src/daemon/conversation-agent-loop.ts +++ b/assistant/src/daemon/conversation-agent-loop.ts @@ -484,6 +484,26 @@ function buildPluginTurnContext( // ── Context Interface ──────────────────────────────────────────────── +/** + * Per-surface entry tracked on the current turn. Inline shape kept stable so + * routes and persistence helpers can consume it via a named import instead of + * `infer`-extracting from {@link AgentLoopConversationContext}. + */ +export interface AssistantSurface { + surfaceId: string; + surfaceType: SurfaceType; + title?: string; + data: SurfaceData; + actions?: Array<{ + id: string; + label: string; + style?: string; + data?: Record; + }>; + display?: string; + persistent?: boolean; +} + export interface AgentLoopConversationContext { readonly conversationId: string; messages: Message[]; @@ -533,20 +553,14 @@ export interface AgentLoopConversationContext { pendingSurfaceActions: Map; surfaceActionRequestIds: Set; approvedViaPromptThisTurn?: boolean; - currentTurnSurfaces: Array<{ - surfaceId: string; - surfaceType: SurfaceType; - title?: string; - data: SurfaceData; - actions?: Array<{ - id: string; - label: string; - style?: string; - data?: Record; - }>; - display?: string; - persistent?: boolean; - }>; + currentTurnSurfaces: AssistantSurface[]; + /** + * Running mirror of the in-flight assistant message's content, used by + * partial-persistence flushes (see `conversation-agent-loop-handlers.ts`). + * Mid-turn snapshot of what `event.message.content` will be at + * `message_complete`. Reset on llm_call_started and finalize. + */ + currentMessageContent: ContentBlock[]; workingDir: string; workspaceTopLevelContext: string | null; diff --git a/assistant/src/daemon/conversation.ts b/assistant/src/daemon/conversation.ts index 1f4b2ac0d39..4b45526d316 100644 --- a/assistant/src/daemon/conversation.ts +++ b/assistant/src/daemon/conversation.ts @@ -61,7 +61,7 @@ import { PermissionPrompter } from "../permissions/prompter.js"; import { SecretPrompter } from "../permissions/secret-prompter.js"; import type { UserDecision } from "../permissions/types.js"; import { buildSystemPrompt } from "../prompts/system-prompt.js"; -import type { Message } from "../providers/types.js"; +import type { ContentBlock, Message } from "../providers/types.js"; import type { Provider } from "../providers/types.js"; import type { TrustClass } from "../runtime/actor-trust-resolver.js"; import { broadcastMessage } from "../runtime/assistant-event-hub.js"; @@ -74,6 +74,7 @@ import type { OnboardingContext } from "../types/onboarding-context.js"; import type { AbortReason } from "../util/abort-reasons.js"; import { getLogger } from "../util/logger.js"; import type { AssistantAttachmentDraft } from "./assistant-attachments.js"; +import type { AssistantSurface } from "./conversation-agent-loop.js"; import { applyCompactionResult, runAgentLoopImpl, @@ -327,20 +328,13 @@ export class Conversation { ReturnType >(); /** @internal */ withSurface = createSurfaceMutex(); - /** @internal */ currentTurnSurfaces: Array<{ - surfaceId: string; - surfaceType: SurfaceType; - title?: string; - data: SurfaceData; - actions?: Array<{ - id: string; - label: string; - style?: string; - data?: Record; - }>; - display?: string; - persistent?: boolean; - }> = []; + /** @internal */ currentTurnSurfaces: AssistantSurface[] = []; + /** + * Running mirror of the in-flight assistant message's content (see + * {@link AgentLoopConversationContext.currentMessageContent}). + * @internal + */ + currentMessageContent: ContentBlock[] = []; /** @internal */ workspaceTopLevelContext: string | null = null; /** @internal */ workspaceTopLevelDirty = true; /** From 0a526018cbc4fc4a531b7b495916edb67e55f74f Mon Sep 17 00:00:00 2001 From: Apollo Bot Date: Sat, 30 May 2026 12:13:14 +0000 Subject: [PATCH 5/5] refactor(agent-loop): host running content on EventHandlerState MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidates currentMessageContent to a single home — per-turn EventHandlerState — eliminating the redundant field on Conversation (only needed for structural compatibility with AgentLoopConversationContext). The running mirror is per-turn state with no external readers, so it belongs alongside other per-turn agent-loop state (pendingPartialFlushTimer, toolUseIdToName, etc) rather than on the long-lived context. Per PR #32602 review feedback. --- .../__tests__/conversation-agent-loop.test.ts | 5 +-- .../conversation-agent-loop-handlers.ts | 35 +++++++++---------- .../src/daemon/conversation-agent-loop.ts | 7 ---- assistant/src/daemon/conversation.ts | 8 +---- 4 files changed, 18 insertions(+), 37 deletions(-) diff --git a/assistant/src/__tests__/conversation-agent-loop.test.ts b/assistant/src/__tests__/conversation-agent-loop.test.ts index aadd161a27a..32823fa0568 100644 --- a/assistant/src/__tests__/conversation-agent-loop.test.ts +++ b/assistant/src/__tests__/conversation-agent-loop.test.ts @@ -613,7 +613,6 @@ function makeCtx( pendingSurfaceActions: new Map(), surfaceActionRequestIds: new Set(), currentTurnSurfaces: [], - currentMessageContent: [], workingDir: "/tmp", workspaceTopLevelContext: null, @@ -3709,9 +3708,7 @@ describe("session-agent-loop", () => { ...messages, { role: "assistant" as const, - content: [ - { type: "text", text: "hello world" }, - ] as ContentBlock[], + content: [{ type: "text", text: "hello world" }] as ContentBlock[], }, ]; }; diff --git a/assistant/src/daemon/conversation-agent-loop-handlers.ts b/assistant/src/daemon/conversation-agent-loop-handlers.ts index 97c6f70958e..d28ed866a60 100644 --- a/assistant/src/daemon/conversation-agent-loop-handlers.ts +++ b/assistant/src/daemon/conversation-agent-loop-handlers.ts @@ -251,6 +251,8 @@ export interface EventHandlerState { pendingPartialFlushTimer: ReturnType | undefined; /** In-flight partial flush write awaited at finalize to avoid overwrite races. */ pendingPartialFlushPromise: Promise | undefined; + /** Running mirror of the in-flight assistant message's content. */ + currentMessageContent: ContentBlock[]; } /** Immutable context shared across event handlers within a single agent loop run. */ @@ -313,6 +315,7 @@ export function createEventHandlerState(): EventHandlerState { serverToolInputs: new Map(), pendingPartialFlushTimer: undefined, pendingPartialFlushPromise: undefined, + currentMessageContent: [], }; } @@ -347,46 +350,40 @@ function buildPersistedAssistantContent( }); } -/** Append a streamed text chunk to `ctx.currentMessageContent`, fusing into tail text block. */ +/** Append a streamed text chunk to `state.currentMessageContent`, fusing into tail text block. */ function appendTextToCurrentMessage( - ctx: AgentLoopConversationContext, + state: EventHandlerState, text: string, ): void { if (text.length === 0) return; - const tail = ctx.currentMessageContent.at(-1); + const tail = state.currentMessageContent.at(-1); if (tail && tail.type === "text") { tail.text = tail.text + text; } else { - ctx.currentMessageContent.push({ type: "text", text }); + state.currentMessageContent.push({ type: "text", text }); } } /** Reset partial-persist accumulator and any pending flush state. Idempotent. */ -function resetPartialPersistAccumulator( - state: EventHandlerState, - ctx: AgentLoopConversationContext, -): void { +function resetPartialPersistAccumulator(state: EventHandlerState): void { if (state.pendingPartialFlushTimer !== undefined) { clearTimeout(state.pendingPartialFlushTimer); state.pendingPartialFlushTimer = undefined; } - ctx.currentMessageContent = []; + state.currentMessageContent = []; state.pendingPartialFlushPromise = undefined; } -/** Flush `ctx.currentMessageContent` to the row via the persistence pipeline. */ +/** Flush `state.currentMessageContent` to the row via the persistence pipeline. */ async function flushAccumulatedContent( state: EventHandlerState, deps: EventHandlerDeps, ): Promise { const messageId = state.lastAssistantMessageId; if (messageId === undefined) return; - if (deps.ctx.currentMessageContent.length === 0) return; + if (state.currentMessageContent.length === 0) return; - const built = buildPersistedAssistantContent( - deps.ctx.currentMessageContent, - [], - ); + const built = buildPersistedAssistantContent(state.currentMessageContent, []); const contentJson = JSON.stringify(built); try { @@ -676,7 +673,7 @@ export async function handleLlmCallStarted( // `assistantRowAwaitingFinalization` cleanup above already deleted // the orphan row, so the accumulator content would point at a // non-existent id. Reset here so the new row starts from zero. - resetPartialPersistAccumulator(state, deps.ctx); + resetPartialPersistAccumulator(state); deps.onEvent({ type: "assistant_turn_start", messageId: reserveResult.message.id, @@ -715,9 +712,9 @@ function handleTextDelta( messageId: state.lastAssistantMessageId, }); if (deps.shouldGenerateTitle) state.firstAssistantText += drained.emitText; - // Mirror the drained delta into ctx.currentMessageContent so partial + // Mirror the drained delta into state.currentMessageContent so partial // flushes mid-turn see the same content the user is watching live. - appendTextToCurrentMessage(deps.ctx, drained.emitText); + appendTextToCurrentMessage(state, drained.emitText); schedulePartialFlush(state, deps); } } @@ -1441,7 +1438,7 @@ export async function handleMessageComplete( state.assistantRowAwaitingFinalization = false; // Reset the partial-persist mirror so subsequent calls in this turn // start with an empty running view. - deps.ctx.currentMessageContent = []; + state.currentMessageContent = []; // ── Indexing + attention projection (restored from the pre-B3 `add` path) ── // `reserveMessage` + `updateMessageContent` are CRUD-only: they don't run diff --git a/assistant/src/daemon/conversation-agent-loop.ts b/assistant/src/daemon/conversation-agent-loop.ts index 216a2969611..9982108b188 100644 --- a/assistant/src/daemon/conversation-agent-loop.ts +++ b/assistant/src/daemon/conversation-agent-loop.ts @@ -554,13 +554,6 @@ export interface AgentLoopConversationContext { surfaceActionRequestIds: Set; approvedViaPromptThisTurn?: boolean; currentTurnSurfaces: AssistantSurface[]; - /** - * Running mirror of the in-flight assistant message's content, used by - * partial-persistence flushes (see `conversation-agent-loop-handlers.ts`). - * Mid-turn snapshot of what `event.message.content` will be at - * `message_complete`. Reset on llm_call_started and finalize. - */ - currentMessageContent: ContentBlock[]; workingDir: string; workspaceTopLevelContext: string | null; diff --git a/assistant/src/daemon/conversation.ts b/assistant/src/daemon/conversation.ts index 4b45526d316..3fb8e59cbe0 100644 --- a/assistant/src/daemon/conversation.ts +++ b/assistant/src/daemon/conversation.ts @@ -61,7 +61,7 @@ import { PermissionPrompter } from "../permissions/prompter.js"; import { SecretPrompter } from "../permissions/secret-prompter.js"; import type { UserDecision } from "../permissions/types.js"; import { buildSystemPrompt } from "../prompts/system-prompt.js"; -import type { ContentBlock, Message } from "../providers/types.js"; +import type { Message } from "../providers/types.js"; import type { Provider } from "../providers/types.js"; import type { TrustClass } from "../runtime/actor-trust-resolver.js"; import { broadcastMessage } from "../runtime/assistant-event-hub.js"; @@ -329,12 +329,6 @@ export class Conversation { >(); /** @internal */ withSurface = createSurfaceMutex(); /** @internal */ currentTurnSurfaces: AssistantSurface[] = []; - /** - * Running mirror of the in-flight assistant message's content (see - * {@link AgentLoopConversationContext.currentMessageContent}). - * @internal - */ - currentMessageContent: ContentBlock[] = []; /** @internal */ workspaceTopLevelContext: string | null = null; /** @internal */ workspaceTopLevelDirty = true; /**