@averyyy/pi-coding-agent 0.82.1-piclient.2 → 0.82.1-piclient.4

package/dist/core/agent-session.js

@@ -12,6 +12,7 @@
  *
  * Modes use this class and add their own I/O layer on top.
  */
+import { randomUUID } from "node:crypto";
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { basename, dirname } from "node:path";
 import { contentText } from "@earendil-works/pi-ai";
@@ -29,11 +30,15 @@ import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
 import { ModelRegistry } from "./model-registry.js";
-import { compactPiServer, syncPiServerTree, } from "./pi-server-client.js";
+import { acknowledgePiServerCompaction, acknowledgePiServerRunMessage, acquirePiServerSessionLease, compactPiServer, getPiServerRunDiagnostic, recordPiServerCompactionApplied, recoverPendingPiServerCompaction, releasePiServerRunLease, releasePiServerSessionLease, syncPiServerTree, } from "./pi-server-client.js";
+import { getPiServerCompactStatePath, readPiServerPendingCompact } from "./pi-server-compact-state.js";
+import { assertPiServerProviderExecution, } from "./pi-server-provider-execution.js";
+import { getPiServerRunStatePath, readPiServerPendingRun } from "./pi-server-run-state.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
 import { createSyntheticSourceInfo } from "./source-info.js";
 import { buildSystemPrompt } from "./system-prompt.js";
+import { beginSessionToolEffect, beginSessionToolEffectFinalization, commitSessionToolEffect, recoverToolEffects, writeSessionToolEffectFinalResult, writeSessionToolEffectResult, } from "./tool-effect-journal.js";
 import { createLocalBashOperations } from "./tools/bash.js";
 import { createAllToolDefinitions } from "./tools/index.js";
 import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js";
@@ -93,16 +98,26 @@ function getStringField(input, key) {
     const value = input[key];
     return typeof value === "string" && value.trim().length > 0 ? value : undefined;
 }
-function getPiServerFailurePhase(message) {
+function getPiServerFailureMetadata(message) {
     for (const diagnostic of message.diagnostics ?? []) {
         if (diagnostic.type !== "pi_server_failure")
             continue;
         const phase = diagnostic.details?.phase;
-        if (typeof phase === "string")
-            return phase;
+        const retryable = diagnostic.details?.retryable;
+        return {
+            phase: typeof phase === "string" ? phase : undefined,
+            retryable: typeof retryable === "boolean" ? retryable : undefined,
+        };
     }
     return undefined;
 }
+function isPiServerControlPlaneFailure(message) {
+    const hasFailure = message.diagnostics?.some((diagnostic) => diagnostic.type === "pi_server_failure") === true;
+    return hasFailure && getPiServerRunDiagnostic(message) === undefined;
+}
+function isPiServerRestartUncertainty(message) {
+    return getPiServerRunDiagnostic(message)?.restartUnknown === true;
+}
 function toolTextContent(content) {
     return content
         .filter((part) => part.type === "text")
@@ -165,6 +180,7 @@ export class AgentSession {
     _isAgentRunActive = false;
     _idleWaitPromise;
     _resolveIdleWait;
+    _toolResultsInFinalization = new WeakSet();
     /** Tracks pending steering messages for UI display. Removed when delivered. */
     _steeringMessages = [];
     /** Tracks pending follow-up messages for UI display. Removed when delivered. */
@@ -176,6 +192,7 @@ export class AgentSession {
     _autoCompactionAbortController = undefined;
     _overflowRecoveryAttempted = false;
     _lastServerCompactionTimestamp = undefined;
+    _piServerSessionLeaseOwner;
     // Branch summarization state
     _branchSummaryAbortController = undefined;
     // Retry state
@@ -208,6 +225,7 @@ export class AgentSession {
     _extensionErrorUnsubscriber;
     _autoSessionNameEnabled;
     _autoSessionNameAttempted = false;
+    _auxiliaryStreamFunction;
     _modelRuntime;
     // Tool registry for extension getTools/setTools
     _toolRegistry = new Map();
@@ -227,6 +245,7 @@ export class AgentSession {
         this._customTools = config.customTools ?? [];
         this._cwd = config.cwd;
         this._modelRuntime = config.modelRuntime;
+        this._auxiliaryStreamFunction = config.auxiliaryStreamFunction;
         this._extensionRunnerRef = config.extensionRunnerRef;
         this._initialActiveToolNames = config.initialActiveToolNames;
         this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined;
@@ -234,6 +253,10 @@ export class AgentSession {
         this._baseToolsOverride = config.baseToolsOverride;
         this._sessionStartEvent = config.sessionStartEvent ?? { type: "session_start", reason: "startup" };
         this._autoSessionNameEnabled = config.autoSessionName ?? true;
+        const toolEffectRecovery = recoverToolEffects(this.sessionManager);
+        if (toolEffectRecovery.recoveredToolCallIds.length > 0) {
+            this.agent.state.messages = this.sessionManager.buildSessionContext().messages;
+        }
         if (isPiServerMode()) {
             this._detachPiServerTerminalAssistantFailure();
         }
@@ -250,6 +273,9 @@ export class AgentSession {
     get modelRuntime() {
         return this._modelRuntime;
     }
+    _getAuxiliaryStreamFunction() {
+        return this._auxiliaryStreamFunction ?? this.agent.streamFunction;
+    }
     reconcilePiServerHistory(snapshot) {
         this.sessionManager.replaceTree(snapshot.entries, snapshot.leafId);
         const sessionContext = this.sessionManager.buildSessionContext();
@@ -260,6 +286,105 @@ export class AgentSession {
             messageCount: snapshot.messages.length,
         });
     }
+    async _syncAndAcknowledgePiServerAssistantMessage(message, signal) {
+        const persistenceSignal = signal.aborted ? undefined : signal;
+        const sessionFile = this.sessionManager.getSessionFile();
+        const runStatePath = sessionFile ? getPiServerRunStatePath(sessionFile) : undefined;
+        try {
+            await syncPiServerTree(this.sessionId, {
+                systemPrompt: this.systemPrompt,
+                messages: this.agent.state.messages,
+                tools: this.agent.state.tools,
+            }, { entries: this.sessionManager.getEntries(), leafId: this.sessionManager.getLeafId() }, {
+                signal: persistenceSignal,
+                onHistoryReconciled: (snapshot) => this.reconcilePiServerHistory(snapshot),
+            });
+            this.sessionManager.flushSessionFile();
+            await acknowledgePiServerRunMessage(this.sessionId, message, runStatePath, persistenceSignal);
+        }
+        catch (error) {
+            releasePiServerRunLease(this.sessionId, runStatePath);
+            throw error;
+        }
+    }
+    _acquirePiServerOperationLease(reentrantOwner) {
+        if (!isPiServerMode())
+            return { path: undefined, acquired: false, owner: undefined };
+        const sessionFile = this.sessionManager.getSessionFile();
+        const path = sessionFile ? getPiServerRunStatePath(sessionFile) : undefined;
+        if (this._piServerSessionLeaseOwner) {
+            if (reentrantOwner !== this._piServerSessionLeaseOwner) {
+                throw new Error(`Cannot start a concurrent durable pi-server operation for session ${this.sessionId}`);
+            }
+            return { path, acquired: false, owner: this._piServerSessionLeaseOwner };
+        }
+        const owner = Symbol("pi-server-session-lease");
+        if (path)
+            acquirePiServerSessionLease(this.sessionId, path);
+        this._piServerSessionLeaseOwner = owner;
+        return { path, acquired: true, owner };
+    }
+    _releasePiServerOperationLease(path, acquired, owner) {
+        if (!acquired)
+            return;
+        if (!owner || this._piServerSessionLeaseOwner !== owner) {
+            throw new Error(`Cannot release a durable pi-server operation lease owned by another operation`);
+        }
+        if (path)
+            releasePiServerSessionLease(this.sessionId, path);
+        this._piServerSessionLeaseOwner = undefined;
+    }
+    async _recoverPendingPiServerRunBeforeControlOperation(signal) {
+        if (!isPiServerMode())
+            return;
+        const sessionFile = this.sessionManager.getSessionFile();
+        if (!sessionFile) {
+            throw new Error("Durable pi-server operations require a persisted session file");
+        }
+        const runStatePath = getPiServerRunStatePath(sessionFile);
+        const pending = readPiServerPendingRun(runStatePath);
+        if (!pending)
+            return;
+        let terminal;
+        const entries = this.sessionManager.getEntries();
+        for (let index = entries.length - 1; index >= 0; index--) {
+            const entry = entries[index];
+            if (entry.type !== "message" || entry.message.role !== "assistant")
+                continue;
+            const message = entry.message;
+            const diagnostic = getPiServerRunDiagnostic(message);
+            if (!diagnostic || diagnostic.runId !== pending.runId)
+                continue;
+            if (diagnostic.sessionId !== pending.sessionId ||
+                diagnostic.requestMac !== pending.requestHash ||
+                pending.sessionId !== this.sessionId) {
+                throw new Error("Pending pi-server run terminal does not match its durable client marker");
+            }
+            terminal = message;
+            break;
+        }
+        if (!terminal) {
+            throw new Error(`Cannot start compaction while pi-server run ${pending.runId} is pending without a durable terminal assistant`);
+        }
+        await this._syncAndAcknowledgePiServerAssistantMessage(terminal, signal ?? new AbortController().signal);
+    }
+    async _recoverPendingPiServerCompaction(reason, willRetry, signal) {
+        if (!isPiServerMode())
+            return undefined;
+        const sessionFile = this.sessionManager.getSessionFile();
+        if (!sessionFile)
+            return undefined;
+        const compactStatePath = getPiServerCompactStatePath(sessionFile);
+        if (!readPiServerPendingCompact(compactStatePath))
+            return undefined;
+        const result = await recoverPendingPiServerCompaction(this.sessionId, { entries: this.sessionManager.getEntries(), leafId: this.sessionManager.getLeafId() }, compactStatePath, { signal });
+        if (!result)
+            return undefined;
+        return this._applyPiServerCompactionResult(result, reason, willRetry);
+    }
+    _assertPiServerProviderExecution(model) {
+        return assertPiServerProviderExecution(this._modelRuntime, model, this._extensionRunner);
+    }
     async _getRequiredRequestAuth(model) {
         let result;
         try {
@@ -288,7 +413,7 @@ export class AgentSession {
         throw new Error(formatNoApiKeyFoundMessage(model.provider));
     }
     async _getSummarizationRequestAuth(model) {
-        if (this.agent.streamFunction === streamSimple) {
+        if (this._getAuxiliaryStreamFunction() === streamSimple) {
             return this._getRequiredRequestAuth(model);
         }
         try {
@@ -480,18 +605,33 @@ export class AgentSession {
     }
     async _emitAgentSettled() {
         this._isAgentRunActive = false;
-        try {
-            await this._extensionRunner.emit({ type: "agent_settled" });
-            this._emit({ type: "agent_settled" });
-        }
-        finally {
-            this._resolveIdleWaitIfIdle();
-        }
+        await this._extensionRunner.emit({ type: "agent_settled" });
+        this._emit({ type: "agent_settled" });
     }
     // Track last assistant message for auto-compaction check
     _lastAssistantMessage = undefined;
     /** Internal handler for agent events - shared by subscribe and reconnect */
     _handleAgentEvent = async (event, signal) => {
+        if (event.type === "tool_execution_start") {
+            beginSessionToolEffect(this.sessionManager, event);
+        }
+        else if (event.type === "tool_execution_end") {
+            writeSessionToolEffectResult(this.sessionManager, event.toolResult, { finalizationPending: true });
+            this._toolResultsInFinalization.add(event.toolResult);
+        }
+        else if (event.type === "message_start" &&
+            event.message.role === "toolResult" &&
+            !this._toolResultsInFinalization.has(event.message)) {
+            beginSessionToolEffectFinalization(this.sessionManager, event.message);
+            this._toolResultsInFinalization.add(event.message);
+        }
+        else if (event.type === "message_end" &&
+            event.message.role === "toolResult" &&
+            !this._toolResultsInFinalization.has(event.message)) {
+            beginSessionToolEffectFinalization(this.sessionManager, event.message);
+            this._toolResultsInFinalization.add(event.message);
+        }
+        const finalizingToolResult = event.type === "message_end" && this._toolResultsInFinalization.has(event.message);
         // When a user message starts, check if it's from either queue and remove it BEFORE emitting
         // This ensures the UI sees the updated queue state
         if (event.type === "message_start" && event.message.role === "user") {
@@ -518,6 +658,13 @@ export class AgentSession {
         await this._emitExtensionEvent(event);
         // Notify all listeners
         this._emit(event.type === "agent_end" ? { ...event, willRetry: this._willRetryAfterAgentEnd(event) } : event);
+        if (finalizingToolResult) {
+            if (event.message.role !== "toolResult") {
+                throw new Error("A toolResult message_end replacement must remain a toolResult");
+            }
+            writeSessionToolEffectFinalResult(this.sessionManager, event.message);
+            this._toolResultsInFinalization.delete(event.message);
+        }
         // Handle session persistence
         if (event.type === "message_end") {
             // Check if this is a custom message from extensions
@@ -529,7 +676,11 @@ export class AgentSession {
                 event.message.role === "assistant" ||
                 event.message.role === "toolResult") {
                 // Regular LLM message - persist as SessionMessageEntry
-                this.sessionManager.appendMessage(event.message);
+                const sessionEntryId = this.sessionManager.appendMessage(event.message);
+                if (event.message.role === "toolResult") {
+                    this.sessionManager.flushSessionFile();
+                    commitSessionToolEffect(this.sessionManager, event.message, sessionEntryId);
+                }
             }
             // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
             // Track assistant message for auto-compaction (checked on agent_end)
@@ -549,12 +700,10 @@ export class AgentSession {
                     });
                     this._retryAttempt = 0;
                 }
-                if (isPiServerMode() && !this._isTerminalAssistantFailure(assistantMsg)) {
-                    await syncPiServerTree(this.sessionId, {
-                        systemPrompt: this.systemPrompt,
-                        messages: this.agent.state.messages,
-                        tools: this.agent.state.tools,
-                    }, { entries: this.sessionManager.getEntries(), leafId: this.sessionManager.getLeafId() }, { signal, onHistoryReconciled: (snapshot) => this.reconcilePiServerHistory(snapshot) });
+                if (isPiServerMode() &&
+                    !this._isPostAssistantPiServerSyncFailure(assistantMsg) &&
+                    !isPiServerControlPlaneFailure(assistantMsg)) {
+                    await this._syncAndAcknowledgePiServerAssistantMessage(assistantMsg, signal);
                 }
             }
         }
@@ -820,7 +969,7 @@ export class AgentSession {
     }
     /** Whether the session has no active agent run, retry, auto-compaction, or queued continuation. */
     get isIdle() {
-        return !this._isAgentRunActive;
+        return !this._isAgentRunActive && !this._piServerSessionLeaseOwner;
     }
     /** Current effective system prompt (includes any per-turn extension modifications) */
     get systemPrompt() {
@@ -972,9 +1121,14 @@ export class AgentSession {
     // Prompting
     // =========================================================================
     async _runAgentPrompt(messages) {
+        if (isPiServerMode() && this.model) {
+            this._assertPiServerProviderExecution(this.model);
+        }
+        const piServerLease = this._acquirePiServerOperationLease();
         this._resetValidationHintState();
         this._isAgentRunActive = true;
         try {
+            await this._recoverPendingPiServerCompaction("threshold", false);
             await this.agent.prompt(messages);
             while (await this._handlePostAgentRun()) {
                 try {
@@ -997,9 +1151,15 @@ export class AgentSession {
             await this._maybeGenerateSessionName();
         }
         finally {
-            this._systemPromptOverride = undefined;
-            this._flushPendingBashMessages();
-            await this._emitAgentSettled();
+            try {
+                this._systemPromptOverride = undefined;
+                this._flushPendingBashMessages();
+                await this._emitAgentSettled();
+            }
+            finally {
+                this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
+                this._resolveIdleWaitIfIdle();
+            }
         }
     }
     _sessionHasInfoEntry() {
@@ -1036,7 +1196,7 @@ export class AgentSession {
         if (!model || !transcript)
             return;
         try {
-            const stream = await this.agent.streamFunction(model, {
+            const stream = await this._getAuxiliaryStreamFunction()(model, {
                 systemPrompt: this.systemPrompt,
                 messages: [
                     {
@@ -1049,7 +1209,7 @@ export class AgentSession {
             }, {
                 maxTokens: 32,
                 reasoning: undefined,
-                sessionId: this.sessionId,
+                sessionId: randomUUID(),
                 transport: this.agent.transport,
                 thinkingBudgets: this.agent.thinkingBudgets,
                 maxRetryDelayMs: this.agent.maxRetryDelayMs,
@@ -1160,6 +1320,9 @@ export class AgentSession {
             if (!this.model) {
                 throw new Error(formatNoModelSelectedMessage());
             }
+            if (isPiServerMode()) {
+                this._assertPiServerProviderExecution(this.model);
+            }
             const hasConfiguredAuth = this._modelRuntime.hasConfiguredAuth(this.model.provider) ||
                 (await this._modelRuntime.checkAuth(this.model.provider)) !== undefined;
             if (!hasConfiguredAuth) {
@@ -1682,9 +1845,26 @@ export class AgentSession {
     // Compaction
     // =========================================================================
     async _applyPiServerCompactionResult(result, reason, willRetry) {
-        this.sessionManager.replaceTree(result.entries, result.leafId);
+        const compactionEntry = result.compactionEntry;
+        const operation = compactionEntry.piServerCompactOperation;
+        if (!operation ||
+            !Number.isSafeInteger(operation.baseEntryCount) ||
+            operation.baseEntryCount < 0 ||
+            (operation.baseLeafId !== null &&
+                (typeof operation.baseLeafId !== "string" || operation.baseLeafId.length === 0)) ||
+            typeof operation.baseTreeHash !== "string" ||
+            !/^[a-f0-9]{64}$/u.test(operation.baseTreeHash)) {
+            throw new Error("Pi-server compaction result is missing a valid durable base identity");
+        }
+        this.sessionManager.appendCompactionEntry(compactionEntry, {
+            baseEntryCount: operation.baseEntryCount,
+            baseLeafId: operation.baseLeafId,
+            baseTreeHash: operation.baseTreeHash,
+        });
+        recordPiServerCompactionApplied(result);
         const sessionContext = this.sessionManager.buildSessionContext();
         this.agent.state.messages = sessionContext.messages;
+        await acknowledgePiServerCompaction(result);
         const firstKeptEntryId = result.compaction.firstKeptEntryId;
         if (!firstKeptEntryId) {
             throw new Error("Pi-server compaction result is missing firstKeptEntryId");
@@ -1703,7 +1883,7 @@ export class AgentSession {
             await this._extensionRunner.emit({
                 type: "session_compact",
                 compactionEntry: savedCompactionEntry,
-                fromExtension: false,
+                fromExtension: savedCompactionEntry.fromHook === true,
                 reason,
                 willRetry,
             });
@@ -1716,42 +1896,36 @@ export class AgentSession {
      * @param customInstructions Optional instructions for the compaction summary
      */
     async compact(customInstructions) {
+        const piServerExecution = isPiServerMode() && this.model ? this._assertPiServerProviderExecution(this.model) : undefined;
         this._disconnectFromAgent();
         await this.abort();
         this._compactionAbortController = new AbortController();
         this._emit({ type: "compaction_start", reason: "manual" });
+        let piServerLease = {
+            path: undefined,
+            acquired: false,
+            owner: undefined,
+        };
         try {
-            if (!this.model) {
-                throw new Error(formatNoModelSelectedMessage());
-            }
-            const { apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model);
-            if (isPiServerMode()) {
-                this._detachPiServerTerminalAssistantFailure();
-                const piServerResult = await compactPiServer(this.model, {
-                    systemPrompt: this.systemPrompt,
-                    messages: this.agent.state.messages,
-                    tools: this.agent.state.tools,
-                }, {
-                    sessionId: this.sessionId,
-                    apiKey,
-                    headers,
-                    customInstructions,
-                    settings: this.settingsManager.getCompactionSettings(),
-                    sessionTree: { entries: this.sessionManager.getEntries(), leafId: this.sessionManager.getLeafId() },
-                    signal: this._compactionAbortController.signal,
-                    reasoning: toProviderReasoning(this.thinkingLevel),
-                    onHistoryReconciled: (snapshot) => this.reconcilePiServerHistory(snapshot),
-                });
-                const result = await this._applyPiServerCompactionResult(piServerResult, "manual", false);
-                this._lastServerCompactionTimestamp = Date.now();
+            piServerLease = this._acquirePiServerOperationLease();
+            await this._recoverPendingPiServerRunBeforeControlOperation(this._compactionAbortController.signal);
+            const recoveredCompaction = await this._recoverPendingPiServerCompaction("manual", false, this._compactionAbortController.signal);
+            if (recoveredCompaction) {
                 this._emit({
                     type: "compaction_end",
                     reason: "manual",
-                    result,
+                    result: recoveredCompaction,
                     aborted: false,
                     willRetry: false,
                 });
-                return result;
+                return recoveredCompaction;
+            }
+            if (!this.model) {
+                throw new Error(formatNoModelSelectedMessage());
+            }
+            const { apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model);
+            if (isPiServerMode()) {
+                this._detachPiServerTerminalAssistantFailure();
             }
             const pathEntries = this.sessionManager.getBranch();
             const settings = this.settingsManager.getCompactionSettings();
@@ -1784,6 +1958,43 @@ export class AgentSession {
                     fromExtension = true;
                 }
             }
+            if (isPiServerMode()) {
+                const sessionFile = this.sessionManager.getSessionFile();
+                if (!sessionFile) {
+                    throw new Error("Durable pi-server compaction requires a persisted session file");
+                }
+                const piServerResult = await compactPiServer(this.model, {
+                    systemPrompt: this.systemPrompt,
+                    messages: this.agent.state.messages,
+                    tools: this.agent.state.tools,
+                }, {
+                    sessionId: this.sessionId,
+                    apiKey,
+                    headers,
+                    env,
+                    customInstructions,
+                    settings,
+                    preparation: { firstKeptEntryId: preparation.firstKeptEntryId },
+                    extensionCompaction,
+                    retry: this.settingsManager.getRetrySettings(),
+                    sessionTree: { entries: this.sessionManager.getEntries(), leafId: this.sessionManager.getLeafId() },
+                    providerExecutionFingerprint: piServerExecution?.providerExecutionFingerprint,
+                    piServerCompactStatePath: getPiServerCompactStatePath(sessionFile),
+                    signal: this._compactionAbortController.signal,
+                    reasoning: toProviderReasoning(this.thinkingLevel),
+                    onHistoryReconciled: (snapshot) => this.reconcilePiServerHistory(snapshot),
+                });
+                const result = await this._applyPiServerCompactionResult(piServerResult, "manual", false);
+                this._lastServerCompactionTimestamp = Date.now();
+                this._emit({
+                    type: "compaction_end",
+                    reason: "manual",
+                    result,
+                    aborted: false,
+                    willRetry: false,
+                });
+                return result;
+            }
             let summary;
             let firstKeptEntryId;
             let tokensBefore;
@@ -1799,7 +2010,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this.agent.streamFunction, env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason: "manual" }));
+                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this._getAuxiliaryStreamFunction(), env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason: "manual" }));
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
@@ -1856,6 +2067,7 @@ export class AgentSession {
             throw error;
         }
         finally {
+            this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
             this._compactionAbortController = undefined;
             this._reconnectToAgent();
         }
@@ -1912,6 +2124,9 @@ export class AgentSession {
         // but must not retry: the assistant answer already completed and agent.continue() cannot
         // continue from an assistant message.
         if (sameModel && isContextOverflow(assistantMessage, contextWindow)) {
+            if (isPiServerMode() && this.model) {
+                this._assertPiServerProviderExecution(this.model);
+            }
             const willRetry = assistantMessage.stopReason !== "stop";
             if (!willRetry) {
                 return await this._runAutoCompaction("overflow", false);
@@ -1970,6 +2185,9 @@ export class AgentSession {
             contextTokens = directContextTokens;
         }
         if (shouldCompact(contextTokens, contextWindow, settings)) {
+            if (isPiServerMode() && this.model) {
+                this._assertPiServerProviderExecution(this.model);
+            }
             return await this._runAutoCompaction("threshold", false);
         }
         return false;
@@ -1979,45 +2197,34 @@ export class AgentSession {
      */
     async _runAutoCompaction(reason, willRetry, preparationOverride, preparationOptions) {
         const settings = this.settingsManager.getCompactionSettings();
+        const model = this.model;
+        const piServerExecution = isPiServerMode() && model ? this._assertPiServerProviderExecution(model) : undefined;
         let started = false;
+        let piServerLease = {
+            path: undefined,
+            acquired: false,
+            owner: undefined,
+        };
         try {
-            if (!this.model) {
-                return false;
+            piServerLease = this._acquirePiServerOperationLease(this._isAgentRunActive ? this._piServerSessionLeaseOwner : undefined);
+            if (!this._isAgentRunActive) {
+                await this._recoverPendingPiServerRunBeforeControlOperation();
             }
-            let apiKey;
-            let headers;
-            let env;
-            if (this.agent.streamFunction === streamSimple) {
-                ({ apiKey, headers, env } = await this._getRequiredRequestAuth(this.model));
-            }
-            else {
-                ({ apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model));
-            }
-            if (isPiServerMode()) {
+            const recoveredCompaction = await this._recoverPendingPiServerCompaction(reason, willRetry);
+            if (recoveredCompaction) {
                 this._emit({ type: "compaction_start", reason });
-                this._autoCompactionAbortController = new AbortController();
                 started = true;
-                const signal = this._autoCompactionAbortController.signal;
-                const piServerResult = await compactPiServer(this.model, {
-                    systemPrompt: this.systemPrompt,
-                    messages: this.agent.state.messages,
-                    tools: this.agent.state.tools,
-                }, {
-                    sessionId: this.sessionId,
-                    apiKey,
-                    headers,
-                    settings,
-                    preparation: preparationOptions,
-                    sessionTree: { entries: this.sessionManager.getEntries(), leafId: this.sessionManager.getLeafId() },
-                    signal,
-                    reasoning: toProviderReasoning(this.thinkingLevel),
-                    onHistoryReconciled: (snapshot) => this.reconcilePiServerHistory(snapshot),
-                });
-                const result = await this._applyPiServerCompactionResult(piServerResult, reason, willRetry);
                 this._lastServerCompactionTimestamp = Date.now();
-                this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
+                this._emit({ type: "compaction_end", reason, result: recoveredCompaction, aborted: false, willRetry });
                 return willRetry ? true : this.agent.hasQueuedMessages();
             }
+            if (!model) {
+                return false;
+            }
+            let apiKey;
+            let headers;
+            let env;
+            ({ apiKey, headers, env } = await this._getSummarizationRequestAuth(model));
             const pathEntries = this.sessionManager.getBranch();
             const preparation = preparationOverride ?? prepareCompaction(pathEntries, settings);
             if (!preparation) {
@@ -2026,6 +2233,7 @@ export class AgentSession {
             this._emit({ type: "compaction_start", reason });
             this._autoCompactionAbortController = new AbortController();
             started = true;
+            const signal = this._autoCompactionAbortController.signal;
             let extensionCompaction;
             let fromExtension = false;
             if (this._extensionRunner.hasHandlers("session_before_compact")) {
@@ -2036,7 +2244,7 @@ export class AgentSession {
                     customInstructions: undefined,
                     reason,
                     willRetry,
-                    signal: this._autoCompactionAbortController.signal,
+                    signal,
                 }));
                 if (extensionResult?.cancel) {
                     this._emit({
@@ -2053,6 +2261,36 @@ export class AgentSession {
                     fromExtension = true;
                 }
             }
+            if (isPiServerMode()) {
+                const sessionFile = this.sessionManager.getSessionFile();
+                if (!sessionFile) {
+                    throw new Error("Durable pi-server compaction requires a persisted session file");
+                }
+                const piServerResult = await compactPiServer(model, {
+                    systemPrompt: this.systemPrompt,
+                    messages: this.agent.state.messages,
+                    tools: this.agent.state.tools,
+                }, {
+                    sessionId: this.sessionId,
+                    apiKey,
+                    headers,
+                    env,
+                    settings,
+                    retry: this.settingsManager.getRetrySettings(),
+                    preparation: preparationOptions ?? { firstKeptEntryId: preparation.firstKeptEntryId },
+                    extensionCompaction,
+                    sessionTree: { entries: this.sessionManager.getEntries(), leafId: this.sessionManager.getLeafId() },
+                    providerExecutionFingerprint: piServerExecution?.providerExecutionFingerprint,
+                    piServerCompactStatePath: getPiServerCompactStatePath(sessionFile),
+                    signal,
+                    reasoning: toProviderReasoning(this.thinkingLevel),
+                    onHistoryReconciled: (snapshot) => this.reconcilePiServerHistory(snapshot),
+                });
+                const result = await this._applyPiServerCompactionResult(piServerResult, reason, willRetry);
+                this._lastServerCompactionTimestamp = Date.now();
+                this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
+                return willRetry ? true : this.agent.hasQueuedMessages();
+            }
             let summary;
             let firstKeptEntryId;
             let tokensBefore;
@@ -2068,7 +2306,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this.agent.streamFunction, env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason }));
+                const compactResult = await compact(preparation, model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this._getAuxiliaryStreamFunction(), env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason }));
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
@@ -2139,6 +2377,7 @@ export class AgentSession {
             return false;
         }
         finally {
+            this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
             this._autoCompactionAbortController = undefined;
         }
     }
@@ -2189,6 +2428,9 @@ export class AgentSession {
         }
         const settings = this.settingsManager.getCompactionSettings();
         const markerParentId = this.sessionManager.getLeafId();
+        if (isPiServerMode() && this.model) {
+            this._assertPiServerProviderExecution(this.model);
+        }
         const markerId = this._appendIntraTurnCompactionMarker();
         const preparationOptions = { firstKeptEntryId: markerId };
         const branchWithMarker = this.sessionManager.getBranch();
@@ -2554,8 +2796,12 @@ export class AgentSession {
             return false;
         if (this._isPostAssistantPiServerSyncFailure(message))
             return false;
-        const piServerPhase = getPiServerFailurePhase(message);
-        if (piServerPhase && piServerPhase !== "provider_stream")
+        if (isPiServerRestartUncertainty(message))
+            return false;
+        const piServerFailure = getPiServerFailureMetadata(message);
+        if (piServerFailure?.phase && piServerFailure.phase !== "provider_stream")
+            return false;
+        if (piServerFailure?.retryable === false)
             return false;
         if (message.errorMessage && this._isNonRetryableProviderLimitError(message.errorMessage))
             return false;
@@ -2793,6 +3039,9 @@ export class AgentSession {
      * @returns Result with editorText (if user message) and cancelled status
      */
     async navigateTree(targetId, options = {}) {
+        if (isPiServerMode() && options.summarize && this.model) {
+            this._assertPiServerProviderExecution(this.model);
+        }
         const oldLeafId = this.sessionManager.getLeafId();
         // No-op if already at target
         if (targetId === oldLeafId) {
@@ -2822,6 +3071,7 @@ export class AgentSession {
             replaceInstructions,
             label,
         };
+        const piServerLease = this._acquirePiServerOperationLease();
         // Set up abort controller for summarization
         this._branchSummaryAbortController = new AbortController();
         try {
@@ -2869,7 +3119,7 @@ export class AgentSession {
                     customInstructions,
                     replaceInstructions,
                     reserveTokens: branchSummarySettings.reserveTokens,
-                    streamFn: this.agent.streamFunction,
+                    streamFn: this._getAuxiliaryStreamFunction(),
                     retry: this.settingsManager.getRetrySettings(),
                     callbacks: this._summarizationRetryCallbacks({ source: "branchSummary" }),
                 });
@@ -2954,6 +3204,7 @@ export class AgentSession {
             return { editorText, cancelled: false, summaryEntry };
         }
         finally {
+            this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
             this._branchSummaryAbortController = undefined;
         }
     }