@opengeni/worker-bundle 2.0.0 → 2.0.1-canary.35916963870001

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@opengeni/worker-bundle",
-  "version": "2.0.0",
+  "version": "2.0.1-canary.35916963870001",
   "description": "OpenGeni worker entry and reusable embedded lifecycle, shipped with a release-coherent pre-bundled Temporal workflow artifact.",
   "license": "Apache-2.0",
   "repository": {
@@ -54,24 +54,24 @@
   },
   "dependencies": {
     "@llamaindex/liteparse": "2.14.2",
-    "@opengeni/agent-proto": "^0.6.0",
-    "@opengeni/capabilities": "^0.3.4",
-    "@opengeni/codemode": "^0.6.0",
-    "@opengeni/codex": "^0.2.24",
-    "@opengeni/config": "^2.0.0",
-    "@opengeni/contracts": "^5.0.0",
-    "@opengeni/core": "^4.0.0",
-    "@opengeni/db": "^6.0.0",
-    "@opengeni/documents": "^0.8.31",
-    "@opengeni/events": "^0.4.29",
-    "@opengeni/github": "^0.7.14",
-    "@opengeni/network": "^0.3.1",
-    "@opengeni/observability": "^0.8.30",
-    "@opengeni/runtime": "^4.0.0",
-    "@opengeni/sdk": "^7.0.0",
-    "@opengeni/storage": "^0.2.131",
-    "@opengeni/tool-gateway": "^0.1.11",
-    "@opengeni/xai-subscription": "^0.1.4",
+    "@opengeni/agent-proto": "^0.6.0-canary.35916963870001",
+    "@opengeni/capabilities": "^0.3.4-canary.35916963870001",
+    "@opengeni/codemode": "^0.6.1-canary.35916963870001",
+    "@opengeni/codex": "^0.2.25-canary.35916963870001",
+    "@opengeni/config": "^2.1.0-canary.35916963870001",
+    "@opengeni/contracts": "^5.1.0-canary.35916963870001",
+    "@opengeni/core": "^4.0.1-canary.35916963870001",
+    "@opengeni/db": "^6.0.1-canary.35916963870001",
+    "@opengeni/documents": "^0.8.32-canary.35916963870001",
+    "@opengeni/events": "^0.4.30-canary.35916963870001",
+    "@opengeni/github": "^0.7.15-canary.35916963870001",
+    "@opengeni/network": "^0.3.1-canary.35916963870001",
+    "@opengeni/observability": "^0.8.31-canary.35916963870001",
+    "@opengeni/runtime": "^4.0.1-canary.35916963870001",
+    "@opengeni/sdk": "^7.1.0-canary.35916963870001",
+    "@opengeni/storage": "^0.2.132-canary.35916963870001",
+    "@opengeni/tool-gateway": "^0.1.12-canary.35916963870001",
+    "@opengeni/xai-subscription": "^0.1.4-canary.35916963870001",
     "@temporalio/activity": "^1.17.0",
     "@temporalio/client": "^1.17.0",
     "@temporalio/worker": "^1.17.0",

package/src/activities/agent-turn/agent-build.ts

@@ -79,6 +79,7 @@ import { resolveVideoReferenceSandboxAccess } from "./video-reference-sandbox";
 
 export type BuildTurnAgentDeps = {
   skillCatalog: NonNullable<BuildAgentOptions["skillCatalog"]>;
+  mcpServers: Settings["mcpServers"];
   input: RunAgentTurnInput;
   db: ActivityServices["db"];
   runtime: ActivityServices["runtime"];
@@ -139,6 +140,7 @@ export type BuildTurnAgentDeps = {
 
 export async function buildTurnAgent(deps: BuildTurnAgentDeps) {
   const {
+    mcpServers,
     input,
     db,
     runtime,
@@ -621,7 +623,9 @@ export async function buildTurnAgent(deps: BuildTurnAgentDeps) {
     const agentConstructionStartedAt = performance.now();
     let agentConstructionOutcome: "completed" | "failed" = "completed";
     try {
-      return runtime.buildAgent(eventing.modelRunSettings, runtimeResources, {
+      // Approval policy must use the same accepted account identities as tool
+      // preparation. Keep the separately resolved model/sandbox settings intact.
+      return runtime.buildAgent({ ...eventing.modelRunSettings, mcpServers }, runtimeResources, {
         ...(linkedToolAuthority
           ? {
               authorizeAttemptExecution: async () => {

package/src/activities/agent-turn/claim.ts

@@ -19,6 +19,7 @@ import { linkCurrentSpanToAdmission, turnExecutionTelemetryKey } from "@opengeni
 import { deliverChildRequiresActionToParent } from "../parent-wake";
 import {
   assertTurnExecutionPolicyMatchesConfigV1,
+  settingsForAcceptedSubscriptionTurn,
   resolveTurnExecutionPolicyV1,
   type Settings,
 } from "@opengeni/config";
@@ -277,7 +278,7 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
     gatewaySettings,
     claimedPolicy.kind === "valid" ? claimedPolicy.policy.productModelId : turn.model,
   );
-  const capabilitySettings = await settingsWithOrganizationProviderCredentials(
+  let capabilitySettings = await settingsWithOrganizationProviderCredentials(
     db,
     input.accountId,
     input.workspaceId,
@@ -287,7 +288,6 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   const codexAppsCredentialId = capabilitySettings.codexConnectedAppsEnabled
     ? await resolveCodexAppsCredentialIdForRun(db, input.workspaceId)
     : null;
-  runtime.configure(capabilitySettings);
   const policyForAbsent =
     claimedPolicy.kind === "valid"
       ? claimedPolicy.policy
@@ -304,6 +304,16 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   if (!installedPolicy.accepted) {
     throw new TurnAttemptFencedError(`turn execution policy was fenced: ${installedPolicy.reason}`);
   }
+  capabilitySettings = settingsForAcceptedSubscriptionTurn(
+    capabilitySettings,
+    installedPolicy.policy,
+    {
+      modelId: turn.model,
+      reasoningEffort: turn.reasoningEffort,
+      latencyMode: turn.latencyMode,
+    },
+  );
+  runtime.configure(capabilitySettings);
   const verifiedExecutionPolicy = assertTurnExecutionPolicyMatchesConfigV1(
     capabilitySettings,
     installedPolicy.policy,

package/src/activities/agent-turn/errors.ts

@@ -119,6 +119,7 @@ export function providerRecoveryResult(input: {
           input.failureCode === "sandbox_command_start_unavailable" ||
           input.failureCode === "mcp_transport_timeout" ||
           input.failureCode === "mcp_transport_unavailable" ||
+          input.failureCode === "turn_execution_policy_definition_mismatch" ||
           input.failureCode === POST_COMPACTION_CONTINUATION_EMPTY_CODE
         ? Math.max(
             providerDelay ?? 0,

package/src/activities/agent-turn/failure-settlement.ts

@@ -16,14 +16,14 @@ import {
 } from "@opengeni/db";
 import { publishDurableSessionEvents } from "@opengeni/events";
 import { maxTurnsExceededRunState } from "@opengeni/runtime";
-import { CancelledFailure } from "@temporalio/activity";
+import { ApplicationFailure, CancelledFailure } from "@temporalio/activity";
 import {
   authoritativeCodexCapacityResetAt,
   classifyCodexPin,
   selectCodexCredentialLeaseForTurn,
   type CodexRotationStrategy,
 } from "../codex-rotation";
-import type { Settings } from "@opengeni/config";
+import { TurnExecutionPolicyDefinitionMismatchError, type Settings } from "@opengeni/config";
 import {
   classifyCodexEncryptedArtifactRejection,
   classifyCodexUsageLimitError,
@@ -1499,10 +1499,28 @@ export async function settleTurnFailure(deps: TurnFailureDeps): Promise<RunAgent
   // truth, recover this SAME accepted turn, then let the workflow re-claim
   // it after a pacing delay. This is independent of goal state and never
   // relies on a synthetic continuation prompt.
-  let failure = agentRunFailurePayload(error, {
-    isCodexTurn: billingState.isCodexTurn,
-  }) as ReturnType<typeof agentRunFailurePayload>;
-  if (failure.retryable && eventing.publish && attempt.turnId && eventing.turnStartedPublished) {
+  // A rolling-deployment definition mismatch is a separate configuration
+  // class: only the exact typed setup error can use this checkpoint before
+  // eventing exists. No generic setup/credential failure gains retry authority.
+  const earlyDefinitionMismatch =
+    error instanceof TurnExecutionPolicyDefinitionMismatchError &&
+    !attempt.modelRequestStarted &&
+    !eventing.turnStartedPublished &&
+    !!attempt.turnId &&
+    !!attempt.triggerEventId &&
+    attempt.executionGeneration > 0;
+  let failure = (
+    earlyDefinitionMismatch
+      ? { error: error.message, code: error.code, retryable: true }
+      : agentRunFailurePayload(error, {
+          isCodexTurn: billingState.isCodexTurn,
+        })
+  ) as ReturnType<typeof agentRunFailurePayload>;
+  if (
+    attempt.turnId &&
+    (earlyDefinitionMismatch ||
+      (failure.retryable && eventing.publish && eventing.turnStartedPublished))
+  ) {
     const nextProviderRecoveryCount = attempt.providerRecoveryCount + 1;
     const recoveryResult = providerRecoveryResult({
       failureCode: failure.code,
@@ -1511,8 +1529,10 @@ export async function settleTurnFailure(deps: TurnFailureDeps): Promise<RunAgent
     });
     try {
       if (recoveryResult.status === "recovering") {
-        await flushRuntimeBatcher();
-        await historySink.reconcileConversationTruth({ requireDurable: true });
+        if (!earlyDefinitionMismatch) {
+          await flushRuntimeBatcher();
+          await historySink.reconcileConversationTruth({ requireDurable: true });
+        }
         const recovery = await requestSessionTurnRecovery(db, input.workspaceId, {
           sessionId: input.sessionId,
           turnId: attempt.turnId,
@@ -1540,6 +1560,18 @@ export async function settleTurnFailure(deps: TurnFailureDeps): Promise<RunAgent
         return claimedResult(recoveryResult);
       }
       failure = providerRecoveryExhaustedFailure(failure, recoveryResult);
+      if (earlyDefinitionMismatch) {
+        // Setup has no eventing sink yet. Carry only the fixed, safe diagnostic
+        // through Temporal into exact-attempt workflow failure settlement.
+        control.activityStatus = "failed";
+        control.turnMetricOutcome = "failed";
+        control.activityError = error;
+        throw ApplicationFailure.create({
+          message: `${error.message}. Automatic same-turn configuration recovery exhausted after ${recoveryResult.providerRecoveryCount} retries.`,
+          type: "TurnExecutionPolicyDefinitionMismatchError",
+          nonRetryable: true,
+        });
+      }
     } catch (recoveryError) {
       const escaped =
         recoveryResult.status === "recovering"

package/src/activities/agent-turn/run.ts

@@ -1302,6 +1302,7 @@ export function createRunAgentTurnActivity(services: () => Promise<ActivityServi
 
           const builtAgent = await buildTurnAgent({
             skillCatalog: toolRuntime.skillCatalog,
+            mcpServers: toolRuntime.mcpServers,
             input,
             db,
             runtime,

package/src/activities/agent-turn/sandbox-runtime.ts

@@ -1,7 +1,7 @@
 import {
   advanceWorkspaceGeneration,
   verifyWorkspaceMutationSettlement,
-  heartbeatLeaseHolder,
+  heartbeatLeaseHolderStatus,
   readLease,
   accrueWarmSeconds,
   SandboxWorkspaceMutationFencedError,
@@ -526,12 +526,15 @@ export function createSandboxTurnRuntime(deps: SandboxTurnRuntimeDeps) {
     // home turn that degraded to the cloud group box (swap-away / flag-off), that
     // is the deployment default (modal), so the fallback box is warm-metered at
     // the cloud rate instead of selfhosted's rate-0 (which would under-bill).
-    const warmRate = sandboxWarmRateMicrosPerSecond(
-      settings,
-      warmBackend ?? (sandbox.established.backendId as Settings["sandboxBackend"]),
-    );
+    const warmRate =
+      settings.sandboxWarmBillingMode === "usage_only"
+        ? 0
+        : sandboxWarmRateMicrosPerSecond(
+            settings,
+            warmBackend ?? (sandbox.established.backendId as Settings["sandboxBackend"]),
+          );
     sandboxState.leaseHeartbeatTimer = setInterval(() => {
-      void heartbeatLeaseHolder(db, {
+      void heartbeatLeaseHolderStatus(db, {
         accountId: input.accountId,
         workspaceId: input.workspaceId,
         sandboxGroupId: heartbeatGroupId,
@@ -539,9 +542,17 @@ export function createSandboxTurnRuntime(deps: SandboxTurnRuntimeDeps) {
         holderId: heartbeatHolderId,
         leaseTtlMs: settings.sandboxLeaseTtlMs,
         expectedEpoch: heartbeatEpoch,
+        billingMode: settings.sandboxWarmBillingMode,
       })
-        .then(async (alive) => {
-          if (alive) return;
+        .then(async (status) => {
+          if (status.fence === "funding") {
+            stopLeaseHeartbeat();
+            sandboxRotationController.abort(
+              new Error("Insufficient OpenGeni credits to extend paid sandbox compute"),
+            );
+            return;
+          }
+          if (status.leaseExtended) return;
           const rotation = await beginRotationPreemption(sandbox, heartbeatEpoch, heartbeatGroupId);
           if (rotation === "not_rotating") {
             // The holder was reaped, the exact attempt closed, the epoch was
@@ -557,9 +568,14 @@ export function createSandboxTurnRuntime(deps: SandboxTurnRuntimeDeps) {
         sandboxGroupId: heartbeatGroupId,
         expectedEpoch: heartbeatEpoch,
         warmRateMicrosPerSecond: warmRate,
+        billingMode: settings.sandboxWarmBillingMode,
         subjectId: input.sessionId,
       })
-        .then((result) => recordCreditMicros(observability, "usage", result.costMicros))
+        .then((result) => {
+          if (settings.sandboxWarmBillingMode === "credits") {
+            recordCreditMicros(observability, "usage", result.costMicros);
+          }
+        })
         .catch(() => undefined);
       // MID-SESSION snapshot (sandbox-file-persistence): while the turn holds
       // the box, fold a fresh /workspace snapshot onto the lease every

package/src/activities/agent-turn/stream-attempt.ts

@@ -17,6 +17,7 @@ import { publishDurableSessionEvents } from "@opengeni/events";
 import {
   normalizeModelCallUsage,
   normalizeSdkEvent,
+  withMcpToolDisplayMetadata,
   extractOpenSuffixFromRunState,
   assertOpenSuffixResumable,
   interruptionKindForCallItem,
@@ -1304,6 +1305,11 @@ export async function runTurnStreamAttempt(
               : {},
           );
         for (const event of normalized) {
+          if (event.type === "agent.toolCall.created")
+            event.payload = withMcpToolDisplayMetadata(
+              eventing.preparedTools?.mcpServers ?? [],
+              event.payload,
+            );
           streamTiming.onEvent(event.type);
           await eventing.batcher.push(event);
         }
@@ -1641,7 +1647,14 @@ export async function runTurnStreamAttempt(
               ? [
                   {
                     type: "session.requiresAction" as const,
-                    payload: { approvals },
+                    payload: {
+                      approvals: approvals.map((approval) =>
+                        withMcpToolDisplayMetadata(
+                          eventing.preparedTools?.mcpServers ?? [],
+                          approval,
+                        ),
+                      ),
+                    },
                   },
                 ]
               : []),

package/src/activities/agent-turn/tool-environment.ts

@@ -31,6 +31,7 @@ import {
   type ConnectorAttachmentMaterializationRequest,
   type ConnectorActionPolicyHooks,
   createFirstPartyInteractionAttemptToolDefinitions,
+  mcpToolDisplayMetadata,
 } from "@opengeni/runtime";
 import {
   createGoogleDrivePublicationAttemptTool,
@@ -1095,6 +1096,9 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
         executionGeneration: attempt.executionGeneration,
       },
       cancellationSignal,
+      undefined,
+      {},
+      (name) => mcpToolDisplayMetadata(tools.mcpServers, name),
     );
     eventing.codemodeDispatcher.start();
   };
@@ -1113,6 +1117,7 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
   return {
     attemptConnectorActionBindings,
     connectorActionPolicy,
+    mcpServers: runSettings.mcpServers,
     generateSessionTitleInParallel: titleToolPlan.generateTitleInParallel,
     postToolPreparationStartedAt,
     preparationIndependentToolNames: [

package/src/activities/codemode-dispatcher.ts

@@ -16,6 +16,7 @@ import {
   CODEMODE_MAX_CONCURRENT_CALLS_PER_ATTEMPT,
   type CodemodeOperation,
   type SessionEvent,
+  type ToolDisplayMetadata,
 } from "@opengeni/contracts";
 import {
   cancelQueuedCodemodeOperationsForAttempt,
@@ -70,6 +71,7 @@ export class CodemodeAttemptDispatcher {
     private readonly turnSignal?: AbortSignal,
     private readonly maxConcurrentCalls = CODEMODE_MAX_CONCURRENT_CALLS_PER_ATTEMPT,
     timings: CodemodeDispatcherTimings = {},
+    private readonly toolDisplayMetadata?: (modelName: string) => ToolDisplayMetadata | undefined,
   ) {
     if (
       environment.catalog.accountId !== scope.accountId ||
@@ -371,6 +373,10 @@ export class CodemodeAttemptDispatcher {
             payload: {
               id: operation.operationId,
               name: entry.modelName,
+              display: this.toolDisplayMetadata?.(entry.modelName) ?? {
+                toolName: entry.identity.toolName,
+                ...(entry.title ? { title: entry.title } : {}),
+              },
               arguments: operation.arguments,
               origin: "codemode",
               subjectId: operation.caller.subjectId,

package/src/activities/knowledge-indexing.ts

@@ -1,15 +1,21 @@
 import { configuredStaticUsageLimits } from "@opengeni/config";
+import { documentEmbeddingCostMicros, paidDocumentEmbedding } from "@opengeni/core";
 import {
+  applyCreditDebitAfterUse,
   appendKnowledgeIndexChunks,
   claimKnowledgeIndexJobs,
   completeKnowledgeIndexJob,
   continueKnowledgeIndexJob,
   deferKnowledgeIndexJob,
+  freezeKnowledgeIndexBillingMode,
   getBillingBalance,
+  guardPaidKnowledgeIndexPublication,
+  knowledgeIndexBillingActivationTime,
   readKnowledgeIndexSource,
   recordUsageEvent,
   sumUsageQuantity,
   withWorkspaceUsageLock,
+  waitKnowledgeIndexForFunding,
 } from "@opengeni/db";
 import type { DocumentServices } from "@opengeni/documents";
 import type { ControlActivityServices } from "./types";
@@ -18,11 +24,26 @@ export function createKnowledgeIndexingActivities(
   services: () => Promise<ControlActivityServices>,
   resolveDocumentServices?: () => Promise<DocumentServices>,
 ) {
+  // Unpaid modes retain a conservative first-poll cutoff; paid mode uses the
+  // operator's explicit timestamp so process restarts cannot change eligibility.
+  let activationTime: Promise<Date> | undefined;
   return {
     indexKnowledge: async () => {
       const result = { completed: 0, advanced: 0, deferred: 0, unavailable: 0 };
       if (!resolveDocumentServices) return result;
       const { db, settings, observability } = await services();
+      let policyActivatedAt: Date;
+      if (paidDocumentEmbedding(settings)) {
+        if (!settings.documentEmbeddingCreditsActivatedAt)
+          throw new Error("paid Knowledge embedding requires an activation cutoff");
+        policyActivatedAt = new Date(settings.documentEmbeddingCreditsActivatedAt);
+      } else {
+        activationTime ??= knowledgeIndexBillingActivationTime(db).catch((error) => {
+          activationTime = undefined;
+          throw error;
+        });
+        policyActivatedAt = await activationTime;
+      }
       const { embedder } = await resolveDocumentServices();
       const { knowledgeIndexChunks } = await import("@opengeni/documents");
       const claims = await claimKnowledgeIndexJobs(db, {
@@ -37,9 +58,9 @@ export function createKnowledgeIndexingActivities(
             result.unavailable++;
             continue;
           }
-          // Reuse the existing embedding budget and usage ledger. Each batch's
-          // checkpoint and charge commit together, so a restart never charges
-          // twice for an accepted projection. Canonical text is never changed.
+          // The checkpoint, usage and post-use debit commit together. A new
+          // generation requires funding once; committed batches may finish even
+          // if their accumulated cost takes the balance below zero.
           await withWorkspaceUsageLock(db, source.billingWorkspaceId, async (lockedDb) => {
             const current = await readKnowledgeIndexSource(lockedDb, claim);
             if (!current) {
@@ -57,9 +78,34 @@ export function createKnowledgeIndexingActivities(
               chunks.push(chunk);
             }
             if (chunks.length) {
-              if (settings.billingMode === "stripe" || settings.usageLimitsMode === "managed") {
+              const frozenPolicy = await freezeKnowledgeIndexBillingMode(
+                lockedDb,
+                claim,
+                // Deterministic embeddings incur no provider charge; a valid
+                // credits-mode config with zero tariff must still index them.
+                // OpenAI shadow keeps its price snapshot for internal estimates.
+                settings.documentEmbeddingProvider === "openai"
+                  ? (settings.documentEmbeddingBillingMode ?? "usage_only")
+                  : "usage_only",
+                policyActivatedAt,
+                settings.documentEmbeddingRateMicrosPerMillionBytes ?? 0,
+              );
+              if (frozenPolicy.mode === "awaiting_review") {
+                result.deferred++;
+                return;
+              }
+              if (frozenPolicy.mode === "obsolete") {
+                result.unavailable++;
+                return;
+              }
+              const paid = frozenPolicy.mode === "credits" && paidDocumentEmbedding(settings);
+              if (paid && current.nextIndex === 0) {
                 const balance = await getBillingBalance(lockedDb, claim.accountId);
-                if (balance.balanceMicros <= 0) throw new Error("insufficient OpenGeni credits");
+                if (balance.balanceMicros <= 0) {
+                  await waitKnowledgeIndexForFunding(lockedDb, claim);
+                  result.deferred++;
+                  return;
+                }
               }
               if (settings.usageLimitsMode === "static" || settings.usageLimitsMode === "managed") {
                 const limit =
@@ -75,9 +121,24 @@ export function createKnowledgeIndexingActivities(
                     throw new Error("monthly document indexing limit reached");
                 }
               }
-              const vectors = await embedder.embedMany(chunks.map((chunk) => chunk.embeddingInput));
+              const inputs = chunks.map((chunk) => chunk.embeddingInput);
+              const bytes = inputs.reduce(
+                (sum, input) => sum + Buffer.byteLength(input, "utf8"),
+                0,
+              );
+              const vectors = await embedder.embedMany(inputs);
               if (vectors.length !== chunks.length)
                 throw new Error("Incomplete Knowledge embeddings");
+              // A reviewer may have rejected this revision during the provider
+              // call. The DB guard holds its publication row through settlement.
+              if (paid) {
+                const publication = await guardPaidKnowledgeIndexPublication(lockedDb, claim);
+                if (publication !== "published") {
+                  if (publication === "obsolete") result.unavailable++;
+                  else result.deferred++;
+                  return;
+                }
+              }
               const appended = await appendKnowledgeIndexChunks(
                 lockedDb,
                 claim,
@@ -98,6 +159,63 @@ export function createKnowledgeIndexingActivities(
                 sourceResourceId: claim.revisionId,
                 idempotencyKey: `knowledge.indexed:${claim.revisionId}:${claim.generation}:${current.nextIndex}`,
               });
+              await recordUsageEvent(lockedDb, {
+                accountId: claim.accountId,
+                workspaceId: current.billingWorkspaceId,
+                eventType: "document.embedding_bytes",
+                quantity: bytes,
+                unit: "byte",
+                sourceResourceType: "knowledge_revision",
+                sourceResourceId: claim.revisionId,
+                idempotencyKey: `knowledge.embedding_bytes:${claim.revisionId}:${claim.generation}:${current.nextIndex}`,
+              });
+              if (frozenPolicy.mode === "shadow" && frozenPolicy.rateMicrosPerMillionBytes > 0) {
+                const estimate = documentEmbeddingCostMicros(
+                  {
+                    ...settings,
+                    documentEmbeddingRateMicrosPerMillionBytes:
+                      frozenPolicy.rateMicrosPerMillionBytes,
+                  },
+                  bytes,
+                );
+                if (estimate > 0)
+                  await recordUsageEvent(lockedDb, {
+                    accountId: claim.accountId,
+                    workspaceId: current.billingWorkspaceId,
+                    eventType: "document.embedding_shadow_estimate",
+                    quantity: estimate,
+                    unit: "micro_usd",
+                    sourceResourceType: "knowledge_revision",
+                    sourceResourceId: claim.revisionId,
+                    idempotencyKey: `knowledge.embedding_shadow:${claim.revisionId}:${claim.generation}:${current.nextIndex}`,
+                  });
+              }
+              if (paid) {
+                const cost = documentEmbeddingCostMicros(
+                  {
+                    ...settings,
+                    documentEmbeddingRateMicrosPerMillionBytes:
+                      frozenPolicy.rateMicrosPerMillionBytes,
+                  },
+                  bytes,
+                );
+                if (cost > 0)
+                  await applyCreditDebitAfterUse(lockedDb, {
+                    accountId: claim.accountId,
+                    workspaceId: current.billingWorkspaceId,
+                    type: "document_embedding_debit",
+                    amountMicros: cost,
+                    sourceType: "knowledge_revision",
+                    sourceId: claim.revisionId,
+                    idempotencyKey: `knowledge.embedding:${claim.revisionId}:${claim.generation}:${current.nextIndex}`,
+                    metadata: {
+                      model: claim.model,
+                      bytes,
+                      chunks: chunks.length,
+                      rateMicrosPerMillionBytes: frozenPolicy.rateMicrosPerMillionBytes,
+                    },
+                  });
+              }
             }
             if (more) {
               await continueKnowledgeIndexJob(lockedDb, claim);

package/src/activities/retained-screenshots.ts

@@ -13,6 +13,7 @@ import {
 import {
   RetainedScreenshotQuotaExceededError,
   getRetainedScreenshotArtifact,
+  getRetainedScreenshotArtifactForToolCall,
   isDatabasePersistenceFailure,
   isSessionEventPersistenceError,
   prepareRetainedScreenshotArtifact,
@@ -874,6 +875,26 @@ async function materializeRetainedScreenshotHistoryWithCache(
   cache: Map<string, string>,
 ): Promise<Array<Record<string, unknown>>> {
   const now = input.now ?? new Date();
+  const receiptForMarker = async (
+    marker: unknown,
+    callId: string | null,
+  ): Promise<RetainedArtifactMetadata | null> => {
+    if (!isRetainedImageMarker(marker)) return null;
+    const receipt = retainedReceipt(marker.artifact);
+    if (receipt) return receipt;
+    // Older canonical history could truncate the artifact UUID and reason as
+    // ordinary text. Recover only from this exact session and an unambiguous
+    // tool call; otherwise fail before constructing a malformed model image.
+    if (!callId) throw new Error("Retained screenshot receipt has no tool call identity");
+    const artifact = await getRetainedScreenshotArtifactForToolCall(
+      input.db,
+      input.workspaceId,
+      input.sessionId,
+      callId,
+    );
+    if (!artifact) throw new Error("Retained screenshot receipt cannot be recovered");
+    return reference(artifact);
+  };
   const dataUrlForReceipt = async (receipt: RetainedArtifactMetadata): Promise<string> => {
     let dataUrl = cache.get(receipt.artifactId);
     if (!dataUrl) {
@@ -921,8 +942,12 @@ async function materializeRetainedScreenshotHistoryWithCache(
     }
     return dataUrl;
   };
-  const materializeEntry = async (entry: unknown): Promise<unknown> => {
-    const receipt = retainedReceiptFromImageContent(entry);
+  const materializeEntry = async (entry: unknown, callId: string | null): Promise<unknown> => {
+    const image =
+      entry && typeof entry === "object" && !Array.isArray(entry)
+        ? (entry as Record<string, unknown>).image
+        : null;
+    const receipt = await receiptForMarker(image, callId);
     if (!receipt) return entry;
     const dataUrl = await dataUrlForReceipt(receipt);
     return { ...(entry as Record<string, unknown>), image: dataUrl };
@@ -930,7 +955,8 @@ async function materializeRetainedScreenshotHistoryWithCache(
 
   const materialized: Array<Record<string, unknown>> = [];
   for (const item of input.history) {
-    const directReceipt = retainedReceiptFromDirectOutput(item.output);
+    const callId = historyCallId(item);
+    const directReceipt = await receiptForMarker(item.output, callId);
     if (directReceipt) {
       materialized.push({
         ...item,
@@ -954,7 +980,7 @@ async function materializeRetainedScreenshotHistoryWithCache(
     let changed = false;
     const output: unknown[] = [];
     for (const entry of outputEntries) {
-      const next = await materializeEntry(entry);
+      const next = await materializeEntry(entry, callId);
       changed ||= next !== entry;
       output.push(next);
     }
@@ -1110,6 +1136,18 @@ function retainedReceiptFromImageContent(entry: unknown): RetainedArtifactMetada
   return retainedReceipt(marker.artifact);
 }
 
+function isRetainedImageMarker(value: unknown): value is {
+  type: "retained_artifact";
+  artifact: unknown;
+} {
+  return Boolean(
+    value &&
+    typeof value === "object" &&
+    !Array.isArray(value) &&
+    (value as Record<string, unknown>).type === RETAINED_IMAGE_MARKER,
+  );
+}
+
 function retainedReceiptFromDirectOutput(output: unknown): RetainedArtifactMetadata | null {
   if (!output || typeof output !== "object" || Array.isArray(output)) return null;
   const marker = output as Record<string, unknown>;