@opengeni/worker-bundle 1.0.0-canary.1 → 1.0.0

package/src/activities/knowledge-indexing.ts

@@ -1,28 +1,140 @@
 import { configuredStaticUsageLimits } from "@opengeni/config";
+import { documentEmbeddingCostMicros, paidDocumentEmbedding } from "@opengeni/core";
 import {
+  applyCreditDebitAfterUse,
   appendKnowledgeIndexChunks,
   claimKnowledgeIndexJobs,
+  checkWorkspaceAllowance,
   completeKnowledgeIndexJob,
   continueKnowledgeIndexJob,
+  creditDebitAttributionMetadata,
   deferKnowledgeIndexJob,
+  freezeKnowledgeIndexBillingMode,
   getBillingBalance,
+  guardPaidKnowledgeIndexPublication,
+  knowledgeIndexBillingActivationTime,
   readKnowledgeIndexSource,
   recordUsageEvent,
   sumUsageQuantity,
   withWorkspaceUsageLock,
+  waitKnowledgeIndexForFunding,
 } from "@opengeni/db";
 import type { DocumentServices } from "@opengeni/documents";
 import type { ControlActivityServices } from "./types";
 
+/** The configured monthly indexed-chunk limit, not a provider failure. */
+export class KnowledgeIndexUsageLimitError extends Error {
+  constructor() {
+    super("monthly document indexing limit reached");
+    this.name = "KnowledgeIndexUsageLimitError";
+  }
+}
+
+export type KnowledgeIndexFailureStage = "embedding" | "processing";
+
+/**
+ * Content-free classification for a deferred Knowledge index batch. Only
+ * protocol constants, an HTTP status, and a PostgreSQL SQLSTATE are retained;
+ * provider messages, bodies, SQL, and identifiers never leave the process.
+ * Outside the embedding call, only a PostgreSQL error in the cause chain is
+ * attributed to the database; any other failure stays a worker failure.
+ */
+export function knowledgeIndexFailureDiagnostic(
+  stage: KnowledgeIndexFailureStage,
+  error: unknown,
+): {
+  errorClass: "KnowledgeIndexOperationError";
+  errorCode:
+    | "knowledge_index_usage_limit_reached"
+    | "knowledge_index_embedding_failed"
+    | "knowledge_index_persistence_failed"
+    | "knowledge_index_failed";
+  origin: "worker" | "db";
+  status?: number;
+  sqlState?: string;
+} {
+  if (error instanceof KnowledgeIndexUsageLimitError) {
+    return {
+      errorClass: "KnowledgeIndexOperationError",
+      errorCode: "knowledge_index_usage_limit_reached",
+      origin: "worker",
+    };
+  }
+  if (stage === "embedding") {
+    const status = ownValue(error, "status");
+    return {
+      errorClass: "KnowledgeIndexOperationError",
+      errorCode: "knowledge_index_embedding_failed",
+      origin: "worker",
+      ...(typeof status === "number" && Number.isInteger(status) && status >= 100 && status <= 599
+        ? { status }
+        : {}),
+    };
+  }
+  const postgres = postgresErrorState(error);
+  if (!postgres) {
+    return {
+      errorClass: "KnowledgeIndexOperationError",
+      errorCode: "knowledge_index_failed",
+      origin: "worker",
+    };
+  }
+  return {
+    errorClass: "KnowledgeIndexOperationError",
+    errorCode: "knowledge_index_persistence_failed",
+    origin: "db",
+    ...(postgres.sqlState ? { sqlState: postgres.sqlState } : {}),
+  };
+}
+
+function ownValue(value: unknown, key: string): unknown {
+  try {
+    if (!value || typeof value !== "object") return undefined;
+    const descriptor = Object.getOwnPropertyDescriptor(value, key);
+    return descriptor && "value" in descriptor ? descriptor.value : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+/** The nearest PostgreSQL error in a short cause chain, with its SQLSTATE
+ * when that is a well-formed five-character code. */
+function postgresErrorState(error: unknown): { sqlState?: string } | undefined {
+  let current = error;
+  for (let depth = 0; depth < 4 && current && typeof current === "object"; depth += 1) {
+    if (ownValue(current, "name") === "PostgresError") {
+      const code = ownValue(current, "code");
+      return typeof code === "string" && /^[0-9A-Z]{5}$/.test(code) ? { sqlState: code } : {};
+    }
+    current = ownValue(current, "cause");
+  }
+  return undefined;
+}
+
 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, {
@@ -31,15 +143,16 @@ export function createKnowledgeIndexingActivities(
         limit: 2,
       });
       for (const claim of claims) {
+        let stage: KnowledgeIndexFailureStage = "processing";
         try {
           const source = await readKnowledgeIndexSource(db, claim);
           if (!source) {
             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 +170,69 @@ 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 && claim.billingAttribution.kind === "unknown") {
+                // Unknown legacy/source-preparation causality is not service.
+                // Keep the checkpoint without making an uncountable paid call.
+                observability.warn("Paid Knowledge indexing awaits initiating attribution", {
+                  errorCode: "knowledge_index_attribution_unavailable",
+                });
+                await deferKnowledgeIndexJob(lockedDb, claim);
+                result.deferred++;
+                return;
+              }
+              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 (paid) {
+                const refusal = await checkWorkspaceAllowance(lockedDb, {
+                  accountId: claim.accountId,
+                  workspaceId: current.billingWorkspaceId,
+                  subjectId:
+                    claim.billingAttribution.kind === "turn" ||
+                    claim.billingAttribution.kind === "human"
+                      ? claim.billingAttribution.initiatingHumanSubjectId
+                      : null,
+                });
+                if (refusal) {
+                  await deferKnowledgeIndexJob(lockedDb, claim);
+                  result.deferred++;
+                  return;
+                }
               }
               if (settings.usageLimitsMode === "static" || settings.usageLimitsMode === "managed") {
                 const limit =
@@ -71,13 +244,29 @@ export function createKnowledgeIndexingActivities(
                     eventType: "document.indexed",
                     since: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)),
                   });
-                  if (used + chunks.length > limit)
-                    throw new Error("monthly document indexing limit reached");
+                  if (used + chunks.length > limit) throw new KnowledgeIndexUsageLimitError();
                 }
               }
-              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,
+              );
+              stage = "embedding";
+              const vectors = await embedder.embedMany(inputs);
               if (vectors.length !== chunks.length)
                 throw new Error("Incomplete Knowledge embeddings");
+              stage = "processing";
+              // 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 +287,64 @@ 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: {
+                      ...creditDebitAttributionMetadata(claim.billingAttribution),
+                      model: claim.model,
+                      bytes,
+                      chunks: chunks.length,
+                      rateMicrosPerMillionBytes: frozenPolicy.rateMicrosPerMillionBytes,
+                    },
+                  });
+              }
             }
             if (more) {
               await continueKnowledgeIndexJob(lockedDb, claim);
@@ -112,16 +359,23 @@ export function createKnowledgeIndexingActivities(
               else result.unavailable++;
             }
           });
-        } catch {
+        } catch (error) {
           // Provider failures retain the last completed projection. The durable
           // queue owns retry/backoff; do not retry an entire activity implicitly.
-          await deferKnowledgeIndexJob(db, claim).catch(() => undefined);
-          result.deferred++;
-          observability.warn("Knowledge indexing batch deferred", {
-            accountId: claim.accountId,
-            entryId: claim.entryId,
-            revisionId: claim.revisionId,
+          // The stored reason stays the SQL lifecycle's fixed code; the log
+          // carries the content-free class/code of the actual cause.
+          observability.warn(
+            "Knowledge indexing batch deferred",
+            knowledgeIndexFailureDiagnostic(stage, error),
+          );
+          await deferKnowledgeIndexJob(db, claim).catch((deferError: unknown) => {
+            const deferDiagnostic = knowledgeIndexFailureDiagnostic("processing", deferError);
+            observability.warn("Knowledge indexing batch deferral failed", {
+              ...deferDiagnostic,
+              errorCode: "knowledge_index_defer_failed",
+            });
           });
+          result.deferred++;
         }
       }
       return result;

package/src/activities/knowledge-source-sync.ts

@@ -1,3 +1,7 @@
+import {
+  ATLASSIAN_NATIVE_RETIRED_REASON,
+  isRetiredNativeAtlassianSource,
+} from "@opengeni/contracts/atlassian-native-retirement";
 import { isDeepStrictEqual } from "node:util";
 import { AsyncLocalStorage } from "node:async_hooks";
 import { scheduledTaskKnowledgeSource } from "@opengeni/contracts";
@@ -243,6 +247,12 @@ export function createKnowledgeSourceSyncActivities(
       ) {
         return { action: "failed", bufferedWake: false };
       }
+      if (isRetiredNativeAtlassianSource(action))
+        return {
+          action: "failed",
+          bufferedWake: false,
+          errorCode: ATLASSIAN_NATIVE_RETIRED_REASON,
+        };
       boundSource = action;
       await authorize();
       await ensureKnowledgeSourceSyncState(db, task);

package/src/activities/mcp-account-routes.ts

@@ -0,0 +1,17 @@
+import type { McpConnectionAccountBinding, ToolAuthNeededPayload } from "@opengeni/contracts";
+
+export function accountRouteAuthNeededPayload(
+  payload: ToolAuthNeededPayload,
+  bindings: readonly McpConnectionAccountBinding[] | null | undefined,
+): ToolAuthNeededPayload {
+  const binding = bindings?.find((candidate) => candidate.serverId === payload.serverId);
+  return binding
+    ? {
+        ...payload,
+        canonicalServerId: binding.canonicalServerId,
+        connectionSubjectScope: binding.subjectScope,
+      }
+    : payload;
+}
+
+export { expandMcpAccountRoutes, expandApiIntegrationAccountRoutes } from "@opengeni/core";