瀏覽代碼

fix(core): batch streamed session deltas

Dax 1 天之前
父節點
當前提交
4fee4d7d86

+ 57 - 39
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<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":

+ 81 - 0
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(

+ 31 - 12
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))