pi-openai-toolkit 0.7.1 → 0.8.0

package/src/payload-rewrite.ts

@@ -9,31 +9,16 @@ import type {
 import type { ResponsesCompatibleRequestPayload } from "./runtime";
 import { rewritePayloadWithDeferredToolCarryover } from "./deferred-tool-carryover";
 import type { NativeCompactionEntry } from "./types";
-import {
-	compareResponsesInputParity,
-	serializeMessagesToResponsesInput,
-	type ResponsesInputContentItem,
-	type ResponsesInputItem,
-	type ResponsesInputMessageItem,
-} from "./serializer";
-
-export type SerializedReplaySlice = {
-	entries: SessionEntry[];
-	messages: AgentMessage[];
-	input: ResponsesInputItem[];
-};
+import { serializeMessagesToResponsesInput, type ResponsesInputContentItem, type ResponsesInputItem, type ResponsesInputMessageItem } from "./serializer";
 
 export type NativeReplaySegments = {
 	boundaryIndex: number;
 	firstKeptEntryIndex: number;
 	instructions?: string;
-	freshPreamble: ResponsesInputMessageItem[];
-	trailingPreamble: ResponsesInputMessageItem[];
+	leading: ResponsesInputItem[];
 	compactionSummary: ResponsesInputItem[];
-	preCompactionKeptWindow: SerializedReplaySlice;
 	compactedWindow: unknown[];
-	postCompactionTail: SerializedReplaySlice;
-	originalPiReplayInput: ResponsesInputItem[];
+	post: ResponsesInputItem[];
 	replayInput: unknown[];
 };
 
@@ -48,7 +33,6 @@ export type NativeReplayPayloadRewriteFailureReason =
 	| "first-kept-entry-not-found"
 	| "invalid-compacted-window"
 	| "unexpected-compaction-after-boundary"
-	| "expected-pi-replay-mismatch"
 	| "compaction-summary-not-found";
 
 export type NativeReplayPayloadRewriteFailure = {
@@ -89,19 +73,6 @@ function isResponsesInputMessageRole(value: unknown): value is ResponsesInputMes
 	return value === "user" || value === "developer" || value === "system";
 }
 
-function isPreambleRole(value: ResponsesInputMessageItem["role"]): value is "developer" | "system" {
-	return value === "developer" || value === "system";
-}
-
-function isResponsesInputMessageItem(value: unknown): value is ResponsesInputMessageItem {
-	if (!isRecord(value) || !isResponsesInputMessageRole(value.role)) {
-		return false;
-	}
-
-	const { content } = value;
-	return typeof content === "string" || (Array.isArray(content) && content.every(isResponsesInputContentItem));
-}
-
 function cloneResponsesInputContentItem(item: ResponsesInputContentItem): ResponsesInputContentItem {
 	return item.type === "input_text"
 		? {
@@ -294,20 +265,6 @@ function findCompactionSummaryIndex(input: readonly unknown[], summaryMarker: st
 	});
 }
 
-/** Count consecutive provider-authored developer/system items at the end of the input. */
-function countTrailingPreambleItems(items: readonly unknown[]): number {
-	let count = 0;
-	for (let index = items.length - 1; index >= 0; index--) {
-		const item = items[index];
-		if (isResponsesInputMessageItem(item) && isPreambleRole(item.role)) {
-			count += 1;
-			continue;
-		}
-		break;
-	}
-	return count;
-}
-
 function collectReplayMessages(entries: readonly SessionEntry[]): AgentMessage[] {
 	const messages: AgentMessage[] = [];
 
@@ -321,27 +278,6 @@ function collectReplayMessages(entries: readonly SessionEntry[]): AgentMessage[]
 	return messages;
 }
 
-function createCompactionSummaryAgentMessage(entry: NativeCompactionEntry): AgentMessage {
-	return {
-		role: "compactionSummary",
-		summary: entry.summary,
-		tokensBefore: entry.tokensBefore,
-		timestamp: new Date(entry.timestamp).getTime(),
-	} as AgentMessage;
-}
-
-function createReplaySlice(
-	entries: readonly SessionEntry[],
-	messages: readonly AgentMessage[],
-	input: readonly ResponsesInputItem[],
-): SerializedReplaySlice {
-	return {
-		entries: [...entries],
-		messages: [...messages],
-		input: [...input],
-	};
-}
-
 function findEntryIndexByIdBeforeBoundary(
 	entries: readonly SessionEntry[],
 	entryId: string,
@@ -436,14 +372,13 @@ function buildNativeReplaySegmentsInternal<TApi extends Api>(args: {
 
 	const preCompactionEntries = args.branchEntries.slice(firstKeptEntryIndex, boundaryIndex);
 	const postCompactionEntries = args.branchEntries.slice(boundaryIndex + 1);
-	const preCompactionKeptMessages = collectReplayMessages(preCompactionEntries);
-	const postCompactionTailMessages = collectReplayMessages(postCompactionEntries);
-	const compactionSummaryMessage = createCompactionSummaryAgentMessage(args.compactionEntry);
-	const keptWindowSerialized = serializeMessagesToResponsesInput(args.model, preCompactionKeptMessages);
-
-	// Pi-authored compaction summary item. Everything outside the replaced region
-	// is preserved verbatim, so this content anchor is the only location probe
-	// needed; the marker text is unique to this compaction entry.
+
+	// Pi-authored compaction summary item. Pi 0.84.3 renders the context after a
+	// compaction as [leading prompts][summary item][all post-compaction messages]:
+	// pre-compaction messages never appear in the payload (entry.retainedTail is
+	// never persisted by appendCompaction), so the summary item itself is the only
+	// anchor and the only thing that gets replaced. Everything else is preserved
+	// verbatim, so transient messages injected by other extensions stay intact.
 	const summaryIndex = findCompactionSummaryIndex(args.payload.input, args.compactionEntry.summary);
 	if (summaryIndex < 0) {
 		return {
@@ -452,66 +387,19 @@ function buildNativeReplaySegmentsInternal<TApi extends Api>(args: {
 		};
 	}
 
-	const keptStart = summaryIndex + 1;
-	const keptEnd = keptStart + keptWindowSerialized.length;
-	const actualKeptWindow = cloneResponsesInputSlice(args.payload.input.slice(keptStart, keptEnd));
-	if (!actualKeptWindow) {
-		return {
-			ok: false,
-			reason: "expected-pi-replay-mismatch",
-		};
-	}
-
-	// The only hard verification: the region immediately after the sentinel must be
-	// exactly the kept pre-compaction window derived from session entries. Content
-	// after this region is preserved as-is, so transient messages injected by other
-	// extensions (which have no session entries) no longer break replay.
-	if (!areEquivalentValues(actualKeptWindow, keptWindowSerialized)) {
-		const parity = compareResponsesInputParity(
-			args.payload.input.slice(keptStart, keptEnd),
-			keptWindowSerialized,
-		);
-		return {
-			ok: false,
-			reason: "expected-pi-replay-mismatch",
-			parity: {
-				actual: parity.actual,
-				expected: parity.expected,
-				mismatches: parity.mismatches,
-			},
-		};
-	}
-
-	const actualCompactionSummary = cloneResponsesInputSlice([args.payload.input[summaryIndex]]);
-	const actualPostCompactionTail = cloneResponsesInputSlice(args.payload.input.slice(keptEnd));
-	if (!actualCompactionSummary || !actualPostCompactionTail) {
-		return {
-			ok: false,
-			reason: "expected-pi-replay-mismatch",
-		};
-	}
-
-	// Keep provider-authored trailing developer/system prompts in place.
-	const trailingCount = countTrailingPreambleItems(actualPostCompactionTail);
-	const trailingPreamble = actualPostCompactionTail.slice(actualPostCompactionTail.length - trailingCount);
-	const liveTail = actualPostCompactionTail.slice(0, actualPostCompactionTail.length - trailingCount);
-
 	const leadingInput = cloneResponsesInputSlice(args.payload.input.slice(0, summaryIndex));
-	if (!leadingInput) {
+	const actualCompactionSummary = cloneResponsesInputSlice([args.payload.input[summaryIndex]]);
+	const actualPostInput = cloneResponsesInputSlice(args.payload.input.slice(summaryIndex + 1));
+	if (!leadingInput || !actualCompactionSummary || !actualPostInput) {
 		return {
 			ok: false,
-			reason: "expected-pi-replay-mismatch",
+			reason: "compaction-summary-not-found",
 		};
 	}
 
 	const baseRewrittenPayload: ResponsesCompatibleRequestPayload = {
 		...args.payload,
-		input: [
-			...leadingInput,
-			...compactedWindow,
-			...liveTail,
-			...trailingPreamble,
-		],
+		input: [...leadingInput, ...compactedWindow, ...actualPostInput],
 	};
 	const carryoverRewrite = rewritePayloadWithDeferredToolCarryover({
 		payload: baseRewrittenPayload,
@@ -523,33 +411,16 @@ function buildNativeReplaySegmentsInternal<TApi extends Api>(args: {
 			| undefined,
 	});
 
-	const sentinelSerialized = serializeMessagesToResponsesInput(args.model, [compactionSummaryMessage]);
-	const serializedTailInput = serializeMessagesToResponsesInput(args.model, postCompactionTailMessages);
-	const originalPiReplayInput: ResponsesInputItem[] = [
-		...leadingInput,
-		...sentinelSerialized,
-		...keptWindowSerialized,
-		...serializedTailInput,
-		...trailingPreamble,
-	];
-
 	return {
 		ok: true,
 		segments: {
 			boundaryIndex,
 			firstKeptEntryIndex,
 			instructions: typeof args.payload.instructions === "string" ? args.payload.instructions : undefined,
-			freshPreamble: leadingInput as ResponsesInputMessageItem[],
-			trailingPreamble: trailingPreamble as ResponsesInputMessageItem[],
+			leading: leadingInput,
 			compactionSummary: actualCompactionSummary,
-			preCompactionKeptWindow: createReplaySlice(
-				preCompactionEntries,
-				preCompactionKeptMessages,
-				actualKeptWindow,
-			),
 			compactedWindow,
-			postCompactionTail: createReplaySlice(postCompactionEntries, postCompactionTailMessages, liveTail),
-			originalPiReplayInput,
+			post: actualPostInput,
 			replayInput: carryoverRewrite.payload.input,
 		},
 		rewrittenPayload: carryoverRewrite.payload,

package/src/remote-v2-client.ts

@@ -1,5 +1,5 @@
 import { writeDebugArtifact } from "./debug";
-import { mergeProviderHeaders } from "./provider-headers";
+import { buildResponsesRequestHeaders } from "./responses-headers";
 import type { NativeCompactionRuntime } from "./runtime";
 import type { NativeCompactionRequestBody, ResponsesInputItem } from "./serializer";
 import type { ArtifactContext, CompactionConfig } from "./types";
@@ -114,57 +114,13 @@ function normalizeResponseTimestamp(value: unknown): string | undefined {
 	return Number.isNaN(parsed) ? trimmed : new Date(parsed).toISOString();
 }
 
-function decodeJwtPayload(token: string): Record<string, unknown> | undefined {
-	const parts = token.split(".");
-	if (parts.length !== 3) {
-		return undefined;
-	}
-
-	try {
-		const payloadText = Buffer.from(parts[1]!, "base64url").toString("utf8");
-		const payload = JSON.parse(payloadText);
-		return isRecord(payload) ? payload : undefined;
-	} catch {
-		return undefined;
-	}
-}
-
-function extractCodexAccountId(token: string): string | undefined {
-	const payload = decodeJwtPayload(token);
-	const authClaims = payload?.["https://api.openai.com/auth"];
-	if (!isRecord(authClaims)) {
-		return undefined;
-	}
-
-	const accountId = authClaims.chatgpt_account_id;
-	return isNonEmptyString(accountId) ? accountId.trim() : undefined;
-}
-
-function buildCodexUserAgent(): string {
-	const platform = typeof process !== "undefined" ? process.platform : "browser";
-	const arch = typeof process !== "undefined" ? process.arch : "unknown";
-	return `pi (${platform}; ${arch})`;
-}
-
 function toHeaders(runtime: NativeCompactionRuntime): Record<string, string> {
-	const headers = new Headers(mergeProviderHeaders(runtime.currentModel.headers, runtime.headers));
-	headers.set("accept", SSE_CONTENT_TYPE);
-	headers.set("content-type", JSON_CONTENT_TYPE);
-	if (!headers.has("authorization")) {
-		headers.set("authorization", `Bearer ${runtime.apiKey}`);
-	}
-
-	if (runtime.api === "openai-codex-responses") {
-		const accountId = extractCodexAccountId(runtime.apiKey);
-		if (accountId) {
-			headers.set("chatgpt-account-id", accountId);
-		}
-		headers.set("originator", "pi");
-		headers.set("user-agent", buildCodexUserAgent());
-		headers.set("openai-beta", "responses=experimental");
-	}
-
-	return Object.fromEntries(headers.entries());
+	return Object.fromEntries(
+		buildResponsesRequestHeaders(runtime, {
+			accept: SSE_CONTENT_TYPE,
+			contentType: JSON_CONTENT_TYPE,
+		}).entries(),
+	);
 }
 
 export function buildRemoteV2CompactionRequest(

package/src/responses-headers.ts

@@ -0,0 +1,67 @@
+import type { ProviderHeaders } from "@earendil-works/pi-ai";
+import { mergeProviderHeaders } from "./provider-headers";
+
+export type ResponsesHeaderRuntime = {
+	api: string;
+	apiKey: string;
+	headers?: ProviderHeaders;
+	currentModel: {
+		headers?: ProviderHeaders;
+	};
+};
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return !!value && typeof value === "object" && !Array.isArray(value);
+}
+
+function decodeJwtPayload(token: string): Record<string, unknown> | undefined {
+	const parts = token.split(".");
+	if (parts.length !== 3) return undefined;
+	try {
+		const payload = JSON.parse(Buffer.from(parts[1]!, "base64url").toString("utf8")) as unknown;
+		return isRecord(payload) ? payload : undefined;
+	} catch {
+		return undefined;
+	}
+}
+
+function extractCodexAccountId(token: string): string | undefined {
+	const authClaims = decodeJwtPayload(token)?.["https://api.openai.com/auth"];
+	if (!isRecord(authClaims)) return undefined;
+	const accountId = authClaims.chatgpt_account_id;
+	return typeof accountId === "string" && accountId.trim() ? accountId.trim() : undefined;
+}
+
+function buildCodexUserAgent(): string {
+	const platform = typeof process !== "undefined" ? process.platform : "browser";
+	const arch = typeof process !== "undefined" ? process.arch : "unknown";
+	return `pi (${platform}; ${arch})`;
+}
+
+export function buildResponsesRequestHeaders(
+	runtime: ResponsesHeaderRuntime,
+	options: { accept: string; contentType?: string },
+): Headers {
+	const headers = new Headers(
+		mergeProviderHeaders(runtime.currentModel.headers, runtime.headers),
+	);
+	headers.set("accept", options.accept);
+	headers.set("content-type", options.contentType ?? "application/json");
+	if (!headers.has("authorization")) {
+		headers.set("authorization", `Bearer ${runtime.apiKey}`);
+	}
+
+	if (runtime.api === "openai-codex-responses") {
+		const accountId = extractCodexAccountId(runtime.apiKey);
+		if (accountId) headers.set("chatgpt-account-id", accountId);
+		headers.set("originator", "pi");
+		headers.set("user-agent", buildCodexUserAgent());
+		headers.set("openai-beta", "responses=experimental");
+	}
+
+	return headers;
+}
+
+export const _responsesHeadersTest = {
+	extractCodexAccountId,
+};

package/src/runtime.ts

@@ -7,8 +7,9 @@ const CODEX_RESPONSES_PATH = "codex/responses";
 const OPENAI_COMPACT_PATH = "responses/compact";
 const CODEX_COMPACT_PATH = "codex/responses/compact";
 
-type ResponsesCompactApi = (typeof RESPONSES_COMPACT_CAPABLE_APIS)[number];
+export type ResponsesApi = (typeof RESPONSES_COMPACT_CAPABLE_APIS)[number];
 
+type ResponsesCompactApi = ResponsesApi;
 type RuntimeModel = Model<Api>;
 
 export type ParsedModelSpec = {
@@ -35,6 +36,8 @@ export type NativeCompactionSupportOptions = {
 	responsesApis?: readonly string[];
 };
 
+export type ResponsesSupportOptions = NativeCompactionSupportOptions;
+
 export type ResponsesCompatibleRequestPayload = {
 	model: string;
 	input: unknown[];
@@ -42,19 +45,22 @@ export type ResponsesCompatibleRequestPayload = {
 	[key: string]: unknown;
 };
 
-export type NativeCompactionRuntime = {
+export type ResponsesRuntime = {
 	provider: string;
-	api: ResponsesCompactApi;
+	api: ResponsesApi;
 	model: string;
 	baseUrl: string;
 	apiKey: string;
 	headers?: ProviderHeaders;
 	responsesPath: string;
 	responsesUrl: string;
+	currentModel: RuntimeModel;
+};
+
+export type NativeCompactionRuntime = ResponsesRuntime & {
 	compactPath: string;
 	compactUrl: string;
 	payload?: ResponsesCompatibleRequestPayload;
-	currentModel: RuntimeModel;
 };
 
 export type NativeCompactionEnvironmentFailure = {
@@ -77,6 +83,15 @@ export type NativeCompactionEnvironmentResolution =
 	| NativeCompactionEnvironmentFailure
 	| NativeCompactionEnvironmentSuccess;
 
+export type ResponsesEnvironmentSuccess = {
+	ok: true;
+	runtime: ResponsesRuntime;
+};
+
+export type ResponsesEnvironmentResolution =
+	| NativeCompactionEnvironmentFailure
+	| ResponsesEnvironmentSuccess;
+
 export type RemoteCompactionExecution = {
 	/** Active session model. Its identity owns replay matching and is never mutated. */
 	consumer: NativeCompactionRuntime;
@@ -327,6 +342,30 @@ async function resolveNativeCompactionEnvironmentForModel(
 	};
 }
 
+export async function resolveResponsesEnvironment(
+	ctx: ExtensionContext,
+	options: ResponsesSupportOptions = {},
+): Promise<ResponsesEnvironmentResolution> {
+	const resolution = await resolveNativeCompactionEnvironmentForModel(ctx, ctx.model, options);
+	if (!resolution.ok) return resolution;
+
+	const runtime = resolution.runtime;
+	return {
+		ok: true,
+		runtime: {
+			provider: runtime.provider,
+			api: runtime.api,
+			model: runtime.model,
+			baseUrl: runtime.baseUrl,
+			apiKey: runtime.apiKey,
+			headers: runtime.headers,
+			responsesPath: runtime.responsesPath,
+			responsesUrl: runtime.responsesUrl,
+			currentModel: runtime.currentModel,
+		},
+	};
+}
+
 export async function resolveNativeCompactionEnvironment(
 	ctx: ExtensionContext,
 	options: NativeCompactionSupportOptions = {},

package/src/types.ts

@@ -4,6 +4,7 @@ import type { CompactionEntry, CompactionResult, ExtensionContext, SessionEntry
 export const TOOLKIT_ID = "pi-openai-toolkit";
 export const COMPACTION_EXTENSION_ID = `${TOOLKIT_ID}:compaction`;
 export const WEB_SEARCH_EXTENSION_ID = `${TOOLKIT_ID}:web-search`;
+export const IMAGE_GENERATION_EXTENSION_ID = `${TOOLKIT_ID}:image-generation`;
 export const DEFAULT_ARTIFACT_ROOT = "~/.pi/agent/artifacts/pi-openai-toolkit/compaction";
 export const REDACTED_VALUE = "[REDACTED]";
 /**
@@ -82,9 +83,16 @@ export type WebSearchConfig = {
 	models: string[];
 };
 
+export type ImageGenerationConfig = {
+	enabled: boolean;
+	/** Exact provider/model keys allowed to use toolkit-native image generation. */
+	models: string[];
+};
+
 export type ToolkitConfig = {
 	compaction: CompactionConfig;
 	webSearch: WebSearchConfig;
+	imageGeneration: ImageGenerationConfig;
 };
 
 export type LoadedToolkitConfig = {
@@ -438,7 +446,13 @@ export const DEFAULT_WEB_SEARCH_CONFIG: WebSearchConfig = {
 	models: [],
 };
 
+export const DEFAULT_IMAGE_GENERATION_CONFIG: ImageGenerationConfig = {
+	enabled: true,
+	models: [],
+};
+
 export const DEFAULT_TOOLKIT_CONFIG: ToolkitConfig = {
 	compaction: DEFAULT_COMPACTION_CONFIG,
 	webSearch: DEFAULT_WEB_SEARCH_CONFIG,
+	imageGeneration: DEFAULT_IMAGE_GENERATION_CONFIG,
 };

package/src/web-search/types.ts

@@ -1,3 +1,4 @@
+import { getExactModelKey, isExactModelAllowed } from "../model-scope";
 import type { WebSearchConfig } from "../types";
 
 export const WEB_SEARCH_CAPABLE_APIS = ["openai-responses", "openai-codex-responses"] as const;
@@ -28,20 +29,17 @@ export type WebSearchPayloadTransform = {
 };
 
 export function getWebSearchModelKey(model: WebSearchModel | undefined): string | undefined {
-	if (!model?.provider || !model.id) return undefined;
-	return `${model.provider}/${model.id}`;
+	return getExactModelKey(model);
 }
 
 export function isWebSearchEnabledForModel(
 	model: WebSearchModel | undefined,
 	config: WebSearchConfig,
 ): boolean {
-	const modelKey = getWebSearchModelKey(model);
 	return (
 		config.enabled &&
 		typeof model?.api === "string" &&
 		(WEB_SEARCH_CAPABLE_APIS as readonly string[]).includes(model.api) &&
-		modelKey !== undefined &&
-		config.models.includes(modelKey)
+		isExactModelAllowed(model, config.models)
 	);
 }

package/tsconfig.check.json

@@ -14,6 +14,7 @@
   "include": [
     "extensions/compaction.ts",
     "extensions/web-search.ts",
+    "extensions/image-generation.ts",
     "src/compact-client.ts",
     "src/config.ts",
     "src/debug.ts",
@@ -21,9 +22,20 @@
     "src/deferred-tool-carryover.ts",
     "src/inline-compaction.ts",
     "src/extension-runtime.ts",
+    "src/image-generation/artifacts.ts",
+    "src/image-generation/client.ts",
+    "src/image-generation/eligibility.ts",
+    "src/image-generation/extension.ts",
+    "src/image-generation/protocol.ts",
+    "src/image-generation/references.ts",
+    "src/image-generation/render.ts",
+    "src/image-generation/service.ts",
+    "src/image-generation/types.ts",
+    "src/model-scope.ts",
     "src/native-fallback.ts",
     "src/payload-rewrite.ts",
     "src/provider-headers.ts",
+    "src/responses-headers.ts",
     "src/remote-v2-client.ts",
     "src/request-context-cache.ts",
     "src/runtime.ts",