pi-agent-browser-native 0.8.3 → 0.9.1

package/dist/extensions/agent-browser/index.js

@@ -1,7 +1,6 @@
-import { AsyncLocalStorage } from "node:async_hooks";
 import { randomUUID } from "node:crypto";
 import { existsSync, readFileSync } from "node:fs";
-import { dirname, join } from "node:path";
+import { dirname, isAbsolute, join } from "node:path";
 import { fileURLToPath } from "node:url";
 import { Text } from "@earendil-works/pi-tui";
 import { batchHasSuccessfulCloseAll, getSuccessfulBatchCloseLifecycle } from "./lib/batch-lifecycle.js";
@@ -13,17 +12,19 @@ import { buildExecutionPlan, canUseHeadlessCompatibilityUserAgent, createEphemer
 import { deleteIdentityKeysInNamespace, extractExplicitNamespace, extractExplicitSessionName, getAgentBrowserSessionIdentityKey, isAgentBrowserSessionIdentityKeyInNamespace, isUpstreamEnvFlagEnabled, resolveAgentBrowserNamespace, scanUpstreamGlobalFlagOccurrences } from "./lib/argv-grammar.js";
 import { parseArgvDescriptor } from "./lib/argv-descriptor.js";
 import { needsManagedSession } from "./lib/command-policy.js";
-import { ManagedSessionRestoreState, resolveOwnedManagedSessionContext, withOwnedManagedSessionContext } from "./lib/managed-session-restore.js";
+import { ManagedSessionRestoreState } from "./lib/managed-session-restore.js";
 import { isRecord } from "./lib/parsing.js";
 import { getAgentBrowserProcessTimeoutMs, resolveAgentBrowserSocketDir, runAgentBrowserProcess } from "./lib/process.js";
 import { getCommandAwareProcessTimeoutMs } from "./lib/orchestration/browser-run/prepare/wait-timeouts.js";
 import { getAgentBrowserProcessEnvironment, withAgentBrowserProcessEnvironment, withIsolatedAgentBrowserEnvironment } from "./lib/process-environment.js";
 import { withNativeSessionDefaults } from "./lib/orchestration/native-session-defaults.js";
-import { getBrowserCwdError, resolveExecutionCwd, restoreManagedSessionCwd } from "./lib/execution-cwd.js";
+import { getBrowserCwdError, resolveExecutionCwd } from "./lib/execution-cwd.js";
 import { resolveOperationPaths } from "./lib/orchestration/operation-paths.js";
+import { AsyncExecutionQueue, KeyedAsyncExecutionQueue } from "./lib/orchestration/execution-queue.js";
+import { mergeBrowserRunArtifactManifest } from "./lib/orchestration/browser-run/artifact-merge.js";
 import { closeManagedSession, inspectManagedSessionDaemon } from "./lib/orchestration/browser-run/managed-session-daemon-policy.js";
 import { MINIMUM_AGENT_BROWSER_VERSION, SUPPORTED_AGENT_BROWSER_VERSION_LABEL, TARGET_AGENT_BROWSER_VERSION, getAgentBrowserVersionValidationError, parseAgentBrowserVersionOutput, } from "./lib/upstream-version.js";
-import { buildPromptPolicy, getLatestUserPrompt } from "./lib/prompt-policy.js";
+import { buildPromptPolicy, getLatestUserMessage, getMessageText } from "./lib/prompt-policy.js";
 import { isCloseAllCommand, isCloseCommand } from "./lib/command-taxonomy.js";
 import { hasLaunchScopedFlagToken } from "./lib/launch-scoped-flags.js";
 import { cleanupSecureTempArtifacts } from "./lib/temp.js";
@@ -37,14 +38,14 @@ import { applyAgentBrowserOutputPath, canWriteAgentBrowserOutput, normalizeReque
 import { appendScriptSessionLease, createBrowserCodeOutput, getScriptSessionLeasesFromBranch } from "./lib/orchestration/script-mode.js";
 import { resolveBrowserExecutionIdentity, withBrowserExecutionLock, withBrowserExecutionLocks } from "./lib/managed-session-policy-lock.js";
 import { AGENT_BROWSER_INSTRUCTION_GROUP, AGENT_BROWSER_TOOL_INVENTORY, registerAgentBrowserToolSurface } from "./lib/tool-surface.js";
-import { formatSessionArtifactRetentionSummary, getSessionArtifactManifestEntryKey, isPendingRecordingCommand, isSessionArtifactManifest, mergeSessionArtifactManifest, retirePendingRecordingManifestEntries } from "./lib/results/artifact-manifest.js";
+import { formatSessionArtifactRetentionSummary, getSessionArtifactManifestEntryKey, isSessionArtifactManifest, retirePendingRecordingManifestEntries } from "./lib/results/artifact-manifest.js";
 import { appendUniqueAgentBrowserNextActions, applyNamespaceToNextActions, applySessionToNextActions, buildNextToolAction } from "./lib/results/next-actions.js";
 import { canRegisterWebSearchTool, loadAgentBrowserConfigSync } from "./lib/config.js";
 import { appendRecordingReservationTransition, applyRecordingArtifactsToReservations, restoreRecordingReservationStateFromBranch, retireRecordingReservation, } from "./lib/recording-reservations.js";
 import { createAgentBrowserWebSearchTool } from "./lib/web-search.js";
 import { scopeReadConfirmationArgs } from "./lib/read-confirmation.js";
 import { isDirectAgentBrowserBashAllowed, isHarmlessAgentBrowserInspectionCommand, looksLikeDirectAgentBrowserBash, } from "./lib/bash-guard.js";
-import { AgentBrowserResultComponent, buildAgentBrowserToolResultPatch, formatAgentBrowserRenderCall, formatAgentBrowserRenderResult, } from "./lib/pi-tool-rendering.js";
+import { AgentBrowserResultComponent, formatAgentBrowserRenderCall, } from "./lib/pi-tool-rendering.js";
 function browserExecutionFailure(error, signal, fallbackCategory = "upstream-error") {
     const timeout = error instanceof Error && error.name === "TimeoutError";
     const summary = redactSensitiveText(error instanceof Error ? error.message : String(error));
@@ -230,84 +231,6 @@ function getRecognizedCompatibilityWorkaround(value) {
         ? { id: workaround.id, reason: workaround.reason }
         : undefined;
 }
-function restoreManagedSessionCompatibilityWorkaroundFromBranch(branch, sessionName, namespace) {
-    let restored;
-    const targetKey = getSessionContextKey(sessionName, namespace);
-    for (const entry of branch) {
-        const message = getBrowserResultMessage(entry);
-        if (!message)
-            continue;
-        const details = isRecord(message.details) ? message.details : undefined;
-        if (!details)
-            continue;
-        if (getSessionContextKey(typeof details.sessionName === "string" ? details.sessionName : undefined, typeof details.namespace === "string" ? details.namespace : undefined) !== targetKey)
-            continue;
-        const recognizedWorkaround = getRecognizedCompatibilityWorkaround(details.compatibilityWorkaround);
-        const succeeded = getSuccessfulToolResult(details, message);
-        const outcome = getManagedSessionOutcome(details);
-        const activeAfterFailure = recognizedWorkaround
-            && outcome?.activeAfter === true
-            && typeof outcome.currentSessionName === "string"
-            && getSessionContextKey(outcome.currentSessionName, typeof outcome.currentSessionNamespace === "string" ? outcome.currentSessionNamespace : undefined) === targetKey
-            && (outcome.status === "created" || outcome.status === "replaced" || outcome.status === "unchanged");
-        if (!succeeded && !activeAfterFailure)
-            continue;
-        if (recognizedWorkaround) {
-            restored = recognizedWorkaround;
-        }
-        else if (!canUseHeadlessCompatibilityUserAgent(getToolResultArgs(details))) {
-            restored = undefined;
-        }
-    }
-    return restored;
-}
-function restoreManagedSessionHeadedAutosaveDisabledFromBranch(branch, sessionName, namespace) {
-    let restored = false;
-    const targetKey = getSessionContextKey(sessionName, namespace);
-    for (const entry of branch) {
-        const message = getBrowserResultMessage(entry);
-        if (!message)
-            continue;
-        const details = isRecord(message.details) ? message.details : undefined;
-        if (!details)
-            continue;
-        if (getSessionContextKey(typeof details.sessionName === "string" ? details.sessionName : undefined, typeof details.namespace === "string" ? details.namespace : undefined) !== targetKey)
-            continue;
-        const outcome = getManagedSessionOutcome(details);
-        const activeAfterFailure = outcome?.activeAfter === true
-            && typeof outcome.currentSessionName === "string"
-            && getSessionContextKey(outcome.currentSessionName, typeof outcome.currentSessionNamespace === "string" ? outcome.currentSessionNamespace : undefined) === targetKey;
-        if ((getSuccessfulToolResult(details, message) || activeAfterFailure) && typeof details.managedSessionHeadedAutosaveDisabled === "boolean") {
-            restored = details.managedSessionHeadedAutosaveDisabled;
-        }
-    }
-    return restored;
-}
-function restoreManagedSessionHeadedAutosaveIntervalFromBranch(branch, sessionName, namespace) {
-    let restored;
-    const targetKey = getSessionContextKey(sessionName, namespace);
-    for (const entry of branch) {
-        const message = getBrowserResultMessage(entry);
-        if (!message)
-            continue;
-        const details = isRecord(message.details) ? message.details : undefined;
-        if (!details)
-            continue;
-        if (getSessionContextKey(typeof details.sessionName === "string" ? details.sessionName : undefined, typeof details.namespace === "string" ? details.namespace : undefined) !== targetKey)
-            continue;
-        const outcome = getManagedSessionOutcome(details);
-        const activeAfterFailure = outcome?.activeAfter === true
-            && typeof outcome.currentSessionName === "string"
-            && getSessionContextKey(outcome.currentSessionName, typeof outcome.currentSessionNamespace === "string" ? outcome.currentSessionNamespace : undefined) === targetKey;
-        if (!getSuccessfulToolResult(details, message) && !activeAfterFailure)
-            continue;
-        if (typeof details.managedSessionHeadedAutosaveInterval === "string")
-            restored = details.managedSessionHeadedAutosaveInterval;
-        else if (details.managedSessionHeadedAutosaveDisabled === true)
-            restored = "0";
-    }
-    return restored;
-}
 function getToolResultArgs(details) {
     if (Array.isArray(details.args) && details.args.every((arg) => typeof arg === "string"))
         return details.args;
@@ -385,18 +308,6 @@ function trackOwnedManagedSession(sessions, sessionName, cwd, options = {}) {
     const headedManagedAutosaveInterval = options.headedManagedAutosaveInterval ?? existing?.headedManagedAutosaveInterval;
     sessions.set(key, { branchOwned, compatibilityWorkaround, cwd, headedManagedAutosaveDisabled, headedManagedAutosaveInterval, namespace: options.namespace, sessionName, socketDir: options.socketDir ?? existing?.socketDir });
 }
-function restoreManagedSessionSocketDir(branch, sessionName, namespace) {
-    const key = getAgentBrowserSessionIdentityKey(sessionName, namespace);
-    for (let index = branch.length - 1; index >= 0; index -= 1) {
-        const details = getBrowserResultMessage(branch[index])?.details;
-        if (!isRecord(details) || typeof details.sessionName !== "string")
-            continue;
-        if (getAgentBrowserSessionIdentityKey(details.sessionName, typeof details.namespace === "string" ? details.namespace : undefined) === key
-            && typeof details.managedSessionSocketDir === "string")
-            return details.managedSessionSocketDir;
-    }
-    return undefined;
-}
 function untrackOwnedManagedSession(sessions, sessionName, namespace) {
     if (!sessionName)
         return;
@@ -590,6 +501,7 @@ function collectBranchManagedResourceEvents(branch) {
         managedSessionActiveIdentities: new Map(),
         managedSessionActiveRanks: new Map(),
         managedSessionCloseRanks: new Map(),
+        managedSessionLaunchState: new Map(),
     };
     let eventRank = 0;
     for (const entry of branch) {
@@ -611,6 +523,32 @@ function collectBranchManagedResourceEvents(branch) {
         const batchCloseLifecycle = getSuccessfulBatchCloseLifecycle(details.batchSteps);
         const closeAllApplied = detailsReportCloseAllApplied(details, succeeded);
         const outcome = getManagedSessionOutcome(details);
+        if (sessionName) {
+            const key = getSessionContextKey(sessionName, namespace) ?? sessionName;
+            const launch = events.managedSessionLaunchState.get(key) ?? { compatibilityWorkaround: undefined, headedManagedAutosaveDisabled: false };
+            if (typeof details.managedSessionCwd === "string" && isAbsolute(details.managedSessionCwd))
+                launch.cwd = details.managedSessionCwd;
+            if (typeof details.managedSessionSocketDir === "string")
+                launch.socketDir = details.managedSessionSocketDir;
+            const retainedAfterFailure = outcome?.activeAfter === true && typeof outcome.currentSessionName === "string"
+                && getSessionContextKey(outcome.currentSessionName, typeof outcome.currentSessionNamespace === "string" ? outcome.currentSessionNamespace : undefined) === key;
+            const workaround = getRecognizedCompatibilityWorkaround(details.compatibilityWorkaround);
+            if (succeeded || workaround && retainedAfterFailure && ["created", "replaced", "unchanged"].includes(String(outcome?.status))) {
+                if (workaround)
+                    launch.compatibilityWorkaround = workaround;
+                else if (!canUseHeadlessCompatibilityUserAgent(getToolResultArgs(details)))
+                    launch.compatibilityWorkaround = undefined;
+            }
+            if (succeeded || retainedAfterFailure) {
+                if (typeof details.managedSessionHeadedAutosaveDisabled === "boolean")
+                    launch.headedManagedAutosaveDisabled = details.managedSessionHeadedAutosaveDisabled;
+                if (typeof details.managedSessionHeadedAutosaveInterval === "string")
+                    launch.headedManagedAutosaveInterval = details.managedSessionHeadedAutosaveInterval;
+                else if (details.managedSessionHeadedAutosaveDisabled === true)
+                    launch.headedManagedAutosaveInterval = "0";
+            }
+            events.managedSessionLaunchState.set(key, launch);
+        }
         const outcomeSucceeded = outcome?.succeeded === true;
         const outcomeStatus = typeof outcome?.status === "string" ? outcome.status : undefined;
         const outcomeCurrentSessionName = typeof outcome?.currentSessionName === "string" ? outcome.currentSessionName : undefined;
@@ -710,90 +648,6 @@ function shouldSerializeBrowserCommand(options) {
         return true;
     return getActiveElectronRecords(options.ownedElectronLaunchRecords).some((record) => record.sessionName === options.explicitSessionName);
 }
-// Serializes managed-session read/modify/write work so overlapping tool calls cannot promote stale state or close an in-use session.
-class AsyncExecutionQueue {
-    tail = Promise.resolve();
-    active = new AsyncLocalStorage();
-    isCurrent() { return this.active.getStore()?.active === true; }
-    run(work, signal, barrier = Promise.resolve()) {
-        if (this.isCurrent()) {
-            signal?.throwIfAborted();
-            return work();
-        }
-        const previous = this.tail;
-        let release;
-        this.tail = new Promise((resolve) => {
-            release = resolve;
-        });
-        let rejectWaiting;
-        const cancelled = new Promise((_resolve, reject) => { rejectWaiting = reject; });
-        const abortWaiting = () => rejectWaiting(signal?.reason);
-        signal?.addEventListener("abort", abortWaiting, { once: true });
-        if (signal?.aborted)
-            abortWaiting();
-        const execution = (async () => {
-            await Promise.all([previous, barrier]);
-            signal?.removeEventListener("abort", abortWaiting);
-            try {
-                signal?.throwIfAborted();
-                const scope = { active: true };
-                try {
-                    return await this.active.run(scope, work);
-                }
-                finally {
-                    scope.active = false;
-                }
-            }
-            finally {
-                release();
-            }
-        })();
-        return signal ? Promise.race([execution, cancelled]) : execution;
-    }
-}
-export class KeyedAsyncExecutionQueue {
-    barriers = new Map();
-    entries = new Map();
-    async run(key, namespace, work, signal) {
-        const entry = this.entries.get(key) ?? { queue: new AsyncExecutionQueue(), users: 0 };
-        if (entry.queue.isCurrent())
-            return work();
-        const barrier = this.barriers.get(getAgentBrowserSessionIdentityKey("", namespace)) ?? Promise.resolve();
-        entry.users += 1;
-        this.entries.set(key, entry);
-        try {
-            return await entry.queue.run(work, signal, barrier);
-        }
-        finally {
-            entry.users -= 1;
-            if (entry.users === 0 && this.entries.get(key) === entry)
-                this.entries.delete(key);
-        }
-    }
-    async runExclusive(namespace, work) {
-        const namespaceKey = getAgentBrowserSessionIdentityKey("", namespace);
-        const previous = this.barriers.get(namespaceKey) ?? Promise.resolve();
-        let release;
-        const blocked = new Promise((resolve) => {
-            release = resolve;
-        });
-        const barrier = previous.then(() => blocked);
-        this.barriers.set(namespaceKey, barrier);
-        const drains = [...this.entries]
-            .filter(([key]) => isAgentBrowserSessionIdentityKeyInNamespace(key, namespace))
-            .map(([, { queue }]) => queue.run(async () => undefined));
-        await previous;
-        await Promise.all(drains);
-        try {
-            return await work();
-        }
-        finally {
-            release();
-            if (this.barriers.get(namespaceKey) === barrier)
-                this.barriers.delete(namespaceKey);
-        }
-    }
-}
 function mergeBrowserRunMap(current, initial, updated) {
     if (updated === initial)
         return current;
@@ -808,27 +662,6 @@ function mergeBrowserRunMap(current, initial, updated) {
     }
     return merged;
 }
-export function mergeBrowserRunArtifactManifest(current, initial, updated) {
-    if (!updated || updated === initial)
-        return current;
-    if (current === initial)
-        return updated;
-    const initialEntries = new Map((initial?.entries ?? []).map((entry) => [getSessionArtifactManifestEntryKey(entry), entry]));
-    const changedEntries = updated.entries
-        .map((entry, index) => ({ entry, index }))
-        .filter(({ entry }) => initialEntries.get(getSessionArtifactManifestEntryKey(entry)) !== entry)
-        .sort((left, right) => left.entry.createdAtMs - right.entry.createdAtMs
-        || Number(isPendingRecordingCommand(left.entry.command, left.entry.subcommand, left.entry.kind)) - Number(isPendingRecordingCommand(right.entry.command, right.entry.subcommand, right.entry.kind))
-        || left.index - right.index)
-        .map(({ entry }) => entry);
-    return changedEntries.length === 0
-        ? current
-        : mergeSessionArtifactManifest({
-            base: current,
-            entries: changedEntries,
-            nowMs: Math.max(Date.now(), (current?.updatedAtMs ?? 0) + 1, updated.updatedAtMs),
-        });
-}
 function findPackageRoot(startDir) {
     let currentDir = startDir;
     while (true) {
@@ -858,7 +691,7 @@ function hasArgvFlag(argv, longFlag, shortFlag) {
 function shouldIncludeProjectConfig(ctx, argv = process.argv) {
     if (hasArgvFlag(argv, "--no-approve", "-na"))
         return false;
-    return ctx?.isProjectTrusted?.() ?? true;
+    return ctx.isProjectTrusted();
 }
 export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
     const ephemeralSessionSeed = createEphemeralSessionSeed();
@@ -868,6 +701,21 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
     });
     const implicitSessionIdleTimeoutMs = String(getImplicitSessionIdleTimeoutMs());
     const implicitSessionCloseTimeoutMs = getImplicitSessionCloseTimeoutMs();
+    let promptState;
+    const restorePromptState = (ctx) => {
+        return promptState = { sessionId: ctx.sessionManager.getSessionId(), message: getLatestUserMessage(ctx.sessionManager) };
+    };
+    const getPromptPolicy = (ctx) => {
+        const state = promptState?.sessionId === ctx.sessionManager.getSessionId() ? promptState : restorePromptState(ctx);
+        // Native message_end precedes later transforms/persistence; retain the message
+        // reference so replacement user text cannot leave policy on the earlier text.
+        const text = getMessageText(state.message?.content);
+        if (!state.policy || state.text !== text) {
+            state.text = text;
+            state.policy = buildPromptPolicy(text);
+        }
+        return state.policy;
+    };
     let webSearchToolRegistered = false;
     let managedSessionActive = false;
     let managedSessionBaseName = createImplicitSessionName(undefined, process.cwd(), ephemeralSessionSeed);
@@ -1121,11 +969,11 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
         clearSessionScopedBrowserState(sessionName, AGENT_BROWSER_SCRIPT_NAMESPACE);
         return undefined;
     };
-    const recoverScriptSessionLeasesWithinQueue = async (ctx) => {
+    const recoverScriptSessionLeasesWithinQueue = async (ctx, branch) => {
         const pendingSessionNames = new Set([...ownedManagedSessions.values()]
             .map((session) => session.sessionName)
             .filter(isAgentBrowserScriptSessionName));
-        for (const lease of getScriptSessionLeasesFromBranch(await readBrowserEntries(ctx.sessionManager), ctx.sessionManager.getSessionId()).values()) {
+        for (const lease of getScriptSessionLeasesFromBranch(branch, ctx.sessionManager.getSessionId()).values()) {
             if (lease.cleanup !== "closed")
                 pendingSessionNames.add(lease.sessionName);
         }
@@ -1138,6 +986,7 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
     const restoreBranchBackedState = async (ctx, options) => {
         branchRestoreGeneration += 1;
         branchStateGeneration += 1;
+        restorePromptState(ctx);
         const previousManagedSessionActive = managedSessionActive;
         const previousManagedSessionName = managedSessionName;
         const previousFreshSessionOrdinal = freshSessionOrdinal;
@@ -1179,15 +1028,11 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
                 : createFreshSessionName(managedSessionBaseName, ephemeralSessionSeed, nextFreshSessionOrdinal)
             : restoredState.sessionName;
         managedSessionNamespace = shouldReservePostCloseSession ? undefined : restoredState.namespace;
-        managedSessionCompatibilityWorkaround = managedSessionActive
-            ? restoreManagedSessionCompatibilityWorkaroundFromBranch(branch, managedSessionName, managedSessionNamespace)
-            : undefined;
-        managedSessionHeadedAutosaveDisabled = managedSessionActive
-            && restoreManagedSessionHeadedAutosaveDisabledFromBranch(branch, managedSessionName, managedSessionNamespace);
-        managedSessionHeadedAutosaveInterval = managedSessionActive
-            ? restoreManagedSessionHeadedAutosaveIntervalFromBranch(branch, managedSessionName, managedSessionNamespace)
-            : undefined;
-        managedSessionCwd = restoreManagedSessionCwd(branch, managedSessionName, managedSessionNamespace, ctx.cwd);
+        const currentLaunch = branchResourceEvents.managedSessionLaunchState.get(getSessionContextKey(managedSessionName, managedSessionNamespace) ?? managedSessionName);
+        managedSessionCompatibilityWorkaround = managedSessionActive ? currentLaunch?.compatibilityWorkaround : undefined;
+        managedSessionHeadedAutosaveDisabled = managedSessionActive && currentLaunch?.headedManagedAutosaveDisabled === true;
+        managedSessionHeadedAutosaveInterval = managedSessionActive ? currentLaunch?.headedManagedAutosaveInterval : undefined;
+        managedSessionCwd = currentLaunch?.cwd ?? ctx.cwd;
         freshSessionOrdinal = nextFreshSessionOrdinal;
         sessionPageState = SessionPageState.fromBranch(branch);
         traceOwners = new Map();
@@ -1238,13 +1083,14 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
                 continue;
             if (!isRestorableManagedSessionName(identity.sessionName, managedSessionBaseName))
                 continue;
-            trackOwnedManagedSession(ownedManagedSessions, identity.sessionName, restoreManagedSessionCwd(branch, identity.sessionName, identity.namespace, ctx.cwd), {
+            const launch = branchResourceEvents.managedSessionLaunchState.get(sessionKey);
+            trackOwnedManagedSession(ownedManagedSessions, identity.sessionName, launch?.cwd ?? ctx.cwd, {
                 branchOwned: true,
-                compatibilityWorkaround: restoreManagedSessionCompatibilityWorkaroundFromBranch(branch, identity.sessionName, identity.namespace),
-                headedManagedAutosaveDisabled: restoreManagedSessionHeadedAutosaveDisabledFromBranch(branch, identity.sessionName, identity.namespace),
-                headedManagedAutosaveInterval: restoreManagedSessionHeadedAutosaveIntervalFromBranch(branch, identity.sessionName, identity.namespace),
+                compatibilityWorkaround: launch?.compatibilityWorkaround,
+                headedManagedAutosaveDisabled: launch?.headedManagedAutosaveDisabled ?? false,
+                headedManagedAutosaveInterval: launch?.headedManagedAutosaveInterval,
                 namespace: identity.namespace,
-                socketDir: restoreManagedSessionSocketDir(branch, identity.sessionName, identity.namespace),
+                socketDir: launch?.socketDir,
             });
         }
         if (restoredState.active) {
@@ -1254,7 +1100,7 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
                 headedManagedAutosaveDisabled: managedSessionHeadedAutosaveDisabled,
                 headedManagedAutosaveInterval: managedSessionHeadedAutosaveInterval,
                 namespace: restoredState.namespace,
-                socketDir: restoreManagedSessionSocketDir(branch, restoredState.sessionName, restoredState.namespace),
+                socketDir: currentLaunch?.socketDir,
             });
         }
         const pendingOwnedBegins = new Map();
@@ -1290,11 +1136,12 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
             const closeRank = branchResourceEvents.managedSessionCloseRanks.get(sessionKey);
             if (activeRank === undefined || (closeRank !== undefined && closeRank >= activeRank))
                 continue;
-            trackOwnedManagedSession(ownedManagedSessions, record.sessionName, restoreManagedSessionCwd(branch, record.sessionName, undefined, ctx.cwd), {
+            const launch = branchResourceEvents.managedSessionLaunchState.get(sessionKey);
+            trackOwnedManagedSession(ownedManagedSessions, record.sessionName, launch?.cwd ?? ctx.cwd, {
                 branchOwned: true,
-                headedManagedAutosaveDisabled: restoreManagedSessionHeadedAutosaveDisabledFromBranch(branch, record.sessionName),
-                headedManagedAutosaveInterval: restoreManagedSessionHeadedAutosaveIntervalFromBranch(branch, record.sessionName),
-                socketDir: restoreManagedSessionSocketDir(branch, record.sessionName),
+                headedManagedAutosaveDisabled: launch?.headedManagedAutosaveDisabled ?? false,
+                headedManagedAutosaveInterval: launch?.headedManagedAutosaveInterval,
+                socketDir: launch?.socketDir,
             });
         }
         mergeActiveElectronLaunchRecords(ownedElectronLaunchRecords, new Map([...electronLaunchRecords].filter(([, record]) => record.ownerSessionId === browserOwnerSessionId)), {
@@ -1323,6 +1170,7 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
                     attachedSessionKeys.add(sessionKey);
             }
         }
+        return branch;
     };
     const registerWebSearchToolIfAvailable = (configState) => {
         if (webSearchToolRegistered || !canRegisterWebSearchTool(configState))
@@ -1339,15 +1187,19 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
         });
         webSearchToolRegistered = true;
     };
+    pi.on("message_end", (event, ctx) => {
+        if (event.message.role === "user")
+            promptState = { sessionId: ctx.sessionManager.getSessionId(), message: event.message };
+    });
     pi.on("session_start", async (_event, ctx) => {
-        await restoreBranchBackedState(ctx, { resetRuntimeOwnership: true });
-        electronChildProcesses = new Map();
-        registerWebSearchToolIfAvailable(loadAgentBrowserConfigSync({
-            cwd: ctx.cwd,
-            includeProjectConfig: shouldIncludeProjectConfig(ctx),
-        }));
         await managedSessionExecutionQueue.run(() => artifactExecutionQueue.run(async () => {
-            await recoverScriptSessionLeasesWithinQueue(ctx);
+            const branch = await restoreBranchBackedState(ctx, { resetRuntimeOwnership: true });
+            electronChildProcesses = new Map();
+            registerWebSearchToolIfAvailable(loadAgentBrowserConfigSync({
+                cwd: ctx.cwd,
+                includeProjectConfig: shouldIncludeProjectConfig(ctx),
+            }));
+            await recoverScriptSessionLeasesWithinQueue(ctx, branch);
             flushRecordingReservations();
             notifyRecordingPersistence(ctx);
         }));
@@ -1358,87 +1210,12 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
             controller.abort();
         await Promise.allSettled([...activeScriptExecutions]);
         await managedSessionExecutionQueue.run(() => artifactExecutionQueue.run(async () => {
-            await restoreBranchBackedState(ctx, { resetRuntimeOwnership: false });
-            await recoverScriptSessionLeasesWithinQueue(ctx);
+            const branch = await restoreBranchBackedState(ctx, { resetRuntimeOwnership: false });
+            await recoverScriptSessionLeasesWithinQueue(ctx, branch);
             flushRecordingReservations();
             notifyRecordingPersistence(ctx);
         }));
     });
-    // Fork-only checkpoint event; official Pi validation types do not declare it.
-    // Pi owns awaited tools/events (including their execution queues). Detached code
-    // children and retained browser/recording resources still need extension-side checks.
-    const checkpointPi = pi;
-    checkpointPi.on("session_checkpoint", async (event, ctx) => {
-        const blocked = (reason) => ({ sleepReady: false, reason });
-        if (browserReplayError)
-            return blocked(browserReplayError);
-        if (activeScriptControllers.size || activeScriptExecutions.size)
-            return blocked("Browser code execution is pending");
-        flushRecordingReservations();
-        if (recordingReservationsDirty)
-            return blocked("Browser recording journal persistence is dirty");
-        if (activeRecordingReservations.size)
-            return blocked("Browser recording is pending; finish it explicitly before sleep");
-        const entries = await readBrowserEntries(ctx.sessionManager, true);
-        if ([...getScriptSessionLeasesFromBranch(entries).values()].some(lease => lease.cleanup !== "closed")
-            || [...ownedManagedSessions.values()].some(owner => isAgentBrowserScriptSessionName(owner.sessionName))) {
-            return blocked("Browser script cleanup lease is unresolved");
-        }
-        if (getActiveElectronRecords(ownedElectronLaunchRecords).length || getActiveElectronRecords(electronLaunchRecords).length
-            || [...electronChildProcesses.values()].some(child => child.exitCode === null && child.signalCode === null)) {
-            return blocked("Electron launch is still active");
-        }
-        if (attachedSessionKeys.size || restoreAttachedSessionKeysFromBranch(entries).size)
-            return blocked("Attached browser state is caller-owned and not checkpointed");
-        if (traceOwners.size || [...networkRoutesBySession.values()].some(routes => routes.length))
-            return blocked("Browser trace or network routes are still active");
-        // Include off-branch and caller-owned/root identities, without acquiring
-        // cleanup ownership. Transcript page/ref details do not serialize a browser.
-        const historicalResources = collectBranchManagedResourceEvents(entries);
-        const sessions = new Map(ownedManagedSessions);
-        if (managedSessionActive)
-            trackOwnedManagedSession(sessions, managedSessionName, managedSessionCwd, { namespace: managedSessionNamespace });
-        for (const entry of entries) {
-            const message = getBrowserResultMessage(entry);
-            if (!message)
-                continue;
-            const details = isRecord(message.details) ? message.details : undefined;
-            // Helpers can launch before a main-process failure. Conversely, native
-            // launch flags on sessionless reads can open the unnamed default browser.
-            const sessionName = typeof details?.sessionName === "string" ? details.sessionName
-                : details?.agentBrowserStarted === true ? "default" : undefined;
-            if (!sessionName)
-                continue;
-            const namespace = typeof details?.namespace === "string" ? details.namespace : undefined;
-            const key = getSessionContextKey(sessionName, namespace) ?? sessionName;
-            if (!sessions.has(key))
-                trackOwnedManagedSession(sessions, sessionName, ctx.cwd, { namespace });
-        }
-        for (const [key, owner] of sessions) {
-            event.signal.throwIfAborted();
-            // An abnormal restart can leave a wrapper-created browser only off-branch.
-            // Reuse native ownership events for inspection, never cleanup ownership.
-            // A later close can refer to the same identity in a different socket root.
-            const historicalIdentity = historicalResources.managedSessionActiveIdentities.get(key);
-            const historicallyOwned = historicalIdentity && isRestorableManagedSessionName(historicalIdentity.sessionName, managedSessionBaseName);
-            const context = resolveOwnedManagedSessionContext({
-                ...owner,
-                currentManagedSessionName: managedSessionActive ? managedSessionName : undefined,
-                currentManagedSessionNamespace: managedSessionNamespace,
-                recordedOwnedSession: ownedManagedSessions.get(key) ?? (historicallyOwned ? { ...historicalIdentity, cwd: ctx.cwd } : undefined),
-                restoreState: managedSessionRestoreState,
-            });
-            // Owned routing must not mask an explicit caller's same-name ambient daemon.
-            for (const inspectionContext of context ? [{ ...context, reuseOnly: true }, undefined] : [undefined]) {
-                const daemon = await withAgentBrowserProcessEnvironment(inspectionContext && owner.socketDir ? { PI_AGENT_BROWSER_SOCKET_DIR: owner.socketDir } : {}, () => withOwnedManagedSessionContext(inspectionContext, () => inspectManagedSessionDaemon({ ...owner, signal: event.signal, timeoutMs: 2_000 })));
-                if (daemon.status !== "inactive")
-                    return blocked("Browser daemon is live or unverified; finish browser work explicitly before sleep");
-            }
-        }
-        // No detached writer remains to pause/resume. Native ingress stays held;
-        // checkpoint never closes a browser or changes ordinary shutdown ownership.
-        return { sleepReady: true };
-    });
     pi.on("session_shutdown", async (event, ctx) => {
         for (const controller of activeScriptControllers)
             controller.abort();
@@ -1548,11 +1325,11 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
             event.systemPromptOptions.sections.agent_browser = fullInstructions(ctx);
     });
     pi.on("tool_call", async (event, ctx) => {
-        const promptPolicy = buildPromptPolicy(getLatestUserPrompt(ctx.sessionManager.getBranch()));
-        if (isBashToolCallEvent(event) &&
-            !promptPolicy.allowLegacyAgentBrowserBash &&
-            looksLikeDirectAgentBrowserBash(event.input.command) &&
-            !isHarmlessAgentBrowserInspectionCommand(event.input.command) &&
+        if (!isBashToolCallEvent(event) || !looksLikeDirectAgentBrowserBash(event.input.command)
+            || isHarmlessAgentBrowserInspectionCommand(event.input.command))
+            return;
+        const promptPolicy = getPromptPolicy(ctx);
+        if (!promptPolicy.allowLegacyAgentBrowserBash &&
             !(await isDirectAgentBrowserBashAllowed(ctx.cwd))) {
             return {
                 block: true,
@@ -1560,7 +1337,6 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
             };
         }
     });
-    pi.on("tool_result", async (event) => buildAgentBrowserToolResultPatch(event));
     let emptyBrowserBranch;
     const captureCurrentBrowserBranch = (ctx) => {
         const generation = branchRestoreGeneration;
@@ -1583,7 +1359,7 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
             const component = context.lastComponent instanceof AgentBrowserResultComponent
                 ? context.lastComponent
                 : new AgentBrowserResultComponent();
-            component.setState(formatAgentBrowserRenderResult(result, options, theme, context.isError), options.expanded, theme);
+            component.setResult(result, options, theme, context.isError);
             return component;
         },
         async execute(toolCallId, params, signal, onUpdate, ctx, nativeToolCallId = toolCallId, capturedCwd, modelVisible = true, commandIndex, selectedBranch) {
@@ -1595,7 +1371,7 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
             catch (error) {
                 return buildValidationFailureResult({ kind: "invalid", status: "invalid", redactedArgs: [], toolArgs: [], validationError: error instanceof Error ? error.message : String(error) });
             }
-            const promptPolicy = buildPromptPolicy(getLatestUserPrompt(ctx.sessionManager.getBranch()));
+            const promptPolicy = getPromptPolicy(ctx);
             const outputPath = isRecord(params) && typeof params.outputPath === "string" ? params.outputPath : undefined;
             const admittedInput = resolveAgentBrowserInput({
                 getBatchPreflightValidationError: (args, stdin) => getArtifactPreflightValidationError({ args, cwd: operationCwd, outputPath, stdin }),
@@ -2266,10 +2042,10 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
             }) };
         const rendered = await renderAgentBrowserObservation(renderOptions);
         if (!priorManifest && !rendered.artifactManifest)
-            return { ...result, content: rendered.content, details: { ...details, browserEventVersion: 1 } };
+            return { ...result, content: rendered.content, structuredContent: rendered.structuredContent, details: { ...details, browserEventVersion: 1 } };
         return artifactExecutionQueue.run(async () => {
             if (!branch.isCurrent())
-                return { ...result, content: rendered.content, details: { ...details, browserEventVersion: 1, artifactManifest: rendered.artifactManifest } };
+                return { ...result, content: rendered.content, structuredContent: rendered.structuredContent, details: { ...details, browserEventVersion: 1, artifactManifest: rendered.artifactManifest } };
             const before = artifactManifest;
             const merged = mergeBrowserRunArtifactManifest(artifactManifest, priorManifest, rendered.artifactManifest);
             const changes = artifactChanges(before, merged);
@@ -2277,10 +2053,10 @@ export default function agentBrowserExtension(pi, { beforeExecute } = {}) {
                 const record = { event: { version: 1, phase: "state", operationId: randomUUID(), toolCallId: "observation", commandIndex: 0, isError: result.isError === true, state: {}, artifacts: changes } };
                 if (!await appendBrowserRecord(ctx.sessionManager, () => appendBrowserTransition(pi, record), record, branch)
                     || !branch.isCurrent())
-                    return { ...result, content: rendered.content, details: { ...details, browserEventVersion: 1, artifactManifest: rendered.artifactManifest } };
+                    return { ...result, content: rendered.content, structuredContent: rendered.structuredContent, details: { ...details, browserEventVersion: 1, artifactManifest: rendered.artifactManifest } };
             }
             artifactManifest = mergeBrowserRunArtifactManifest(artifactManifest, priorManifest, rendered.artifactManifest);
-            return { ...result, content: rendered.content, details: { ...details, browserEventVersion: 1, artifactManifest: invocationArtifactManifest(artifactManifest, before, details, priorManifest) } };
+            return { ...result, content: rendered.content, structuredContent: rendered.structuredContent, details: { ...details, browserEventVersion: 1, artifactManifest: invocationArtifactManifest(artifactManifest, before, details, priorManifest) } };
         });
     };
     registerAgentBrowserToolSurface(pi, {

package/dist/extensions/agent-browser/lib/argv-descriptor.js

@@ -76,7 +76,7 @@ function getOpenCommandTarget(commandTokens) {
     }
     return undefined;
 }
-export function parseCommandInfoFromTokens(commandTokens) {
+function parseCommandInfoFromTokens(commandTokens) {
     const upstreamCommandTokens = stripUpstreamGlobalFlags(commandTokens);
     const command = upstreamCommandTokens[0];
     return {