@pasko70/pibo 1.7.12 → 1.8.1

package/dist/core/runtime.js

@@ -14,7 +14,7 @@ import { getMcpAgentContextFile } from "../mcp/agent-context.js";
 import { createPiboSystemPromptTemplateExtension } from "./system-prompt-template.js";
 import { getActivePiboBasePromptPath } from "./base-prompt.js";
 import { createPiboCompactionPromptExtension } from "./compaction-prompt.js";
-import { createPiboAssistantContextGuardExtension } from "./context-guard.js";
+import { cancelPiboAssistantContextGuardRecovery, createPiboAssistantContextGuardExtension, createPiboAssistantContextGuardRecovery, isPiboAssistantContextGuardRecoveryPending, registerPiboAssistantContextGuardRecovery, } from "./context-guard.js";
 import { getPiPackageRuntimeOptions } from "../pi-packages/runtime.js";
 import { getDefaultPiboWorkspace } from "./workspace.js";
 import { DEFAULT_USER_TIMEZONE } from "./user-settings.js";
@@ -166,10 +166,10 @@ function getBuiltinToolAllowlist(profile, customTools) {
         return undefined;
     return [...selectedBuiltinTools, ...customTools.map((tool) => tool.name)];
 }
-function getProfileExtensionFactories(profile, extensionFactories) {
+function getProfileExtensionFactories(profile, extensionFactories, contextGuardRecovery, getSettingsManager) {
     const piboPromptTemplateExtension = createPiboSystemPromptTemplateExtension();
-    const piboCompactionPromptExtension = createPiboCompactionPromptExtension();
-    const piboContextGuardExtension = createPiboAssistantContextGuardExtension();
+    const piboCompactionPromptExtension = createPiboCompactionPromptExtension({ getSettingsManager });
+    const piboContextGuardExtension = createPiboAssistantContextGuardExtension({}, contextGuardRecovery);
     const providerToolExtensions = profile.tools
         .filter((tool) => tool.enabled !== false)
         .filter(isWebSearchProviderTool)
@@ -234,12 +234,14 @@ export async function createPiboRuntime(options = {}) {
     const sessionManager = await createSessionManager(cwd, profile, options.persistSession !== false);
     const authStorage = AuthStorage.create();
     const createRuntime = async ({ cwd: runtimeCwd, agentDir: runtimeAgentDir, sessionManager: runtimeSessionManager, sessionStartEvent, }) => {
+        const contextGuardRecovery = createPiboAssistantContextGuardRecovery();
         const contextFiles = await loadContextFiles(runtimeCwd, profile.contextFiles);
         const sessionContextFile = createSessionContextFile({ piboSessionId: profile.sessionId, ...options.sessionContext });
         const installedToolContextFile = getInstalledCliToolContextFile();
         const mcpAgentContextFile = await getMcpAgentContextFile(profile.mcpServers);
         const skillPaths = getEnabledSkillPaths(runtimeCwd, profile);
         const piPackageOptions = getPiPackageRuntimeOptions(runtimeCwd, profile);
+        let runtimeSettingsManager;
         const services = await createAgentSessionServices({
             cwd: runtimeCwd,
             agentDir: runtimeAgentDir,
@@ -247,7 +249,7 @@ export async function createPiboRuntime(options = {}) {
             resourceLoaderOptions: {
                 ...piPackageOptions.resourceLoaderOptions,
                 additionalSkillPaths: skillPaths,
-                extensionFactories: getProfileExtensionFactories(profile, options.extensionFactories),
+                extensionFactories: getProfileExtensionFactories(profile, options.extensionFactories, contextGuardRecovery, () => runtimeSettingsManager),
                 noExtensions: true,
                 noSkills: true,
                 noPromptTemplates: true,
@@ -264,6 +266,7 @@ export async function createPiboRuntime(options = {}) {
                 }),
             },
         });
+        runtimeSettingsManager = services.settingsManager;
         applyPiboRuntimeRetryDefaults(services.settingsManager, options.retryDefaults);
         registerOpenAiGpt56Models(services.modelRegistry);
         registerMiniMaxProvider(services.modelRegistry);
@@ -293,6 +296,10 @@ export async function createPiboRuntime(options = {}) {
             tools: getBuiltinToolAllowlist(profile, customTools),
         });
         installValidationOutputCompaction(created.session.agent);
+        registerPiboAssistantContextGuardRecovery(created.session, contextGuardRecovery);
+        if (options.contextGuardTuiQueueOrdering === true) {
+            installPiboContextGuardTuiQueueOrdering(created.session);
+        }
         const resourceLoader = services.resourceLoader;
         const diagnostics = [
             ...piPackageOptions.diagnostics,
@@ -303,13 +310,14 @@ export async function createPiboRuntime(options = {}) {
                 message: `Failed to load extension "${path}": ${error}`,
             })),
         ];
-        if (localRuntimeRegistry) {
-            const originalDispose = created.session.dispose.bind(created.session);
-            created.session.dispose = () => {
+        const originalDispose = created.session.dispose.bind(created.session);
+        created.session.dispose = () => {
+            cancelPiboAssistantContextGuardRecovery(created.session, new Error("Context guard recovery cancelled because the Pi session was disposed"));
+            if (localRuntimeRegistry) {
                 void localRuntimeRegistry.closeControllerSessions(profile.sessionId ?? "local", { force: true });
-                originalDispose();
-            };
-        }
+            }
+            originalDispose();
+        };
         return {
             ...created,
             services,
@@ -456,6 +464,39 @@ export async function inspectPiboProfile(options = {}) {
         await runtime.dispose();
     }
 }
+function installPiboContextGuardTuiQueueOrdering(session) {
+    const originalSubscribe = session.subscribe.bind(session);
+    const originalPrompt = session.prompt.bind(session);
+    const originalSteer = session.steer.bind(session);
+    session.subscribe = ((listener) => originalSubscribe((event) => {
+        if (event.type === "compaction_end"
+            && event.result
+            && isPiboAssistantContextGuardRecoveryPending(session)) {
+            listener({ ...event, willRetry: true });
+            return;
+        }
+        listener(event);
+    }));
+    session.prompt = async (text, options) => {
+        if (isPiboAssistantContextGuardRecoveryPending(session)) {
+            if (!session.isStreaming) {
+                await session.followUp(text, options?.images);
+                options?.preflightResult?.(true);
+                return;
+            }
+            await originalPrompt(text, { ...options, streamingBehavior: "followUp" });
+            return;
+        }
+        await originalPrompt(text, options);
+    };
+    session.steer = async (text, images) => {
+        if (isPiboAssistantContextGuardRecoveryPending(session)) {
+            await session.followUp(text, images);
+            return;
+        }
+        await originalSteer(text, images);
+    };
+}
 export async function runPiboTui(options = {}) {
     const profile = options.profile ?? createDefaultPiboProfile();
     const hasEnabledSubagents = profile.subagents.some((subagent) => subagent.enabled !== false);
@@ -465,7 +506,7 @@ export async function runPiboTui(options = {}) {
         process.exitCode = 1;
         return;
     }
-    const runtime = await createPiboRuntime({ ...options, profile });
+    const runtime = await createPiboRuntime({ ...options, profile, contextGuardTuiQueueOrdering: true });
     try {
         const fatal = runtime.diagnostics.find((diagnostic) => diagnostic.type === "error");
         for (const diagnostic of runtime.diagnostics) {

package/dist/core/session-errors.js

@@ -10,6 +10,11 @@ const PROVIDER_NETWORK_ERROR_MARKERS = [
     "reset before headers",
     "socket hang up",
     "socket connection was closed",
+    "eai_again",
+    "enotfound",
+    "econnreset",
+    "econnrefused",
+    "etimedout",
 ];
 export function classifySessionErrorMessage(message, options = {}) {
     const normalized = message.toLowerCase();
@@ -20,7 +25,15 @@ export function classifySessionErrorMessage(message, options = {}) {
         return { category: "provider_transport", errorClass: "provider_transport", code: "websocket_error", origin: "provider", retryable: true, userMessage: "The provider WebSocket connection failed." };
     }
     if (normalized.includes("request was aborted") || normalized.includes("aborted")) {
-        return { category: "runtime_abort", errorClass: "runtime_abort", code: "request_aborted", origin: "runtime", retryable: true, userMessage: "The active model request was aborted." };
+        return { category: "runtime_abort", errorClass: "runtime_abort", code: "request_aborted", origin: "runtime", retryable: false, userMessage: "The active model request was aborted." };
+    }
+    if (normalized.includes("insufficient_quota")
+        || normalized.includes("quota exceeded")
+        || normalized.includes("out of budget")
+        || normalized.includes("billing")
+        || normalized.includes("usage limit")
+        || normalized.includes("available balance")) {
+        return { category: "quota_exhausted", errorClass: "provider_rate_limit", code: "quota_exhausted", origin: "provider", retryable: false, userMessage: "The provider quota or billing limit was reached." };
     }
     if (normalized.includes("rate limit") || normalized.includes("429")) {
         return { category: "rate_limit", errorClass: "provider_rate_limit", code: "rate_limited", origin: "provider", retryable: true, userMessage: "The provider rate limit was reached." };

package/dist/core/session-router.js

@@ -20,6 +20,7 @@ import { withWorkflowSessionKind } from "../sessions/workflow-session-kind.js";
 import { PiboRuntimeTelemetryRecorder } from "./runtime-telemetry.js";
 import { createPiboProviderTelemetryExtension } from "./provider-telemetry.js";
 const DEFAULT_SUBAGENT_REPLY_TIMEOUT_MS = 10 * 60 * 1000;
+const DEFAULT_ROUTED_SESSION_IDLE_TIMEOUT_MS = 30 * 60 * 1000;
 export const RALPH_RUNTIME_RETRY_DEFAULTS = {
     enabled: true,
     maxRetries: 7,
@@ -132,6 +133,8 @@ export class PiboSessionRouter {
     signalRegistry;
     runtimeRegistry;
     scheduledRunReminders = new Map();
+    idleSessionTimers = new Map();
+    routedSessionIdleTimeoutMs;
     baseProfile;
     pluginRegistry;
     sessionStore;
@@ -146,6 +149,12 @@ export class PiboSessionRouter {
         this.telemetryRecorder = this.telemetryStore
             ? new PiboRuntimeTelemetryRecorder(this.telemetryStore, undefined, { providerEventMode: providerEventTelemetryModeFromEnv() })
             : undefined;
+        const idleTimeoutMs = options.routedSessionIdleTimeoutMs;
+        this.routedSessionIdleTimeoutMs = idleTimeoutMs === false
+            ? false
+            : typeof idleTimeoutMs === "number" && Number.isFinite(idleTimeoutMs) && idleTimeoutMs > 0
+                ? idleTimeoutMs
+                : DEFAULT_ROUTED_SESSION_IDLE_TIMEOUT_MS;
         const defaultProfileName = selectDefaultPiboProfileName(this.pluginRegistry);
         this.baseProfile = options.profile ?? createPiboProfileFromRegistryOrDefault(this.pluginRegistry, defaultProfileName);
         this.reliabilityStore = options.reliabilityStore ?? (options.persistSession === false ? undefined : createDefaultPiboReliabilityStore());
@@ -165,20 +174,29 @@ export class PiboSessionRouter {
     }
     async emit(event) {
         const session = await this.getOrCreateSession(event.piboSessionId);
-        if (event.type === "message") {
-            return session.enqueueMessage(event);
-        }
-        const output = await session.executeAction(event);
-        if (event.action === "abort") {
-            this.signalRegistry.project({ type: "session_interrupted", piboSessionId: event.piboSessionId, reason: "abort action" });
-        }
-        if (event.action === "dispose") {
-            this.disposeSignalSubtree(event.piboSessionId, "dispose action");
+        this.clearIdleSessionTimer(event.piboSessionId);
+        try {
+            if (event.type === "message") {
+                return session.enqueueMessage(event);
+            }
+            const output = await session.executeAction(event);
+            if (event.action === "abort") {
+                this.signalRegistry.project({ type: "session_interrupted", piboSessionId: event.piboSessionId, reason: "abort action" });
+            }
+            if (event.action === "dispose") {
+                await this.disposeSessionSubtree(event.piboSessionId, "dispose action", { cancelRuns: true });
+            }
+            else if (event.action === "kill" || event.action === "kill_all") {
+                await this.disposeSessionSubtree(event.piboSessionId, `${event.action} action`, { cancelRuns: event.action === "kill_all" });
+            }
+            else if (shouldResetSessionAfterAction(event.action)) {
+                await this.resetCachedSession(event.piboSessionId, "provider auth changed");
+            }
+            return output;
         }
-        else if (shouldResetSessionAfterAction(event.action)) {
-            await this.resetCachedSession(event.piboSessionId, "provider auth changed");
+        finally {
+            this.scheduleIdleSessionEvictionIfIdle(event.piboSessionId);
         }
-        return output;
     }
     async killSession(piboSessionId, options) {
         const killed = [];
@@ -190,23 +208,32 @@ export class PiboSessionRouter {
                 const runs = this.runRegistry.cancelControllerRuns(piboSessionId);
                 cancelledRuns.push(...runs.map((r) => r.runId));
             }
-            await this.runtimeRegistry.closeControllerSessions(piboSessionId, { force: true });
             this.signalRegistry.project({ type: "session_interrupted", piboSessionId, reason: "kill" });
             const children = await this.killChildSessions(piboSessionId, options);
             killed.push(...children.killed);
             cancelledRuns.push(...children.cancelledRuns);
+            await this.disposeSessionSubtree(piboSessionId, "kill", { cancelRuns: false });
         }
         return { killed, cancelledRuns };
     }
-    disposeSignalSubtree(piboSessionId, reason) {
-        const descendants = this.descendantSessionIds(piboSessionId);
-        for (const id of [piboSessionId, ...descendants]) {
-            this.runRegistry.cancelControllerRuns(id);
-            void this.runtimeRegistry.closeControllerSessions(id, { force: true });
-            this.signalRegistry.project({ type: "session_disposed", piboSessionId: id, reason });
+    async disposeSessionSubtree(piboSessionId, reason, options) {
+        const ids = [piboSessionId, ...this.descendantSessionIds(piboSessionId)];
+        const sessions = [];
+        for (const id of ids) {
+            if (options.cancelRuns)
+                this.runRegistry.cancelControllerRuns(id);
+            this.clearIdleSessionTimer(id);
             this.scheduledRunReminders.delete(id);
+            const cached = this.sessions.get(id);
+            if (cached)
+                sessions.push(cached);
             this.sessions.delete(id);
         }
+        await Promise.all(ids.map((id) => this.runtimeRegistry.closeControllerSessions(id, { force: true })));
+        await Promise.all(sessions.map((session) => session.dispose()));
+        for (const id of ids) {
+            this.signalRegistry.project({ type: "session_disposed", piboSessionId: id, reason });
+        }
     }
     descendantSessionIds(parentId) {
         const output = [];
@@ -292,6 +319,8 @@ export class PiboSessionRouter {
         const eventWithId = { ...event, id: event.id ?? randomUUID() };
         return await new Promise((resolve, reject) => {
             let settled = false;
+            let lastAssistantMessage;
+            let timeout;
             const unsubscribe = this.subscribe((output) => {
                 if (output.piboSessionId !== eventWithId.piboSessionId ||
                     !("eventId" in output) ||
@@ -299,20 +328,21 @@ export class PiboSessionRouter {
                     return;
                 }
                 if (output.type === "assistant_message") {
-                    finish(output);
+                    lastAssistantMessage = output;
+                }
+                else if (output.type === "message_finished") {
+                    finish(lastAssistantMessage ?? new Error(`Pibo session "${eventWithId.piboSessionId}" finished without an assistant reply`));
                 }
                 else if (output.type === "session_error") {
                     finish(new Error(output.error));
                 }
             });
-            const timeout = setTimeout(() => {
-                finish(new Error(`Timed out waiting for assistant reply from Pibo session "${eventWithId.piboSessionId}"`));
-            }, timeoutMs);
             const finish = (result) => {
                 if (settled)
                     return;
                 settled = true;
-                clearTimeout(timeout);
+                if (timeout)
+                    clearTimeout(timeout);
                 unsubscribe();
                 if (result instanceof Error) {
                     reject(result);
@@ -321,12 +351,29 @@ export class PiboSessionRouter {
                     resolve(result);
                 }
             };
+            timeout = setTimeout(() => {
+                if (settled)
+                    return;
+                settled = true;
+                unsubscribe();
+                const timeoutError = new Error(`Timed out waiting for assistant reply from Pibo session "${eventWithId.piboSessionId}"`);
+                reject(timeoutError);
+                void this.emit({
+                    type: "execution",
+                    piboSessionId: eventWithId.piboSessionId,
+                    action: "abort",
+                    id: randomUUID(),
+                }).catch(() => { });
+            }, timeoutMs);
             this.emit(eventWithId).catch(finish);
         });
     }
     async disposeAll() {
         const sessions = [...this.sessions.values()];
         this.sessions.clear();
+        for (const timer of this.idleSessionTimers.values())
+            clearTimeout(timer);
+        this.idleSessionTimers.clear();
         this.runRegistry.cancelAll("Pibo session router was disposed.");
         for (const session of sessions)
             this.signalRegistry.project({ type: "session_disposed", piboSessionId: session.getStatus().piboSessionId, reason: "router disposed" });
@@ -334,10 +381,54 @@ export class PiboSessionRouter {
         await this.runtimeRegistry.closeAll({ force: true });
         await Promise.all(sessions.map((session) => session.dispose()));
     }
+    clearIdleSessionTimer(piboSessionId) {
+        const timer = this.idleSessionTimers.get(piboSessionId);
+        if (timer)
+            clearTimeout(timer);
+        this.idleSessionTimers.delete(piboSessionId);
+    }
+    scheduleIdleSessionEvictionIfIdle(piboSessionId) {
+        if (this.routedSessionIdleTimeoutMs === false)
+            return;
+        const session = this.sessions.get(piboSessionId);
+        if (!session)
+            return;
+        const status = session.getStatus();
+        if (status.disposed || status.processing || status.streaming || status.queuedMessages > 0) {
+            this.clearIdleSessionTimer(piboSessionId);
+            return;
+        }
+        this.clearIdleSessionTimer(piboSessionId);
+        const timer = setTimeout(() => {
+            this.idleSessionTimers.delete(piboSessionId);
+            void this.evictIdleSession(piboSessionId, session).catch((error) => {
+                const message = error instanceof Error ? error.message : String(error);
+                this.emitOutput({
+                    type: "session_error",
+                    piboSessionId,
+                    error: `Failed to dispose idle routed runtime: ${message}`,
+                    errorDetails: runtimeSessionErrorDetails(message),
+                });
+            });
+        }, this.routedSessionIdleTimeoutMs);
+        timer.unref();
+        this.idleSessionTimers.set(piboSessionId, timer);
+    }
+    async evictIdleSession(piboSessionId, expected) {
+        const current = this.sessions.get(piboSessionId);
+        if (current !== expected)
+            return;
+        const status = current.getStatus();
+        if (status.disposed || status.processing || status.streaming || status.queuedMessages > 0)
+            return;
+        await this.resetCachedSession(piboSessionId, "routed runtime idle timeout");
+    }
     async getOrCreateSession(piboSessionId) {
         const existing = this.sessions.get(piboSessionId);
-        if (existing)
+        if (existing) {
+            this.clearIdleSessionTimer(piboSessionId);
             return existing;
+        }
         const pending = this.pendingSessions.get(piboSessionId);
         if (pending)
             return pending;
@@ -388,7 +479,16 @@ export class PiboSessionRouter {
         const initialFastMode = resolvePiboSessionInitialFastMode(piboSession) ?? selectRequestedFastMode(profileForSession(profile, piboSession.piSessionId, parentPiSessionId), modelDefaults) ?? false;
         const session = new RoutedSession(piboSession.id, runtime, this.emitOutput, this.pluginRegistry, this.options.forwardPiEvents ?? false, this.telemetryRecorder
             ? (id, event, context) => this.telemetryRecorder?.recordPiEvent(id, event, { session: this.sessionStore.get(id), status: context.status, activeEventId: context.activeEventId })
-            : undefined, initialFastMode, (result, event) => this.handleSessionOperation(result, event), (id, opts) => this.killChildSessions(id, opts), (state) => this.signalRegistry.project({ type: "session_processing_changed", piboSessionId: piboSession.id, processing: state.processing, queuedMessages: state.queuedMessages }));
+            : undefined, initialFastMode, (result, event) => this.handleSessionOperation(result, event), (id, opts) => this.killChildSessions(id, opts), (state) => {
+            this.signalRegistry.project({ type: "session_processing_changed", piboSessionId: piboSession.id, processing: state.processing, queuedMessages: state.queuedMessages });
+            if (state.disposed || state.processing || state.queuedMessages > 0)
+                this.clearIdleSessionTimer(piboSession.id);
+            else
+                this.scheduleIdleSessionEvictionIfIdle(piboSession.id);
+        }, (messages, reason) => this.telemetryRecorder?.recordMessagesInterrupted(messages, {
+            session: this.sessionStore.get(piboSession.id),
+            status: this.sessions.get(piboSession.id)?.getStatus(),
+        }, reason));
         this.sessions.set(piboSession.id, session);
         return session;
     }
@@ -447,6 +547,7 @@ export class PiboSessionRouter {
     }
     async resetCachedSession(piboSessionId, reason) {
         const cached = this.sessions.get(piboSessionId);
+        this.clearIdleSessionTimer(piboSessionId);
         this.sessions.delete(piboSessionId);
         await this.runtimeRegistry.closeControllerSessions(piboSessionId, { force: true });
         await cached?.dispose();

package/dist/data/telemetry.js

@@ -17,6 +17,61 @@ export class TelemetryStore {
     getTurnTimeline(turnIdOrEventId, input = {}) {
         return getTelemetryTurnTimeline(this.db, turnIdOrEventId, input);
     }
+    getOpenPhaseForTurn(turnId, name) {
+        const row = this.db.prepare(`
+			SELECT * FROM telemetry_phases
+			WHERE turn_id = ? AND name = ? AND status = 'open'
+			ORDER BY COALESCE(last_progress_at, started_at) DESC, created_at DESC
+			LIMIT 1
+		`).get(turnId, name);
+        return row ? phaseFromRow(row) : undefined;
+    }
+    countPhasesForTurn(turnId, name) {
+        const row = this.db.prepare("SELECT COUNT(*) AS count FROM telemetry_phases WHERE turn_id = ? AND name = ?").get(turnId, name);
+        return Number(row.count);
+    }
+    listOpenPhasesForTurn(turnId) {
+        const rows = this.db.prepare(`
+			SELECT * FROM telemetry_phases
+			WHERE turn_id = ? AND status = 'open'
+			ORDER BY started_at ASC, created_at ASC
+		`).all(turnId);
+        return rows.map(phaseFromRow);
+    }
+    getLatestProviderRequestForTurn(turnId) {
+        const row = this.db.prepare(`
+			SELECT * FROM telemetry_provider_requests
+			WHERE turn_id = ?
+			ORDER BY started_at DESC, created_at DESC
+			LIMIT 1
+		`).get(turnId);
+        return row ? providerRequestFromRow(row) : undefined;
+    }
+    getActiveProviderRequestForTurn(turnId) {
+        const row = this.db.prepare(`
+			SELECT * FROM telemetry_provider_requests
+			WHERE turn_id = ? AND status NOT IN ('completed', 'error', 'aborted', 'timeout')
+			ORDER BY started_at DESC, created_at DESC
+			LIMIT 1
+		`).get(turnId);
+        return row ? providerRequestFromRow(row) : undefined;
+    }
+    listActiveProviderRequestsForTurn(turnId) {
+        const rows = this.db.prepare(`
+			SELECT * FROM telemetry_provider_requests
+			WHERE turn_id = ? AND status NOT IN ('completed', 'error', 'aborted', 'timeout')
+			ORDER BY started_at ASC, created_at ASC
+		`).all(turnId);
+        return rows.map(providerRequestFromRow);
+    }
+    listActiveToolCallsForTurn(turnId) {
+        const rows = this.db.prepare(`
+			SELECT * FROM telemetry_tool_calls
+			WHERE turn_id = ? AND status NOT IN ('ok', 'error', 'aborted', 'timeout')
+			ORDER BY created_at ASC
+		`).all(turnId);
+        return rows.map(toolCallFromRow);
+    }
     listProviderEventsPage(providerRequestId, input = {}) {
         return listTelemetryProviderEventsPage(this.db, providerRequestId, input);
     }
@@ -176,6 +231,53 @@ export class TelemetryStore {
         const normalizedDelta = input.normalizedEventDelta ?? (input.normalizedType ? 1 : 0);
         this.incrementProviderCounters(input.providerRequestId, input.eventType, receivedAt, byteSize, parseStatus, normalizedDelta);
     }
+    recordProviderProgress(input) {
+        const existing = this.getProviderRequest(input.providerRequestId);
+        if (!existing)
+            return undefined;
+        const eventTypeCounts = { ...existing.eventTypeCounts };
+        for (const [eventType, delta] of Object.entries(input.eventTypeCounts ?? {})) {
+            if (!Number.isFinite(delta) || delta <= 0)
+                continue;
+            const current = typeof eventTypeCounts[eventType] === "number" ? eventTypeCounts[eventType] : 0;
+            eventTypeCounts[eventType] = current + delta;
+        }
+        const rawEventCount = Math.max(0, input.rawEventCount ?? 0);
+        const normalizedEventCount = Math.max(0, input.normalizedEventCount ?? 0);
+        const parseErrorCount = Math.max(0, input.parseErrorCount ?? 0);
+        const unknownEventCount = Math.max(0, input.unknownEventCount ?? 0);
+        const bytesReceived = Math.max(0, input.bytesReceived ?? 0);
+        const updatedAt = input.updatedAt ?? input.lastNormalizedEventAt ?? input.lastRawEventAt ?? new Date().toISOString();
+        this.db.prepare(`
+			UPDATE telemetry_provider_requests SET
+				status = CASE WHEN status IN ('completed', 'error', 'aborted', 'timeout') THEN status ELSE COALESCE(?, status) END,
+				last_raw_event_at = COALESCE(?, last_raw_event_at),
+				last_normalized_event_at = COALESCE(?, last_normalized_event_at),
+				upstream_response_id = COALESCE(?, upstream_response_id),
+				raw_event_count = raw_event_count + ?,
+				normalized_event_count = normalized_event_count + ?,
+				parse_error_count = parse_error_count + ?,
+				unknown_event_count = unknown_event_count + ?,
+				bytes_received = CASE WHEN ? = 0 THEN bytes_received ELSE COALESCE(bytes_received, 0) + ? END,
+				event_type_counts_json = ?,
+				updated_at = ?
+			WHERE provider_request_id = ?
+		`).run(input.status ?? null, input.lastRawEventAt ?? null, input.lastNormalizedEventAt ?? null, input.upstreamResponseId ?? null, rawEventCount, normalizedEventCount, parseErrorCount, unknownEventCount, bytesReceived, bytesReceived, JSON.stringify(eventTypeCounts), updatedAt, input.providerRequestId);
+        return {
+            ...existing,
+            status: isTerminalProviderRequestStatus(existing.status) ? existing.status : input.status ?? existing.status,
+            lastRawEventAt: input.lastRawEventAt ?? existing.lastRawEventAt,
+            lastNormalizedEventAt: input.lastNormalizedEventAt ?? existing.lastNormalizedEventAt,
+            upstreamResponseId: input.upstreamResponseId ?? existing.upstreamResponseId,
+            rawEventCount: existing.rawEventCount + rawEventCount,
+            normalizedEventCount: existing.normalizedEventCount + normalizedEventCount,
+            parseErrorCount: existing.parseErrorCount + parseErrorCount,
+            unknownEventCount: existing.unknownEventCount + unknownEventCount,
+            bytesReceived: bytesReceived > 0 ? (existing.bytesReceived ?? 0) + bytesReceived : existing.bytesReceived,
+            eventTypeCounts,
+            updatedAt,
+        };
+    }
     appendProviderEventSummary(input) {
         const now = input.updatedAt ?? new Date().toISOString();
         const receivedAt = input.receivedAt ?? now;
@@ -293,34 +395,15 @@ export class TelemetryStore {
         return row.next_sequence;
     }
     incrementProviderCounters(providerRequestId, eventType, receivedAt, byteSize, parseStatus, normalizedDelta) {
-        const existing = this.getProviderRequest(providerRequestId);
-        if (!existing)
-            return;
-        const eventTypeCounts = { ...existing.eventTypeCounts };
-        const currentCount = typeof eventTypeCounts[eventType] === "number" ? eventTypeCounts[eventType] : 0;
-        eventTypeCounts[eventType] = currentCount + 1;
-        this.upsertProviderRequest({
+        this.recordProviderProgress({
             providerRequestId,
-            piboSessionId: existing.piboSessionId,
-            rootSessionId: existing.rootSessionId,
-            roomId: existing.roomId,
-            turnId: existing.turnId,
-            phaseId: existing.phaseId,
-            provider: existing.provider,
-            api: existing.api,
-            model: existing.model,
-            transport: existing.transport,
-            serviceTier: existing.serviceTier,
-            status: existing.status,
             lastRawEventAt: receivedAt,
-            rawEventCount: existing.rawEventCount + 1,
-            normalizedEventCount: existing.normalizedEventCount + normalizedDelta,
-            parseErrorCount: existing.parseErrorCount + (parseStatus === "invalid_json" ? 1 : 0),
-            unknownEventCount: existing.unknownEventCount + (parseStatus === "unknown_type" ? 1 : 0),
-            bytesReceived: (existing.bytesReceived ?? 0) + byteSize,
-            eventTypeCounts,
-            captureMode: existing.captureMode,
-            retentionClass: existing.retentionClass,
+            rawEventCount: 1,
+            normalizedEventCount: normalizedDelta,
+            parseErrorCount: parseStatus === "invalid_json" ? 1 : 0,
+            unknownEventCount: parseStatus === "unknown_type" ? 1 : 0,
+            bytesReceived: byteSize,
+            eventTypeCounts: { [eventType]: 1 },
             updatedAt: receivedAt,
         });
     }
@@ -347,6 +430,9 @@ export class BestEffortTelemetryService {
     recordProviderEventSummary(input) {
         this.safe(() => this.store?.recordProviderEventSummary(input));
     }
+    recordProviderProgress(input) {
+        return this.safe(() => this.store?.recordProviderProgress(input));
+    }
     appendProviderEventSummary(input) {
         return this.safe(() => this.store?.appendProviderEventSummary(input));
     }
@@ -363,6 +449,9 @@ export class BestEffortTelemetryService {
         }
     }
 }
+function isTerminalProviderRequestStatus(status) {
+    return status === "completed" || status === "error" || status === "aborted" || status === "timeout";
+}
 function fail(message) {
     throw new Error(message);
 }

package/dist/gateway/request.js

@@ -75,6 +75,7 @@ export async function sendGatewayMessageAndWaitForReply(event, options = {}) {
         let settled = false;
         let response;
         let reply;
+        let messageFinished = false;
         const timeout = setTimeout(() => {
             finish(new Error(`Timed out waiting for assistant reply from session "${event.piboSessionId}"`));
         }, timeoutMs);
@@ -91,14 +92,21 @@ export async function sendGatewayMessageAndWaitForReply(event, options = {}) {
                 resolve(result);
             }
         };
+        const finishCompletedReply = () => {
+            if (!response?.ok || !messageFinished)
+                return;
+            finish(reply
+                ? { response, reply }
+                : new Error(`Session "${eventWithId.piboSessionId}" finished without an assistant reply`));
+        };
         const handleFrame = (frame) => {
             if (frame.type === "res" && frame.id === id) {
                 response = frame;
                 if (!frame.ok) {
                     finish(new Error(frame.error?.message ?? "Gateway rejected the message"));
                 }
-                else if (reply) {
-                    finish({ response, reply });
+                else {
+                    finishCompletedReply();
                 }
                 return;
             }
@@ -111,13 +119,14 @@ export async function sendGatewayMessageAndWaitForReply(event, options = {}) {
                 finish(new Error(output.error));
                 return;
             }
-            if (output.type === "assistant_message" &&
-                output.piboSessionId === eventWithId.piboSessionId &&
-                output.eventId === eventWithId.id) {
+            if (output.piboSessionId !== eventWithId.piboSessionId || !("eventId" in output) || output.eventId !== eventWithId.id)
+                return;
+            if (output.type === "assistant_message") {
                 reply = output;
-                if (response?.ok) {
-                    finish({ response, reply });
-                }
+            }
+            else if (output.type === "message_finished") {
+                messageFinished = true;
+                finishCompletedReply();
             }
         };
         socket.once("connect", () => {