fix(ai): use unique Gemini block ids (#46279)

This commit is contained in:
Aiden Cline 2026-08-30 13:21:35 -05:00 committed by GitHub
parent d323b34826
commit c746ea3210
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 158 additions and 24 deletions

View file

@ -240,6 +240,10 @@ interface ParserState {
readonly lifecycle: Lifecycle.State
readonly reasoningSignature?: string
readonly textSignature?: string
readonly reasoningId?: string
readonly textId?: string
readonly nextReasoningId: number
readonly nextTextId: number
readonly seenCallIds?: ReadonlySet<string>
}
@ -571,19 +575,23 @@ const finish = (state: ParserState): ReadonlyArray<LLMEvent> => {
const events: LLMEvent[] = []
let lifecycle = state.lifecycle
if (state.reasoningSignature !== undefined)
if (state.reasoningId !== undefined)
lifecycle = Lifecycle.reasoningEnd(
lifecycle,
events,
"reasoning-0",
providerMetadata(state.providerMetadataKey, { thoughtSignature: state.reasoningSignature }),
state.reasoningId,
state.reasoningSignature === undefined
? undefined
: providerMetadata(state.providerMetadataKey, { thoughtSignature: state.reasoningSignature }),
)
if (state.textSignature !== undefined)
if (state.textId !== undefined)
lifecycle = Lifecycle.textEnd(
lifecycle,
events,
"text-0",
providerMetadata(state.providerMetadataKey, { thoughtSignature: state.textSignature }),
state.textId,
state.textSignature === undefined
? undefined
: providerMetadata(state.providerMetadataKey, { thoughtSignature: state.textSignature }),
)
Lifecycle.finish(lifecycle, events, {
reason: {
@ -632,6 +640,10 @@ const step = (state: ParserState, event: GeminiEvent) => {
let lifecycle = nextState.lifecycle
let reasoningSignature = nextState.reasoningSignature
let textSignature = nextState.textSignature
let reasoningId = nextState.reasoningId
let textId = nextState.textId
let nextReasoningId = nextState.nextReasoningId
let nextTextId = nextState.nextTextId
// Supplier ids must be tracked across chunks of the same response, not just within one event's parts.
const seenCallIds = new Set(nextState.seenCallIds)
@ -657,27 +669,51 @@ const step = (state: ParserState, event: GeminiEvent) => {
else if (signature !== undefined && "text" in part) textSignature = signature
if ("text" in part && part.text.length > 0) {
if (part.thought) {
if (textId !== undefined) {
lifecycle = Lifecycle.textEnd(
lifecycle,
events,
textId,
textSignature
? providerMetadata(state.providerMetadataKey, { thoughtSignature: textSignature })
: undefined,
)
textId = undefined
textSignature = undefined
}
if (reasoningId === undefined) {
reasoningId = `reasoning-${nextReasoningId}`
nextReasoningId += 1
}
lifecycle = Lifecycle.reasoningDelta(
lifecycle,
events,
"reasoning-0",
reasoningId,
part.text,
signature ? providerMetadata(state.providerMetadataKey, { thoughtSignature: signature }) : undefined,
)
continue
}
lifecycle = Lifecycle.reasoningEnd(
lifecycle,
events,
"reasoning-0",
reasoningSignature
? providerMetadata(state.providerMetadataKey, { thoughtSignature: reasoningSignature })
: undefined,
)
if (reasoningId !== undefined) {
lifecycle = Lifecycle.reasoningEnd(
lifecycle,
events,
reasoningId,
reasoningSignature
? providerMetadata(state.providerMetadataKey, { thoughtSignature: reasoningSignature })
: undefined,
)
reasoningId = undefined
reasoningSignature = undefined
}
if (textId === undefined) {
textId = `text-${nextTextId}`
nextTextId += 1
}
lifecycle = Lifecycle.textDelta(
lifecycle,
events,
"text-0",
textId,
part.text,
textSignature ? providerMetadata(state.providerMetadataKey, { thoughtSignature: textSignature }) : undefined,
)
@ -695,14 +731,18 @@ const step = (state: ParserState, event: GeminiEvent) => {
const duplicate = supplied !== undefined && seenCallIds.has(supplied)
if (supplied !== undefined) seenCallIds.add(supplied)
const id = supplied !== undefined && !duplicate ? supplied : `tool_${crypto.randomUUID().replaceAll("-", "")}`
lifecycle = Lifecycle.reasoningEnd(
lifecycle,
events,
"reasoning-0",
reasoningSignature
? providerMetadata(state.providerMetadataKey, { thoughtSignature: reasoningSignature })
: undefined,
)
if (reasoningId !== undefined) {
lifecycle = Lifecycle.reasoningEnd(
lifecycle,
events,
reasoningId,
reasoningSignature
? providerMetadata(state.providerMetadataKey, { thoughtSignature: reasoningSignature })
: undefined,
)
reasoningId = undefined
reasoningSignature = undefined
}
lifecycle = Lifecycle.stepStart(lifecycle, events)
events.push(
LLMEvent.toolCall({
@ -725,6 +765,10 @@ const step = (state: ParserState, event: GeminiEvent) => {
lifecycle,
reasoningSignature,
textSignature,
reasoningId,
textId,
nextReasoningId,
nextTextId,
seenCallIds,
finishReason: candidate.finishReason ?? nextState.finishReason,
},
@ -752,6 +796,8 @@ export const protocol = Protocol.make({
providerMetadataKey: request.model.route.providerMetadataKey ?? String(request.model.provider),
hasToolCalls: false,
lifecycle: Lifecycle.initial(),
nextReasoningId: 0,
nextTextId: 0,
}),
step,
onHalt: (state) => Effect.succeed(finish(state)),

View file

@ -906,6 +906,94 @@ describe("Gemini route", () => {
}),
)
it.effect("assigns unique ids to separated reasoning blocks", () =>
Effect.gen(function* () {
const body = sseEvents(
{
candidates: [
{
content: {
role: "model",
parts: [{ text: "A", thought: true, thoughtSignature: "reasoning_sig_a" }],
},
},
],
},
{
candidates: [
{
content: { role: "model", parts: [{ text: "X", thoughtSignature: "text_sig_x" }] },
},
],
},
{
candidates: [
{
content: {
role: "model",
parts: [{ text: "B", thought: true, thoughtSignature: "reasoning_sig_b" }],
},
},
],
},
{
candidates: [
{
content: { role: "model", parts: [{ text: "Y", thoughtSignature: "text_sig_y" }] },
finishReason: "STOP",
},
],
},
)
const response = yield* LLMClient.generate(request).pipe(Effect.provide(fixedResponse(body)))
const starts = response.events.filter((event) => event.type === "reasoning-start")
const deltas = response.events.filter((event) => event.type === "reasoning-delta")
const ends = response.events.filter((event) => event.type === "reasoning-end")
expect(starts.map((event) => event.id)).toEqual(["reasoning-0", "reasoning-1"])
expect(starts[0]?.id).not.toBe(starts[1]?.id)
expect(deltas.map((event) => ({ id: event.id, text: event.text }))).toEqual([
{ id: "reasoning-0", text: "A" },
{ id: "reasoning-1", text: "B" },
])
expect(ends.map((event) => event.id)).toEqual(["reasoning-0", "reasoning-1"])
expect(response.events.filter((event) => event.type === "text-start").map((event) => event.id)).toEqual([
"text-0",
"text-1",
])
expect(response.events.filter((event) => event.type === "text-delta").map((event) => event.id)).toEqual([
"text-0",
"text-1",
])
expect(response.events.filter((event) => event.type === "text-end").map((event) => event.id)).toEqual([
"text-0",
"text-1",
])
expect(response.message.content).toEqual([
{
type: "reasoning",
text: "A",
providerMetadata: { google: { thoughtSignature: "reasoning_sig_a" } },
},
{
type: "text",
text: "X",
providerMetadata: { google: { thoughtSignature: "text_sig_x" } },
},
{
type: "reasoning",
text: "B",
providerMetadata: { google: { thoughtSignature: "reasoning_sig_b" } },
},
{
type: "text",
text: "Y",
providerMetadata: { google: { thoughtSignature: "text_sig_y" } },
},
])
}),
)
it.effect("ignores unknown response parts", () =>
Effect.gen(function* () {
const response = yield* LLMClient.generate(request).pipe(