From c5bf4edb105d62e1cf24eda2d01d8732c6b03980 Mon Sep 17 00:00:00 2001 From: Aiden Cline <63023139+rekram1-node@users.noreply.github.com> Date: Sat, 25 Jul 2026 14:06:00 -0500 Subject: [PATCH] fix(ai): preserve response message phases (#38777) --- packages/ai/src/protocols/open-responses.ts | 164 +++++++++++++++--- packages/ai/src/protocols/openai-responses.ts | 21 ++- .../openai-compatible-responses.test.ts | 27 ++- .../ai/test/provider/openai-responses.test.ts | 115 ++++++++++++ .../client/src/promise/generated/types.ts | 46 +++-- packages/core/src/session/message-updater.ts | 5 +- .../src/session/runner/publish-llm-event.ts | 9 +- .../core/src/session/runner/to-llm-message.ts | 9 +- .../core/test/session-runner-message.test.ts | 31 ++++ packages/core/test/session-runner.test.ts | 41 +++++ packages/schema/src/session-event.ts | 1 + packages/schema/src/session-message.ts | 1 + 12 files changed, 416 insertions(+), 54 deletions(-) diff --git a/packages/ai/src/protocols/open-responses.ts b/packages/ai/src/protocols/open-responses.ts index 249f6a93b43..3372145087c 100644 --- a/packages/ai/src/protocols/open-responses.ts +++ b/packages/ai/src/protocols/open-responses.ts @@ -55,6 +55,9 @@ const OpenResponsesOutputText = Schema.Struct({ text: Schema.String, }) +export const MessagePhase = Schema.Literals(["commentary", "final_answer"]) +type MessagePhase = Schema.Schema.Type + const OpenResponsesReasoningSummaryText = Schema.Struct({ type: Schema.tag("summary_text"), text: Schema.String, @@ -86,10 +89,14 @@ const OpenResponsesFunctionCallOutput = Schema.Union([ Schema.Array(OpenResponsesFunctionCallOutputContent), ]) -const OpenResponsesInputItem = Schema.Union([ +export const InputItem = Schema.Union([ Schema.Struct({ role: Schema.tag("system"), content: Schema.String }), Schema.Struct({ role: Schema.tag("user"), content: Schema.Array(OpenResponsesInputContent) }), - Schema.Struct({ role: Schema.tag("assistant"), content: Schema.Array(OpenResponsesOutputText) }), + Schema.Struct({ + role: Schema.tag("assistant"), + content: Schema.Array(OpenResponsesOutputText), + phase: Schema.optionalKey(MessagePhase), + }), OpenResponsesReasoningItem, OpenResponsesItemReference, Schema.Struct({ @@ -104,7 +111,14 @@ const OpenResponsesInputItem = Schema.Union([ output: OpenResponsesFunctionCallOutput, }), ]) -type OpenResponsesInputItem = Schema.Schema.Type +type OpenResponsesInputItem = Schema.Schema.Type +type LoweredInputItem = + | OpenResponsesInputItem + | { + readonly role: "assistant" + readonly content: ReadonlyArray<{ readonly type: "output_text"; readonly text: string }> + readonly phase?: MessagePhase | null + } // Mutable counterpart of the schema reasoning item so `lowerMessages` can fold // multiple streamed summary parts into the same item before flushing. @@ -135,7 +149,7 @@ export const ToolChoice = Schema.Union([ // transports in sync without a destructure-and-strip dance. export const coreFields = { model: Schema.String, - input: Schema.Array(OpenResponsesInputItem), + input: Schema.Array(InputItem), instructions: Schema.optional(Schema.String), tools: optionalArray(Tool), tool_choice: Schema.optional(ToolChoice), @@ -206,6 +220,7 @@ export const Event = Schema.StructWithRest( Schema.Struct({ type: Schema.String, delta: Schema.optional(Schema.String), + text: Schema.optional(Schema.String), item_id: Schema.optional(Schema.String), summary_index: Schema.optional(Schema.Number), item: Schema.optional(StreamItem), @@ -238,6 +253,7 @@ export interface Extension { readonly media: ProviderShared.ValidatedMedia readonly request: LLMRequest }) => MediaInput | undefined + readonly messagePhase?: (value: unknown) => MessagePhase | null | undefined } const BASE: Extension = { id: ADAPTER, name: NAME } @@ -249,6 +265,9 @@ export interface ParserState { readonly tools: ToolStream.State readonly hasFunctionCall: boolean readonly lifecycle: Lifecycle.State + readonly messageItems: ReadonlySet + readonly messagePhase: (value: unknown) => MessagePhase | null | undefined + readonly messagePhases: Readonly> readonly reasoningItems: Readonly> readonly store: boolean | undefined } @@ -378,9 +397,9 @@ const lowerToolResultOutput = Effect.fn("OpenResponses.lowerToolResultOutput")(f }) const lowerMessages = Effect.fn("OpenResponses.lowerMessages")(function* (request: LLMRequest, extension: Extension) { - const system: OpenResponsesInputItem[] = + const system: LoweredInputItem[] = request.system.length === 0 ? [] : [{ role: "system", content: ProviderShared.joinText(request.system) }] - const input: OpenResponsesInputItem[] = [...system] + const input: LoweredInputItem[] = [...system] const store = OpenResponsesOptions.resolve(request).store const providerMetadataKey = request.model.route.providerMetadataKey ?? "openresponses" @@ -412,7 +431,27 @@ const lowerMessages = Effect.fn("OpenResponses.lowerMessages")(function* (reques const hostedToolReferences = new Set() const flushText = () => { if (content.length === 0) return - input.push({ role: "assistant", content: content.map((part) => ({ type: "output_text", text: part.text })) }) + const groups = content.reduce>( + (groups, part) => { + const metadata = part.providerMetadata?.[providerMetadataKey] + const phase = + ProviderShared.isRecord(metadata) + ? messagePhase(metadata.phase, extension) + : undefined + const group = groups.at(-1) + if (group && group.phase === phase) group.parts.push(part) + else groups.push({ phase, parts: [part] }) + return groups + }, + [], + ) + input.push( + ...groups.map((group) => ({ + role: "assistant" as const, + content: group.parts.map((part) => ({ type: "output_text" as const, text: part.text })), + ...(group.phase === undefined ? {} : { phase: group.phase }), + })), + ) content.splice(0, content.length) } for (const part of message.content) { @@ -513,9 +552,9 @@ const lowerOptions = (request: LLMRequest) => { } } -export const fromRequest = Effect.fn("OpenResponses.fromRequest")(function* ( +export const fromRequestWithExtension = Effect.fn("OpenResponses.fromRequestWithExtension")(function* ( request: LLMRequest, - extension: Extension = BASE, + extension: Extension, ) { const generation = request.generation const toolSchemaCompatibility = request.model.compatibility?.toolSchema @@ -541,6 +580,12 @@ export const fromRequest = Effect.fn("OpenResponses.fromRequest")(function* ( } }) +const decodeBody = ProviderShared.validateWith(Schema.decodeUnknownEffect(OpenResponsesBody)) + +export const fromRequest = Effect.fn("OpenResponses.fromRequest")(function* (request: LLMRequest) { + return yield* decodeBody(yield* fromRequestWithExtension(request, BASE)) +}) + // ============================================================================= // Stream Parsing // ============================================================================= @@ -595,24 +640,30 @@ const NO_EVENTS: StepResult["1"] = [] const TERMINAL_TYPES = new Set(["response.completed", "response.incomplete", "response.failed"]) export const terminal = (event: Event) => TERMINAL_TYPES.has(event.type) -const onOutputTextDelta = (state: ParserState, event: Event): StepResult => { +const onOutputTextDelta = (state: ParserState, event: Event, id: string): StepResult => { if (!event.delta) return [state, NO_EVENTS] const events: LLMEvent[] = [] + const phase = state.messagePhases[id] + const metadata = phase === undefined ? undefined : providerMetadata(state, { phase }) + const lifecycle = Lifecycle.textStart(state.lifecycle, events, id, metadata) return [ - { ...state, lifecycle: Lifecycle.textDelta(state.lifecycle, events, event.item_id ?? "text-0", event.delta) }, + { ...state, lifecycle: Lifecycle.textDelta(lifecycle, events, id, event.delta) }, events, ] } -const onOutputTextDone = (state: ParserState, event: Event): StepResult => { +const onOutputTextDone = (state: ParserState, event: Event, id: string): StepResult => { + if (state.messageItems.has(id)) { + if (state.lifecycle.text.has(id) || event.text === undefined) return [state, NO_EVENTS] + return onOutputTextDelta(state, { ...event, delta: event.text }, id) + } const events: LLMEvent[] = [] - return [{ ...state, lifecycle: Lifecycle.textEnd(state.lifecycle, events, event.item_id ?? "text-0") }, events] + return [{ ...state, lifecycle: Lifecycle.textEnd(state.lifecycle, events, id) }, events] } -export const onReasoningDelta = (state: ParserState, event: Event): StepResult => { +export const onReasoningDelta = (state: ParserState, event: Event, itemID: string): StepResult => { if (!event.delta) return [state, NO_EVENTS] const events: LLMEvent[] = [] - const itemID = event.item_id ?? "reasoning-0" const id = event.summary_index !== undefined || state.reasoningItems[itemID] ? `${itemID}:${event.summary_index ?? 0}` : itemID return [ @@ -643,6 +694,18 @@ const reasoningMetadata = (state: ParserState, item: StreamItem & { id: string } // best-effort, not guaranteed. const onOutputItemAdded = (state: ParserState, event: Event): StepResult => { const item = event.item + if (item?.type === "message" && item.id) + return [ + { + ...state, + messageItems: new Set([...state.messageItems, item.id]), + messagePhases: (() => { + const phase = state.messagePhase(item.phase) + return phase === undefined ? state.messagePhases : { ...state.messagePhases, [item.id]: phase } + })(), + }, + NO_EVENTS, + ] if (item && isReasoningItem(item)) { const events: LLMEvent[] = [] return [ @@ -799,7 +862,28 @@ const onOutputItemDone = Effect.fn("OpenResponses.onOutputItemDone")(function* ( const item = event.item if (!item) return [state, NO_EVENTS] satisfies StepResult - if (item.type === "message" && item.id) return onOutputTextDone(state, { ...event, item_id: item.id }) + if (item.type === "message" && item.id) { + const itemPhase = state.messagePhase(item.phase) + const phase = itemPhase === undefined ? state.messagePhases[item.id] : itemPhase + const events: LLMEvent[] = [] + const messageItems = new Set(state.messageItems) + messageItems.delete(item.id) + const { [item.id]: _phase, ...messagePhases } = state.messagePhases + return [ + { + ...state, + lifecycle: Lifecycle.textEnd( + state.lifecycle, + events, + item.id, + phase === undefined ? undefined : providerMetadata(state, { phase }), + ), + messageItems, + messagePhases, + }, + events, + ] satisfies StepResult + } if (item.type === "function_call") { if (!item.id || !item.call_id || !item.name) return [state, NO_EVENTS] satisfies StepResult @@ -899,19 +983,41 @@ const providerError = (state: ParserState, event: Event, fallback: string) => { } export const step = (state: ParserState, event: Event) => { - if (event.type === "response.output_text.delta") return Effect.succeed(onOutputTextDelta(state, event)) - if (event.type === "response.output_text.done") return Effect.succeed(onOutputTextDone(state, event)) - if (event.type === "response.reasoning.delta" || event.type === "response.reasoning_summary_text.delta") - return Effect.succeed(onReasoningDelta(state, event)) - if (event.type === "response.reasoning.done" || event.type === "response.reasoning_summary_text.done") + if (event.type === "response.output_text.delta" || event.type === "response.output_text.done") { + if (!event.item_id) return ProviderShared.eventError(state.id, `${event.type} is missing item_id`) + return Effect.succeed( + event.type === "response.output_text.delta" + ? onOutputTextDelta(state, event, event.item_id) + : onOutputTextDone(state, event, event.item_id), + ) + } + if (event.type === "response.reasoning.delta" || event.type === "response.reasoning_summary_text.delta") { + if (!event.item_id) return ProviderShared.eventError(state.id, `${event.type} is missing item_id`) + return Effect.succeed(onReasoningDelta(state, event, event.item_id)) + } + if (event.type === "response.reasoning.done" || event.type === "response.reasoning_summary_text.done") { + if (!event.item_id) return ProviderShared.eventError(state.id, `${event.type} is missing item_id`) return Effect.succeed(onReasoningDone(state, event)) + } if (event.type === "response.reasoning_summary_part.added") - return Effect.succeed(onReasoningSummaryPartAdded(state, event)) + return event.item_id + ? Effect.succeed(onReasoningSummaryPartAdded(state, event)) + : ProviderShared.eventError(state.id, `${event.type} is missing item_id`) if (event.type === "response.reasoning_summary_part.done") - return Effect.succeed(onReasoningSummaryPartDone(state, event)) - if (event.type === "response.output_item.added") return Effect.succeed(onOutputItemAdded(state, event)) + return event.item_id + ? Effect.succeed(onReasoningSummaryPartDone(state, event)) + : ProviderShared.eventError(state.id, `${event.type} is missing item_id`) + if (event.type === "response.output_item.added") { + if (event.item?.type === "message" && !event.item.id) + return ProviderShared.eventError(state.id, `${event.type} message is missing id`) + return Effect.succeed(onOutputItemAdded(state, event)) + } if (event.type === "response.function_call_arguments.delta") return onFunctionCallArgumentsDelta(state, event) - if (event.type === "response.output_item.done") return onOutputItemDone(state, event) + if (event.type === "response.output_item.done") { + if (event.item?.type === "message" && !event.item.id) + return ProviderShared.eventError(state.id, `${event.type} message is missing id`) + return onOutputItemDone(state, event) + } if (event.type === "response.completed" || event.type === "response.incomplete") return Effect.succeed(onResponseFinish(state, event)) if (event.type === "response.failed") return providerError(state, event, `${state.name} response failed`) @@ -933,10 +1039,18 @@ export const initial = (request: LLMRequest, extension: Extension = BASE): Parse hasFunctionCall: false, tools: ToolStream.empty(), lifecycle: Lifecycle.initial(), + messageItems: new Set(), + messagePhase: (value) => messagePhase(value, extension), + messagePhases: {}, reasoningItems: {}, store: OpenResponsesOptions.resolve(request).store, }) +const messagePhase = (value: unknown, extension: Extension): MessagePhase | null | undefined => { + if (value === "commentary" || value === "final_answer") return value + return extension.messagePhase?.(value) +} + export const protocol = Protocol.make({ id: ADAPTER, body: { diff --git a/packages/ai/src/protocols/openai-responses.ts b/packages/ai/src/protocols/openai-responses.ts index f8655c67294..43237e89f60 100644 --- a/packages/ai/src/protocols/openai-responses.ts +++ b/packages/ai/src/protocols/openai-responses.ts @@ -35,8 +35,18 @@ const OpenAIResponsesToolChoice = Schema.Union([ Schema.Struct({ type: Schema.tag("image_generation") }), ]) +const OpenAIResponsesInputItem = Schema.Union([ + Schema.Struct({ + role: Schema.tag("assistant"), + content: Schema.Array(Schema.Struct({ type: Schema.tag("output_text"), text: Schema.String })), + phase: Schema.optionalKey(Schema.NullOr(OpenResponses.MessagePhase)), + }), + OpenResponses.InputItem, +]) + const OpenAIResponsesCoreFields = { ...OpenResponses.coreFields, + input: Schema.Array(OpenAIResponsesInputItem), tools: optionalArray(OpenAIResponsesTools), tool_choice: Schema.optional(OpenAIResponsesToolChoice), } @@ -60,6 +70,7 @@ const encodeWebSocketMessage = Schema.encodeSync(Schema.fromJsonString(OpenAIRes const extension = { id: ADAPTER, name: NAME, + messagePhase: (value: unknown) => (value === null ? null : undefined), lowerMedia: ({ part, media, request }) => { if (request.model.provider !== "xai" || media.mime !== "application/pdf") return undefined return { @@ -102,7 +113,7 @@ const lowerToolChoice = (toolChoice: NonNullable, tool }) const fromRequest = Effect.fn("OpenAIResponses.fromRequest")(function* (request: LLMRequest) { - const body = yield* OpenResponses.fromRequest( + const body = yield* OpenResponses.fromRequestWithExtension( LLMRequest.update(request, { tools: [], toolChoice: undefined }), extension, ) @@ -208,9 +219,13 @@ const onHostedToolDone = Effect.fn("OpenAIResponses.onHostedToolDone")(function* const step = (state: OpenResponses.ParserState, event: OpenResponses.Event) => { if (event.type === "response.reasoning_text.delta" || event.type === "response.reasoning_summary.delta") - return Effect.succeed(OpenResponses.onReasoningDelta(state, event)) + return event.item_id + ? Effect.succeed(OpenResponses.onReasoningDelta(state, event, event.item_id)) + : ProviderShared.eventError(ADAPTER, `${event.type} is missing item_id`) if (event.type === "response.reasoning_text.done" || event.type === "response.reasoning_summary.done") - return Effect.succeed(OpenResponses.onReasoningDone(state, event)) + return event.item_id + ? Effect.succeed(OpenResponses.onReasoningDone(state, event)) + : ProviderShared.eventError(ADAPTER, `${event.type} is missing item_id`) if (event.type === "response.output_item.done" && event.item && isHostedToolItem(event.item)) return onHostedToolDone(state, event.item) return OpenResponses.step(state, event) diff --git a/packages/ai/test/provider/openai-compatible-responses.test.ts b/packages/ai/test/provider/openai-compatible-responses.test.ts index 39819821097..f6591f50d2c 100644 --- a/packages/ai/test/provider/openai-compatible-responses.test.ts +++ b/packages/ai/test/provider/openai-compatible-responses.test.ts @@ -1,6 +1,6 @@ import { describe, expect } from "bun:test" import { Effect } from "effect" -import { LLM, LLMEvent } from "../../src" +import { LLM, LLMEvent, Message } from "../../src" import { configure } from "../../src/providers/openai-compatible-responses" import { OpenAI } from "../../src/providers" import { OpenResponses } from "../../src/protocols/open-responses" @@ -70,6 +70,31 @@ describe("Open Responses-compatible route", () => { }), ) + it.effect("omits OpenAI-only nullable phases from the Open Responses baseline", () => + Effect.gen(function* () { + const model = configure({ + apiKey: "test-key", + baseURL: "https://responses.example.test/v1", + }).model("example-model") + const prepared = yield* LLMClient.prepare( + LLM.request({ + model, + messages: [ + Message.assistant({ + type: "text", + text: "Unclassified.", + providerMetadata: { openresponses: { phase: null } }, + }), + ], + }), + ) + + expect(prepared.body).toMatchObject({ + input: [{ role: "assistant", content: [{ type: "output_text", text: "Unclassified." }] }], + }) + }), + ) + it.effect("reads standard options from the Open Responses namespace", () => Effect.gen(function* () { const model = configure({ diff --git a/packages/ai/test/provider/openai-responses.test.ts b/packages/ai/test/provider/openai-responses.test.ts index 43d19ed42ec..a3d75a7eebf 100644 --- a/packages/ai/test/provider/openai-responses.test.ts +++ b/packages/ai/test/provider/openai-responses.test.ts @@ -882,6 +882,121 @@ describe("OpenAI Responses route", () => { }), ) + it.effect("preserves and replays assistant message phases", () => + Effect.gen(function* () { + const response = yield* LLMClient.generate(request).pipe( + Effect.provide( + fixedResponse( + sseEvents( + { + type: "response.output_item.added", + item: { type: "message", id: "msg_commentary" }, + }, + { type: "response.output_text.delta", item_id: "msg_commentary", delta: "Checking." }, + { type: "response.output_text.done", item_id: "msg_commentary" }, + { + type: "response.output_item.done", + item: { type: "message", id: "msg_commentary", phase: "commentary" }, + }, + { + type: "response.output_item.added", + item: { type: "message", id: "msg_final", phase: "final_answer" }, + }, + { type: "response.output_text.done", item_id: "msg_final", text: "Finished." }, + { + type: "response.output_item.done", + item: { type: "message", id: "msg_final", phase: "final_answer" }, + }, + { type: "response.output_item.added", item: { type: "message", id: "msg_null", phase: null } }, + { type: "response.output_text.delta", item_id: "msg_null", delta: "Unclassified." }, + { type: "response.output_item.done", item: { type: "message", id: "msg_null", phase: null } }, + { type: "response.completed", response: { id: "resp_1" } }, + ), + ), + ), + ) + + expect(response.message.content).toEqual([ + { + type: "text", + text: "Checking.", + providerMetadata: { openai: { phase: "commentary" } }, + }, + { + type: "text", + text: "Finished.", + providerMetadata: { openai: { phase: "final_answer" } }, + }, + { + type: "text", + text: "Unclassified.", + providerMetadata: { openai: { phase: null } }, + }, + ]) + + const prepared = yield* LLMClient.prepare( + LLM.request({ model, messages: [response.message] }), + ) + expect(prepared.body.input).toEqual([ + { + role: "assistant", + content: [{ type: "output_text", text: "Checking." }], + phase: "commentary", + }, + { + role: "assistant", + content: [{ type: "output_text", text: "Finished." }], + phase: "final_answer", + }, + { + role: "assistant", + content: [{ type: "output_text", text: "Unclassified." }], + phase: null, + }, + ]) + }), + ) + + it.effect("rejects output text events without the spec-required item id", () => + Effect.gen(function* () { + const error = yield* LLMClient.generate(request).pipe( + Effect.provide( + fixedResponse( + sseEvents( + { type: "response.output_text.delta", delta: "orphaned" }, + { type: "response.completed", response: { id: "resp_1" } }, + ), + ), + ), + Effect.flip, + ) + + expect(error.reason._tag).toBe("InvalidProviderOutput") + expect(error.message).toContain("response.output_text.delta is missing item_id") + }), + ) + + it.effect("rejects reasoning events without the spec-required item id", () => + Effect.gen(function* () { + const events = [ + { type: "response.reasoning_summary_part.added", summary_index: 0 }, + { type: "response.reasoning_summary_part.done", summary_index: 0 }, + { type: "response.reasoning_text.done" }, + ] + + for (const event of events) { + const error = yield* LLMClient.generate(request).pipe( + Effect.provide( + fixedResponse(sseEvents(event, { type: "response.completed", response: { id: "resp_1" } })), + ), + Effect.flip, + ) + expect(error.reason._tag).toBe("InvalidProviderOutput") + expect(error.message).toContain(`${event.type} is missing item_id`) + } + }), + ) + it.effect("maps incomplete response reasons", () => Effect.gen(function* () { const generate = (incompleteDetails: object) => diff --git a/packages/client/src/promise/generated/types.ts b/packages/client/src/promise/generated/types.ts index a1b2367c8b4..ef614aa1bbd 100644 --- a/packages/client/src/promise/generated/types.ts +++ b/packages/client/src/promise/generated/types.ts @@ -98,8 +98,6 @@ export type SessionMessageShell = { output?: { output: string; cursor: number; size: number; truncated: boolean } } -export type SessionMessageAssistantText = { type: "text"; text: string } - export type SessionMessageProviderState = { [x: string]: JsonValue } export type SessionMessageToolStateStreaming = { status: "streaming"; input: string } @@ -157,8 +155,6 @@ export type ShellInfo = { time: { started: number; completed?: number } } -export type SessionMessageProviderState3 = { [x: string]: any } - export type SessionMessageProviderState4 = { [x: string]: any } export type SessionMessageProviderState5 = { [x: string]: any } @@ -167,6 +163,10 @@ export type SessionMessageProviderState6 = { [x: string]: any } export type SessionMessageProviderState7 = { [x: string]: any } +export type SessionMessageProviderState8 = { [x: string]: any } + +export type SessionMessageProviderState9 = { [x: string]: any } + export type EventLogSynced = { type: "log.synced"; aggregateID: string; seq?: number } export type ModelReasoningField = "reasoning" | "reasoning_content" | "reasoning_text" | (string & {}) @@ -733,16 +733,6 @@ export type SessionTextStarted = { data: { sessionID: string; assistantMessageID: string; ordinal: number } } -export type SessionTextEnded = { - id: string - created: number - metadata?: { [x: string]: any } - type: "session.text.ended" - durable: { aggregateID: string; seq: number; version: 1 } - location?: LocationRef - data: { sessionID: string; assistantMessageID: string; ordinal: number; text: string } -} - export type SessionToolInputStarted = { id: string created: number @@ -1249,6 +1239,8 @@ export type SessionPendingSynthetic = { delivery: "steer" | "queue" } +export type SessionMessageAssistantText = { type: "text"; text: string; state?: SessionMessageProviderState } + export type SessionMessageAssistantReasoning = { type: "reasoning" text: string @@ -1359,6 +1351,22 @@ export type ShellCreated = { data: { info: ShellInfo } } +export type SessionTextEnded = { + id: string + created: number + metadata?: { [x: string]: any } + type: "session.text.ended" + durable: { aggregateID: string; seq: number; version: 1 } + location?: LocationRef + data: { + sessionID: string + assistantMessageID: string + ordinal: number + text: string + state?: SessionMessageProviderState4 + } +} + export type SessionReasoningStarted = { id: string created: number @@ -1366,7 +1374,7 @@ export type SessionReasoningStarted = { type: "session.reasoning.started" durable: { aggregateID: string; seq: number; version: 1 } location?: LocationRef - data: { sessionID: string; assistantMessageID: string; ordinal: number; state?: SessionMessageProviderState3 } + data: { sessionID: string; assistantMessageID: string; ordinal: number; state?: SessionMessageProviderState5 } } export type SessionReasoningEnded = { @@ -1381,7 +1389,7 @@ export type SessionReasoningEnded = { assistantMessageID: string ordinal: number text: string - state?: SessionMessageProviderState4 + state?: SessionMessageProviderState6 } } @@ -1398,7 +1406,7 @@ export type SessionToolCalled = { callID: string input: { [x: string]: any } executed: boolean - state?: SessionMessageProviderState5 + state?: SessionMessageProviderState7 } } @@ -1853,7 +1861,7 @@ export type SessionToolSuccess = { content: [LLMToolContent, ...Array] metadata?: { [x: string]: JsonValue } executed: boolean - resultState?: SessionMessageProviderState6 + resultState?: SessionMessageProviderState8 } } @@ -1872,7 +1880,7 @@ export type SessionToolFailed = { content?: [LLMToolContent, ...Array] metadata?: { [x: string]: JsonValue } executed: boolean - resultState?: SessionMessageProviderState7 + resultState?: SessionMessageProviderState9 } } diff --git a/packages/core/src/session/message-updater.ts b/packages/core/src/session/message-updater.ts index 3a256e043a8..6d9c4723ff2 100644 --- a/packages/core/src/session/message-updater.ts +++ b/packages/core/src/session/message-updater.ts @@ -319,7 +319,10 @@ export function update(adapter: Adapter, event: SessionEvent.Event) { "session.text.ended": (event) => { return updateOwnedAssistant(event.data.assistantMessageID, (draft) => { const match = latestText(draft) - if (match) match.text = event.data.text + if (match) { + match.text = event.data.text + match.state = castDraft(event.data.state) + } }) }, "session.tool.input.started": (event) => { diff --git a/packages/core/src/session/runner/publish-llm-event.ts b/packages/core/src/session/runner/publish-llm-event.ts index fcec89352cf..440f6779b86 100644 --- a/packages/core/src/session/runner/publish-llm-event.ts +++ b/packages/core/src/session/runner/publish-llm-event.ts @@ -149,13 +149,14 @@ export const createLLMEventPublisher = (events: Pick + (_textID, value, ordinal, state) => Effect.gen(function* () { yield* events.publish(SessionEvent.Text.Ended, { sessionID: input.sessionID, assistantMessageID: yield* currentAssistantMessageID(), ordinal, text: value, + state, }) }), true, @@ -326,7 +327,7 @@ export const createLLMEventPublisher = (events: Pick { - if (item.type === "text") return [{ type: "text", text: item.text }] + if (item.type === "text") + return [ + { + type: "text", + text: item.text, + providerMetadata: sameProvider ? providerMetadata(providerMetadataKey, item.state) : undefined, + }, + ] if (item.type === "reasoning") return reuseProviderMetadata ? [ diff --git a/packages/core/test/session-runner-message.test.ts b/packages/core/test/session-runner-message.test.ts index 3cebdcb985d..de33bcd69dc 100644 --- a/packages/core/test/session-runner-message.test.ts +++ b/packages/core/test/session-runner-message.test.ts @@ -767,4 +767,35 @@ Recent work }, ]) }) + + test("preserves assistant text provider state across same-provider model changes and failures", () => { + const messages = toLLMMessages( + [ + SessionMessage.Assistant.make({ + id: id("assistant-phase"), + type: "assistant", + agent: build, + model: { id: ModelV2.ID.make("old"), providerID: ProviderV2.ID.make("provider") }, + content: [ + SessionMessage.AssistantText.make({ + type: "text", + text: "Checking.", + state: { phase: "commentary" }, + }), + ], + error: { type: "provider.unknown", message: "Interrupted after commentary" }, + time: { created, completed: created }, + }), + ], + ModelV2.Ref.make({ id: ModelV2.ID.make("new"), providerID: ProviderV2.ID.make("provider") }), + ) + + expect(messages[0]?.content).toEqual([ + { + type: "text", + text: "Checking.", + providerMetadata: { provider: { phase: "commentary" } }, + }, + ]) + }) }) diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 5e07b3551af..1f96f41cbc7 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -2672,6 +2672,47 @@ describe("SessionRunnerLLM", () => { }), ) + it.effect("restores durable text provider metadata in the next request", () => + Effect.gen(function* () { + const session = yield* setup + yield* admit(session, "Check first") + + response = [ + LLMEvent.stepStart({ index: 0 }), + LLMEvent.textStart({ id: "commentary", providerMetadata: { openai: { phase: "commentary" } } }), + LLMEvent.textDelta({ id: "commentary", text: "Checking." }), + LLMEvent.textEnd({ + id: "commentary", + providerMetadata: { openai: { phase: "commentary" }, anthropic: { ignored: true } }, + }), + LLMEvent.stepFinish({ index: 0, reason: { normalized: "stop" } }), + LLMEvent.finish({ reason: { normalized: "stop" } }), + ] + yield* session.resume(sessionID) + yield* replaySessionProjection(sessionID) + + expect(yield* session.context(sessionID)).toMatchObject([ + { type: "user", text: "Check first" }, + { + type: "assistant", + content: [{ type: "text", text: "Checking.", state: { phase: "commentary" } }], + }, + ]) + + yield* admit(session, "Continue") + response = [] + yield* session.resume(sessionID) + + expect(requests[1]?.messages[1]?.content).toEqual([ + { + type: "text", + text: "Checking.", + providerMetadata: { openai: { phase: "commentary" } }, + }, + ]) + }), + ) + it.effect("replays durable provider-executed tool results inline in the next request", () => Effect.gen(function* () { const session = yield* setup diff --git a/packages/schema/src/session-event.ts b/packages/schema/src/session-event.ts index faa4853aded..3a31643bbbf 100644 --- a/packages/schema/src/session-event.ts +++ b/packages/schema/src/session-event.ts @@ -314,6 +314,7 @@ export namespace Text { assistantMessageID: SessionMessage.ID, ordinal: NonNegativeInt, text: Schema.String, + state: SessionMessage.ProviderState.pipe(optional), }, }) export type Ended = typeof Ended.Type diff --git a/packages/schema/src/session-message.ts b/packages/schema/src/session-message.ts index c30ffc6aa46..4e0858e4055 100644 --- a/packages/schema/src/session-message.ts +++ b/packages/schema/src/session-message.ts @@ -155,6 +155,7 @@ export interface AssistantText extends Schema.Schema.Type export const AssistantText = Schema.Struct({ type: Schema.tag("text"), text: Schema.String, + state: ProviderState.pipe(optional), }).annotate({ identifier: "Session.Message.Assistant.Text" }) export interface AssistantReasoning extends Schema.Schema.Type {}