@automatalabs/pi-acp 0.1.3 → 0.2.1

package/dist/agent.js

@@ -1,11 +1,100 @@
-import { existsSync, readFileSync, statSync } from "node:fs";
+import { existsSync, statSync } from "node:fs";
+import { randomUUID } from "node:crypto";
 import { isAbsolute } from "node:path";
+import { DefaultResourceLoader, SettingsManager, createBashToolDefinition, getAgentDir, } from "@earendil-works/pi-coding-agent";
 import { AUTH_METHODS, authenticateMethod } from "./auth.js";
-import { adapterError, isRequestError, unexpectedError } from "./errors.js";
-import { bridgeMcpServers, disposeMcpBridge, } from "./mcp-bridge.js";
+import { adapterError, isChildCleanupError, isRequestError, unexpectedError } from "./errors.js";
+import { bridgeMcpServers, } from "./mcp-bridge.js";
 import { PiSession } from "./session.js";
-import { StructuredOutputState, STRUCTURED_TOOL_NAME } from "./structured-output.js";
-export const PKG_VERSION = String(JSON.parse(readFileSync(new URL("../package.json", import.meta.url), "utf8")).version);
+import { ChildProcessRegistrySlot, createTrackedBashOperations } from "./child-process-registry.js";
+import { PKG_VERSION } from "./version.js";
+export { PKG_VERSION } from "./version.js";
+/** Retry owner for the narrow interval after Pi exists but before PiSession is
+ * publishable.  It gives failed-open rollback the same abort/tree barrier and
+ * hidden-record ownership as a fully constructed session. */
+class FailedOpenCleanup {
+    pi;
+    bridge;
+    children;
+    lifecycle;
+    deps;
+    resourceDispose;
+    constructor(pi, bridge, children, lifecycle, deps) {
+        this.pi = pi;
+        this.bridge = bridge;
+        this.children = children;
+        this.lifecycle = lifecycle;
+        this.deps = deps;
+    }
+    get remainingChildren() { return this.children.remainingChildren; }
+    startResources(bridgeClose) {
+        this.resourceDispose ??= (async () => {
+            await this.bridge.drainRefreshes().catch((error) => {
+                console.error("pi-acp failed-open refresh drain error:", error);
+            });
+            const results = await Promise.allSettled([
+                Promise.resolve().then(() => this.pi.dispose()),
+                bridgeClose,
+            ]);
+            for (const result of results) {
+                if (result.status === "rejected")
+                    console.error("pi-acp failed-open resource disposal error:", result.reason);
+            }
+        })();
+        return this.resourceDispose;
+    }
+    async dispose() {
+        const deadline = new AbortController();
+        const timer = new AbortController();
+        const expiry = this.deps.sleep(this.deps.graceMs, timer.signal).then(() => {
+            const failure = adapterError("child_cleanup_error", {
+                details: { remainingChildren: this.children.remainingChildren },
+            });
+            deadline.abort(failure);
+            throw failure;
+        });
+        expiry.catch(() => undefined);
+        // Failed-open cleanup has the same synchronous prefix as a published
+        // session: close spawn admission first, logically close every MCP client
+        // before aborting the binding lifetime, then start Pi abort before its
+        // one-shot resource disposal can run.
+        const captured = this.children.closeEpoch(deadline.signal);
+        this.bridge.startDisposal();
+        if (!this.lifecycle.signal.aborted)
+            this.lifecycle.abort(new Error("failed open disposed"));
+        this.bridge.abortRefreshes();
+        const bridgeClose = this.bridge.close();
+        bridgeClose.catch(() => undefined);
+        let abort;
+        try {
+            abort = this.pi.abort();
+        }
+        catch (error) {
+            abort = Promise.reject(error);
+        }
+        abort.catch(() => undefined);
+        const settled = Promise.allSettled([abort, captured.drain]);
+        let cleanupError;
+        await Promise.race([settled, expiry]).then((results) => {
+            if (!Array.isArray(results))
+                return;
+            const failure = results.find((result) => result.status === "rejected");
+            if (failure) {
+                cleanupError = adapterError("child_cleanup_error", {
+                    details: { remainingChildren: this.children.remainingChildren },
+                });
+            }
+        }, (error) => { cleanupError = error; });
+        timer.abort();
+        // Pi disposal is defense in depth only after the abort/tree barrier has
+        // committed success or failure. MCP physical close was already started
+        // by the synchronous disposal prefix above.
+        const resources = this.startResources(bridgeClose);
+        await resources;
+        if (cleanupError)
+            throw cleanupError;
+    }
+}
 function validateCwd(cwd) {
     if (!isAbsolute(cwd))
         throw adapterError("invalid_cwd");
@@ -35,7 +124,11 @@ export class PiAcpAgent {
     openingControllers = new Set();
     openingTasks = new Set();
     tombstones = new Set();
+    cleanupRecords = new Map();
+    mcpOwnerToken = {};
     disposed = false;
+    disposePromise;
+    disposeSucceeded = false;
     constructor(deps) {
         this.deps = deps;
     }
@@ -50,9 +143,8 @@ export class PiAcpAgent {
             agentCapabilities: {
                 loadSession: true,
                 promptCapabilities: { image: true },
-                mcpCapabilities: {},
+                mcpCapabilities: { http: true, sse: true },
                 sessionCapabilities: { resume: {}, fork: {}, list: {}, close: {} },
-                _meta: { "@automatalabs/pi-acp": { outputSchema: true } },
             },
             authMethods: AUTH_METHODS,
         };
@@ -68,16 +160,29 @@ export class PiAcpAgent {
         if (this.live.has(id) || this.opening.has(id))
             throw adapterError("session_already_open");
     }
-    reserve(id, controller) {
+    reserve(id, opening) {
         this.ensureMayOpen(id);
-        this.opening.set(id, controller);
+        this.opening.set(id, opening);
     }
     beginOpening(requestSignal) {
         if (this.disposed)
             throw adapterError("internal_error");
         const controller = openSignal(requestSignal);
+        let settled = false;
+        let resolveSettlement;
+        const settlement = new Promise((resolve) => { resolveSettlement = resolve; });
+        const opening = {
+            controller,
+            settlement,
+            settle: () => {
+                if (settled)
+                    return;
+                settled = true;
+                resolveSettlement();
+            },
+        };
         this.openingControllers.add(controller);
-        return { controller };
+        return opening;
     }
     track(task) {
         this.openingTasks.add(task);
@@ -90,78 +195,213 @@ export class PiAcpAgent {
             throw opening.controller.signal.reason ?? adapterError("internal_error");
         }
     }
+    async connectMcp(opening, sessionId, cwd, client, mcpServers) {
+        const lifecycle = new AbortController();
+        const state = { published: false };
+        const binding = {
+            sessionId,
+            cwd,
+            client,
+            sessionSignal: lifecycle.signal,
+            getPi: () => state.pi,
+            getTurnSignal: () => state.wrapper?.activeTurnSignal(),
+            isPublished: () => state.published,
+            emitDiagnostic: (text) => {
+                if (state.wrapper)
+                    state.wrapper.emitMcpDiagnostic(text);
+                else
+                    console.error(text);
+            },
+            poison: () => state.wrapper?.poison(),
+            ownerToken: this.mcpOwnerToken,
+            modelRuntime: this.deps.modelRuntime,
+        };
+        try {
+            const bridge = await bridgeMcpServers(mcpServers, opening.controller.signal, this.deps, binding);
+            return { bridge, lifecycle, state };
+        }
+        catch (error) {
+            lifecycle.abort(new Error("MCP opening failed"));
+            throw error;
+        }
+    }
     async construct(opening, manager, cwd, client, mcpServers, replay, preconnected) {
         const id = manager.getSessionId();
         if (opening.id === undefined) {
-            this.reserve(id, opening.controller);
+            this.reserve(id, opening);
             opening.id = id;
         }
-        let bridge = preconnected;
+        let prepared = preconnected;
+        let bridge = prepared?.bridge;
         let pi;
         let wrapper;
+        let lifecycle = prepared?.lifecycle;
+        let bindingState = prepared?.state;
+        const childRegistry = new ChildProcessRegistrySlot(this.deps);
         try {
             this.gate(opening);
-            bridge ??= await bridgeMcpServers(mcpServers, opening.controller.signal, this.deps);
+            prepared ??= await this.connectMcp(opening, id, cwd, client, mcpServers);
+            bridge = prepared.bridge;
+            lifecycle = prepared.lifecycle;
+            bindingState = prepared.state;
             this.gate(opening);
-            const structured = new StructuredOutputState();
+            const settingsManager = SettingsManager.create(cwd, getAgentDir());
+            const instructionFactory = bridge.instructionsExtension;
+            const controlExtension = {
+                name: "agentprism-pi-acp-control",
+                factory: async (api) => {
+                    if (typeof instructionFactory === "function")
+                        await instructionFactory(api);
+                    else
+                        await instructionFactory.factory(api);
+                    api.registerTool(createBashToolDefinition(cwd, {
+                        commandPrefix: settingsManager.getShellCommandPrefix(),
+                        operations: createTrackedBashOperations(childRegistry, settingsManager.getShellPath(), this.deps, () => wrapper?.childCleanupFailure()),
+                    }));
+                },
+            };
+            const resourceLoader = new DefaultResourceLoader({
+                cwd,
+                agentDir: getAgentDir(),
+                settingsManager,
+                extensionFactories: [bridge.inlineExtension, controlExtension],
+                extensionsOverride: (base) => {
+                    const matches = base.extensions.filter(({ path }) => path === "<inline:agentprism-pi-acp-mcp>");
+                    const controls = base.extensions.filter(({ path }) => path === "<inline:agentprism-pi-acp-control>");
+                    if (matches.length !== 1 || controls.length !== 1)
+                        throw adapterError("extension_setup_error");
+                    const reserved = matches[0];
+                    const control = controls[0];
+                    const configured = base.extensions.filter(({ path }) => !path.startsWith("<inline:"));
+                    const configuredBash = configured.some(({ tools }) => tools.has("bash"));
+                    if (configuredBash)
+                        control.tools.delete("bash");
+                    // Pi records conflict diagnostics before this override can impose
+                    // the adapter's deliberate precedence. Remove only the conflicts
+                    // that this transaction has actually resolved: the reserved MCP
+                    // extension is moved first, and core bash is omitted when a
+                    // configured extension already owns it. User/user conflicts and
+                    // every loader/factory error remain fatal and retain their order.
+                    const reservedNames = new Set(reserved.tools.keys());
+                    for (let index = base.errors.length - 1; index >= 0; index -= 1) {
+                        const issue = base.errors[index];
+                        const resolvedReserved = issue.path === reserved.path
+                            && [...reservedNames].some((name) => issue.error.startsWith(`Tool "${name}" conflicts with `));
+                        const resolvedBash = configuredBash
+                            && issue.path === control.path
+                            && issue.error.startsWith('Tool "bash" conflicts with ');
+                        if (resolvedReserved || resolvedBash)
+                            base.errors.splice(index, 1);
+                    }
+                    const result = { ...base, extensions: [reserved, ...base.extensions.filter((item) => item !== reserved)] };
+                    if (result.runtime !== base.runtime || result.errors !== base.errors)
+                        throw adapterError("extension_setup_error");
+                    return result;
+                },
+            });
+            await resourceLoader.reload();
+            if (resourceLoader.getExtensions().errors.length > 0)
+                throw adapterError("extension_setup_error");
             const created = await this.deps.createAgentSession({
                 cwd,
                 sessionManager: manager,
                 modelRuntime: this.deps.modelRuntime,
-                customTools: [...bridge.tools, structured.tool],
+                resourceLoader,
+                settingsManager,
             });
             pi = created.session;
-            this.gate(opening);
+            bindingState.pi = pi;
             wrapper = new PiSession({
                 sessionId: id,
                 session: pi,
                 manager,
                 client,
                 deps: this.deps,
-                mcpClients: bridge.clients,
+                mcpBridge: bridge,
                 failedMcpResults: bridge.failedResults,
-                structured,
-                onWedged: (sessionId, session) => this.terminateWedged(sessionId, session),
+                availableModels: [],
+                childRegistry,
+                lifecycleController: lifecycle,
+                onWedged: (sessionId, session, cleanupRetryRequired) => this.terminateWedged(sessionId, session, cleanupRetryRequired),
             });
-            structured.install(pi);
-            const names = new Set(pi.getAllTools().map(({ name }) => name));
-            if (!names.has(STRUCTURED_TOOL_NAME))
-                throw adapterError("structured_tool_collision");
+            bindingState.wrapper = wrapper;
+            await pi.bindExtensions({});
+            const availableModels = [...await this.deps.modelRuntime.getAvailable()];
+            wrapper.publishAvailableModels(availableModels);
+            this.gate(opening);
+            bridge.bindSession(pi);
+            const toolInfos = pi.getAllTools();
+            const names = new Set(toolInfos.map(({ name }) => name));
             const missingAlias = bridge.aliases.find((alias) => !names.has(alias));
             if (missingAlias) {
                 throw adapterError("mcp_init_error", {
                     server: bridge.aliasServers.get(missingAlias) ?? "unknown",
                 });
             }
+            for (const alias of bridge.aliases) {
+                const info = toolInfos.find(({ name }) => name === alias);
+                if (info?.sourceInfo.path !== "<inline:agentprism-pi-acp-mcp>") {
+                    throw adapterError("mcp_init_error", { server: bridge.aliasServers.get(alias) ?? "unknown" });
+                }
+            }
+            const bash = toolInfos.find(({ name }) => name === "bash");
+            if (!bash || bash.sourceInfo.path === "<builtin:bash>")
+                throw adapterError("extension_setup_error");
             if (replay)
                 await wrapper.replay(manager.getBranch());
+            bridge.assertReady();
             this.gate(opening);
             this.live.set(id, wrapper);
+            bindingState.published = true;
             this.opening.delete(id);
             return wrapper;
         }
         catch (error) {
-            if (wrapper)
-                await wrapper.disposeResources();
+            let cleanupError;
+            if (wrapper) {
+                try {
+                    await wrapper.dispose();
+                }
+                catch (candidate) {
+                    if (isChildCleanupError(candidate)) {
+                        this.cleanupRecords.set(id, wrapper);
+                        cleanupError = candidate;
+                    }
+                }
+            }
             else {
                 if (pi) {
+                    const rollback = new FailedOpenCleanup(pi, bridge, childRegistry, lifecycle, this.deps);
                     try {
-                        await pi.dispose();
+                        await rollback.dispose();
                     }
-                    catch (disposeError) {
-                        console.error("pi-acp rollback dispose error:", disposeError);
+                    catch (candidate) {
+                        if (isChildCleanupError(candidate)) {
+                            this.cleanupRecords.set(id, rollback);
+                            cleanupError = candidate;
+                        }
                     }
                 }
-                if (bridge)
-                    await disposeMcpBridge(bridge.clients, this.deps);
+                else if (bridge) {
+                    bridge.startDisposal();
+                    if (lifecycle && !lifecycle.signal.aborted)
+                        lifecycle.abort(new Error("failed open disposed"));
+                    bridge.abortRefreshes();
+                    await bridge.close();
+                }
+            }
+            if (cleanupError) {
+                opening.cleanupError = cleanupError;
+                throw cleanupError;
             }
             throw error;
         }
         finally {
-            if (opening.id !== undefined && this.opening.get(opening.id) === opening.controller) {
+            if (opening.id !== undefined && this.opening.get(opening.id) === opening) {
                 this.opening.delete(opening.id);
             }
             this.openingControllers.delete(opening.controller);
+            opening.settle();
         }
     }
     openingError(error) {
@@ -180,11 +420,12 @@ export class PiAcpAgent {
         }
         const opening = this.beginOpening(context.signal);
         try {
-            this.reserve(manager.getSessionId(), opening.controller);
+            this.reserve(manager.getSessionId(), opening);
             opening.id = manager.getSessionId();
         }
         catch (error) {
             this.openingControllers.delete(opening.controller);
+            opening.settle();
             throw error;
         }
         const task = this.construct(opening, manager, context.params.cwd, context.client, context.params.mcpServers, false).then((session) => ({ sessionId: session.sessionId, configOptions: session.configOptions(), modes: null }))
@@ -195,11 +436,12 @@ export class PiAcpAgent {
         validateCwd(context.params.cwd);
         const opening = this.beginOpening(context.signal);
         try {
-            this.reserve(context.params.sessionId, opening.controller);
+            this.reserve(context.params.sessionId, opening);
             opening.id = context.params.sessionId;
         }
         catch (error) {
             this.openingControllers.delete(opening.controller);
+            opening.settle();
             throw error;
         }
         const task = (async () => {
@@ -221,10 +463,11 @@ export class PiAcpAgent {
                 return { configOptions: session.configOptions(), modes: null };
             }
             catch (error) {
-                if (opening.id !== undefined && this.opening.get(opening.id) === opening.controller) {
+                if (opening.id !== undefined && this.opening.get(opening.id) === opening) {
                     this.opening.delete(opening.id);
                 }
                 this.openingControllers.delete(opening.controller);
+                opening.settle();
                 return this.openingError(error);
             }
         })();
@@ -245,7 +488,8 @@ export class PiAcpAgent {
         validateCwd(context.params.cwd);
         const opening = this.beginOpening(context.signal);
         const task = (async () => {
-            let bridge;
+            let prepared;
+            let preparedTransferred = false;
             try {
                 let sourcePath;
                 if (liveSource) {
@@ -260,32 +504,43 @@ export class PiAcpAgent {
                         throw adapterError("unknown_session");
                 }
                 this.gate(opening);
-                bridge = await bridgeMcpServers(context.params.mcpServers ?? [], opening.controller.signal, this.deps);
-                this.gate(opening);
                 if (this.tombstones.has(context.params.sessionId))
                     throw adapterError("session_terminated");
                 if (liveSource?.busy)
                     throw adapterError("session_busy");
+                // Pin the target id before the irreversible journal write so the MCP
+                // binding can be fully connected and owned first. SessionManager
+                // validates and uses this exact id when forkFrom eventually writes.
+                const targetId = randomUUID();
+                this.reserve(targetId, opening);
+                opening.id = targetId;
+                prepared = await this.connectMcp(opening, targetId, context.params.cwd, context.client, context.params.mcpServers ?? []);
+                this.gate(opening);
                 let manager;
                 try {
-                    manager = this.deps.sessions.forkFrom(sourcePath, context.params.cwd, this.deps.sessionDir);
+                    manager = this.deps.sessions.forkFrom(sourcePath, context.params.cwd, this.deps.sessionDir, { id: targetId });
                 }
                 catch {
                     throw adapterError("session_corrupt");
                 }
-                this.reserve(manager.getSessionId(), opening.controller);
-                opening.id = manager.getSessionId();
-                const session = await this.construct(opening, manager, context.params.cwd, context.client, context.params.mcpServers ?? [], false, bridge);
-                bridge = undefined;
+                preparedTransferred = true;
+                const session = await this.construct(opening, manager, context.params.cwd, context.client, context.params.mcpServers ?? [], false, prepared);
                 return { sessionId: session.sessionId, configOptions: session.configOptions(), modes: null };
             }
             catch (error) {
-                if (bridge)
-                    await disposeMcpBridge(bridge.clients, this.deps);
-                if (opening.id !== undefined && this.opening.get(opening.id) === opening.controller) {
+                if (prepared && !preparedTransferred) {
+                    prepared.bridge.startDisposal();
+                    prepared.lifecycle.abort(new Error("fork failed before construction"));
+                    prepared.bridge.abortRefreshes();
+                    await prepared.bridge.close().catch((closeError) => {
+                        console.error("pi-acp fork MCP rollback error:", closeError);
+                    });
+                }
+                if (opening.id !== undefined && this.opening.get(opening.id) === opening) {
                     this.opening.delete(opening.id);
                 }
                 this.openingControllers.delete(opening.controller);
+                opening.settle();
                 return this.openingError(error);
             }
         })();
@@ -333,18 +588,27 @@ export class PiAcpAgent {
         };
     }
     async closeSession(context) {
-        const controller = this.opening.get(context.params.sessionId);
-        if (controller) {
-            controller.abort(new Error("session closed while opening"));
+        const opening = this.opening.get(context.params.sessionId);
+        if (opening) {
+            opening.controller.abort(new Error("session closed while opening"));
+            await opening.settlement;
+            if (opening.cleanupError)
+                throw opening.cleanupError;
             return {};
         }
-        const session = this.live.get(context.params.sessionId);
+        const session = this.live.get(context.params.sessionId) ?? this.cleanupRecords.get(context.params.sessionId);
         if (!session)
             return {};
         try {
             await session.dispose();
+            this.cleanupRecords.delete(context.params.sessionId);
         }
         catch (error) {
+            if (isChildCleanupError(error)) {
+                this.tombstones.add(context.params.sessionId);
+                this.cleanupRecords.set(context.params.sessionId, session);
+                throw error;
+            }
             console.error("pi-acp close error:", error);
         }
         finally {
@@ -377,22 +641,72 @@ export class PiAcpAgent {
     cancel(context) {
         this.live.get(context.params.sessionId)?.cancel();
     }
-    async terminateWedged(id, session) {
+    async terminateWedged(id, session, cleanupRetryRequired) {
         this.tombstones.add(id);
         if (this.live.get(id) === session)
             this.live.delete(id);
-        await session.disposeResources();
+        if (cleanupRetryRequired)
+            this.cleanupRecords.set(id, session);
+        try {
+            if (cleanupRetryRequired) {
+                await session.disposeAfterCleanupFailure();
+            }
+            else {
+                await session.dispose();
+                this.cleanupRecords.delete(id);
+            }
+        }
+        catch (error) {
+            if (isChildCleanupError(error) || session.cleanupRetryRequired) {
+                this.cleanupRecords.set(id, session);
+            }
+            else {
+                console.error("pi-acp wedged-session resource disposal error:", error);
+            }
+        }
     }
-    async dispose() {
-        if (this.disposed)
+    dispose() {
+        if (this.disposePromise)
+            return this.disposePromise;
+        if (this.disposeSucceeded)
+            return Promise.resolve();
+        this.disposePromise = this.disposeGeneration()
+            .then(() => { this.disposeSucceeded = true; })
+            .finally(() => {
+            if (!this.disposeSucceeded)
+                this.disposePromise = undefined;
+        });
+        return this.disposePromise;
+    }
+    async disposeGeneration() {
+        if (this.disposed && this.cleanupRecords.size === 0)
             return;
-        this.disposed = true;
-        for (const controller of this.openingControllers)
-            controller.abort(new Error("agent disposed"));
-        await Promise.allSettled([...this.openingTasks]);
-        const sessions = [...this.live.values()];
+        if (!this.disposed) {
+            this.disposed = true;
+            for (const controller of this.openingControllers)
+                controller.abort(new Error("agent disposed"));
+            await Promise.allSettled([...this.openingTasks]);
+        }
+        const records = new Map(this.cleanupRecords);
+        for (const [id, session] of this.live)
+            records.set(id, session);
         this.live.clear();
-        await Promise.allSettled(sessions.map((session) => session.dispose()));
-        this.tombstones.clear();
+        const entries = [...records.entries()];
+        const results = await Promise.allSettled(entries.map(([, session]) => session.dispose()));
+        this.cleanupRecords.clear();
+        let sawChildFailure = false;
+        for (let index = 0; index < results.length; index += 1) {
+            const result = results[index];
+            if (result.status === "rejected" && isChildCleanupError(result.reason)) {
+                const [id, session] = entries[index];
+                this.cleanupRecords.set(id, session);
+                sawChildFailure = true;
+            }
+        }
+        if (sawChildFailure) {
+            const remainingChildren = [...this.cleanupRecords.values()]
+                .reduce((sum, session) => sum + session.remainingChildren, 0);
+            throw adapterError("child_cleanup_error", { details: { remainingChildren } });
+        }
     }
 }