@opengeni/worker-bundle 2.0.1 → 2.0.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@opengeni/worker-bundle",
-  "version": "2.0.1",
+  "version": "2.0.3",
   "description": "OpenGeni worker entry and reusable embedded lifecycle, shipped with a release-coherent pre-bundled Temporal workflow artifact.",
   "license": "Apache-2.0",
   "repository": {
@@ -56,21 +56,21 @@
     "@llamaindex/liteparse": "2.14.2",
     "@opengeni/agent-proto": "^0.6.0",
     "@opengeni/capabilities": "^0.3.4",
-    "@opengeni/codemode": "^0.6.1",
-    "@opengeni/codex": "^0.2.25",
-    "@opengeni/config": "^2.1.0",
-    "@opengeni/contracts": "^5.1.0",
-    "@opengeni/core": "^4.0.1",
-    "@opengeni/db": "^6.0.1",
-    "@opengeni/documents": "^0.8.32",
-    "@opengeni/events": "^0.4.30",
-    "@opengeni/github": "^0.7.15",
+    "@opengeni/codemode": "^0.6.2",
+    "@opengeni/codex": "^0.2.26",
+    "@opengeni/config": "^2.1.1",
+    "@opengeni/contracts": "^5.1.1",
+    "@opengeni/core": "^4.0.3",
+    "@opengeni/db": "^6.0.3",
+    "@opengeni/documents": "^0.8.34",
+    "@opengeni/events": "^0.4.32",
+    "@opengeni/github": "^0.7.16",
     "@opengeni/network": "^0.3.1",
-    "@opengeni/observability": "^0.8.31",
-    "@opengeni/runtime": "^4.0.1",
-    "@opengeni/sdk": "^7.1.0",
-    "@opengeni/storage": "^0.2.132",
-    "@opengeni/tool-gateway": "^0.1.12",
+    "@opengeni/observability": "^0.8.32",
+    "@opengeni/runtime": "^4.0.3",
+    "@opengeni/sdk": "^7.1.1",
+    "@opengeni/storage": "^0.2.133",
+    "@opengeni/tool-gateway": "^0.1.13",
     "@opengeni/xai-subscription": "^0.1.4",
     "@temporalio/activity": "^1.17.0",
     "@temporalio/client": "^1.17.0",

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

@@ -16,7 +16,8 @@ import {
   type CompactionProviderRejection,
   isMcpRequestTimeoutError,
   isMcpTransportConnectivityError,
-  isModalTaskExecStartDnsResolutionError,
+  isModalTaskExecStartPreDispatchUnavailableError,
+  isRoutingMutationOutcomeUnknownError,
   RoutingWorkspaceRootChangedError,
   SandboxMaterializationVerificationError,
   materializationVerificationDiagnostic,
@@ -119,6 +120,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,
@@ -1021,10 +1023,13 @@ function baseAgentRunFailurePayload(
       retryable: true,
     };
   }
-  if (isModalTaskExecStartDnsResolutionError(error)) {
+  if (
+    !isRoutingMutationOutcomeUnknownError(error) &&
+    isModalTaskExecStartPreDispatchUnavailableError(error)
+  ) {
     return {
       error:
-        "The managed sandbox command transport was temporarily unreachable before the command started. The same turn will retry after a short delay.",
+        "The managed sandbox command router was not ready before the command was sent. The same turn will retry after a short delay.",
       code: "sandbox_command_start_unavailable",
       retryable: true,
     };

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/history.ts

@@ -186,7 +186,20 @@ export function toolCallProducesRetainableSessionImage(name: string | null): boo
   return (
     name === "computer_screenshot" ||
     name === "view_image" ||
-    name === "interaction__computer_observe"
+    name === "interaction__computer_observe" ||
+    retainableBrowserScreenshotToolCall(name)
+  );
+}
+
+/** Browser interactions can carry an image block, including opt-in observe frames. */
+export function retainableBrowserScreenshotToolCall(name: string | null): boolean {
+  return (
+    name === "browser_screenshot" ||
+    name === "interaction__browser_screenshot" ||
+    name === "browser_observe" ||
+    name === "interaction__browser_observe" ||
+    name === "browser_act" ||
+    name === "interaction__browser_act"
   );
 }
 

package/src/activities/agent-turn/media-artifacts.ts

@@ -12,7 +12,7 @@ import {
   type TurnToolCancellationFence,
 } from "@opengeni/runtime";
 import type { Settings } from "@opengeni/config";
-import type { RetainedArtifactMetadata } from "@opengeni/contracts";
+import type { RetainedArtifactMetadata, RetainedSessionScreenshotKind } from "@opengeni/contracts";
 import type { ResumedTurnSandbox } from "../../sandbox-resume";
 import type { SharedActivityServices } from "../types";
 import {
@@ -79,6 +79,7 @@ export class TurnMediaArtifacts {
   sandboxFileDownloadBackend: Settings["sandboxBackend"];
   nativeImageGenerationRetention: NativeImageGenerationRetention | null = null;
   readonly retainedSessionImageCallIds = new Set<string>();
+  readonly retainedSessionImageKindsByCallId = new Map<string, RetainedSessionScreenshotKind>();
 
   constructor(private readonly deps: TurnMediaArtifactDeps) {
     this.sandboxFileDownloadBackend = deps.getModelRunSettings().sandboxBackend;

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

@@ -61,7 +61,7 @@ import {
 import {
   compactRetainedScreenshotHistory,
   sdkEventContainsInlineImage,
-  retainComputerScreenshot,
+  retainSessionScreenshot,
   typedScreenshotFromSdkEvent,
   unavailableRetainedSessionImage,
 } from "../retained-screenshots";
@@ -108,6 +108,7 @@ import {
 import {
   pendingToolCallFromSdkEvent,
   toolCallProducesRetainableSessionImage,
+  retainableBrowserScreenshotToolCall,
   completedToolCallFromSdkEvent,
 } from "./history";
 import { checkpointHistoryBeforeProviderDispatch } from "./provider-dispatch-barrier";
@@ -1125,6 +1126,12 @@ export async function runTurnStreamAttempt(
           currentToolBatchCallIds.add(pendingToolCall.callId);
           if (toolCallProducesRetainableSessionImage(pendingToolCall.callName)) {
             media.retainedSessionImageCallIds.add(pendingToolCall.callId);
+            if (retainableBrowserScreenshotToolCall(pendingToolCall.callName)) {
+              media.retainedSessionImageKindsByCallId.set(
+                pendingToolCall.callId,
+                "browser_screenshot",
+              );
+            }
           }
         }
         const completedToolCall = completedToolCallFromSdkEvent(durableSdkEvent);
@@ -1147,6 +1154,7 @@ export async function runTurnStreamAttempt(
               toolCallId: completedToolCall.callId,
               toolOutputId: completedToolCall.callId,
               reason: "unsupported",
+              kind: media.retainedSessionImageKindsByCallId.get(completedToolCall.callId),
             });
             media.retainedScreenshotReceiptsByCallId.set(
               completedToolCall.callId,
@@ -1169,12 +1177,13 @@ export async function runTurnStreamAttempt(
               toolCallId: typedScreenshot.callId,
               toolOutputId: typedScreenshot.toolOutputId,
               reason: "pending",
+              kind: media.retainedSessionImageKindsByCallId.get(completedToolCall.callId),
             });
             media.retainedScreenshotReceiptsByCallId.set(
               completedToolCall.callId,
               retainedScreenshotMetadata,
             );
-            retainedScreenshotMetadata = await retainComputerScreenshot({
+            retainedScreenshotMetadata = await retainSessionScreenshot({
               db,
               objectStorage,
               accountId: input.accountId,
@@ -1183,6 +1192,7 @@ export async function runTurnStreamAttempt(
               turnId: activeTurnId,
               attemptId: input.attemptId,
               output: typedScreenshot,
+              kind: media.retainedSessionImageKindsByCallId.get(completedToolCall.callId),
             });
             media.retainedScreenshotReceiptsByCallId.set(
               completedToolCall.callId,
@@ -1339,6 +1349,7 @@ export async function runTurnStreamAttempt(
           }
           for (const callId of stableToolCallIdsToClear) {
             media.retainedSessionImageCallIds.delete(callId);
+            media.retainedSessionImageKindsByCallId.delete(callId);
           }
         }
       }

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,43 @@ 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;
+              }
+              // An already-priced generation cannot become free when the
+              // operator turns off paid embedding. Retain its checkpoint and
+              // retry only after paid mode is restored; never call the provider
+              // or append a batch that the live policy would not debit.
+              if (frozenPolicy.mode === "credits" && !paidDocumentEmbedding(settings)) {
+                await deferKnowledgeIndexJob(lockedDb, claim);
+                result.deferred++;
+                return;
+              }
+              const paid = frozenPolicy.mode === "credits";
+              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 +130,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 +168,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);