From 4fee4d7d86820b3c812421dbcc6ec3a05de52550 Mon Sep 17 00:00:00 2001 From: Dax Date: Sat, 15 Aug 2026 18:16:01 -0400 Subject: [PATCH] fix(core): batch streamed session deltas --- .../src/session/runner/publish-llm-event.ts | 96 +++++++++++-------- .../test/session-runner-tool-events.test.ts | 81 ++++++++++++++++ packages/core/test/session-runner.test.ts | 43 ++++++--- 3 files changed, 169 insertions(+), 51 deletions(-) diff --git a/packages/core/src/session/runner/publish-llm-event.ts b/packages/core/src/session/runner/publish-llm-event.ts index b4c84db9228..1527a0f961e 100644 --- a/packages/core/src/session/runner/publish-llm-event.ts +++ b/packages/core/src/session/runner/publish-llm-event.ts @@ -1,5 +1,5 @@ import { type LLMEvent, type ProviderMetadata, type ToolResultValue } from "@opencode-ai/ai" -import { Effect } from "effect" +import { Clock, Effect } from "effect" import { Bus } from "../../bus.js" import { Model } from "../../model.js" import { SessionEvent } from "../event.js" @@ -81,6 +81,7 @@ const hostedContent = (result: ToolResultValue): NonEmptyContent => { * and consumers fold by id/ordinal rather than global position. */ export const createLLMEventPublisher = (bus: Pick, input: Input) => { + const deltaBatchInterval = 100 const tools = new Map< string, { @@ -123,32 +124,52 @@ export const createLLMEventPublisher = (bus: Pick, inp const fragments = ( name: string, ended: (id: string, value: string, ordinal: number, state?: Record) => Effect.Effect, + delta?: (id: string, value: string, ordinal: number) => Effect.Effect, single = false, ) => { - const chunks = new Map< - string, - { readonly ordinal: number; readonly values: string[]; state?: Record } - >() + type Fragment = { + readonly ordinal: number + readonly values: string[] + pending: string + publishedAt?: number + state?: Record + } + const chunks = new Map() let nextOrdinal = 0 const start = (id: string, state?: Record) => Effect.suspend(() => { if (chunks.has(id)) return Effect.die(new Error(`Duplicate ${name} start: ${id}`)) if (single && chunks.size > 0) return Effect.die(new Error(`${name} start before end: ${id}`)) const ordinal = nextOrdinal++ - chunks.set(id, { ordinal, values: [], state }) + chunks.set(id, { ordinal, values: [], pending: "", state }) return Effect.succeed(ordinal) }) - const append = (id: string, value: string, state?: Record) => - Effect.suspend(() => { - const current = chunks.get(id) - if (!current) return Effect.die(new Error(`${name} delta before start: ${id}`)) - current.values.push(value) - if (state !== undefined) current.state = { ...current.state, ...state } - return Effect.succeed(current.ordinal) - }) + const publishDelta = Effect.fnUntraced(function* (id: string, force = false) { + if (!delta) return undefined + const current = chunks.get(id) + if (!current) return yield* Effect.die(new Error(`${name} delta before start: ${id}`)) + if (!current.pending) return undefined + const now = yield* Clock.currentTimeMillis + if (!force && current.publishedAt !== undefined && now - current.publishedAt < deltaBatchInterval) + return undefined + yield* delta(id, current.pending, current.ordinal) + current.pending = "" + current.publishedAt = now + return undefined + }) + const append = Effect.fnUntraced(function* (id: string, value: string, state?: Record) { + const current = chunks.get(id) + if (!current) return yield* Effect.die(new Error(`${name} delta before start: ${id}`)) + current.values.push(value) + if (delta) current.pending += value + if (state !== undefined) current.state = { ...current.state, ...state } + yield* publishDelta(id) + return current.ordinal + }) const end = Effect.fnUntraced(function* (id: string, state?: Record, value?: string) { const current = chunks.get(id) if (!current) return yield* Effect.die(new Error(`${name} end before start: ${id}`)) + yield* publishDelta(id, true) yield* ended( id, value ?? current.values.join(""), @@ -156,9 +177,10 @@ export const createLLMEventPublisher = (bus: Pick, inp state === undefined ? current.state : { ...current.state, ...state }, ) chunks.delete(id) + return undefined }) const flush = Effect.fnUntraced(function* () { - for (const id of chunks.keys()) yield* end(id) + for (const id of Array.from(chunks.keys())) yield* end(id) }) return { start, append, end, flush, has: (id: string) => chunks.has(id) } } @@ -175,6 +197,15 @@ export const createLLMEventPublisher = (bus: Pick, inp state, }) }), + (_textID, value, ordinal) => + Effect.gen(function* () { + yield* bus.publish(SessionEvent.Text.Delta, { + sessionID: input.sessionID, + assistantMessageID: yield* currentAssistantMessageID(), + ordinal, + delta: value, + }) + }), true, ) const reasoning = fragments( @@ -189,6 +220,15 @@ export const createLLMEventPublisher = (bus: Pick, inp state, }) }), + (_reasoningID, value, ordinal) => + Effect.gen(function* () { + yield* bus.publish(SessionEvent.Reasoning.Delta, { + sessionID: input.sessionID, + assistantMessageID: yield* currentAssistantMessageID(), + ordinal, + delta: value, + }) + }), true, ) const toolInput = fragments("tool input", (id, value) => @@ -351,13 +391,7 @@ export const createLLMEventPublisher = (bus: Pick, inp }) return case "text-delta": - const deltaTextOrdinal = yield* text.append(event.id, event.text, providerState(event.providerMetadata)) - yield* bus.publish(SessionEvent.Text.Delta, { - sessionID: input.sessionID, - assistantMessageID: yield* currentAssistantMessageID(), - ordinal: deltaTextOrdinal, - delta: event.text, - }) + yield* text.append(event.id, event.text, providerState(event.providerMetadata)) return case "text-end": yield* text.end(event.id, providerState(event.providerMetadata)) @@ -373,17 +407,7 @@ export const createLLMEventPublisher = (bus: Pick, inp }) return case "reasoning-delta": - const deltaReasoningOrdinal = yield* reasoning.append( - event.id, - event.text, - providerState(event.providerMetadata), - ) - yield* bus.publish(SessionEvent.Reasoning.Delta, { - sessionID: input.sessionID, - assistantMessageID: yield* currentAssistantMessageID(), - ordinal: deltaReasoningOrdinal, - delta: event.text, - }) + yield* reasoning.append(event.id, event.text, providerState(event.providerMetadata)) return case "reasoning-end": yield* reasoning.end(event.id, providerState(event.providerMetadata)) @@ -399,12 +423,6 @@ export const createLLMEventPublisher = (bus: Pick, inp return yield* Effect.die(new Error(`Tool input name changed for ${event.id}: ${tool.name} -> ${event.name}`)) if (!toolInput.has(event.id)) return yield* Effect.die(new Error(`Tool input delta after end: ${event.id}`)) yield* toolInput.append(event.id, event.text) - yield* bus.publish(SessionEvent.Tool.Input.Delta, { - sessionID: input.sessionID, - assistantMessageID: tool.assistantMessageID, - id: event.id, - delta: event.text, - }) return } case "tool-input-end": diff --git a/packages/core/test/session-runner-tool-events.test.ts b/packages/core/test/session-runner-tool-events.test.ts index 831f289e62f..ed6146465e7 100644 --- a/packages/core/test/session-runner-tool-events.test.ts +++ b/packages/core/test/session-runner-tool-events.test.ts @@ -13,6 +13,8 @@ import { Provider } from "@opencode-ai/core/provider" import { RelativePath } from "@opencode-ai/core/schema" import { Snapshot } from "@opencode-ai/core/snapshot" import { createLLMEventPublisher } from "@opencode-ai/core/session/runner/publish-llm-event" +import { it } from "./lib/effect" +import { TestClock } from "effect/testing" const sessionID = Session.ID.make("ses_tool_event_test") const base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB" @@ -201,6 +203,85 @@ test("reasoning state from start, empty delta, and end is merged", async () => { }) }) +it.effect("batches text deltas and flushes pending text before the terminal event", () => + Effect.gen(function* () { + const { published, publisher } = capture() + yield* Effect.forEach( + [ + LLMEvent.textStart({ id: "text" }), + LLMEvent.textDelta({ id: "text", text: "one" }), + LLMEvent.textDelta({ id: "text", text: " two" }), + LLMEvent.textDelta({ id: "text", text: " three" }), + ], + publisher.publish, + { discard: true }, + ) + + expect(published.filter((event) => event.type === "session.text.delta").map((event) => event.data)).toMatchObject([ + { delta: "one" }, + ]) + yield* TestClock.adjust("99 millis") + expect(published.filter((event) => event.type === "session.text.delta")).toHaveLength(1) + yield* TestClock.adjust("1 millis") + yield* publisher.publish(LLMEvent.textDelta({ id: "text", text: " four" })) + expect(published.filter((event) => event.type === "session.text.delta").map((event) => event.data)).toMatchObject([ + { delta: "one" }, + { delta: " two three four" }, + ]) + + yield* publisher.publish(LLMEvent.textDelta({ id: "text", text: " five" })) + yield* publisher.publish(LLMEvent.textEnd({ id: "text" })) + expect(published.slice(-2).map((event) => event.type)).toEqual(["session.text.delta", "session.text.ended.1"]) + expect(published.at(-2)?.data).toMatchObject({ delta: " five" }) + }), +) + +it.effect("batches reasoning deltas and flushes pending reasoning before the terminal event", () => + Effect.gen(function* () { + const { published, publisher } = capture() + yield* Effect.forEach( + [ + LLMEvent.reasoningStart({ id: "reasoning" }), + LLMEvent.reasoningDelta({ id: "reasoning", text: "one" }), + LLMEvent.reasoningDelta({ id: "reasoning", text: " two" }), + LLMEvent.reasoningDelta({ id: "reasoning", text: " three" }), + LLMEvent.reasoningEnd({ id: "reasoning" }), + ], + publisher.publish, + { discard: true }, + ) + + expect( + published.filter((event) => event.type === "session.reasoning.delta").map((event) => event.data), + ).toMatchObject([{ delta: "one" }, { delta: " two three" }]) + expect(published.slice(-2).map((event) => event.type)).toEqual([ + "session.reasoning.delta", + "session.reasoning.ended.1", + ]) + }), +) + +test("tool input deltas are accumulated without being published", async () => { + const { published, publisher } = capture() + await Effect.runPromise( + Effect.forEach( + [ + LLMEvent.toolInputStart({ id: "call", name: "read" }), + LLMEvent.toolInputDelta({ id: "call", name: "read", text: '{"path":' }), + LLMEvent.toolInputDelta({ id: "call", name: "read", text: '"file.txt"}' }), + LLMEvent.toolInputEnd({ id: "call", name: "read" }), + ], + publisher.publish, + { discard: true }, + ), + ) + + expect(published.some((event) => event.type === "session.tool.input.delta")).toBe(false) + expect(published.find((event) => event.type === "session.tool.input.ended.1")?.data).toMatchObject({ + text: '{"path":"file.txt"}', + }) +}) + test("provider-executed tool metadata is flattened using the route key", async () => { const { published, publisher } = capture("openai") await Effect.runPromise( diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index b3a3adb030e..36c12752946 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -686,7 +686,7 @@ const replaySessionProjection = (id: Session.ID) => type FragmentKind = "text" | "reasoning" | "tool input" type FragmentFixture = { - readonly delta: Event.Definition + readonly delta?: Event.Definition readonly completeEvents: LLMEvent[] readonly partialEvents: LLMEvent[] readonly expectedAssistant: unknown @@ -748,7 +748,6 @@ const fragmentFixture = (kind: FragmentKind, id: string, chunks: readonly string ] const expectedContent = { type: "tool", id, state: { status: "streaming", input: text } } return { - delta: SessionEvent.Tool.Input.Delta, partialEvents, completeEvents: [...partialEvents, LLMEvent.toolInputEnd({ id, name: "echo" })], expectedAssistant: { type: "assistant", content: [expectedContent] }, @@ -767,20 +766,37 @@ const verifyEphemeralDeltas = (kind: FragmentKind) => const expectedContext = [{ type: "user", text: prompt }, fixture.expectedAssistant] yield* admit(session, prompt) const bus = yield* Bus.Service - const live = yield* bus.subscribe(fixture.delta).pipe(Stream.take(32), Stream.runCollect, Effect.forkScoped) + const live = fixture.delta + ? yield* bus.subscribe(fixture.delta).pipe(Stream.take(2), Stream.runCollect, Effect.forkScoped) + : undefined yield* Effect.yieldNow yield* TestLLM.push(fixture.completeEvents) yield* session.resume(sessionID) const { db } = yield* Database.Service - const deltas = yield* db - .select({ type: EventTable.type }) - .from(EventTable) - .where(eq(EventTable.type, Bus.versionedType(fixture.delta.type, 1))) - .all() - .pipe(Effect.orDie) - expect(Array.from(yield* Fiber.join(live))).toHaveLength(32) + const deltas = fixture.delta + ? yield* db + .select({ type: EventTable.type }) + .from(EventTable) + .where(eq(EventTable.type, Bus.versionedType(fixture.delta.type, 1))) + .all() + .pipe(Effect.orDie) + : [] + if (live) { + const streamed = Array.from(yield* Fiber.join(live)) + expect(streamed).toHaveLength(2) + expect( + streamed + .map((event) => { + if (!event.data || typeof event.data !== "object" || !("delta" in event.data)) + throw new Error("Expected delta event") + if (typeof event.data.delta !== "string") throw new Error("Expected string delta") + return event.data.delta + }) + .join(""), + ).toBe(chunks.join("")) + } expect(deltas).toHaveLength(0) expect(yield* session.context(sessionID)).toMatchObject(expectedContext) @@ -5219,8 +5235,11 @@ describe("SessionRunnerLLM", () => { ) for (const kind of fragmentKinds) { - it.effect(`broadcasts provider ${kind} deltas without storing projection rewrites`, () => - verifyEphemeralDeltas(kind), + it.effect( + kind === "tool input" + ? "does not broadcast provider tool input deltas" + : `batches provider ${kind} deltas without storing projection rewrites`, + () => verifyEphemeralDeltas(kind), ) it.effect(`durably closes partial ${kind} when the provider stream fails`, () => verifyPartialFlushOnFailure(kind))