pi-openai-toolkit 0.18.0 → 0.20.0

package/src/image-generation/protocol.ts

@@ -1,6 +1,6 @@
 import { DEFAULT_IMAGE_GENERATION_MODELS } from "../types";
 import {
-	IMAGE_GENERATION_MIME_TYPE,
+	IMAGE_GENERATION_MIME_TYPES,
 	IMAGE_GENERATION_QUALITIES,
 	IMAGE_GENERATION_SIZES,
 	MAX_GENERATED_IMAGE_BYTES,
@@ -15,6 +15,7 @@ import {
 	ImageGenerationError,
 	sanitizeImageDiagnostic,
 	type GenerateImageParams,
+	type ImageGenerationMimeType,
 	type NormalizedGenerateImageParams,
 	type ParsedGeneratedImage,
 	type PreparedReferenceImage,
@@ -269,35 +270,7 @@ function hasPngEndChunk(buffer: Buffer): boolean {
 }
 
 export function detectReferenceImageMimeType(bytes: Uint8Array): ReferenceImageMimeType | undefined {
-	const buffer = Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength);
-	if (readPngDimensions(buffer) && hasPngEndChunk(buffer)) return "image/png";
-	if (
-		buffer.length >= 4 &&
-		buffer[0] === 0xff &&
-		buffer[1] === 0xd8 &&
-		buffer[2] === 0xff &&
-		buffer.lastIndexOf(Buffer.from([0xff, 0xd9])) >= 3
-	) {
-		return "image/jpeg";
-	}
-	if (
-		buffer.length >= 20 &&
-		buffer.toString("ascii", 0, 4) === "RIFF" &&
-		buffer.toString("ascii", 8, 12) === "WEBP"
-	) {
-		const declaredSize = buffer.readUInt32LE(4) + 8;
-		const chunkType = buffer.toString("ascii", 12, 16);
-		const chunkSize = buffer.readUInt32LE(16);
-		const paddedChunkSize = chunkSize + (chunkSize % 2);
-		if (
-			declaredSize === buffer.length &&
-			(chunkType === "VP8 " || chunkType === "VP8L" || chunkType === "VP8X") &&
-			20 + paddedChunkSize <= buffer.length
-		) {
-			return "image/webp";
-		}
-	}
-	return undefined;
+	return readImageMetadata(bytes)?.mimeType;
 }
 
 export function isValidPng(bytes: Uint8Array): boolean {
@@ -319,14 +292,145 @@ export function readPngDimensions(bytes: Uint8Array): { width: number; height: n
 	return { width, height };
 }
 
+function boundedDimensions(width: number, height: number): { width: number; height: number } | undefined {
+	if (
+		!Number.isInteger(width) ||
+		!Number.isInteger(height) ||
+		width <= 0 ||
+		height <= 0 ||
+		width > MAX_IMAGE_DIMENSION ||
+		height > MAX_IMAGE_DIMENSION
+	) {
+		return undefined;
+	}
+	return { width, height };
+}
+
+const JPEG_DIMENSION_MARKERS = new Set([
+	0xc0,
+	0xc1,
+	0xc2,
+	0xc3,
+	0xc5,
+	0xc6,
+	0xc7,
+	0xc9,
+	0xca,
+	0xcb,
+	0xcd,
+	0xce,
+	0xcf,
+]);
+
+function readJpegDimensions(bytes: Uint8Array): { width: number; height: number } | undefined {
+	const buffer = Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+	if (buffer.length < 4 || buffer[0] !== 0xff || buffer[1] !== 0xd8) return undefined;
+	let offset = 2;
+	while (offset < buffer.length) {
+		if (buffer[offset] !== 0xff) {
+			offset += 1;
+			continue;
+		}
+		while (offset < buffer.length && buffer[offset] === 0xff) offset += 1;
+		if (offset >= buffer.length) return undefined;
+		const marker = buffer[offset++];
+		if (marker === 0xd9 || marker === 0xda) return undefined;
+		if (marker === 0x01 || (marker >= 0xd0 && marker <= 0xd7)) continue;
+		if (offset + 2 > buffer.length) return undefined;
+		const segmentLength = buffer.readUInt16BE(offset);
+		if (segmentLength < 2 || offset + segmentLength > buffer.length) return undefined;
+		if (JPEG_DIMENSION_MARKERS.has(marker) && segmentLength >= 7) {
+			const height = buffer.readUInt16BE(offset + 3);
+			const width = buffer.readUInt16BE(offset + 5);
+			return boundedDimensions(width, height);
+		}
+		offset += segmentLength;
+	}
+	return undefined;
+}
+
+function hasJpegEndMarker(bytes: Uint8Array): boolean {
+	const buffer = Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+	return buffer.length >= 2 && buffer[buffer.length - 2] === 0xff && buffer[buffer.length - 1] === 0xd9;
+}
+
+function readWebpDimensions(bytes: Uint8Array): { width: number; height: number } | undefined {
+	const buffer = Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+	if (
+		buffer.length < 20 ||
+		buffer.toString("ascii", 0, 4) !== "RIFF" ||
+		buffer.toString("ascii", 8, 12) !== "WEBP"
+	) {
+		return undefined;
+	}
+	const declaredSize = buffer.readUInt32LE(4) + 8;
+	if (declaredSize !== buffer.length) return undefined;
+	let offset = 12;
+	while (offset + 8 <= buffer.length) {
+		const chunkType = buffer.toString("ascii", offset, offset + 4);
+		const chunkSize = buffer.readUInt32LE(offset + 4);
+		const dataStart = offset + 8;
+		const dataEnd = dataStart + chunkSize;
+		if (dataEnd > buffer.length) return undefined;
+		if (chunkType === "VP8X" && chunkSize >= 10) {
+			const width = 1 + buffer.readUIntLE(dataStart + 4, 3);
+			const height = 1 + buffer.readUIntLE(dataStart + 7, 3);
+			return boundedDimensions(width, height);
+		}
+		if (chunkType === "VP8 " && chunkSize >= 10) {
+			if (
+				buffer[dataStart + 3] === 0x9d &&
+				buffer[dataStart + 4] === 0x01 &&
+				buffer[dataStart + 5] === 0x2a
+			) {
+				const width = buffer.readUInt16LE(dataStart + 6) & 0x3fff;
+				const height = buffer.readUInt16LE(dataStart + 8) & 0x3fff;
+				return boundedDimensions(width, height);
+			}
+		}
+		if (chunkType === "VP8L" && chunkSize >= 5 && buffer[dataStart] === 0x2f) {
+			const width = 1 + ((buffer[dataStart + 1] | (buffer[dataStart + 2] << 8)) & 0x3fff);
+			const height = 1 + ((buffer[dataStart + 2] >> 6) | (buffer[dataStart + 3] << 2) | (buffer[dataStart + 4] << 10));
+			return boundedDimensions(width, height);
+		}
+		offset = dataEnd + (chunkSize % 2);
+	}
+	return undefined;
+}
+
 function extractBase64(value: string): string {
 	const trimmed = value.trim();
 	const dataUrl = DATA_URL_PATTERN.exec(trimmed);
 	return (dataUrl?.[1] ?? trimmed).trim();
 }
 
-export function decodeGeneratedPng(value: string):
-	| { ok: true; bytes: Buffer; width: number; height: number }
+function normalizeImageMimeType(value: unknown): ImageGenerationMimeType | undefined {
+	if (typeof value !== "string") return undefined;
+	const normalized = value.trim().toLowerCase();
+	if (normalized === "image/jpg") return "image/jpeg";
+	return (IMAGE_GENERATION_MIME_TYPES as readonly string[]).includes(normalized)
+		? normalized as ImageGenerationMimeType
+		: undefined;
+}
+
+export function readImageMetadata(bytes: Uint8Array):
+	| { mimeType: ImageGenerationMimeType; width: number; height: number }
+	| undefined {
+	if (readPngDimensions(bytes) && hasPngEndChunk(Buffer.from(bytes))) {
+		const dimensions = readPngDimensions(bytes)!;
+		return { mimeType: "image/png", ...dimensions };
+	}
+	if (hasJpegEndMarker(bytes)) {
+		const dimensions = readJpegDimensions(bytes);
+		if (dimensions) return { mimeType: "image/jpeg", ...dimensions };
+	}
+	const dimensions = readWebpDimensions(bytes);
+	if (dimensions) return { mimeType: "image/webp", ...dimensions };
+	return undefined;
+}
+
+export function decodeGeneratedImage(value: string, hintedMimeType?: unknown):
+	| { ok: true; bytes: Buffer; mimeType: ImageGenerationMimeType; width: number; height: number }
 	| { ok: false; reason: "malformed-response" | "oversized-response"; errorMessage: string } {
 	const base64 = extractBase64(value);
 	const maxBase64Chars = Math.ceil(MAX_GENERATED_IMAGE_BYTES / 3) * 4;
@@ -358,16 +462,40 @@ export function decodeGeneratedPng(value: string):
 			errorMessage: "Image generation returned invalid or oversized image data.",
 		};
 	}
-	const dimensions = readPngDimensions(bytes);
-	if (!dimensions) {
+	const metadata = readImageMetadata(bytes);
+	if (!metadata) {
 		bytes.fill(0);
 		return {
 			ok: false,
 			reason: "malformed-response",
-			errorMessage: "Image generation returned data that is not a valid PNG image.",
+			errorMessage: "Image generation returned data that is not a supported PNG, JPEG, or WebP image.",
 		};
 	}
-	return { ok: true, bytes, ...dimensions };
+	if (typeof hintedMimeType === "string" && hintedMimeType.trim()) {
+		const hinted = normalizeImageMimeType(hintedMimeType);
+		if (!hinted || hinted !== metadata.mimeType) {
+			bytes.fill(0);
+			return {
+				ok: false,
+				reason: "malformed-response",
+				errorMessage: "Image generation returned an image with inconsistent MIME metadata.",
+			};
+		}
+	}
+	return { ok: true, bytes, ...metadata };
+}
+
+export function decodeGeneratedPng(value: string):
+	| { ok: true; bytes: Buffer; width: number; height: number }
+	| { ok: false; reason: "malformed-response" | "oversized-response"; errorMessage: string } {
+	const decoded = decodeGeneratedImage(value, "image/png");
+	if (!decoded.ok) return decoded;
+	return {
+		ok: true,
+		bytes: decoded.bytes,
+		width: decoded.width,
+		height: decoded.height,
+	};
 }
 
 export function extractProviderErrorMessage(value: unknown, fallback: string): string {
@@ -419,13 +547,22 @@ export function parseImagesApiResponse(value: unknown): ParsedImageResponseResul
 			errorMessage: "Image generation response did not contain base64 image data.",
 		};
 	}
-	const decoded = decodeGeneratedPng(item.b64_json);
+	const mimeHint = item.mime_type ?? item.mimeType;
+	if (mimeHint !== undefined && typeof mimeHint !== "string") {
+		return {
+			ok: false,
+			reason: "malformed-response",
+			errorMessage: "Image generation returned invalid MIME metadata.",
+		};
+	}
+	const decoded = decodeGeneratedImage(item.b64_json, mimeHint);
 	if (!decoded.ok) return decoded;
 
 	return {
 		ok: true,
 		image: {
 			bytes: decoded.bytes,
+			mimeType: decoded.mimeType,
 			imageCallId:
 				normalizedOptionalString(item.id, MAX_IMAGE_IDENTIFIER_CHARS) ?? "image_generation_0",
 			responseId: normalizedOptionalString(value.id, MAX_IMAGE_IDENTIFIER_CHARS),
@@ -438,5 +575,5 @@ export function parseImagesApiResponse(value: unknown): ParsedImageResponseResul
 
 export const _protocolTest = {
 	PNG_SIGNATURE,
-	IMAGE_GENERATION_MIME_TYPE,
+	IMAGE_GENERATION_MIME_TYPES,
 };

package/src/image-generation/render.ts

@@ -15,12 +15,12 @@ import {
 	type Component,
 } from "@earendil-works/pi-tui";
 import { getGeneratedImagesRoot } from "./artifacts";
-import { isValidPng, readPngDimensions } from "./protocol";
+import { readImageMetadata } from "./protocol";
 import {
-	IMAGE_GENERATION_MIME_TYPE,
 	MAX_GENERATED_IMAGE_BYTES,
 	isImageGenerationDetails,
 	type GenerateImageParams,
+	type ImageGenerationMimeType,
 	type ImageGenerationToolDetails,
 } from "./types";
 
@@ -66,6 +66,7 @@ class DiskBackedImageResult extends Container {
 	readonly artifactPath?: string;
 	readonly fileSize?: number;
 	readonly mtimeMs?: number;
+	readonly mimeType?: ImageGenerationMimeType;
 	readonly showImages: boolean;
 	readonly textValue: string;
 
@@ -74,6 +75,7 @@ class DiskBackedImageResult extends Container {
 		artifactPath?: string;
 		fileSize?: number;
 		mtimeMs?: number;
+		mimeType?: ImageGenerationMimeType;
 		showImages: boolean;
 		base64?: string;
 		theme: Theme;
@@ -83,6 +85,7 @@ class DiskBackedImageResult extends Container {
 		this.artifactPath = args.artifactPath;
 		this.fileSize = args.fileSize;
 		this.mtimeMs = args.mtimeMs;
+		this.mimeType = args.mimeType;
 		this.showImages = args.showImages;
 		this.textValue = `${args.text}\n${args.previewWarning ?? ""}`;
 		this.addChild(new Text(args.theme.fg("toolOutput", args.text), 0, 0));
@@ -94,7 +97,7 @@ class DiskBackedImageResult extends Container {
 			this.addChild(
 				new Image(
 					args.base64,
-					IMAGE_GENERATION_MIME_TYPE,
+					args.mimeType ?? "image/png",
 					{ fallbackColor: (text) => args.theme.fg("toolOutput", text) },
 					{ maxWidthCells: 60, filename: args.artifactPath },
 				),
@@ -107,6 +110,7 @@ class DiskBackedImageResult extends Container {
 		artifactPath?: string;
 		fileSize?: number;
 		mtimeMs?: number;
+		mimeType?: ImageGenerationMimeType;
 		showImages: boolean;
 		previewWarning?: string;
 	}): boolean {
@@ -115,6 +119,7 @@ class DiskBackedImageResult extends Container {
 			this.artifactPath === args.artifactPath &&
 			this.fileSize === args.fileSize &&
 			this.mtimeMs === args.mtimeMs &&
+			this.mimeType === args.mimeType &&
 			this.showImages === args.showImages
 		);
 	}
@@ -204,6 +209,7 @@ export function createImageGenerationResultRenderer(
 				artifactPath: realArtifact,
 				fileSize: info.size,
 				mtimeMs: info.mtimeMs,
+				mimeType: result.details.mimeType,
 				showImages: true,
 			})
 		) {
@@ -221,13 +227,14 @@ export function createImageGenerationResultRenderer(
 				previewWarning: "Image preview unavailable: the artifact could not be read.",
 			});
 		}
-		if (bytes.length !== info.size || !isValidPng(bytes) || !readPngDimensions(bytes)) {
+		const metadata = readImageMetadata(bytes);
+		if (bytes.length !== info.size || !metadata || metadata.mimeType !== result.details.mimeType) {
 			bytes.fill(0);
 			return textOnly({
 				text,
 				theme,
 				context,
-				previewWarning: "Image preview unavailable: the artifact is not a valid PNG.",
+				previewWarning: "Image preview unavailable: the artifact is not a valid supported image.",
 			});
 		}
 		const base64 = bytes.toString("base64");
@@ -237,6 +244,7 @@ export function createImageGenerationResultRenderer(
 			artifactPath: realArtifact,
 			fileSize: info.size,
 			mtimeMs: info.mtimeMs,
+			mimeType: metadata.mimeType,
 			showImages: true,
 			base64,
 			theme,

package/src/image-generation/service.ts

@@ -1,7 +1,7 @@
 import { getAgentDir, type ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { assertConfigValid, loadToolkitConfig, resolveToolkitConfig } from "../config";
 import { notifyConfigIssues } from "../config/notifications";
-import { resolveResponsesEnvironment } from "../runtime";
+import { resolveEffectiveModel, resolveResponsesEnvironment, VIRTUAL_MODEL_API } from "../runtime";
 import {
 	copyImageToExplicitPath,
 	prepareExplicitOutputPath,
@@ -21,7 +21,6 @@ import {
 } from "./references";
 import {
 	IMAGE_GENERATION_CAPABLE_APIS,
-	IMAGE_GENERATION_MIME_TYPE,
 	MAX_IMAGE_IDENTIFIER_CHARS,
 	ImageGenerationError,
 	sanitizeImageDiagnostic,
@@ -124,16 +123,28 @@ function asImageGenerationError(value: unknown): ImageGenerationError {
 	);
 }
 
+function formatImageMimeType(mimeType: ImageGenerationDetails["mimeType"]): string {
+	switch (mimeType) {
+		case "image/png":
+			return "PNG";
+		case "image/jpeg":
+			return "JPEG";
+		case "image/webp":
+			return "WebP";
+	}
+}
+
 function formatResultText(args: {
 	artifactPath: string;
 	outputPath?: string;
 	warning?: string;
 	edited: boolean;
+	mimeType: ImageGenerationDetails["mimeType"];
 	width: number;
 	height: number;
 }): string {
 	const lines = [
-		`${args.edited ? "Edited" : "Generated"} PNG image (${args.width}x${args.height}).`,
+		`${args.edited ? "Edited" : "Generated"} ${formatImageMimeType(args.mimeType)} image (${args.width}x${args.height}).`,
 		`Artifact: ${args.artifactPath}`,
 	];
 	if (args.outputPath) lines.push(`Copied to: ${args.outputPath}`);
@@ -142,7 +153,7 @@ function formatResultText(args: {
 }
 
 function formatBatchResultText(details: ImageGenerationBatchDetails): string {
-	const lines = [`Generated ${details.succeededCount}/${details.requestedCount} PNG images.`];
+	const lines = [`Generated ${details.succeededCount}/${details.requestedCount} images.`];
 	const failedIndexes = new Set((details.failures ?? []).map((failure) => failure.index));
 	let imageIndex = 0;
 	for (let batchIndex = 0; batchIndex < details.requestedCount; batchIndex += 1) {
@@ -183,13 +194,18 @@ async function persistGeneratedImage(args: {
 			agentDir: args.agentDir,
 			sessionId: args.sessionId,
 			imageCallId,
+			mimeType: args.generated.mimeType,
 		});
 
 		let outputPath: string | undefined;
 		let warning: string | undefined;
 		if (args.explicitOutput) {
 			try {
-				await args.deps.copyExplicit({ bytes: args.generated.bytes, plan: args.explicitOutput });
+				await args.deps.copyExplicit({
+					bytes: args.generated.bytes,
+					plan: args.explicitOutput,
+					mimeType: args.generated.mimeType,
+				});
 				outputPath = args.explicitOutput.path;
 			} catch (error) {
 				warning = sanitizeImageDiagnostic(
@@ -206,7 +222,7 @@ async function persistGeneratedImage(args: {
 			imageModel: args.imageModel,
 			imageCallId,
 			...(args.generated.responseId ? { responseId: args.generated.responseId } : {}),
-			mimeType: IMAGE_GENERATION_MIME_TYPE,
+			mimeType: args.generated.mimeType,
 			byteCount: args.generated.bytes.length,
 			width: args.generated.width,
 			height: args.generated.height,
@@ -235,7 +251,18 @@ export function createImageGenerationExecutor(
 		notifyConfigIssues(args.ctx, resolved);
 		assertConfigValid(resolved, "imageGeneration", "compatibility");
 		const { config } = resolved;
-		if (!isImageGenerationEnabledForModel(args.ctx.model, config.imageGeneration)) {
+		let effectiveModel = args.ctx.model;
+		if (args.ctx.model?.api === VIRTUAL_MODEL_API) {
+			const effective = await resolveEffectiveModel(args.ctx);
+			if (!effective.ok) {
+				throw new ImageGenerationError(
+					"unsupported-model",
+					"Image generation is unavailable because the virtual model could not be resolved to a physical model.",
+				);
+			}
+			effectiveModel = effective.model;
+		}
+		if (!isImageGenerationEnabledForModel(args.ctx.model, config.imageGeneration, effectiveModel)) {
 			throw new ImageGenerationError(
 				"unsupported-model",
 				"Image generation is not enabled for the current provider/model-id.",
@@ -332,6 +359,7 @@ export function createImageGenerationExecutor(
 						outputPath: details.outputPath,
 						warning: details.warning,
 						edited: details.edited,
+						mimeType: details.mimeType,
 						width: details.width,
 						height: details.height,
 					}),

package/src/image-generation/types.ts

@@ -1,4 +1,7 @@
 export const IMAGE_GENERATION_TOOL_NAME = "openai_generate_image";
+/** Image formats accepted from independent Images API responses. */
+export const IMAGE_GENERATION_MIME_TYPES = ["image/png", "image/jpeg", "image/webp"] as const;
+/** Backward-compatible default format for callers that construct PNG-only fixtures. */
 export const IMAGE_GENERATION_MIME_TYPE = "image/png" as const;
 export const IMAGE_GENERATION_CAPABLE_APIS = [
 	"openai-responses",
@@ -32,7 +35,19 @@ export type ImageGenerationCapableApi = (typeof IMAGE_GENERATION_CAPABLE_APIS)[n
 export type ImageGenerationSize = (typeof IMAGE_GENERATION_SIZES)[number];
 export type ImageGenerationQuality = (typeof IMAGE_GENERATION_QUALITIES)[number];
 export type ImageGenerationAction = "generate" | "edit";
-export type ReferenceImageMimeType = "image/png" | "image/jpeg" | "image/webp";
+export type ImageGenerationMimeType = (typeof IMAGE_GENERATION_MIME_TYPES)[number];
+export type ReferenceImageMimeType = ImageGenerationMimeType;
+
+export function imageFileExtension(mimeType: ImageGenerationMimeType): "png" | "jpg" | "webp" {
+	switch (mimeType) {
+		case "image/png":
+			return "png";
+		case "image/jpeg":
+			return "jpg";
+		case "image/webp":
+			return "webp";
+	}
+}
 
 export type ImageGenerationModel = {
 	api?: string;
@@ -72,6 +87,7 @@ export type PreparedReferenceImage = {
 
 export type ParsedGeneratedImage = {
 	bytes: Buffer;
+	mimeType: ImageGenerationMimeType;
 	imageCallId: string;
 	responseId?: string;
 	revisedPrompt?: string;
@@ -86,7 +102,7 @@ export type ImageGenerationDetails = {
 	imageModel: string;
 	imageCallId: string;
 	responseId?: string;
-	mimeType: typeof IMAGE_GENERATION_MIME_TYPE;
+	mimeType: ImageGenerationMimeType;
 	byteCount: number;
 	width: number;
 	height: number;
@@ -188,6 +204,10 @@ function isBoundedString(value: unknown, maxChars: number): value is string {
 	return typeof value === "string" && value.trim().length > 0 && value.length <= maxChars;
 }
 
+function isImageGenerationMimeType(value: unknown): value is ImageGenerationMimeType {
+	return (IMAGE_GENERATION_MIME_TYPES as readonly string[]).includes(value as string);
+}
+
 export function isImageGenerationDetails(value: unknown): value is ImageGenerationDetails {
 	if (!isRecord(value)) return false;
 	return (
@@ -197,7 +217,7 @@ export function isImageGenerationDetails(value: unknown): value is ImageGenerati
 		isBoundedString(value.imageModel, MAX_IMAGE_MODEL_ID_CHARS) &&
 		isBoundedString(value.imageCallId, MAX_IMAGE_IDENTIFIER_CHARS) &&
 		(value.responseId === undefined || isBoundedString(value.responseId, MAX_IMAGE_IDENTIFIER_CHARS)) &&
-		value.mimeType === IMAGE_GENERATION_MIME_TYPE &&
+		isImageGenerationMimeType(value.mimeType) &&
 		typeof value.byteCount === "number" &&
 		Number.isInteger(value.byteCount) &&
 		value.byteCount > 0 &&

package/src/payload-rewrite.ts

@@ -9,7 +9,13 @@ import type {
 	NativeCompactionEntry,
 	NativeCompactionInputProvenance,
 } from "./types";
-import { serializeMessagesToResponsesInput, type ResponsesInputContentItem, type ResponsesInputItem, type ResponsesInputMessageItem } from "./serializer";
+import {
+	serializeMessagesToResponsesInput,
+	serializeMessagesToResponsesInputWithLimits,
+	type ResponsesInputContentItem,
+	type ResponsesInputItem,
+	type ResponsesInputMessageItem,
+} from "./serializer";
 
 export type NativeReplaySegments = {
 	boundaryIndex: number;
@@ -369,6 +375,15 @@ export function serializeLiveTailToResponsesInput<TApi extends Api>(args: {
 	});
 }
 
+export async function serializeLiveTailToResponsesInputWithLimits<TApi extends Api>(args: {
+	model: Model<TApi>;
+	entries: readonly SessionEntry[];
+}): Promise<ResponsesInputItem[]> {
+	return serializeMessagesToResponsesInputWithLimits(args.model, collectReplayMessages(args.entries), {
+		firstSystemMessageIsUpdate: true,
+	});
+}
+
 function buildNativeReplaySegmentsInternal<TApi extends Api>(args: {
 	model: Model<TApi>;
 	payload: ResponsesCompatibleRequestPayload;

package/src/responses-headers.ts

@@ -127,7 +127,12 @@ export function buildResponsesRequestHeaders(
 	} else {
 		headers.set("content-type", options.contentType ?? "application/json");
 	}
-	if (!headers.has("authorization")) {
+	if (runtime.api === "azure-openai-responses") {
+		// Azure OpenAI authenticates Responses requests with `api-key`; do not
+		// forward an inherited bearer token from a model/provider header.
+		headers.delete("authorization");
+		headers.set("api-key", runtime.apiKey);
+	} else if (!headers.has("authorization")) {
 		headers.set("authorization", `Bearer ${runtime.apiKey}`);
 	}
 

package/src/responses-websocket/extension.ts

@@ -0,0 +1,32 @@
+import { getApiProvider, registerApiProvider, unregisterApiProviders } from "@earendil-works/pi-ai/compat";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { loadToolkitConfig, resolveToolkitConfig } from "../config";
+import { createResponsesWebSocketProvider } from "./transport";
+
+const SOURCE_ID = "pi-openai-toolkit:openai-responses-websocket";
+
+/**
+ * Add the transport branch at the API-adapter boundary instead of registering
+ * a provider. Existing provider ids, models, credentials, and base URLs remain
+ * untouched; only openai-responses models eligible for transport=auto/websocket
+ * are handled by the wrapper.
+ */
+export function installResponsesWebSocketTransport(
+	loadConfig: typeof loadToolkitConfig = loadToolkitConfig,
+): boolean {
+	// Re-installation can happen during Pi reloads. Remove only this extension's prior
+	// adapter so a config change cannot leave a stale enabled wrapper behind.
+	unregisterApiProviders(SOURCE_ID);
+	const resolved = resolveToolkitConfig(loadConfig());
+	if (!resolved.policy.responsesWebSocket.enabled || resolved.invalidFeatures.includes("responsesWebSocket")) {
+		return false;
+	}
+	const original = getApiProvider("openai-responses");
+	if (!original) return false;
+	registerApiProvider(createResponsesWebSocketProvider(original, true), SOURCE_ID);
+	return true;
+}
+
+export default function responsesWebSocketExtension(_pi: ExtensionAPI): void {
+	installResponsesWebSocketTransport();
+}