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

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

@@ -1,3 +1,6 @@
+import { withDirectModelProviders } from "@opengeni/config";
+import { loadDirectModelProviderConnection } from "@opengeni/db";
+import { FILESYSTEM_DISCONTINUITY_PROTOCOL } from "./recovery-warning";
 import {
   applySessionTurnSettlement,
   claimSessionWorkForAttempt,
@@ -14,9 +17,11 @@ import {
   type SessionTurnForExecution,
 } from "@opengeni/db";
 import { appendAndPublishTurnEventsFenced, publishDurableSessionEvents } from "@opengeni/events";
+import { linkCurrentSpanToAdmission, turnExecutionTelemetryKey } from "@opengeni/observability";
 import { deliverChildRequiresActionToParent } from "../parent-wake";
 import {
   assertTurnExecutionPolicyMatchesConfigV1,
+  settingsForAcceptedSubscriptionTurn,
   resolveTurnExecutionPolicyV1,
   type Settings,
 } from "@opengeni/config";
@@ -47,18 +52,22 @@ import {
   recordSessionEventAppendPhase,
   recordSessionEventPublishLatency,
   recordTurnStartupPhase,
+  measureTurnStartupPhase,
   recordTurnStartupMilestone,
   turnLifecycleMetricsFor,
 } from "../../observability-metrics";
 import { createTurnCredentialLeases } from "./credential-leases";
 import { createTurnMediaArtifacts } from "./media-artifacts";
 import { readTurnExecutionPolicyV1 } from "@opengeni/contracts";
+import { turnCredentialRestriction } from "./credential-restriction";
 
 import {
   credentialSubjectIdForTurnInitiator,
   turnExecutionPolicyBillingIdentity,
   legacyTurnExecutionPolicyInput,
   ensureRunAllowed,
+  AllowanceExhaustedError,
+  type AllowanceRefusal,
 } from "./admission";
 import { providerRecoveryCountFromMetadata, isWorkerShutdownCancellation } from "./errors";
 import { throwIfTurnOperationCancelled, waitForTurnOperation } from "./sandbox-provision";
@@ -178,8 +187,17 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
     acknowledgeLostAttemptOwnership,
   } = deps;
 
-  const deploymentCatalogSettings = (await resolveCatalogSettings(db, catalogSourceSettings))
-    .settings;
+  const deploymentCatalogSettings = (
+    await measureTurnStartupPhase(
+      observability,
+      {
+        phase: "claim_catalog_read",
+        provider: "unresolved",
+        backend: "unresolved",
+      },
+      () => resolveCatalogSettings(db, catalogSourceSettings),
+    )
+  ).settings;
 
   const validatePendingSystemUpdateAuthority: NonNullable<
     ClaimSessionWorkForAttemptInput["validatePendingSystemUpdateAuthority"]
@@ -191,15 +209,25 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
       sessionId: input.sessionId,
       update,
     });
-  const claim = await claimSessionWorkForAttempt(db, input.workspaceId, {
-    sessionId: input.sessionId,
-    workflowId: input.workflowId,
-    workflowRunId: input.workflowRunId,
-    attemptId: input.attemptId,
-    dispatchId,
-    trigger: input.trigger,
-    validatePendingSystemUpdateAuthority,
-  });
+  const claim = await measureTurnStartupPhase(
+    observability,
+    {
+      phase: "claim_atomic",
+      provider: "unresolved",
+      backend: "unresolved",
+    },
+    () =>
+      claimSessionWorkForAttempt(db, input.workspaceId, {
+        filesystemDiscontinuityProtocol: FILESYSTEM_DISCONTINUITY_PROTOCOL,
+        sessionId: input.sessionId,
+        workflowId: input.workflowId,
+        workflowRunId: input.workflowRunId,
+        attemptId: input.attemptId,
+        dispatchId,
+        trigger: input.trigger,
+        validatePendingSystemUpdateAuthority,
+      }),
+  );
   if (claim.action === "unclaimed") {
     control.activityStatus = "unclaimed";
     return { exit: { status: "unclaimed", reason: claim.reason } };
@@ -209,6 +237,21 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   attempt.dispatchId = dispatchId;
   attempt.executionGeneration = turn.executionGeneration;
   attempt.providerRecoveryCount = providerRecoveryCountFromMetadata(turn.metadata);
+  const authRecovery = turn.metadata?.claudeAuthRecovery;
+  attempt.claudeAuthRecovery =
+    authRecovery &&
+    typeof authRecovery === "object" &&
+    "credentialId" in authRecovery &&
+    typeof authRecovery.credentialId === "string" &&
+    "credentialVersion" in authRecovery &&
+    typeof authRecovery.credentialVersion === "number" &&
+    Number.isSafeInteger(authRecovery.credentialVersion) &&
+    authRecovery.credentialVersion > 0
+      ? {
+          credentialId: authRecovery.credentialId,
+          credentialVersion: authRecovery.credentialVersion,
+        }
+      : undefined;
   attempt.triggerEventId = turn.triggerEventId;
   // The durable attempt UUID is stable for a Temporal retry of this activity
   // input and freshly generated for worker-death redispatch/continue-as-new.
@@ -223,25 +266,42 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   // Therefore every failure with no turnId came from the one atomic claim
   // transaction and can be classified without conflating ordinary runtime
   // or transport failures with admission failures.
-  const session = await requireSession(db, input.workspaceId, input.sessionId);
   let installedApiIntegrations: readonly ApiIntegrationRuntime[] = [];
   const credentialSubjectId = credentialSubjectIdForTurnInitiator(turn);
   const fileAuthoritySubjectId = turn.initiatingHumanSubjectId ?? null;
-  const mcpSettings = await settingsWithEnabledCapabilityMcpServers(
-    db,
-    input.workspaceId,
-    deploymentCatalogSettings,
-    {
-      ...(credentialSubjectId
-        ? { subjectId: credentialSubjectId }
-        : {
-            personalConnectionDelegations: turn.personalConnectionDelegations,
-          }),
-      onResolvedApiIntegrations: (integrations) => {
-        installedApiIntegrations = integrations;
+  // Both are fresh scoped reads on the root pool after exact claim ownership.
+  // Neither consumes the other's result; retain the capability helper's own
+  // subject/delegation authority and await both before credential/policy gates.
+  const [session, mcpSettings] = await Promise.all([
+    measureTurnStartupPhase(
+      observability,
+      {
+        phase: "claim_session_read",
+        provider: "unresolved",
+        backend: turn.sandboxBackend,
       },
-    },
-  );
+      () => requireSession(db, input.workspaceId, input.sessionId),
+    ),
+    measureTurnStartupPhase(
+      observability,
+      {
+        phase: "claim_capability_settings",
+        provider: "unresolved",
+        backend: turn.sandboxBackend,
+      },
+      () =>
+        settingsWithEnabledCapabilityMcpServers(db, input.workspaceId, deploymentCatalogSettings, {
+          ...(credentialSubjectId
+            ? { subjectId: credentialSubjectId }
+            : {
+                personalConnectionDelegations: turn.personalConnectionDelegations,
+              }),
+          onResolvedApiIntegrations: (integrations) => {
+            installedApiIntegrations = integrations;
+          },
+        }),
+    ),
+  ]);
   // Read the active-credential flag once for the runtime capability overlay.
   // Accepted billing/provider identity comes from the turn policy below,
   // never from this mutable health snapshot.
@@ -249,6 +309,7 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
     db,
     mcpSettings,
     input.workspaceId,
+    turn.id,
   );
   const codexSettings = await settingsWithCodexCredential(
     db,
@@ -273,21 +334,42 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
     gatewaySettings,
     claimedPolicy.kind === "valid" ? claimedPolicy.policy.productModelId : turn.model,
   );
-  const capabilitySettings = await settingsWithOrganizationProviderCredentials(
+  let capabilitySettings = await settingsWithOrganizationProviderCredentials(
     db,
     input.accountId,
     input.workspaceId,
     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;
-  runtime.configure(capabilitySettings);
-  const policyForAbsent =
+  const candidatePolicy =
     claimedPolicy.kind === "valid"
       ? claimedPolicy.policy
       : resolveTurnExecutionPolicyV1(capabilitySettings, legacyTurnExecutionPolicyInput(turn));
+  // This context is frozen by the accepted-turn writer from the exact source
+  // turn. Its reserved restriction field cannot come from public service JSON.
+  const credentialRestriction =
+    claimedPolicy.kind === "absent"
+      ? turn.initiatorContext?.credentialRestriction === "developer_setup"
+        ? "developer_setup"
+        : turnCredentialRestriction(candidatePolicy, session.metadata)
+      : undefined;
+  const policyForAbsent = credentialRestriction
+    ? { ...candidatePolicy, credentialRestriction }
+    : candidatePolicy;
   const installedPolicy = await installOrReadTurnExecutionPolicyForAttempt(db, {
     accountId: input.accountId,
     workspaceId: input.workspaceId,
@@ -300,6 +382,16 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   if (!installedPolicy.accepted) {
     throw new TurnAttemptFencedError(`turn execution policy was fenced: ${installedPolicy.reason}`);
   }
+  capabilitySettings = settingsForAcceptedSubscriptionTurn(
+    capabilitySettings,
+    installedPolicy.policy,
+    {
+      modelId: turn.model,
+      reasoningEffort: turn.reasoningEffort,
+      latencyMode: turn.latencyMode,
+    },
+  );
+  runtime.configure(capabilitySettings);
   const verifiedExecutionPolicy = assertTurnExecutionPolicyMatchesConfigV1(
     capabilitySettings,
     installedPolicy.policy,
@@ -317,10 +409,14 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   billingState.countsTowardTokenCap = billingIdentity.countsTowardTokenCap;
   billingState.isCodexTurn = billingIdentity.codexSubscription;
   billingState.isXaiTurn = billingIdentity.xaiSubscription;
+  billingState.isClaudeTurn =
+    verifiedExecutionPolicy.provider.kind === "claude-subscription-workspace" ||
+    verifiedExecutionPolicy.provider.kind === "claude-subscription-organization";
   const trigger = await getSessionEvent(db, input.workspaceId, attempt.triggerEventId);
   if (!trigger) {
     throw new Error(`Trigger event not found: ${attempt.triggerEventId}`);
   }
+  if (trigger.type === "user.message") linkCurrentSpanToAdmission(trigger.id);
   const humanInputResume = await getHumanInputResumeForEvent(
     db,
     input.workspaceId,
@@ -336,6 +432,7 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   attempt.triggerType = trigger.type;
   const attachPendingUpdatesAfterOpenSuffix = async (): Promise<boolean> => {
     const attached = await claimSessionWorkForAttempt(db, input.workspaceId, {
+      filesystemDiscontinuityProtocol: FILESYSTEM_DISCONTINUITY_PROTOCOL,
       sessionId: input.sessionId,
       workflowId: input.workflowId,
       workflowRunId: input.workflowRunId,
@@ -352,20 +449,27 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   // bypass OpenGeni credit/token gates)
   // AND the optional host `entitlements` port (when bound, its admitRun replaces
   // the local credit read). Unset port → today's local-ledger path.
-  await waitForTurnOperation(
-    ensureRunAllowed(
-      capabilitySettings,
-      db,
-      input.accountId,
-      input.workspaceId,
-      billingState.isExternallyBilledTurn,
-      entitlements,
-      billingState.chargesOpenGeniCredits,
-      billingState.countsTowardTokenCap,
-    ),
-    cancellationSignal,
-    undefined,
-  );
+  let allowanceRefusal: AllowanceRefusal | null = null;
+  try {
+    await waitForTurnOperation(
+      ensureRunAllowed(
+        capabilitySettings,
+        db,
+        input.accountId,
+        input.workspaceId,
+        billingState.isExternallyBilledTurn,
+        entitlements,
+        billingState.chargesOpenGeniCredits,
+        billingState.countsTowardTokenCap,
+        turn.initiatingHumanSubjectId,
+      ),
+      cancellationSignal,
+      undefined,
+    );
+  } catch (error) {
+    if (!(error instanceof AllowanceExhaustedError)) throw error;
+    allowanceRefusal = error.refusal;
+  }
   // Setup (variableSet load, MCP connects, sandbox restore) does not
   // stream and so never observes cancellation on its own; these explicit
   // checks let a graceful shutdown checkpoint the turn before the worker is
@@ -387,6 +491,7 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
     turnId: attempt.turnId,
     opAcks: {},
   };
+  eventing.heartbeatDetails = heartbeatDetails;
   const opJournal = makeTurnOpJournal(activityContext, heartbeatDetails);
   eventing.heartbeatTimer = startActivityHeartbeat(activityContext, heartbeatDetails);
   let producerSeq = 0;
@@ -497,6 +602,9 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
       ...(inputSettlement.suppressGoalContinuation !== undefined
         ? { suppressGoalContinuation: inputSettlement.suppressGoalContinuation }
         : {}),
+      ...(inputSettlement.allowanceGoalPause
+        ? { allowanceGoalPause: inputSettlement.allowanceGoalPause }
+        : {}),
       events: inputs,
       ...(runState ? { runState } : {}),
       ...(compactionRequestFailure ? { compactionRequestFailure } : {}),
@@ -547,12 +655,47 @@ export async function claimTurnAttempt(deps: ClaimTurnDeps): Promise<ClaimTurnOu
   throwIfTurnCancelled();
   recordTurnStartupPhase(observability, {
     phase: "claim_and_policy",
+    executionCorrelationId: turnExecutionTelemetryKey(
+      input.workspaceId,
+      input.sessionId,
+      input.attemptId,
+    ),
     provider: turnExecutionPolicy.providerId,
     backend: turn.sandboxBackend,
     outcome: "completed",
     durationSeconds: (performance.now() - activityStarted) / 1_000,
   });
   const turnStartSettlementStartedAt = performance.now();
+  if (allowanceRefusal) {
+    // Claim has frozen the exact human, but no provider or sandbox has started.
+    // Use the same terminal valve as a post-response stop, retaining a usable
+    // session and a visible typed refusal instead of manufacturing a failure.
+    if (
+      !(await eventing.settle({
+        events: [
+          { type: "usage.exhausted", payload: allowanceRefusal },
+          {
+            type: "turn.completed",
+            payload: {
+              output: "",
+              segmentLimit: "budget_exhausted",
+              ...allowanceRefusal,
+            },
+          },
+          { type: "session.status.changed", payload: { status: "idle" } },
+        ],
+        turnStatus: "completed",
+        sessionStatus: "idle",
+        activeTurnId: null,
+        allowanceGoalPause: { rationale: allowanceRefusal.message },
+      }))
+    ) {
+      return { exit: claimedResult({ status: "cancelled" }) };
+    }
+    control.turnMetricOutcome = "completed";
+    control.activityStatus = "idle";
+    return { exit: claimedResult({ status: "idle" }) };
+  }
   if (
     !(await eventing.settle({
       events: [

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

@@ -0,0 +1,221 @@
+import { createHash } from "node:crypto";
+import {
+  emptyClaudeUsage,
+  mergeClaudeUsage,
+  parseClaudeUsageHeaders,
+  type parseModelProvidersJson,
+  type ClaudeUsageObservation,
+} 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;
+  observation?: ClaudeUsageObservation;
+  responseStatus?: number;
+  upstreamModelId?: string;
+  requestId?: string;
+  refresh?: { status: "reconnect"; checkedAt: string };
+};
+
+/** Capture the exact credential before requests; replacement fences late responses. */
+export async function createClaudeUsageObserver(
+  providers: ReturnType<typeof parseModelProvidersJson>,
+  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(
+        (provider) =>
+          provider.kind === "claude-subscription-workspace" ||
+          provider.kind === "claude-subscription-organization",
+      )
+      .map(async (provider) => {
+        const scope: Scope =
+          provider.kind === "claude-subscription-workspace" ? "workspace" : "organization";
+        const binding = provider.anthropic?.credentialBinding;
+        if (binding && provider.apiKey)
+          return [
+            provider.id,
+            {
+              scope,
+              token: provider.apiKey,
+              expectedConnectionId: binding.connectionId,
+              expectedCredentialVersion: binding.credentialVersion,
+            },
+          ] as const;
+        const credential = await readCredential(scope).catch(() => null);
+        if (!credential || credential.token !== provider.apiKey) return null;
+        return [
+          provider.id,
+          {
+            scope,
+            token: credential.token,
+            expectedConnectionId: credential.connectionId,
+            expectedCredentialVersion: credential.credentialVersion,
+          },
+        ] as const;
+      }),
+  );
+  const captured = new Map(bindings.filter((binding) => binding !== null));
+  let anonymousReceipt = 0;
+  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, ...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 rawRequestId = response.headers.get("request-id");
+    const requestId = rawRequestId && rawRequestId.length <= 256 ? rawRequestId : undefined;
+    // Retain exact physical responses, including HTTP 200 streams which can
+    // terminate with typed authentication/rate-limit errors. Parallel titles
+    // must never overwrite the serving request receipt.
+    const receiptKey =
+      requestId ??
+      (response.status === 401 || response.status === 429
+        ? `anonymous-${++anonymousReceipt}`
+        : "aggregate");
+    const captureKey = `${receiptKey}:${scope}:${identity.expectedConnectionId}:${identity.expectedCredentialVersion}:${tokenKey}:${upstreamModelId ?? "unknown"}`;
+    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),
+        observation,
+      );
+      observation = {
+        windows: merged.windows,
+        observedAt: merged.observedAt!,
+        source: merged.source!,
+        requestStatus: merged.requestStatus ?? null,
+        requestRestrictions: merged.requestRestrictions ?? [],
+      };
+    }
+    if (requestId || observation || response.status === 401 || response.status === 429)
+      latest.set(captureKey, {
+        scope,
+        responseStatus: response.status,
+        ...(response.headers.get("request-id") && response.headers.get("request-id")!.length <= 256
+          ? { requestId: response.headers.get("request-id")! }
+          : {}),
+        ...(upstreamModelId ? { upstreamModelId } : {}),
+        ...identity,
+        ...(previous?.observation && !observation ? { observation: previous.observation } : {}),
+        ...(observation ? { observation } : {}),
+        ...(response.status === 401
+          ? { refresh: { status: "reconnect", checkedAt: new Date().toISOString() } }
+          : {}),
+      });
+  };
+  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) {
+      const binding = captured.get(providerId);
+      return binding ? { ...binding } : undefined;
+    },
+    renew(
+      providerId: string,
+      credential: {
+        token: string;
+        connectionId: string;
+        credentialVersion: number;
+      },
+    ) {
+      const binding = captured.get(providerId);
+      if (
+        binding &&
+        binding.expectedConnectionId === credential.connectionId &&
+        binding.expectedCredentialVersion === credential.credentialVersion
+      )
+        captured.set(providerId, { ...binding, token: credential.token });
+    },
+  });
+}