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

package/src/activities/agent-turn/code-search.ts

@@ -0,0 +1,275 @@
+import type { AttemptToolDefinition } from "@opengeni/codemode";
+import { usableJevApiKey, type Settings } from "@opengeni/config";
+import {
+  CODE_SEARCH_TOOL_DESCRIPTION,
+  CODE_SEARCH_TOOL_NAME,
+  CodeSearchArgumentError,
+  CodeSearchRipgrepMissingError,
+  CodeSearchWorkspaceError,
+  JevCircuitBreaker,
+  JevClient,
+  JevRequestError,
+  JevUnavailableError,
+  codeSearchInputSchema,
+  parseCodeSearchArguments,
+  renderCodeSearchError,
+  runCodeSearch,
+  type CodeSearchWorkspace,
+  type JevCircuitLease,
+} from "@opengeni/jev";
+import type { Observability } from "@opengeni/observability";
+import {
+  ChannelANotFoundError,
+  ChannelAUnavailableError,
+  ChannelAUnsupportedError,
+  ChannelAValidationError,
+  isWindowsConnectedMachinePath,
+  type SandboxChannelAService,
+} from "@opengeni/runtime/sandbox";
+import { recordCodeSearchCall, type CodeSearchCallOutcome } from "../../observability-metrics";
+
+/**
+ * One breaker per worker process. Repeated Jev outages (including an exhausted
+ * account) make `code_search` calls on this worker fail at once for a cooldown
+ * instead of each running a full retry cycle. The breaker never changes which
+ * tools a turn is offered: the tool list and instructions are the start of the
+ * model's cached prompt, and sessions move between workers whose breakers
+ * disagree.
+ */
+export const codeSearchCircuitBreaker = new JevCircuitBreaker();
+
+/** Ripgrep output kept on the box per call before framing. */
+const CODE_SEARCH_RIPGREP_MAX_BYTES = 32 * 1024 * 1024;
+
+type CodeSearchChannel = Pick<
+  SandboxChannelAService,
+  "codeSearchRipgrep" | "codeSearchPathKinds" | "fsRead"
+>;
+
+/** Adapt the turn's sandbox or Connected Machine to the code search engine. */
+export function codeSearchWorkspaceFromChannel(channel: CodeSearchChannel): CodeSearchWorkspace {
+  return {
+    ripgrep: async (args, options) => {
+      options.signal?.throwIfAborted();
+      const outcome = await channel.codeSearchRipgrep(args, {
+        timeoutMs: options.timeoutMs,
+        maxBytes: CODE_SEARCH_RIPGREP_MAX_BYTES,
+      });
+      if (!outcome.available) throw new CodeSearchRipgrepMissingError();
+      return {
+        stdout: outcome.stdout,
+        exitCode: outcome.exitCode,
+        truncated: outcome.truncated,
+        timedOut: outcome.timedOut,
+      };
+    },
+    readText: async (path, options) => {
+      options.signal?.throwIfAborted();
+      try {
+        const read = await channel.fsRead({ path, encoding: "utf8", maxBytes: options.maxBytes });
+        return { text: read.content, truncated: read.truncated, binary: read.isBinary };
+      } catch (error) {
+        if (error instanceof ChannelANotFoundError) return null;
+        throw error;
+      }
+    },
+    pathKinds: async (paths, options) => {
+      options.signal?.throwIfAborted();
+      return await channel.codeSearchPathKinds(paths);
+    },
+  };
+}
+
+/** Jev work done by one completed `code_search` call, for per-workspace usage records. */
+export type CodeSearchUsage = {
+  operationId: string;
+  jevRequests: number;
+  jevInputTokens: number;
+  jevCostUsd: number;
+};
+
+function textResult(text: string, isError: boolean) {
+  return { isError, content: [{ type: "text" as const, text }] };
+}
+
+/**
+ * The model-facing `code_search` tool. Jev scores candidates inside the worker;
+ * the sandbox only runs read-only ripgrep and file reads, so the Jev key never
+ * leaves this process. A Jev failure is reported to the model instead of
+ * degrading to keyword-only ranking, which lowered answer quality in testing.
+ */
+export function createCodeSearchAttemptToolDefinition(input: {
+  settings: Pick<Settings, "jevApiKey" | "jevBaseUrl" | "jevModel" | "jevRequestTimeoutMs">;
+  apiKey: string;
+  workspace: () => Promise<CodeSearchWorkspace>;
+  observability: Observability;
+  /** Records Jev usage against the workspace. Failures are logged, never surfaced. */
+  recordUsage?: (usage: CodeSearchUsage) => Promise<void>;
+  breaker?: JevCircuitBreaker;
+  fetch?: typeof fetch;
+}): AttemptToolDefinition {
+  const breaker = input.breaker ?? codeSearchCircuitBreaker;
+  const jev = new JevClient({
+    apiKey: input.apiKey,
+    baseUrl: input.settings.jevBaseUrl,
+    model: input.settings.jevModel,
+    timeoutMs: input.settings.jevRequestTimeoutMs,
+    ...(input.fetch ? { fetch: input.fetch } : {}),
+  });
+  return {
+    identity: { serverId: "opengeni", toolName: CODE_SEARCH_TOOL_NAME },
+    modelName: CODE_SEARCH_TOOL_NAME,
+    codemodePath: ["opengeni", CODE_SEARCH_TOOL_NAME],
+    title: "Search code",
+    description: CODE_SEARCH_TOOL_DESCRIPTION,
+    inputSchema: codeSearchInputSchema,
+    annotations: {
+      title: "Search code",
+      readOnlyHint: true,
+      destructiveHint: false,
+      idempotentHint: true,
+      openWorldHint: false,
+    },
+    source: "opengeni",
+    approval: "none",
+    execute: async (args, context) => {
+      const startedAt = performance.now();
+      let outcome: CodeSearchCallOutcome = "failed";
+      let jevRequests = 0;
+      let jevCostUsd = 0;
+      // The breaker lease (the single trial while half-open) is settled exactly
+      // once: by what Jev did, or released when Jev was never judged. Settling
+      // takes it, so the finally block releases only a lease still held.
+      let lease: JevCircuitLease | null = null;
+      try {
+        const request = parseCodeSearchArguments(args);
+        lease = breaker.tryAcquire(Date.now());
+        if (!lease) {
+          outcome = "breaker_open";
+          return textResult(
+            renderCodeSearchError(
+              new JevUnavailableError("Jev is not responding; calls are paused for a few minutes"),
+            ),
+            true,
+          );
+        }
+        const workspace = await input.workspace();
+        const result = await runCodeSearch({
+          ...request,
+          workspace,
+          jev,
+          ...(context.signal ? { signal: context.signal } : {}),
+        });
+        const held = lease;
+        lease = null;
+        // A pack whose final status check hit an outage still counts against Jev.
+        if (result.statusCheckError instanceof JevUnavailableError) {
+          breaker.recordFailure(result.statusCheckError, Date.now(), held);
+        } else if (result.stats.jev.requests > 0) {
+          breaker.recordSuccess(held);
+        } else {
+          breaker.release(held);
+        }
+        outcome = "completed";
+        jevRequests = result.stats.jev.requests;
+        jevCostUsd = result.stats.jev.costUsd;
+        if (input.recordUsage && jevRequests > 0) {
+          await input
+            .recordUsage({
+              operationId: context.operationId,
+              jevRequests,
+              jevInputTokens: result.stats.jev.inputTokens,
+              jevCostUsd,
+            })
+            .catch((error: unknown) => {
+              input.observability.warn("code_search usage record failed", {
+                error: error instanceof Error ? error.message : String(error),
+              });
+            });
+        }
+        return textResult(result.text, false);
+      } catch (error) {
+        if (context.signal?.aborted) {
+          outcome = "cancelled";
+          throw error;
+        }
+        if (error instanceof CodeSearchArgumentError) {
+          outcome = "invalid_arguments";
+          return textResult(error.message, true);
+        }
+        if (error instanceof JevUnavailableError) {
+          const held = lease;
+          lease = null;
+          breaker.recordFailure(error, Date.now(), held);
+          outcome = "jev_unavailable";
+          return textResult(renderCodeSearchError(error), true);
+        }
+        if (error instanceof JevRequestError) {
+          outcome = "jev_rejected";
+          return textResult(renderCodeSearchError(error), true);
+        }
+        if (error instanceof CodeSearchWorkspaceError) {
+          outcome = "workspace_unavailable";
+          return textResult(renderCodeSearchError(error), true);
+        }
+        if (
+          error instanceof ChannelAUnavailableError ||
+          error instanceof ChannelAUnsupportedError ||
+          error instanceof ChannelAValidationError
+        ) {
+          outcome = "workspace_unavailable";
+          return textResult(
+            renderCodeSearchError(new CodeSearchWorkspaceError(error.message)),
+            true,
+          );
+        }
+        throw error;
+      } finally {
+        if (lease) breaker.release(lease);
+        recordCodeSearchCall(input.observability, {
+          outcome,
+          durationSeconds: (performance.now() - startedAt) / 1_000,
+          jevRequests,
+          jevCostUsd,
+        });
+      }
+    },
+  };
+}
+
+/**
+ * The `code_search` definition for one turn, or none. It is offered only when
+ * the session's frozen decision and the deployment enable it, a usable Jev key
+ * exists, and the turn has compute that can run its POSIX shell commands (not
+ * a Windows Connected Machine). Every input is durable, so the tool list stays
+ * the same from turn to turn and on every worker. Transient Jev health never
+ * hides the tool; the breaker only refuses calls.
+ */
+export function codeSearchToolDefinitions(input: {
+  enabled: boolean;
+  settings: Pick<Settings, "jevApiKey" | "jevBaseUrl" | "jevModel" | "jevRequestTimeoutMs">;
+  backend: Settings["sandboxBackend"];
+  /** The turn's Connected Machine workspace root, when a machine is primary. */
+  machineWorkspaceRoot?: string | null;
+  observability: Observability;
+  workspace: () => Promise<CodeSearchWorkspace>;
+  recordUsage?: (usage: CodeSearchUsage) => Promise<void>;
+  breaker?: JevCircuitBreaker;
+}): AttemptToolDefinition[] {
+  const apiKey = usableJevApiKey(input.settings);
+  const breaker = input.breaker ?? codeSearchCircuitBreaker;
+  if (!input.enabled || !apiKey || input.backend === "none") return [];
+  if (input.machineWorkspaceRoot && isWindowsConnectedMachinePath(input.machineWorkspaceRoot)) {
+    return [];
+  }
+  return [
+    createCodeSearchAttemptToolDefinition({
+      settings: input.settings,
+      apiKey,
+      workspace: input.workspace,
+      observability: input.observability,
+      ...(input.recordUsage ? { recordUsage: input.recordUsage } : {}),
+      breaker,
+    }),
+  ];
+}

package/src/activities/agent-turn/codex-capacity.ts

@@ -2,6 +2,7 @@ import {
   connectionModelAllowed,
   getSessionGoal,
   acquireCodexCredentialLease,
+  recheckCodexCredentialPlan,
   CodexCredentialLeaseAttemptFencedError,
   CodexCredentialFailoverExhaustedError,
   CODEX_CREDENTIAL_LEASE_TTL_MS,
@@ -14,7 +15,7 @@ import {
   type CodexCredentialLeaseSelectionContext,
 } from "@opengeni/db";
 import { type Settings } from "@opengeni/config";
-import { CodexReloginRequired } from "@opengeni/codex";
+import { CodexReloginRequired, codexPlanKey } from "@opengeni/codex";
 import { publishDurableSessionEvents } from "@opengeni/events";
 import {
   authoritativeCodexCapacityResetAt,
@@ -43,8 +44,15 @@ import { recordTurnStartupPhase } from "../../observability-metrics";
 import { createTurnCredentialLeases } from "./credential-leases";
 import { deliverFailedChildTurnToParent } from "../parent-wake";
 import { randomUUID } from "node:crypto";
+import { createLogThrottle, type LogThrottle } from "@opengeni/observability";
 
 import { refreshCappedCodexUsageRows } from "./codex";
+import {
+  CodexPlanEntitlementError,
+  codexAccountDisplayLabel,
+  codexPlanEntitlementAdmissionBlock,
+  codexPlanEntitlementFailurePayload,
+} from "./codex-plan-entitlement";
 import { codexUsageLimitFailurePayload, CODEX_USAGE_LIMIT_MAX_RESUME_MS } from "./errors";
 import type { ClaimTurnOk } from "./claim";
 import type {
@@ -56,6 +64,39 @@ import type {
   TurnControlState,
 } from "./turn-context";
 
+/** The eligible-pool gauge and `opengeni_codex_pool_low_total` stay per turn;
+ * the warning line is the first observation per workspace pool depth, then at
+ * most one per interval with the count it hid. The public log projection drops
+ * the identifiers and counts, so the closed `reason` keeps the depth visible. */
+export const CODEX_POOL_LOW_WARNING_INTERVAL_MS = 10 * 60_000;
+const codexPoolLowWarningThrottle = createLogThrottle({
+  intervalMs: CODEX_POOL_LOW_WARNING_INTERVAL_MS,
+  maxKeys: 1_024,
+});
+
+export function warnCodexPoolLow(
+  observability: Pick<ActivityServices["observability"], "warn">,
+  input: {
+    workspaceKey: string;
+    workspaceId: string;
+    eligibleCount: number;
+    connectedCount: number;
+    depth: "zero" | "one";
+  },
+  throttle: LogThrottle = codexPoolLowWarningThrottle,
+): void {
+  const admission = throttle.admit(`${input.workspaceKey}:${input.depth}`);
+  if (!admission) return;
+  observability.warn("Codex eligible credential pool is low", {
+    workspaceId: input.workspaceId,
+    eligibleCount: input.eligibleCount,
+    connectedCount: input.connectedCount,
+    depth: input.depth,
+    reason: input.depth === "zero" ? "eligible_pool_zero" : "eligible_pool_one",
+    ...(admission.suppressedCount > 0 ? { suppressedCount: admission.suppressedCount } : {}),
+  });
+}
+
 export type CapacityPhaseDeps = {
   input: RunAgentTurnInput;
   settings: Settings;
@@ -142,7 +183,13 @@ export async function selectCodexTurnCapacity(
         )
           throw new Error("This model is disabled for the pinned Codex subscription");
         return selectCodexCredentialLeaseForTurn({
-          context: { ...context, accounts: allowed },
+          // The accepted product model also scopes proven plan entitlement:
+          // an account whose current plan excludes it is not a candidate.
+          context: {
+            ...context,
+            accounts: allowed,
+            modelId: deps.turnExecutionPolicy.productModelId,
+          },
           sessionId: input.sessionId,
           sessionPinnedCredentialId: lockedSessionCodexState.pinnedCredentialId,
           sessionPinSource: lockedSessionCodexState.pinSource,
@@ -173,10 +220,17 @@ export async function selectCodexTurnCapacity(
         );
       if (leased.decision.kind === "allCapped") {
         // Bounded self-heal of stale usage cache, then ONE new atomic selection.
-        await refreshCappedCodexUsageRows(db, settings, input.workspaceId, leased.accounts, {
-          signalCodexCapacityWorkflow,
-          wakeSessionWorkflow,
-        });
+        await refreshCappedCodexUsageRows(
+          db,
+          settings,
+          input.workspaceId,
+          leased.accounts,
+          {
+            signalCodexCapacityWorkflow,
+            wakeSessionWorkflow,
+          },
+          turn.id,
+        );
         leaseAcquisitionStartedAtMs = performance.now();
         leased = await acquireCodexCredentialLease(
           db,
@@ -197,6 +251,99 @@ export async function selectCodexTurnCapacity(
           selectForTurn,
         );
       }
+      // No account can serve the model because the only candidates' CURRENT
+      // plans were proven not to include it. Re-read those plans once (an
+      // upgrade may not have been observed yet), then either select again or
+      // fail with typed copy instead of an indefinite capacity wait.
+      const productModelId = deps.turnExecutionPolicy.productModelId;
+      const planBlock = (current: typeof leased) =>
+        codexPlanEntitlementAdmissionBlock({
+          accounts: current.accounts,
+          modelId: productModelId,
+          credentialId: current.credentialId,
+          rotationEnabled: current.rotationEnabled,
+          activeCredentialId: current.activeCredentialId,
+          pinnedCredentialId: current.sessionCodexState.pinnedCredentialId,
+          pinSource: current.sessionCodexState.pinSource,
+          now: new Date(),
+        });
+      const blocked = planBlock(leased);
+      if (blocked) {
+        const rechecks = await Promise.all(
+          blocked.slice(0, 4).map((account) =>
+            recheckCodexCredentialPlan(db, settings, input.workspaceId, account.id, {
+              turnId,
+              purpose: "capacity_refresh",
+            }).catch(() => null),
+          ),
+        );
+        const planMoved = rechecks.some(
+          (recheck) =>
+            recheck?.planType != null &&
+            codexPlanKey(recheck.planType) !== codexPlanKey(recheck.previousPlanType),
+        );
+        if (planMoved) {
+          leaseAcquisitionStartedAtMs = performance.now();
+          leased = await acquireCodexCredentialLease(
+            db,
+            {
+              accountId: input.accountId,
+              workspaceId: input.workspaceId,
+              sessionId: input.sessionId,
+              turnId,
+              attemptId: input.attemptId,
+              executionGeneration: attempt.executionGeneration,
+              workflowId: input.workflowId,
+              workflowRunId: input.workflowRunId,
+              dispatchId,
+              expectedRedispatches: attempt.redispatchesAtDispatch,
+              holderId,
+              advanceActivePointer: true,
+            },
+            selectForTurn,
+          );
+        }
+        const stillBlocked = planMoved ? planBlock(leased) : blocked;
+        if (stillBlocked) {
+          const account = stillBlocked[0]!;
+          // Name the plan only when this admission just observed it.
+          const observedPlan =
+            rechecks[blocked.findIndex((candidate) => candidate.id === account.id)]?.planType ??
+            null;
+          const payload = codexPlanEntitlementFailurePayload({
+            accountLabel: stillBlocked.length === 1 ? codexAccountDisplayLabel(account) : null,
+            planType: observedPlan,
+            planChanged: false,
+            modelId: productModelId,
+          });
+          if (turn.source === "compaction") {
+            if (
+              !(await eventing.settle!({
+                events: [
+                  {
+                    type: "turn.cancelled",
+                    payload: {
+                      maintenance: "context_compaction",
+                      reason: payload.code,
+                      requestPreserved: true,
+                    },
+                  },
+                  { type: "session.status.changed", payload: { status: "idle" } },
+                ],
+                turnStatus: "cancelled",
+                sessionStatus: "idle",
+                activeTurnId: null,
+              }))
+            ) {
+              return { exit: claimedResult({ status: "cancelled" }) };
+            }
+            control.turnMetricOutcome = "cancelled";
+            control.activityStatus = "idle";
+            return { exit: claimedResult({ status: "idle", deferredUntilWake: true }) };
+          }
+          throw new CodexPlanEntitlementError(payload);
+        }
+      }
       const lockedSessionCodexState = leased.sessionCodexState;
       providerTurn.codexPolicySnapshot = leased.codexPolicySnapshot;
       const sessionPin = lockedSessionCodexState.pinnedCredentialId;
@@ -276,6 +423,7 @@ export async function selectCodexTurnCapacity(
       }
       providerTurn.effectiveCodexCredentialId = leased.credentialId;
       providerTurn.codexCredentialFailoverLimit = leased.failoverLimit;
+      providerTurn.codexProductModelId = deps.turnExecutionPolicy.productModelId;
       leases.codex.generation = leased.generation;
       leases.codex.confirmedUntilMs =
         leased.leasedUntil && leaseAcquisitionStartedAtMs !== null
@@ -395,11 +543,12 @@ export async function selectCodexTurnCapacity(
           help: "Alert signal emitted when the eligible Codex pool is zero or one.",
           labels: { workspace_key: codexWorkspaceKey, depth: poolDepth },
         });
-        observability.warn("Codex eligible credential pool is low", {
+        warnCodexPoolLow(observability, {
+          workspaceKey: codexWorkspaceKey,
           workspaceId: input.workspaceId,
           eligibleCount,
           connectedCount: leased.accounts.length,
-          depth: poolDepth,
+          depth: poolDepth === "zero" ? "zero" : "one",
         });
       }
 
@@ -466,6 +615,11 @@ export async function selectCodexTurnCapacity(
             },
           },
         );
+        if (evaluated.action === "stopped") {
+          control.turnMetricOutcome = "failed";
+          control.activityStatus = evaluated.sessionStatus === "queued" ? "idle" : "failed";
+          return { exit: claimedResult({ status: control.activityStatus }) };
+        }
         if (evaluated.action === "resumed") {
           control.turnMetricOutcome = "recovering";
           control.activityStatus = "recovering";
@@ -583,6 +737,11 @@ export async function selectCodexTurnCapacity(
             },
           },
         );
+        if (evaluated.action === "stopped") {
+          control.turnMetricOutcome = "failed";
+          control.activityStatus = evaluated.sessionStatus === "queued" ? "idle" : "failed";
+          return { exit: claimedResult({ status: control.activityStatus }) };
+        }
         if (evaluated.action === "resumed") {
           control.turnMetricOutcome = "recovering";
           control.activityStatus = "recovering";
@@ -682,6 +841,11 @@ export async function selectCodexTurnCapacity(
             failurePayload,
           },
         );
+        if (evaluated.action === "stopped") {
+          control.turnMetricOutcome = "failed";
+          control.activityStatus = evaluated.sessionStatus === "queued" ? "idle" : "failed";
+          return { exit: claimedResult({ status: control.activityStatus }) };
+        }
         if (evaluated.action === "resumed") {
           control.turnMetricOutcome = "recovering";
           control.activityStatus = "recovering";