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 deleted file mode 100644 index 0086a7fec24e..000000000000 --- a/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.test.ts +++ /dev/null @@ -1,167 +0,0 @@ -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-submit.js", () => ({ - prepareEmbeddedAttemptPromptExecution: hoisted.prepareEmbeddedAttemptPromptExecution, - submitEmbeddedAttemptPrompt: hoisted.submitEmbeddedAttemptPrompt, -})); -vi.mock("./attempt-prompt-support.js", () => ({ - observeEmbeddedAttemptPrompt: hoisted.observeEmbeddedAttemptPrompt, -})); -vi.mock("./attempt-prompt-preflight.js", () => ({ - prepareEmbeddedAttemptPromptPreflight: hoisted.prepareEmbeddedAttemptPromptPreflight, -})); -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" }], - imageFactIndexes: [null], - detectedRefs: [], - failedMediaCount: 0, - 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" }], - imageFactIndexes: [null], - detectedRefs: [], - failedMediaCount: 0, - 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 deleted file mode 100644 index b67c68db6729..000000000000 --- a/src/agents/embedded-agent-runner/run/attempt-prompt-dispatch.ts +++ /dev/null @@ -1,138 +0,0 @@ -/** Runs prompt-local image preparation, observability, preflight, and provider dispatch. */ -import type { prepareEmbeddedAttemptPromptContext } from "./attempt-prompt-build.js"; -import { prepareEmbeddedAttemptPromptPreflight } from "./attempt-prompt-preflight.js"; -import { - prepareEmbeddedAttemptPromptExecution, - submitEmbeddedAttemptPrompt, -} from "./attempt-prompt-submit.js"; -import { observeEmbeddedAttemptPrompt } from "./attempt-prompt-support.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: 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, - 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-prompt-phase.test.ts b/src/agents/embedded-agent-runner/run/attempt-prompt-phase.test.ts index 20b98f175a5c..5a13b3314a8b 100644 --- a/src/agents/embedded-agent-runner/run/attempt-prompt-phase.test.ts +++ b/src/agents/embedded-agent-runner/run/attempt-prompt-phase.test.ts @@ -1,16 +1,20 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ + applyPromptToolsAllow: vi.fn(), beforeAgentRun: vi.fn(), - dispatchPrompt: vi.fn(), handlePromptError: vi.fn(), handleMidTurnPrecheck: vi.fn(), + observePrompt: vi.fn(), prepareGooglePromptCache: vi.fn(), preparePromptAssembly: vi.fn(), preparePromptContext: vi.fn(), + preparePromptExecution: vi.fn(), + preparePromptPreflight: vi.fn(), releasePendingSteering: vi.fn(), removeTrailingPrecheckError: vi.fn(), resolveApiKey: vi.fn(), + submitPrompt: vi.fn(), debug: vi.fn(), warn: vi.fn(), })); @@ -36,35 +40,34 @@ vi.mock("./attempt-prompt-build.js", () => ({ prepareEmbeddedAttemptPromptAssembly: mocks.preparePromptAssembly, prepareEmbeddedAttemptPromptContext: mocks.preparePromptContext, })); -vi.mock("./attempt-prompt-dispatch.js", () => ({ - dispatchEmbeddedAttemptPrompt: mocks.dispatchPrompt, -})); vi.mock("./attempt-prompt-submit.js", () => ({ handleEmbeddedAttemptPromptError: mocks.handlePromptError, + prepareEmbeddedAttemptPromptExecution: mocks.preparePromptExecution, + submitEmbeddedAttemptPrompt: mocks.submitPrompt, })); vi.mock("./attempt-prompt-preflight.js", () => ({ handleEmbeddedAttemptMidTurnPrecheck: mocks.handleMidTurnPrecheck, + prepareEmbeddedAttemptPromptPreflight: mocks.preparePromptPreflight, +})); +vi.mock("./attempt-prompt-support.js", () => ({ + applyPromptBuildToolsAllow: mocks.applyPromptToolsAllow, + observeEmbeddedAttemptPrompt: mocks.observePrompt, })); vi.mock("./attempt-transcript-helpers.js", () => ({ removeTrailingMidTurnPrecheckAssistantError: mocks.removeTrailingPrecheckError, })); import { runEmbeddedAttemptPromptPhase } from "./attempt-prompt-phase.js"; +import type { prepareEmbeddedAttemptPromptPreflight } from "./attempt-prompt-preflight.js"; +import type { submitEmbeddedAttemptPrompt } from "./attempt-prompt-submit.js"; type PromptPhaseInput = Parameters[0]; type PromptPhaseState = ReturnType; type AssemblyCall = { setLeasedSteering: (lease: { leaseId: string; runIds: string[] }) => void; }; -type DispatchCall = { - getCompactionReserveTokens: () => number; - publishState: (state: PromptPhaseState & { skipPromptSubmission: boolean }) => void; - state: PromptPhaseState & { skipPromptSubmission: boolean }; - submission: { - onFinalPromptText: (prompt: string) => void; - onSteeringAcknowledged: () => void; - }; -}; +type PromptPreflightCall = Parameters[0]; +type PromptSubmissionCall = Parameters[0]; type PromptErrorCall = { error: unknown; markYieldAborted: () => void; @@ -128,7 +131,7 @@ function createFixture() { promptSubmission: { prompt: "hello", runtimeOnly: false }, promptToolResultAggregateMaxChars: 2_000, promptToolResultMaxChars: 1_000, - runtimeContextMessageForCurrentTurn: undefined, + runtimeContextMessageForCurrentTurn: { role: "custom", content: "runtime" }, systemPromptForHook: "system", }; }); @@ -141,15 +144,31 @@ function createFixture() { order.push("google-cache"); return undefined; }); - mocks.dispatchPrompt.mockImplementation(async (input: DispatchCall) => { - order.push("dispatch"); - expect(input.getCompactionReserveTokens()).toBe(77); - input.submission.onFinalPromptText("hello"); - input.submission.onSteeringAcknowledged(); - const nextState = { ...input.state, skipPromptSubmission: false }; - input.publishState(nextState); - return nextState; + mocks.preparePromptExecution.mockImplementation(async () => { + order.push("images"); + return { + images: [{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" }], + imageFactIndexes: [null], + detectedRefs: [], + failedMediaCount: 0, + loadedCount: 1, + skippedCount: 0, + }; }); + mocks.observePrompt.mockImplementation(() => { + order.push("observe"); + return { skipPromptSubmission: false }; + }); + mocks.preparePromptPreflight.mockImplementation(async (preflightInput: PromptPreflightCall) => { + order.push("preflight"); + return preflightInput.state; + }); + mocks.submitPrompt.mockImplementation(async (submissionInput: PromptSubmissionCall) => { + order.push("submit"); + submissionInput.onFinalPromptText("hello"); + submissionInput.onSteeringAcknowledged(); + }); + mocks.handlePromptError.mockResolvedValue({}); const activeSession = { messages: [], @@ -231,7 +250,10 @@ function createFixture() { }, lifecycle: { readState: () => state, - writeState: (nextState: PromptPhaseState) => Object.assign(state, nextState), + writeState: (nextState: PromptPhaseState) => { + order.push("publish"); + Object.assign(state, nextState); + }, getPrePromptMessageCount: () => prePromptMessageCount, setPrePromptMessageCount, setCurrentUserTimestampOverride: vi.fn(), @@ -274,19 +296,39 @@ describe("runEmbeddedAttemptPromptPhase", () => { "context", "before-agent-run", "google-cache", - "dispatch", + "images", + "observe", + "publish", + "preflight", + "publish", + "submit", + "publish", "stop-steering", ]); expect(fixture.setPrePromptMessageCount).toHaveBeenCalledWith(2); expect(fixture.setPromptCacheChangesForTurn).toHaveBeenCalledWith([]); expect(fixture.setFinalPromptText).toHaveBeenCalledWith("hello"); - expect(mocks.dispatchPrompt).toHaveBeenCalledWith( + expect(mocks.preparePromptExecution).toHaveBeenCalledWith( expect.objectContaining({ - observation: expect.objectContaining({ transcriptLeafId: "leaf-1" }), - submission: expect.objectContaining({ - leasedSteering: { leaseId: "lease-1", runIds: ["run-1"] }, - transcriptLeafId: "leaf-1", - }), + prompt: "hello", + skipPromptSubmission: false, + }), + ); + expect(mocks.observePrompt).toHaveBeenCalledWith( + expect.objectContaining({ + imageCount: 1, + reserveTokens: 77, + transcriptLeafId: "leaf-1", + }), + ); + expect(mocks.submitPrompt).toHaveBeenCalledWith( + expect.objectContaining({ + images: [expect.objectContaining({ type: "image" })], + leasedSteering: { leaseId: "lease-1", runIds: ["run-1"] }, + modelPrompt: "hello", + runtimeContextMessage: expect.objectContaining({ content: "runtime" }), + transcriptLeafId: "leaf-1", + transcriptPrompt: "hello", }), ); expect(mocks.releasePendingSteering).not.toHaveBeenCalled(); @@ -303,7 +345,13 @@ describe("runEmbeddedAttemptPromptPhase", () => { "assembly", "context", "google-cache", - "dispatch", + "images", + "observe", + "publish", + "preflight", + "publish", + "submit", + "publish", "stop-steering", ]); }); @@ -318,23 +366,18 @@ describe("runEmbeddedAttemptPromptPhase", () => { await runEmbeddedAttemptPromptPhase(fixture.input); - expect(mocks.dispatchPrompt).toHaveBeenCalledWith( - expect.objectContaining({ - state: expect.objectContaining({ - skipPromptSubmission: false, - promptError: null, - promptErrorSource: null, - }), - }), + expect(mocks.preparePromptExecution).toHaveBeenCalledWith( + expect.objectContaining({ skipPromptSubmission: false }), ); + expect(mocks.submitPrompt).toHaveBeenCalledOnce(); }); it("reads yield state after submission fails and publishes abort state before recovery", async () => { const fixture = createFixture(); const submissionError = new Error("submission failed"); const yieldAbortSettled = Promise.resolve(); - mocks.dispatchPrompt.mockImplementation(async () => { - fixture.order.push("dispatch"); + mocks.submitPrompt.mockImplementation(async () => { + fixture.order.push("submit"); fixture.yieldState.yieldDetected = true; fixture.yieldState.yieldAbortSettled = yieldAbortSettled; fixture.yieldState.yieldMessage = "yield context"; @@ -355,7 +398,7 @@ describe("runEmbeddedAttemptPromptPhase", () => { }); expect(fixture.order.slice(-4)).toEqual([ - "dispatch", + "submit", "prompt-error", "yield-aborted", "stop-steering", @@ -365,4 +408,71 @@ describe("runEmbeddedAttemptPromptPhase", () => { expect.objectContaining({ leaseId: "lease-1", runIds: ["run-1"] }), ); }); + + it("releases steering when preflight skips provider submission", async () => { + const fixture = createFixture(); + const promptError = new Error("preflight rejected"); + mocks.observePrompt.mockImplementationOnce(() => { + fixture.order.push("observe"); + return { skipPromptSubmission: true }; + }); + mocks.preparePromptPreflight.mockImplementationOnce( + async (preflightInput: PromptPreflightCall) => { + fixture.order.push("preflight"); + return { + ...preflightInput.state, + promptError, + promptErrorSource: "precheck", + }; + }, + ); + + await runEmbeddedAttemptPromptPhase(fixture.input); + + expect(fixture.state.promptError).toBe(promptError); + expect(mocks.releasePendingSteering).toHaveBeenCalledWith( + expect.objectContaining({ + error: "preflight rejected", + leaseId: "lease-1", + runIds: ["run-1"], + }), + ); + expect(mocks.submitPrompt).not.toHaveBeenCalled(); + }); + + it("publishes preflight state before a submission failure", async () => { + const fixture = createFixture(); + const promptError = new Error("admission warning"); + const submitError = new Error("provider failed"); + mocks.preparePromptPreflight.mockImplementationOnce( + async (preflightInput: PromptPreflightCall) => { + fixture.order.push("preflight"); + return { + ...preflightInput.state, + promptError, + promptErrorSource: "precheck", + }; + }, + ); + mocks.submitPrompt.mockImplementationOnce(async () => { + fixture.order.push("submit"); + expect(fixture.state.promptError).toBe(promptError); + throw submitError; + }); + mocks.handlePromptError.mockImplementationOnce(async (errorInput: PromptErrorCall) => { + fixture.order.push("prompt-error"); + expect(errorInput.error).toBe(submitError); + return {}; + }); + + await runEmbeddedAttemptPromptPhase(fixture.input); + + expect(fixture.order.slice(-5)).toEqual([ + "preflight", + "publish", + "submit", + "prompt-error", + "stop-steering", + ]); + }); }); diff --git a/src/agents/embedded-agent-runner/run/attempt-prompt-phase.ts b/src/agents/embedded-agent-runner/run/attempt-prompt-phase.ts index 246fec1ed22f..720b3632d170 100644 --- a/src/agents/embedded-agent-runner/run/attempt-prompt-phase.ts +++ b/src/agents/embedded-agent-runner/run/attempt-prompt-phase.ts @@ -13,10 +13,19 @@ import { prepareEmbeddedAttemptPromptAssembly, prepareEmbeddedAttemptPromptContext, } from "./attempt-prompt-build.js"; -import { dispatchEmbeddedAttemptPrompt } from "./attempt-prompt-dispatch.js"; -import { handleEmbeddedAttemptMidTurnPrecheck } from "./attempt-prompt-preflight.js"; -import { handleEmbeddedAttemptPromptError } from "./attempt-prompt-submit.js"; -import { applyPromptBuildToolsAllow } from "./attempt-prompt-support.js"; +import { + handleEmbeddedAttemptMidTurnPrecheck, + prepareEmbeddedAttemptPromptPreflight, +} from "./attempt-prompt-preflight.js"; +import { + handleEmbeddedAttemptPromptError, + prepareEmbeddedAttemptPromptExecution, + submitEmbeddedAttemptPrompt, +} from "./attempt-prompt-submit.js"; +import { + applyPromptBuildToolsAllow, + observeEmbeddedAttemptPrompt, +} from "./attempt-prompt-support.js"; import { removeTrailingMidTurnPrecheckAssistantError } from "./attempt-transcript-helpers.js"; import type { MidTurnPrecheckRequest } from "./midturn-precheck.js"; @@ -24,12 +33,15 @@ type PromptAssemblyInput = Parameters>; type PromptContextInput = Parameters[0]; type PromptContextResult = ReturnType; -type PromptDispatchInput = Parameters[0]; type PromptErrorInput = Parameters[0]; +type PromptExecutionInput = Parameters[0]; +type PromptObservationInput = Parameters[0]; +type PromptPreflightInput = Parameters[0]; +type PromptSubmissionInput = Parameters[0]; type BeforeAgentRunOutcome = NonNullable< Awaited> >; -type PromptPhaseState = Omit; +type PromptPhaseState = Omit; type PromptAssemblyPhaseInput = Omit< PromptAssemblyInput, @@ -43,14 +55,44 @@ type PromptContextPhaseInput = Omit< PromptContextInput, "attempt" | "messages" | "prompt" | "replaceSessionMessages" >; -type PromptExecutionPhaseInput = PromptDispatchInput["execution"]; -type PromptObservationPhaseInput = Omit; +type PromptExecutionPhaseInput = Omit< + PromptExecutionInput, + "attempt" | "prompt" | "skipPromptSubmission" +>; +type PromptObservationPhaseInput = Omit< + PromptObservationInput, + | "attempt" + | "contextTokenBudget" + | "effectivePrompt" + | "hookMessagesForCurrentPrompt" + | "imageCount" + | "llmBoundaryPromptForPrecheck" + | "promptForModel" + | "promptSubmissionRuntimeOnly" + | "reserveTokens" + | "sessionMessages" + | "skipPromptSubmission" + | "systemPromptForHook" + | "transcriptLeafId" +>; type PromptToolSurface = ReturnType; -type PromptPreflightPhaseInput = PromptDispatchInput["preflight"] & { - activeContextEngine?: PromptDispatchInput["activeContextEngine"]; +type PromptPreflightPhaseInput = Omit< + PromptPreflightInput, + | "attempt" + | "activeContextEngine" + | "contextTokenBudget" + | "hookMessagesForCurrentPrompt" + | "promptForPrecheck" + | "reserveTokens" + | "sessionMessageCount" + | "state" + | "systemPrompt" + | "toolResultMaxChars" +> & { + activeContextEngine?: PromptPreflightInput["activeContextEngine"]; }; type PromptSubmissionPhaseInput = Pick< - PromptDispatchInput["submission"], + PromptSubmissionInput, | "promptActiveSession" | "sessionPromptState" | "toolResultPromptProjectionState" @@ -113,7 +155,7 @@ export async function runEmbeddedAttemptPromptPhase(input: { const patchState = (patch: Partial) => { input.lifecycle.writeState({ ...input.lifecycle.readState(), ...patch }); }; - const publishDispatchState = (state: PromptDispatchInput["state"]) => { + const publishDispatchState = (state: PromptPreflightInput["state"]) => { const { skipPromptSubmission: nextSkipPromptSubmission, ...phaseState } = state; skipPromptSubmission = nextSkipPromptSubmission; input.lifecycle.writeState(phaseState); @@ -264,21 +306,17 @@ export async function runEmbeddedAttemptPromptPhase(input: { } } - const { activeContextEngine, ...preflight } = input.preflight; - const dispatchState = await dispatchEmbeddedAttemptPrompt({ + const imageResult = await prepareEmbeddedAttemptPromptExecution({ + ...input.execution, attempt, - ...(activeContextEngine ? { activeContextEngine } : {}), - activeSession, - promptContext, - getCompactionReserveTokens: input.getCompactionReserveTokens, - publishState: publishDispatchState, - releaseLeasedSteering, - state: { - ...input.lifecycle.readState(), - skipPromptSubmission, - }, - execution: input.execution, - observation: { + prompt: promptContext.promptSubmission.prompt, + skipPromptSubmission, + }); + + const reserveTokens = input.getCompactionReserveTokens(); + let state: PromptPreflightInput["state"] = { + ...input.lifecycle.readState(), + skipPromptSubmission: observeEmbeddedAttemptPrompt({ ...input.observation, ...(promptToolSurface ? { @@ -287,22 +325,69 @@ export async function runEmbeddedAttemptPromptPhase(input: { uncompactedEffectiveTools: promptToolSurface.uncompactedEffectiveTools, } : {}), + 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, + systemPromptForHook: promptContext.systemPromptForHook, transcriptLeafId, - }, - preflight, - submission: { + }).skipPromptSubmission, + }; + // Publish each admission transition before the next fallible phase so outer cleanup sees it. + publishDispatchState(state); + + const { activeContextEngine, ...preflight } = input.preflight; + state = await prepareEmbeddedAttemptPromptPreflight({ + ...preflight, + attempt, + ...(activeContextEngine ? { activeContextEngine } : {}), + contextTokenBudget: promptContext.contextTokenBudget, + hookMessagesForCurrentPrompt: promptContext.hookMessagesForCurrentPrompt, + promptForPrecheck: promptContext.llmBoundaryPromptForPrecheck, + reserveTokens, + sessionMessageCount: activeSession.messages.length, + state, + systemPrompt: promptContext.systemPromptForHook, + toolResultMaxChars: promptContext.promptToolResultMaxChars, + }); + publishDispatchState(state); + + if (!state.skipPromptSubmission) { + await submitEmbeddedAttemptPrompt({ ...(promptBuildAppendContext ? { appendContext: promptBuildAppendContext } : {}), + attempt, + activeSession, + contextTokenBudget: promptContext.contextTokenBudget, + images: imageResult.images, ...(leasedSteering ? { leasedSteering } : {}), + modelPrompt: promptContext.promptForModel, onFinalPromptText: input.lifecycle.setFinalPromptText, onSteeringAcknowledged: () => { leasedSteering = undefined; }, ...(promptBuildPrependContext ? { prependContext: promptBuildPrependContext } : {}), + ...(promptContext.runtimeContextMessageForCurrentTurn + ? { runtimeContextMessage: promptContext.runtimeContextMessageForCurrentTurn } + : {}), + runtimeOnly: promptContext.promptSubmission.runtimeOnly === true, + systemPrompt: promptContext.systemPromptForHook, + toolResultAggregateMaxChars: promptContext.promptToolResultAggregateMaxChars, + toolResultMaxChars: promptContext.promptToolResultMaxChars, transcriptLeafId, + transcriptPrompt: promptContext.promptForSession, ...input.submission, - }, - }); - publishDispatchState(dispatchState); + }); + } else { + releaseLeasedSteering(state.promptError ?? "prompt submission skipped"); + } + publishDispatchState(state); } catch (error) { const promptErrorOutcome = await handleEmbeddedAttemptPromptError({ activeSession,