Explorar o código

fix(core): retain reasoning delta state (#35758)

Kit Langton hai 1 mes
pai
achega
521ea87192

+ 20 - 7
packages/core/src/session/runner/publish-llm-event.ts

@@ -101,27 +101,36 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
     ended: (id: string, value: string, ordinal: number, state?: Record<string, unknown>) => Effect.Effect<void>,
     ended: (id: string, value: string, ordinal: number, state?: Record<string, unknown>) => Effect.Effect<void>,
     single = false,
     single = false,
   ) => {
   ) => {
-    const chunks = new Map<string, { readonly ordinal: number; readonly values: string[] }>()
+    const chunks = new Map<
+      string,
+      { readonly ordinal: number; readonly values: string[]; state?: Record<string, unknown> }
+    >()
     let nextOrdinal = 0
     let nextOrdinal = 0
-    const start = (id: string) =>
+    const start = (id: string, state?: Record<string, unknown>) =>
       Effect.suspend(() => {
       Effect.suspend(() => {
         if (chunks.has(id)) return Effect.die(new Error(`Duplicate ${name} start: ${id}`))
         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}`))
         if (single && chunks.size > 0) return Effect.die(new Error(`${name} start before end: ${id}`))
         const ordinal = nextOrdinal++
         const ordinal = nextOrdinal++
-        chunks.set(id, { ordinal, values: [] })
+        chunks.set(id, { ordinal, values: [], state })
         return Effect.succeed(ordinal)
         return Effect.succeed(ordinal)
       })
       })
-    const append = (id: string, value: string) =>
+    const append = (id: string, value: string, state?: Record<string, unknown>) =>
       Effect.suspend(() => {
       Effect.suspend(() => {
         const current = chunks.get(id)
         const current = chunks.get(id)
         if (!current) return Effect.die(new Error(`${name} delta before start: ${id}`))
         if (!current) return Effect.die(new Error(`${name} delta before start: ${id}`))
         current.values.push(value)
         current.values.push(value)
+        if (state !== undefined) current.state = { ...current.state, ...state }
         return Effect.succeed(current.ordinal)
         return Effect.succeed(current.ordinal)
       })
       })
     const end = Effect.fnUntraced(function* (id: string, state?: Record<string, unknown>) {
     const end = Effect.fnUntraced(function* (id: string, state?: Record<string, unknown>) {
       const current = chunks.get(id)
       const current = chunks.get(id)
       if (!current) return yield* Effect.die(new Error(`${name} end before start: ${id}`))
       if (!current) return yield* Effect.die(new Error(`${name} end before start: ${id}`))
-      yield* ended(id, current.values.join(""), current.ordinal, state)
+      yield* ended(
+        id,
+        current.values.join(""),
+        current.ordinal,
+        state === undefined ? current.state : { ...current.state, ...state },
+      )
       chunks.delete(id)
       chunks.delete(id)
     })
     })
     const flush = Effect.fnUntraced(function* () {
     const flush = Effect.fnUntraced(function* () {
@@ -288,7 +297,7 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
         return
         return
       case "reasoning-start":
       case "reasoning-start":
         retryEvidence = true
         retryEvidence = true
-        const startedReasoningOrdinal = yield* reasoning.start(event.id)
+        const startedReasoningOrdinal = yield* reasoning.start(event.id, providerState(event.providerMetadata))
         yield* events.publish(SessionEvent.Reasoning.Started, {
         yield* events.publish(SessionEvent.Reasoning.Started, {
           sessionID: input.sessionID,
           sessionID: input.sessionID,
           assistantMessageID: yield* startAssistant(),
           assistantMessageID: yield* startAssistant(),
@@ -297,7 +306,11 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
         })
         })
         return
         return
       case "reasoning-delta":
       case "reasoning-delta":
-        const deltaReasoningOrdinal = yield* reasoning.append(event.id, event.text)
+        const deltaReasoningOrdinal = yield* reasoning.append(
+          event.id,
+          event.text,
+          providerState(event.providerMetadata),
+        )
         yield* events.publish(SessionEvent.Reasoning.Delta, {
         yield* events.publish(SessionEvent.Reasoning.Delta, {
           sessionID: input.sessionID,
           sessionID: input.sessionID,
           assistantMessageID: yield* currentAssistantMessageID(),
           assistantMessageID: yield* currentAssistantMessageID(),

+ 19 - 0
packages/core/test/session-runner-tool-events.test.ts

@@ -110,6 +110,25 @@ test("provider state uses the route provider instead of the catalog provider", a
   })
   })
 })
 })
 
 
+test("reasoning state from an empty delta is retained at reasoning end", async () => {
+  const { published, publisher } = capture()
+  await Effect.runPromise(publisher.publish(LLMEvent.reasoningStart({ id: "reasoning" })))
+  await Effect.runPromise(
+    publisher.publish(
+      LLMEvent.reasoningDelta({
+        id: "reasoning",
+        text: "",
+        providerMetadata: { openai: { signature: "signed" } },
+      }),
+    ),
+  )
+  await Effect.runPromise(publisher.publish(LLMEvent.reasoningEnd({ id: "reasoning" })))
+
+  expect(published.find((event) => event.type === "session.reasoning.ended.1")?.data).toMatchObject({
+    state: { signature: "signed" },
+  })
+})
+
 test("binary failure emits no success event", async () => {
 test("binary failure emits no success event", async () => {
   const { published, publisher } = capture()
   const { published, publisher } = capture()
   await Effect.runPromise(publisher.publish(call))
   await Effect.runPromise(publisher.publish(call))