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

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

@@ -1,16 +1,20 @@
 import {
   CODEX_CREDENTIAL_LEASE_TTL_MS,
   XAI_CREDENTIAL_LEASE_TTL_MS,
+  CLAUDE_CREDENTIAL_LEASE_TTL_MS,
+  heartbeatClaudeCredentialLeaseUntil,
   heartbeatCodexCredentialLeaseUntil,
   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,229 +38,128 @@ 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;
+/** Scoped subscriptions share the existing holder and monotonic deadline implementation. */
+class ScopedSubscriptionTurnLease 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,
+    provider: {
+      name: string;
+      code: string;
+      ttlMs: number;
+      heartbeat: typeof heartbeatXaiCredentialLeaseUntil;
+    },
+  ) {
+    super({
+      ttlMs: provider.ttlMs,
+      getTurnId: deps.getTurnId,
+      heartbeat: ({ turnId, holderId, generation }) =>
+        this.subjectId
+          ? provider.heartbeat(deps.db, {
+              workspaceId: deps.workspaceId,
+              subjectId: this.subjectId,
+              turnId,
+              holderId,
+              generation,
+              leaseTtlMs: provider.ttlMs,
+            })
+          : Promise.resolve(null),
+      lostError: (reason) =>
+        Object.assign(
+          new Error(provider.name + " credential lease is not usable for provider dispatch"),
+          { code: provider.code, reason },
+        ),
+      onLost: (reason) =>
+        deps.observability.warn(
+          provider.name + " credential lease was lost during an active turn",
+          { workspaceId: deps.workspaceId, turnId: deps.getTurnId(), reason },
+        ),
+      onError: (error) =>
+        deps.observability.warn(provider.name + " 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;
-  };
+  }
+}
+export class XaiTurnLease extends ScopedSubscriptionTurnLease {
+  constructor(deps: TurnCredentialLeaseDeps) {
+    super(deps, {
+      name: "SuperGrok",
+      code: "xai_credential_lease_lost",
+      ttlMs: XAI_CREDENTIAL_LEASE_TTL_MS,
+      heartbeat: heartbeatXaiCredentialLeaseUntil,
+    });
+  }
+}
+export class ClaudeTurnLease extends ScopedSubscriptionTurnLease {
+  constructor(deps: TurnCredentialLeaseDeps) {
+    super(deps, {
+      name: "Claude",
+      code: "claude_credential_lease_lost",
+      ttlMs: CLAUDE_CREDENTIAL_LEASE_TTL_MS,
+      heartbeat: heartbeatClaudeCredentialLeaseUntil,
+    });
+  }
 }
 
-/** Both serving-credential leases for one turn attempt plus their composite views. */
+/** Serving-credential leases for one turn attempt plus their composite views. */
 export type TurnCredentialLeases = {
   codex: CodexTurnLease;
   xai: XaiTurnLease;
+  claude: ClaudeTurnLease;
   renewServing: (reason: LeaseRenewReason) => Promise<void>;
   servingLost: () => boolean;
 };
@@ -264,13 +167,16 @@ export type TurnCredentialLeases = {
 export function createTurnCredentialLeases(deps: TurnCredentialLeaseDeps): TurnCredentialLeases {
   const codex = new CodexTurnLease(deps);
   const xai = new XaiTurnLease(deps);
+  const claude = new ClaudeTurnLease(deps);
   return {
     codex,
     xai,
+    claude,
     renewServing: async (reason) => {
       await codex.renew(reason);
-      await xai.renew();
+      await xai.renew(reason);
+      await claude.renew(reason);
     },
-    servingLost: () => codex.lost || xai.lost,
+    servingLost: () => codex.lost || xai.lost || claude.lost,
   };
 }

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

@@ -0,0 +1,14 @@
+import { readTurnExecutionPolicyV1, type TurnExecutionPolicyV1 } from "@opengeni/contracts";
+
+/** Only accepted policies establish the ceiling; ordinary metadata is not authority. */
+export function turnCredentialRestriction(
+  acceptedPolicy: TurnExecutionPolicyV1,
+  initialSessionMetadata: unknown,
+): "developer_setup" | undefined {
+  const initialPolicy = readTurnExecutionPolicyV1(initialSessionMetadata);
+  return acceptedPolicy.credentialRestriction === "developer_setup" ||
+    (initialPolicy.kind === "valid" &&
+      initialPolicy.policy.credentialRestriction === "developer_setup")
+    ? "developer_setup"
+    : undefined;
+}