backfill encrypted_content from response.completed for missing reasoning blocks (#6608)

fixes #6409

packages/ai/src/api/openai-responses-shared.ts

@@ -337,6 +337,7 @@ export async function processResponsesStream<TApi extends Api>(
 ): Promise<void> {
 	let sawTerminalResponseEvent = false;
 	const outputSlots = new Map<number, ResponsesOutputSlot>();
+	const reasoningBlocksById = new Map<string, ThinkingContent>();
 	const getSlot = <TType extends ResponsesOutputSlot["type"]>(
 		outputIndex: number,
 		type: TType,
@@ -388,10 +389,29 @@ export async function processResponsesStream<TApi extends Api>(
 	const getOrCreateSlot = (outputIndex: number, item: ResponseOutputItem): ResponsesOutputSlot | undefined => {
 		return outputSlots.get(outputIndex) ?? createSlot(outputIndex, item);
 	};
+	// Azure OpenAI can omit reasoning.encrypted_content from response.output_item.done
+	// and provide it only in response.completed.response.output. Backfill the
+	// persisted reasoning signature from the terminal response to keep store:false
+	// multi-turn replay stateless. See https://github.com/earendil-works/pi/issues/6409.
+	const backfillReasoningSignatures = (responseOutput: ResponseOutputItem[]): void => {
+		for (const item of responseOutput) {
+			if (item.type !== "reasoning" || !item.encrypted_content) continue;
+			const block = reasoningBlocksById.get(item.id);
+			if (!block?.thinkingSignature) continue;
+
+			const storedItem = JSON.parse(block.thinkingSignature) as ResponseReasoningItem;
+			if (storedItem.encrypted_content) continue;
+			block.thinkingSignature = JSON.stringify({
+				...storedItem,
+				encrypted_content: item.encrypted_content,
+			});
+		}
+	};
 	const finalizeResponse = (
 		response: Extract<ResponseStreamEvent, { type: "response.completed" | "response.incomplete" }>["response"],
 	): void => {
 		sawTerminalResponseEvent = true;
+		backfillReasoningSignatures(response.output ?? []);
 		if (response?.id) {
 			output.responseId = response.id;
 		}
@@ -519,6 +539,7 @@ export async function processResponsesStream<TApi extends Api>(
 				const contentText = item.content?.map((c) => c.text).join("\n\n") || "";
 				slot.block.thinking = summaryText || contentText || slot.block.thinking;
 				slot.block.thinkingSignature = JSON.stringify(item);
+				reasoningBlocksById.set(item.id, slot.block);
 				stream.push({
 					type: "thinking_end",
 					contentIndex: slot.contentIndex,

packages/ai/test/azure-openai-responses-reasoning-replay.test.ts

@@ -0,0 +1,136 @@
+import type { ResponseReasoningItem, ResponseStreamEvent } from "openai/resources/responses/responses.js";
+import { describe, expect, it } from "vitest";
+import { convertResponsesMessages, processResponsesStream } from "../src/api/openai-responses-shared.ts";
+import type { AssistantMessage, Context, Model } from "../src/types.ts";
+import { AssistantMessageEventStream } from "../src/utils/event-stream.ts";
+
+function createModel(): Model<"azure-openai-responses"> {
+	return {
+		id: "gpt-5-mini",
+		name: "GPT-5 Mini",
+		api: "azure-openai-responses",
+		provider: "azure-openai-responses",
+		baseUrl: "https://example.invalid",
+		reasoning: true,
+		input: ["text"],
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+		contextWindow: 400000,
+		maxTokens: 128000,
+	};
+}
+
+function createOutput(model: Model<"azure-openai-responses">): AssistantMessage {
+	return {
+		role: "assistant",
+		content: [],
+		api: model.api,
+		provider: model.provider,
+		model: model.id,
+		usage: {
+			input: 0,
+			output: 0,
+			cacheRead: 0,
+			cacheWrite: 0,
+			totalTokens: 0,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		stopReason: "stop",
+		timestamp: Date.now(),
+	};
+}
+
+async function* createEvents(
+	doneItem: ResponseReasoningItem,
+	completedItem: ResponseReasoningItem,
+): AsyncIterable<ResponseStreamEvent> {
+	yield {
+		type: "response.output_item.added",
+		output_index: 0,
+		sequence_number: 0,
+		item: { type: "reasoning", id: doneItem.id, summary: [] },
+	} as ResponseStreamEvent;
+	yield {
+		type: "response.output_item.done",
+		output_index: 0,
+		sequence_number: 1,
+		item: doneItem,
+	} as ResponseStreamEvent;
+	yield {
+		type: "response.completed",
+		sequence_number: 2,
+		response: {
+			id: "resp_test",
+			status: "completed",
+			output: [completedItem],
+		},
+	} as ResponseStreamEvent;
+}
+
+function getReplayedReasoning(model: Model<"azure-openai-responses">, assistant: AssistantMessage) {
+	const context: Context = {
+		messages: [
+			{ role: "user", content: "first", timestamp: Date.now() - 1 },
+			assistant,
+			{ role: "user", content: "follow-up", timestamp: Date.now() },
+		],
+	};
+	const input = convertResponsesMessages(model, context, new Set(["azure-openai-responses"]));
+	return input.find((item) => item.type === "reasoning");
+}
+
+describe("Azure OpenAI Responses reasoning replay", () => {
+	it("preserves existing encrypted_content from output_item.done", async () => {
+		const model = createModel();
+		const output = createOutput(model);
+		const doneItem: ResponseReasoningItem = {
+			type: "reasoning",
+			id: "rs_done",
+			summary: [],
+			encrypted_content: "from-output-item-done",
+		};
+		const completedItem: ResponseReasoningItem = {
+			...doneItem,
+			encrypted_content: "from-response-completed",
+		};
+
+		await processResponsesStream(
+			createEvents(doneItem, completedItem),
+			output,
+			new AssistantMessageEventStream(),
+			model,
+		);
+
+		expect(getReplayedReasoning(model, output)).toMatchObject({
+			type: "reasoning",
+			id: "rs_done",
+			encrypted_content: "from-output-item-done",
+		});
+	});
+
+	it("fills encrypted_content when output_item.done omitted it", async () => {
+		const model = createModel();
+		const output = createOutput(model);
+		const doneItem: ResponseReasoningItem = {
+			type: "reasoning",
+			id: "rs_missing",
+			summary: [],
+		};
+		const completedItem: ResponseReasoningItem = {
+			...doneItem,
+			encrypted_content: "from-response-completed",
+		};
+
+		await processResponsesStream(
+			createEvents(doneItem, completedItem),
+			output,
+			new AssistantMessageEventStream(),
+			model,
+		);
+
+		expect(getReplayedReasoning(model, output)).toMatchObject({
+			type: "reasoning",
+			id: "rs_missing",
+			encrypted_content: "from-response-completed",
+		});
+	});
+});