diff --git a/src/agents/embedded-agent-runner/run/attempt-session-runtime-prepare.test.ts b/src/agents/embedded-agent-runner/run/attempt-session-runtime-prepare.test.ts new file mode 100644 index 000000000000..35ee53c5c871 --- /dev/null +++ b/src/agents/embedded-agent-runner/run/attempt-session-runtime-prepare.test.ts @@ -0,0 +1,300 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + createAnthropicPayloadLogger: vi.fn(), + createCacheTrace: vi.fn(), + createSessionSettleTracker: vi.fn(), + getSessionPromptState: vi.fn(), + installContextGuards: vi.fn(), + prepareAgentSession: vi.fn(), + prepareSessionBoundary: vi.fn(), + prepareSessionManager: vi.fn(), + prepareTrajectory: vi.fn(), + prepareTransport: vi.fn(), +})); + +vi.mock("../../anthropic-payload-log.js", () => ({ + createAnthropicPayloadLogger: mocks.createAnthropicPayloadLogger, +})); +vi.mock("../../cache-trace.js", () => ({ createCacheTrace: mocks.createCacheTrace })); +vi.mock("../session-prompt-state.js", () => ({ + getEmbeddedSessionPromptState: mocks.getSessionPromptState, +})); +vi.mock("./attempt-context-guards.js", () => ({ + installEmbeddedAttemptContextGuards: mocks.installContextGuards, +})); +vi.mock("./attempt-session-boundary.js", () => ({ + prepareEmbeddedAttemptSessionBoundary: mocks.prepareSessionBoundary, +})); +vi.mock("./attempt-session-manager-prepare.js", () => ({ + prepareEmbeddedAttemptSessionManager: mocks.prepareSessionManager, +})); +vi.mock("./attempt-session-settle.js", () => ({ + createEmbeddedAttemptSessionSettleTracker: mocks.createSessionSettleTracker, +})); +vi.mock("./attempt-session.js", () => ({ + prepareEmbeddedAttemptAgentSession: mocks.prepareAgentSession, +})); +vi.mock("./attempt-stream-transport.js", () => ({ + prepareEmbeddedAttemptTransport: mocks.prepareTransport, +})); +vi.mock("./attempt-trajectory.js", () => ({ + prepareEmbeddedAttemptTrajectory: mocks.prepareTrajectory, +})); + +import { prepareEmbeddedAttemptSessionRuntime } from "./attempt-session-runtime-prepare.js"; + +type PrepareInput = Parameters[0]; + +function createFixture() { + const order: string[] = []; + const sessionManager = { kind: "manager" }; + const activeSession = { + messages: [{ role: "user" }, { role: "assistant" }], + sessionId: "active-session", + }; + const settingsManager = { kind: "settings" }; + const setActiveSessionSystemPrompt = vi.fn(); + const agentSession = { + activeSession, + clientToolDefs: [{ name: "read" }, { name: "write" }], + setActiveSessionSystemPrompt, + settingsManager, + }; + const boundary = { setCurrentUserTimestampOverride: vi.fn() }; + const promptState = { toolResults: { projected: true } }; + const abortActiveSession = vi.fn(async () => undefined); + const buildAbortSettlePromise = vi.fn(() => null); + const trackPromptSettlePromise = vi.fn((promise: Promise) => promise); + const settleTracker = { + abortActiveSession, + buildAbortSettlePromise, + trackPromptSettlePromise, + }; + const contextGuards = { + getAfterTurnCheckpoint: vi.fn(() => null), + remove: vi.fn(), + takePendingMidTurnPrecheckRequest: vi.fn(() => null), + }; + const cacheTrace = { kind: "cache-trace" }; + const anthropicPayloadLogger = { kind: "payload-logger" }; + const trajectoryRecorder = { kind: "trajectory" }; + const transport = { + effectiveAgentTransport: "sse", + effectiveExtraParams: { cacheRetention: "long" }, + effectivePromptCacheRetention: "long", + providerTextTransforms: undefined, + streamStrategy: "provider", + }; + const transcriptPolicy = { repairToolUseResultPairing: true }; + + mocks.prepareSessionManager.mockImplementation(async (input) => { + order.push("manager"); + input.onSessionManagerCreated(sessionManager); + return { + isOpenAIResponsesApi: true, + preparedUserTurnMessage: { role: "user", content: "hello" }, + sessionManager, + transcriptPolicy, + }; + }); + mocks.prepareAgentSession.mockImplementation(async (input) => { + order.push("agent-session"); + input.onSessionCreated(activeSession); + input.onSystemPromptChanged("runtime prompt"); + return agentSession; + }); + mocks.prepareSessionBoundary.mockImplementation(() => { + order.push("boundary"); + return boundary; + }); + mocks.getSessionPromptState.mockImplementation(() => { + order.push("prompt-state"); + return promptState; + }); + mocks.createSessionSettleTracker.mockImplementation(() => { + order.push("settle-tracker"); + return settleTracker; + }); + mocks.installContextGuards.mockImplementation(() => { + order.push("context-guards"); + return contextGuards; + }); + mocks.createCacheTrace.mockImplementation(() => { + order.push("cache-trace"); + return cacheTrace; + }); + mocks.createAnthropicPayloadLogger.mockImplementation(() => { + order.push("payload-logger"); + return anthropicPayloadLogger; + }); + mocks.prepareTrajectory.mockImplementation(async () => { + order.push("trajectory"); + return trajectoryRecorder; + }); + mocks.prepareTransport.mockImplementation(async () => { + order.push("transport"); + return transport; + }); + + const lifecycle = { + onContextGuardsInstalled: vi.fn(() => order.push("own-context-guards")), + onSessionCreated: vi.fn(() => order.push("own-session")), + onSessionManagerCreated: vi.fn(() => order.push("own-manager")), + onSessionSettleTrackerReady: vi.fn(() => order.push("own-settle-tracker")), + onSessionYieldReady: vi.fn(() => order.push("own-yield")), + onTrajectoryRecorderCreated: vi.fn(() => order.push("own-trajectory")), + }; + const externalAbortController = { + setActiveSessionAbort: vi.fn(() => order.push("arm-session-abort")), + }; + const input = { + attempt: { + model: { api: "openai-responses" }, + modelId: "gpt-5", + provider: "openai", + runId: "run-1", + sessionId: "session-1", + workspaceDir: "/workspace", + }, + agentDir: "/agent", + effectiveCwd: "/workspace", + effectiveWorkspace: "/workspace", + initialSystemPrompt: "initial prompt", + isRawModelRun: false, + sessionManager: { + replayAllowedToolNames: new Set(["read"]), + resolveActiveContextEnginePluginId: vi.fn(), + sessionAgentId: "main", + sessionLockController: {}, + withOwnedSessionWriteLock: vi.fn(), + }, + agentSession: { + agentCoreThinkingLevel: "medium", + clientToolPreparation: {}, + getCurrentAttemptPluginMetadataSnapshot: vi.fn(), + markStage: vi.fn(), + runAbortSignal: new AbortController().signal, + }, + contextGuards: { computerContextEpoch: { value: 0 } }, + trajectory: { effectiveToolCount: 4, localModelLeanEnabled: false }, + transport: { + abortSignal: new AbortController().signal, + codeModeControlsEnabled: false, + getProviderRuntimeHandle: vi.fn(), + providerThinkingLevel: "medium", + sandboxSessionKey: "sandbox-1", + }, + externalAbortController, + lifecycle, + } as unknown as PrepareInput; + + return { + abortActiveSession, + activeSession, + anthropicPayloadLogger, + boundary, + buildAbortSettlePromise, + cacheTrace, + contextGuards, + externalAbortController, + input, + lifecycle, + order, + promptState, + sessionManager, + settingsManager, + trajectoryRecorder, + transport, + }; +} + +beforeEach(() => { + vi.clearAllMocks(); +}); + +describe("prepareEmbeddedAttemptSessionRuntime", () => { + it("prepares the session runtime in ownership-safe order and keeps prompt state live", async () => { + const fixture = createFixture(); + + const result = await prepareEmbeddedAttemptSessionRuntime(fixture.input); + + expect(fixture.order).toEqual([ + "manager", + "own-manager", + "agent-session", + "own-session", + "boundary", + "prompt-state", + "settle-tracker", + "arm-session-abort", + "own-settle-tracker", + "own-yield", + "context-guards", + "own-context-guards", + "cache-trace", + "payload-logger", + "trajectory", + "own-trajectory", + "transport", + ]); + expect(result).toEqual( + expect.objectContaining({ + anthropicPayloadLogger: fixture.anthropicPayloadLogger, + boundary: fixture.boundary, + cacheTrace: fixture.cacheTrace, + contextGuards: fixture.contextGuards, + sessionManager: fixture.sessionManager, + sessionPromptState: fixture.promptState, + toolResultPromptProjectionState: fixture.promptState.toolResults, + trajectoryRecorder: fixture.trajectoryRecorder, + transport: fixture.transport, + }), + ); + expect(result.state).toEqual({ + prePromptMessageCount: 2, + promptCache: undefined, + systemPromptText: "runtime prompt", + }); + expect(fixture.externalAbortController.setActiveSessionAbort).toHaveBeenCalledWith( + fixture.abortActiveSession, + ); + expect(fixture.lifecycle.onSessionSettleTrackerReady).toHaveBeenCalledWith( + fixture.buildAbortSettlePromise, + ); + expect(fixture.lifecycle.onSessionYieldReady).toHaveBeenCalledWith({ + abortActiveSession: fixture.abortActiveSession, + activeSession: fixture.activeSession, + }); + + result.state.prePromptMessageCount = 7; + result.state.promptCache = { cacheRead: 3 } as never; + result.state.systemPromptText = "updated prompt"; + const guardInput = mocks.installContextGuards.mock.calls[0]?.[0]; + expect(guardInput.getPrePromptMessageCount()).toBe(7); + expect(guardInput.getPromptCache()).toEqual({ cacheRead: 3 }); + expect(guardInput.getPromptCacheRetention()).toBe("long"); + expect(guardInput.getSystemPrompt()).toBe("updated prompt"); + }); + + it("publishes every cleanup owner before a later transport failure", async () => { + const fixture = createFixture(); + mocks.prepareTransport.mockRejectedValueOnce(new Error("transport failed")); + + await expect(prepareEmbeddedAttemptSessionRuntime(fixture.input)).rejects.toThrow( + "transport failed", + ); + + expect(fixture.lifecycle.onSessionManagerCreated).toHaveBeenCalledWith(fixture.sessionManager); + expect(fixture.lifecycle.onSessionCreated).toHaveBeenCalledWith(fixture.activeSession); + expect(fixture.lifecycle.onContextGuardsInstalled).toHaveBeenCalledWith( + fixture.contextGuards.remove, + ); + expect(fixture.lifecycle.onSessionSettleTrackerReady).toHaveBeenCalledWith( + fixture.buildAbortSettlePromise, + ); + expect(fixture.lifecycle.onTrajectoryRecorderCreated).toHaveBeenCalledWith( + fixture.trajectoryRecorder, + ); + }); +}); diff --git a/src/agents/embedded-agent-runner/run/attempt-session-runtime-prepare.ts b/src/agents/embedded-agent-runner/run/attempt-session-runtime-prepare.ts new file mode 100644 index 000000000000..63ebe85a97dc --- /dev/null +++ b/src/agents/embedded-agent-runner/run/attempt-session-runtime-prepare.ts @@ -0,0 +1,250 @@ +/** Prepares the session-owned runtime used by one embedded attempt. */ +import { createAnthropicPayloadLogger } from "../../anthropic-payload-log.js"; +import { createCacheTrace } from "../../cache-trace.js"; +import type { guardSessionManager } from "../../session-tool-result-guard-wrapper.js"; +import type { AgentSession } from "../../sessions/index.js"; +import { getEmbeddedSessionPromptState } from "../session-prompt-state.js"; +import type { createEmbeddedAttemptExternalAbortController } from "./attempt-abort.js"; +import { installEmbeddedAttemptContextGuards } from "./attempt-context-guards.js"; +import { prepareEmbeddedAttemptSessionBoundary } from "./attempt-session-boundary.js"; +import { prepareEmbeddedAttemptSessionManager } from "./attempt-session-manager-prepare.js"; +import { createEmbeddedAttemptSessionSettleTracker } from "./attempt-session-settle.js"; +import { prepareEmbeddedAttemptAgentSession } from "./attempt-session.js"; +import { prepareEmbeddedAttemptTransport } from "./attempt-stream-transport.js"; +import { prepareEmbeddedAttemptTrajectory } from "./attempt-trajectory.js"; +import type { EmbeddedRunAttemptParams, EmbeddedRunAttemptResult } from "./types.js"; + +type SessionManagerInput = Parameters[0]; +type AgentSessionInput = Parameters[0]; +type ContextGuardsInput = Parameters[0]; +type TransportInput = Parameters[0]; +type TrajectoryInput = Parameters[0]; +type AttemptSessionManager = ReturnType; +type SessionSettleTracker = ReturnType; +type TrajectoryRecorder = Awaited>; +type ExternalAbortController = Pick< + ReturnType, + "setActiveSessionAbort" +>; + +type EmbeddedAttemptSessionRuntimeState = { + prePromptMessageCount: number; + promptCache: EmbeddedRunAttemptResult["promptCache"]; + systemPromptText: string; +}; + +export async function prepareEmbeddedAttemptSessionRuntime(input: { + attempt: EmbeddedRunAttemptParams; + activeContextEngine?: SessionManagerInput["activeContextEngine"]; + agentDir: string; + effectiveCwd: string; + effectiveWorkspace: string; + initialSystemPrompt: string; + isRawModelRun: boolean; + sessionManager: Pick< + SessionManagerInput, + | "replayAllowedToolNames" + | "resolveActiveContextEnginePluginId" + | "sessionAgentId" + | "sessionLockController" + | "withOwnedSessionWriteLock" + >; + agentSession: Pick< + AgentSessionInput, + | "agentCoreThinkingLevel" + | "clientToolPreparation" + | "getCurrentAttemptPluginMetadataSnapshot" + | "markStage" + | "runAbortSignal" + >; + contextGuards: Pick; + trajectory: Pick< + TrajectoryInput, + "effectiveToolCount" | "localModelLeanEnabled" | "systemPromptReport" + >; + transport: Pick< + TransportInput, + | "abortSignal" + | "codeModeControlsEnabled" + | "getProviderRuntimeHandle" + | "providerThinkingLevel" + | "sandbox" + | "sandboxSessionKey" + >; + externalAbortController: ExternalAbortController; + lifecycle: { + onContextGuardsInstalled: (remove: () => void) => void; + onSessionCreated: (session: AgentSession) => void; + onSessionManagerCreated: (sessionManager: AttemptSessionManager) => void; + onSessionSettleTrackerReady: ( + buildAbortSettlePromise: SessionSettleTracker["buildAbortSettlePromise"], + ) => void; + onSessionYieldReady: (input: { + abortActiveSession: SessionSettleTracker["abortActiveSession"]; + activeSession: AgentSession; + }) => void; + onTrajectoryRecorderCreated: (recorder: TrajectoryRecorder) => void; + }; +}) { + const { attempt } = input; + const preparedSessionManager = await prepareEmbeddedAttemptSessionManager({ + attempt, + ...(input.activeContextEngine ? { activeContextEngine: input.activeContextEngine } : {}), + agentDir: input.agentDir, + effectiveCwd: input.effectiveCwd, + effectiveWorkspace: input.effectiveWorkspace, + onSessionManagerCreated: input.lifecycle.onSessionManagerCreated, + replayAllowedToolNames: input.sessionManager.replayAllowedToolNames, + resolveActiveContextEnginePluginId: input.sessionManager.resolveActiveContextEnginePluginId, + sessionAgentId: input.sessionManager.sessionAgentId, + sessionLockController: input.sessionManager.sessionLockController, + withOwnedSessionWriteLock: input.sessionManager.withOwnedSessionWriteLock, + }); + const { isOpenAIResponsesApi, preparedUserTurnMessage, sessionManager, transcriptPolicy } = + preparedSessionManager; + + const state: EmbeddedAttemptSessionRuntimeState = { + prePromptMessageCount: 0, + promptCache: undefined, + systemPromptText: input.initialSystemPrompt, + }; + const preparedAgentSession = await prepareEmbeddedAttemptAgentSession({ + attempt, + ...(input.activeContextEngine + ? { activeContextEngineInfo: input.activeContextEngine.info } + : {}), + agentCoreThinkingLevel: input.agentSession.agentCoreThinkingLevel, + agentDir: input.agentDir, + clientToolPreparation: input.agentSession.clientToolPreparation, + effectiveCwd: input.effectiveCwd, + getCurrentAttemptPluginMetadataSnapshot: + input.agentSession.getCurrentAttemptPluginMetadataSnapshot, + initialSystemPrompt: state.systemPromptText, + markStage: input.agentSession.markStage, + onSessionCreated: input.lifecycle.onSessionCreated, + onSystemPromptChanged: (systemPromptText) => { + state.systemPromptText = systemPromptText; + }, + runAbortSignal: input.agentSession.runAbortSignal, + sessionAgentId: input.sessionManager.sessionAgentId, + sessionLockController: input.sessionManager.sessionLockController, + sessionManager, + }); + const { activeSession, setActiveSessionSystemPrompt, settingsManager } = preparedAgentSession; + const boundary = prepareEmbeddedAttemptSessionBoundary({ + activeSession, + attempt, + isRawModelRun: input.isRawModelRun, + preparedUserTurnMessage, + sessionManager, + setActiveSessionSystemPrompt, + }); + state.prePromptMessageCount = activeSession.messages.length; + + // Session-owned projections survive attempt teardown so already-sent tool results + // cannot rewrite the provider prompt-cache tail between turns (#99495). + const sessionPromptState = getEmbeddedSessionPromptState(attempt.sessionId); + const toolResultPromptProjectionState = sessionPromptState.toolResults; + const settleTracker = createEmbeddedAttemptSessionSettleTracker(activeSession); + input.externalAbortController.setActiveSessionAbort(settleTracker.abortActiveSession); + input.lifecycle.onSessionSettleTrackerReady(settleTracker.buildAbortSettlePromise); + input.lifecycle.onSessionYieldReady({ + abortActiveSession: settleTracker.abortActiveSession, + activeSession, + }); + + // Guard hooks run during prompt submission, after transport setup fills this value. + const promptCacheRetentionRef: { + current: Awaited< + ReturnType + >["effectivePromptCacheRetention"]; + } = { current: undefined }; + const contextGuards = installEmbeddedAttemptContextGuards({ + ...(input.activeContextEngine ? { activeContextEngine: input.activeContextEngine } : {}), + activeSession, + agentDir: input.agentDir, + attempt, + computerContextEpoch: input.contextGuards.computerContextEpoch, + effectiveCwd: input.effectiveCwd, + effectiveWorkspace: input.effectiveWorkspace, + getPrePromptMessageCount: () => state.prePromptMessageCount, + getPromptCache: () => state.promptCache, + getPromptCacheRetention: () => promptCacheRetentionRef.current, + getSystemPrompt: () => state.systemPromptText, + isOpenAIResponsesApi, + repairToolUseResultPairing: transcriptPolicy.repairToolUseResultPairing, + sessionAgentId: input.sessionManager.sessionAgentId, + sessionManager, + settingsManager, + }); + input.lifecycle.onContextGuardsInstalled(contextGuards.remove); + + const cacheTrace = createCacheTrace({ + cfg: attempt.config, + env: process.env, + runId: attempt.runId, + sessionId: activeSession.sessionId, + sessionKey: attempt.sessionKey, + provider: attempt.provider, + modelId: attempt.modelId, + modelApi: attempt.model.api, + workspaceDir: attempt.workspaceDir, + }); + const anthropicPayloadLogger = createAnthropicPayloadLogger({ + env: process.env, + runId: attempt.runId, + sessionId: activeSession.sessionId, + sessionKey: attempt.sessionKey, + provider: attempt.provider, + modelId: attempt.modelId, + modelApi: attempt.model.api, + workspaceDir: attempt.workspaceDir, + }); + const trajectoryRecorder = await prepareEmbeddedAttemptTrajectory({ + activeSession, + attempt, + clientToolCount: preparedAgentSession.clientToolDefs.length, + effectiveToolCount: input.trajectory.effectiveToolCount, + effectiveWorkspace: input.effectiveWorkspace, + localModelLeanEnabled: input.trajectory.localModelLeanEnabled, + sessionAgentId: input.sessionManager.sessionAgentId, + ...(input.trajectory.systemPromptReport + ? { systemPromptReport: input.trajectory.systemPromptReport } + : {}), + }); + input.lifecycle.onTrajectoryRecorderCreated(trajectoryRecorder); + + const transport = await prepareEmbeddedAttemptTransport({ + attempt, + session: activeSession, + settingsManager, + providerThinkingLevel: input.transport.providerThinkingLevel, + sessionAgentId: input.sessionManager.sessionAgentId, + workspaceDir: input.effectiveWorkspace, + agentDir: input.agentDir, + abortSignal: input.transport.abortSignal, + getProviderRuntimeHandle: input.transport.getProviderRuntimeHandle, + sandboxSessionKey: input.transport.sandboxSessionKey, + ...(input.transport.sandbox !== undefined ? { sandbox: input.transport.sandbox } : {}), + codeModeControlsEnabled: input.transport.codeModeControlsEnabled, + }); + promptCacheRetentionRef.current = transport.effectivePromptCacheRetention; + + return { + agentSession: preparedAgentSession, + anthropicPayloadLogger, + boundary, + cacheTrace, + contextGuards, + isOpenAIResponsesApi, + preparedUserTurnMessage, + sessionManager, + sessionPromptState, + settleTracker, + state, + toolResultPromptProjectionState, + trajectoryRecorder, + transcriptPolicy, + transport, + }; +} diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts index b87b8ec2be82..6a4a3d4c42b8 100644 --- a/src/agents/embedded-agent-runner/run/attempt.ts +++ b/src/agents/embedded-agent-runner/run/attempt.ts @@ -10,8 +10,6 @@ import type { AssistantMessage } from "../../../llm/types.js"; import { createBundleLspToolRuntime } from "../../agent-bundle-lsp-runtime.js"; import { materializeBundleMcpToolsForRun } from "../../agent-bundle-mcp-tools.js"; import { resolveAgentDir, resolveSessionAgentIds } from "../../agent-scope.js"; -import { createAnthropicPayloadLogger } from "../../anthropic-payload-log.js"; -import { createCacheTrace } from "../../cache-trace.js"; import type { AgentMessage } from "../../runtime/index.js"; import type { guardSessionManager } from "../../session-tool-result-guard-wrapper.js"; import type { AgentSession } from "../../sessions/index.js"; @@ -24,22 +22,17 @@ import type { NormalizedUsage } from "../../usage.js"; import { log } from "../logger.js"; import type { PromptCacheBreak, PromptCacheChange } from "../prompt-cache-observability.js"; import { clearActiveEmbeddedRun } from "../runs.js"; -import { getEmbeddedSessionPromptState } from "../session-prompt-state.js"; import { createEmbeddedAttemptExternalAbortController, type EmbeddedAttemptAbortStatePort, } from "./attempt-abort.js"; import { prepareEmbeddedAttemptBootstrap } from "./attempt-bootstrap-prepare.js"; import { prepareEmbeddedAttemptBundleTools } from "./attempt-bundle-tools.js"; -import { installEmbeddedAttemptContextGuards } from "./attempt-context-guards.js"; import { runEmbeddedAttemptPromptPhase } from "./attempt-prompt-phase.js"; import { completeEmbeddedAttemptResult } from "./attempt-result.js"; -import { prepareEmbeddedAttemptSessionBoundary } from "./attempt-session-boundary.js"; import { cleanupEmbeddedAttemptSessionPhase } from "./attempt-session-cleanup.js"; import { prepareEmbeddedAttemptSessionLock } from "./attempt-session-lock-prepare.js"; -import { prepareEmbeddedAttemptSessionManager } from "./attempt-session-manager-prepare.js"; -import { createEmbeddedAttemptSessionSettleTracker } from "./attempt-session-settle.js"; -import { prepareEmbeddedAttemptAgentSession } from "./attempt-session.js"; +import { prepareEmbeddedAttemptSessionRuntime } from "./attempt-session-runtime-prepare.js"; import { prepareEmbeddedAttemptSetup } from "./attempt-setup.js"; import { createEmbeddedRunStageTracker } from "./attempt-stage-timing.js"; import { @@ -49,11 +42,9 @@ import { } from "./attempt-startup.js"; import { finalizeEmbeddedAttemptStreamPhase } from "./attempt-stream-finalize.js"; import { prepareEmbeddedAttemptStreamRuntime } from "./attempt-stream-runtime-prepare.js"; -import { prepareEmbeddedAttemptTransport } from "./attempt-stream-transport.js"; import { prepareEmbeddedAttemptSystemPrompt } from "./attempt-system-prompt-prepare.js"; import { prepareEmbeddedAttemptToolBase } from "./attempt-tool-base-prepare.js"; import { prepareEmbeddedAttemptToolCatalog } from "./attempt-tool-catalog.js"; -import { prepareEmbeddedAttemptTrajectory } from "./attempt-trajectory.js"; import type { EmbeddedAttemptSessionFileOwner } from "./attempt.session-lock.js"; import { queueSessionsYieldInterruptMessage, @@ -321,7 +312,7 @@ export async function runEmbeddedAttempt( toolSearchCatalogRef, }); const { runtimeChannel, runtimeInfo, systemPromptReport } = preparedSystemPrompt; - let systemPromptText = preparedSystemPrompt.systemPromptText; + const initialSystemPromptText = preparedSystemPrompt.systemPromptText; let sessionManager: ReturnType | undefined; const { @@ -343,181 +334,132 @@ export async function runEmbeddedAttempt( let session: AgentSession | undefined; let removeToolResultContextGuard: (() => void) | undefined; - let trajectoryRecorder: Awaited> = null; + let trajectoryRecorder: Awaited< + ReturnType + >["trajectoryRecorder"] = null; let trajectoryEndRecorded = false; let buildAbortSettlePromise: () => Promise | null = () => null; let cleanupYieldAborted = false; let repairedRejectedThinkingReplay = false; try { - const preparedSessionManager = await prepareEmbeddedAttemptSessionManager({ + const preparedSessionRuntime = await prepareEmbeddedAttemptSessionRuntime({ attempt: params, - activeContextEngine, + ...(activeContextEngine ? { activeContextEngine } : {}), agentDir, effectiveCwd, effectiveWorkspace, - onSessionManagerCreated: (createdSessionManager) => { - sessionManager = createdSessionManager; - }, - replayAllowedToolNames, - resolveActiveContextEnginePluginId, - sessionAgentId, - sessionLockController, - withOwnedSessionWriteLock, - }); - const { isOpenAIResponsesApi, preparedUserTurnMessage, transcriptPolicy } = - preparedSessionManager; - sessionManager = preparedSessionManager.sessionManager; - - const { - activeSession, - allCustomTools, - builtinToolNames, - clientToolCallSlots, - clientToolDefs, - clientToolLoopDetection, - hasDeliveredSourceReply, - hookRunner, - markSourceReplyDelivered, - replaySafeToolNames, - replaySafeTools, - setActiveSessionSystemPrompt, - settingsManager, - } = await prepareEmbeddedAttemptAgentSession({ - attempt: params, - activeContextEngineInfo: activeContextEngine?.info, - agentCoreThinkingLevel, - agentDir, - clientToolPreparation: { - catalogToolHookContext, - clientTools, - codeModeControlsEnabledForRun, - deferredDirectoryToolsCallable, - effectiveTools, - replaySafetyOptions, - sandboxEnabled: Boolean(sandbox?.enabled), - sandboxSessionKey, - sessionAgentId, - toolSearchCatalogRef, - toolSearchRuntimeConfig, - uncompactedEffectiveTools, - }, - effectiveCwd, - getCurrentAttemptPluginMetadataSnapshot, - initialSystemPrompt: systemPromptText, - markStage: (stage) => prepStages.mark(stage), - onSessionCreated: (createdSession) => { - session = createdSession; - }, - onSystemPromptChanged: (nextSystemPrompt) => { - systemPromptText = nextSystemPrompt; - }, - runAbortSignal: runAbortController.signal, - sessionAgentId, - sessionLockController, - sessionManager, - }); - const sessionBoundary = prepareEmbeddedAttemptSessionBoundary({ - activeSession, - attempt: params, + initialSystemPrompt: initialSystemPromptText, isRawModelRun, - preparedUserTurnMessage, - sessionManager, - setActiveSessionSystemPrompt, + sessionManager: { + replayAllowedToolNames, + resolveActiveContextEnginePluginId, + sessionAgentId, + sessionLockController, + withOwnedSessionWriteLock, + }, + agentSession: { + agentCoreThinkingLevel, + clientToolPreparation: { + catalogToolHookContext, + clientTools, + codeModeControlsEnabledForRun, + deferredDirectoryToolsCallable, + effectiveTools, + replaySafetyOptions, + sandboxEnabled: Boolean(sandbox?.enabled), + sandboxSessionKey, + sessionAgentId, + toolSearchCatalogRef, + toolSearchRuntimeConfig, + uncompactedEffectiveTools, + }, + getCurrentAttemptPluginMetadataSnapshot, + markStage: (stage) => prepStages.mark(stage), + runAbortSignal: runAbortController.signal, + }, + contextGuards: { computerContextEpoch }, + trajectory: { + effectiveToolCount: effectiveTools.length, + localModelLeanEnabled, + ...(systemPromptReport ? { systemPromptReport } : {}), + }, + transport: { + abortSignal: runAbortController.signal, + codeModeControlsEnabled: codeModeControlsEnabledForRun, + getProviderRuntimeHandle, + providerThinkingLevel, + ...(sandbox !== undefined ? { sandbox } : {}), + sandboxSessionKey, + }, + externalAbortController, + lifecycle: { + onContextGuardsInstalled: (remove) => { + removeToolResultContextGuard = remove; + }, + onSessionCreated: (createdSession) => { + session = createdSession; + }, + onSessionManagerCreated: (createdSessionManager) => { + sessionManager = createdSessionManager; + }, + onSessionSettleTrackerReady: (build) => { + buildAbortSettlePromise = build; + }, + onSessionYieldReady: ({ abortActiveSession, activeSession }) => { + abortSessionForYield = () => { + yieldAbortSettled = abortActiveSession(SESSIONS_YIELD_ABORT_REASON); + }; + queueYieldInterruptForSession = () => { + queueSessionsYieldInterruptMessage(activeSession); + }; + }, + onTrajectoryRecorderCreated: (recorder) => { + trajectoryRecorder = recorder; + }, + }, }); - const { boundaryTimezone, includeBoundaryTimestamp, orphanRepair } = sessionBoundary; - let prePromptMessageCount = activeSession.messages.length; - // Session-owned projections survive attempt teardown so already-sent tool results - // cannot rewrite the provider prompt-cache tail between turns (#99495). - const sessionPromptState = getEmbeddedSessionPromptState(params.sessionId); - const toolResultPromptProjectionState = sessionPromptState.toolResults; - let promptCache: EmbeddedRunAttemptResult["promptCache"]; - const sessionSettleTracker = createEmbeddedAttemptSessionSettleTracker(activeSession); - const { abortActiveSession, trackPromptSettlePromise } = sessionSettleTracker; - externalAbortController.setActiveSessionAbort(abortActiveSession); - buildAbortSettlePromise = sessionSettleTracker.buildAbortSettlePromise; - abortSessionForYield = () => { - yieldAbortSettled = abortActiveSession(SESSIONS_YIELD_ABORT_REASON); - }; - queueYieldInterruptForSession = () => { - queueSessionsYieldInterruptMessage(activeSession); - }; - const contextGuards = installEmbeddedAttemptContextGuards({ - activeContextEngine, - activeSession, - agentDir, - attempt: params, - computerContextEpoch, - effectiveCwd, - effectiveWorkspace, - getPrePromptMessageCount: () => prePromptMessageCount, - getPromptCache: () => promptCache, - getPromptCacheRetention: () => effectivePromptCacheRetention, - getSystemPrompt: () => systemPromptText, - isOpenAIResponsesApi, - repairToolUseResultPairing: transcriptPolicy.repairToolUseResultPairing, - sessionAgentId, - sessionManager, - settingsManager, - }); - removeToolResultContextGuard = contextGuards.remove; - const cacheTrace = createCacheTrace({ - cfg: params.config, - env: process.env, - runId: params.runId, - sessionId: activeSession.sessionId, - sessionKey: params.sessionKey, - provider: params.provider, - modelId: params.modelId, - modelApi: params.model.api, - workspaceDir: params.workspaceDir, - }); - const anthropicPayloadLogger = createAnthropicPayloadLogger({ - env: process.env, - runId: params.runId, - sessionId: activeSession.sessionId, - sessionKey: params.sessionKey, - provider: params.provider, - modelId: params.modelId, - modelApi: params.model.api, - workspaceDir: params.workspaceDir, - }); - trajectoryRecorder = await prepareEmbeddedAttemptTrajectory({ - activeSession, - attempt: params, - clientToolCount: clientToolDefs.length, - effectiveToolCount: effectiveTools.length, - effectiveWorkspace, - localModelLeanEnabled, - sessionAgentId, - ...(systemPromptReport ? { systemPromptReport } : {}), - }); - const { - effectiveAgentTransport, - effectiveExtraParams, - effectivePromptCacheRetention, - providerTextTransforms, - streamStrategy, - } = await prepareEmbeddedAttemptTransport({ - attempt: params, - session: activeSession, - settingsManager, - providerThinkingLevel, - sessionAgentId, - workspaceDir: effectiveWorkspace, - agentDir, - abortSignal: runAbortController.signal, - getProviderRuntimeHandle, - sandboxSessionKey, - sandbox, - codeModeControlsEnabled: codeModeControlsEnabledForRun, - }); + agentSession: { + activeSession, + allCustomTools, + builtinToolNames, + clientToolCallSlots, + clientToolLoopDetection, + hasDeliveredSourceReply, + hookRunner, + markSourceReplyDelivered, + replaySafeToolNames, + replaySafeTools, + setActiveSessionSystemPrompt, + settingsManager, + }, + anthropicPayloadLogger, + boundary: sessionBoundary, + cacheTrace, + contextGuards, + isOpenAIResponsesApi, + preparedUserTurnMessage, + sessionManager: activeSessionManager, + sessionPromptState, + settleTracker: { abortActiveSession, trackPromptSettlePromise }, + state: sessionRuntimeState, + toolResultPromptProjectionState, + transcriptPolicy, + transport: { + effectiveAgentTransport, + effectiveExtraParams, + effectivePromptCacheRetention, + providerTextTransforms, + streamStrategy, + }, + } = preparedSessionRuntime; + const { boundaryTimezone, includeBoundaryTimestamp, orphanRepair } = sessionBoundary; let yieldAborted = false; const hookAgentId = sessionAgentId; const preparedStreamRuntime = await prepareEmbeddedAttemptStreamRuntime({ attempt: params, activeSession, - sessionManager, + sessionManager: activeSessionManager, sessionLockController, ownedTranscriptWriteContext, runAbortController, @@ -530,7 +472,7 @@ export async function runEmbeddedAttempt( sessionAgentId, cacheTrace, allCustomTools, - systemPromptText, + systemPromptText: sessionRuntimeState.systemPromptText, transcriptPolicy, isOpenAIResponsesApi, replayAllowedToolNames, @@ -552,7 +494,7 @@ export async function runEmbeddedAttempt( replayAllowedToolNames, sessionAgentId, settingsManager, - systemPromptText, + systemPromptText: sessionRuntimeState.systemPromptText, transcriptPolicy, setActiveSessionSystemPrompt, }, @@ -637,7 +579,6 @@ export async function runEmbeddedAttempt( let sessionIdUsed = activeSession.sessionId; let sessionFileUsed: string | undefined = params.sessionFile; - const activeSessionManager = sessionManager; let preflightRecovery: EmbeddedRunAttemptResult["preflightRecovery"]; let promptErrorSource: EmbeddedRunAttemptResult["promptErrorSource"] = null; try { @@ -657,7 +598,7 @@ export async function runEmbeddedAttempt( ...(orphanRepair ? { orphanRepair } : {}), sessionAgentId, runtimeModel: runtimeInfo.model, - systemPromptText, + systemPromptText: sessionRuntimeState.systemPromptText, setActiveSessionSystemPrompt, cache: { observabilityEnabled: cacheObservabilityEnabled, @@ -678,7 +619,7 @@ export async function runEmbeddedAttempt( sessionAgentId, setActiveSessionSystemPrompt, ...(systemPromptReport ? { systemPromptReport } : {}), - systemPromptText, + systemPromptText: sessionRuntimeState.systemPromptText, toolResultPromptProjectionState, }, execution: { @@ -699,7 +640,7 @@ export async function runEmbeddedAttempt( isRawModelRun, runTrace, streamStrategy, - systemPromptText, + systemPromptText: sessionRuntimeState.systemPromptText, toolSearchCompacted: toolSearch.compacted, tools, trajectoryRecorder, @@ -736,9 +677,9 @@ export async function runEmbeddedAttempt( promptError = state.promptError; promptErrorSource = state.promptErrorSource; }, - getPrePromptMessageCount: () => prePromptMessageCount, + getPrePromptMessageCount: () => sessionRuntimeState.prePromptMessageCount, setPrePromptMessageCount: (count) => { - prePromptMessageCount = count; + sessionRuntimeState.prePromptMessageCount = count; }, setCurrentUserTimestampOverride: (override) => { sessionBoundary.setCurrentUserTimestampOverride(override); @@ -791,7 +732,7 @@ export async function runEmbeddedAttempt( currentAttemptAssistant = settledStream.currentAttemptAssistant; attemptUsage = settledStream.attemptUsage; cacheBreak = settledStream.cacheBreak; - promptCache = settledStream.promptCache; + sessionRuntimeState.promptCache = settledStream.promptCache; }, getState: () => ({ promptError, @@ -814,7 +755,7 @@ export async function runEmbeddedAttempt( isProbeSession, onBlockReplyFlush, abortable, - prePromptMessageCount, + prePromptMessageCount: sessionRuntimeState.prePromptMessageCount, toolSearchTargetTranscriptProjections, cache: { observabilityEnabled: cacheObservabilityEnabled, @@ -904,7 +845,7 @@ export async function runEmbeddedAttempt( lastAssistant, currentAttemptAssistant, attemptUsage, - promptCache, + promptCache: sessionRuntimeState.promptCache, contextBudgetStatus, yieldDetected, didDeliverSourceReplyViaMessageTool: hasDeliveredSourceReply(),