@opengeni/worker-bundle 2.0.1 → 2.0.3-canary.36079201993001

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);

package/src/activities/retained-screenshots.ts

@@ -6,13 +6,16 @@ import {
   COMPUTER_SCREENSHOT_WORKSPACE_QUOTA_BYTES,
   RetainedArtifactMetadataSchema,
   retainedScreenshotReferenceFromFile,
+  retainedSessionScreenshotKindFromObjectKey,
   type RetainedArtifactMetadata,
   type RetainedArtifactReference,
   type RetainedOutputUnavailableReason,
+  type RetainedSessionScreenshotKind,
 } from "@opengeni/contracts";
 import {
   RetainedScreenshotQuotaExceededError,
   getRetainedScreenshotArtifact,
+  getRetainedScreenshotArtifactForToolCall,
   isDatabasePersistenceFailure,
   isSessionEventPersistenceError,
   prepareRetainedScreenshotArtifact,
@@ -209,14 +212,26 @@ export function toolOutputContainsInlineImage(output: unknown): boolean {
       return (
         (record.type === "input_image" &&
           sdkImageSourceContainsInlineImage(record.image ?? record.image_url ?? record.imageUrl)) ||
-        (record.type === "image" && decodeInlineMcpImage(record.data, record.mimeType) !== null)
+        (record.type === "image" && isDeclaredInlineMcpImage(record))
       );
     });
   }
   if (!output || typeof output !== "object") return false;
   const record = output as Record<string, unknown>;
   if (Array.isArray(record.content)) return toolOutputContainsInlineImage(record.content);
-  return record.type === "image" && sdkImageSourceContainsInlineImage(record.image);
+  return (
+    record.type === "image" &&
+    (sdkImageSourceContainsInlineImage(record.image) || isDeclaredInlineMcpImage(record))
+  );
+}
+
+/** Detect image-bearing MCP blocks even when decoding rejects oversized or malformed bytes. */
+function isDeclaredInlineMcpImage(record: Record<string, unknown>): boolean {
+  return (
+    typeof record.data === "string" &&
+    typeof record.mimeType === "string" &&
+    record.mimeType.toLowerCase().startsWith("image/")
+  );
 }
 
 function decodeInlineMcpImage(
@@ -520,9 +535,14 @@ export function retainedScreenshotIdentity(input: {
   attemptId: string;
   toolCallId: string;
   toolOutputId: string;
+  kind?: RetainedSessionScreenshotKind | undefined;
 }): { artifactId: string; settlementKey: string } {
   const settlementKey = createHash("sha256")
-    .update("opengeni:computer-screenshot:v1\0")
+    .update(
+      input.kind === "browser_screenshot"
+        ? "opengeni:browser-screenshot:v1\0"
+        : "opengeni:computer-screenshot:v1\0",
+    )
     .update(input.sessionId)
     .update("\0")
     .update(input.turnId)
@@ -551,12 +571,13 @@ export function unavailableRetainedSessionImage(input: {
   toolCallId: string;
   toolOutputId: string;
   reason: RetainedOutputUnavailableReason;
+  kind?: RetainedSessionScreenshotKind | undefined;
 }): RetainedArtifactMetadata {
   const identity = retainedScreenshotIdentity(input);
   return unavailable(identity.artifactId, input.reason);
 }
 
-export async function retainComputerScreenshot(input: {
+export async function retainSessionScreenshot(input: {
   db: Database;
   objectStorage: ObjectStorage | null;
   accountId: string;
@@ -565,6 +586,7 @@ export async function retainComputerScreenshot(input: {
   turnId: string;
   attemptId: string;
   output: TypedScreenshotToolOutput;
+  kind?: RetainedSessionScreenshotKind | undefined;
   now?: Date;
   retentionMs?: number;
   workspaceQuotaBytes?: number;
@@ -575,6 +597,7 @@ export async function retainComputerScreenshot(input: {
     attemptId: input.attemptId,
     toolCallId: input.output.callId,
     toolOutputId: input.output.toolOutputId,
+    kind: input.kind,
   });
   let screenshot: ValidatedSessionImage;
   try {
@@ -594,7 +617,8 @@ export async function retainComputerScreenshot(input: {
   const retentionExpiresAt = new Date(
     now.getTime() + (input.retentionMs ?? COMPUTER_SCREENSHOT_RETENTION_MS),
   );
-  const objectKey = `workspaces/${input.workspaceId}/files/${identity.artifactId}/retained/session-image.${screenshot.extension}`;
+  const objectName = input.kind === "browser_screenshot" ? "browser-screenshot" : "session-image";
+  const objectKey = `workspaces/${input.workspaceId}/files/${identity.artifactId}/retained/${objectName}.${screenshot.extension}`;
   let prepared: RetainedScreenshotArtifact;
   try {
     prepared = (
@@ -697,6 +721,9 @@ export async function retainComputerScreenshot(input: {
   }
 }
 
+/** Existing computer callers keep their stable entry point. */
+export const retainComputerScreenshot = retainSessionScreenshot;
+
 /** Replace only new screenshot bytes with a compact receipt before first persistence. */
 export function compactRetainedScreenshotHistory(
   history: Array<Record<string, unknown>>,
@@ -874,6 +901,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 +968,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 +981,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 +1006,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);
     }
@@ -1059,12 +1111,15 @@ function reference(artifact: RetainedScreenshotArtifact): RetainedArtifactRefere
   if (!artifact.sessionId) {
     throw new Error(`Ready retained screenshot is detached: ${artifact.artifactId}`);
   }
+  const kind = retainedSessionScreenshotKindFromObjectKey(artifact.file.objectKey);
+  if (!kind) throw new Error(`Ready retained screenshot has unknown key: ${artifact.artifactId}`);
   const value = retainedScreenshotReferenceFromFile({
     ...artifact.file,
     sessionId: artifact.sessionId,
     width: artifact.width,
     height: artifact.height,
     expiresAt: artifact.retentionExpiresAt.toISOString(),
+    kind,
   });
   if (!value)
     throw new Error(`Ready retained screenshot receipt is invalid: ${artifact.artifactId}`);
@@ -1110,6 +1165,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>;
@@ -1120,7 +1187,11 @@ function retainedReceiptFromDirectOutput(output: unknown): RetainedArtifactMetad
 function retainedReceipt(value: unknown): RetainedArtifactMetadata | null {
   const parsed = RetainedArtifactMetadataSchema.safeParse(value);
   if (!parsed.success) return null;
-  return parsed.data.available && parsed.data.kind !== "computer_screenshot" ? null : parsed.data;
+  return parsed.data.available &&
+    parsed.data.kind !== "computer_screenshot" &&
+    parsed.data.kind !== "browser_screenshot"
+    ? null
+    : parsed.data;
 }
 
 function isInlineImageContent(entry: unknown): boolean {
@@ -1130,7 +1201,7 @@ function isInlineImageContent(entry: unknown): boolean {
     (record.type === "input_image" &&
       typeof record.image === "string" &&
       record.image.startsWith("data:image/")) ||
-    (record.type === "image" && decodeInlineMcpImage(record.data, record.mimeType) !== null)
+    (record.type === "image" && isDeclaredInlineMcpImage(record))
   );
 }