@opengeni/worker-bundle 2.3.1 → 2.3.2-canary.37090756480001

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@opengeni/worker-bundle",
-  "version": "2.3.1",
+  "version": "2.3.2-canary.37090756480001",
   "description": "OpenGeni worker entry and reusable embedded lifecycle, shipped with a release-coherent pre-bundled Temporal workflow artifact.",
   "license": "Apache-2.0",
   "repository": {
@@ -54,26 +54,26 @@
   },
   "dependencies": {
     "@llamaindex/liteparse": "2.14.2",
-    "@opengeni/agent-proto": "^0.6.1",
-    "@opengeni/capabilities": "^0.3.5",
-    "@opengeni/codemode": "^0.6.8",
-    "@opengeni/codex": "^0.2.31",
-    "@opengeni/config": "^3.2.1",
-    "@opengeni/contracts": "^5.6.0",
-    "@opengeni/core": "^5.1.1",
-    "@opengeni/db": "^6.5.0",
-    "@opengeni/documents": "^0.8.40",
-    "@opengeni/events": "^0.4.38",
-    "@opengeni/github": "^0.8.2",
-    "@opengeni/interaction": "^0.4.45",
-    "@opengeni/jev": "^0.2.1",
-    "@opengeni/network": "^0.3.1",
-    "@opengeni/observability": "^0.8.37",
-    "@opengeni/runtime": "^4.6.0",
-    "@opengeni/sdk": "^7.6.1",
-    "@opengeni/storage": "^0.2.138",
-    "@opengeni/tool-gateway": "^0.1.18",
-    "@opengeni/xai-subscription": "^0.1.5",
+    "@opengeni/agent-proto": "^0.6.1-canary.37090756480001",
+    "@opengeni/capabilities": "^0.3.6-canary.37090756480001",
+    "@opengeni/codemode": "^0.7.0-canary.37090756480001",
+    "@opengeni/codex": "^0.2.32-canary.37090756480001",
+    "@opengeni/config": "^3.3.0-canary.37090756480001",
+    "@opengeni/contracts": "^5.7.0-canary.37090756480001",
+    "@opengeni/core": "^5.2.0-canary.37090756480001",
+    "@opengeni/db": "^6.6.0-canary.37090756480001",
+    "@opengeni/documents": "^0.8.41-canary.37090756480001",
+    "@opengeni/events": "^0.4.39-canary.37090756480001",
+    "@opengeni/github": "^0.8.3-canary.37090756480001",
+    "@opengeni/interaction": "^0.4.46-canary.37090756480001",
+    "@opengeni/jev": "^0.2.1-canary.37090756480001",
+    "@opengeni/network": "^0.3.1-canary.37090756480001",
+    "@opengeni/observability": "^0.8.38-canary.37090756480001",
+    "@opengeni/runtime": "^4.7.0-canary.37090756480001",
+    "@opengeni/sdk": "^7.7.0-canary.37090756480001",
+    "@opengeni/storage": "^0.2.139-canary.37090756480001",
+    "@opengeni/tool-gateway": "^0.1.19-canary.37090756480001",
+    "@opengeni/xai-subscription": "^0.1.5-canary.37090756480001",
     "@temporalio/activity": "^1.17.0",
     "@temporalio/client": "^1.17.0",
     "@temporalio/worker": "^1.17.0",

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

@@ -1,3 +1,5 @@
+import { withDirectModelProviders } from "@opengeni/config";
+import { loadDirectModelProviderConnection } from "@opengeni/db";
 import { FILESYSTEM_DISCONTINUITY_PROTOCOL } from "./recovery-warning";
 import {
   applySessionTurnSettlement,
@@ -288,6 +290,17 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
     workspaceProviderSettings,
     claimedPolicy.kind === "valid" ? claimedPolicy.policy.productModelId : turn.model,
   );
+  const selectedDirectConnection = await loadDirectModelProviderConnection(
+    db,
+    capabilitySettings,
+    input.workspaceId,
+    claimedPolicy.kind === "valid" ? claimedPolicy.policy.productModelId : (turn.model ?? ""),
+  );
+  // Execution only needs the selected customer connection. Ordinary turns
+  // must not load or install unrelated workspace provider configurations.
+  if (selectedDirectConnection) {
+    capabilitySettings = withDirectModelProviders(capabilitySettings, [selectedDirectConnection]);
+  }
   const codexAppsCredentialId = capabilitySettings.codexConnectedAppsEnabled
     ? await resolveCodexAppsCredentialIdForRun(db, input.workspaceId)
     : null;

package/src/activities/agent-turn/claude-usage-observer.ts

@@ -1,3 +1,4 @@
+import { createHash } from "node:crypto";
 import {
   emptyClaudeUsage,
   mergeClaudeUsage,
@@ -7,7 +8,22 @@ import {
 } from "@opengeni/config";
 
 type Scope = "workspace" | "organization";
+type ClaudeRequestCredential = {
+  token: string;
+  connectionId: string;
+  credentialVersion: number;
+};
+export class ClaudeSubscriptionConnectionUnavailable extends Error {
+  readonly status = 409;
+  readonly code = "claude_subscription_connection_changed";
+  constructor() {
+    super(
+      "Claude connection changed or was disconnected. Start a new turn with the current connection.",
+    );
+  }
+}
 export type CapturedClaudeUsage = {
+  scope: Scope;
   token: string;
   expectedConnectionId: string;
   expectedCredentialVersion: number;
@@ -18,13 +34,22 @@ export type CapturedClaudeUsage = {
 /** Capture the exact credential before requests; replacement fences late responses. */
 export async function createClaudeUsageObserver(
   providers: ReturnType<typeof parseModelProvidersJson>,
-  latest: Map<Scope, CapturedClaudeUsage>,
+  latest: Map<string, CapturedClaudeUsage>,
   readCredential: (scope: Scope) => Promise<{
     token: string;
     connectionId: string;
     credentialVersion: number;
   } | null>,
 ) {
+  const managedProviderIds = new Set(
+    providers
+      .filter(
+        (provider) =>
+          provider.kind === "claude-subscription-workspace" ||
+          provider.kind === "claude-subscription-organization",
+      )
+      .map((provider) => provider.id),
+  );
   const bindings = await Promise.all(
     providers
       .filter(
@@ -60,12 +85,27 @@ export async function createClaudeUsageObserver(
       }),
   );
   const captured = new Map(bindings.filter((binding) => binding !== null));
-  const observe = (providerId: string, response: Response) => {
-    const binding = captured.get(providerId);
+  const observeBinding = (
+    binding:
+      | Pick<
+          CapturedClaudeUsage,
+          "scope" | "token" | "expectedConnectionId" | "expectedCredentialVersion"
+        >
+      | undefined,
+    response: Response,
+    upstreamModelId?: string,
+    requestToken?: string | null,
+  ) => {
+    if (requestToken === null) return;
     if (!binding) return;
-    const { scope, ...identity } = binding;
-    const previous = latest.get(scope);
-    let observation = parseClaudeUsageHeaders(response.headers);
+    const { scope, ...capturedIdentity } = binding;
+    const identity = { ...capturedIdentity, token: requestToken ?? capturedIdentity.token };
+    // Same-generation OAuth renewal can overlap an older request. Never attach
+    // its authentication failure to the newly renewed token or merge the two.
+    const tokenKey = createHash("sha256").update(identity.token).digest("hex");
+    const captureKey = `${scope}:${identity.expectedConnectionId}:${identity.expectedCredentialVersion}:${tokenKey}`;
+    const previous = latest.get(captureKey);
+    let observation = parseClaudeUsageHeaders(response.headers, new Date(), upstreamModelId);
     if (observation && previous?.observation) {
       const merged = mergeClaudeUsage(
         mergeClaudeUsage(emptyClaudeUsage(binding.expectedCredentialVersion), previous.observation),
@@ -75,10 +115,13 @@ export async function createClaudeUsageObserver(
         windows: merged.windows,
         observedAt: merged.observedAt!,
         source: merged.source!,
+        requestStatus: merged.requestStatus ?? null,
+        requestRestrictions: merged.requestRestrictions ?? [],
       };
     }
     if (observation || response.status === 401)
-      latest.set(scope, {
+      latest.set(captureKey, {
+        scope,
         ...identity,
         ...(previous?.observation && !observation ? { observation: previous.observation } : {}),
         ...(observation ? { observation } : {}),
@@ -87,9 +130,57 @@ export async function createClaudeUsageObserver(
           : {}),
       });
   };
+  const observe = (
+    providerId: string,
+    response: Response,
+    upstreamModelId?: string,
+    requestToken?: string | null,
+  ) => observeBinding(captured.get(providerId), response, upstreamModelId, requestToken);
+  const prepareRequestWithObserver = async (
+    providerId: string,
+    headers: Headers,
+    resolve: (binding: {
+      scope: Scope;
+      expectedConnectionId: string;
+      expectedCredentialVersion: number;
+    }) => Promise<ClaudeRequestCredential | null>,
+  ) => {
+    if (!managedProviderIds.has(providerId)) return { headers, observe };
+    const binding = captured.get(providerId);
+    if (!binding) throw new ClaudeSubscriptionConnectionUnavailable();
+    const credential = await resolve({
+      scope: binding.scope,
+      expectedConnectionId: binding.expectedConnectionId,
+      expectedCredentialVersion: binding.expectedCredentialVersion,
+    });
+    if (
+      !credential ||
+      credential.connectionId !== binding.expectedConnectionId ||
+      credential.credentialVersion !== binding.expectedCredentialVersion
+    )
+      throw new ClaudeSubscriptionConnectionUnavailable();
+    const dispatched = Object.freeze({ ...binding, token: credential.token });
+    captured.set(providerId, dispatched);
+    headers.set("authorization", `Bearer ${credential.token}`);
+    headers.delete("x-api-key");
+    return {
+      headers,
+      observe: (
+        _providerId: string,
+        response: Response,
+        upstreamModelId?: string,
+        requestToken?: string | null,
+      ) => observeBinding(dispatched, response, upstreamModelId, requestToken),
+    };
+  };
   return Object.assign(observe, {
+    prepareRequestWithObserver,
+    async prepareRequest(...args: Parameters<typeof prepareRequestWithObserver>) {
+      return (await prepareRequestWithObserver(...args)).headers;
+    },
     binding(providerId: string) {
-      return captured.get(providerId);
+      const binding = captured.get(providerId);
+      return binding ? { ...binding } : undefined;
     },
     renew(
       providerId: string,

package/src/activities/agent-turn/credential-leases.ts

@@ -5,12 +5,14 @@ import {
   heartbeatXaiCredentialLeaseUntil,
 } from "@opengeni/db";
 import type { SharedActivityServices } from "../types";
-import { codexCredentialLeaseDeadlineExpired } from "./codex";
+import {
+  SubscriptionTurnLease,
+  type LeaseRenewReason,
+  type LeaseLossReason,
+} from "./subscription-lease";
+export type { LeaseRenewReason, LeaseLossReason } from "./subscription-lease";
 import { safeErrorDiagnostic } from "./errors";
 
-export type LeaseRenewReason = "timer" | "runtime_event" | "model_usage";
-export type LeaseLossReason = "deadline" | "not_found";
-
 export class CodexCredentialLeaseLostError extends Error {
   readonly code = "codex_credential_lease_lost";
 
@@ -34,223 +36,94 @@ export type TurnCredentialLeaseDeps = {
  * cross-replica fairness primitive; the heartbeat here only extends its short
  * TTL. A killed worker stops heartbeating and the holder self-expires.
  */
-export class CodexTurnLease {
-  held = false;
-  lost = false;
-  lossReason: LeaseLossReason | null = null;
-  holderId: string | null = null;
-  generation: number | null = null;
-  // Monotonic worker deadline, not a comparison between the Postgres and
-  // worker wall clocks. It is advanced only after a database renewal confirms,
-  // from the request START + TTL; slow queries therefore shorten (never extend)
-  // the conservative ownership window.
-  confirmedUntilMs: number | null = null;
-  private heartbeatInFlight = false;
-  private heartbeatTimer: ReturnType<typeof setInterval> | undefined;
-
-  constructor(private readonly deps: TurnCredentialLeaseDeps) {}
-
-  markLost = (reason: LeaseLossReason): void => {
-    if (this.lost) return;
-    this.lost = true;
-    this.lossReason = reason;
-    this.deps.observability.incrementCounter({
-      name: "opengeni_codex_lease_renewals_total",
-      help: "Codex lease renewal checkpoints by outcome and reason.",
-      labels: { workspace_key: this.deps.codexWorkspaceKey, outcome: "lost", reason },
-    });
-    this.deps.observability.warn("Codex credential lease was lost during an active turn", {
-      workspaceId: this.deps.workspaceId,
-      turnId: this.deps.getTurnId(),
-      reason,
-    });
-  };
-
-  /** Fail closed at the last worker-confirmed lease deadline before dispatch. */
-  assertUsable = (): void => {
-    if (this.lost) {
-      throw new CodexCredentialLeaseLostError(this.lossReason ?? "not_found");
-    }
-    if (
-      !this.held ||
-      this.holderId === null ||
-      this.generation === null ||
-      codexCredentialLeaseDeadlineExpired(this.confirmedUntilMs)
-    ) {
-      this.markLost("deadline");
-      throw new CodexCredentialLeaseLostError("deadline");
-    }
-  };
-
-  renew = async (reason: LeaseRenewReason): Promise<void> => {
-    const turnId = this.deps.getTurnId();
-    if (!turnId || !this.held || !this.holderId || this.generation === null || this.lost) {
-      return;
-    }
-    if (codexCredentialLeaseDeadlineExpired(this.confirmedUntilMs)) {
-      this.markLost("deadline");
-      return;
-    }
-    if (this.heartbeatInFlight) return;
-    this.heartbeatInFlight = true;
-    const priorConfirmedUntilMs = this.confirmedUntilMs;
-    const renewalStartedAtMs = performance.now();
-    try {
-      const renewedUntil = await heartbeatCodexCredentialLeaseUntil(
-        this.deps.db,
-        this.deps.accountId,
-        this.deps.workspaceId,
-        turnId,
-        this.holderId,
-        this.generation,
-        CODEX_CREDENTIAL_LEASE_TTL_MS,
-      );
-      if (!renewedUntil) {
-        this.markLost("not_found");
-      } else if (
-        codexCredentialLeaseDeadlineExpired(priorConfirmedUntilMs) ||
-        codexCredentialLeaseDeadlineExpired(renewalStartedAtMs + CODEX_CREDENTIAL_LEASE_TTL_MS)
-      ) {
-        // The database call may have waited on a lock until the prior lease
-        // expired. A late successful UPDATE is not proof that this worker
-        // owned the credential continuously; do not accept or extend it.
-        this.markLost("deadline");
-      } else {
-        this.confirmedUntilMs = renewalStartedAtMs + CODEX_CREDENTIAL_LEASE_TTL_MS;
-        this.deps.observability.incrementCounter({
+export class CodexTurnLease extends SubscriptionTurnLease {
+  constructor(deps: TurnCredentialLeaseDeps) {
+    super({
+      ttlMs: CODEX_CREDENTIAL_LEASE_TTL_MS,
+      getTurnId: deps.getTurnId,
+      heartbeat: ({ turnId, holderId, generation }) =>
+        heartbeatCodexCredentialLeaseUntil(
+          deps.db,
+          deps.accountId,
+          deps.workspaceId,
+          turnId,
+          holderId,
+          generation,
+          CODEX_CREDENTIAL_LEASE_TTL_MS,
+        ),
+      lostError: (reason) => new CodexCredentialLeaseLostError(reason),
+      onLost: (reason) => {
+        deps.observability.incrementCounter({
           name: "opengeni_codex_lease_renewals_total",
           help: "Codex lease renewal checkpoints by outcome and reason.",
-          labels: {
-            workspace_key: this.deps.codexWorkspaceKey,
-            outcome: "completed",
-            reason,
-          },
+          labels: { workspace_key: deps.codexWorkspaceKey, outcome: "lost", reason },
         });
-      }
-    } catch (error) {
-      if (codexCredentialLeaseDeadlineExpired(this.confirmedUntilMs)) {
-        this.markLost("deadline");
-        return;
-      }
-      this.deps.observability.warn("Codex credential lease heartbeat failed", {
-        workspaceId: this.deps.workspaceId,
-        turnId,
-        reason,
-        ...safeErrorDiagnostic(error),
-      });
-      this.deps.observability.incrementCounter({
-        name: "opengeni_codex_lease_renewals_total",
-        help: "Codex lease renewal checkpoints by outcome and reason.",
-        labels: {
-          workspace_key: this.deps.codexWorkspaceKey,
-          outcome: "error",
+        deps.observability.warn("Codex credential lease was lost during an active turn", {
+          workspaceId: deps.workspaceId,
+          turnId: deps.getTurnId(),
           reason,
-        },
-      });
-    } finally {
-      this.heartbeatInFlight = false;
-    }
-  };
-
-  startHeartbeat = (): void => {
-    if (!this.deps.getTurnId() || this.heartbeatTimer) return;
-    this.heartbeatTimer = setInterval(() => {
-      void this.renew("timer");
-    }, 60_000);
-    this.heartbeatTimer.unref?.();
-  };
-
-  stopHeartbeat = (): void => {
-    if (!this.heartbeatTimer) return;
-    clearInterval(this.heartbeatTimer);
-    this.heartbeatTimer = undefined;
-  };
+        });
+      },
+      onRenewed: (reason) =>
+        deps.observability.incrementCounter({
+          name: "opengeni_codex_lease_renewals_total",
+          help: "Codex lease renewal checkpoints by outcome and reason.",
+          labels: { workspace_key: deps.codexWorkspaceKey, outcome: "completed", reason },
+        }),
+      onError: (error, reason) => {
+        deps.observability.warn("Codex credential lease heartbeat failed", {
+          workspaceId: deps.workspaceId,
+          turnId: deps.getTurnId(),
+          reason,
+          ...safeErrorDiagnostic(error),
+        });
+        deps.observability.incrementCounter({
+          name: "opengeni_codex_lease_renewals_total",
+          help: "Codex lease renewal checkpoints by outcome and reason.",
+          labels: { workspace_key: deps.codexWorkspaceKey, outcome: "error", reason },
+        });
+      },
+    });
+  }
 }
 
-/** The xAI/SuperGrok credential holder for one running turn; same shape. */
-export class XaiTurnLease {
-  held = false;
-  lost = false;
+/** SuperGrok uses the same deadline and holder fences as Codex. */
+export class XaiTurnLease extends SubscriptionTurnLease {
   subjectId: string | null = null;
-  holderId: string | null = null;
-  generation: number | null = null;
-  confirmedUntilMs: number | null = null;
-  private heartbeatInFlight = false;
-  private heartbeatTimer: ReturnType<typeof setInterval> | undefined;
-
-  constructor(private readonly deps: TurnCredentialLeaseDeps) {}
-
-  markLost = (reason: LeaseLossReason): void => {
-    if (this.lost) return;
-    this.lost = true;
-    this.deps.observability.warn("xAI credential lease was lost during an active turn", {
-      workspaceId: this.deps.workspaceId,
-      turnId: this.deps.getTurnId(),
-      reason,
+  constructor(deps: TurnCredentialLeaseDeps) {
+    super({
+      ttlMs: XAI_CREDENTIAL_LEASE_TTL_MS,
+      getTurnId: deps.getTurnId,
+      heartbeat: ({ turnId, holderId, generation }) =>
+        this.subjectId
+          ? heartbeatXaiCredentialLeaseUntil(deps.db, {
+              workspaceId: deps.workspaceId,
+              subjectId: this.subjectId,
+              turnId,
+              holderId,
+              generation,
+              leaseTtlMs: XAI_CREDENTIAL_LEASE_TTL_MS,
+            })
+          : Promise.resolve(null),
+      lostError: (reason) =>
+        Object.assign(new Error("SuperGrok credential lease is not usable for provider dispatch"), {
+          code: "xai_credential_lease_lost",
+          reason,
+        }),
+      onLost: (reason) =>
+        deps.observability.warn("xAI credential lease was lost during an active turn", {
+          workspaceId: deps.workspaceId,
+          turnId: deps.getTurnId(),
+          reason,
+        }),
+      onError: (error) =>
+        deps.observability.warn("xAI credential lease heartbeat failed", {
+          workspaceId: deps.workspaceId,
+          turnId: deps.getTurnId(),
+          ...safeErrorDiagnostic(error),
+        }),
     });
-  };
-
-  renew = async (): Promise<void> => {
-    const turnId = this.deps.getTurnId();
-    if (
-      !turnId ||
-      !this.held ||
-      !this.subjectId ||
-      !this.holderId ||
-      this.generation === null ||
-      this.lost
-    ) {
-      return;
-    }
-    if (codexCredentialLeaseDeadlineExpired(this.confirmedUntilMs)) {
-      this.markLost("deadline");
-      return;
-    }
-    if (this.heartbeatInFlight) return;
-    this.heartbeatInFlight = true;
-    const renewalStartedAtMs = performance.now();
-    try {
-      const renewedUntil = await heartbeatXaiCredentialLeaseUntil(this.deps.db, {
-        workspaceId: this.deps.workspaceId,
-        subjectId: this.subjectId,
-        turnId,
-        holderId: this.holderId,
-        generation: this.generation,
-        leaseTtlMs: XAI_CREDENTIAL_LEASE_TTL_MS,
-      });
-      if (!renewedUntil) {
-        this.markLost("not_found");
-      } else {
-        this.confirmedUntilMs = renewalStartedAtMs + XAI_CREDENTIAL_LEASE_TTL_MS;
-      }
-    } catch (error) {
-      if (codexCredentialLeaseDeadlineExpired(this.confirmedUntilMs)) {
-        this.markLost("deadline");
-        return;
-      }
-      this.deps.observability.warn("xAI credential lease heartbeat failed", {
-        workspaceId: this.deps.workspaceId,
-        turnId,
-        ...safeErrorDiagnostic(error),
-      });
-    } finally {
-      this.heartbeatInFlight = false;
-    }
-  };
-
-  startHeartbeat = (): void => {
-    if (!this.deps.getTurnId() || this.heartbeatTimer) return;
-    this.heartbeatTimer = setInterval(() => {
-      void this.renew();
-    }, 60_000);
-    this.heartbeatTimer.unref?.();
-  };
-
-  stopHeartbeat = (): void => {
-    if (!this.heartbeatTimer) return;
-    clearInterval(this.heartbeatTimer);
-    this.heartbeatTimer = undefined;
-  };
+  }
 }
 
 /** Both serving-credential leases for one turn attempt plus their composite views. */