refactor(agents): extract attempt prompt dispatch phase

This commit is contained in:
Peter Steinberger
2026-07-14 05:03:21 +01:00
parent c853064a42
commit 632b17618e
3 changed files with 360 additions and 102 deletions

View File

@@ -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<typeof dispatchEmbeddedAttemptPrompt>[0];
type PreflightMockInput = { state: DispatchInput["state"] };
function createInput(overrides: Partial<DispatchInput> = {}): 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);
});
});

View File

@@ -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<typeof prepareEmbeddedAttemptPromptContext>;
type PromptExecutionInput = Parameters<typeof prepareEmbeddedAttemptPromptExecution>[0];
type PromptObservationInput = Parameters<typeof observeEmbeddedAttemptPrompt>[0];
type PromptPreflightInput = Parameters<typeof prepareEmbeddedAttemptPromptPreflight>[0];
type PromptSubmissionInput = Parameters<typeof submitEmbeddedAttemptPrompt>[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<PromptExecutionInput, "attempt" | "prompt" | "session" | "skipPromptSubmission">;
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<PromptDispatchState> {
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;
}

View File

@@ -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,