@opengeni/worker-bundle 1.3.0 → 1.4.2

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

@@ -0,0 +1,381 @@
+import type { AttemptToolDefinition } from "@opengeni/codemode";
+import {
+  WEB_FETCH_TOOL_NAME,
+  WEB_SEARCH_TOOL_NAME,
+  webSearchCallPricing,
+  webSearchCreditMicros,
+  webSearchProviderConfig,
+  webSearchToolPlan,
+  type WebSearchProviderConfig,
+  type WebSearchProviderToolName,
+  type WebSearchSettings,
+} from "@opengeni/config";
+import {
+  WebSearchBillingRefusedError,
+  type WebSearchBilling,
+  type WebSearchCallCost,
+  type WebSearchCallScope,
+} from "@opengeni/core";
+import type { Observability } from "@opengeni/observability";
+import {
+  WEB_FETCH_TOOL_DESCRIPTION,
+  WEB_FETCH_RETAINED_MAX_CHARS,
+  WEB_SEARCH_TOOL_DESCRIPTION,
+  WebSearchProviderError,
+  WebToolArgumentError,
+  createWebFetchProvider,
+  createWebSearchProvider,
+  parseWebFetchArguments,
+  parseWebSearchArguments,
+  renderWebPageWindow,
+  renderWebSearchResults,
+  webFetchInputSchema,
+  webSearchInputSchema,
+  type WebPage,
+} from "@opengeni/runtime/web-search";
+import { hostedWebSearchForTurn } from "./tool-policy";
+
+/**
+ * Which web search one turn gets: the provider's hosted tool, Opengeni's
+ * provider tools, or neither. Tool preparation and agent construction both
+ * call this with the same accepted model and settings, so they agree.
+ */
+export function turnWebSearchPlan(
+  resolvedModel: {
+    configured: { hostedWebSearch: boolean };
+    provider: { kind?: string };
+  } | null,
+  settings: WebSearchSettings,
+) {
+  return webSearchToolPlan(settings, {
+    hostedWebSearch: hostedWebSearchForTurn(resolvedModel, settings.webSearchEnabled),
+    // SuperGrok's native search is added by its transport, not as an agent
+    // tool, so provider tools would only duplicate it.
+    transportHostedSearch: resolvedModel?.provider.kind === "xai-subscription",
+  });
+}
+
+/** Pages kept per attempt so paging through one page is not billed again. */
+const PAGE_CACHE_ENTRIES = 16;
+
+function textResult(text: string, isError: boolean) {
+  return { isError, content: [{ type: "text" as const, text }] };
+}
+
+type Pricing = ReturnType<typeof webSearchCallPricing>;
+
+function settledCost(
+  operationId: string,
+  operation: "search" | "fetch",
+  provider: string,
+  pricing: Pricing,
+  reportedCostMicros: number | undefined,
+): WebSearchCallCost {
+  // An explicit operator price wins; then the provider's own reported cost;
+  // then the built-in list price.
+  const basis: WebSearchCallCost["basis"] = pricing.explicit
+    ? "configured_price"
+    : reportedCostMicros !== undefined
+      ? "provider_reported"
+      : "list_price";
+  const providerMicros =
+    basis === "provider_reported" ? reportedCostMicros! : pricing.providerMicros;
+  return {
+    operationId,
+    operation,
+    provider,
+    providerMicros,
+    creditMicros: webSearchCreditMicros(providerMicros, pricing.marginBps),
+    marginBps: pricing.marginBps,
+    basis,
+  };
+}
+
+function providerFailure(error: WebSearchProviderError): string {
+  return `${error.message}${error.retryable ? ". You may retry shortly." : "."}`;
+}
+
+/** Closed outcome set for `opengeni_web_search_calls_total`. */
+export type WebSearchCallOutcome =
+  | "ok"
+  | "provider_error"
+  | "provider_retryable"
+  | "billing_refused"
+  | "error";
+
+type WebSearchObservability = Pick<Observability, "warn" | "incrementCounter" | "observeHistogram">;
+
+const WEB_SEARCH_DURATION_BUCKETS = [0.25, 0.5, 1, 2, 4, 8, 15, 30];
+
+/**
+ * One provider call: a counter by fixed outcome plus its latency. Labels are
+ * the operation, the configured provider id, and the outcome only, never the
+ * query, URL, or session.
+ */
+function recordWebSearchCall(
+  observability: WebSearchObservability,
+  operation: "search" | "fetch",
+  provider: string,
+  outcome: WebSearchCallOutcome,
+  startedAt: number,
+): void {
+  observability.incrementCounter({
+    name: "opengeni_web_search_calls_total",
+    help: "Provider web search/fetch calls made by the worker, by operation, provider and outcome.",
+    labels: { operation, provider, outcome },
+  });
+  observability.observeHistogram({
+    name: "opengeni_web_search_call_duration_seconds",
+    help: "Provider web search/fetch call duration in seconds.",
+    buckets: WEB_SEARCH_DURATION_BUCKETS,
+    value: (performance.now() - startedAt) / 1000,
+    labels: { operation, provider },
+  });
+}
+
+function providerFailureOutcome(error: WebSearchProviderError): WebSearchCallOutcome {
+  return error.retryable ? "provider_retryable" : "provider_error";
+}
+
+/**
+ * The `web_search` and `web_fetch` attempt tools for one turn. The caller
+ * decides which names the turn gets ({@link webSearchToolPlan}); this only
+ * builds them. Provider keys stay in the worker and never reach a sandbox.
+ */
+export function webSearchToolDefinitions(input: {
+  settings: WebSearchSettings;
+  tools: readonly WebSearchProviderToolName[];
+  scope: WebSearchCallScope;
+  billing: WebSearchBilling;
+  observability: WebSearchObservability;
+  fetch?: typeof fetch;
+}): AttemptToolDefinition[] {
+  const config: WebSearchProviderConfig | null = webSearchProviderConfig(input.settings);
+  if (!config || input.tools.length === 0) return [];
+  const adapterInput = {
+    timeoutMs: config.timeoutMs,
+    ...(input.fetch ? { fetch: input.fetch } : {}),
+  };
+  const settle = async (cost: WebSearchCallCost) => {
+    try {
+      await input.billing.settle(input.scope, cost);
+    } catch (error) {
+      // The provider already answered. Keep the model's result; the missing
+      // receipt is an operator-visible accounting fault, never a tool failure.
+      input.observability.warn("web search usage settlement failed", {
+        operation: cost.operation,
+        provider: cost.provider,
+        error: error instanceof Error ? error.message : String(error),
+      });
+    }
+  };
+  const definitions: AttemptToolDefinition[] = [];
+
+  if (input.tools.includes(WEB_SEARCH_TOOL_NAME)) {
+    const provider = createWebSearchProvider({
+      ...adapterInput,
+      endpoint: config.search,
+    });
+    const pricing = webSearchCallPricing(config, "search");
+    definitions.push({
+      identity: { serverId: "opengeni", toolName: WEB_SEARCH_TOOL_NAME },
+      modelName: WEB_SEARCH_TOOL_NAME,
+      codemodePath: ["opengeni", WEB_SEARCH_TOOL_NAME],
+      title: "Search the web",
+      description: WEB_SEARCH_TOOL_DESCRIPTION,
+      inputSchema: webSearchInputSchema as unknown as AttemptToolDefinition["inputSchema"],
+      annotations: {
+        title: "Search the web",
+        readOnlyHint: true,
+        destructiveHint: false,
+        idempotentHint: true,
+        openWorldHint: true,
+      },
+      source: "opengeni",
+      approval: "none",
+      execute: async (args, context) => {
+        let request: ReturnType<typeof parseWebSearchArguments>;
+        try {
+          request = parseWebSearchArguments(args);
+          await input.billing.admit(input.scope, pricing.providerMicros);
+        } catch (error) {
+          if (error instanceof WebSearchBillingRefusedError) {
+            recordWebSearchCall(
+              input.observability,
+              "search",
+              config.search.provider,
+              "billing_refused",
+              performance.now(),
+            );
+          }
+          if (
+            error instanceof WebToolArgumentError ||
+            error instanceof WebSearchBillingRefusedError
+          )
+            return textResult(error.message, true);
+          throw error;
+        }
+        const startedAt = performance.now();
+        try {
+          const response = await provider.search(request, {
+            signal: context.signal,
+          });
+          recordWebSearchCall(
+            input.observability,
+            "search",
+            config.search.provider,
+            "ok",
+            startedAt,
+          );
+          await settle(
+            settledCost(
+              context.operationId,
+              "search",
+              config.search.provider,
+              pricing,
+              response.reportedCostMicros,
+            ),
+          );
+          return textResult(renderWebSearchResults(request.query, response.results), false);
+        } catch (error) {
+          if (context.signal?.aborted) throw error;
+          if (error instanceof WebSearchProviderError) {
+            recordWebSearchCall(
+              input.observability,
+              "search",
+              config.search.provider,
+              providerFailureOutcome(error),
+              startedAt,
+            );
+            input.observability.warn("web search provider call failed", {
+              operation: "search",
+              provider: config.search.provider,
+              status: error.status ?? undefined,
+              retryable: error.retryable,
+              error: error.message,
+            });
+            return textResult(`Web search failed: ${providerFailure(error)}`, true);
+          }
+          recordWebSearchCall(
+            input.observability,
+            "search",
+            config.search.provider,
+            "error",
+            startedAt,
+          );
+          throw error;
+        }
+      },
+    });
+  }
+
+  if (input.tools.includes(WEB_FETCH_TOOL_NAME) && config.fetch) {
+    const fetchEndpoint = config.fetch;
+    const provider = createWebFetchProvider({
+      ...adapterInput,
+      endpoint: fetchEndpoint,
+    });
+    const pricing = webSearchCallPricing(config, "fetch");
+    const pages = new Map<string, WebPage>();
+    definitions.push({
+      identity: { serverId: "opengeni", toolName: WEB_FETCH_TOOL_NAME },
+      modelName: WEB_FETCH_TOOL_NAME,
+      codemodePath: ["opengeni", WEB_FETCH_TOOL_NAME],
+      title: "Read a web page",
+      description: WEB_FETCH_TOOL_DESCRIPTION,
+      inputSchema: webFetchInputSchema as unknown as AttemptToolDefinition["inputSchema"],
+      annotations: {
+        title: "Read a web page",
+        readOnlyHint: true,
+        destructiveHint: false,
+        idempotentHint: true,
+        openWorldHint: true,
+      },
+      source: "opengeni",
+      approval: "none",
+      execute: async (args, context) => {
+        let request: ReturnType<typeof parseWebFetchArguments>;
+        try {
+          request = parseWebFetchArguments(args);
+        } catch (error) {
+          if (error instanceof WebToolArgumentError) return textResult(error.message, true);
+          throw error;
+        }
+        const cached = pages.get(request.url);
+        if (cached) return textResult(renderWebPageWindow(cached, request), false);
+        try {
+          await input.billing.admit(input.scope, pricing.providerMicros);
+        } catch (error) {
+          if (error instanceof WebSearchBillingRefusedError) {
+            recordWebSearchCall(
+              input.observability,
+              "fetch",
+              fetchEndpoint.provider,
+              "billing_refused",
+              performance.now(),
+            );
+            return textResult(error.message, true);
+          }
+          throw error;
+        }
+        const startedAt = performance.now();
+        try {
+          const page = await provider.fetch({ url: request.url }, { signal: context.signal });
+          recordWebSearchCall(
+            input.observability,
+            "fetch",
+            fetchEndpoint.provider,
+            "ok",
+            startedAt,
+          );
+          await settle(
+            settledCost(
+              context.operationId,
+              "fetch",
+              fetchEndpoint.provider,
+              pricing,
+              page.reportedCostMicros,
+            ),
+          );
+          if (pages.size >= PAGE_CACHE_ENTRIES) {
+            pages.delete(pages.keys().next().value!);
+          }
+          // Bound per-attempt memory; one extra character keeps the cut visible.
+          pages.set(request.url, {
+            ...page,
+            content: page.content.slice(0, WEB_FETCH_RETAINED_MAX_CHARS + 1),
+          });
+          return textResult(renderWebPageWindow(page, request), false);
+        } catch (error) {
+          if (context.signal?.aborted) throw error;
+          if (error instanceof WebSearchProviderError) {
+            recordWebSearchCall(
+              input.observability,
+              "fetch",
+              fetchEndpoint.provider,
+              providerFailureOutcome(error),
+              startedAt,
+            );
+            input.observability.warn("web search provider call failed", {
+              operation: "fetch",
+              provider: fetchEndpoint.provider,
+              status: error.status ?? undefined,
+              retryable: error.retryable,
+              error: error.message,
+            });
+            return textResult(`Web fetch failed: ${providerFailure(error)}`, true);
+          }
+          recordWebSearchCall(
+            input.observability,
+            "fetch",
+            fetchEndpoint.provider,
+            "error",
+            startedAt,
+          );
+          throw error;
+        }
+      },
+    });
+  }
+  return definitions;
+}

package/src/activities/codemode-dispatcher.ts

@@ -61,6 +61,12 @@ export type CodemodeDispatcherTimings = {
   claimHeartbeatMs?: number;
 };
 
+const RESUMABLE_STATES: readonly CodemodeOperation["state"][] = [
+  "waiting_for_approval",
+  "queued",
+  "running",
+];
+
 /**
  * Owns the only Codemode execution edge for one exact attempt. NATS carries a
  * wake-up only; the durable row is claimed before any side effect and the call
@@ -190,11 +196,20 @@ export class CodemodeAttemptDispatcher {
     );
   }
 
-  async resumeApproved(
+  /**
+   * Join-on-use recovery before the first model request. The journal read needs
+   * no tool environment, so a turn with nothing unfinished returns at once and
+   * never waits on lazy MCP preparation. Only a stored waiting/queued/running
+   * operation awaits the prepared dispatcher to adopt and execute it.
+   */
+  static async resumeApproved(
+    db: Database,
+    scope: CodemodeDispatcherScope,
     callerSubjectId: string,
-    decisionOperationId?: string,
+    decisionOperationId: string | undefined,
+    prepared: () => Promise<CodemodeAttemptDispatcher | null>,
   ): Promise<CodemodeOperation[]> {
-    const operations = await listTurnCodemodeApprovals(this.db, this.scope);
+    const operations = await listTurnCodemodeApprovals(db, scope);
     if (
       decisionOperationId &&
       // Model-tool approval IDs are opaque strings. Only programmatic UUIDs
@@ -202,15 +217,38 @@ export class CodemodeAttemptDispatcher {
       CodemodeOperation.shape.operationId.safeParse(decisionOperationId).success &&
       !operations.some((operation) => operation.operationId === decisionOperationId)
     ) {
-      const decided = await readTurnCodemodeOperation(this.db, {
-        ...this.scope,
+      const decided = await readTurnCodemodeOperation(db, {
+        ...scope,
         callerSubjectId,
         operationId: decisionOperationId,
       });
       if (decided) operations.push(decided);
     }
+    if (!operations.some((operation) => RESUMABLE_STATES.includes(operation.state))) {
+      return operations;
+    }
+    return (await (await prepared())?.resumeOperations(operations, callerSubjectId)) ?? [];
+  }
+
+  async resumeApproved(
+    callerSubjectId: string,
+    decisionOperationId?: string,
+  ): Promise<CodemodeOperation[]> {
+    return await CodemodeAttemptDispatcher.resumeApproved(
+      this.db,
+      this.scope,
+      callerSubjectId,
+      decisionOperationId,
+      async () => this,
+    );
+  }
+
+  private async resumeOperations(
+    operations: CodemodeOperation[],
+    callerSubjectId: string,
+  ): Promise<CodemodeOperation[]> {
     for (const operation of operations) {
-      if (!["waiting_for_approval", "queued", "running"].includes(operation.state)) continue;
+      if (!RESUMABLE_STATES.includes(operation.state)) continue;
       let effectDigest: string | null = null;
       try {
         effectDigest = this.environment.effectDigest(operation.identity);
@@ -380,7 +418,17 @@ export class CodemodeAttemptDispatcher {
     try {
       if (operation.durableApproval) {
         if (!this.approvalWait) throw new Error("Durable approval yield is unavailable");
-        releaseWait = this.approvalWait.beginWait();
+        // The wait gate is sealed while a model request is in flight or the
+        // stream is settling. A call arriving then cannot durably wait, but it
+        // must not fail outright: prepare it like a non-waiting client, so an
+        // allowed call still executes and an Ask settles approval_required.
+        try {
+          releaseWait = this.approvalWait.beginWait();
+        } catch {
+          releaseWait = undefined;
+        }
+      }
+      if (releaseWait) {
         if (
           !(await bindCodemodePreparation(this.db, {
             ...this.scope,
@@ -407,10 +455,11 @@ export class CodemodeAttemptDispatcher {
         },
         {
           signal,
-          ...(operation.durableApproval ? { transportMeta: { durableApproval: true } } : {}),
+          ...(releaseWait ? { transportMeta: { durableApproval: true } } : {}),
         },
       );
       if (preparedCall.waitingForApproval) {
+        if (!releaseWait) throw new AttemptToolApprovalRequiredError();
         const waiting = await waitForCodemodeApproval(this.db, {
           ...this.scope,
           operationId: operation.operationId,

package/src/activities/connector-attachments.ts

@@ -12,6 +12,7 @@ import {
 import {
   ChannelAPartialMutationError,
   RoutingMutationOutcomeUnknownError,
+  isRoutingMutationOutputRejectedError,
   SandboxChannelAService,
 } from "@opengeni/runtime/sandbox";
 
@@ -159,7 +160,8 @@ export async function materializeConnectorAttachmentsInChannel(
   } catch (error) {
     if (
       error instanceof ChannelAPartialMutationError ||
-      error instanceof RoutingMutationOutcomeUnknownError
+      error instanceof RoutingMutationOutcomeUnknownError ||
+      isRoutingMutationOutputRejectedError(error)
     ) {
       throw error;
     }

package/src/activities/sandbox-lease.ts

@@ -53,6 +53,7 @@ import {
   recordRecoveredModalProviderCreate,
   markWarmBillingStopCutoff,
   listMeterableWarmLeases,
+  countInteractionOnlyWarmLeasesByIdle,
   listSandboxViewerForceDrainWorkspaceIds,
   markSandboxCheckpointArtifactDeletePending,
   persistDrainSnapshot,
@@ -174,6 +175,8 @@ import {
   type RetainedProcessReconciliationOutcome,
   type SandboxInventoryProjectionDomain,
   recordSandboxLeaseGauges,
+  recordModalSandboxInventoryGauges,
+  recordInteractionOnlyLeaseGauges,
   recordSandboxOrphansTerminated,
   recordSandboxCommandContainment,
   recordSandboxProviderMissingBeforeCapture,
@@ -2688,6 +2691,17 @@ async function refreshQueueLeaseAndCreditGauges(
         );
       },
     ),
+    refreshSandboxInventoryGauge(
+      observability,
+      "interaction_idle",
+      "interaction-only-lease",
+      async () => {
+        recordInteractionOnlyLeaseGauges(
+          observability,
+          await countInteractionOnlyWarmLeasesByIdle(db, await listMeterableWarmLeases(db)),
+        );
+      },
+    ),
     refreshSandboxInventoryGauge(observability, "expired_drains", "expired-draining", async () => {
       recordExpiredDrainingSandboxLeaseGauges(
         observability,
@@ -2876,6 +2890,16 @@ export function modalOrphanTerminationStillEligible(
   if (latest.some((lease) => lease.liveness === "warming" && lease.instanceId === null)) {
     return false;
   }
+  return modalOrphanCandidateStillUnowned(latest, candidate);
+}
+
+/** Whether no live lease owns the candidate box by exact instance or by its
+ * attribution tags. Unlike termination eligibility this ignores the pending
+ * create postponement: a postponed orphan is still running without a lease. */
+export function modalOrphanCandidateStillUnowned(
+  latest: Awaited<ReturnType<typeof listLiveModalSandboxLeaseAttributions>>,
+  candidate: ModalOrphanSweepTermination,
+): boolean {
   if (latest.some((lease) => lease.instanceId === candidate.sandboxId)) {
     return false;
   }
@@ -2925,6 +2949,30 @@ async function sweepModalOrphansForConfiguredBackend(
       return modalOrphanTerminationStillEligible(latest, candidate);
     },
   });
+  // Only a complete listing is an inventory; a pass cut short by the
+  // termination budget leaves the previous projection to age out as stale.
+  // Re-read durable ownership after the listing so a lease that started or
+  // finished while the app was being listed is not reported as a leak.
+  if (result.inventory.complete) {
+    try {
+      const latest = await listLiveModalSandboxLeaseAttributions(db);
+      recordModalSandboxInventoryGauges(observability, {
+        running: result.inventory.running,
+        unleased: result.inventory.unterminated.filter((candidate) =>
+          modalOrphanCandidateStillUnowned(latest, candidate),
+        ).length,
+        liveLeaseInstancesMissing: result.inventory.missingLiveLeaseInstanceIds.filter(
+          (instanceId) => latest.some((lease) => lease.instanceId === instanceId),
+        ).length,
+      });
+      recordSandboxInventoryProjectionSuccess(observability, "modal_provider");
+    } catch (error) {
+      recordSandboxInventoryProjectionFailure(observability, "modal_provider");
+      observability.warn("sandbox reaper: Modal inventory ownership re-read failed", {
+        error: error instanceof Error ? error.message : String(error),
+      });
+    }
+  }
   for (const terminated of result.terminated) {
     observability.warn("sandbox reaper: terminated Modal orphan sandbox", {
       sandboxId: terminated.sandboxId,
@@ -3524,15 +3572,18 @@ async function terminateDrainableBox(
   // with draining->cold. Until this succeeds, arrivals remain fenced by that
   // exact claim; a timestamp or a failed provider call can never reopen a box
   // while termination may still be in flight.
-  const { wentCold, backgroundCommandEvents } = await confirmDrainCold(db, {
-    accountId,
-    workspaceId: row.workspaceId,
-    sandboxGroupId: row.sandboxGroupId,
-    expectedEpoch: row.leaseEpoch,
-    ...(captureClaim ? { expectedCaptureId: captureClaim.id } : {}),
-    providerMissingBeforeCapture: providerMissing,
-    idleCommandContainmentMs: settings.sandboxIdleCommandContainmentMs,
-  });
+  const { wentCold, unpublishedProviderLost, backgroundCommandEvents } = await confirmDrainCold(
+    db,
+    {
+      accountId,
+      workspaceId: row.workspaceId,
+      sandboxGroupId: row.sandboxGroupId,
+      expectedEpoch: row.leaseEpoch,
+      ...(captureClaim ? { expectedCaptureId: captureClaim.id } : {}),
+      providerMissingBeforeCapture: providerMissing,
+      idleCommandContainmentMs: settings.sandboxIdleCommandContainmentMs,
+    },
+  );
   // The command terminal events and agent inputs are already durable in the
   // cold commit; this is only best-effort live fanout.
   await publishDurableSessionEvents(bus, row.workspaceId, backgroundCommandEvents, (error) => {
@@ -3545,7 +3596,9 @@ async function terminateDrainableBox(
     // Only the exact successful cold commit counts provider loss. A missing
     // probe, a stale capture, a failed commit, or a retried child is not another
     // observed loss. Keep this outside the best-effort session event writer.
-    if (providerMissing) {
+    // An unpublished warming replacement (its creator died before warm
+    // publication) held no workspace; its absence is not a capture loss.
+    if (providerMissing && !unpublishedProviderLost) {
       recordSandboxProviderMissingBeforeCapture(observability, backend);
     }
     if (lease.unobservableCommandDrainIds?.length) {

package/src/activity-services.ts

@@ -1,5 +1,6 @@
 import {
   codeSearchDeploymentPolicy,
+  resolveWebSearchProvider,
   dbSearchPath,
   getSettings,
   resolveNatsControlPlaneAuth,
@@ -52,6 +53,13 @@ export function createSharedActivityServices(
       // Child, scheduled and automation sessions created here freeze their
       // code_search decision from the same deployment policy as the API.
       configureCodeSearchDeploymentPolicy(codeSearchDeploymentPolicy(settings));
+      // A named but unusable web search provider withholds the tools; say why.
+      const webSearchProvider = resolveWebSearchProvider(settings);
+      if (webSearchProvider.status === "invalid") {
+        observability.warn("web search provider is misconfigured; web_search stays off", {
+          reason: webSearchProvider.reason,
+        });
+      }
       return {
         settings: resolvedSettings,
         catalogSourceSettings: settings,