@oai404iao/pi-codex-runtime 0.2.0 → 0.3.1

package/README.md

@@ -4,12 +4,17 @@ Alpha bootstrap candidate; initial npm publication is still pending.
 It has no Pi extension entry
 and does not automatically register tools or providers.
 
-Peer floor: Pi 0.85.1; tested against 0.85.1.
+Peer floor: Pi 0.86.1; tested against 0.86.1.
 
 Owns shared Codex authentication/headers, wire identity, settings/catalog,
 Responses replay contracts and the session-scoped composition broker.
 Capability-specific HTTP clients, storage and transports live in other packages.
 
+Responses helpers encode/decode generic grammar tools and replay calls/results
+according to the current declaration, preserving legacy custom patch behavior.
+An internal structural event adapter lets owners offer optional Code Mode
+contributions without importing or depending on the private Code Mode package.
+
 Configuration paths remain
 `<agentDir>/extensions/pi-codex-minimal-tools/{config,models}.json`.
 This package ships the canonical schemas and default catalog, including the

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@oai404iao/pi-codex-runtime",
-  "version": "0.2.0",
+  "version": "0.3.1",
   "private": false,
   "description": "Unpublished pi-codex-runtime workspace for the staged Codex split",
   "type": "module",
@@ -23,8 +23,8 @@
     "./subagent-inline": "./src/subagent-inline.ts"
   },
   "peerDependencies": {
-    "@earendil-works/pi-ai": ">=0.85.1",
-    "@earendil-works/pi-coding-agent": ">=0.85.1"
+    "@earendil-works/pi-ai": ">=0.86.1",
+    "@earendil-works/pi-coding-agent": ">=0.86.1"
   },
   "peerDependenciesMeta": {
     "@earendil-works/pi-ai": {
@@ -36,8 +36,8 @@
   },
   "dependencies": {},
   "devDependencies": {
-    "@earendil-works/pi-ai": "0.85.1",
-    "@earendil-works/pi-coding-agent": "0.85.1",
+    "@earendil-works/pi-ai": "0.86.1",
+    "@earendil-works/pi-coding-agent": "0.86.1",
     "@types/node": "^26.2.0",
     "tsx": "^4.20.6",
     "typescript": "^7.0.2"
@@ -64,5 +64,5 @@
   "engines": {
     "node": ">=22.19.0"
   },
-  "gitHead": "61f1d82471fff63e3e9d5c0276b5e069c9fc8f1f"
+  "gitHead": "c791e1192d0547834c3e2f2bae7853fb214f9808"
 }

package/src/broker.ts

@@ -2,6 +2,7 @@ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { createRequire } from "node:module";
 import type { PackageToolName } from "./capabilities.js";
 import type { ProviderPresentation } from "./extension/provider-presentation.js";
+import type { OwnerState } from "./code-mode-owner.js";
 
 export const CODEX_BROKER_CHANNEL = "@oai404iao/pi-codex:broker";
 export const CODEX_RUNTIME_VERSION: string = createRequire(import.meta.url)("../package.json").version;
@@ -10,6 +11,7 @@ const CACHE = Symbol.for("@oai404iao/pi-codex/broker/v1");
 export interface InstalledTool {
 	register(): void;
 	registered: boolean;
+	codeModeOwner?: OwnerState;
 }
 
 export interface CodexBroker {
@@ -18,6 +20,7 @@ export interface CodexBroker {
 	readonly closed: boolean;
 	readonly tools: Map<PackageToolName, InstalledTool>;
 	coreEnabled: boolean;
+	codeModeDefinitions?: Map<string, { parameters: unknown; description: string; providerId: string }>;
 	claim(name: string): boolean;
 	addPresentation(name: string, presentation: ProviderPresentation): void;
 	presentation: ProviderPresentation;

package/src/code-mode-contributions.ts

@@ -0,0 +1,139 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { getCodexBroker } from "./broker.js";
+import { codeModeOwner, type DirectBinding } from "./code-mode-owner.js";
+import type { PackageToolName } from "./capabilities.js";
+import { randomUUID } from "node:crypto";
+
+/** Unique public schema reference proves which registration won Pi's registry.
+ * If a future Pi clones metadata, cooperation fails closed instead of claiming
+ * a name/source belonging to another extension. */
+export function registerCodeModeOwnedTool(pi: ExtensionAPI, definition: Record<string, unknown>, providerId: string): void {
+	if (typeof definition.name !== "string" || typeof definition.description !== "string"
+		|| !definition.parameters || typeof definition.parameters !== "object"
+		|| !/^[a-z][a-z0-9_]{0,39}$/.test(providerId)) throw new Error("Invalid owned Code Mode tool definition");
+	const broker = getCodexBroker(pi);
+	const definitions = broker.codeModeDefinitions ??= new Map();
+	if (definitions.has(definition.name)) throw new Error(`Duplicate Code Mode tool definition: ${definition.name}`);
+	const owned = { ...definition, name: definition.name, description: definition.description, parameters: { ...definition.parameters } };
+	pi.registerTool(owned as never);
+	definitions.set(owned.name, { ...owned, providerId });
+}
+
+export interface NestedCodeModeContext {
+	readonly cellId: string;
+	readonly toolCallId: string;
+	readonly cwd: string;
+	readonly signal: AbortSignal;
+	readonly pi?: ExtensionContext;
+}
+export interface OwnedCodeModeTool {
+	name: string;
+	description: string;
+	parameters: unknown;
+	effect: "read" | "write";
+	parallel?: boolean;
+	requires?: readonly string[];
+	requiredPolicies?: readonly string[];
+	approval?: string;
+	availability?: { state: "available" | "unavailable" | "not-ready" | "failed"; reason?: string };
+	direct?: DirectBinding;
+	invoke(input: unknown, context: NestedCodeModeContext): Promise<{ value: unknown }>;
+}
+const unavailableInvoke = async (): Promise<{ value: unknown }> => { throw new Error("Code Mode tool requirements were not negotiated"); };
+
+/** Structural client for the optional v2/v1 bus contract. Neither side installs
+ * the other package; all authority still comes from Code Mode's exact grants. */
+export function registerCodeModeContribution(pi: ExtensionAPI, id: string,
+	tools: (context: ExtensionContext) => readonly OwnedCodeModeTool[]): void {
+	let context: ExtensionContext | undefined;
+	let disposed = false;
+	let resolved: readonly OwnedCodeModeTool[] | undefined;
+	let registration: Readonly<{ owner: string; instanceId: string; revision: number }> = Object.freeze({ owner: id, instanceId: randomUUID(), revision: 1 });
+	let accepted = new WeakSet<object>();
+	const generations = new Map<string, number>();
+	const changed = (phase: "withdrawn" | "ready" | "disposed") => {
+		const change = Object.freeze({ protocol: 2, registration, kind: "execution", phase });
+		pi.events.emit("@oai404iao/pi-code-mode:changed/v2", change);
+		pi.events.emit("@oai404iao/pi-code-mode:changed/v1", { version: 1, change });
+	};
+	const snapshot = (features: readonly unknown[] = [], legacy = false) => {
+		const broker = getCodexBroker(pi);
+		const offered = context ? (resolved ??= tools(context).map((tool) => ({ ...tool }))) : [];
+		if (offered.length > 64) throw new Error("Code Mode declaration budget exceeded");
+		const declarations = offered.map((tool) => {
+			if (tool.requires !== undefined && (!Array.isArray(tool.requires) || tool.requires.length > 32))
+				throw new Error("Invalid Code Mode feature requirements");
+			const requires = [...new Set([...(tool.requires ?? []), ...(tool.approval ? ["approval/1"] : []),
+				...(tool.requiredPolicies?.length ? ["required-policies/1"] : [])])];
+			if (requires.length > 32 || requires.some((feature) => typeof feature !== "string" || feature.length > 128
+				|| !/^[a-z][a-z0-9-]*\/[1-9][0-9]*$/.test(feature)))
+				throw new Error("Invalid Code Mode feature requirements");
+			return { ...tool, requires };
+		});
+		return { id, tools: declarations.filter((tool) => !legacy || (!tool.requires.length && !tool.requiredPolicies?.length
+			&& tool.approval === undefined && (!tool.availability || tool.availability.state === "available"))).map((tool) => {
+			if (tool.requires.some((feature) => !features.includes(feature)) || (tool.availability && tool.availability.state !== "available"))
+				return { ...tool, direct: undefined, invoke: unavailableInvoke };
+			const owned = broker.tools.get(tool.name as PackageToolName);
+			const definition = broker.codeModeDefinitions?.get(tool.name);
+			return { ...tool, direct: owned && definition?.providerId === id
+				? codeModeOwner(pi, tool.name, owned, definition)?.binding : undefined };
+		}) };
+	};
+	const offV2 = pi.events.on("@oai404iao/pi-code-mode:discover/v2", (value) => {
+		const request = value as { hello?: { protocol?: number; instanceId?: string; generation?: number; features?: unknown[] };
+			offer?: (offer: unknown) => { status: string } } | undefined;
+		const hello = request?.hello;
+		if (disposed || !hello || hello.protocol !== 2 || !hello.instanceId || hello.instanceId.length > 128
+			|| !Number.isSafeInteger(hello.generation) || hello.generation! < 1 || !Array.isArray(hello.features)
+			|| hello.features.length > 32 || typeof request?.offer !== "function") return;
+		const previous = generations.get(hello.instanceId);
+		if ((previous !== undefined && hello.generation! < previous) || (previous === undefined && generations.size >= 64)) return;
+		generations.set(hello.instanceId, hello.generation!);
+		try {
+			const provider = snapshot(hello.features);
+			const availability = { state: context ? "available" : "not-ready" };
+			if (request.offer({ kind: "provider", registration, requires: [], availability, provider }).status === "compatible") accepted.add(hello);
+		} catch {
+			request.offer({ kind: "provider", registration, requires: [], availability: { state: "failed", reason: "owner-not-ready" }, provider: { id, tools: [] } });
+		}
+	});
+	const off = pi.events.on("@oai404iao/pi-code-mode:discover/v1", (value) => {
+		const discovery = value as { version?: number; provider?: (provider: unknown) => void;
+			consumer?: object; receipts?: { registration: object; consumer: object }[] } | undefined;
+		if (!disposed && context && discovery?.version === 1 && typeof discovery.provider === "function") {
+			if (discovery.consumer) {
+				const hello = discovery.consumer as { protocol?: number; instanceId?: string; generation?: number };
+				if (hello.protocol !== 2 || typeof hello.instanceId !== "string" || !hello.instanceId || hello.instanceId.length > 128
+					|| !Number.isSafeInteger(hello.generation) || hello.generation! < 1) return;
+				const previous = generations.get(hello.instanceId);
+				if ((previous !== undefined && hello.generation! < previous) || (previous === undefined && generations.size >= 64)) return;
+				generations.set(hello.instanceId, hello.generation!);
+			}
+			if (discovery.consumer && accepted.has(discovery.consumer) && Array.isArray(discovery.receipts) && discovery.receipts.length <= 80
+				&& discovery.receipts.some((receipt) => receipt.consumer === discovery.consumer && receipt.registration === registration)) return;
+			let provider: ReturnType<typeof snapshot>;
+			// Validate the whole declaration before filtering legacy tools:
+			// a failed v2 provider must not reappear as a partial v1 mirror.
+			try { provider = snapshot([], true); } catch { return; }
+			if (provider.tools.length) discovery.provider(provider);
+		}
+	});
+	const update = (_event: unknown, ctx: ExtensionContext) => {
+		if (disposed) return;
+		context = undefined;
+		resolved = undefined;
+		changed("withdrawn");
+		context = ctx;
+		registration = Object.freeze({ ...registration, revision: registration.revision + 1 });
+		accepted = new WeakSet();
+		changed("ready");
+	};
+	pi.on("session_start", update);
+	pi.on("model_select", update);
+	pi.on("session_tree", update);
+	pi.on("session_shutdown", () => {
+		if (disposed) return;
+		disposed = true; context = undefined; off(); offV2(); changed("disposed");
+	});
+}

package/src/code-mode-owner.ts

@@ -0,0 +1,96 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+export interface DirectBinding {
+	readonly version: 1; readonly name: string;
+	acquire(): { reconcile(): boolean; release(): void } | undefined;
+}
+interface Control {
+	binding: DirectBinding; readonly activeIntent: boolean | undefined;
+	projectActive(active: boolean): boolean | undefined;
+	reconcile(): boolean; dispose(): void;
+}
+interface Factory {
+	version: 1; create(pi: ExtensionAPI, options: { name: string; sourcePath: string }): Control;
+}
+export interface OwnerState {
+	identity: string; replaced?: boolean; factory?: Factory; control?: Control;
+	transitioning?: boolean; pendingDispose?: () => void; closed?: boolean;
+}
+export const OWNER_CHANGED = "@oai404iao/pi-code-mode:direct-owner-changed/v1";
+
+function retire(state: OwnerState): void {
+	if (state.control && !state.pendingDispose) {
+		const old = state.control;
+		state.pendingDispose = () => old.dispose();
+	}
+	state.pendingDispose?.();
+	state.pendingDispose = undefined;
+	state.control = undefined;
+	state.factory = undefined;
+}
+export function disposeCodeModeOwner(tool: { codeModeOwner?: OwnerState }, close = true): void {
+	const state = tool.codeModeOwner;
+	if (!state) return;
+	if (close) state.closed = true;
+	if (state.transitioning) return;
+	state.transitioning = true;
+	try { retire(state); } finally { state.transitioning = false; }
+}
+
+/** Optional v1 bus client. No private-package dependency or global registry. */
+export function codeModeOwner(pi: ExtensionAPI, name: string, tool: { registered: boolean; codeModeOwner?: OwnerState },
+	expected?: { parameters: unknown; description: string }): Control | undefined {
+	if (!["apply_patch", "web_search"].includes(name) || tool.codeModeOwner?.transitioning || tool.codeModeOwner?.closed) return;
+	if (!tool.registered) { disposeCodeModeOwner(tool, false); return; }
+	const info = pi.getAllTools?.().find((item) => item.name === name);
+	if (!info?.sourceInfo || ["builtin", "sdk"].includes(info.sourceInfo.source)
+		|| !expected || info.parameters !== expected.parameters || info.description !== expected.description) {
+		(tool.codeModeOwner ??= { identity: "" }).replaced = true;
+		disposeCodeModeOwner(tool, false);
+		return;
+	}
+	const identity = JSON.stringify([info.sourceInfo, info.description, info.parameters, info.promptGuidelines]);
+	const stillOwns = () => {
+		const current = pi.getAllTools().find((item) => item.name === name);
+		return Boolean(current && current.parameters === expected.parameters
+			&& JSON.stringify([current.sourceInfo, current.description, current.parameters, current.promptGuidelines]) === identity);
+	};
+	const state = tool.codeModeOwner ??= { identity };
+	// Never rebind a replaced registration, even if the optional factory reloads.
+	if (state.identity !== identity) state.replaced = true;
+	if (state.replaced) { disposeCodeModeOwner(tool, false); return; }
+	const factories: Factory[] = [];
+	pi.events.emit("@oai404iao/pi-code-mode:direct-owner/v1", { version: 1, accept(value: Factory) {
+		if (factories.length < 2) factories.push(value);
+	} });
+	const factory = factories.length === 1 && factories[0]?.version === 1
+		&& typeof factories[0].create === "function" ? factories[0] : undefined;
+	if (state.closed) return;
+	if (!stillOwns()) {
+		state.replaced = true;
+		disposeCodeModeOwner(tool, false);
+		return;
+	}
+	if (state.pendingDispose || factory !== state.factory || (factory && !state.control)) {
+		state.transitioning = true;
+		try {
+			retire(state);
+			if (!stillOwns()) state.replaced = true;
+			if (state.closed || state.replaced) return;
+			if (factory) {
+				const control = factory.create(pi, { name, sourcePath: info.sourceInfo.path });
+				if (!control || control.binding?.version !== 1 || control.binding.name !== name
+					|| typeof control.binding.acquire !== "function" || typeof control.projectActive !== "function"
+					|| typeof control.reconcile !== "function" || typeof control.dispose !== "function") {
+					if (typeof control?.dispose === "function") state.pendingDispose = () => control.dispose();
+					throw new Error("Invalid Code Mode owner control");
+				}
+				state.control = control;
+				state.factory = factory;
+				if (!stillOwns()) state.replaced = true;
+				if (state.closed || state.replaced) { retire(state); return; }
+			}
+		} finally { state.transitioning = false; }
+	}
+	return state.control;
+}

package/src/providers/responses/grammar.ts

@@ -0,0 +1,52 @@
+import type { Api, AssistantMessage, AssistantMessageEventStream, Model, Tool } from "@earendil-works/pi-ai";
+import { grammarInputDelta, resolveGrammarSampling } from "./sampling.js";
+import type { CustomToolCallState, OpenAIResponsesStreamOptions } from "./types.js";
+import { localToolName } from "./text.js";
+
+export function supportsGrammar(model: Model<Api>): boolean {
+	return (model.compat as { supportsOpenAIGrammarTools?: boolean } | undefined)?.supportsOpenAIGrammarTools === true;
+}
+
+/** Capture the actual request's custom declarations, including Lite namespaces.
+ * Legacy custom tools without sampling metadata keep their established {input}. */
+export function responseGrammarProperties(body: { tools?: unknown; input?: unknown }, tools?: Tool[], includeLegacy = false): ReadonlyMap<string, string> {
+	const properties = new Map<string, string>();
+	const visit = (items: unknown, namespace?: string): void => {
+		if (!Array.isArray(items)) return;
+		for (const item of items) {
+			if (!item || typeof item !== "object") continue;
+			if (item.type === "namespace" && typeof item.name === "string") { visit(item.tools, item.name); continue; }
+			if (item.type !== "custom" || typeof item.name !== "string") continue;
+			const name = localToolName(namespace, item.name);
+			const tool = tools?.find((candidate) => candidate.name === name);
+			const grammar = tool && resolveGrammarSampling(tool, true);
+			if (grammar) properties.set(name, grammar.inputProperty);
+			else if (includeLegacy) properties.set(name, "input");
+		}
+	};
+	visit(body.tools);
+	if (Array.isArray(body.input)) for (const item of body.input) {
+		if (item?.type === "additional_tools") visit(item.tools);
+	}
+	return properties;
+}
+
+export function customArguments(name: string, input: string, options?: OpenAIResponsesStreamOptions): Record<string, string> {
+	return { [options?.grammarToolInputProperties?.get(name) ?? "input"]: input };
+}
+
+export function updateCustomInput(state: CustomToolCallState, input: string, close: boolean,
+	output: AssistantMessage, stream: AssistantMessageEventStream, options?: OpenAIResponsesStreamOptions): void {
+	const property = options?.grammarToolInputProperties?.get(state.block.name);
+	const previous = state.input;
+	state.input = input;
+	state.block.partialInput = input;
+	state.block.arguments = customArguments(state.block.name, input, options);
+	if (property !== undefined) {
+		state.inputJson ??= { input: "", started: false, closed: false };
+		const delta = grammarInputDelta(state.inputJson, property, input, close);
+		if (delta) stream.push({ type: "toolcall_delta", contentIndex: state.blockIndex, delta, partial: output });
+	} else if (!close && input.startsWith(previous) && input.length > previous.length) {
+		stream.push({ type: "toolcall_delta", contentIndex: state.blockIndex, delta: input.slice(previous.length), partial: output });
+	}
+}

package/src/providers/responses/messages.ts

@@ -6,6 +6,7 @@ import { decodeWebSearchActivityTextSignature, isWebSearchActivityTextSignature,
 import { sanitizeSurrogates, shortHash } from "./text.js";
 import { wireToolIdentity } from "./tool-identity.js";
 import { type ConvertResponsesMessagesOptions, type InternalAssistantContent, type Message } from "./types.js";
+import { grammarToolInput } from "./sampling.js";
 
 export function convertResponsesMessages<TApi extends Api>(
 	model: Model<TApi>,
@@ -14,6 +15,7 @@ export function convertResponsesMessages<TApi extends Api>(
 	options?: ConvertResponsesMessagesOptions,
 ): ResponseInput {
 	const messages: ResponseInput = [];
+	const customCalls = new Map<string, boolean>();
 	const normalizeIdPart = (part: string) => {
 		const sanitized = part.replace(/[^a-zA-Z0-9_-]/g, "_");
 		const normalized = sanitized.length > 64 ? sanitized.slice(0, 64) : sanitized;
@@ -92,9 +94,14 @@ export function convertResponsesMessages<TApi extends Api>(
 					assistantBlockIndex++;
 				} else if (block.type === "toolCall") {
 					const [callId, itemIdRaw] = block.id.split("|");
-					const custom = itemIdRaw?.startsWith("ctc_") === true;
+					const property = options?.grammarToolInputProperties?.get(block.name);
+					const declared = options?.grammarToolInputProperties !== undefined && context.tools?.some((tool) => tool.name === block.name);
+					const custom = property !== undefined || (!declared && itemIdRaw?.startsWith("ctc_") === true && typeof block.arguments.input === "string");
+					customCalls.set(callId, custom);
 					let itemId: string | undefined = itemIdRaw;
 					if (isDifferentModel && (itemId?.startsWith("fc_") || itemId?.startsWith("ctc_"))) itemId = undefined;
+					if (!custom && itemId?.startsWith("ctc_")) itemId = undefined;
+					if (custom && itemId?.startsWith("fc_")) itemId = undefined;
 					const wireIdentity = wireToolIdentity(block.name, block.thoughtSignature);
 					if (custom) {
 						output.push({
@@ -103,7 +110,7 @@ export function convertResponsesMessages<TApi extends Api>(
 							call_id: callId,
 							name: wireIdentity.name,
 							...(wireIdentity.namespace ? { namespace: wireIdentity.namespace } : {}),
-							input: typeof block.arguments.input === "string" ? block.arguments.input : "",
+							input: sanitizeSurrogates(grammarToolInput(block.name, block.arguments, property ?? "input")),
 						} as ResponseInput[number]);
 					} else {
 						output.push({
@@ -157,7 +164,7 @@ export function convertResponsesMessages<TApi extends Api>(
 					]
 				: sanitizeSurrogates(hasText ? textResult : "(see attached image)");
 			messages.push({
-				type: itemId?.startsWith("ctc_") ? "custom_tool_call_output" : "function_call_output",
+				type: (customCalls.get(callId) ?? (options?.grammarToolInputProperties?.has(msg.toolName) || itemId?.startsWith("ctc_"))) ? "custom_tool_call_output" : "function_call_output",
 				call_id: callId,
 				output,
 			} as ResponseInput[number]);

package/src/providers/responses/sampling.ts

@@ -0,0 +1,41 @@
+import type { Tool } from "@earendil-works/pi-ai";
+
+export interface GrammarInputBuffer { input: string; started: boolean; closed: boolean }
+export interface GrammarSampling { format: "lark" | "regex"; definition: string; inputProperty: string }
+
+// Pi's extension loader aliases SDK roots to files, so api/* subpath imports
+// do not resolve in production extensions. Implement the small wire contract
+// locally rather than resolving private SDK dist paths.
+export function resolveGrammarSampling(tool: Tool, supported: boolean): GrammarSampling | undefined {
+	const sampling = tool.constrainedSampling;
+	if (!supported || !sampling || sampling.type !== "grammar") return undefined;
+	const schema = tool.parameters as { type?: unknown; required?: unknown; properties?: Record<string, { type?: unknown }> };
+	const required = schema.required;
+	if (schema.type !== "object" || !Array.isArray(required) || required.length !== 1
+		|| typeof required[0] !== "string" || schema.properties?.[required[0]]?.type !== "string") {
+		throw new Error(`Tool ${tool.name} grammar requires exactly one required string property`);
+	}
+	for (const [variant, format] of [["openai_lark", "lark"], ["openai_regex", "regex"]] as const) {
+		const definition = sampling.variants[variant];
+		if (typeof definition === "string" && definition.trim()) return { format, definition, inputProperty: required[0] };
+	}
+	throw new Error(`Tool ${tool.name} has no supported grammar variant`);
+}
+export function grammarToolInput(name: string, args: Record<string, unknown>, property: string): string {
+	const value = args[property];
+	if (typeof value !== "string") throw new Error(`Grammar tool ${name} requires string argument ${property}`);
+	return value;
+}
+export function grammarInputDelta(buffer: GrammarInputBuffer, property: string, next: string, close: boolean): string | undefined {
+	if (buffer.closed) {
+		if (next !== buffer.input) throw new Error("Grammar input changed after it was closed");
+		return undefined;
+	}
+	if (!next.startsWith(buffer.input)) throw new Error("Grammar input is not monotonic");
+	const suffix = JSON.stringify(next.slice(buffer.input.length)).slice(1, -1);
+	const delta = (buffer.started ? "" : `{${JSON.stringify(property)}:"`) + suffix + (close ? '"}' : "");
+	buffer.input = next;
+	buffer.started = true;
+	buffer.closed = close;
+	return delta || undefined;
+}

package/src/providers/responses/stream.ts

@@ -5,7 +5,8 @@ import { createResponsesStreamState } from "./stream-state.js";
 import { createResponseTextRenderer } from "./text-renderer.js";
 import { localToolName, parseStreamingJson } from "./text.js";
 import { encodeToolNamespaceSignature } from "./tool-identity.js";
-import type { TextBlock, ThinkingBlock, ToolCallBlock } from "./types.js";
+import type { CustomToolCallState, TextBlock, ThinkingBlock, ToolCallBlock } from "./types.js";
+import { customArguments, updateCustomInput } from "./grammar.js";
 import { type OpenAIResponsesStreamOptions } from "./types.js";
 import { finalizeResponseUsage } from "./usage.js";
 
@@ -97,12 +98,12 @@ export async function processResponsesStream<TApi extends Api>(
 					type: "toolCall",
 					id: `${customItem.call_id}|${itemId}`,
 					name: localToolName(customItem.namespace, customItem.name),
-					arguments: { input },
+					arguments: customArguments(localToolName(customItem.namespace, customItem.name), input, options),
 					...(thoughtSignature ? { thoughtSignature } : {}),
 					partialInput: input,
 				};
 				pushResponseBlock(currentBlock, event.output_index);
-				outputStates.set(event.output_index, {
+				const state: CustomToolCallState = {
 					kind: "custom_tool_call",
 					blockIndex: blockIndex(),
 					block: currentBlock,
@@ -110,8 +111,10 @@ export async function processResponsesStream<TApi extends Api>(
 					sourceItemId: customItem.id,
 					callId: customItem.call_id,
 					input,
-				});
+				};
+				outputStates.set(event.output_index, state);
 				stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
+				if (options?.grammarToolInputProperties?.has(currentBlock.name)) updateCustomInput(state, input, false, output, stream, options);
 			}
 		} else if (event.type === "response.reasoning_summary_part.added") {
 			const state = outputStates.get(event.output_index);
@@ -193,11 +196,12 @@ export async function processResponsesStream<TApi extends Api>(
 			const customEvent = event as unknown as { output_index?: number; item_id?: string; call_id?: string; delta?: string };
 			const state = findCustomToolCallState(customEvent);
 			if (state && typeof customEvent.delta === "string") {
-				state.input += customEvent.delta;
-				state.block.partialInput = state.input;
-				state.block.arguments = { input: state.input };
-				stream.push({ type: "toolcall_delta", contentIndex: state.blockIndex, delta: customEvent.delta, partial: output });
+				updateCustomInput(state, state.input + customEvent.delta, false, output, stream, options);
 			}
+		} else if ((event as { type?: string }).type === "response.custom_tool_call_input.done") {
+			const customEvent = event as unknown as { output_index?: number; item_id?: string; call_id?: string; input?: string };
+			const state = findCustomToolCallState(customEvent);
+			if (state && typeof customEvent.input === "string") updateCustomInput(state, customEvent.input, true, output, stream, options);
 		} else if (event.type === "response.function_call_arguments.delta") {
 			const state = outputStates.get(event.output_index);
 			if (state?.kind === "function_call") {
@@ -309,9 +313,8 @@ export async function processResponsesStream<TApi extends Api>(
 				const thoughtSignature = encodeToolNamespaceSignature(customItem.namespace, customItem.name);
 				const toolCall = state
 					? (() => {
-						state.input = input;
 						state.block.name = localToolName(customItem.namespace, customItem.name);
-						state.block.arguments = { input };
+						updateCustomInput(state, input, true, output, stream, options);
 						if (thoughtSignature) state.block.thoughtSignature = thoughtSignature;
 						delete state.block.partialInput;
 						return state.block;
@@ -321,7 +324,7 @@ export async function processResponsesStream<TApi extends Api>(
 							type: "toolCall",
 							id: `${customItem.call_id}|${customItemId(customItem.id, customItem.call_id)}`,
 							name: localToolName(customItem.namespace, customItem.name),
-							arguments: { input },
+							arguments: customArguments(localToolName(customItem.namespace, customItem.name), input, options),
 							...(thoughtSignature ? { thoughtSignature } : {}),
 						};
 						pushResponseBlock(fallbackToolCall, event.output_index);

package/src/providers/responses/text.ts

@@ -11,10 +11,11 @@ export function shortHash(str: string): string {
 	return (h2 >>> 0).toString(36) + (h1 >>> 0).toString(36);
 }
 
-export function parseStreamingJson(partialJson: string): Record<string, unknown> {
+export function parseStreamingJson(partialJson: string): JsonObject {
 	if (!partialJson || partialJson.trim() === "") return {};
 	try {
-		return JSON.parse(partialJson) as Record<string, unknown>;
+		const value: unknown = JSON.parse(partialJson);
+		return value !== null && typeof value === "object" && !Array.isArray(value) ? value as JsonObject : {};
 	} catch {
 		return {};
 	}
@@ -30,3 +31,4 @@ export function localToolName(namespace: unknown, name: string): string {
 	if (namespace === "image_gen" && name === "imagegen") return "image_generation";
 	return name;
 }
+import type { JsonObject } from "@earendil-works/pi-ai";

package/src/providers/responses/tools.ts

@@ -1,14 +1,19 @@
 import { type Tool } from "@earendil-works/pi-ai";
 import { type Tool as OpenAITool } from "openai/resources/responses/responses.js";
 import { type ConvertResponsesToolsOptions } from "./types.js";
+import { resolveGrammarSampling } from "./sampling.js";
 
 export function convertResponsesTools(tools: Tool[], options?: ConvertResponsesToolsOptions): OpenAITool[] {
 	const strict = options?.strict === undefined ? false : options.strict;
-	return tools.map((tool) => ({
-		type: "function",
-		name: tool.name,
-		description: tool.description,
-		parameters: tool.parameters as unknown as Record<string, unknown>,
-		strict,
-	}));
+	return tools.map((tool) => {
+		const grammar = resolveGrammarSampling(tool, options?.supportsOpenAIGrammarTools === true);
+		if (grammar) return {
+			type: "custom", name: tool.name, description: tool.description,
+			format: { type: "grammar", syntax: grammar.format, definition: grammar.definition },
+		};
+		return {
+			type: "function", name: tool.name, description: tool.description,
+			parameters: tool.parameters as unknown as Record<string, unknown>, strict,
+		};
+	});
 }

package/src/providers/responses/types.ts

@@ -1,6 +1,7 @@
 import type { AssistantMessage } from "@earendil-works/pi-ai";
 import { type Context, type Usage } from "@earendil-works/pi-ai";
 import { type ResponseCreateParamsStreaming } from "openai/resources/responses/responses.js";
+import type { GrammarInputBuffer } from "./sampling.js";
 
 type MessageRole = Context["messages"][number]["role"];
 
@@ -52,6 +53,7 @@ export interface ReplayableResponseMessageItem {
 export type InternalAssistantContent = Extract<Message, { role: "assistant" }>["content"][number] | ImageGenerationCallBlock;
 
 export interface OpenAIResponsesStreamOptions {
+	grammarToolInputProperties?: ReadonlyMap<string, string>;
 	serviceTier?: ResponseCreateParamsStreaming["service_tier"];
 	resolveServiceTier?: (
 		responseServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
@@ -66,10 +68,12 @@ export type TextSignaturePhase = "commentary" | "final_answer";
 
 export interface ConvertResponsesMessagesOptions {
 	includeSystemPrompt?: boolean;
+	grammarToolInputProperties?: ReadonlyMap<string, string>;
 }
 
 export interface ConvertResponsesToolsOptions {
 	strict?: boolean | null;
+	supportsOpenAIGrammarTools?: boolean;
 }
 
 export type ThinkingBlock = Extract<AssistantMessage["content"][number], { type: "thinking" }>;
@@ -108,6 +112,7 @@ export type CustomToolCallState = {
 	sourceItemId?: string;
 	callId: string;
 	input: string;
+	inputJson?: GrammarInputBuffer;
 };
 
 export type OutputState = ReasoningState | MessageState | FunctionCallState | CustomToolCallState;

package/src/tool-activation.ts

@@ -8,6 +8,7 @@ import {
 import { installCodexIdentityLifecycle } from "./codex-identity-extension.js";
 import { loadModelSettings } from "./model-catalog/runtime.js";
 import { loadSettings } from "./settings.js";
+import { codeModeOwner, disposeCodeModeOwner, OWNER_CHANGED } from "./code-mode-owner.js";
 
 function enableDefinitions(broker: CodexBroker) {
 	for (const tool of broker.tools.values()) {
@@ -37,10 +38,15 @@ export function ensureCodexServices(pi: ExtensionAPI): CodexBroker {
 		for (const [name, index] of [...suppressed].sort(([, a], [, b]) => a - b)) {
 			if (!active.includes(name)) active.splice(Math.min(index, active.length), 0, name);
 		}
-		suppressed.clear();
 		return active;
 	};
-	const sync = (ctx: ExtensionContext) => {
+	let latest: ExtensionContext | undefined;
+	let syncing = false;
+	const sync = (ctx: ExtensionContext, fromTree = false) => {
+		latest = ctx;
+		if (syncing) return;
+		syncing = true;
+		try {
 		const settings = loadSettings(ctx.cwd);
 		const available = settings.enabled && hasConfiguredModelsLoaded(ctx, settings);
 		if (available) enableDefinitions(broker);
@@ -48,42 +54,78 @@ export function ensureCodexServices(pi: ExtensionAPI): CodexBroker {
 		const active = new Set(current);
 		const capabilities = computeToolCapabilities(ctx.model as ModelLike | undefined, settings);
 		const model = loadModelSettings(ctx.model as ModelLike | undefined, ctx.cwd, settings);
+		const controls = new Map<string, NonNullable<ReturnType<typeof codeModeOwner>>>();
 		for (const name of PACKAGE_TOOL_NAMES) {
 			const owned = broker.tools.get(name);
 			if (!owned) continue; // Other extensions retain ownership of uninstalled names.
+			const control = codeModeOwner(pi, name, owned, broker.codeModeDefinitions?.get(name));
+			if (owned.codeModeOwner?.replaced) continue;
+			if (control) {
+				controls.set(name, control);
+				if (control.activeIntent === undefined) continue; // foreign replacement
+				if (control.activeIntent) active.add(name); else active.delete(name);
+			}
 			const hostedWithoutCore = !broker.coreEnabled && (
 				(name === "web_search" && model.webSearchImplementation === "hosted")
 				|| (name === "image_generation" && model.imageGenerationImplementation === "hosted" && !settings.directImageApiFallback)
 			);
 			const desired = available && owned.registered && capabilities[name].enabled && !hostedWithoutCore;
 			if (!desired) active.delete(name);
-			else if (settings.autoEnable) active.add(name);
+			else if (settings.autoEnable && !(fromTree && control?.activeIntent !== undefined)) active.add(name);
 		}
-		if (broker.tools.has("apply_patch") && active.has("apply_patch")) {
+		const ownsPatch = broker.tools.has("apply_patch") && !broker.tools.get("apply_patch")?.codeModeOwner?.replaced;
+		if (ownsPatch && active.has("apply_patch")) {
 			for (const name of NATIVE_MUTATION_TOOL_NAMES) {
 				if (active.delete(name) && !suppressed.has(name)) suppressed.set(name, current.indexOf(name));
 			}
 		}
-		const next = current.filter(name => active.has(name));
-		for (const name of active) if (!next.includes(name)) next.push(name);
-		if (!broker.tools.has("apply_patch") || !active.has("apply_patch")) restore(next);
+		const physical = new Set(active);
+		for (const [name, control] of controls) {
+			const projected = control.projectActive(active.has(name));
+			if (projected === undefined) { if (current.includes(name)) physical.add(name); else physical.delete(name); }
+			else if (projected) physical.add(name); else physical.delete(name);
+		}
+		const next = current.filter(name => physical.has(name));
+		for (const name of physical) if (!next.includes(name)) next.push(name);
+		if (!ownsPatch || !active.has("apply_patch")) restore(next);
 		if (next.join("\0") !== current.join("\0")) pi.setActiveTools(next);
+		if (!ownsPatch || !active.has("apply_patch")) suppressed.clear();
+		for (const control of controls.values()) control.reconcile();
+		} finally { syncing = false; }
 	};
+	const offOwner = pi.events.on(OWNER_CHANGED, (message) => {
+		if ((message as { version?: number })?.version === 1 && latest && !broker.closed) sync(latest);
+	});
 	pi.on("session_start", (_event, ctx) => {
 		broker.presentation.clear();
+		suppressed.clear();
 		sync(ctx);
 	});
 	pi.on("model_select", (_event, ctx) => sync(ctx));
 	pi.on("thinking_level_select", (_event, ctx) => sync(ctx));
+	pi.on("session_tree", (_event, ctx) => {
+		// Restoration receipts belong to the previous physical loadout, not
+		// to tools absent from the newly selected branch.
+		suppressed.clear();
+		sync(ctx, true);
+	});
 	pi.on("agent_end", () => broker.presentation.scheduleFlush());
 	pi.on("session_shutdown", () => {
-		try { broker.presentation.flush(); }
-		finally {
-			broker.presentation.clear();
+		offOwner();
+		const errors: unknown[] = [];
+		for (const tool of broker.tools.values()) if (tool.codeModeOwner) tool.codeModeOwner.closed = true;
+		for (const tool of broker.tools.values()) {
+			try { disposeCodeModeOwner(tool); } catch (error) { errors.push(error); }
+		}
+		try { broker.presentation.flush(); } catch (error) { errors.push(error); }
+		try { broker.presentation.clear(); } catch (error) { errors.push(error); }
+		try {
 			const current = pi.getActiveTools?.() ?? [];
 			const next = restore([...current]);
 			if (next.join("\0") !== current.join("\0")) pi.setActiveTools(next);
-		}
+			suppressed.clear();
+		} catch (error) { errors.push(error); }
+		if (errors.length) throw new AggregateError(errors, "Codex owner shutdown failed");
 	});
 	return broker;
 }