@opengeni/worker-bundle 1.3.0 → 1.4.2

package/src/activities/agent-turn/sandbox-establish.ts

@@ -66,6 +66,7 @@ import {
   classifySandboxLogicalProvisionFailure,
   isLazySandboxProvisionRetryable,
   createTurnSandboxProvisioner,
+  trackPhysicalSandboxResume,
 } from "./sandbox-provision";
 import {
   resolveActiveSandboxBackend,
@@ -645,40 +646,43 @@ export async function establishTurnSandbox(deps: EstablishTurnSandboxDeps): Prom
           const lazyHolderId = sandboxState.sandboxHolderId;
           const lazyGroupId = sandboxState.sandboxGroupId;
           sandboxState.resumeManagedGroupBox = () =>
-            resumeBoxForTurn(
-              {
-                db,
-                settings: runSettings,
-                logicalFallbackSettings: logicalSandboxSettings,
-                cancellationSignal: sandboxResumeSignal,
-                sandboxMetrics: runtimeMetricsHooksForObservability(observability),
-                observability,
-                freshSandboxReadinessReplacementBudget,
-                onSandboxLost: publishSandboxLost,
-                objectStorage,
-                bus,
-              },
-              {
-                accountId: input.accountId,
-                workspaceId: input.workspaceId,
-                sandboxGroupId: lazyGroupId,
-                sessionId: input.sessionId,
-                backend: groupBoxBackend,
-                os: groupBoxOs,
-                environment: sandboxEnvironment,
-                ...(groupBoxImage
-                  ? {
-                      image: groupBoxImage,
-                      imagePolicy: groupBoxImagePolicy,
-                    }
-                  : {}),
-                // The lazy acquire must enforce the same frozen rig authority
-                // as the eager acquire; otherwise a warm box for another rig
-                // can bypass the shared-state conflict/rotation fence.
-                ...(rigVersion ? { rigVersionId: rigVersion.id } : {}),
-              },
-              "turn",
-              lazyHolderId,
+            trackPhysicalSandboxResume(
+              sandboxState.inFlightSandboxResumes,
+              resumeBoxForTurn(
+                {
+                  db,
+                  settings: runSettings,
+                  logicalFallbackSettings: logicalSandboxSettings,
+                  cancellationSignal: sandboxResumeSignal,
+                  sandboxMetrics: runtimeMetricsHooksForObservability(observability),
+                  observability,
+                  freshSandboxReadinessReplacementBudget,
+                  onSandboxLost: publishSandboxLost,
+                  objectStorage,
+                  bus,
+                },
+                {
+                  accountId: input.accountId,
+                  workspaceId: input.workspaceId,
+                  sandboxGroupId: lazyGroupId,
+                  sessionId: input.sessionId,
+                  backend: groupBoxBackend,
+                  os: groupBoxOs,
+                  environment: sandboxEnvironment,
+                  ...(groupBoxImage
+                    ? {
+                        image: groupBoxImage,
+                        imagePolicy: groupBoxImagePolicy,
+                      }
+                    : {}),
+                  // The lazy acquire must enforce the same frozen rig authority
+                  // as the eager acquire; otherwise a warm box for another rig
+                  // can bypass the shared-state conflict/rotation fence.
+                  ...(rigVersion ? { rigVersionId: rigVersion.id } : {}),
+                },
+                "turn",
+                lazyHolderId,
+              ),
             );
           if (
             shouldPrefetchManagedSandbox({
@@ -725,52 +729,55 @@ export async function establishTurnSandbox(deps: EstablishTurnSandboxDeps): Prom
         );
         try {
           sandboxState.resolvedSandbox = await waitForTurnOperation(
-            resumeBoxForTurn(
-              {
-                db,
-                settings: runSettings,
-                logicalFallbackSettings: logicalSandboxSettings,
-                cancellationSignal: sandboxResumeSignal,
-                sandboxMetrics: runtimeMetricsHooksForObservability(observability),
-                observability,
-                freshSandboxReadinessReplacementBudget,
-                onSandboxLost: publishSandboxLost,
-                objectStorage,
-                bus,
-              },
-              {
-                accountId: input.accountId,
-                workspaceId: input.workspaceId,
-                sandboxGroupId: session.sandboxGroupId,
-                sessionId: input.sessionId,
-                // groupBoxBackend, not turn.sandboxBackend: a machine-home turn that
-                // is not machine-primary resumes a real cloud group box (the
-                // deployment default), never a "selfhosted" box (which would throw
-                // for lack of a bound agentId).
-                backend: groupBoxBackend,
-                os: groupBoxOs,
-                environment: sandboxEnvironment,
-                // Deployment/workspace pins select NEW creates only. The lease
-                // retains a live group's image through between-turn repins, even
-                // with other holders; only rotation/reaping elects a new image.
-                // Select the image for the actual group-box backend; a configured Modal image must
-                // never override a Docker run. The selfhosted branch
-                // (establishSelfhostedTurnSession) NEVER passes
-                // an image — B3 lives only on this managed-box branch.
-                ...(groupBoxImage
-                  ? {
-                      image: groupBoxImage,
-                      imagePolicy: groupBoxImagePolicy,
-                    }
-                  : {}),
-                // RIG IS SHARED STATE (M3): stamp the frozen rig version so the lease
-                // conflicts on a live shared box set up under a different rig (solo
-                // durable rotation / N-holders SandboxRigConflictError). Omitted for a
-                // rig-less turn -> never stamped or enforced (shares exactly as today).
-                ...(rigVersion ? { rigVersionId: rigVersion.id } : {}),
-              },
-              "turn",
-              managedOwnership!.holderId,
+            trackPhysicalSandboxResume(
+              sandboxState.inFlightSandboxResumes,
+              resumeBoxForTurn(
+                {
+                  db,
+                  settings: runSettings,
+                  logicalFallbackSettings: logicalSandboxSettings,
+                  cancellationSignal: sandboxResumeSignal,
+                  sandboxMetrics: runtimeMetricsHooksForObservability(observability),
+                  observability,
+                  freshSandboxReadinessReplacementBudget,
+                  onSandboxLost: publishSandboxLost,
+                  objectStorage,
+                  bus,
+                },
+                {
+                  accountId: input.accountId,
+                  workspaceId: input.workspaceId,
+                  sandboxGroupId: session.sandboxGroupId,
+                  sessionId: input.sessionId,
+                  // groupBoxBackend, not turn.sandboxBackend: a machine-home turn that
+                  // is not machine-primary resumes a real cloud group box (the
+                  // deployment default), never a "selfhosted" box (which would throw
+                  // for lack of a bound agentId).
+                  backend: groupBoxBackend,
+                  os: groupBoxOs,
+                  environment: sandboxEnvironment,
+                  // Deployment/workspace pins select NEW creates only. The lease
+                  // retains a live group's image through between-turn repins, even
+                  // with other holders; only rotation/reaping elects a new image.
+                  // Select the image for the actual group-box backend; a configured Modal image must
+                  // never override a Docker run. The selfhosted branch
+                  // (establishSelfhostedTurnSession) NEVER passes
+                  // an image — B3 lives only on this managed-box branch.
+                  ...(groupBoxImage
+                    ? {
+                        image: groupBoxImage,
+                        imagePolicy: groupBoxImagePolicy,
+                      }
+                    : {}),
+                  // RIG IS SHARED STATE (M3): stamp the frozen rig version so the lease
+                  // conflicts on a live shared box set up under a different rig (solo
+                  // durable rotation / N-holders SandboxRigConflictError). Omitted for a
+                  // rig-less turn -> never stamped or enforced (shares exactly as today).
+                  ...(rigVersion ? { rigVersionId: rigVersion.id } : {}),
+                },
+                "turn",
+                managedOwnership!.holderId,
+              ),
             ),
             sandboxResumeSignal,
             releaseLateSandbox,

package/src/activities/agent-turn/sandbox-provision.ts

@@ -153,6 +153,57 @@ export async function waitForTurnOperation<T>(
   }
 }
 
+/**
+ * Upper bound for joining a cancelled attempt's in-flight provider establish.
+ * Worker shutdown cancels activities 5s after SIGTERM and force-stops 100s
+ * after it; this leaves room for the remaining finalization steps.
+ */
+export const PHYSICAL_SANDBOX_RESUME_DRAIN_TIMEOUT_MS = 60_000;
+
+/** Register the PHYSICAL lease acquire/provider establish promise, not the
+ * cancellable wrapper awaiting it. */
+export function trackPhysicalSandboxResume<T>(
+  inFlight: Set<Promise<unknown>>,
+  operation: Promise<T>,
+): Promise<T> {
+  inFlight.add(operation);
+  const settle = (): void => {
+    inFlight.delete(operation);
+  };
+  operation.then(settle, settle);
+  return operation;
+}
+
+/**
+ * Join every lease acquire/provider establish this attempt started, even after
+ * the turn-facing wrapper rejected on cancellation. An establish has no
+ * portable abort: after cancellation it still finishes its own exact cleanup
+ * (terminate an unpublished box and roll warming back to cold, or keep a
+ * published box and drop its holder). If the activity returned first, a
+ * shutting-down worker exited mid-cleanup and stranded a warming lease naming a
+ * terminated box; the replacement attempt then waited out the warming TTL and
+ * inherited a lost-workspace verdict. This wait is deliberately independent of
+ * the finalizer cancellation signal (cancellation is what started it) and is
+ * bounded so an uncooperative provider call cannot pin the activity.
+ */
+export async function drainPhysicalSandboxResumes(
+  inFlight: Set<Promise<unknown>>,
+  timeoutMs: number = PHYSICAL_SANDBOX_RESUME_DRAIN_TIMEOUT_MS,
+): Promise<"none" | "settled" | "timed_out"> {
+  if (inFlight.size === 0) return "none";
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  try {
+    return await Promise.race([
+      Promise.allSettled([...inFlight]).then(() => "settled" as const),
+      new Promise<"timed_out">((resolve) => {
+        timer = setTimeout(() => resolve("timed_out"), Math.max(1, timeoutMs));
+      }),
+    ]);
+  } finally {
+    if (timer) clearTimeout(timer);
+  }
+}
+
 export function sleep(ms: number): Promise<void> {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }

package/src/activities/agent-turn/sandbox-runtime.ts

@@ -5,6 +5,7 @@ import {
   readLease,
   accrueWarmSeconds,
   SandboxWorkspaceMutationFencedError,
+  SandboxWorkspaceMutationOutputRejectedError,
   retainWorkspaceMutationProcess,
   retainedProviderCommandPersistence,
   SandboxRetainedProcessPromotionFencedError,
@@ -12,6 +13,7 @@ import {
 import { createProviderCommandRetainer } from "@opengeni/db/retained-provider-commands";
 import {
   RoutingMutationOutcomeUnknownError,
+  RoutingMutationOutputRejectedError,
   ProviderCommandStartOutcomeUnknownError,
   ProviderCommandObservationUnavailableError,
   isModalCommandStartOutcomeUnknownError,
@@ -622,6 +624,20 @@ export function createSandboxTurnRuntime(deps: SandboxTurnRuntimeDeps) {
     try {
       await settleMutation("resolved");
     } catch (settlementError) {
+      if (
+        settlementError instanceof SandboxWorkspaceMutationOutputRejectedError &&
+        settlementError.matchesPhysicalSettlement({
+          accountId: identity.accountId,
+          workspaceId: identity.workspaceId,
+          admission,
+          operation,
+          outcome: "resolved",
+        })
+      ) {
+        throw new RoutingMutationOutputRejectedError(operation, settlementError.code, {
+          cause: settlementError,
+        });
+      }
       throw new RoutingMutationOutcomeUnknownError(
         operation,
         `Platform workspace mutation "${operation}" returned from the provider but lost its durable settlement fence; its outcome is unknown and it was not replayed`,

package/src/activities/agent-turn/stream-attempt.ts

@@ -79,6 +79,7 @@ import {
   isCompletedGeneratedImageSdkEvent,
 } from "../generated-images";
 import { programmaticApproval, programmaticContinuationNote } from "../programmatic-approvals";
+import { CodemodeAttemptDispatcher } from "../codemode-dispatcher";
 import { ToolResultSpill } from "./tool-result-spill";
 import { ownedTurnSandboxForAgent } from "./turn-sandbox-access";
 import { createTurnCredentialLeases } from "./credential-leases";
@@ -661,6 +662,16 @@ export async function runTurnStreamAttempt(
     const providerPublishesNativeRequestEvents =
       resolvedModel?.provider.kind === "codex-subscription" ||
       resolvedModel?.provider.kind === "xai-subscription";
+    // Bounded provider label for the streaming SLIs — the resolved registry
+    // provider id (or the built-in OpenAI/Azure provider), never a raw
+    // user-supplied model string. Created before the stream starts so the
+    // first request's model entry and provider dispatch are observed.
+    const streamTiming = new StreamTimingMetrics(observability, {
+      provider: streamProvider,
+    });
+    // Native Codex/SuperGrok transports report their literal dispatch through
+    // the shared turn context; generic transports use the wire-dispatch clock.
+    eventing.providerDispatchObserver = () => streamTiming.onProviderDispatch();
     let fallbackProviderRequestStartedAt: number | null = null;
     let fallbackProviderRequestLifecycleStartedAt: number | null = null;
     const recordFallbackProviderDispatchAtWire = async (): Promise<void> => {
@@ -853,7 +864,10 @@ export async function runTurnStreamAttempt(
         }
         attempt.modelRequestStarted = true;
         return await runtime.runStream(agent, runInput!, eventing.modelRunSettings, {
-          beforeModelRequest: modelCallAdmission.beforeModelRequest,
+          beforeModelRequest: () => {
+            streamTiming.onModelRequestEntry();
+            return modelCallAdmission.beforeModelRequest();
+          },
           onModelResponse: (event) => {
             responseCreditPolicyRevisions.add(creditPolicyRevision);
             return modelCallAdmission.onModelResponse(event);
@@ -965,6 +979,7 @@ export async function runTurnStreamAttempt(
             ? {
                 onModelTransportStarted: recordFallbackProviderDispatchAtWire,
                 onModelTransportDispatched: (clock) => {
+                  streamTiming.onProviderDispatch();
                   eventing.initialModelWireDispatch.record(
                     { provider: streamProvider, dispatchId },
                     clock,
@@ -1006,12 +1021,6 @@ export async function runTurnStreamAttempt(
       modelCallAdmission.close();
       throw error;
     }
-    // Bounded provider label for the streaming SLIs — the resolved registry
-    // provider id (or the built-in OpenAI/Azure provider), never a raw
-    // user-supplied model string.
-    const streamTiming = new StreamTimingMetrics(observability, {
-      provider: streamProvider,
-    });
     eventing.batcher = createRuntimeBatcher(
       async (events) => {
         await eventing.publish!(events);
@@ -1960,14 +1969,26 @@ export async function runTurnStreamAttempt(
 
   // Recover committed waits/results before any model dispatch, including crash replacement.
   // A new attempt uses its own signed caller; origin provenance never changes.
-  await eventing.toolPreparationReady;
-  const codemodeOperations =
-    (await eventing.codemodeDispatcher?.resumeApproved(
-      `sandbox:${input.attemptId}`,
-      trigger.type === "user.approvalDecision"
-        ? (trigger.payload as { approvalId: string }).approvalId
-        : undefined,
-    )) ?? [];
+  // Join-on-use: lazy MCP preparation is awaited only when stored work must resume.
+  const codemodeOperations = await CodemodeAttemptDispatcher.resumeApproved(
+    db,
+    {
+      accountId: input.accountId,
+      workspaceId: input.workspaceId,
+      sessionId: input.sessionId,
+      turnId: activeTurnId,
+      attemptId: input.attemptId,
+      executionGeneration: attempt.executionGeneration,
+    },
+    `sandbox:${input.attemptId}`,
+    trigger.type === "user.approvalDecision"
+      ? (trigger.payload as { approvalId: string }).approvalId
+      : undefined,
+    async () => {
+      await eventing.toolPreparationReady;
+      return eventing.codemodeDispatcher;
+    },
+  );
   codemodeContinuationNote = programmaticContinuationNote(codemodeOperations);
   programmaticApprovalAcknowledged =
     trigger.type === "user.approvalDecision" &&

package/src/activities/agent-turn/subscription-lease.ts

@@ -1,6 +1,8 @@
 export type LeaseRenewReason = "timer" | "runtime_event" | "model_usage";
 export type LeaseLossReason = "deadline" | "not_found";
 
+const HEARTBEAT_INTERVAL_MS = 60_000;
+
 type LeaseIdentity = { turnId: string; holderId: string; generation: number };
 export type SubscriptionLeaseDeps = {
   ttlMs: number;
@@ -24,11 +26,13 @@ export class SubscriptionTurnLease {
   private heartbeatInFlight = false;
   private heartbeatTimer: ReturnType<typeof setInterval> | undefined;
   private readonly now: () => number;
+  private readonly renewalIntervalMs: number;
 
   constructor(private readonly leaseDeps: SubscriptionLeaseDeps) {
     if (!Number.isFinite(leaseDeps.ttlMs) || leaseDeps.ttlMs <= 0)
       throw new Error("Subscription lease TTL must be positive");
     this.now = leaseDeps.now ?? (() => performance.now());
+    this.renewalIntervalMs = Math.min(HEARTBEAT_INTERVAL_MS, leaseDeps.ttlMs / 5);
   }
 
   private expired = (deadline: number | null): boolean =>
@@ -62,6 +66,11 @@ export class SubscriptionTurnLease {
       return;
     }
     if (this.heartbeatInFlight) return;
+    // Acquisition and successful renewal confirm ownership from request start.
+    // Reuse that proof between heartbeats; a failed renewal leaves it unchanged
+    // so the next checkpoint can retry immediately. Expiry is checked first.
+    const renewAt = this.confirmedUntilMs! - this.leaseDeps.ttlMs + this.renewalIntervalMs;
+    if (this.now() < renewAt) return;
     this.heartbeatInFlight = true;
     const priorDeadline = this.confirmedUntilMs;
     const startedAt = this.now();
@@ -103,7 +112,7 @@ export class SubscriptionTurnLease {
 
   startHeartbeat = (): void => {
     if (!this.leaseDeps.getTurnId() || this.heartbeatTimer) return;
-    this.heartbeatTimer = setInterval(() => void this.renew("timer"), 60_000);
+    this.heartbeatTimer = setInterval(() => void this.renew("timer"), HEARTBEAT_INTERVAL_MS);
     this.heartbeatTimer.unref?.();
   };
 

package/src/activities/agent-turn/tool-environment.ts

@@ -67,6 +67,7 @@ import {
   scheduledTurnMcpServerIds,
   hasPermission,
   requireExplicitPermissionDelegation,
+  createWebSearchBilling,
 } from "@opengeni/core";
 import { loadWorkspaceEnvironmentForRunWithCredentials } from "../environment";
 import { withFirstPartyTools } from "../goals";
@@ -116,6 +117,7 @@ import { resolveTurnSandboxAccess } from "./turn-sandbox-access";
 import { createListModelsAttemptToolDefinition } from "./list-models";
 import { createRefreshCredentialsAttemptToolDefinition } from "./refresh-credentials";
 import { codeSearchToolDefinitions, codeSearchWorkspaceFromChannel } from "./code-search";
+import { turnWebSearchPlan, webSearchToolDefinitions } from "./web-search";
 import { createWorkspaceSkillTools } from "./skill-tools";
 import { loadConfiguredBundledSkills } from "./skill-selection";
 import { guardSkillFilesystem } from "./skill-transfer";
@@ -886,6 +888,22 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
   const attemptToolFamilies = resolveAgentToolFamilies(session.agent, {
     hasSkills: skillCatalog.length > 0,
   });
+  // Provider web search: only where the turn has no hosted search (or the
+  // operator chose `replace`), only when a provider is configured, and only
+  // when the session's agent configuration allows web search (filtered below).
+  const webSearchTools = webSearchToolDefinitions({
+    settings: runSettings,
+    tools: turnWebSearchPlan(resolvedModel, runSettings).providerTools,
+    scope: {
+      accountId: input.accountId,
+      workspaceId: input.workspaceId,
+      sessionId: input.sessionId,
+      turnId: turn.id,
+      attemptId: input.attemptId,
+    },
+    billing: createWebSearchBilling({ db, settings: runSettings }),
+    observability,
+  });
   const attemptToolDefinitions = [
     ...(operationReadStore
       ? [
@@ -976,6 +994,7 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
       ? [googleDrivePublicationTool]
       : []),
     ...codeSearchTools,
+    ...webSearchTools,
   ].filter((tool) => attemptToolFamilies.allowsFunctionTool(tool.modelName));
   recordTurnStartupPhase(observability, {
     phase: "tool_context_preparation",

package/src/activities/agent-turn/tool-policy.ts

@@ -2,8 +2,11 @@ import {
   desktopCapableBackend,
   type ComputerToolMode,
   type LazyToolTransport,
+  type OpenGeniRuntime,
 } from "@opengeni/runtime";
+import type { TurnExecutionPolicyV1 } from "@opengeni/contracts";
 import {
+  configuredModelForAcceptedTurnExecutionPolicy,
   isDirectOpenAiApiBaseUrl,
   type ModelProviderApi,
   type ResolvedModelProvider,
@@ -208,6 +211,28 @@ export function shouldDeferNonEagerToolPreparation(args: {
   );
 }
 
+/**
+ * Resolve the provider routing/gating shape for an accepted turn. The current
+ * definition is projected back onto the verified frozen policy: a turn
+ * accepted before hosted web search was enabled keeps its frozen tool set on
+ * every recovery attempt (stable tool prefix, exact accepted definition); the
+ * next accepted logical turn resolves the newly enabled tool.
+ */
+export function resolveAcceptedTurnModel(
+  runtime: Pick<OpenGeniRuntime, "resolveTurnModel">,
+  settings: Settings,
+  policy: TurnExecutionPolicyV1,
+): ReturnType<OpenGeniRuntime["resolveTurnModel"]> {
+  const current = runtime.resolveTurnModel(settings, policy.productModelId);
+  if (!current) return null;
+  const configured = configuredModelForAcceptedTurnExecutionPolicy(
+    current.configured,
+    current.provider,
+    policy,
+  );
+  return configured === current.configured ? current : { ...current, configured };
+}
+
 /**
  * Native web search is a runtime capability, not part of the session's MCP
  * allow-list. Attach it whenever the resolved provider advertises runnable

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

@@ -103,6 +103,9 @@ export type SandboxRuntimeState = {
   }>;
   turnSandboxProvisioner: TurnSandboxProvisioner<ResumedTurnSandbox> | null;
   resumeManagedGroupBox: (() => Promise<ResumedTurnSandbox>) | null;
+  /** Physical resumeBoxForTurn promises still running for this attempt; the
+   * finalizer joins them so a cancelled establish completes its own cleanup. */
+  inFlightSandboxResumes: Set<Promise<unknown>>;
   prefetchedManagedBox: Promise<ResumedTurnSandbox> | null;
   prefetchedManagedBoxResult: ResumedTurnSandbox | null;
   setupBoxSession: unknown;
@@ -153,6 +156,9 @@ export type EventingState = {
   firstModelRequestPreparationRecorded: boolean;
   firstModelRequestCheckpointAt: number | null;
   initialModelWireDispatch: InitialModelWireDispatchClock;
+  /** Diagnostic only: the current stream's timing hook for a native transport's
+   *  literal provider dispatch. Never joins, fences or fails the request. */
+  providerDispatchObserver: (() => void) | null;
   companyBrainContextContributions: readonly ModelContextContributionSummary[] | null;
   /** Skill ids in this turn's frozen, model-visible Skill index; telemetry only. */
   modelVisibleSkillIds: ReadonlySet<string> | null;
@@ -259,6 +265,7 @@ export function createTurnContext(input: {
       firstModelPreparationNestedSandboxPhases: [],
       turnSandboxProvisioner: null,
       resumeManagedGroupBox: null,
+      inFlightSandboxResumes: new Set(),
       prefetchedManagedBox: null,
       prefetchedManagedBoxResult: null,
       setupBoxSession: null,
@@ -301,6 +308,7 @@ export function createTurnContext(input: {
       firstModelRequestPreparationRecorded: false,
       firstModelRequestCheckpointAt: null,
       initialModelWireDispatch: new InitialModelWireDispatchClock(),
+      providerDispatchObserver: null,
       companyBrainContextContributions: null,
       modelVisibleSkillIds: null,
     },