@opengeni/worker-bundle 2.0.1 → 2.0.3-canary.36199476632001

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

@@ -61,7 +61,7 @@ import {
 import {
   compactRetainedScreenshotHistory,
   sdkEventContainsInlineImage,
-  retainComputerScreenshot,
+  retainSessionScreenshot,
   typedScreenshotFromSdkEvent,
   unavailableRetainedSessionImage,
 } from "../retained-screenshots";
@@ -108,6 +108,7 @@ import {
 import {
   pendingToolCallFromSdkEvent,
   toolCallProducesRetainableSessionImage,
+  retainableBrowserScreenshotToolCall,
   completedToolCallFromSdkEvent,
 } from "./history";
 import { checkpointHistoryBeforeProviderDispatch } from "./provider-dispatch-barrier";
@@ -125,7 +126,10 @@ import {
   recordModelUsageAndDebitCredits,
   recordAuthoritativeModelCallFact,
 } from "./model-usage";
-import { startParallelSessionTitleGeneration } from "./session-title";
+import {
+  sessionTitleGenerationOptions,
+  startParallelSessionTitleGeneration,
+} from "./session-title";
 import {
   assertAgentStreamNotCancelled,
   assertSuccessfulAgentStreamCompletion,
@@ -1125,6 +1129,12 @@ export async function runTurnStreamAttempt(
           currentToolBatchCallIds.add(pendingToolCall.callId);
           if (toolCallProducesRetainableSessionImage(pendingToolCall.callName)) {
             media.retainedSessionImageCallIds.add(pendingToolCall.callId);
+            if (retainableBrowserScreenshotToolCall(pendingToolCall.callName)) {
+              media.retainedSessionImageKindsByCallId.set(
+                pendingToolCall.callId,
+                "browser_screenshot",
+              );
+            }
           }
         }
         const completedToolCall = completedToolCallFromSdkEvent(durableSdkEvent);
@@ -1147,6 +1157,7 @@ export async function runTurnStreamAttempt(
               toolCallId: completedToolCall.callId,
               toolOutputId: completedToolCall.callId,
               reason: "unsupported",
+              kind: media.retainedSessionImageKindsByCallId.get(completedToolCall.callId),
             });
             media.retainedScreenshotReceiptsByCallId.set(
               completedToolCall.callId,
@@ -1169,12 +1180,13 @@ export async function runTurnStreamAttempt(
               toolCallId: typedScreenshot.callId,
               toolOutputId: typedScreenshot.toolOutputId,
               reason: "pending",
+              kind: media.retainedSessionImageKindsByCallId.get(completedToolCall.callId),
             });
             media.retainedScreenshotReceiptsByCallId.set(
               completedToolCall.callId,
               retainedScreenshotMetadata,
             );
-            retainedScreenshotMetadata = await retainComputerScreenshot({
+            retainedScreenshotMetadata = await retainSessionScreenshot({
               db,
               objectStorage,
               accountId: input.accountId,
@@ -1183,6 +1195,7 @@ export async function runTurnStreamAttempt(
               turnId: activeTurnId,
               attemptId: input.attemptId,
               output: typedScreenshot,
+              kind: media.retainedSessionImageKindsByCallId.get(completedToolCall.callId),
             });
             media.retainedScreenshotReceiptsByCallId.set(
               completedToolCall.callId,
@@ -1339,6 +1352,7 @@ export async function runTurnStreamAttempt(
           }
           for (const callId of stableToolCallIdsToClear) {
             media.retainedSessionImageCallIds.delete(callId);
+            media.retainedSessionImageKindsByCallId.delete(callId);
           }
         }
       }
@@ -1775,18 +1789,16 @@ export async function runTurnStreamAttempt(
       signal: runtimeCancellationSignal,
       generate: async (signal) =>
         await withSessionTitleProviderRequestContext(() =>
-          runtime.generateSessionTitle!(runSettings, sessionTitlePrompt, {
-            ...(resolvedModel
-              ? {
-                  client: resolvedModel.client,
-                  provider: resolvedModel.provider,
-                  model: resolvedModel.model,
-                }
-              : {}),
-            modelName: turnExecutionPolicy.upstreamModelId,
-            ...(serviceTier ? { serviceTier } : {}),
-            signal,
-          }),
+          runtime.generateSessionTitle!(
+            runSettings,
+            sessionTitlePrompt,
+            sessionTitleGenerationOptions({
+              resolvedModel,
+              modelName: turnExecutionPolicy.upstreamModelId,
+              serviceTier,
+              signal,
+            }),
+          ),
         ),
       onError: (error) => {
         observability.warn("parallel session title generation failed", {

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

@@ -15,6 +15,7 @@ import {
   organizationModelProviderConnectionActiveForWorkspace,
   persistAttemptToolCatalog,
   prepareConnectorActionApproval,
+  recordUsageEvent,
   previewConnectorActionApproval,
   namedSubjectHasLiveWorkspaceAuthority,
   updateSessionTitleWithEvent,
@@ -101,11 +102,13 @@ import type {
 } from "./turn-context";
 import {
   createSessionTitleAttemptToolDefinition,
+  routeAllowsSessionTitleRequests,
   sessionTitleToolPlan,
   shouldRequestMissingSessionTitle,
 } from "./session-title";
 import { resolveTurnSandboxAccess } from "./turn-sandbox-access";
 import { createListModelsAttemptToolDefinition } from "./list-models";
+import { codeSearchToolDefinitions, codeSearchWorkspaceFromChannel } from "./code-search";
 import { createWorkspaceSkillTools } from "./skill-tools";
 import { loadConfiguredBundledSkills } from "./skill-selection";
 import { guardSkillFilesystem } from "./skill-transfer";
@@ -152,6 +155,7 @@ export type PrepareTurnToolRuntimeDeps = {
   turnExecutionPolicy: ClaimTurnOk["turnExecutionPolicy"];
   trigger: ClaimTurnOk["trigger"];
   runSettings: GovernanceModelOk["runSettings"];
+  resolvedModel: GovernanceModelOk["resolvedModel"];
   lazyToolTransport: GovernanceModelOk["lazyToolTransport"];
   turnTools: ReturnType<typeof withFirstPartyTools>;
   connectionScope: { accountId: string; workspaceId: string };
@@ -163,6 +167,8 @@ export type PrepareTurnToolRuntimeDeps = {
   credentialSubjectId: ClaimTurnOk["credentialSubjectId"];
   interactionInterventionResume: ClaimTurnOk["interactionInterventionResume"];
   runWorkspaceMutationForSandbox: SandboxTurnRuntime["runWorkspaceMutationForSandbox"];
+  /** Deployment and workspace allow the Jev-backed code_search tool. */
+  codeSearchEnabled: boolean;
   throwIfWorkerShuttingDown: () => void;
   throwIfTurnCancelled: () => void;
 };
@@ -371,6 +377,7 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
     turnExecutionPolicy,
     trigger,
     runSettings: canonicalRunSettings,
+    resolvedModel,
     lazyToolTransport,
     turnTools: canonicalTurnTools,
     sandboxArtifactRuntime,
@@ -381,6 +388,7 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
     credentialSubjectId,
     interactionInterventionResume,
     runWorkspaceMutationForSandbox,
+    codeSearchEnabled,
     throwIfWorkerShuttingDown,
     throwIfTurnCancelled,
   } = deps;
@@ -573,6 +581,7 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
       firstPartyMcpPermissions: effectiveFirstPartyPermissions,
     }),
     parallelGenerationAvailable: typeof runtime.generateSessionTitle === "function",
+    routeAllowsTitleRequests: routeAllowsSessionTitleRequests(resolvedModel),
   });
   const googleDrivePublicationAllowed =
     selectedFirstPartyMcpTools.includes("editable_artifact_export") &&
@@ -762,6 +771,60 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
       return ((await prepared.ready) ?? prepared).attemptToolEnvironment;
     },
   });
+  const codeSearchTools = codeSearchToolDefinitions({
+    enabled: codeSearchEnabled,
+    settings: runSettings,
+    backend: activeSandboxBackend ?? groupBoxBackend,
+    machineWorkspaceRoot: sandboxState.machinePrimarySession?.workspaceRoot ?? null,
+    observability,
+    // OpenGeni's Jev key pays for these calls whatever model billing the
+    // workspace uses; record them per workspace so the cost stays visible.
+    recordUsage: async (usage) => {
+      const shared = {
+        accountId: input.accountId,
+        workspaceId: input.workspaceId,
+        sourceResourceType: "code_search",
+        sourceResourceId: usage.operationId,
+        sessionId: input.sessionId,
+        turnId: turn.id,
+        turnAttemptId: input.attemptId,
+      };
+      await recordUsageEvent(db, {
+        ...shared,
+        eventType: "code_search.jev_input_tokens",
+        quantity: usage.jevInputTokens,
+        unit: "tokens",
+        idempotencyKey: `usage:code_search.jev_input_tokens:${input.attemptId}:${usage.operationId}`,
+      });
+      await recordUsageEvent(db, {
+        ...shared,
+        eventType: "code_search.jev_cost",
+        quantity: Math.round(usage.jevCostUsd * 1_000_000),
+        unit: "usd_micros",
+        idempotencyKey: `usage:code_search.jev_cost:${input.attemptId}:${usage.operationId}`,
+      });
+    },
+    workspace: async () => {
+      throwIfWorkerShuttingDown();
+      throwIfTurnCancelled();
+      const access = await resolveTurnSandboxAccess(
+        sandboxState,
+        media.sdkOwnedSandboxSession,
+        "code_search requires a sandbox or Connected Machine.",
+      );
+      const machineRoot = sandboxState.machinePrimarySession?.workspaceRoot;
+      const runAs = sandboxRunAs(runSettings);
+      return codeSearchWorkspaceFromChannel(
+        new SandboxChannelAService({
+          session: access.session,
+          workspaceRoot: machineRoot ?? "/workspace",
+          ...(machineRoot ? { providerPathMode: "workspace-relative" as const } : {}),
+          leaseEpoch: access.leaseEpoch,
+          ...(runAs ? { runAs } : {}),
+        }),
+      );
+    },
+  });
   const attemptToolDefinitions = [
     ...(operationReadStore
       ? [
@@ -903,6 +966,7 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
     ...(googleDrivePublicationTool && googleDrivePublicationAllowed
       ? [googleDrivePublicationTool]
       : []),
+    ...codeSearchTools,
   ];
   recordTurnStartupPhase(observability, {
     phase: "tool_context_preparation",
@@ -1124,6 +1188,7 @@ export async function prepareTurnToolRuntime(deps: PrepareTurnToolRuntimeDeps) {
       ...titleToolPlan.preparationIndependentToolNames,
       "skill_read",
     ],
+    codeSearchAvailable: codeSearchTools.length > 0,
     skillCatalog,
   };
 }

package/src/activities/context-compaction.ts

@@ -9,7 +9,8 @@ import {
 import {
   EmptyCompactionSummaryError,
   REMOTE_COMPACTION_V2_IMPLEMENTATION,
-  SUMMARY_BUFFER_TOKENS,
+  compactionSummaryOutputTokens,
+  buildSummaryItem,
   buildCompactionReplacementHistory,
   buildRemoteV2ReplacementHistory,
   compactionThresholdTokens,
@@ -58,7 +59,13 @@ export type MaybeCompactResult =
  * Codex, call Codex `/codex/responses` with `compaction_trigger` and persist the
  * opaque compaction item. Fail closed — never silently fall back to portable.
  */
-export type CompactionSummarizer = (settings: Settings, input: CompactionItem[]) => Promise<string>;
+export type CompactionSummarizer = ((
+  settings: Settings,
+  input: CompactionItem[],
+) => Promise<string>) & {
+  /** Model-visible instructions and tool schemas outside the history estimate. */
+  estimatePrefixTokens?: () => number;
+};
 
 /** Returns the opaque Codex remote compaction v2 item. */
 export type RemoteCompactionV2Requester = (
@@ -81,7 +88,7 @@ export async function maybeCompactContext(
   // Injectable for tests; defaults to the real provider-aware model call.
   summarize: CompactionSummarizer = (s, m) =>
     summarizeForCompaction(s, m, {
-      maxOutputTokens: SUMMARY_BUFFER_TOKENS,
+      maxOutputTokens: compactionSummaryOutputTokens(s.contextWindowTokens),
     }),
   // Operator-forced (the /compact command): bypass the budget trigger and
   // compact now if there is anything to summarize. Structural guards still hold.
@@ -300,6 +307,7 @@ async function settleSkippedAfterStart(
   reason:
     | "no_history"
     | "replacement_not_smaller"
+    | "replacement_exceeds_model_budget"
     | "replacement_unchanged"
     | "summarization_failed",
 ): Promise<Extract<MaybeCompactResult, { compacted: false }>> {
@@ -455,12 +463,23 @@ async function compactContextPortable(
   const summarized = await summarizeWithCodexOverflowTrimming(summarize, settings, items);
   const summaryBody = summarized.summaryBody;
   const retainedTokens = await retentionTokenCounts(canonicalItems, projectForWire);
+  const outputReserve = compactionSummaryOutputTokens(settings.contextWindowTokens);
+  const structuralBudget = Math.min(
+    contextInputBudgetTokens(settings) || settings.contextWindowTokens - outputReserve,
+    settings.contextWindowTokens - outputReserve,
+  );
+  const prefixTokens = Math.max(0, Math.ceil(summarize.estimatePrefixTokens?.() ?? 0));
+  const summaryTokens = estimateTokens([buildSummaryItem(summaryBody)]);
   const replacementHistory = buildCompactionReplacementHistory(
     canonicalItems,
     summaryBody,
     (item) => retainedTokens.get(item) ?? estimateTokens([item]),
+    Math.min(outputReserve, Math.max(0, structuralBudget - prefixTokens - summaryTokens)),
   );
   const estimatedTokensAfter = estimateTokens(await projectForWire(replacementHistory));
+  if (estimatedTokensAfter + prefixTokens > structuralBudget) {
+    return await settleSkippedAfterStart(db, scope, options, "replacement_exceeds_model_budget");
+  }
   const replacementFingerprint = compactionReplacementFingerprint(replacementHistory);
   const previousReplacementFingerprint = latestCompactionReplacementFingerprint(canonicalItems);
   const summaryIndex =
@@ -520,7 +539,7 @@ async function compactContextPortable(
   };
 }
 
-async function summarizeWithCodexOverflowTrimming(
+export async function summarizeWithCodexOverflowTrimming(
   summarize: CompactionSummarizer,
   settings: Settings,
   activeHistory: CompactionItem[],
@@ -531,17 +550,32 @@ async function summarizeWithCodexOverflowTrimming(
 }> {
   // Codex's estimator is intentionally coarse. Keep the explicit checkpoint
   // request below both the effective input window and the raw window minus the
-  // requested summary, then leave 15% estimator headroom. This changes only the
-  // temporary summarizer input; durable active history remains untouched until
-  // applyContextCompaction succeeds under the attempt fence.
-  const summaryAwareBudget = Math.max(0, settings.contextWindowTokens - SUMMARY_BUFFER_TOKENS);
+  // requested summary. Preserve the full portable history copy on the first
+  // call whenever it fits; only trim further after an actual provider overflow.
+  // Durable active history remains untouched until applyContextCompaction.
+  const summaryAwareBudget = Math.max(
+    0,
+    settings.contextWindowTokens - compactionSummaryOutputTokens(settings.contextWindowTokens),
+  );
   const configuredInputBudget = contextInputBudgetTokens(settings);
   const structuralBudget = Math.min(
     configuredInputBudget > 0 ? configuredInputBudget : summaryAwareBudget,
     summaryAwareBudget,
   );
-  const initialBudget = Math.floor(structuralBudget * 0.85);
+  const prefixTokens = Math.max(0, Math.ceil(summarize.estimatePrefixTokens?.() ?? 0));
+  const initialBudget = Math.max(0, structuralBudget - prefixTokens);
   let preparation = prepareCompactionPromptInput(activeHistory, initialBudget);
+  // A checkpoint prompt without source history cannot summarize that history.
+  // Do not let a plausible-sounding reply replace durable active context.
+  const requireHistory = () => {
+    if (activeHistory.length > 0 && preparation.input.length === 1) {
+      throw new EmptyCompactionSummaryError({
+        stage: "portable_input_budget",
+        reason: "no_history_fit",
+      });
+    }
+  };
+  requireHistory();
   try {
     return {
       summaryBody: await summarize(settings, preparation.input),
@@ -551,15 +585,16 @@ async function summarizeWithCodexOverflowTrimming(
   } catch (error) {
     if (!isContextWindowExceeded(error)) throw error;
     // The provider is more authoritative than the byte/4 estimate. Refit once
-    // to 70% of both the configured target and the actual prepared estimate;
+    // to 40% of both the available target and the actual prepared estimate;
     // then fail terminally with prior history intact. Never issue one failing
-    // request per oldest item. The production incident proved that the
-    // provider can count slightly more than twice the byte/4 estimate, so a
-    // half-size retry is the smallest honest bound for that observed skew.
+    // request per oldest item. The retry is bounded, not a guarantee: the
+    // provider can count slightly more than twice the byte/4 estimate. The
+    // prepared prefix has already been reserved from the available budget.
     const retryBudget = Math.floor(
-      Math.min(initialBudget * 0.5, preparation.estimatedInputTokens * 0.5),
+      Math.min(initialBudget * 0.4, preparation.estimatedInputTokens * 0.4),
     );
     preparation = prepareCompactionPromptInput(activeHistory, retryBudget);
+    requireHistory();
     return {
       summaryBody: await summarize(settings, preparation.input),
       preparation,