@opengeni/worker-bundle 1.1.0 → 1.2.0

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

@@ -11,6 +11,7 @@ import {
   safeDatabaseErrorFacts,
   isRetryableDatabaseTransportFailure,
   isSessionEventPersistenceError,
+  DatabaseTransactionError,
   SandboxLeaseTransitionError,
 } from "@opengeni/db";
 import {
@@ -419,6 +420,12 @@ function retryableDatabaseFailureCode(
     const boundaries = new Set<object>();
     const ownDatabaseNodes = new Set<object>();
     const codes = new Set<PostClaimDatabaseRecoveryDetail["code"]>();
+    const ownChildren = (node: object): object[] =>
+      node instanceof DatabaseTransactionError
+        ? // Callback failures retained beside a failed rollback supply vetoes,
+          // never driver provenance for a provider's connection-looking error.
+          [...graph.get(node)!].filter((child) => child === node.cause)
+        : [...graph.get(node)!];
     // Ask the canonical transport predicate about ONLY this node's facts. Its
     // recursive search must not pair a DB sibling with an unrelated provider.
     for (const node of graph.keys()) {
@@ -441,7 +448,11 @@ function retryableDatabaseFailureCode(
           : isDatabaseConnectionSqlState(sqlState)
       )
         transports.add(node);
-      if (node instanceof DrizzleQueryError || isSessionEventPersistenceError(node)) {
+      if (
+        node instanceof DrizzleQueryError ||
+        node instanceof DatabaseTransactionError ||
+        isSessionEventPersistenceError(node)
+      ) {
         // Only actual errors raised at our ORM/typed persistence boundary own
         // their driver subtree. A PostgresError name, SDK wrapper or provider
         // socket by itself is never own-client provenance for a running turn.
@@ -449,7 +460,7 @@ function retryableDatabaseFailureCode(
         for (const source of queue) {
           if (ownDatabaseNodes.has(source)) continue;
           ownDatabaseNodes.add(source);
-          queue.push(...graph.get(source)!);
+          queue.push(...ownChildren(source));
         }
       }
     }
@@ -460,7 +471,7 @@ function retryableDatabaseFailureCode(
         if (seen.has(node)) continue;
         seen.add(node);
         if (transports.has(node)) return true;
-        queue.push(...graph.get(node)!);
+        queue.push(...ownChildren(node));
       }
       return false;
     };
@@ -469,7 +480,7 @@ function retryableDatabaseFailureCode(
     // order; a deeper transport/reset cannot override it either.
     for (const node of graph.keys()) {
       const record = node as Record<string, unknown>;
-      if (node instanceof DrizzleQueryError) {
+      if (node instanceof DrizzleQueryError || node instanceof DatabaseTransactionError) {
         boundaries.add(node);
         if (hasOwnTransport(node)) codes.add("db_failure");
       }
@@ -513,6 +524,12 @@ function retryableDatabaseFailureCode(
 }
 
 const RUNNING_TURN_DATABASE_TRANSPORT_CODES = new Set([
+  // postgres.js reports these for a physically lost connection, including
+  // transaction cleanup after its socket has already closed. Own-client
+  // provenance and the unknown-outcome veto remain mandatory above.
+  "CONNECTION_CLOSED",
+  "CONNECTION_DESTROYED",
+  "CONNECTION_ENDED",
   "ECONNREFUSED",
   "ECONNRESET",
   "CONNECT_TIMEOUT",

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

@@ -277,12 +277,20 @@ export async function settleTurnFailure(deps: TurnFailureDeps): Promise<RunAgent
   try {
     return await settleTurnFailureInAttempt(deps);
   } catch (error) {
+    if (
+      deps.control.activityStatus === "recovering" &&
+      error instanceof ApplicationFailure &&
+      error.type === "OpenGeniPostClaimDatabaseRecovery"
+    )
+      throw error;
     // Connectivity can disappear while settling an unrelated run error too.
     // Do not overwrite a possibly committed settlement; the control lane
     // re-reads exact ownership and becomes a stale no-op if it already closed.
     if (deps.attempt.turnId && deps.attempt.triggerEventId) {
       const recovery = postClaimDatabaseRecoveryFailure({
-        error,
+        // A failed rollback/terminal write must not erase no-replay evidence
+        // from the failure we were settling (notably unknown tool effects).
+        error: new AggregateError([error, deps.error], "Turn failure settlement failed"),
         turnId: deps.attempt.turnId,
         triggerEventId: deps.attempt.triggerEventId,
         executionGeneration: deps.attempt.executionGeneration,

package/src/activities/agent-turn/finalization-monitor.ts

@@ -1,4 +1,4 @@
-import type { Observability } from "@opengeni/observability";
+import { turnExecutionTelemetryKey, type Observability } from "@opengeni/observability";
 import type { TurnHeartbeatDetails } from "../../op-journal";
 import { armTurnQuiescenceWatchdog } from "./quiescence";
 
@@ -42,10 +42,21 @@ export function startTurnFinalizationMonitor(input: {
   details: TurnHeartbeatDetails;
   heartbeat: (details: TurnHeartbeatDetails) => void;
   requestWorkerDrain: () => void;
+  execution?: { workspaceId: string; sessionId: string; attemptId: string };
   timeoutMs?: number;
   slowAfterMs?: number;
 }) {
   const { observability, details } = input;
+  const correlation: { correlationId?: string } = {};
+  diagnose(() => {
+    if (input.execution) {
+      correlation.correlationId = turnExecutionTelemetryKey(
+        input.execution.workspaceId,
+        input.execution.sessionId,
+        input.execution.attemptId,
+      );
+    }
+  });
   if (!initialized.has(observability)) {
     for (const stage of TURN_FINALIZATION_STAGES) {
       diagnose(() => observability.incrementGauge({ ...INFLIGHT, labels: { stage }, amount: 0 }));
@@ -88,6 +99,7 @@ export function startTurnFinalizationMonitor(input: {
               surface: "turn_finalization",
               outcome: "containment",
               reason: next,
+              ...correlation,
             });
           } catch {
             // Diagnostic failure must never disable physical containment.
@@ -104,6 +116,7 @@ export function startTurnFinalizationMonitor(input: {
               surface: "turn_finalization",
               outcome: "slow",
               reason: next,
+              ...correlation,
             });
           }),
         input.slowAfterMs ?? 30_000,

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

@@ -128,6 +128,11 @@ export async function finalizeTurnAttempt(deps: TurnFinalizationDeps): Promise<v
       }
     },
     requestWorkerDrain: deps.requestWorkerDrain,
+    execution: {
+      workspaceId: deps.input.workspaceId,
+      sessionId: deps.input.sessionId,
+      attemptId: deps.input.attemptId,
+    },
   });
   try {
     monitor.enter("tool_writers");

package/src/activities/agent-turn/model-usage.ts

@@ -1,3 +1,4 @@
+import { canonicalizeConfiguredModelId } from "@opengeni/config";
 import {
   applyCreditDebitUpToBalance,
   recordUsageEvent,
@@ -55,6 +56,27 @@ export function modelUsageSourceKey(input: {
   return input.dispatchId ? `${input.dispatchId}:${input.positionalKey}` : input.positionalKey;
 }
 
+/** Legacy aggregate usage may only debit a policy shared by its completed calls. */
+export function aggregateCreditPolicyRevision(input: {
+  responseRevisions: ReadonlySet<number | undefined>;
+  lastAdmittedRevision: number | undefined;
+  chargesOpenGeniCredits: boolean;
+  totalTokens: number | null;
+}): number | undefined {
+  if (
+    input.chargesOpenGeniCredits &&
+    input.responseRevisions.size > 1 &&
+    (input.totalTokens ?? 0) > 0
+  ) {
+    throw new Error("Aggregate model usage spans different credit policy revisions");
+  }
+  // Bind the completed response, even if later preparation changed admission.
+  // Older runtimes without response callbacks retain their admitted snapshot.
+  return input.responseRevisions.size === 1
+    ? input.responseRevisions.values().next().value
+    : input.lastAdmittedRevision;
+}
+
 export function providerContextTokens(
   usage:
     | {
@@ -209,6 +231,7 @@ export async function processModelResponseTerminalEvent(input: {
   externallyBilled: boolean;
   chargesOpenGeniCredits?: boolean;
   countsTowardTokenCap?: boolean;
+  creditPolicyRevision?: number | undefined;
   servingCredentialId: string | null;
   priorSessionCredentialId: string | null;
   emittedSourceKeys: Set<string>;
@@ -231,6 +254,15 @@ export async function processModelResponseTerminalEvent(input: {
   if (!terminal) {
     return { status: "not_response" };
   }
+  // Some providers mirror a terminal response before normalized usage arrives.
+  // Claiming that empty raw mirror would discard the SDK's billable response.
+  if (
+    !terminal.usage &&
+    input.event.type === "raw_model_stream_event" &&
+    input.event.data.type !== "response_done"
+  ) {
+    return { status: "not_response" };
+  }
 
   const responseOrdinal = input.state.responseCount + 1;
   const sourceKey = modelUsageSourceKey({
@@ -268,6 +300,7 @@ export async function processModelResponseTerminalEvent(input: {
         turnId: input.turnId,
         turnAttemptId: input.turnAttemptId,
         model: input.model,
+        creditPolicyRevision: input.creditPolicyRevision,
         externallyBilled: input.externallyBilled,
         ...(input.chargesOpenGeniCredits !== undefined
           ? { chargesOpenGeniCredits: input.chargesOpenGeniCredits }
@@ -374,6 +407,7 @@ export async function processCompactionModelUsageEvent(input: {
   externallyBilled: boolean;
   chargesOpenGeniCredits?: boolean;
   countsTowardTokenCap?: boolean;
+  creditPolicyRevision?: number | undefined;
   servingCredentialId: string | null;
   priorSessionCredentialId: string | null;
   emittedSourceKeys: Set<string>;
@@ -416,6 +450,7 @@ export async function processCompactionModelUsageEvent(input: {
         turnId: input.turnId,
         turnAttemptId: input.turnAttemptId,
         model: input.model,
+        creditPolicyRevision: input.creditPolicyRevision,
         externallyBilled: input.externallyBilled,
         ...(input.chargesOpenGeniCredits !== undefined
           ? { chargesOpenGeniCredits: input.chargesOpenGeniCredits }
@@ -651,6 +686,7 @@ export async function recordModelUsageAndDebitCredits(
     externallyBilled: boolean;
     chargesOpenGeniCredits?: boolean;
     countsTowardTokenCap?: boolean;
+    creditPolicyRevision?: number | undefined;
     gatewayBilling?: ModelResponseUsage["gatewayBilling"];
     usage?: ModelUsageInput | ModelCallUsageInput | null;
     normalizedUsage?: ModelCallUsageNormalization;
@@ -842,6 +878,8 @@ export async function recordModelUsageAndDebitCredits(
       workspaceId: input.workspaceId,
       type: "model_usage_debit",
       requestedAmountMicros: costMicros,
+      modelId: canonicalizeConfiguredModelId(settings, input.model),
+      creditPolicyRevision: input.creditPolicyRevision,
       sourceType: "model_response",
       sourceId: `${input.turnId}:${input.sourceKey}`,
       idempotencyKey: `credit:model_usage_debit:${input.turnId}:${input.sourceKey}`,

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

@@ -1090,6 +1090,7 @@ export function createRunAgentTurnActivity(services: () => Promise<ActivityServi
               ),
             );
           const compactionPrep = await prepareCompaction({
+            entitlements,
             input,
             settings: capabilitySettings,
             db,

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

@@ -122,6 +122,7 @@ import {
   completedToolCallFromSdkEvent,
 } from "./history";
 import {
+  aggregateCreditPolicyRevision,
   modelUsageSourceKey,
   recordCompletedModelCallBeforeOwnershipFences,
   TurnEventPublisher,
@@ -384,6 +385,8 @@ export async function runTurnStreamAttempt(
   );
   let parallelSessionTitle: ReturnType<typeof startParallelSessionTitleGeneration> | null = null;
   let parallelSessionTitleFinished = false;
+  let creditPolicyRevision: number | undefined;
+  let titleCreditPolicyRevision: number | undefined;
   const finishParallelSessionTitle = async (): Promise<void> => {
     if (parallelSessionTitleFinished) return;
     parallelSessionTitleFinished = true;
@@ -393,6 +396,7 @@ export async function runTurnStreamAttempt(
     if (generated.usage) {
       await processSessionTitleModelUsageEvent({
         usage: generated.usage,
+        creditPolicyRevision: titleCreditPolicyRevision,
         state: sessionTitleUsageState,
         dispatchId: modelUsageDispatchId,
         settings,
@@ -595,7 +599,8 @@ export async function runTurnStreamAttempt(
   let finalReplyNudged = false;
   const revalidateModelCallAdmission = async () => {
     await historySink.reconcileConversationTruth({ requireDurable: true });
-    await ensureRunAllowedBetweenModelCalls({
+    creditPolicyRevision = await ensureRunAllowedBetweenModelCalls({
+      modelId: resolvedModel?.configured.id ?? turn.model,
       settings,
       db,
       accountId: input.accountId,
@@ -625,6 +630,7 @@ export async function runTurnStreamAttempt(
       signal: runtimeCancellationSignal,
       admit: revalidateModelCallAdmission,
     });
+    const responseCreditPolicyRevisions = new Set<number | undefined>();
     const responseCountBeforeStream = modelResponseState.responseCount;
     eventing.batcher = null;
     // The SDK emits every processed call item for one model response before
@@ -842,7 +848,10 @@ export async function runTurnStreamAttempt(
         attempt.modelRequestStarted = true;
         return await runtime.runStream(agent, runInput!, eventing.modelRunSettings, {
           beforeModelRequest: modelCallAdmission.beforeModelRequest,
-          onModelResponse: modelCallAdmission.onModelResponse,
+          onModelResponse: (event) => {
+            responseCreditPolicyRevisions.add(creditPolicyRevision);
+            return modelCallAdmission.onModelResponse(event);
+          },
           signal: runtimeCancellationSignal,
           sandboxEnvironment,
           onModelVisibleContext: async (snapshot) => {
@@ -1068,6 +1077,7 @@ export async function runTurnStreamAttempt(
           ? await media.retainNativeGeneratedImage(generatedImage)
           : null;
         const responseResult = await processModelResponseTerminalEvent({
+          creditPolicyRevision,
           event: next.value,
           state: modelResponseState,
           dispatchId: modelUsageDispatchId,
@@ -1158,21 +1168,9 @@ export async function runTurnStreamAttempt(
           await historySink.reconcileConversationTruth();
           turnLifecycleMetricsFor(observability).progress({ attemptId: input.attemptId });
           modelCheckpointMemoryCollector.schedule(observability);
-          await ensureRunAllowedBetweenModelCalls({
-            settings,
-            db,
-            accountId: input.accountId,
-            workspaceId: input.workspaceId,
-            isExternallyBilledTurn: billingState.isExternallyBilledTurn,
-            entitlements,
-            chargesOpenGeniCredits: billingState.chargesOpenGeniCredits,
-            countsTowardTokenCap: billingState.countsTowardTokenCap,
-            initiatingHumanSubjectId: turn.initiatingHumanSubjectId,
-            serializedRunState: () =>
-              media.compactMediaRunState(String(eventing.stream!.state.toString())),
-          });
         }
-        // Release only after both the debit and frozen-human admission finish.
+        // Release after settlement. The producer checks fresh admission before
+        // its next request; a final response keeps the revision it ran under.
         modelCallAdmission.settle(next.value);
         const durableSdkEvent = generatedImageReceipt
           ? compactGeneratedImageSdkEvent(next.value, generatedImageReceipt)
@@ -1561,6 +1559,12 @@ export async function runTurnStreamAttempt(
     if (!streamSawPerResponseUsage) {
       const aggregateUsage = eventing.stream.state.usage;
       const normalizedAggregateUsage = normalizeModelCallUsage(aggregateUsage);
+      const aggregatePolicyRevision = aggregateCreditPolicyRevision({
+        responseRevisions: responseCreditPolicyRevisions,
+        lastAdmittedRevision: creditPolicyRevision,
+        chargesOpenGeniCredits: billingState.chargesOpenGeniCredits,
+        totalTokens: normalizedAggregateUsage.totalTokens,
+      });
       const aggregateInput = normalizedAggregateUsage.telemetry.inputTokens;
       const aggregateSourceKey = modelUsageSourceKey({
         responseId: null,
@@ -1583,6 +1587,7 @@ export async function runTurnStreamAttempt(
           leaseLostMessage: "Provider credential lease expired during the active turn",
           recordUsage: async () => {
             const billing = await recordModelUsageAndDebitCredits(settings, db, {
+              creditPolicyRevision: aggregatePolicyRevision,
               accountId: input.accountId,
               workspaceId: input.workspaceId,
               sessionId: input.sessionId,
@@ -1975,6 +1980,7 @@ export async function runTurnStreamAttempt(
       turnExecutionPolicy.providerId,
       turnExecutionPolicy.latencyMode,
     );
+    titleCreditPolicyRevision = creditPolicyRevision;
     parallelSessionTitle = startParallelSessionTitleGeneration({
       signal: runtimeCancellationSignal,
       generate: async (signal) =>

package/src/activities/knowledge-indexing.ts

@@ -10,7 +10,7 @@ import {
   creditDebitAttributionMetadata,
   deferKnowledgeIndexJob,
   freezeKnowledgeIndexBillingMode,
-  getBillingBalance,
+  getSpendableCreditBalance,
   guardPaidKnowledgeIndexPublication,
   knowledgeIndexBillingActivationTime,
   readKnowledgeIndexSource,
@@ -211,7 +211,7 @@ export function createKnowledgeIndexingActivities(
                 return;
               }
               if (paid && current.nextIndex === 0) {
-                const balance = await getBillingBalance(lockedDb, claim.accountId);
+                const balance = await getSpendableCreditBalance(lockedDb, claim.accountId);
                 if (balance.balanceMicros <= 0) {
                   await waitKnowledgeIndexForFunding(lockedDb, claim);
                   result.deferred++;

package/src/observability-metrics.ts

@@ -1308,6 +1308,7 @@ export function recordSandboxDeadlineRotationsRequested(
 export const SANDBOX_COMMAND_CONTAINMENT_OUTCOMES = [
   "idle_enrolled",
   "deadline_enrolled",
+  "quiescence_enrolled",
   "resumed_enrolled",
   "not_eligible",
   "inspection_failed",