mirror of
https://github.com/anomalyco/opencode.git
synced 2026-11-14 03:56:55 +03:00
fix(core): batch streamed session deltas
This commit is contained in:
@@ -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<Bus.Interface, "publish">, input: Input) => {
|
||||
const deltaBatchInterval = 100
|
||||
const tools = new Map<
|
||||
string,
|
||||
{
|
||||
@@ -123,32 +124,52 @@ export const createLLMEventPublisher = (bus: Pick<Bus.Interface, "publish">, inp
|
||||
const fragments = (
|
||||
name: string,
|
||||
ended: (id: string, value: string, ordinal: number, state?: Record<string, unknown>) => Effect.Effect<void>,
|
||||
delta?: (id: string, value: string, ordinal: number) => Effect.Effect<void>,
|
||||
single = false,
|
||||
) => {
|
||||
const chunks = new Map<
|
||||
string,
|
||||
{ readonly ordinal: number; readonly values: string[]; state?: Record<string, unknown> }
|
||||
>()
|
||||
type Fragment = {
|
||||
readonly ordinal: number
|
||||
readonly values: string[]
|
||||
pending: string
|
||||
publishedAt?: number
|
||||
state?: Record<string, unknown>
|
||||
}
|
||||
const chunks = new Map<string, Fragment>()
|
||||
let nextOrdinal = 0
|
||||
const start = (id: string, state?: Record<string, unknown>) =>
|
||||
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<string, unknown>) =>
|
||||
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<string, unknown>) {
|
||||
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<string, unknown>, 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<Bus.Interface, "publish">, 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<Bus.Interface, "publish">, 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<Bus.Interface, "publish">, 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<Bus.Interface, "publish">, 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<Bus.Interface, "publish">, 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<Bus.Interface, "publish">, 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":
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user