@opengeni/worker-bundle 1.2.1-canary.37222965366001 → 1.4.0

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

@@ -0,0 +1,252 @@
+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." : "."}`;
+}
+
+/**
+ * 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: Pick<Observability, "warn">;
+  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 WebToolArgumentError ||
+            error instanceof WebSearchBillingRefusedError
+          )
+            return textResult(error.message, true);
+          throw error;
+        }
+        try {
+          const response = await provider.search(request, { signal: context.signal });
+          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) {
+            return textResult(`Web search failed: ${providerFailure(error)}`, true);
+          }
+          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) return textResult(error.message, true);
+          throw error;
+        }
+        try {
+          const page = await provider.fetch({ url: request.url }, { signal: context.signal });
+          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) {
+            return textResult(`Web fetch failed: ${providerFailure(error)}`, true);
+          }
+          throw error;
+        }
+      },
+    });
+  }
+  return definitions;
+}

package/src/activities/codemode-dispatcher.ts

@@ -1,5 +1,10 @@
 import { randomUUID } from "node:crypto";
-import { bindMcpTelemetry, measureMcpPhase, withMcpCallIdentity } from "@opengeni/observability";
+import {
+  bindMcpTelemetry,
+  measureMcpPhase,
+  withMcpCallIdentity,
+  recordToolApproval,
+} from "@opengeni/observability";
 import {
   AttemptToolApprovalRequiredError,
   AttemptToolCatalogStaleError,
@@ -15,13 +20,18 @@ import {
   CODEMODE_CLAIM_LEASE_MS,
   CODEMODE_CLAIM_HEARTBEAT_MS,
   CODEMODE_MAX_CONCURRENT_CALLS_PER_ATTEMPT,
-  type CodemodeOperation,
+  CodemodeOperation,
   type SessionEvent,
   type ToolDisplayMetadata,
   type ToolFamily,
 } from "@opengeni/contracts";
 import {
   cancelQueuedCodemodeOperationsForAttempt,
+  waitForCodemodeApproval,
+  bindCodemodePreparation,
+  listTurnCodemodeApprovals,
+  adoptCodemodeApprovalWithOutput,
+  readTurnCodemodeOperation,
   claimCodemodeOperation,
   failCodemodeOperation,
   getSessionEventByClientEventId,
@@ -76,6 +86,7 @@ export class CodemodeAttemptDispatcher {
     private readonly toolDisplayMetadata?: (modelName: string) => ToolDisplayMetadata | undefined,
     /** Content-free analytics family for a catalog entry; event enrichment only. */
     private readonly toolFamily?: (entry: AttemptToolCatalogEntry) => ToolFamily | null,
+    private readonly approvalWait?: { beginWait(): (accepted: boolean) => void },
   ) {
     if (
       environment.catalog.accountId !== scope.accountId ||
@@ -179,6 +190,94 @@ export class CodemodeAttemptDispatcher {
     );
   }
 
+  async resumeApproved(
+    callerSubjectId: string,
+    decisionOperationId?: string,
+  ): Promise<CodemodeOperation[]> {
+    const operations = await listTurnCodemodeApprovals(this.db, this.scope);
+    if (
+      decisionOperationId &&
+      // Model-tool approval IDs are opaque strings. Only programmatic UUIDs
+      // can address the Codemode operation journal.
+      CodemodeOperation.shape.operationId.safeParse(decisionOperationId).success &&
+      !operations.some((operation) => operation.operationId === decisionOperationId)
+    ) {
+      const decided = await readTurnCodemodeOperation(this.db, {
+        ...this.scope,
+        callerSubjectId,
+        operationId: decisionOperationId,
+      });
+      if (decided) operations.push(decided);
+    }
+    for (const operation of operations) {
+      if (!["waiting_for_approval", "queued", "running"].includes(operation.state)) continue;
+      let effectDigest: string | null = null;
+      try {
+        effectDigest = this.environment.effectDigest(operation.identity);
+      } catch {
+        /* Removed tool: stale. */
+      }
+      const continuation = await adoptCodemodeApprovalWithOutput(this.db, {
+        ...this.scope,
+        operationId: operation.operationId,
+        callerSubjectId,
+        catalogDigest: this.environment.catalog.digest,
+        effectDigest,
+      });
+      await publishDurableSessionEvents(
+        this.bus,
+        this.scope.workspaceId,
+        this.scope.sessionId,
+        continuation.events,
+      );
+      const adopted = continuation.operation;
+      if (adopted && adopted.state !== operation.state) {
+        const outcome =
+          adopted.state === "queued"
+            ? "resumed"
+            : adopted.state === "outcome_unknown"
+              ? "unknown"
+              : adopted.errorCode === "approval_stale"
+                ? "stale"
+                : "rejected";
+        recordToolApproval(
+          outcome,
+          "continuation",
+          operation.state === "waiting_for_approval"
+            ? Math.max(0, Date.now() - Date.parse(operation.updatedAt))
+            : undefined,
+        );
+      }
+      if (adopted?.state !== "queued") continue;
+      const claim = await claimCodemodeOperation(this.db, {
+        ...this.scope,
+        operationId: adopted.operationId,
+        claimId: randomUUID(),
+        catalogDigest: this.environment.catalog.digest,
+        claimLeaseMs: this.claimLeaseMs,
+      });
+      if (claim.status === "claimed") {
+        const execution = this.execute(claim.operation, claim.claimId, true);
+        this.inFlight.set(adopted.operationId, execution);
+        try {
+          await execution;
+        } finally {
+          this.inFlight.delete(adopted.operationId);
+        }
+      }
+    }
+    const receipts: CodemodeOperation[] = [];
+    for (const operation of operations) {
+      const receipt = await readTurnCodemodeOperation(this.db, {
+        ...this.scope,
+        callerSubjectId,
+        operationId: operation.operationId,
+      });
+      if (receipt) receipts.push(receipt);
+    }
+    return receipts;
+  }
+
   async close(reason = "Execution attempt ended"): Promise<void> {
     if (this.closing) {
       await Promise.allSettled([...this.inFlight.values()]);
@@ -277,18 +376,56 @@ export class CodemodeAttemptDispatcher {
     }
 
     let preparedCall: Awaited<ReturnType<AttemptToolEnvironment["prepareCall"]>>;
+    let releaseWait: ((accepted: boolean) => void) | undefined;
     try {
+      if (operation.durableApproval) {
+        if (!this.approvalWait) throw new Error("Durable approval yield is unavailable");
+        releaseWait = this.approvalWait.beginWait();
+        if (
+          !(await bindCodemodePreparation(this.db, {
+            ...this.scope,
+            operationId: operation.operationId,
+            claimId,
+            effectDigest: this.environment.effectDigest(operation.identity),
+          }))
+        )
+          throw new Error("Operation preparation is stale");
+      }
       preparedCall = await this.environment.prepareCall(
         {
           operationId: operation.operationId,
-          catalogDigest: operation.catalogDigest,
+          catalogDigest: this.environment.catalog.digest,
           identity: operation.identity,
           arguments: operation.arguments,
-          caller: operation.caller,
+          caller: {
+            kind: "codemode",
+            subjectId:
+              operation.caller.subjectId === `sandbox:${operation.attemptId}`
+                ? `sandbox:${this.scope.attemptId}`
+                : operation.caller.subjectId,
+          },
+        },
+        {
+          signal,
+          ...(operation.durableApproval ? { transportMeta: { durableApproval: true } } : {}),
         },
-        { signal },
       );
+      if (preparedCall.waitingForApproval) {
+        const waiting = await waitForCodemodeApproval(this.db, {
+          ...this.scope,
+          operationId: operation.operationId,
+          claimId,
+          ...preparedCall.waitingForApproval,
+          effectDigest: preparedCall.effectDigest!,
+        });
+        if (!waiting) throw new Error("Approval could not be bound to the operation");
+        recordToolApproval("waiting");
+        releaseWait?.(true);
+        return;
+      }
+      releaseWait?.(false);
     } catch (error) {
+      releaseWait?.(false);
       await this.settleWithOutput(operation, claimId, {
         state: "failed",
         errorCode: errorCode(error, signal, false),
@@ -448,6 +585,7 @@ export class CodemodeAttemptDispatcher {
       settlement,
     });
     if (!result.committed || result.events.length === 0) return;
+    recordToolApproval(settlement.state === "outcome_unknown" ? "unknown" : settlement.state);
     await publishDurableSessionEvents(
       this.bus,
       this.scope.workspaceId,
@@ -490,8 +628,8 @@ function assertMatchingCodemodeToolCallCreated(
   if (
     event.type !== "agent.toolCall.created" ||
     event.turnId !== scope.turnId ||
-    event.turnGeneration !== scope.executionGeneration ||
-    event.turnAttemptId !== scope.attemptId ||
+    event.turnGeneration !== operation.executionGeneration ||
+    event.turnAttemptId !== operation.attemptId ||
     event.turnAssociation !== "current" ||
     payload?.id !== operation.operationId ||
     payload?.name !== entry.modelName ||

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;
     }