@averyyy/pi-coding-agent 0.82.1-piclient.5 → 0.82.1-piclient.7

package/dist/core/agent-session.js

@@ -12,7 +12,6 @@
  *
  * 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";
@@ -30,14 +29,11 @@ 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 { 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 { getPiServerRunStatePath, readPiServerPendingRun } from "./pi-server-run-state.js";
+import { compactPiServer, syncPiServerTree, } from "./pi-server-client.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";
@@ -97,26 +93,16 @@ function getStringField(input, key) {
     const value = input[key];
     return typeof value === "string" && value.trim().length > 0 ? value : undefined;
 }
-function getPiServerFailureMetadata(message) {
+function getPiServerFailurePhase(message) {
     for (const diagnostic of message.diagnostics ?? []) {
         if (diagnostic.type !== "pi_server_failure")
             continue;
         const phase = diagnostic.details?.phase;
-        const retryable = diagnostic.details?.retryable;
-        return {
-            phase: typeof phase === "string" ? phase : undefined,
-            retryable: typeof retryable === "boolean" ? retryable : undefined,
-        };
+        if (typeof phase === "string")
+            return phase;
     }
     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")
@@ -179,7 +165,6 @@ 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. */
@@ -191,7 +176,6 @@ export class AgentSession {
     _autoCompactionAbortController = undefined;
     _overflowRecoveryAttempted = false;
     _lastServerCompactionTimestamp = undefined;
-    _piServerSessionLeaseOwner;
     // Branch summarization state
     _branchSummaryAbortController = undefined;
     // Retry state
@@ -224,7 +208,6 @@ export class AgentSession {
     _extensionErrorUnsubscriber;
     _autoSessionNameEnabled;
     _autoSessionNameAttempted = false;
-    _auxiliaryStreamFunction;
     _modelRuntime;
     // Tool registry for extension getTools/setTools
     _toolRegistry = new Map();
@@ -244,7 +227,6 @@ 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;
@@ -252,10 +234,6 @@ 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();
         }
@@ -272,9 +250,6 @@ 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();
@@ -285,102 +260,6 @@ 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);
-    }
     async _getRequiredRequestAuth(model) {
         let result;
         try {
@@ -409,7 +288,7 @@ export class AgentSession {
         throw new Error(formatNoApiKeyFoundMessage(model.provider));
     }
     async _getSummarizationRequestAuth(model) {
-        if (this._getAuxiliaryStreamFunction() === streamSimple) {
+        if (this.agent.streamFunction === streamSimple) {
             return this._getRequiredRequestAuth(model);
         }
         try {
@@ -601,33 +480,18 @@ export class AgentSession {
     }
     async _emitAgentSettled() {
         this._isAgentRunActive = false;
-        await this._extensionRunner.emit({ type: "agent_settled" });
-        this._emit({ type: "agent_settled" });
+        try {
+            await this._extensionRunner.emit({ type: "agent_settled" });
+            this._emit({ type: "agent_settled" });
+        }
+        finally {
+            this._resolveIdleWaitIfIdle();
+        }
     }
     // 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") {
@@ -654,13 +518,6 @@ 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
@@ -672,11 +529,7 @@ export class AgentSession {
                 event.message.role === "assistant" ||
                 event.message.role === "toolResult") {
                 // Regular LLM message - persist as SessionMessageEntry
-                const sessionEntryId = this.sessionManager.appendMessage(event.message);
-                if (event.message.role === "toolResult") {
-                    this.sessionManager.flushSessionFile();
-                    commitSessionToolEffect(this.sessionManager, event.message, sessionEntryId);
-                }
+                this.sessionManager.appendMessage(event.message);
             }
             // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
             // Track assistant message for auto-compaction (checked on agent_end)
@@ -696,10 +549,12 @@ export class AgentSession {
                     });
                     this._retryAttempt = 0;
                 }
-                if (isPiServerMode() &&
-                    !this._isPostAssistantPiServerSyncFailure(assistantMsg) &&
-                    !isPiServerControlPlaneFailure(assistantMsg)) {
-                    await this._syncAndAcknowledgePiServerAssistantMessage(assistantMsg, signal);
+                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) });
                 }
             }
         }
@@ -965,7 +820,7 @@ export class AgentSession {
     }
     /** Whether the session has no active agent run, retry, auto-compaction, or queued continuation. */
     get isIdle() {
-        return !this._isAgentRunActive && !this._piServerSessionLeaseOwner;
+        return !this._isAgentRunActive;
     }
     /** Current effective system prompt (includes any per-turn extension modifications) */
     get systemPrompt() {
@@ -1117,11 +972,9 @@ export class AgentSession {
     // Prompting
     // =========================================================================
     async _runAgentPrompt(messages) {
-        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 {
@@ -1144,15 +997,9 @@ export class AgentSession {
             await this._maybeGenerateSessionName();
         }
         finally {
-            try {
-                this._systemPromptOverride = undefined;
-                this._flushPendingBashMessages();
-                await this._emitAgentSettled();
-            }
-            finally {
-                this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
-                this._resolveIdleWaitIfIdle();
-            }
+            this._systemPromptOverride = undefined;
+            this._flushPendingBashMessages();
+            await this._emitAgentSettled();
         }
     }
     _sessionHasInfoEntry() {
@@ -1189,7 +1036,7 @@ export class AgentSession {
         if (!model || !transcript)
             return;
         try {
-            const stream = await this._getAuxiliaryStreamFunction()(model, {
+            const stream = await this.agent.streamFunction(model, {
                 systemPrompt: this.systemPrompt,
                 messages: [
                     {
@@ -1202,7 +1049,7 @@ export class AgentSession {
             }, {
                 maxTokens: 32,
                 reasoning: undefined,
-                sessionId: randomUUID(),
+                sessionId: this.sessionId,
                 transport: this.agent.transport,
                 thinkingBudgets: this.agent.thinkingBudgets,
                 maxRetryDelayMs: this.agent.maxRetryDelayMs,
@@ -1835,26 +1682,9 @@ export class AgentSession {
     // Compaction
     // =========================================================================
     async _applyPiServerCompactionResult(result, reason, willRetry) {
-        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);
+        this.sessionManager.replaceTree(result.entries, result.leafId);
         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");
@@ -1873,7 +1703,7 @@ export class AgentSession {
             await this._extensionRunner.emit({
                 type: "session_compact",
                 compactionEntry: savedCompactionEntry,
-                fromExtension: savedCompactionEntry.fromHook === true,
+                fromExtension: false,
                 reason,
                 willRetry,
             });
@@ -1890,31 +1720,38 @@ export class AgentSession {
         await this.abort();
         this._compactionAbortController = new AbortController();
         this._emit({ type: "compaction_start", reason: "manual" });
-        let piServerLease = {
-            path: undefined,
-            acquired: false,
-            owner: undefined,
-        };
         try {
-            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: recoveredCompaction,
-                    aborted: false,
-                    willRetry: false,
-                });
-                return recoveredCompaction;
-            }
             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();
+                this._emit({
+                    type: "compaction_end",
+                    reason: "manual",
+                    result,
+                    aborted: false,
+                    willRetry: false,
+                });
+                return result;
             }
             const pathEntries = this.sessionManager.getBranch();
             const settings = this.settingsManager.getCompactionSettings();
@@ -1947,42 +1784,6 @@ 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() },
-                    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;
@@ -1998,7 +1799,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                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" }));
+                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" }));
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
@@ -2055,7 +1856,6 @@ export class AgentSession {
             throw error;
         }
         finally {
-            this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
             this._compactionAbortController = undefined;
             this._reconnectToAgent();
         }
@@ -2179,33 +1979,45 @@ export class AgentSession {
      */
     async _runAutoCompaction(reason, willRetry, preparationOverride, preparationOptions) {
         const settings = this.settingsManager.getCompactionSettings();
-        const model = this.model;
         let started = false;
-        let piServerLease = {
-            path: undefined,
-            acquired: false,
-            owner: undefined,
-        };
         try {
-            piServerLease = this._acquirePiServerOperationLease(this._isAgentRunActive ? this._piServerSessionLeaseOwner : undefined);
-            if (!this._isAgentRunActive) {
-                await this._recoverPendingPiServerRunBeforeControlOperation();
+            if (!this.model) {
+                return false;
             }
-            const recoveredCompaction = await this._recoverPendingPiServerCompaction(reason, willRetry);
-            if (recoveredCompaction) {
+            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()) {
                 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: recoveredCompaction, aborted: false, willRetry });
+                this._emit({ type: "compaction_end", reason, result, 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) {
@@ -2214,7 +2026,6 @@ 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")) {
@@ -2225,7 +2036,7 @@ export class AgentSession {
                     customInstructions: undefined,
                     reason,
                     willRetry,
-                    signal,
+                    signal: this._autoCompactionAbortController.signal,
                 }));
                 if (extensionResult?.cancel) {
                     this._emit({
@@ -2242,35 +2053,6 @@ 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() },
-                    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;
@@ -2286,7 +2068,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                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 }));
+                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 }));
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
@@ -2357,7 +2139,6 @@ export class AgentSession {
             return false;
         }
         finally {
-            this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
             this._autoCompactionAbortController = undefined;
         }
     }
@@ -2773,12 +2554,8 @@ export class AgentSession {
             return false;
         if (this._isPostAssistantPiServerSyncFailure(message))
             return false;
-        if (isPiServerRestartUncertainty(message))
-            return false;
-        const piServerFailure = getPiServerFailureMetadata(message);
-        if (piServerFailure?.phase && piServerFailure.phase !== "provider_stream")
-            return false;
-        if (piServerFailure?.retryable === false)
+        const piServerPhase = getPiServerFailurePhase(message);
+        if (piServerPhase && piServerPhase !== "provider_stream")
             return false;
         if (message.errorMessage && this._isNonRetryableProviderLimitError(message.errorMessage))
             return false;
@@ -3045,7 +2822,6 @@ export class AgentSession {
             replaceInstructions,
             label,
         };
-        const piServerLease = this._acquirePiServerOperationLease();
         // Set up abort controller for summarization
         this._branchSummaryAbortController = new AbortController();
         try {
@@ -3093,7 +2869,7 @@ export class AgentSession {
                     customInstructions,
                     replaceInstructions,
                     reserveTokens: branchSummarySettings.reserveTokens,
-                    streamFn: this._getAuxiliaryStreamFunction(),
+                    streamFn: this.agent.streamFunction,
                     retry: this.settingsManager.getRetrySettings(),
                     callbacks: this._summarizationRetryCallbacks({ source: "branchSummary" }),
                 });
@@ -3178,7 +2954,6 @@ export class AgentSession {
             return { editorText, cancelled: false, summaryEntry };
         }
         finally {
-            this._releasePiServerOperationLease(piServerLease.path, piServerLease.acquired, piServerLease.owner);
             this._branchSummaryAbortController = undefined;
         }
     }