@pasko70/pibo 1.7.12 → 1.8.1

package/dist/core/routed-session.js

@@ -2,6 +2,8 @@ import { SessionManager, shouldCompact } from "@earendil-works/pi-coding-agent";
 import { getOpenAiCodexProviderUsageForActiveModel } from "../auth/openai-codex-usage.js";
 import { normalizeSessionErrorDetails, runtimeSessionErrorDetails } from "./session-errors.js";
 import { expandInlineSkills } from "./skill-expansion.js";
+import { PIBO_CONTEXT_GUARD_RESUME_MESSAGE_TYPE, PIBO_CONTEXT_GUARD_RESUME_PROMPT, cancelPiboAssistantContextGuardRecovery, claimPiboAssistantContextGuardRecovery, waitForPiboAssistantContextGuardRecovery, } from "./context-guard.js";
+import { PIBO_PROVIDER_RECOVERY_MESSAGE_TYPE, PIBO_PROVIDER_RECOVERY_PROMPT, PiboProviderRecoveryCancelledError, isRetryablePiboAssistantError, isRetryablePiboProviderError, resolvePiboProviderRecoverySettings, waitForPiboProviderRecovery, } from "./provider-recovery.js";
 const FAST_SERVICE_TIER = "priority";
 function modelSupportsFastServiceTier(model) {
     if (!model)
@@ -21,6 +23,9 @@ function withFastServiceTierOption(options) {
 function errorMessage(error) {
     return error instanceof Error ? error.message : String(error);
 }
+function isAssistantMessage(message) {
+    return Boolean(message && typeof message === "object" && message.role === "assistant");
+}
 function numberValue(value) {
     return typeof value === "number" && Number.isFinite(value) ? value : undefined;
 }
@@ -388,6 +393,7 @@ export class RoutedSession {
     onSessionOperation;
     onKillChildren;
     onStateChange;
+    onMessagesInterrupted;
     queue = [];
     processing = false;
     disposed = false;
@@ -398,9 +404,15 @@ export class RoutedSession {
     nextAssistantIndex = 0;
     activeThinkingIndex;
     nextThinkingIndex = 0;
+    pendingAssistantError;
+    pendingAssistantErrorRetryable = false;
+    activeMessageFailed = false;
+    providerRecoveryCancelled = false;
+    providerRecoveryAbortController;
     unsubscribe;
+    recoverySession;
     isContinuePatched = false;
-    constructor(piboSessionId, runtime, emit, pluginRegistry, forwardPiEvents, onPiEventTelemetry, initialFastMode, onSessionOperation, onKillChildren, onStateChange) {
+    constructor(piboSessionId, runtime, emit, pluginRegistry, forwardPiEvents, onPiEventTelemetry, initialFastMode, onSessionOperation, onKillChildren, onStateChange, onMessagesInterrupted) {
         this.piboSessionId = piboSessionId;
         this.runtime = runtime;
         this.emit = emit;
@@ -410,6 +422,7 @@ export class RoutedSession {
         this.onSessionOperation = onSessionOperation;
         this.onKillChildren = onKillChildren;
         this.onStateChange = onStateChange;
+        this.onMessagesInterrupted = onMessagesInterrupted;
         this.fastMode = initialFastMode && this.fastModeSupported();
         this.bindRuntimeSession();
         this.patchFastModeProviderRequest();
@@ -485,12 +498,32 @@ export class RoutedSession {
     }
     bindRuntimeSession() {
         this.unsubscribe?.();
-        this.unsubscribe = this.runtime.session.subscribe((event) => {
+        const session = this.runtime.session;
+        if (this.recoverySession && this.recoverySession !== session) {
+            this.cancelProviderRecovery();
+            cancelPiboAssistantContextGuardRecovery(this.recoverySession, new Error("Context guard recovery cancelled because the Pi session changed"));
+        }
+        this.recoverySession = session;
+        claimPiboAssistantContextGuardRecovery(session);
+        this.unsubscribe = session.subscribe((event) => {
             this.onPiEventTelemetry?.(this.piboSessionId, event, { status: this.getStatus(), activeEventId: this.activeMessage?.id });
             const model = this.runtime.session.model;
             const normalized = normalizePiEvent(this.piboSessionId, event, { contextWindow: numberValue(model?.contextWindow) });
-            if (normalized) {
-                this.emit(this.withActiveMessage(normalized));
+            const candidate = event && typeof event === "object" ? event : undefined;
+            const assistantMessageEnded = candidate?.type === "message_end" && isAssistantMessage(candidate.message);
+            // Pi gets the first chance to recover through its short retry/compaction loop.
+            // Keep the final error pending so the routed turn can continue durable recovery.
+            if (assistantMessageEnded && normalized?.type === "session_error") {
+                this.pendingAssistantError = this.withActiveMessage(normalized);
+                this.pendingAssistantErrorRetryable = isRetryablePiboAssistantError(candidate.message);
+            }
+            else {
+                if (assistantMessageEnded) {
+                    this.pendingAssistantError = undefined;
+                    this.pendingAssistantErrorRetryable = false;
+                }
+                if (normalized)
+                    this.emit(this.withActiveMessage(normalized));
             }
             if (this.forwardPiEvents) {
                 this.emit({ type: "pi_event", piboSessionId: this.piboSessionId, event });
@@ -498,6 +531,87 @@ export class RoutedSession {
             this.handleCompactionEvent(event);
         });
     }
+    flushPendingAssistantError() {
+        if (!this.pendingAssistantError)
+            return;
+        this.activeMessageFailed = true;
+        this.emit(this.pendingAssistantError);
+        this.pendingAssistantError = undefined;
+        this.pendingAssistantErrorRetryable = false;
+    }
+    cancelProviderRecovery() {
+        this.providerRecoveryCancelled = true;
+        this.providerRecoveryAbortController?.abort();
+        this.providerRecoveryAbortController = undefined;
+    }
+    async resumeContextGuardRecovery(session) {
+        while (await waitForPiboAssistantContextGuardRecovery(session)) {
+            try {
+                await session.sendCustomMessage({
+                    customType: PIBO_CONTEXT_GUARD_RESUME_MESSAGE_TYPE,
+                    content: [{ type: "text", text: PIBO_CONTEXT_GUARD_RESUME_PROMPT }],
+                    display: false,
+                }, { triggerTurn: true });
+            }
+            catch (error) {
+                const resumeError = error instanceof Error ? error : new Error(String(error));
+                cancelPiboAssistantContextGuardRecovery(session, resumeError);
+                throw resumeError;
+            }
+        }
+    }
+    async recoverTransientProviderErrors(session) {
+        let attempt = 0;
+        while (this.pendingAssistantError && this.pendingAssistantErrorRetryable) {
+            if (this.providerRecoveryCancelled)
+                throw new PiboProviderRecoveryCancelledError();
+            const settings = resolvePiboProviderRecoverySettings(session.settingsManager);
+            if (!settings.enabled)
+                return;
+            attempt += 1;
+            this.pendingAssistantError = undefined;
+            this.pendingAssistantErrorRetryable = false;
+            const controller = new AbortController();
+            this.providerRecoveryAbortController = controller;
+            try {
+                await waitForPiboProviderRecovery(attempt, settings, controller.signal);
+            }
+            finally {
+                if (this.providerRecoveryAbortController === controller) {
+                    this.providerRecoveryAbortController = undefined;
+                }
+            }
+            if (this.providerRecoveryCancelled || this.disposed || this.runtime.session !== session || !this.activeMessage) {
+                throw new PiboProviderRecoveryCancelledError();
+            }
+            try {
+                await session.sendCustomMessage({
+                    customType: PIBO_PROVIDER_RECOVERY_MESSAGE_TYPE,
+                    content: [{ type: "text", text: PIBO_PROVIDER_RECOVERY_PROMPT }],
+                    display: false,
+                    details: { attempt },
+                }, { triggerTurn: true });
+                await this.resumeContextGuardRecovery(session);
+                if (this.providerRecoveryCancelled)
+                    throw new PiboProviderRecoveryCancelledError();
+            }
+            catch (error) {
+                if (error instanceof PiboProviderRecoveryCancelledError)
+                    throw error;
+                if (!isRetryablePiboProviderError(error))
+                    throw error;
+                const message = errorMessage(error);
+                this.pendingAssistantError = {
+                    type: "session_error",
+                    piboSessionId: this.piboSessionId,
+                    eventId: this.activeMessage.id,
+                    error: message,
+                    errorDetails: runtimeSessionErrorDetails(message),
+                };
+                this.pendingAssistantErrorRetryable = true;
+            }
+        }
+    }
     handleCompactionEvent(event) {
         if (!event || typeof event !== "object")
             return;
@@ -602,15 +716,17 @@ export class RoutedSession {
     }
     removeQueuedMessages(predicate) {
         this.assertActive();
-        let removed = 0;
+        const removedMessages = [];
         for (let index = this.queue.length - 1; index >= 0; index -= 1) {
             const item = this.queue[index];
             if (item.kind !== "message" || !predicate(item.event))
                 continue;
             this.queue.splice(index, 1);
-            removed += 1;
+            removedMessages.push(item.event);
         }
-        return removed;
+        removedMessages.reverse();
+        this.notifyMessagesInterrupted(removedMessages, "queued message removed");
+        return removedMessages.length;
     }
     getCurrentSession() {
         return this.createSessionSnapshot();
@@ -733,16 +849,25 @@ export class RoutedSession {
     async dispose() {
         if (this.disposed)
             return;
+        this.notifyMessagesInterrupted(this.activeAndQueuedMessages(), "session disposed");
+        this.cancelProviderRecovery();
         this.queue.length = 0;
         this.onStateChange?.({ processing: this.processing, queuedMessages: this.queue.length, disposed: true });
         this.unsubscribe?.();
         this.unsubscribe = undefined;
+        if (this.recoverySession) {
+            this.cancelContextGuardRecovery("Context guard recovery cancelled because the routed session was disposed");
+            this.recoverySession = undefined;
+        }
         this.disposed = true;
         await this.runtime.dispose();
     }
     async kill() {
+        this.notifyMessagesInterrupted(this.activeAndQueuedMessages(), "session killed");
+        this.cancelProviderRecovery();
         this.queue.length = 0;
         this.onStateChange?.({ processing: this.processing, queuedMessages: this.queue.length, disposed: this.disposed });
+        this.cancelContextGuardRecovery("Context guard recovery cancelled because the routed session was killed");
         await this.runtime.session.abort();
         return this.piboSessionId;
     }
@@ -750,11 +875,16 @@ export class RoutedSession {
         this.assertActive();
         const queuedIndex = this.queue.findIndex((item) => item.event.id === eventId);
         if (queuedIndex >= 0) {
-            this.queue.splice(queuedIndex, 1);
+            const [removed] = this.queue.splice(queuedIndex, 1);
+            if (removed?.kind === "message")
+                this.notifyMessagesInterrupted([removed.event], "message cancelled");
             this.onStateChange?.({ processing: this.processing, queuedMessages: this.queue.length, disposed: this.disposed });
             return true;
         }
         if (this.activeMessage?.id === eventId) {
+            this.notifyMessagesInterrupted([this.activeMessage], "message cancelled");
+            this.cancelProviderRecovery();
+            this.cancelContextGuardRecovery("Context guard recovery cancelled with the active message");
             await this.runtime.session.abort();
             return true;
         }
@@ -792,20 +922,32 @@ export class RoutedSession {
         });
         try {
             this.activeMessage = event;
+            this.providerRecoveryCancelled = false;
+            this.pendingAssistantError = undefined;
+            this.pendingAssistantErrorRetryable = false;
+            this.activeMessageFailed = false;
             this.activeAssistantIndex = undefined;
             this.nextAssistantIndex = 0;
             this.activeThinkingIndex = undefined;
             this.nextThinkingIndex = 0;
-            const expandedText = expandInlineSkills(event.text, this.runtime.session.resourceLoader.getSkills().skills);
-            await this.runtime.session.prompt(expandedText, { source: promptSource(event.source) });
-            this.emit({
-                type: "message_finished",
-                piboSessionId: this.piboSessionId,
-                eventId: event.id,
-                source: event.source,
-            });
+            const session = this.runtime.session;
+            const expandedText = expandInlineSkills(event.text, session.resourceLoader.getSkills().skills);
+            await session.prompt(expandedText, { source: promptSource(event.source) });
+            await this.resumeContextGuardRecovery(session);
+            await this.recoverTransientProviderErrors(session);
+            this.flushPendingAssistantError();
+            if (!this.activeMessageFailed) {
+                this.emit({
+                    type: "message_finished",
+                    piboSessionId: this.piboSessionId,
+                    eventId: event.id,
+                    source: event.source,
+                });
+            }
         }
         catch (error) {
+            if (error instanceof PiboProviderRecoveryCancelledError)
+                return;
             const message = errorMessage(error);
             this.emit({
                 type: "session_error",
@@ -817,6 +959,10 @@ export class RoutedSession {
         }
         finally {
             this.activeMessage = undefined;
+            this.providerRecoveryCancelled = false;
+            this.pendingAssistantError = undefined;
+            this.pendingAssistantErrorRetryable = false;
+            this.activeMessageFailed = false;
             this.activeAssistantIndex = undefined;
             this.nextAssistantIndex = 0;
             this.activeThinkingIndex = undefined;
@@ -873,6 +1019,9 @@ export class RoutedSession {
             getProviderUsage: () => this.getProviderUsage(),
             clearQueue: () => this.clearQueue(),
             abort: async () => {
+                if (this.activeMessage)
+                    this.notifyMessagesInterrupted([this.activeMessage], "abort requested");
+                this.cancelContextGuardRecovery("Context guard recovery cancelled by abort");
                 await this.runtime.session.abort();
             },
             dispose: () => this.dispose(),
@@ -913,6 +1062,11 @@ export class RoutedSession {
             },
         }, event);
     }
+    cancelContextGuardRecovery(message) {
+        const session = this.recoverySession ?? this.runtime.session;
+        cancelPiboAssistantContextGuardRecovery(session, new Error(message));
+        session.abortCompaction?.();
+    }
     assertActive() {
         if (this.disposed) {
             throw new Error(`Session "${this.piboSessionId}" has been disposed`);
@@ -934,10 +1088,21 @@ export class RoutedSession {
     }
     clearQueue() {
         const cleared = this.queue.length;
+        const removedMessages = this.queue.flatMap((item) => item.kind === "message" ? [item.event] : []);
         this.queue.length = 0;
+        this.notifyMessagesInterrupted(removedMessages, "queue cleared");
         this.onStateChange?.({ processing: this.processing, queuedMessages: this.queue.length, disposed: this.disposed });
         return cleared;
     }
+    activeAndQueuedMessages() {
+        const messages = this.queue.flatMap((item) => item.kind === "message" ? [item.event] : []);
+        return this.activeMessage ? [this.activeMessage, ...messages] : messages;
+    }
+    notifyMessagesInterrupted(messages, reason) {
+        if (messages.length === 0)
+            return;
+        this.onMessagesInterrupted?.(messages, reason);
+    }
     createSessionSnapshot() {
         const session = this.runtime.session;
         const manager = session.sessionManager;