@superblocksteam/vite-plugin-file-sync 2.0.156-next.3 → 2.0.156

package/dist/ai-service/index.js

@@ -2,9 +2,9 @@ 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 { hasToolCall, stepCountIs, } from "ai";
+import { hasToolCall, streamText as aiStreamText, } from "ai";
 import { z } from "zod";
-import { AiGenerationState, AiMode, resolveAttachmentFileName, } from "@superblocksteam/library-shared/types";
+import { AiGenerationState, AiMode, MAX_OUTPUT_TOKENS, resolveAttachmentFileName, } from "@superblocksteam/library-shared/types";
 import { isToolPermissionClarification } from "@superblocksteam/library-shared/types/tool-permission-response";
 import { addTracingToMethods, AiSystemTriggerType, classifyAttachmentDelivery, FactAccessType, FactCreatedSource, FactEntityScope, TracedEventEmitter, getAiQuotaPaywallReasonFromMessage, } from "@superblocksteam/shared";
 import { npmRegistryEmitter } from "@superblocksteam/telemetry";
@@ -14,8 +14,6 @@ import { POLICY_GATE_RUN_IN_PROGRESS_ERROR, policyGateRunRegistry, } from "../po
 import { getErrorMeta, getLogger, setDefaultLogger } from "../util/logger.js";
 import { composeSecurityAgentPrompt, } from "./agent/prompts/build-security-scan-prompt.js";
 import { getToolCallArguments, getToolErrorCode, getToolOutput, } from "./agent/tool-message-utils.js";
-import { recordToolPermissionDeliveryOutcome } from "./agent/tool-permission-delivery-metrics.js";
-import { createExpiredPermissionPromptDeliveryError, flushPendingToolPermissionPrompt, inFlightToolPermissionReplay, } from "./agent/tool-permission-delivery.js";
 import { buildSecurityScanTools, isBackgroundTaskFrameworkEnabled, } from "./agent/tools.js";
 import { hasPendingEscalations, flushPendingEscalations, } from "./agent/tools/apis/api-validation-orchestrator.js";
 import { safeSampleJson } from "./agent/tools/apis/sample-json.js";
@@ -24,7 +22,7 @@ import { applyDevDatabaseMigrationsTaskMetadata, createDevDatabaseTaskTypes, wir
 import { debugCache } from "./agent/tools/debug-cache.js";
 import { securityScanResultSchema, } from "./agent/tools/report-security-findings.js";
 import { SessionEntityPermissionStore, } from "./agent/tools2/entity-permissions.js";
-import { ToolRegistry, } from "./agent/tools2/registry.js";
+import { ToolRegistry } from "./agent/tools2/registry.js";
 import { FileSystemInterface } from "./app-interface/file-system-interface.js";
 import { createDefaultVirtualFileSystem, createPathValidator, } from "./app-interface/filesystem/index.js";
 import { NpmPackageLookup } from "./app-interface/npm-package-lookup.js";
@@ -52,7 +50,6 @@ import { LiveEditEventForwarder, readForwarderConfigFromEnv, } from "./live-edit
 import { createChaosFetch } from "./llm/chaos-fetch.js";
 import { createClarkStreamTelemetrySink } from "./llm/clark-stream-telemetry-sink.js";
 import { LLMClient } from "./llm/client.js";
-import { resolveContextOptionsV2 } from "./llm/context-v2/config.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";
@@ -60,7 +57,8 @@ 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 { getJwtTraceTags, traceLLM } from "./llmobs/helpers.js";
+import { traceLLM } from "./llmobs/helpers.js";
+import { getJwtTraceTags } from "./llmobs/helpers.js";
 import llmobs from "./llmobs/index.js";
 import { PlaywrightMcpServerManager } from "./mcp/playwright-server.js";
 import { evaluateMigrationAutoDispatchReadiness } from "./migration-auto-dispatch-readiness.js";
@@ -89,10 +87,10 @@ import { doRuntimeReviewing } from "./state-machine/handlers/runtime-reviewing.j
 import { classificationHelper } from "./state-machine/helpers/classification.js";
 import { getContextId } from "./state-machine/helpers/context-id.js";
 import { bindHeadlessTurnLifecycle } from "./state-machine/helpers/headless-turn-lifecycle.js";
-import { clearPendingToolPermissionRequest, sendUserGenerationStateChannel, sendUserMessageChannel, } from "./state-machine/helpers/peer.js";
-import { finalizePolicyGateUsageMetrics, policyGateFailureFromStream, recordPolicyGateAttempt, } from "./state-machine/helpers/policy-gate-usage-metrics.js";
+import { sendUserGenerationStateChannel } from "./state-machine/helpers/peer.js";
+import { sendUserMessageChannel, clearPendingToolPermissionRequest, } from "./state-machine/helpers/peer.js";
+import { collectPolicyGateLlmUsage, finalizePolicyGateUsageMetrics, policyGateFailureFromStream, recordPolicyGateAttempt, } from "./state-machine/helpers/policy-gate-usage-metrics.js";
 import { finalizePromptBillingMetrics } from "./state-machine/helpers/prompt-billing-metrics.js";
-import { setPromptDegradationMetricIdentity } from "./state-machine/helpers/prompt-degradation-metrics.js";
 import { paywallStopAdvice } from "./state-machine/helpers/prompt-failure.js";
 import { CLARK_USAGE_ORIGIN, promptFailureFromStreamAbort, promptFailurePaywall, promptFailurePlatform, promptFailureUserCanceled, PROMPT_ATTEMPT_ENTRYPOINT, 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";
@@ -102,7 +100,6 @@ import { TemplateRenderer } from "./template-renderer.js";
 import { describeMigrationInferenceJwt, isUsableMigrationInferenceJwt, } from "./util/inference-jwt.js";
 import { buildModelConfigAttributes, processLLMConfig, } from "./util/llm-config-utils.js";
 import { parseJwt } from "./util/parse-jwt.js";
-import { PERMISSION_PROMPT_DELIVERY_TIMEOUT_MS } from "./util/rpc-timeout.js";
 import { getToolCallSignature } from "./util/tool-signature.js";
 export { AiServiceFeatureFlags } from "./features.js";
 export { getManagedNativeDbGateDenialReason, hasEditLifecycleCapability, isManagedNativeDbAvailableForClark, } from "./managed-native-db-availability.js";
@@ -145,9 +142,6 @@ 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;
-// Ceiling above the historical 30-step cap; 100 is an existing policy-agent
-// histogram bucket. Tighten from production p99 if scans complete far below this.
-const POLICY_GATE_MAX_STEPS = 100;
 // 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]";
@@ -360,9 +354,6 @@ const POLICY_GATE_USAGE_METADATA = {
     origin: CLARK_USAGE_ORIGIN.policy_gate,
     surface: "policy_gate",
 };
-const ERROR_JSON_OUTPUT_TYPE = "error-json";
-const TOOL_MESSAGE_ROLE = "tool";
-const TOOL_RESULT_CONTENT_TYPE = "tool-result";
 export class AiService extends TracedEventEmitter {
     config;
     templateRenderer;
@@ -505,6 +496,10 @@ export class AiService extends TracedEventEmitter {
         this._runtimeSyncService = syncService;
         this._runtimeGitOperationLock = gitOperationLock;
     }
+    /**
+     * Sets the sync context provider used by the remote context restore fallback.
+     * Called by the socket manager once the sync service is initialized.
+     */
     setSyncContextProvider(provider) {
         this._syncContextProvider = provider;
     }
@@ -569,6 +564,7 @@ export class AiService extends TracedEventEmitter {
         this.tryAutoDispatchMigrationTurnIfReady();
         return true;
     }
+    // ===== Recording API =====
     /**
      * Check if recording is currently enabled.
      */
@@ -625,10 +621,6 @@ export class AiService extends TracedEventEmitter {
             applicationId: config.applicationId,
             organizationId: config.organizationId,
         });
-        setPromptDegradationMetricIdentity({
-            applicationId: config.applicationId,
-            organizationId: config.organizationId,
-        });
         this.startLiveEditEventForwarding();
         // Store the wrapped logger so getLogger() can return it
         // The wrapped logger forwards logs to the browser console when Clark logs are enabled
@@ -912,12 +904,6 @@ export class AiService extends TracedEventEmitter {
             organizationId: config.organizationId,
             applicationId: config.applicationId,
             getJwt: () => this.requireProviderInferenceAuthToken(),
-            stepTokenUsageLabels: () => ({
-                coreModel: this.clark?.context.llmConfig?.coreModel,
-                mode: this.clark?.context.currentMode === AiMode.PLAN ? "plan" : "build",
-                surface: this.clark?.context.usageMetadata?.surface ??
-                    APP_CLARK_GENERATION_METADATA.surface,
-            }),
             onUsageRecorded: (usage) => {
                 // Accumulate workflow-level metrics
                 this.clark.updateContext((context) => {
@@ -951,7 +937,9 @@ export class AiService extends TracedEventEmitter {
             createTagProvider: this.createClarkTagProviderFactory(),
             sessionId: this.sessionId,
         });
+        // Initialize judge service
         this.mcpServerManager = new PlaywrightMcpServerManager(this.clark.logger);
+        // Initialize recording manager for capturing AI interactions
         this.recordingManager = new RecordingManager({
             recordingsDir: path.join(config.appRootDirPath, ".superblocks", "recordings"),
             appRootDirPath: config.appRootDirPath,
@@ -1319,7 +1307,6 @@ export class AiService extends TracedEventEmitter {
                     llmClient: this.llmClient,
                     toolRegistry: this.toolRegistry,
                     entityPermissionStore: this.entityPermissionStore,
-                    onToolPermissionParked: (messageId) => this.scheduleToolPermissionReplayWindow(messageId),
                     mcpServerManager: this.mcpServerManager,
                     viteServer: this.viteServer,
                     recordingManager: this.recordingManager,
@@ -2272,6 +2259,7 @@ export class AiService extends TracedEventEmitter {
                     useMockGeneration: request.mockGeneration ?? false,
                     planContext: mergePlanContexts(context.planContext, request.planContext),
                     browserContext: request.browserContext,
+                    // Initialize workflow metrics tracking
                     workflowMetrics: {
                         toolCallCount: 0,
                         stepCount: 0,
@@ -2621,6 +2609,7 @@ export class AiService extends TracedEventEmitter {
         const timeoutController = new AbortController();
         const timeoutId = setTimeout(() => timeoutController.abort(), timeout);
         try {
+            // Build the appropriate content part based on attachment type
             const contentPart = this.attachmentToContentPart(attachment);
             // Route labeling through the shared LLMClient (APPS-5061) as a stateless
             // one-shot, mirroring the APPS-1213 (#20263) migrations. No safety config:
@@ -2674,6 +2663,7 @@ export class AiService extends TracedEventEmitter {
             };
         }
         else if (attachment.type === "pdf" || attachment.type === "archive") {
+            // Parse data URL to extract base64 data
             const match = attachment.data.match(/^data:([^;]+);base64,(.+)$/);
             const data = match?.[2] ?? attachment.data;
             const defaultMediaType = attachment.type === "pdf"
@@ -2773,10 +2763,6 @@ export class AiService extends TracedEventEmitter {
         await this.clark.recordSystemMessage(params, "Environment event", "User closed the integration setup form. The integration setup might not have been completed successfully.");
     }
     async handleUserCanceled() {
-        // Drop parked/in-flight permission prompts first so interval replay and
-        // grace expiry cannot keep pushing the dialog or inject a continuation
-        // after the user has already stopped the turn.
-        clearPendingToolPermissionRequest(this.clark, "user-canceled");
         const awaitingUser = this.isAwaitingUser();
         if (!this.isBusy() && !awaitingUser) {
             return;
@@ -3035,9 +3021,11 @@ export class AiService extends TracedEventEmitter {
             logger.warn(`[ai-service] Skip request for testRunId=${testRunId} but current is ${this.clark.context.currentTestRunId}`);
             return;
         }
+        // Update context with skip request
         this.clark.updateContext({
             skipRequestedForTestCase: testCaseId,
         });
+        // Update the test case status to skipped
         const testCaseStatuses = this.clark.context.testCaseStatuses ??
             new Map();
         testCaseStatuses.set(testCaseId, "skipped");
@@ -3058,127 +3046,24 @@ export class AiService extends TracedEventEmitter {
         }
         logger.info(`[ai-service] Test case ${testCaseId} marked as skipped, flag set for subagent`);
     }
-    scheduleToolPermissionReplayWindow(messageId) {
-        const pending = this.clark.context.pendingToolPermissionRequest;
-        if (pending?.messageId !== messageId ||
-            pending.delivery.status !== "parked") {
-            return;
-        }
-        this.scheduleToolPermissionExpiry(messageId);
-        this.scheduleToolPermissionReplay(messageId);
-    }
-    scheduleToolPermissionExpiry(messageId) {
-        const pending = this.clark.context.pendingToolPermissionRequest;
-        if (pending?.messageId !== messageId ||
-            pending.delivery.status !== "parked") {
-            return;
-        }
-        const expiresInMs = Math.max(0, pending.delivery.expiresAtMs - Date.now());
-        setTimeout(() => {
-            void (async () => {
-                await inFlightToolPermissionReplay(this.clark);
-                const current = this.clark.context.pendingToolPermissionRequest;
-                if (current?.messageId !== messageId ||
-                    current.delivery.status !== "parked") {
-                    return;
-                }
-                if (Date.now() < current.delivery.expiresAtMs) {
-                    this.scheduleToolPermissionExpiry(messageId);
-                    return;
-                }
-                if (!(await this.waitUntilToolPermissionContinuationAllowed(messageId))) {
-                    return;
-                }
-                const expired = this.clark.context.pendingToolPermissionRequest;
-                if (expired?.messageId !== messageId ||
-                    expired.delivery.status !== "parked") {
-                    return;
-                }
-                const error = createExpiredPermissionPromptDeliveryError(expired.toolName, expired.delivery);
-                recordToolPermissionDeliveryOutcome("expired");
-                await this.handleToolPermissionOutcome({
-                    kind: "delivery_error",
-                    messageId,
-                    error,
-                });
-            })().catch((cause) => {
-                this.getLogger().error("[ai-service] Failed to continue after permission delivery expired", getErrorMeta(cause));
-            });
-        }, expiresInMs);
-    }
     /**
-     * `userSentPrompt` is only valid from Idle or AwaitingUser. A parked
-     * permission can expire while a sibling tool is still generating; wait
-     * until the continuation can actually run, or abort if the request is gone.
+     * Handles the user's response to a tool permission request.
+     *
+     * If approved, immediately executes the tool with the original input,
+     * injects the tool call and result into the conversation context,
+     * and restarts the LLM with the updated context.
+     *
+     * If denied, injects a tool error into the context and restarts the LLM.
      */
-    async waitUntilToolPermissionContinuationAllowed(messageId) {
-        const TOOL_PERMISSION_CONTINUATION_POLL_MS = 250;
-        for (;;) {
-            const current = this.clark.context.pendingToolPermissionRequest;
-            if (current?.messageId !== messageId ||
-                current.delivery.status !== "parked") {
-                return false;
-            }
-            switch (this.clark.state) {
-                case ClarkStateNames.Idle:
-                case ClarkStateNames.AwaitingUser:
-                    return true;
-                case ClarkStateNames.Dead:
-                    return false;
-                case ClarkStateNames.AgentPlanning:
-                case ClarkStateNames.LLMGenerating:
-                case ClarkStateNames.PostProcessing:
-                case ClarkStateNames.RuntimeReviewing: {
-                    await new Promise((resolve) => {
-                        setTimeout(resolve, TOOL_PERMISSION_CONTINUATION_POLL_MS);
-                    });
-                    break;
-                }
-                default: {
-                    const exhaustive = this.clark.state;
-                    return exhaustive;
-                }
-            }
-        }
-    }
-    scheduleToolPermissionReplay(messageId) {
-        setTimeout(() => {
-            const pending = this.clark.context.pendingToolPermissionRequest;
-            if (pending?.messageId !== messageId ||
-                pending.delivery.status !== "parked" ||
-                Date.now() >= pending.delivery.expiresAtMs) {
-                return;
-            }
-            const replay = this.stablePeer.isConnected
-                ? flushPendingToolPermissionPrompt(this.clark, { source: "retry" })
-                : Promise.resolve();
-            void replay.finally(() => {
-                const current = this.clark.context.pendingToolPermissionRequest;
-                if (current?.messageId === messageId &&
-                    current.delivery.status === "parked") {
-                    this.scheduleToolPermissionReplay(messageId);
-                }
-            });
-        }, PERMISSION_PROMPT_DELIVERY_TIMEOUT_MS);
-    }
     async handleAiRespondToToolPermission(request) {
-        return this.handleToolPermissionOutcome({
-            kind: "response",
-            ...request,
-        });
-    }
-    async handleToolPermissionOutcome(request) {
         const logger = this.getLogger();
-        const approved = request.kind === "response" && request.approved;
-        logger.info(`[ai-service] handleAiRespondToToolPermission: messageId=${request.messageId}, approved=${approved}`, { outcomeKind: request.kind });
+        logger.info(`[ai-service] handleAiRespondToToolPermission: messageId=${request.messageId}, approved=${request.approved}`);
         // a reconnect can cause the client to re-send this response. If
         // we already applied a response with this requestId, treat the repeat as a
         // no-op so the tool is not executed (or denied) twice for the same user
         // decision. Only a peek here — the id is recorded once the response is
         // accepted below.
-        if (request.kind === "response" &&
-            request.requestId &&
-            this.requestDeduplicator.has(request.requestId)) {
+        if (request.requestId && this.requestDeduplicator.has(request.requestId)) {
             logger.info("[ai-service] Ignoring duplicate tool-permission re-send", {
                 requestId: request.requestId,
             });
@@ -3198,17 +3083,14 @@ export class AiService extends TracedEventEmitter {
         // and honor the result so a re-send dropped here can't be applied twice. A
         // duplicate that loses the atomic `markIfNew` bails before the tool runs.
         // (APPS-4663)
-        if (request.kind === "response" &&
-            request.requestId &&
+        if (request.requestId &&
             !this.requestDeduplicator.markIfNew(request.requestId)) {
             logger.warn("[ai-service] Ignoring duplicate tool-permission re-send (lost dedup race)", { requestId: request.requestId });
             return;
         }
         logger.info(`[ai-service] Clearing pending request`);
         // Clear the pending request and notify UI
-        clearPendingToolPermissionRequest(this.clark, request.kind === "delivery_error"
-            ? "tool-permission-delivery-expired"
-            : "tool-permission-response");
+        clearPendingToolPermissionRequest(this.clark, "tool-permission-response");
         const toolCallId = pendingRequest.messageId;
         const toolName = pendingRequest.toolName;
         const rawInput = pendingRequest.rawInput;
@@ -3216,6 +3098,7 @@ export class AiService extends TracedEventEmitter {
         const llmConfig = this.clark.context.llmConfig;
         logger.info(`[ai-service] Getting context toolCallId=${toolCallId})`);
         const sendUserMessage = sendUserMessageChannel(this.clark, this.chatSessionStore);
+        // Build the assistant message with the NEW tool call ID
         const assistantMessage = {
             role: "assistant",
             content: [
@@ -3227,8 +3110,9 @@ export class AiService extends TracedEventEmitter {
                 },
             ],
         };
+        // Build the tool result message based on approval/execution
         let toolResultMessage;
-        if (approved) {
+        if (request.approved) {
             // Approve in entity permission store
             this.entityPermissionStore.approveToolCall(toolName, rawInput);
             const signature = getToolCallSignature(toolName, rawInput);
@@ -3244,11 +3128,9 @@ export class AiService extends TracedEventEmitter {
                 signature,
             });
             // Store always-allow permissions
-            if (request.kind === "response") {
-                request.alwaysAllowPermissions?.forEach((permission) => {
-                    void this.entityPermissionStore.approve(permission.entityType, permission.entityId, permission.permissionType);
-                });
-            }
+            request.alwaysAllowPermissions?.forEach((permission) => {
+                void this.entityPermissionStore.approve(permission.entityType, permission.entityId, permission.permissionType);
+            });
             // Execute the tool with proper access control context from Clark
             const accessControlContext = this.clark.getAccessControl(this.config.features);
             logger.info(`[ai-service] Executing approved tool: ${toolName}`);
@@ -3340,23 +3222,6 @@ export class AiService extends TracedEventEmitter {
                 });
             }
         }
-        else if (request.kind === "delivery_error") {
-            toolResultMessage = {
-                role: TOOL_MESSAGE_ROLE,
-                content: [
-                    {
-                        type: TOOL_RESULT_CONTENT_TYPE,
-                        toolCallId,
-                        toolName,
-                        output: {
-                            type: ERROR_JSON_OUTPUT_TYPE,
-                            value: { error: request.error.message },
-                        },
-                        isError: true,
-                    },
-                ],
-            };
-        }
         else {
             const denyReason = request.denyMessage || "User denied permission for this tool call";
             if (toolName === "testApi") {
@@ -3431,11 +3296,10 @@ export class AiService extends TracedEventEmitter {
         if (!this.clark.context.llmConfig) {
             throw new Error("Cannot continue after tool permission: llmConfig is missing from context");
         }
-        const isClarification = isToolPermissionClarification(approved, request.kind === "response" ? request.denyMessage : undefined);
+        const isClarification = isToolPermissionClarification(request.approved, request.denyMessage);
         logger.info("[ai-service] tool_permission_response continuation", {
-            approved,
+            approved: request.approved,
             clarification: isClarification,
-            outcomeKind: request.kind,
         });
         trace
             .getActiveSpan()
@@ -3456,9 +3320,9 @@ export class AiService extends TracedEventEmitter {
             outcomeDedupeKey: promptOutcomeDedupeKey,
         });
         const continueRequest = {
-            prompt: approved || request.kind === "delivery_error"
+            prompt: request.approved
                 ? ``
-                : isClarification && request.denyMessage
+                : isClarification
                     ? request.denyMessage.trim()
                     : this.buildDenialContinuationPrompt(toolName, rawInput),
             promptContext: {
@@ -3473,13 +3337,8 @@ export class AiService extends TracedEventEmitter {
             responseMetadata: {
                 type: "tool_permission_response",
                 responseToMessageId: request.messageId,
-                approved,
-                outcome: request.kind === "delivery_error"
-                    ? "delivery_expired"
-                    : "user_response",
-                alwaysAllowPermissions: request.kind === "response"
-                    ? request.alwaysAllowPermissions
-                    : undefined,
+                approved: request.approved,
+                alwaysAllowPermissions: request.alwaysAllowPermissions,
                 ...(isClarification ? { clarification: true } : {}),
             },
         };
@@ -3557,9 +3416,9 @@ export class AiService extends TracedEventEmitter {
      * Returned so the caller can be told which turn it now owns; a turn started
      * without one is reachable by this socket alone.
      */
-    beginHeadlessTurn(peer, peerId, turnId, liveness) {
+    beginHeadlessTurn(peer, peerId, turnId) {
         this.getLogger().info(`[ai-service] headless turn started by gateway peer ${peerId}`);
-        this.stablePeer.beginTurn({ liveness, peer, peerId, turnId });
+        this.stablePeer.beginTurn({ peer, peerId, turnId });
         return turnId;
     }
     /**
@@ -3571,9 +3430,8 @@ export class AiService extends TracedEventEmitter {
      * the caller already owns the turn, or is talking about a turn that is not
      * this one — are silent. Only an actual handover is worth a line.
      */
-    claimHeadlessTurn(peer, peerId, turnId, liveness) {
-        if (this.stablePeer.claimTurn({ liveness, peer, peerId, turnId }) !==
-            "claimed") {
+    claimHeadlessTurn(peer, peerId, turnId) {
+        if (this.stablePeer.claimTurn({ peer, peerId, turnId }) !== "claimed") {
             return;
         }
         this.getLogger().info(`[ai-service] headless turn moved to reattached gateway peer ${peerId}`);
@@ -3594,19 +3452,14 @@ export class AiService extends TracedEventEmitter {
      * it, so a socket that a reattach already replaced cannot strip routing from
      * its successor.
      */
-    endHeadlessTurn(peerId, turnIdOrLiveness) {
-        const turnId = typeof turnIdOrLiveness === "string" ? turnIdOrLiveness : undefined;
-        const liveness = typeof turnIdOrLiveness === "string" ? undefined : turnIdOrLiveness;
-        if (this.stablePeer.turnPeerId !== peerId ||
-            (turnId !== undefined && this.stablePeer.turnId !== turnId) ||
-            (liveness &&
-                this.stablePeer.turnLivenessGeneration !== liveness.generation)) {
+    endHeadlessTurn(peerId) {
+        if (this.stablePeer.turnPeerId !== peerId) {
             return;
         }
         this.getLogger().info(`[ai-service] headless turn peer ${peerId} detached; falling back to the editor binding`);
         this.stablePeer.releaseTurnPeer(peerId);
     }
-    handleUserConnected(peer, peerId, liveness) {
+    handleUserConnected(peer, peerId) {
         const existingPeerId = this.stablePeer.peerId;
         this.getLogger().info(`[ai-service] handleUserConnected: peerId=${peerId}`);
         if (existingPeerId) {
@@ -3623,7 +3476,7 @@ export class AiService extends TracedEventEmitter {
         // `clark.context.peer` is a permanent StablePeer proxy; we swap its
         // inner (and peerId) here. Any queued calls accumulated while
         // disconnected flush in order as part of setInner.
-        this.stablePeer.setInner({ liveness, peer, peerId });
+        this.stablePeer.setInner({ peer, peerId });
         this.getLogger().info(`[ai-service] Peer set: ${peerId}`);
         // A replacement (or freshly reconnecting) peer must re-prove its
         // identity via `notifyAuthenticatedRpc` before live task pushes
@@ -3651,18 +3504,25 @@ export class AiService extends TracedEventEmitter {
                 messageId: request.messageId,
                 toolName: request.toolName,
             });
-            // replaceInFlight also covers an in-progress first delivery
-            // (awaiting_ack). The replacement takes park/ack ownership so a
-            // miss while the tab is still loading still starts the replay
-            // grace window instead of wedging the session.
-            const needsReplayWindow = request.delivery.status === "acked";
-            void flushPendingToolPermissionPrompt(this.clark, {
-                replaceInFlight: true,
-                source: "reconnect",
-            }).then((outcome) => {
-                if (needsReplayWindow && outcome === "parked") {
-                    this.scheduleToolPermissionReplayWindow(request.messageId);
-                }
+            // Note: The promise from the original RPC call will have timed out,
+            // but resending allows the user to see and respond to the request.
+            // The tool execution will fail, but the agent can retry it naturally.
+            void peer.call
+                .aiRequestToolPermission({
+                messageId: request.messageId,
+                toolName: request.toolName,
+                actionName: request.actionName ?? undefined,
+                formattedInput: request.formattedInput ?? undefined,
+                rawInput: request.rawInput,
+                permissions: request.permissions ?? [],
+            })
+                .catch((error) => {
+                this.getLogger().error(`[ai-service] Failed to resend tool permission request: ${error}`);
+                this.getLogger().warn(`[debug] reconnect: resend pending tool permission failed`, {
+                    messageId: request.messageId,
+                    toolName: request.toolName,
+                    error: error instanceof Error ? error.message : String(error),
+                });
             });
         }
         else {
@@ -3761,7 +3621,7 @@ export class AiService extends TracedEventEmitter {
      * will never answer — the case a headless build leaves behind when its
      * gateway socket closes over an open editor tab.
      */
-    handleUserDisconnected(peerId, replacement, liveness) {
+    handleUserDisconnected(peerId, replacement) {
         const currentPeerId = this.stablePeer.peerId;
         this.getLogger().info(`[ai-service] handleUserDisconnected: peerId=${peerId}`);
         if (!currentPeerId) {
@@ -3771,19 +3631,6 @@ export class AiService extends TracedEventEmitter {
             this.getLogger().info(`[ai-service] Stale disconnect ignored: peer ${peerId} disconnecting but current peer is ${currentPeerId}`);
             return; // Don't clear the current peer on stale disconnect
         }
-        if (liveness &&
-            this.stablePeer.editorLivenessGeneration !== liveness.generation) {
-            this.getLogger().info(`[ai-service] Stale disconnect ignored for replaced connection generation`);
-            return;
-        }
-        if (replacement?.peerId === peerId &&
-            replacement.liveness &&
-            liveness &&
-            replacement.liveness.generation !== liveness.generation) {
-            this.getLogger().info(`[ai-service] Rebinding peer ${peerId} to surviving editor connection`);
-            this.handleUserConnected(replacement.peer, replacement.peerId, replacement.liveness);
-            return;
-        }
         this.getLogger().info(`[ai-service] peer disconnecting: ${peerId}`);
         // Drop any deferred state-push registration for this peer so a
         // disconnect that races a never-arrived auth proof doesn't leave
@@ -3803,11 +3650,11 @@ export class AiService extends TracedEventEmitter {
         if (contextId) {
             this.contextManagerV2.clearContext(contextId);
         }
-        if (replacement && (replacement.peerId !== peerId || liveness)) {
+        if (replacement && replacement.peerId !== peerId) {
             this.getLogger().info(`[ai-service] Rebinding peer ${peerId} to surviving editor ${replacement.peerId}`);
             // Same door a fresh tab comes through, so the survivor gets the state
             // resync and re-auth a newly bound peer is owed.
-            this.handleUserConnected(replacement.peer, replacement.peerId, replacement.liveness);
+            this.handleUserConnected(replacement.peer, replacement.peerId);
             return;
         }
         this.clearPeerConnection();
@@ -4091,12 +3938,12 @@ export class AiService extends TracedEventEmitter {
         let textDeltaChars = 0;
         let modelName;
         let promptHash;
-        // Hoisted so the outer catch can attribute a failed run to the same model as
-        // the attempt it belongs to. Usage arrives from the client's usage observer,
-        // which also flushes on abort, on error, and for compaction calls.
+        // Hoisted so the outer catch can attribute a failed run to the same model
+        // and provider as the attempt it belongs to, and can still price whatever
+        // tokens the stream consumed before it threw.
         let coreModel;
-        const llmUsage = [];
-        let hadFallback = false;
+        let providerName = "unknown";
+        let streamResult;
         let attemptRecorded = false;
         const scanRoot = typeof this.appShell.appRootDirPath === "string"
             ? this.appShell.appRootDirPath
@@ -4111,15 +3958,13 @@ export class AiService extends TracedEventEmitter {
             metricsRecorded = true;
             recordPolicyAgentExecution({
                 durationMs: Date.now() - startedAt,
-                llmUsage,
                 outcome,
                 stepCount,
                 toolCallCount: totalToolCalls(),
             });
         };
         const buildRuntimeMetadata = (finalFailureReason) => ({
-            diagnosticsVersion: 3,
-            diffFiles: request.scanScopeJson?.changedFiles ?? [],
+            diagnosticsVersion: 2,
             durationMs: Date.now() - startedAt,
             filePathsListed: Array.from(filePathsListed).slice(0, 100),
             filePathsRead: Array.from(filePathsRead).slice(0, 100),
@@ -4156,8 +4001,8 @@ export class AiService extends TracedEventEmitter {
                 trigger_type: "policy_gate",
             }));
             const { model, providerConfig } = llmProvider.modelForTask("generateBroadEdit");
-            const { model: summarizationModel, providerConfig: summarizationProviderConfig, } = llmProvider.modelForTask("summarizeMessages");
             modelName = typeof model === "string" ? model : model.modelId;
+            providerName = providerConfig?.provider ?? "unknown";
             coreModel = llmConfig?.coreModel;
             // Counted as soon as the model resolves, because everything below is
             // inside this try and so is guaranteed a terminal outcome from the catch.
@@ -4170,105 +4015,37 @@ export class AiService extends TracedEventEmitter {
                 name: i.name,
                 pluginType: i.plugin?.id ?? "unknown",
             }));
-            let scanScopeJson = request.scanScopeJson;
-            if (scanScopeJson?.unifiedDiff) {
-                try {
-                    const scanResult = await scanContentForSecrets(scanScopeJson.unifiedDiff);
-                    if (scanResult.success && scanResult.hasSecrets) {
-                        scanScopeJson = {
-                            ...scanScopeJson,
-                            unifiedDiff: new SecretRedactor(scanResult.findings).redact(scanScopeJson.unifiedDiff),
-                        };
-                    }
-                    else if (!scanResult.success) {
-                        logger.warn("[secret-scan] Secret scan failed on policy gate diff; using the original diff");
-                    }
-                }
-                catch {
-                    logger.warn("[secret-scan] Secret scan threw on policy gate diff; using the original diff");
-                }
-            }
             const appContext = { integrations };
-            const systemPrompt = composeSecurityAgentPrompt(appContext, request.adminPrompt, scanScopeJson);
+            const systemPrompt = composeSecurityAgentPrompt(appContext, request.adminPrompt);
             promptHash = createHash("sha256").update(systemPrompt).digest("hex");
-            const scanContext = {};
-            const { headers, interleavedThinking, maxOutputTokens, providerOptions, thinkingBudgetTokens, thinkingEnabled, } = processLLMConfig(llmConfig);
-            // Scan-owned client: gate JWT, own usage sink, built first so scan tools
-            // (webFetch generateObject) use this client instead of Clark's.
-            const policyGateLlmClient = new LLMClient({
-                appShell: this.appShell,
-                applicationId: request.applicationId,
-                contextManagerV2: this.contextManagerV2,
-                getJwt: () => this.requireProvidedInferenceJwt(request.jwt, "policy gate"),
-                onUsageRecorded: (usage) => {
-                    if (usage.hadProviderFallback === true) {
-                        hadFallback = true;
-                    }
-                    llmUsage.push(...usage.promptCostUsageEntries);
-                },
-                organizationId: this.config.organizationId,
-                stepTokenUsageLabels: () => ({
-                    coreModel,
-                    mode: "build",
-                    surface: POLICY_GATE_USAGE_METADATA.surface,
-                }),
-                retryOptions: streamRetryOptionsFromConfig(this.config.retryOptions, this.config.features.streamTimeoutBudgets),
-                unifiedProviderEnabled: this.config.features.unifiedProviderEnabled ?? false,
-            });
             const scanToolsDeps = {
                 appShell: this.appShell,
                 features: this.config.features ?? {},
                 integrationStore: this.integrationStore,
-                llmClient: policyGateLlmClient,
+                llmClient: this.llmClient,
                 llmProvider: llmProvider,
-                storeSummary: (toolCallId, summary) => {
-                    scanContext.current?.storeSummary(toolCallId, summary);
-                },
             };
             const tools = await buildSecurityScanTools(this.clark, scanToolsDeps);
-            const result = await policyGateLlmClient.streamText({
-                abortSignal: AbortSignal.timeout(POLICY_GATE_MAX_TOTAL_TIME_MS),
-                context: {
-                    contextManagement: {
-                        summarizationModel,
-                        summarizationProviderConfig,
-                        type: "client",
+            const result = aiStreamText({
+                messages: [
+                    { content: systemPrompt, role: "system" },
+                    {
+                        content: "Begin your security review of this application now.",
+                        role: "user",
                     },
-                    // Product defaults: caller llmConfig cannot shrink the scan window to 0.
-                    contextOptionsV2: resolveContextOptionsV2(),
-                },
-                enableFallback: llmConfig?.unifiedProviderEnabled ?? false,
-                fallbackStrategy: llmConfig?.fallbackStrategy,
-                headers,
-                logging: false,
-                // processLLMConfig returns the same MAX_OUTPUT_TOKENS cap Clark uses.
-                maxOutputTokens,
+                ],
+                // Same output cap Clark uses (see processLLMConfig). Without it the
+                // provider default applies, which is model-dependent.
+                maxOutputTokens: MAX_OUTPUT_TOKENS,
                 model,
-                onContextReady: (ctx) => {
-                    scanContext.current = ctx;
-                },
-                providerConfig,
-                providerOptions: {
-                    ...providerOptions,
-                    ...(coreModel ? { unified: { coreModel } } : {}),
-                },
-                recordObsContent: false,
-                system: systemPrompt,
-                thinking: {
-                    budgetTokens: thinkingBudgetTokens,
-                    enabled: thinkingEnabled,
-                    interleaved: interleavedThinking,
-                },
+                // The unified router resolves the model from the request, falling back
+                // to Clark's core-model flag when absent.
+                ...(coreModel ? { providerOptions: { unified: { coreModel } } } : {}),
                 tools,
-                stopWhen: [
-                    hasToolCall("reportReviewRun"),
-                    stepCountIs(POLICY_GATE_MAX_STEPS),
-                ],
-                user: {
-                    content: "Begin your security review of this application now.",
-                    role: "user",
-                },
+                stopWhen: hasToolCall("reportReviewRun"),
+                timeout: { totalMs: POLICY_GATE_MAX_TOTAL_TIME_MS },
             });
+            streamResult = result;
             // 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.
@@ -4299,6 +4076,16 @@ export class AiService extends TracedEventEmitter {
                         : "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,
+                        });
+                    }
                     const textDelta = part.textDelta ??
                         part.text ??
                         part.delta;
@@ -4411,9 +4198,7 @@ export class AiService extends TracedEventEmitter {
             }
             catch (error) {
                 if (error instanceof Error &&
-                    (error.name === "AbortError" ||
-                        error.name === "AbortedError" ||
-                        error.name === "TimeoutError")) {
+                    (error.name === "AbortError" || error.name === "TimeoutError")) {
                     const elapsedMs = Date.now() - startedAt;
                     const isBudgetTimeout = error.name === "TimeoutError" ||
                         elapsedMs + POLICY_GATE_TIMEOUT_CLASSIFICATION_SLACK_MS >=
@@ -4433,18 +4218,6 @@ export class AiService extends TracedEventEmitter {
                     });
                     failPolicyGateExecution("Policy gate agent aborted", "aborted");
                 }
-                if (error instanceof Error &&
-                    (error.name === "LLMProviderError" ||
-                        error.name === "RetriesExhaustedError")) {
-                    const streamError = serializePolicyGateStreamError(error);
-                    if (streamErrors.length < MAX_POLICY_GATE_STREAM_ERRORS) {
-                        streamErrors.push(streamError);
-                    }
-                    logger.warn("[ai-service] Policy gate stream error", {
-                        ...policyRunLogContext,
-                        streamError,
-                    });
-                }
                 throw error;
             }
             if (!scanResult) {
@@ -4466,8 +4239,11 @@ export class AiService extends TracedEventEmitter {
                 aiAgentRunId: request.aiAgentRunId,
                 chatSessionStore: this.chatSessionStore,
                 coreModel,
-                hadFallback,
-                llmUsage,
+                llmUsage: await collectPolicyGateLlmUsage({
+                    modelName: modelName ?? "unknown",
+                    providerName,
+                    result,
+                }),
                 logger,
                 outcome: "success",
             });
@@ -4490,8 +4266,13 @@ export class AiService extends TracedEventEmitter {
                     chatSessionStore: this.chatSessionStore,
                     coreModel,
                     failure: policyGateFailureFromStream(streamErrors),
-                    hadFallback,
-                    llmUsage,
+                    llmUsage: streamResult
+                        ? await collectPolicyGateLlmUsage({
+                            modelName: modelName ?? "unknown",
+                            providerName,
+                            result: streamResult,
+                        })
+                        : [],
                     logger,
                     outcome: "failure",
                 });