@superblocksteam/vite-plugin-file-sync 2.0.151 → 2.0.152-next.0

package/dist/ai-service/index.js

@@ -2,7 +2,7 @@ import { createHash, randomUUID } from "node:crypto";
 import os from "node:os";
 import path from "node:path";
 import { context as otelContext, propagation, SpanKind, SpanStatusCode, trace, } from "@opentelemetry/api";
-import { stepCountIs, streamText as aiStreamText, } from "ai";
+import { hasToolCall, streamText as aiStreamText, } from "ai";
 import { z } from "zod";
 import { AiGenerationState, AiMode, resolveAttachmentFileName, } from "@superblocksteam/library-shared/types";
 import { isToolPermissionClarification } from "@superblocksteam/library-shared/types/tool-permission-response";
@@ -34,6 +34,7 @@ import { ChatSessionStore } from "./chat/chat-session-store.js";
 import { normalizeTextAttachmentForModel } from "./chat/utils.js";
 import { API_MIGRATION_ORIGINS, isApiMigrationOrigin, readPersistedChecklist, updateChecklistItem, } from "./checklist/persisted-checklist-store.js";
 import { downloadContextFromBucketeer, downloadLatestContextFromBucketeer, restoreContextFromCheckpoint, } from "./context-download.js";
+import { restoreConversationToCheckpoint, } from "./context-restore.js";
 import { uploadContextToBucketeer } from "./context-upload.js";
 import { AppContextStore } from "./context/app-context.js";
 import { createDevDatabaseHttpTransport, DevDatabaseClient, } from "./dev-database-client.js";
@@ -47,8 +48,10 @@ import { createChaosFetch } from "./llm/chaos-fetch.js";
 import { createClarkStreamTelemetrySink } from "./llm/clark-stream-telemetry-sink.js";
 import { LLMClient } from "./llm/client.js";
 import { selectContextManagementStrategy, } from "./llm/context-v2/context-management.js";
+import { recordCheckpointRestore } from "./llm/context-v2/context-metrics.js";
 import { ContextManagerV2 } from "./llm/context-v2/manager.js";
 import { JsonlEventStore } from "./llm/context-v2/storage/jsonl-event-store.js";
+import { adoptLegacyPerUserContextDirs } from "./llm/context-v2/storage/legacy-context-adoption.js";
 import { LocalContextStorageV2 } from "./llm/context-v2/storage/local.js";
 import { createLLMProvider, setSpanError } from "./llm/provider.js";
 import { traceLLM } from "./llmobs/helpers.js";
@@ -60,6 +63,7 @@ import { recordMigrationAutoDispatchOutcome } from "./migration-autodispatch-met
 import { buildServerMigrationLlmConfig, isMigrationAutoDispatchEnabled, } from "./migration-llm-config.js";
 import { isMigrationSessionActive } from "./migration-session-active.js";
 import { hasSettledNormalMigrationTurn, recordMigrationTurnDispatched, recordMigrationTurnSettled, } from "./migration-turns-log.js";
+import { recordPolicyAgentExecution, } from "./policy-agent-metrics.js";
 import { policyGateStreamPartToAgentStep } from "./policy-gate-step.js";
 import { ClarkProfiler } from "./profiler/clark-profiler.js";
 import { buildSummarizeApiUsagePrompt } from "./prompts/summarize-api-usage.js";
@@ -83,7 +87,7 @@ import { sendUserGenerationStateChannel } from "./state-machine/helpers/peer.js"
 import { sendUserMessageChannel, clearPendingToolPermissionRequest, } from "./state-machine/helpers/peer.js";
 import { finalizePromptBillingMetrics } from "./state-machine/helpers/prompt-billing-metrics.js";
 import { paywallStopAdvice } from "./state-machine/helpers/prompt-failure.js";
-import { promptFailureFromStreamAbort, promptFailurePaywall, promptFailurePlatform, promptFailureUserCanceled, recordPromptAttempt, recordPromptEntryRejection, recordPromptOutcome, recordAwaitingUserHandoffDismissed, resolvePromptMetricLabels, resolvePromptOutcomeDedupeKey, resolvePromptRequestDedupeKey, setPromptRequestMetricIdentity, } from "./state-machine/helpers/prompt-request-metrics.js";
+import { promptFailureFromStreamAbort, promptFailurePaywall, promptFailurePlatform, promptFailureUserCanceled, recordBaseCommitResolutionFailed, recordPromptAttempt, recordPromptEntryRejection, recordPromptOutcome, recordAwaitingUserHandoffDismissed, resolvePromptMetricLabels, resolvePromptOutcomeDedupeKey, resolvePromptRequestDedupeKey, setPromptRequestMetricIdentity, } from "./state-machine/helpers/prompt-request-metrics.js";
 import { catchBackgroundPeerDisconnect, StablePeer, } from "./state-machine/helpers/stable-peer.js";
 import { transitionFrom } from "./state-machine/helpers/transition.js";
 import { isTaskAccessibleTo, registerBackgroundTaskTypes, resyncActiveTaskStatusesToPeer, TaskStore, wireTaskStatusPushToPeer, } from "./tasks/index.js";
@@ -132,6 +136,9 @@ const MAX_POLICY_GATE_STREAM_ERROR_ARRAY_LENGTH = 10;
 const MAX_POLICY_GATE_STREAM_ERROR_FIELD_LENGTH = 500;
 const MAX_POLICY_GATE_STREAM_ERROR_OBJECT_DEPTH = 3;
 const MAX_POLICY_GATE_STREAM_ERROR_OBJECT_KEYS = 20;
+const POLICY_GATE_MAX_TOTAL_TIME_MS = 15 * 60 * 1000;
+// AbortError near the budget is treated as timeout (SDK may abort slightly early).
+const POLICY_GATE_TIMEOUT_CLASSIFICATION_SLACK_MS = 5_000;
 const POLICY_GATE_STREAM_ERROR_REDACTED_VALUE = "[redacted]";
 const POLICY_GATE_STREAM_ERROR_SECRET_KEY_PATTERN = /(access[-_]?key|api[-_]?key|authorization|connection[-_]?string|credential|dsn|password|private[-_]?key|secret|token)/i;
 const POLICY_GATE_STREAM_ERROR_SECRET_VALUE_PATTERN_SOURCES = [
@@ -284,6 +291,17 @@ const NPM_REGISTRY_JWT_WAIT_MS = 5_000;
  *     via the second-401 branch in `NpmRegistryClient.fetchWithRefresh`.
  */
 const NPM_REGISTRY_REFRESH_WAIT_MS = 5_000;
+/**
+ * Deadline for the chat-store fallback in `resolveBaseCommitId` (2 seconds).
+ *
+ * baseCommitId is a billing annotation, not something the prompt needs to
+ * start -- a slow or hung `chatSessionStore.getLatestCommitId()` call must
+ * not delay the turn. A store that never resolves within this window is
+ * treated the same as a failure: the counter records
+ * `source: "chat_store_timeout"`, a correlated warn is logged, and the
+ * prompt proceeds with no baseCommitId rather than waiting indefinitely.
+ */
+const RESOLVE_BASE_COMMIT_ID_CHAT_STORE_TIMEOUT_MS = 2_000;
 export function truncateString(str, maxBytes) {
     // Rough approximation: assume 2 bytes per character
     const maxChars = Math.floor(maxBytes / 2);
@@ -800,6 +818,9 @@ export class AiService extends TracedEventEmitter {
             rpcClient: config.rpcClient,
             logger: config.logger,
         });
+        // Lazy closure: this.clark is assigned later in the constructor, so the
+        // provider reads the active turn's promptId at append time, not now.
+        this.chatSessionStore.setPromptIdProvider(() => this.clark?.context?.promptId);
         if (this.config.tokenManager) {
             this.config.tokenManager.on("tokenUpdated", this.handleTokenUpdated.bind(this));
         }
@@ -809,6 +830,7 @@ export class AiService extends TracedEventEmitter {
             storage: new LocalContextStorageV2(config.appRootDirPath),
             eventStore: new JsonlEventStore(config.appRootDirPath),
             rootDir: config.appRootDirPath,
+            adoptLegacyContexts: () => adoptLegacyPerUserContextDirs(config.appRootDirPath),
             getModeState: () => {
                 const mode = this.clark?.context?.currentMode;
                 if (mode === undefined)
@@ -1267,10 +1289,11 @@ export class AiService extends TracedEventEmitter {
                         to === ClarkStateNames.AwaitingUser) {
                         // The turn is over: drop per-turn billing attribution so stray
                         // LLM calls before the next prompt bill as default generation.
-                        if (clark.context.usageMetadata) {
+                        if (clark.context.usageMetadata || clark.context.baseCommitId) {
                             clark.updateContext((context) => ({
                                 ...context,
                                 usageMetadata: undefined,
+                                baseCommitId: undefined,
                             }));
                         }
                         this.triggerGenerationComplete();
@@ -1412,8 +1435,80 @@ export class AiService extends TracedEventEmitter {
             summarizeModel: () => createLLMProvider(this.getLLMProviderSettings(llmConfig), () => this.getProviderInferenceAuthToken(), llmConfig, () => this.getLlmAiAttributes(this.clark), this.clark.context.traceHeaders).modelForTask("summarizeMessages"),
         });
     }
+    /**
+     * The checkpoint the app is currently based on: the last non-draft commit
+     * seeded by sync (preferred — same key the post-generation upload uses) or
+     * the most recent chat-history commit. Captured once at prompt start
+     * because the checkpoint an eventual accept will create does not exist yet
+     * while the LLM calls run. Never throws: a billing annotation must not be
+     * able to break generation.
+     */
+    async resolveBaseCommitId(promptId, chatStoreTimeoutMs = RESOLVE_BASE_COMMIT_ID_CHAT_STORE_TIMEOUT_MS) {
+        let preferredSourceFailed = false;
+        try {
+            // Normalize falsy (e.g. "") to undefined so an empty string can't
+            // short-circuit the `??` fallback below and then be silently dropped
+            // by the truthy merge guard downstream in getLlmAiAttributes.
+            const syncCommitId = this._syncContextProvider?.getLastCommitId() || undefined;
+            if (syncCommitId) {
+                return syncCommitId;
+            }
+        }
+        catch (error) {
+            preferredSourceFailed = true;
+            this.getLogger().warn("[ai-service] Failed to resolve sync base commit id; falling back to chat history", {
+                applicationId: this.config.applicationId,
+                sessionId: this.sessionId,
+                promptId,
+                ...getErrorMeta(error),
+            });
+        }
+        let source = "chat_store";
+        try {
+            const chatStoreCommitId = await this.raceChatStoreLatestCommitId(chatStoreTimeoutMs, () => {
+                source = "chat_store_timeout";
+            });
+            return chatStoreCommitId || undefined;
+        }
+        catch (error) {
+            this.recordAndWarnBaseCommitResolutionFailure(source, preferredSourceFailed, promptId, error);
+            return undefined;
+        }
+    }
+    /**
+     * Races the chat-store fallback against `timeoutMs` so a hung store
+     * degrades exactly like a failure instead of delaying prompt start. The
+     * timeout rejects (never resolves) so the caller's catch block handles
+     * both outcomes uniformly; `onTimeout` lets the caller distinguish which
+     * branch fired for metrics/logging without a second code path.
+     * `.unref()` keeps the fallback timer from holding the process open.
+     */
+    raceChatStoreLatestCommitId(timeoutMs, onTimeout) {
+        let timeoutHandle;
+        const timeout = new Promise((_, reject) => {
+            timeoutHandle = setTimeout(() => {
+                onTimeout();
+                reject(new Error(`chatSessionStore.getLatestCommitId timed out after ${timeoutMs}ms`));
+            }, timeoutMs);
+            timeoutHandle.unref?.();
+        });
+        return Promise.race([
+            this.chatSessionStore.getLatestCommitId(),
+            timeout,
+        ]).finally(() => clearTimeout(timeoutHandle));
+    }
+    recordAndWarnBaseCommitResolutionFailure(source, preferredSourceFailed, promptId, error) {
+        recordBaseCommitResolutionFailed({ source, preferredSourceFailed });
+        this.getLogger().warn("[ai-service] Failed to resolve base commit id for usage metadata", {
+            applicationId: this.config.applicationId,
+            sessionId: this.sessionId,
+            promptId,
+            ...getErrorMeta(error),
+        });
+    }
     getLlmAiAttributes(clark) {
         const usageMetadata = clark.context.usageMetadata ?? APP_CLARK_GENERATION_METADATA;
+        const baseCommitId = clark.context.baseCommitId;
         return {
             prompt_id: clark.context.promptId,
             chatmessage_id: clark.context.chatmessageId,
@@ -1422,7 +1517,9 @@ export class AiService extends TracedEventEmitter {
             // Note: retry_count for LLM retries is added by orchestrator via providerOptions
             mode: clark.context.currentMode === AiMode.PLAN ? "plan" : "build",
             plan_size: clark.context.planSize,
-            metadata: usageMetadata,
+            metadata: baseCommitId
+                ? { ...usageMetadata, baseCommitId }
+                : usageMetadata,
         };
     }
     /**
@@ -1626,9 +1723,11 @@ export class AiService extends TracedEventEmitter {
             retryMode,
             triggeredBy,
         };
+        const baseCommitId = await this.resolveBaseCommitId(promptId);
         this.clark.updateContext((context) => ({
             ...context,
             promptId,
+            baseCommitId,
             usageMetadata: migrationUsageMetadata,
             llmConfig: migrationLlmConfig,
         }));
@@ -2072,11 +2171,13 @@ export class AiService extends TracedEventEmitter {
                 dedupeKeySource: request.requestId ? "requestId" : "ephemeral",
                 applicationId: this.config.applicationId,
             });
+            const baseCommitId = await this.resolveBaseCommitId(promptId);
             this.clark.updateContext((context) => {
                 return {
                     ...context,
                     promptId,
                     promptOutcomeDedupeKey,
+                    baseCommitId,
                     isInterruptAbort: undefined,
                     userInitiatedPromptAbort: undefined,
                     // Third-party import turns are billed as app_migration; everything
@@ -2646,6 +2747,7 @@ export class AiService extends TracedEventEmitter {
             promptId: undefined,
             promptOutcomeDedupeKey: undefined,
             usageMetadata: undefined,
+            baseCommitId: undefined,
             awaitingUserHandoffPromptId: undefined,
             awaitingUserHandoffRecorded: undefined,
         }));
@@ -3645,25 +3747,11 @@ export class AiService extends TracedEventEmitter {
     async handlePolicyGateExecution(request) {
         const logger = this.getLogger();
         logger.info(`[ai-service] Policy gate execution requested for app=${request.applicationId}`);
-        const requestLlmConfig = request.llmConfig;
-        const contextLlmConfig = this.clark.context.llmConfig;
-        const llmConfig = requestLlmConfig ?? contextLlmConfig;
-        const llmProvider = createLLMProvider(this.getLLMProviderSettings(llmConfig), () => this.requireProvidedInferenceJwt(request.jwt, "policy gate"), llmConfig, () => ({
-            ai_agent_run_id: request.aiAgentRunId,
-            application_id: request.applicationId,
-            metadata: POLICY_GATE_USAGE_METADATA,
-            session_id: this.sessionId,
-            trigger_type: "policy_gate",
-        }));
-        const { model } = llmProvider.modelForTask("generateBroadEdit");
-        const availableIntegrations = await this.integrationStore.availableIntegrations();
-        const integrations = availableIntegrations.map((i) => ({
-            name: i.name,
-            pluginType: i.plugin?.id ?? "unknown",
-        }));
-        const appContext = { integrations };
-        const systemPrompt = composeSecurityAgentPrompt(appContext, request.adminPrompt);
         const startedAt = Date.now();
+        const policyRunLogContext = {
+            applicationId: request.applicationId,
+            reviewPolicyId: request.reviewPolicyId,
+        };
         const toolCallCounts = {};
         const filePathsRead = new Set();
         const filePathsListed = new Set();
@@ -3671,15 +3759,31 @@ export class AiService extends TracedEventEmitter {
         const grepPatterns = [];
         const grepMatchCounts = [];
         const grepPatternByToolCallId = new Map();
+        const lsPathByToolCallId = new Map();
         const streamErrors = [];
         const streamPartTypeCounts = {};
+        let stepCount = 0;
         let textDeltaChars = 0;
-        const promptHash = createHash("sha256").update(systemPrompt).digest("hex");
-        const modelName = typeof model === "string" ? model : model.modelId;
+        let modelName;
+        let promptHash;
         const scanRoot = typeof this.appShell.appRootDirPath === "string"
             ? this.appShell.appRootDirPath
             : undefined;
         let scanResult;
+        let metricsRecorded = false;
+        const totalToolCalls = () => Object.values(toolCallCounts).reduce((sum, count) => sum + count, 0);
+        const recordExecutionMetrics = (outcome) => {
+            if (metricsRecorded) {
+                return;
+            }
+            metricsRecorded = true;
+            recordPolicyAgentExecution({
+                durationMs: Date.now() - startedAt,
+                outcome,
+                stepCount,
+                toolCallCount: totalToolCalls(),
+            });
+        };
         const buildRuntimeMetadata = (finalFailureReason) => ({
             diagnosticsVersion: 2,
             durationMs: Date.now() - startedAt,
@@ -3694,156 +3798,260 @@ export class AiService extends TracedEventEmitter {
             reportToolCalled: Boolean(scanResult),
             scanComplete: scanResult?.scanComplete === true,
             scanRoot,
+            stepCount,
             streamErrors,
             streamPartTypeCounts,
             textDeltaChars,
             toolCallCounts,
         });
-        const failPolicyGateExecution = (message) => {
+        const failPolicyGateExecution = (message, outcome) => {
+            recordExecutionMetrics(outcome);
             const error = new Error(message);
             error.runtimeMetadataJson = buildRuntimeMetadata(message);
             throw error;
         };
-        const scanToolsDeps = {
-            appShell: this.appShell,
-            features: this.config.features ?? {},
-            integrationStore: this.integrationStore,
-            llmClient: this.llmClient,
-            llmProvider: llmProvider,
-        };
-        const tools = await buildSecurityScanTools(this.clark, scanToolsDeps);
-        const result = aiStreamText({
-            messages: [
-                { content: systemPrompt, role: "system" },
-                {
-                    content: "Begin your security review of this application now.",
-                    role: "user",
-                },
-            ],
-            model,
-            tools,
-            stopWhen: stepCountIs(30),
-        });
-        // Best-effort: surface the agent's latest tool/reasoning step in the editor
-        // so a running policy row shows live activity instead of an opaque spinner.
-        // Only meaningful when an editor session is connected; never blocks the run.
-        const peer = this.clark.context.peer;
-        const emitAgentStep = (part) => {
-            if (!this.hasConnectedPeer() || !peer) {
-                return;
-            }
-            const step = policyGateStreamPartToAgentStep(part);
-            if (!step) {
-                return;
-            }
-            void peer.call
-                .aiPushPolicyAgentStep({
-                applicationId: request.applicationId,
-                policyId: request.reviewPolicyId,
-                step,
-            })
-                .catch((error) => {
-                logger.warn("[ai-service] Failed to push policy agent step", getErrorMeta(error));
+        try {
+            const requestLlmConfig = request.llmConfig;
+            const contextLlmConfig = this.clark.context.llmConfig;
+            const llmConfig = requestLlmConfig ?? contextLlmConfig;
+            const llmProvider = createLLMProvider(this.getLLMProviderSettings(llmConfig), () => this.requireProvidedInferenceJwt(request.jwt, "policy gate"), llmConfig, () => ({
+                ai_agent_run_id: request.aiAgentRunId,
+                application_id: request.applicationId,
+                metadata: POLICY_GATE_USAGE_METADATA,
+                session_id: this.sessionId,
+                trigger_type: "policy_gate",
+            }));
+            const { model } = llmProvider.modelForTask("generateBroadEdit");
+            modelName = typeof model === "string" ? model : model.modelId;
+            const availableIntegrations = await this.integrationStore.availableIntegrations();
+            const integrations = availableIntegrations.map((i) => ({
+                name: i.name,
+                pluginType: i.plugin?.id ?? "unknown",
+            }));
+            const appContext = { integrations };
+            const systemPrompt = composeSecurityAgentPrompt(appContext, request.adminPrompt);
+            promptHash = createHash("sha256").update(systemPrompt).digest("hex");
+            const scanToolsDeps = {
+                appShell: this.appShell,
+                features: this.config.features ?? {},
+                integrationStore: this.integrationStore,
+                llmClient: this.llmClient,
+                llmProvider: llmProvider,
+            };
+            const tools = await buildSecurityScanTools(this.clark, scanToolsDeps);
+            const result = aiStreamText({
+                messages: [
+                    { content: systemPrompt, role: "system" },
+                    {
+                        content: "Begin your security review of this application now.",
+                        role: "user",
+                    },
+                ],
+                model,
+                tools,
+                stopWhen: hasToolCall("reportReviewRun"),
+                timeout: { totalMs: POLICY_GATE_MAX_TOTAL_TIME_MS },
             });
-        };
-        for await (const part of result.fullStream) {
-            emitAgentStep(part);
-            const partType = typeof part.type === "string"
-                ? part.type
-                : "unknown";
-            streamPartTypeCounts[partType] =
-                (streamPartTypeCounts[partType] ?? 0) + 1;
-            if (partType === "error") {
-                const streamError = serializePolicyGateStreamError(part.error, part);
-                if (streamErrors.length < MAX_POLICY_GATE_STREAM_ERRORS) {
-                    streamErrors.push(streamError);
+            // Best-effort: surface the agent's latest tool/reasoning step in the editor
+            // so a running policy row shows live activity instead of an opaque spinner.
+            // Only meaningful when an editor session is connected; never blocks the run.
+            const peer = this.clark.context.peer;
+            const emitAgentStep = (part) => {
+                if (!this.hasConnectedPeer() || !peer) {
+                    return;
                 }
-                logger.warn("[ai-service] Policy gate stream error", { streamError });
-            }
-            const textDelta = part.textDelta ??
-                part.text ??
-                part.delta;
-            if (typeof textDelta === "string") {
-                textDeltaChars += textDelta.length;
-            }
-            const finishReason = part.finishReason;
-            if (typeof finishReason === "string") {
-                finishReasons.push(finishReason);
-            }
-            if (part.type === "tool-call") {
-                toolCallCounts[part.toolName] =
-                    (toolCallCounts[part.toolName] ?? 0) + 1;
-                const input = part.input ??
-                    part.args;
-                if (input && typeof input === "object") {
-                    const record = input;
-                    if (part.toolName === "build_readFile" &&
-                        typeof record.filePath === "string") {
-                        filePathsRead.add(record.filePath);
+                const step = policyGateStreamPartToAgentStep(part);
+                if (!step) {
+                    return;
+                }
+                void peer.call
+                    .aiPushPolicyAgentStep({
+                    applicationId: request.applicationId,
+                    policyId: request.reviewPolicyId,
+                    step,
+                })
+                    .catch((error) => {
+                    logger.warn("[ai-service] Failed to push policy agent step", getErrorMeta(error));
+                });
+            };
+            try {
+                for await (const part of result.fullStream) {
+                    emitAgentStep(part);
+                    const partType = typeof part.type === "string"
+                        ? part.type
+                        : "unknown";
+                    streamPartTypeCounts[partType] =
+                        (streamPartTypeCounts[partType] ?? 0) + 1;
+                    if (partType === "error") {
+                        const streamError = serializePolicyGateStreamError(part.error, part);
+                        if (streamErrors.length < MAX_POLICY_GATE_STREAM_ERRORS) {
+                            streamErrors.push(streamError);
+                        }
+                        logger.warn("[ai-service] Policy gate stream error", {
+                            ...policyRunLogContext,
+                            streamError,
+                        });
                     }
-                    if (part.toolName === "grep" && typeof record.pattern === "string") {
-                        grepPatterns.push(record.pattern);
-                        const toolCallId = typeof part.toolCallId === "string"
-                            ? part.toolCallId
-                            : `grep-${grepPatterns.length}`;
-                        grepPatternByToolCallId.set(toolCallId, record.pattern);
+                    const textDelta = part.textDelta ??
+                        part.text ??
+                        part.delta;
+                    if (typeof textDelta === "string") {
+                        textDeltaChars += textDelta.length;
                     }
-                }
-            }
-            if (part.type === "tool-result") {
-                const output = part.output ??
-                    part.result;
-                if (output && typeof output === "object") {
-                    const record = output;
-                    if (part.toolName === "build_listFiles" &&
-                        Array.isArray(record.files)) {
-                        for (const file of record.files) {
-                            if (typeof file === "string") {
-                                filePathsListed.add(file);
+                    if (partType === "finish-step") {
+                        stepCount += 1;
+                    }
+                    const finishReason = part
+                        .finishReason;
+                    if (typeof finishReason === "string") {
+                        finishReasons.push(finishReason);
+                    }
+                    if (part.type === "tool-call") {
+                        toolCallCounts[part.toolName] =
+                            (toolCallCounts[part.toolName] ?? 0) + 1;
+                        const input = part.input ??
+                            part.args;
+                        if (input && typeof input === "object") {
+                            const record = input;
+                            if (part.toolName === "build_readFile" &&
+                                typeof record.filePath === "string") {
+                                filePathsRead.add(record.filePath);
+                            }
+                            if (part.toolName === "grep" &&
+                                typeof record.pattern === "string") {
+                                grepPatterns.push(record.pattern);
+                                const toolCallId = typeof part.toolCallId ===
+                                    "string"
+                                    ? part.toolCallId
+                                    : `grep-${grepPatterns.length}`;
+                                grepPatternByToolCallId.set(toolCallId, record.pattern);
+                            }
+                            if (part.toolName === "ls") {
+                                const toolCallId = typeof part.toolCallId ===
+                                    "string"
+                                    ? part.toolCallId
+                                    : `ls-${lsPathByToolCallId.size + 1}`;
+                                lsPathByToolCallId.set(toolCallId, typeof record.path === "string" ? record.path : ".");
                             }
                         }
                     }
-                    if (part.toolName === "grep" && typeof record.matches === "number") {
-                        const toolCallId = typeof part.toolCallId === "string"
-                            ? part.toolCallId
-                            : `grep-${grepMatchCounts.length + 1}`;
-                        grepMatchCounts.push({
-                            matches: record.matches,
-                            pattern: grepPatternByToolCallId.get(toolCallId) ??
-                                grepPatterns[grepPatterns.length - 1] ??
-                                "",
-                        });
+                    if (part.type === "tool-result") {
+                        const output = part.output ??
+                            part.result;
+                        if (output && typeof output === "object") {
+                            const record = output;
+                            if ((part.toolName === "build_listFiles" ||
+                                part.toolName === "glob") &&
+                                Array.isArray(record.files)) {
+                                for (const file of record.files) {
+                                    if (typeof file === "string") {
+                                        filePathsListed.add(file);
+                                    }
+                                }
+                            }
+                            if (part.toolName === "ls" && typeof record.output === "string") {
+                                const toolCallId = typeof part.toolCallId ===
+                                    "string"
+                                    ? part.toolCallId
+                                    : undefined;
+                                const dirPath = (toolCallId
+                                    ? lsPathByToolCallId.get(toolCallId)
+                                    : undefined) ?? ".";
+                                for (const line of record.output.split("\n")) {
+                                    const trimmed = line.trim();
+                                    if (!trimmed || trimmed.startsWith("total")) {
+                                        continue;
+                                    }
+                                    const name = trimmed.split(/\s+/).at(-1);
+                                    if (!name) {
+                                        continue;
+                                    }
+                                    filePathsListed.add(dirPath === "."
+                                        ? name
+                                        : `${dirPath.replace(/\/$/, "")}/${name}`);
+                                }
+                            }
+                            if (part.toolName === "grep" &&
+                                typeof record.matches === "number") {
+                                const toolCallId = typeof part.toolCallId ===
+                                    "string"
+                                    ? part.toolCallId
+                                    : `grep-${grepMatchCounts.length + 1}`;
+                                grepMatchCounts.push({
+                                    matches: record.matches,
+                                    pattern: grepPatternByToolCallId.get(toolCallId) ??
+                                        grepPatterns[grepPatterns.length - 1] ??
+                                        "",
+                                });
+                            }
+                        }
+                    }
+                    if (part.type === "tool-call" &&
+                        part.toolName === "reportReviewRun") {
+                        const input = part.input ??
+                            part.args;
+                        const parsedInput = securityScanResultSchema.safeParse(input);
+                        if (!parsedInput.success) {
+                            logger.error("[ai-service] Invalid policy gate reportReviewRun input", {
+                                ...policyRunLogContext,
+                                ...getErrorMeta(parsedInput.error),
+                            });
+                            failPolicyGateExecution("Invalid policy gate reportReviewRun input", "invalid_report");
+                        }
+                        scanResult = parsedInput.data;
                     }
                 }
             }
-            if (part.type === "tool-call" && part.toolName === "reportReviewRun") {
-                const input = part.input ??
-                    part.args;
-                const parsedInput = securityScanResultSchema.safeParse(input);
-                if (!parsedInput.success) {
-                    logger.error("[ai-service] Invalid policy gate reportReviewRun input", getErrorMeta(parsedInput.error));
-                    failPolicyGateExecution("Invalid policy gate reportReviewRun input");
+            catch (error) {
+                if (error instanceof Error &&
+                    (error.name === "AbortError" || error.name === "TimeoutError")) {
+                    const elapsedMs = Date.now() - startedAt;
+                    const isBudgetTimeout = error.name === "TimeoutError" ||
+                        elapsedMs + POLICY_GATE_TIMEOUT_CLASSIFICATION_SLACK_MS >=
+                            POLICY_GATE_MAX_TOTAL_TIME_MS;
+                    if (isBudgetTimeout) {
+                        logger.warn("[ai-service] Policy gate agent timed out", {
+                            ...policyRunLogContext,
+                            ...getErrorMeta(error),
+                            elapsedMs,
+                        });
+                        failPolicyGateExecution("Policy gate agent timed out", "timeout");
+                    }
+                    logger.warn("[ai-service] Policy gate agent aborted", {
+                        ...policyRunLogContext,
+                        ...getErrorMeta(error),
+                        elapsedMs,
+                    });
+                    failPolicyGateExecution("Policy gate agent aborted", "aborted");
                 }
-                scanResult = parsedInput.data;
+                throw error;
             }
+            if (!scanResult) {
+                logger.warn("[ai-service] Policy gate agent did not call reportReviewRun", policyRunLogContext);
+                failPolicyGateExecution("Policy gate agent did not complete a report", "missing_report");
+            }
+            const completedScanResult = scanResult;
+            if (!completedScanResult.scanComplete) {
+                logger.warn("[ai-service] Policy gate agent reported an incomplete scan", {
+                    ...policyRunLogContext,
+                    findingCount: completedScanResult.findings.length,
+                    scannedFiles: completedScanResult.scannedFiles,
+                });
+                failPolicyGateExecution("Policy gate agent did not complete a report", "incomplete_report");
+            }
+            logger.info(`[ai-service] Policy gate execution complete: ${completedScanResult.findings.length} findings`);
+            recordExecutionMetrics("success");
+            return {
+                ...completedScanResult,
+                runtimeMetadataJson: buildRuntimeMetadata(),
+            };
         }
-        if (!scanResult) {
-            logger.warn("[ai-service] Policy gate agent did not call reportReviewRun");
-            failPolicyGateExecution("Policy gate agent did not complete a report");
-        }
-        const completedScanResult = scanResult;
-        if (!completedScanResult.scanComplete) {
-            logger.warn("[ai-service] Policy gate agent reported an incomplete scan", {
-                findingCount: completedScanResult.findings.length,
-                scannedFiles: completedScanResult.scannedFiles,
-            });
-            failPolicyGateExecution("Policy gate agent did not complete a report");
+        catch (error) {
+            if (!metricsRecorded) {
+                recordExecutionMetrics("error");
+            }
+            throw error;
         }
-        logger.info(`[ai-service] Policy gate execution complete: ${completedScanResult.findings.length} findings`);
-        return {
-            ...completedScanResult,
-            runtimeMetadataJson: buildRuntimeMetadata(),
-        };
     }
     handleUserChange(change) {
         this.getLogger().info(`[ai-service] User changed: ${change.type}`, change);
@@ -3867,7 +4075,7 @@ export class AiService extends TracedEventEmitter {
             // instead of the one the user just cleared (APPS-4407).
             if (contextId) {
                 await this.contextManagerV2.persistEmptyContext(contextId);
-                await this.uploadEmptyContextToBucketeer(contextId.userId, commitIdForUpload);
+                await this.uploadEmptyContextToBucketeer(commitIdForUpload);
             }
         }
         catch (error) {
@@ -3898,7 +4106,7 @@ export class AiService extends TracedEventEmitter {
      * commitId is captured before newChat() clears the session.
      * Failures are logged but do not fail clear chat.
      */
-    async uploadEmptyContextToBucketeer(userId, commitId) {
+    async uploadEmptyContextToBucketeer(commitId) {
         const jwt = this.getInferenceJwt();
         const branchName = this._syncContextProvider?.getBranchName();
         if (!jwt || !commitId || !branchName) {
@@ -3910,7 +4118,6 @@ export class AiService extends TracedEventEmitter {
                 commitId,
                 branchName,
                 appRootDirPath: this.config.appRootDirPath,
-                userId,
                 useServerUpload: this.isAiFileUploadEnabled,
                 superblocksBaseUrl: this.config.superblocksBaseUrl,
                 logger: this.getLogger(),
@@ -4018,44 +4225,104 @@ export class AiService extends TracedEventEmitter {
             this.getLogger().warn(`[handleRestoreToCommit] No JWT available, cannot restore or truncate context`);
             return;
         }
-        // 1. Try to restore context from checkpoint
-        let restored = false;
+        // 1. Roll the conversation back to the checkpoint. The local event log is
+        // the source of truth (see restoreConversationToCheckpoint); the remote
+        // checkpoint archive is only a fallback for a recycled pod, because a
+        // post-checkpoint flush may have overwritten it. The orchestrator returns
+        // an outcome and does NOT touch the metric — this method is the single
+        // place restore outcomes are recorded, so every exit is counted exactly
+        // once.
+        let outcome;
         try {
             const contextId = getContextId(this.clark, this.config);
-            const result = await restoreContextFromCheckpoint({
-                applicationId,
-                commitId,
-                branchName,
+            outcome = await restoreConversationToCheckpoint({
+                contextManagerV2: this.contextManagerV2,
                 contextId,
-                appRootDirPath: this.config.appRootDirPath,
-                superblocksBaseUrl: this.config.superblocksBaseUrl,
+                sinceTimestamp,
                 logger: this.getLogger(),
-                jwt,
-                contextManagerV2: this.contextManagerV2,
+                coldRestore: async () => {
+                    const result = await restoreContextFromCheckpoint({
+                        applicationId,
+                        commitId,
+                        branchName,
+                        contextId,
+                        appRootDirPath: this.config.appRootDirPath,
+                        superblocksBaseUrl: this.config.superblocksBaseUrl,
+                        logger: this.getLogger(),
+                        jwt,
+                        contextManagerV2: this.contextManagerV2,
+                    });
+                    if (!result.success) {
+                        const reason = result.savedFiles.length > 0
+                            ? "main.json not found in checkpoint"
+                            : "no checkpoint found";
+                        this.getLogger().info(`[handleRestoreToCommit] Context restore failed: ${reason} for commitId=${commitId}, preserving chat messages`);
+                    }
+                    return result.success;
+                },
+                republish: async () => {
+                    await uploadContextToBucketeer({
+                        applicationId,
+                        commitId,
+                        branchName,
+                        appRootDirPath: this.config.appRootDirPath,
+                        useServerUpload: this.isAiFileUploadEnabled,
+                        superblocksBaseUrl: this.config.superblocksBaseUrl,
+                        logger: this.getLogger(),
+                        jwt,
+                        // The heal may have deleted the legacy snapshot, leaving an
+                        // events-only archive that IS the rolled-back state; upload it so the
+                        // remote /context/latest is healed (the download rebuilds the record
+                        // from the truncated event log) instead of staying poisoned.
+                        allowEventsOnlyArchive: true,
+                    });
+                },
             });
-            if (result.success) {
-                restored = true;
-            }
-            else {
-                const reason = result.savedFiles.length > 0
-                    ? "main.json not found in checkpoint"
-                    : "no checkpoint found";
-                this.getLogger().info(`[handleRestoreToCommit] Context restore failed: ${reason} for commitId=${commitId}, preserving chat messages`);
-            }
         }
         catch (error) {
-            this.getLogger().warn(`[handleRestoreToCommit] Error restoring context from checkpoint, preserving chat messages`, getErrorMeta(error));
-        }
-        if (!restored) {
+            // The orchestrator itself threw (e.g. a cold restore's scoped-token or
+            // network error propagating out of coldRestore). Count it so the restore
+            // metric reflects every attempt, not just the ones that return normally.
+            this.getLogger().error(`[handleRestoreToCommit] Error restoring context from checkpoint, preserving chat messages`, getErrorMeta(error));
+            recordCheckpointRestore("aborted");
             return;
         }
-        // 2. Checkpoint restored — truncate server chat to match
-        try {
-            await this.chatSessionStore.truncateAfterTimestamp(sinceTimestamp);
-            this.getLogger().info(`[handleRestoreToCommit] Successfully restored context and truncated chat for commitId=${commitId}`);
+        if (outcome !== "rolled_back") {
+            // Context was NOT rolled back (aborted / unhealed / no_checkpoint). Record
+            // the outcome and leave the chat intact so it keeps matching the context.
+            recordCheckpointRestore(outcome);
+            this.getLogger().info(`[handleRestoreToCommit] No context restored (${outcome}) for commitId=${commitId}; leaving chat intact`, { applicationId, branchName });
+            return;
         }
-        catch (error) {
-            this.getLogger().warn(`[handleRestoreToCommit] Context restored but failed to truncate chat for commitId=${commitId}`, getErrorMeta(error));
+        // 2. Context is rolled back — truncate the server chat to match. This is a
+        // cross-service RPC and is NOT transactional with the context rollback
+        // above, so a failure here leaves the context rolled back while the chat
+        // still shows the reverted turns (a chat-ahead-of-context divergence).
+        // Retry a few times to ride out a transient blip; if it still fails,
+        // surface it at error level with a metric so the divergence is alertable,
+        // not buried.
+        // truncateAfterTimestamp is idempotent (same cutoff → same result), so the
+        // retries are safe.
+        const CHAT_TRUNCATE_ATTEMPTS = 3;
+        for (let attempt = 1; attempt <= CHAT_TRUNCATE_ATTEMPTS; attempt++) {
+            try {
+                await this.chatSessionStore.truncateAfterTimestamp(sinceTimestamp);
+                this.getLogger().info(`[handleRestoreToCommit] Successfully restored context and truncated chat for commitId=${commitId}`);
+                recordCheckpointRestore("restored");
+                return;
+            }
+            catch (error) {
+                if (attempt < CHAT_TRUNCATE_ATTEMPTS) {
+                    this.getLogger().warn(`[handleRestoreToCommit] Chat truncation attempt ${attempt}/${CHAT_TRUNCATE_ATTEMPTS} failed for commitId=${commitId}, retrying`, getErrorMeta(error));
+                    await new Promise((resolve) => setTimeout(resolve, 200 * attempt));
+                    continue;
+                }
+                // Exhausted retries: context is rolled back but the chat is not. Log at
+                // error (a real, if less-harmful, chat/context divergence) and count it.
+                this.getLogger().error(`[handleRestoreToCommit] Context rolled back but chat truncation failed after ${CHAT_TRUNCATE_ATTEMPTS} attempts for commitId=${commitId} ` +
+                    `(applicationId=${applicationId}, branchName=${branchName}); chat may show reverted turns until the next turn resyncs it`, getErrorMeta(error));
+                recordCheckpointRestore("chat_truncate_failed");
+            }
         }
     }
     async handleApiUpdate(update) {