From 632b17618ec90d5ebb32cea47fd5f6799c9a4882 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 14 Jul 2026 05:03:21 +0100 Subject: [PATCH] refactor(agents): extract attempt prompt dispatch phase --- .../run/attempt-prompt-dispatch.test.ts | 166 ++++++++++++++++++ .../run/attempt-prompt-dispatch.ts | 137 +++++++++++++++ .../embedded-agent-runner/run/attempt.ts | 159 ++++++----------- 3 files changed, 360 insertions(+), 102 deletions(-) create mode 100644 src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.test.ts create mode 100644 src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.ts diff --git a/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.test.ts b/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.test.ts new file mode 100644 index 000000000000..c0f676307854 --- /dev/null +++ b/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.test.ts @@ -0,0 +1,166 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const hoisted = vi.hoisted(() => ({ + observeEmbeddedAttemptPrompt: vi.fn(), + prepareEmbeddedAttemptPromptExecution: vi.fn(), + prepareEmbeddedAttemptPromptPreflight: vi.fn(), + submitEmbeddedAttemptPrompt: vi.fn(), +})); + +vi.mock("./attempt-prompt-execution-prepare.js", () => ({ + prepareEmbeddedAttemptPromptExecution: hoisted.prepareEmbeddedAttemptPromptExecution, +})); +vi.mock("./attempt-prompt-observability.js", () => ({ + observeEmbeddedAttemptPrompt: hoisted.observeEmbeddedAttemptPrompt, +})); +vi.mock("./attempt-prompt-preflight.js", () => ({ + prepareEmbeddedAttemptPromptPreflight: hoisted.prepareEmbeddedAttemptPromptPreflight, +})); +vi.mock("./attempt-prompt-submit.js", () => ({ + submitEmbeddedAttemptPrompt: hoisted.submitEmbeddedAttemptPrompt, +})); + +import { dispatchEmbeddedAttemptPrompt } from "./attempt-prompt-dispatch.js"; + +type DispatchInput = Parameters[0]; +type PreflightMockInput = { state: DispatchInput["state"] }; + +function createInput(overrides: Partial = {}): DispatchInput { + return { + attempt: { runId: "run-1", sessionId: "session-1" }, + activeSession: { messages: [] }, + promptContext: { + contextTokenBudget: 8_000, + effectivePrompt: "effective prompt", + hookMessagesForCurrentPrompt: [], + llmBoundaryPromptForPrecheck: "boundary prompt", + promptForModel: "model prompt", + promptForSession: "session prompt", + promptSubmission: { prompt: "submission prompt", runtimeOnly: false }, + promptToolResultAggregateMaxChars: 8_000, + promptToolResultMaxChars: 4_000, + runtimeContextMessageForCurrentTurn: { role: "custom", content: "runtime" }, + systemPromptForHook: "system prompt", + }, + getCompactionReserveTokens: () => 1_000, + publishState: vi.fn(), + releaseLeasedSteering: vi.fn(), + state: { + contextBudgetStatus: undefined, + preflightRecovery: undefined, + promptError: null, + promptErrorSource: null, + skipPromptSubmission: false, + }, + execution: {}, + observation: {}, + preflight: {}, + submission: {}, + ...overrides, + } as unknown as DispatchInput; +} + +describe("dispatchEmbeddedAttemptPrompt", () => { + beforeEach(() => { + vi.clearAllMocks(); + hoisted.prepareEmbeddedAttemptPromptExecution.mockResolvedValue({ + images: [{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" }], + detectedRefs: [], + loadedCount: 1, + skippedCount: 0, + }); + hoisted.observeEmbeddedAttemptPrompt.mockReturnValue({ skipPromptSubmission: false }); + hoisted.prepareEmbeddedAttemptPromptPreflight.mockImplementation( + async (input: PreflightMockInput) => input.state, + ); + hoisted.submitEmbeddedAttemptPrompt.mockResolvedValue(undefined); + }); + + it("runs image preparation, observability, preflight, and submission in order", async () => { + const order: string[] = []; + hoisted.prepareEmbeddedAttemptPromptExecution.mockImplementationOnce(async () => { + order.push("images"); + return { + images: [{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" }], + detectedRefs: [], + loadedCount: 1, + skippedCount: 0, + }; + }); + hoisted.observeEmbeddedAttemptPrompt.mockImplementationOnce(() => { + order.push("observe"); + return { skipPromptSubmission: false }; + }); + hoisted.prepareEmbeddedAttemptPromptPreflight.mockImplementationOnce( + async (input: PreflightMockInput) => { + order.push("preflight"); + return input.state; + }, + ); + hoisted.submitEmbeddedAttemptPrompt.mockImplementationOnce(async () => { + order.push("submit"); + }); + const publishState = vi.fn(() => { + order.push("publish"); + }); + const input = createInput({ publishState }); + + await expect(dispatchEmbeddedAttemptPrompt(input)).resolves.toEqual(input.state); + + expect(order).toEqual(["images", "observe", "publish", "preflight", "publish", "submit"]); + expect(hoisted.prepareEmbeddedAttemptPromptExecution).toHaveBeenCalledWith( + expect.objectContaining({ prompt: "submission prompt", skipPromptSubmission: false }), + ); + expect(hoisted.observeEmbeddedAttemptPrompt).toHaveBeenCalledWith( + expect.objectContaining({ imageCount: 1, reserveTokens: 1_000 }), + ); + expect(hoisted.submitEmbeddedAttemptPrompt).toHaveBeenCalledWith( + expect.objectContaining({ + images: [expect.objectContaining({ type: "image" })], + modelPrompt: "model prompt", + runtimeContextMessage: expect.objectContaining({ content: "runtime" }), + transcriptPrompt: "session prompt", + }), + ); + }); + + it("releases steering when preflight skips provider submission", async () => { + const promptError = new Error("preflight rejected"); + const releaseLeasedSteering = vi.fn(); + hoisted.observeEmbeddedAttemptPrompt.mockReturnValueOnce({ skipPromptSubmission: true }); + hoisted.prepareEmbeddedAttemptPromptPreflight.mockImplementationOnce( + async (input: PreflightMockInput) => ({ + ...input.state, + promptError, + promptErrorSource: "precheck", + }), + ); + + const result = await dispatchEmbeddedAttemptPrompt(createInput({ releaseLeasedSteering })); + + expect(result.promptError).toBe(promptError); + expect(releaseLeasedSteering).toHaveBeenCalledWith(promptError); + expect(hoisted.submitEmbeddedAttemptPrompt).not.toHaveBeenCalled(); + }); + + it("publishes preflight state before a submission failure", async () => { + const promptError = new Error("admission warning"); + const submitError = new Error("provider failed"); + const admittedState = { + contextBudgetStatus: undefined, + preflightRecovery: undefined, + promptError, + promptErrorSource: "precheck" as const, + skipPromptSubmission: false, + }; + const publishState = vi.fn(); + hoisted.prepareEmbeddedAttemptPromptPreflight.mockResolvedValueOnce(admittedState); + hoisted.submitEmbeddedAttemptPrompt.mockRejectedValueOnce(submitError); + + await expect(dispatchEmbeddedAttemptPrompt(createInput({ publishState }))).rejects.toBe( + submitError, + ); + + expect(publishState).toHaveBeenLastCalledWith(admittedState); + }); +}); diff --git a/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.ts b/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.ts new file mode 100644 index 000000000000..c7bb04dae937 --- /dev/null +++ b/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.ts @@ -0,0 +1,137 @@ +/** Runs prompt-local image preparation, observability, preflight, and provider dispatch. */ +import type { prepareEmbeddedAttemptPromptContext } from "./attempt-prompt-context.js"; +import { prepareEmbeddedAttemptPromptExecution } from "./attempt-prompt-execution-prepare.js"; +import { observeEmbeddedAttemptPrompt } from "./attempt-prompt-observability.js"; +import { prepareEmbeddedAttemptPromptPreflight } from "./attempt-prompt-preflight.js"; +import { submitEmbeddedAttemptPrompt } from "./attempt-prompt-submit.js"; +import type { EmbeddedRunAttemptParams } from "./types.js"; + +type PromptContext = ReturnType; +type PromptExecutionInput = Parameters[0]; +type PromptObservationInput = Parameters[0]; +type PromptPreflightInput = Parameters[0]; +type PromptSubmissionInput = Parameters[0]; +type PromptDispatchState = PromptPreflightInput["state"]; + +export async function dispatchEmbeddedAttemptPrompt(input: { + attempt: EmbeddedRunAttemptParams; + activeContextEngine?: PromptPreflightInput["activeContextEngine"]; + activeSession: PromptExecutionInput["session"] & PromptSubmissionInput["activeSession"]; + promptContext: PromptContext; + getCompactionReserveTokens: () => number; + publishState: (state: PromptDispatchState) => void; + releaseLeasedSteering: (error?: unknown) => void; + state: PromptDispatchState; + execution: Omit; + observation: Omit< + PromptObservationInput, + | "attempt" + | "contextTokenBudget" + | "effectivePrompt" + | "hookMessagesForCurrentPrompt" + | "imageCount" + | "llmBoundaryPromptForPrecheck" + | "promptForModel" + | "promptSubmissionRuntimeOnly" + | "reserveTokens" + | "sessionMessages" + | "skipPromptSubmission" + | "systemPromptForHook" + >; + preflight: Omit< + PromptPreflightInput, + | "attempt" + | "activeContextEngine" + | "contextTokenBudget" + | "hookMessagesForCurrentPrompt" + | "promptForPrecheck" + | "reserveTokens" + | "sessionMessageCount" + | "state" + | "systemPrompt" + | "toolResultMaxChars" + >; + submission: Omit< + PromptSubmissionInput, + | "attempt" + | "activeSession" + | "contextTokenBudget" + | "images" + | "modelPrompt" + | "runtimeContextMessage" + | "runtimeOnly" + | "systemPrompt" + | "toolResultAggregateMaxChars" + | "toolResultMaxChars" + | "transcriptPrompt" + >; +}): Promise { + const { activeSession, attempt, promptContext } = input; + const imageResult = await prepareEmbeddedAttemptPromptExecution({ + ...input.execution, + attempt, + prompt: promptContext.promptSubmission.prompt, + session: activeSession, + skipPromptSubmission: input.state.skipPromptSubmission, + }); + + const reserveTokens = input.getCompactionReserveTokens(); + let state: PromptDispatchState = { + ...input.state, + skipPromptSubmission: observeEmbeddedAttemptPrompt({ + ...input.observation, + attempt, + contextTokenBudget: promptContext.contextTokenBudget, + effectivePrompt: promptContext.effectivePrompt, + hookMessagesForCurrentPrompt: promptContext.hookMessagesForCurrentPrompt, + imageCount: imageResult.images.length, + llmBoundaryPromptForPrecheck: promptContext.llmBoundaryPromptForPrecheck, + promptForModel: promptContext.promptForModel, + promptSubmissionRuntimeOnly: promptContext.promptSubmission.runtimeOnly, + reserveTokens, + sessionMessages: activeSession.messages, + skipPromptSubmission: input.state.skipPromptSubmission, + systemPromptForHook: promptContext.systemPromptForHook, + }).skipPromptSubmission, + }; + // Publish each admission transition before the next fallible phase so outer cleanup sees it. + input.publishState(state); + + state = await prepareEmbeddedAttemptPromptPreflight({ + ...input.preflight, + attempt, + ...(input.activeContextEngine ? { activeContextEngine: input.activeContextEngine } : {}), + contextTokenBudget: promptContext.contextTokenBudget, + hookMessagesForCurrentPrompt: promptContext.hookMessagesForCurrentPrompt, + promptForPrecheck: promptContext.llmBoundaryPromptForPrecheck, + reserveTokens, + sessionMessageCount: activeSession.messages.length, + state, + systemPrompt: promptContext.systemPromptForHook, + toolResultMaxChars: promptContext.promptToolResultMaxChars, + }); + input.publishState(state); + + if (!state.skipPromptSubmission) { + await submitEmbeddedAttemptPrompt({ + ...input.submission, + attempt, + activeSession, + contextTokenBudget: promptContext.contextTokenBudget, + images: imageResult.images, + modelPrompt: promptContext.promptForModel, + ...(promptContext.runtimeContextMessageForCurrentTurn + ? { runtimeContextMessage: promptContext.runtimeContextMessageForCurrentTurn } + : {}), + runtimeOnly: promptContext.promptSubmission.runtimeOnly === true, + systemPrompt: promptContext.systemPromptForHook, + toolResultAggregateMaxChars: promptContext.promptToolResultAggregateMaxChars, + toolResultMaxChars: promptContext.promptToolResultMaxChars, + transcriptPrompt: promptContext.promptForSession, + }); + } else { + input.releaseLeasedSteering(state.promptError ?? "prompt submission skipped"); + } + + return state; +} diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts index 41864ed55e19..71482dea0a89 100644 --- a/src/agents/embedded-agent-runner/run/attempt.ts +++ b/src/agents/embedded-agent-runner/run/attempt.ts @@ -47,14 +47,9 @@ import { installEmbeddedAttemptContextGuards } from "./attempt-context-guards.js import { prepareEmbeddedAttemptHistory } from "./attempt-history-prepare.js"; import { prepareEmbeddedAttemptPromptAssembly } from "./attempt-prompt-assembly.js"; import { prepareEmbeddedAttemptPromptContext } from "./attempt-prompt-context.js"; +import { dispatchEmbeddedAttemptPrompt } from "./attempt-prompt-dispatch.js"; import { handleEmbeddedAttemptPromptError } from "./attempt-prompt-error.js"; -import { prepareEmbeddedAttemptPromptExecution } from "./attempt-prompt-execution-prepare.js"; -import { observeEmbeddedAttemptPrompt } from "./attempt-prompt-observability.js"; -import { - handleEmbeddedAttemptMidTurnPrecheck, - prepareEmbeddedAttemptPromptPreflight, -} from "./attempt-prompt-preflight.js"; -import { submitEmbeddedAttemptPrompt } from "./attempt-prompt-submit.js"; +import { handleEmbeddedAttemptMidTurnPrecheck } from "./attempt-prompt-preflight.js"; import { completeEmbeddedAttemptResult } from "./attempt-result.js"; import { prepareEmbeddedAttemptSessionBoundary } from "./attempt-session-boundary.js"; import { cleanupEmbeddedAttemptSessionPhase } from "./attempt-session-cleanup.js"; @@ -817,16 +812,8 @@ export async function runEmbeddedAttempt( }); const { aggregatePressureEngaged, - contextTokenBudget, - effectivePrompt, hookMessagesForCurrentPrompt, - llmBoundaryPromptForPrecheck, promptForModel, - promptForSession, - promptSubmission, - promptToolResultAggregateMaxChars, - promptToolResultMaxChars, - runtimeContextMessageForCurrentTurn, systemPromptForHook, } = promptContext; prePromptMessageCount = promptContext.prePromptMessageCount; @@ -889,61 +876,26 @@ export async function runEmbeddedAttempt( } } - const imageResult = await prepareEmbeddedAttemptPromptExecution({ - attempt: params, - effectiveFsWorkspaceOnly, - effectiveWorkspace, - prompt: promptSubmission.prompt, - sandbox, - session: activeSession, - sessionLockController, + ({ + contextBudgetStatus, + preflightRecovery, + promptError, + promptErrorSource, skipPromptSubmission, - }); - - const reserveTokens = settingsManager.getCompactionReserveTokens(); - skipPromptSubmission = observeEmbeddedAttemptPrompt({ - attempt: params, - cacheTrace, - contextTokenBudget, - diagnosticTrace, - effectivePrompt, - effectiveTools, - hookAgentId, - hookMessagesForCurrentPrompt, - hookRunner, - imageCount: imageResult.images.length, - isRawModelRun, - llmBoundaryPromptForPrecheck, - promptForModel, - promptSubmissionRuntimeOnly: promptSubmission.runtimeOnly, - reserveTokens, - runTrace, - sessionMessages: activeSession.messages, - skipPromptSubmission, - streamStrategy, - systemPromptForHook, - systemPromptText, - toolSearchCompacted: toolSearch.compacted, - tools, - trajectoryRecorder, - transcriptLeafId, - transport: effectiveAgentTransport, - uncompactedEffectiveTools, - }).skipPromptSubmission; - - const promptPreflight = await prepareEmbeddedAttemptPromptPreflight({ + } = await dispatchEmbeddedAttemptPrompt({ attempt: params, ...(activeContextEngine ? { activeContextEngine } : {}), - contextEngineAssemblySucceeded, - contextEnginePromptAuthority, - contextTokenBudget, - hookMessagesForCurrentPrompt, - includeBoundaryTimestamp, - promptForPrecheck: llmBoundaryPromptForPrecheck, - reserveTokens, - sessionAgentId, - sessionManager: activeSessionManager, - sessionMessageCount: activeSession.messages.length, + activeSession, + promptContext, + getCompactionReserveTokens: () => settingsManager.getCompactionReserveTokens(), + publishState: (state) => { + contextBudgetStatus = state.contextBudgetStatus; + preflightRecovery = state.preflightRecovery; + promptError = state.promptError; + promptErrorSource = state.promptErrorSource; + skipPromptSubmission = state.skipPromptSubmission; + }, + releaseLeasedSteering, state: { contextBudgetStatus, preflightRecovery, @@ -951,31 +903,44 @@ export async function runEmbeddedAttempt( promptErrorSource, skipPromptSubmission, }, - systemPrompt: systemPromptForHook, - ...(boundaryTimezone ? { timezone: boundaryTimezone } : {}), - toolResultMaxChars: promptToolResultMaxChars, - ...(unwindowedContextEngineMessagesForPrecheck - ? { unwindowedContextEngineMessagesForPrecheck } - : {}), - withOwnedSessionWriteLock, - }); - ({ - contextBudgetStatus, - preflightRecovery, - promptError, - promptErrorSource, - skipPromptSubmission, - } = promptPreflight); - - if (!skipPromptSubmission) { - await submitEmbeddedAttemptPrompt({ - attempt: params, - activeSession, + execution: { + effectiveFsWorkspaceOnly, + effectiveWorkspace, + sandbox, + sessionLockController, + }, + observation: { + cacheTrace, + diagnosticTrace, + effectiveTools, + hookAgentId, + hookRunner, + isRawModelRun, + runTrace, + streamStrategy, + systemPromptText, + toolSearchCompacted: toolSearch.compacted, + tools, + trajectoryRecorder, + transcriptLeafId, + transport: effectiveAgentTransport, + uncompactedEffectiveTools, + }, + preflight: { + contextEngineAssemblySucceeded, + contextEnginePromptAuthority, + includeBoundaryTimestamp, + sessionAgentId, + sessionManager: activeSessionManager, + ...(boundaryTimezone ? { timezone: boundaryTimezone } : {}), + ...(unwindowedContextEngineMessagesForPrecheck + ? { unwindowedContextEngineMessagesForPrecheck } + : {}), + withOwnedSessionWriteLock, + }, + submission: { ...(promptBuildAppendContext ? { appendContext: promptBuildAppendContext } : {}), - contextTokenBudget, - images: imageResult.images, ...(leasedSteering ? { leasedSteering } : {}), - modelPrompt: promptForModel, onFinalPromptText: (prompt) => { finalPromptText = prompt; }, @@ -984,22 +949,12 @@ export async function runEmbeddedAttempt( }, ...(promptBuildPrependContext ? { prependContext: promptBuildPrependContext } : {}), promptActiveSession, - ...(runtimeContextMessageForCurrentTurn - ? { runtimeContextMessage: runtimeContextMessageForCurrentTurn } - : {}), - runtimeOnly: promptSubmission.runtimeOnly === true, sessionPromptState, - systemPrompt: systemPromptForHook, - toolResultAggregateMaxChars: promptToolResultAggregateMaxChars, - toolResultMaxChars: promptToolResultMaxChars, toolResultPromptProjectionState, trajectoryRecorder, transcriptLeafId, - transcriptPrompt: promptForSession, - }); - } else { - releaseLeasedSteering(promptError ?? "prompt submission skipped"); - } + }, + })); } catch (err) { const promptErrorOutcome = await handleEmbeddedAttemptPromptError({ activeSession,