@oai404iao/pi-codex-runtime 0.3.0 → 0.3.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@oai404iao/pi-codex-runtime",
-  "version": "0.3.0",
+  "version": "0.3.1",
   "private": false,
   "description": "Unpublished pi-codex-runtime workspace for the staged Codex split",
   "type": "module",
@@ -64,5 +64,5 @@
   "engines": {
     "node": ">=22.19.0"
   },
-  "gitHead": "b02484cecad9081886797bf9b4438135aaae1e07"
+  "gitHead": "c791e1192d0547834c3e2f2bae7853fb214f9808"
 }

package/src/broker.ts

@@ -20,7 +20,7 @@ export interface CodexBroker {
 	readonly closed: boolean;
 	readonly tools: Map<PackageToolName, InstalledTool>;
 	coreEnabled: boolean;
-	codeModeDefinitions?: Map<string, { parameters: unknown; description: string }>;
+	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

@@ -2,17 +2,21 @@ import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-a
 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>): void {
+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") throw new Error("Invalid owned Code Mode tool definition");
+		|| !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 } };
-	(broker.codeModeDefinitions ??= new Map()).set(owned.name, owned);
 	pi.registerTool(owned as never);
+	definitions.set(owned.name, { ...owned, providerId });
 }
 
 export interface NestedCodeModeContext {
@@ -28,36 +32,108 @@ export interface OwnedCodeModeTool {
 	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 v1 bus contract. Neither side installs
+/** 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;
-	const changed = () => pi.events.emit("@oai404iao/pi-code-mode:changed/v1", { version: 1 });
+	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 } | undefined;
+		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") {
-			const broker = getCodexBroker(pi);
-			discovery.provider({ id, tools: tools(context).map((tool) => {
-				const owned = broker.tools.get(tool.name as PackageToolName);
-				return { ...tool, direct: owned ? codeModeOwner(pi, tool.name, owned, broker.codeModeDefinitions?.get(tool.name))?.binding : undefined };
-			}) });
+			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;
-		changed();
+		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(); changed();
+		disposed = true; context = undefined; off(); offV2(); changed("disposed");
 	});
 }

package/src/code-mode-owner.ts

@@ -12,34 +12,85 @@ interface Control {
 interface Factory {
 	version: 1; create(pi: ExtensionAPI, options: { name: string; sourcePath: string }): Control;
 }
-export interface OwnerState { identity: string; replaced?: boolean; factory?: Factory; control?: 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.registered) return;
+	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)) return tool.codeModeOwner?.control;
-	if (!expected || info.parameters !== expected.parameters || info.description !== expected.description) {
-		if (expected) (tool.codeModeOwner ??= { identity: "" }).replaced = true;
+	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; return state.control; }
-	if (state.replaced) return state.control;
+	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 (factory !== state.factory) {
-		const old = state.control;
-		state.control = undefined; state.factory = factory;
-		old?.dispose();
-		if (factory) state.control = factory.create(pi, { name, sourcePath: info.sourceInfo.path });
+	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/tool-activation.ts

@@ -8,7 +8,7 @@ import {
 import { installCodexIdentityLifecycle } from "./codex-identity-extension.js";
 import { loadModelSettings } from "./model-catalog/runtime.js";
 import { loadSettings } from "./settings.js";
-import { codeModeOwner, OWNER_CHANGED } from "./code-mode-owner.js";
+import { codeModeOwner, disposeCodeModeOwner, OWNER_CHANGED } from "./code-mode-owner.js";
 
 function enableDefinitions(broker: CodexBroker) {
 	for (const tool of broker.tools.values()) {
@@ -38,12 +38,11 @@ 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;
 	};
 	let latest: ExtensionContext | undefined;
 	let syncing = false;
-	const sync = (ctx: ExtensionContext) => {
+	const sync = (ctx: ExtensionContext, fromTree = false) => {
 		latest = ctx;
 		if (syncing) return;
 		syncing = true;
@@ -72,7 +71,7 @@ export function ensureCodexServices(pi: ExtensionAPI): CodexBroker {
 			);
 			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);
 		}
 		const ownsPatch = broker.tools.has("apply_patch") && !broker.tools.get("apply_patch")?.codeModeOwner?.replaced;
 		if (ownsPatch && active.has("apply_patch")) {
@@ -90,6 +89,7 @@ export function ensureCodexServices(pi: ExtensionAPI): CodexBroker {
 		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; }
 	};
@@ -98,21 +98,34 @@ export function ensureCodexServices(pi: ExtensionAPI): CodexBroker {
 	});
 	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", () => {
 		offOwner();
-		for (const tool of broker.tools.values()) tool.codeModeOwner?.control?.dispose();
-		try { broker.presentation.flush(); }
-		finally {
-			broker.presentation.clear();
+		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;
 }