mirror of
https://github.com/openclaw/openclaw.git
synced 2026-08-12 21:53:00 -06:00
refactor(agents): extract attempt stream runtime setup
This commit is contained in:
Peter Steinberger
@@ -0,0 +1,249 @@
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const mocks = vi.hoisted(() => ({
|
||||
abortable: vi.fn(),
|
||||
bindOwnedSessionTranscriptWrites: vi.fn(),
|
||||
createRunAbort: vi.fn(),
|
||||
flushPendingToolResultsAfterIdle: vi.fn(),
|
||||
installStreamGuards: vi.fn(),
|
||||
prepareHistory: vi.fn(),
|
||||
prepareStream: vi.fn(),
|
||||
prepareTimeout: vi.fn(),
|
||||
withOwnedSessionTranscriptWrites: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("../../../config/sessions/transcript-write-context.js", () => ({
|
||||
bindOwnedSessionTranscriptWrites: mocks.bindOwnedSessionTranscriptWrites,
|
||||
withOwnedSessionTranscriptWrites: mocks.withOwnedSessionTranscriptWrites,
|
||||
}));
|
||||
vi.mock("../wait-for-idle-before-flush.js", () => ({
|
||||
flushPendingToolResultsAfterIdle: mocks.flushPendingToolResultsAfterIdle,
|
||||
}));
|
||||
vi.mock("./abortable.js", () => ({ abortable: mocks.abortable }));
|
||||
vi.mock("./attempt-abort.js", () => ({
|
||||
createEmbeddedAttemptRunAbort: mocks.createRunAbort,
|
||||
}));
|
||||
vi.mock("./attempt-history-prepare.js", () => ({
|
||||
prepareEmbeddedAttemptHistory: mocks.prepareHistory,
|
||||
}));
|
||||
vi.mock("./attempt-stream-prepare.js", () => ({
|
||||
prepareEmbeddedAttemptStream: mocks.prepareStream,
|
||||
}));
|
||||
vi.mock("./attempt-stream.js", () => ({
|
||||
installEmbeddedAttemptStreamGuards: mocks.installStreamGuards,
|
||||
}));
|
||||
vi.mock("./attempt-timeout-prepare.js", () => ({
|
||||
prepareEmbeddedAttemptTimeout: mocks.prepareTimeout,
|
||||
}));
|
||||
|
||||
import { prepareEmbeddedAttemptStreamRuntime } from "./attempt-stream-runtime-prepare.js";
|
||||
|
||||
type StreamRuntimeInput = Parameters<typeof prepareEmbeddedAttemptStreamRuntime>[0];
|
||||
|
||||
function createFixture(options: { aborted?: boolean } = {}) {
|
||||
const order: string[] = [];
|
||||
const abortController = new AbortController();
|
||||
if (options.aborted) {
|
||||
abortController.abort(new Error("already aborted"));
|
||||
}
|
||||
const runAbort = vi.fn();
|
||||
const toolSearchCatalogExecutor = vi.fn();
|
||||
const subscription = {
|
||||
isCompacting: vi.fn(() => false),
|
||||
};
|
||||
const queueHandle = { kind: "embedded", runId: "run-1" };
|
||||
const streamResult = {
|
||||
subscription,
|
||||
queueHandle,
|
||||
toolSearchCatalogExecutor,
|
||||
getBeforeAgentFinalizeRevisionReason: vi.fn(),
|
||||
stopAcceptingSteerMessages: vi.fn(),
|
||||
};
|
||||
const timeoutResult = {
|
||||
getRunAbortDeadlineAtMs: vi.fn(() => 123),
|
||||
clearTimers: vi.fn(),
|
||||
removeAbortSignalListener: vi.fn(),
|
||||
};
|
||||
const activeSession = {
|
||||
agent: { streamFn: vi.fn() },
|
||||
dispose: vi.fn(),
|
||||
isCompacting: false,
|
||||
messages: [],
|
||||
prompt: vi.fn(async () => undefined),
|
||||
};
|
||||
const sessionManager = {};
|
||||
const externalAbortController = {
|
||||
setRunAbort: vi.fn(() => order.push("set-run-abort")),
|
||||
setCompactionState: vi.fn(() => order.push("set-compaction-state")),
|
||||
};
|
||||
const markIdleTimedOut = vi.fn();
|
||||
const markStreamReady = vi.fn(() => order.push("stream-ready"));
|
||||
const setToolSearchCatalogExecutor = vi.fn(() => order.push("set-catalog"));
|
||||
const trackPromptSettlePromise = vi.fn((promise: Promise<void>) => promise);
|
||||
|
||||
mocks.abortable.mockImplementation((_signal, promise) => promise);
|
||||
mocks.bindOwnedSessionTranscriptWrites.mockImplementation((_context, operation) => operation);
|
||||
mocks.withOwnedSessionTranscriptWrites.mockImplementation(
|
||||
async (_context, operation) => await operation(),
|
||||
);
|
||||
mocks.installStreamGuards.mockImplementation(() => {
|
||||
order.push("guards");
|
||||
return {
|
||||
cacheObservabilityEnabled: true,
|
||||
promptCacheToolNames: new Set(["read"]),
|
||||
};
|
||||
});
|
||||
mocks.prepareHistory.mockImplementation(async () => {
|
||||
order.push("history");
|
||||
return {
|
||||
contextEnginePromptAuthority: "assembled",
|
||||
contextEngineAssemblySucceeded: true,
|
||||
};
|
||||
});
|
||||
mocks.createRunAbort.mockImplementation(() => {
|
||||
order.push("abort");
|
||||
return runAbort;
|
||||
});
|
||||
mocks.prepareStream.mockImplementation(() => {
|
||||
order.push("stream");
|
||||
return streamResult;
|
||||
});
|
||||
mocks.prepareTimeout.mockImplementation(() => {
|
||||
order.push("timeout");
|
||||
return timeoutResult;
|
||||
});
|
||||
|
||||
const input = {
|
||||
attempt: {
|
||||
abortSignal: abortController.signal,
|
||||
onBlockReply: vi.fn(),
|
||||
onBlockReplyFlush: vi.fn(),
|
||||
runId: "run-1",
|
||||
sessionId: "session-1",
|
||||
timeoutMs: 30_000,
|
||||
},
|
||||
activeSession,
|
||||
sessionManager,
|
||||
sessionLockController: {},
|
||||
ownedTranscriptWriteContext: {},
|
||||
runAbortController: new AbortController(),
|
||||
externalAbortController,
|
||||
abortActiveSession: vi.fn(async () => undefined),
|
||||
abortState: {},
|
||||
trackPromptSettlePromise,
|
||||
compactionTimeoutMs: 1_000,
|
||||
guards: {},
|
||||
history: {},
|
||||
stream: {},
|
||||
lifecycle: {
|
||||
isYieldDetected: () => false,
|
||||
markRejectedThinkingReplayRepaired: vi.fn(),
|
||||
markStreamReady,
|
||||
markIdleTimedOut,
|
||||
markExternalAbort: vi.fn(),
|
||||
markTimedOutDuringCompaction: vi.fn(),
|
||||
markTimedOutByRunBudget: vi.fn(),
|
||||
readRunState: () => ({
|
||||
aborted: false,
|
||||
promptError: null,
|
||||
timedOut: false,
|
||||
yieldDetected: false,
|
||||
}),
|
||||
setToolSearchCatalogExecutor,
|
||||
},
|
||||
} as unknown as StreamRuntimeInput;
|
||||
|
||||
return {
|
||||
activeSession,
|
||||
externalAbortController,
|
||||
input,
|
||||
markIdleTimedOut,
|
||||
order,
|
||||
runAbort,
|
||||
sessionManager,
|
||||
streamResult,
|
||||
subscription,
|
||||
timeoutResult,
|
||||
toolSearchCatalogExecutor,
|
||||
trackPromptSettlePromise,
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe("prepareEmbeddedAttemptStreamRuntime", () => {
|
||||
it("prepares guarded history, abort handling, stream subscription, and timeout in order", async () => {
|
||||
const fixture = createFixture();
|
||||
|
||||
const result = await prepareEmbeddedAttemptStreamRuntime(fixture.input);
|
||||
|
||||
expect(fixture.order).toEqual([
|
||||
"guards",
|
||||
"stream-ready",
|
||||
"history",
|
||||
"abort",
|
||||
"set-run-abort",
|
||||
"stream",
|
||||
"set-catalog",
|
||||
"set-compaction-state",
|
||||
"timeout",
|
||||
]);
|
||||
expect(result).toEqual(
|
||||
expect.objectContaining({
|
||||
cache: {
|
||||
observabilityEnabled: true,
|
||||
promptToolNames: new Set(["read"]),
|
||||
},
|
||||
history: expect.objectContaining({ contextEngineAssemblySucceeded: true }),
|
||||
isProbeSession: false,
|
||||
stream: fixture.streamResult,
|
||||
timeout: fixture.timeoutResult,
|
||||
}),
|
||||
);
|
||||
expect(fixture.input.lifecycle.setToolSearchCatalogExecutor).toHaveBeenCalledWith(
|
||||
fixture.toolSearchCatalogExecutor,
|
||||
);
|
||||
expect(fixture.externalAbortController.setCompactionState).toHaveBeenCalledWith({
|
||||
isPendingOrRetrying: fixture.subscription.isCompacting,
|
||||
isInFlight: expect.any(Function),
|
||||
});
|
||||
expect(mocks.prepareTimeout).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
abortRun: fixture.runAbort,
|
||||
compactionState: fixture.subscription,
|
||||
}),
|
||||
);
|
||||
|
||||
const guardInput = mocks.installStreamGuards.mock.calls[0]?.[0];
|
||||
const idleError = new Error("idle timeout");
|
||||
guardInput.onIdleTimeout(idleError);
|
||||
expect(fixture.markIdleTimedOut).toHaveBeenCalledOnce();
|
||||
expect(fixture.runAbort).toHaveBeenCalledWith(true, idleError);
|
||||
|
||||
await result.promptActiveSession("hello");
|
||||
expect(fixture.activeSession.prompt).toHaveBeenCalledWith("hello", undefined);
|
||||
expect(fixture.trackPromptSettlePromise).toHaveBeenCalledOnce();
|
||||
expect(mocks.withOwnedSessionTranscriptWrites).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("flushes pending tool results and disposes the session when history preparation fails", async () => {
|
||||
const fixture = createFixture({ aborted: true });
|
||||
const failure = new Error("history failed");
|
||||
mocks.prepareHistory.mockRejectedValueOnce(failure);
|
||||
mocks.flushPendingToolResultsAfterIdle.mockResolvedValue(undefined);
|
||||
|
||||
await expect(prepareEmbeddedAttemptStreamRuntime(fixture.input)).rejects.toBe(failure);
|
||||
|
||||
expect(mocks.flushPendingToolResultsAfterIdle).toHaveBeenCalledWith({
|
||||
agent: fixture.activeSession.agent,
|
||||
sessionManager: fixture.sessionManager,
|
||||
timeoutMs: 0,
|
||||
});
|
||||
expect(fixture.activeSession.dispose).toHaveBeenCalledOnce();
|
||||
expect(mocks.createRunAbort).not.toHaveBeenCalled();
|
||||
expect(mocks.prepareStream).not.toHaveBeenCalled();
|
||||
expect(mocks.prepareTimeout).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,192 @@
|
||||
/** Prepares guarded history, abort handling, stream subscription, and run deadlines. */
|
||||
import {
|
||||
bindOwnedSessionTranscriptWrites,
|
||||
withOwnedSessionTranscriptWrites,
|
||||
} from "../../../config/sessions/transcript-write-context.js";
|
||||
import { log } from "../logger.js";
|
||||
import type { EmbeddedAgentQueueHandle } from "../runs.js";
|
||||
import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js";
|
||||
import { abortable as abortableWithSignal } from "./abortable.js";
|
||||
import {
|
||||
type createEmbeddedAttemptExternalAbortController,
|
||||
createEmbeddedAttemptRunAbort,
|
||||
} from "./attempt-abort.js";
|
||||
import { prepareEmbeddedAttemptHistory } from "./attempt-history-prepare.js";
|
||||
import { prepareEmbeddedAttemptStream } from "./attempt-stream-prepare.js";
|
||||
import { installEmbeddedAttemptStreamGuards } from "./attempt-stream.js";
|
||||
import { prepareEmbeddedAttemptTimeout } from "./attempt-timeout-prepare.js";
|
||||
import type { EmbeddedRunAttemptParams } from "./types.js";
|
||||
|
||||
type StreamGuardInput = Parameters<typeof installEmbeddedAttemptStreamGuards>[0];
|
||||
type HistoryInput = Parameters<typeof prepareEmbeddedAttemptHistory>[0];
|
||||
type StreamInput = Parameters<typeof prepareEmbeddedAttemptStream>[0];
|
||||
type ToolResultFlushInput = Parameters<typeof flushPendingToolResultsAfterIdle>[0];
|
||||
type ExternalAbortController = Pick<
|
||||
ReturnType<typeof createEmbeddedAttemptExternalAbortController>,
|
||||
"setCompactionState" | "setRunAbort"
|
||||
>;
|
||||
type StreamGuardPhaseInput = Omit<
|
||||
StreamGuardInput,
|
||||
| "abortSignal"
|
||||
| "attempt"
|
||||
| "isYieldDetected"
|
||||
| "onIdleTimeout"
|
||||
| "onRejectedThinkingReplayRepaired"
|
||||
| "session"
|
||||
| "sessionLockController"
|
||||
| "sessionManager"
|
||||
>;
|
||||
type HistoryPhaseInput = Omit<HistoryInput, "activeSession" | "attempt" | "sessionManager">;
|
||||
type StreamPhaseInput = Omit<
|
||||
StreamInput,
|
||||
| "abortRun"
|
||||
| "activeSession"
|
||||
| "attempt"
|
||||
| "getRunState"
|
||||
| "markExternalAbort"
|
||||
| "onBlockReply"
|
||||
| "onBlockReplyFlush"
|
||||
| "runAbortController"
|
||||
>;
|
||||
|
||||
export async function prepareEmbeddedAttemptStreamRuntime(input: {
|
||||
attempt: EmbeddedRunAttemptParams;
|
||||
activeSession: StreamInput["activeSession"];
|
||||
sessionManager: HistoryInput["sessionManager"] &
|
||||
NonNullable<ToolResultFlushInput["sessionManager"]>;
|
||||
sessionLockController: StreamGuardInput["sessionLockController"];
|
||||
ownedTranscriptWriteContext: Parameters<typeof withOwnedSessionTranscriptWrites>[0];
|
||||
runAbortController: AbortController;
|
||||
externalAbortController: ExternalAbortController;
|
||||
abortActiveSession: Parameters<typeof createEmbeddedAttemptRunAbort>[0]["abortActiveSession"];
|
||||
abortState: Parameters<typeof createEmbeddedAttemptRunAbort>[0]["state"];
|
||||
trackPromptSettlePromise: (promise: Promise<void>) => Promise<void>;
|
||||
compactionTimeoutMs: number;
|
||||
guards: StreamGuardPhaseInput;
|
||||
history: HistoryPhaseInput;
|
||||
stream: StreamPhaseInput;
|
||||
lifecycle: {
|
||||
isYieldDetected: StreamGuardInput["isYieldDetected"];
|
||||
markRejectedThinkingReplayRepaired: () => void;
|
||||
markStreamReady: () => void;
|
||||
markIdleTimedOut: () => void;
|
||||
markExternalAbort: () => void;
|
||||
markTimedOutDuringCompaction: () => void;
|
||||
markTimedOutByRunBudget: () => void;
|
||||
readRunState: StreamInput["getRunState"];
|
||||
setToolSearchCatalogExecutor: (
|
||||
executor: ReturnType<typeof prepareEmbeddedAttemptStream>["toolSearchCatalogExecutor"],
|
||||
) => void;
|
||||
};
|
||||
}) {
|
||||
const { activeSession, attempt, sessionManager } = input;
|
||||
const idleTimeoutTriggerRef: { current?: (error: Error) => void } = {};
|
||||
const { cacheObservabilityEnabled, promptCacheToolNames } = installEmbeddedAttemptStreamGuards({
|
||||
...input.guards,
|
||||
attempt,
|
||||
session: activeSession,
|
||||
sessionManager,
|
||||
sessionLockController: input.sessionLockController,
|
||||
isYieldDetected: input.lifecycle.isYieldDetected,
|
||||
onRejectedThinkingReplayRepaired: input.lifecycle.markRejectedThinkingReplayRepaired,
|
||||
onIdleTimeout: (error) => idleTimeoutTriggerRef.current?.(error),
|
||||
abortSignal: input.runAbortController.signal,
|
||||
});
|
||||
input.lifecycle.markStreamReady();
|
||||
|
||||
let preparedHistory: Awaited<ReturnType<typeof prepareEmbeddedAttemptHistory>>;
|
||||
try {
|
||||
preparedHistory = await prepareEmbeddedAttemptHistory({
|
||||
...input.history,
|
||||
attempt,
|
||||
activeSession,
|
||||
sessionManager,
|
||||
});
|
||||
} catch (error) {
|
||||
await flushPendingToolResultsAfterIdle({
|
||||
agent: activeSession.agent,
|
||||
sessionManager,
|
||||
// An already-aborted setup must dispose immediately without orphaning tool calls.
|
||||
...(attempt.abortSignal?.aborted ? { timeoutMs: 0 } : {}),
|
||||
});
|
||||
activeSession.dispose();
|
||||
throw error;
|
||||
}
|
||||
|
||||
const isProbeSession = attempt.sessionId?.startsWith("probe-") ?? false;
|
||||
const queueHandleRef: { current?: EmbeddedAgentQueueHandle } = {};
|
||||
const abortRun = createEmbeddedAttemptRunAbort({
|
||||
abortActiveSession: input.abortActiveSession,
|
||||
activeSession,
|
||||
attempt,
|
||||
getQueueHandle: () => queueHandleRef.current,
|
||||
isProbeSession,
|
||||
log,
|
||||
runAbortController: input.runAbortController,
|
||||
sessionLockController: input.sessionLockController,
|
||||
state: input.abortState,
|
||||
});
|
||||
input.externalAbortController.setRunAbort(abortRun);
|
||||
idleTimeoutTriggerRef.current = (error) => {
|
||||
input.lifecycle.markIdleTimedOut();
|
||||
abortRun(true, error);
|
||||
};
|
||||
const abortable = <T>(promise: Promise<T>): Promise<T> =>
|
||||
abortableWithSignal(input.runAbortController.signal, promise);
|
||||
const promptActiveSession = (
|
||||
prompt: string,
|
||||
options?: Parameters<typeof activeSession.prompt>[1],
|
||||
): Promise<void> =>
|
||||
withOwnedSessionTranscriptWrites(input.ownedTranscriptWriteContext, async () =>
|
||||
abortable(input.trackPromptSettlePromise(activeSession.prompt(prompt, options))),
|
||||
);
|
||||
const onBlockReply = attempt.onBlockReply
|
||||
? bindOwnedSessionTranscriptWrites(input.ownedTranscriptWriteContext, attempt.onBlockReply)
|
||||
: undefined;
|
||||
const onBlockReplyFlush = attempt.onBlockReplyFlush
|
||||
? bindOwnedSessionTranscriptWrites(input.ownedTranscriptWriteContext, attempt.onBlockReplyFlush)
|
||||
: undefined;
|
||||
const preparedStream = prepareEmbeddedAttemptStream({
|
||||
...input.stream,
|
||||
attempt,
|
||||
activeSession,
|
||||
runAbortController: input.runAbortController,
|
||||
abortRun,
|
||||
markExternalAbort: input.lifecycle.markExternalAbort,
|
||||
getRunState: input.lifecycle.readRunState,
|
||||
onBlockReply,
|
||||
onBlockReplyFlush,
|
||||
});
|
||||
input.lifecycle.setToolSearchCatalogExecutor(preparedStream.toolSearchCatalogExecutor);
|
||||
input.externalAbortController.setCompactionState({
|
||||
isPendingOrRetrying: preparedStream.subscription.isCompacting,
|
||||
isInFlight: () => activeSession.isCompacting,
|
||||
});
|
||||
queueHandleRef.current = preparedStream.queueHandle;
|
||||
|
||||
const attemptTimeout = prepareEmbeddedAttemptTimeout({
|
||||
attempt,
|
||||
activeSession,
|
||||
compactionState: preparedStream.subscription,
|
||||
compactionTimeoutMs: input.compactionTimeoutMs,
|
||||
isProbeSession,
|
||||
abortRun,
|
||||
markExternalAbort: input.lifecycle.markExternalAbort,
|
||||
markTimedOutDuringCompaction: input.lifecycle.markTimedOutDuringCompaction,
|
||||
markTimedOutByRunBudget: input.lifecycle.markTimedOutByRunBudget,
|
||||
});
|
||||
|
||||
return {
|
||||
abortable,
|
||||
cache: {
|
||||
observabilityEnabled: cacheObservabilityEnabled,
|
||||
promptToolNames: promptCacheToolNames,
|
||||
},
|
||||
history: preparedHistory,
|
||||
isProbeSession,
|
||||
onBlockReplyFlush,
|
||||
promptActiveSession,
|
||||
stream: preparedStream,
|
||||
timeout: attemptTimeout,
|
||||
};
|
||||
}
|
||||
@@ -1,10 +1,6 @@
|
||||
/**
|
||||
* Orchestrates one embedded-agent attempt from prompt setup through stream result.
|
||||
*/
|
||||
import {
|
||||
bindOwnedSessionTranscriptWrites,
|
||||
withOwnedSessionTranscriptWrites,
|
||||
} from "../../../config/sessions/transcript-write-context.js";
|
||||
import {
|
||||
assertContextEngineHostSupport,
|
||||
OPENCLAW_EMBEDDED_CONTEXT_ENGINE_HOST,
|
||||
@@ -27,19 +23,15 @@ import {
|
||||
import type { NormalizedUsage } from "../../usage.js";
|
||||
import { log } from "../logger.js";
|
||||
import type { PromptCacheBreak, PromptCacheChange } from "../prompt-cache-observability.js";
|
||||
import { clearActiveEmbeddedRun, type EmbeddedAgentQueueHandle } from "../runs.js";
|
||||
import { clearActiveEmbeddedRun } from "../runs.js";
|
||||
import { getEmbeddedSessionPromptState } from "../session-prompt-state.js";
|
||||
import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js";
|
||||
import { abortable as abortableWithSignal } from "./abortable.js";
|
||||
import {
|
||||
createEmbeddedAttemptExternalAbortController,
|
||||
createEmbeddedAttemptRunAbort,
|
||||
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 { prepareEmbeddedAttemptHistory } from "./attempt-history-prepare.js";
|
||||
import { runEmbeddedAttemptPromptPhase } from "./attempt-prompt-phase.js";
|
||||
import { completeEmbeddedAttemptResult } from "./attempt-result.js";
|
||||
import { prepareEmbeddedAttemptSessionBoundary } from "./attempt-session-boundary.js";
|
||||
@@ -56,11 +48,9 @@ import {
|
||||
type EmitDiagnosticRunCompleted,
|
||||
} from "./attempt-startup.js";
|
||||
import { finalizeEmbeddedAttemptStreamPhase } from "./attempt-stream-finalize.js";
|
||||
import { prepareEmbeddedAttemptStream } from "./attempt-stream-prepare.js";
|
||||
import { prepareEmbeddedAttemptStreamRuntime } from "./attempt-stream-runtime-prepare.js";
|
||||
import { prepareEmbeddedAttemptTransport } from "./attempt-stream-transport.js";
|
||||
import { installEmbeddedAttemptStreamGuards } from "./attempt-stream.js";
|
||||
import { prepareEmbeddedAttemptSystemPrompt } from "./attempt-system-prompt-prepare.js";
|
||||
import { prepareEmbeddedAttemptTimeout } from "./attempt-timeout-prepare.js";
|
||||
import { prepareEmbeddedAttemptToolBase } from "./attempt-tool-base-prepare.js";
|
||||
import { prepareEmbeddedAttemptToolCatalog } from "./attempt-tool-catalog.js";
|
||||
import { prepareEmbeddedAttemptTrajectory } from "./attempt-trajectory.js";
|
||||
@@ -522,42 +512,36 @@ export async function runEmbeddedAttempt(
|
||||
sandbox,
|
||||
codeModeControlsEnabled: codeModeControlsEnabledForRun,
|
||||
});
|
||||
const { cacheObservabilityEnabled, promptCacheToolNames } =
|
||||
installEmbeddedAttemptStreamGuards({
|
||||
attempt: params,
|
||||
session: activeSession,
|
||||
let yieldAborted = false;
|
||||
const hookAgentId = sessionAgentId;
|
||||
const preparedStreamRuntime = await prepareEmbeddedAttemptStreamRuntime({
|
||||
attempt: params,
|
||||
activeSession,
|
||||
sessionManager,
|
||||
sessionLockController,
|
||||
ownedTranscriptWriteContext,
|
||||
runAbortController,
|
||||
externalAbortController,
|
||||
abortActiveSession,
|
||||
abortState,
|
||||
trackPromptSettlePromise,
|
||||
compactionTimeoutMs,
|
||||
guards: {
|
||||
sessionAgentId,
|
||||
cacheTrace,
|
||||
allCustomTools,
|
||||
systemPromptText,
|
||||
transcriptPolicy,
|
||||
sessionManager,
|
||||
sessionLockController,
|
||||
isOpenAIResponsesApi,
|
||||
replayAllowedToolNames,
|
||||
liveAllowedToolNames,
|
||||
isYieldDetected: () => yieldDetected,
|
||||
clientToolLoopDetection,
|
||||
anthropicPayloadLogger,
|
||||
onRejectedThinkingReplayRepaired: () => {
|
||||
repairedRejectedThinkingReplay = true;
|
||||
},
|
||||
onIdleTimeout: (error) => idleTimeoutTrigger?.(error),
|
||||
effectiveAgentTransport,
|
||||
providerTextTransforms,
|
||||
abortSignal: runAbortController.signal,
|
||||
runTrace,
|
||||
});
|
||||
prepStages.mark("stream-setup");
|
||||
emitPrepStageSummary("stream-ready");
|
||||
let promptCacheChangesForTurn: PromptCacheChange[] | null = null;
|
||||
|
||||
let preparedHistory: Awaited<ReturnType<typeof prepareEmbeddedAttemptHistory>>;
|
||||
try {
|
||||
preparedHistory = await prepareEmbeddedAttemptHistory({
|
||||
attempt: params,
|
||||
activeSession,
|
||||
sessionManager,
|
||||
},
|
||||
history: {
|
||||
...(activeContextEngine ? { activeContextEngine } : {}),
|
||||
cacheTrace,
|
||||
capabilityToolNames,
|
||||
@@ -571,90 +555,66 @@ export async function runEmbeddedAttempt(
|
||||
systemPromptText,
|
||||
transcriptPolicy,
|
||||
setActiveSessionSystemPrompt,
|
||||
});
|
||||
} catch (err) {
|
||||
await flushPendingToolResultsAfterIdle({
|
||||
agent: activeSession?.agent,
|
||||
sessionManager,
|
||||
// PERF: If the run was aborted during the setup,
|
||||
// skip the idle wait and flush pending results synchronously so we can
|
||||
// immediately dispose the session without orphaning tool calls.
|
||||
...(params.abortSignal?.aborted ? { timeoutMs: 0 } : {}),
|
||||
});
|
||||
activeSession.dispose();
|
||||
throw err;
|
||||
}
|
||||
const {
|
||||
contextEnginePromptAuthority,
|
||||
contextEngineAssemblySucceeded,
|
||||
unwindowedContextEngineMessagesForPrecheck,
|
||||
} = preparedHistory;
|
||||
|
||||
let yieldAborted = false;
|
||||
const isProbeSession = params.sessionId?.startsWith("probe-") ?? false;
|
||||
const queueHandleRef: { current?: EmbeddedAgentQueueHandle } = {};
|
||||
const abortRun = createEmbeddedAttemptRunAbort({
|
||||
abortActiveSession,
|
||||
activeSession,
|
||||
attempt: params,
|
||||
getQueueHandle: () => queueHandleRef.current,
|
||||
isProbeSession,
|
||||
log,
|
||||
runAbortController,
|
||||
sessionLockController,
|
||||
state: abortState,
|
||||
});
|
||||
externalAbortController.setRunAbort(abortRun);
|
||||
const idleTimeoutTrigger: ((error: Error) => void) | undefined = (error) => {
|
||||
idleTimedOut = true;
|
||||
abortRun(true, error);
|
||||
};
|
||||
const abortable = <T>(promise: Promise<T>): Promise<T> =>
|
||||
abortableWithSignal(runAbortController.signal, promise);
|
||||
const promptActiveSession = (
|
||||
prompt: string,
|
||||
options?: Parameters<typeof activeSession.prompt>[1],
|
||||
): Promise<void> =>
|
||||
withOwnedSessionTranscriptWrites(ownedTranscriptWriteContext, async () =>
|
||||
abortable(trackPromptSettlePromise(activeSession.prompt(prompt, options))),
|
||||
);
|
||||
// Hook runner was already obtained earlier before tool creation.
|
||||
const hookAgentId = sessionAgentId;
|
||||
const onBlockReply = params.onBlockReply
|
||||
? bindOwnedSessionTranscriptWrites(ownedTranscriptWriteContext, params.onBlockReply)
|
||||
: undefined;
|
||||
const onBlockReplyFlush = params.onBlockReplyFlush
|
||||
? bindOwnedSessionTranscriptWrites(ownedTranscriptWriteContext, params.onBlockReplyFlush)
|
||||
: undefined;
|
||||
const preparedStream = prepareEmbeddedAttemptStream({
|
||||
attempt: params,
|
||||
activeSession,
|
||||
runtimeChannel,
|
||||
hookRunner,
|
||||
hookAgentId,
|
||||
diagnosticTrace,
|
||||
clientToolCallSlots,
|
||||
toolSearchTargetTranscriptProjections,
|
||||
isReplaySafeTool: (tool) => replaySafeTools.has(tool as never),
|
||||
runAbortController,
|
||||
abortRun,
|
||||
markExternalAbort: () => {
|
||||
externalAbort = true;
|
||||
},
|
||||
getRunState: () => ({
|
||||
aborted,
|
||||
promptError,
|
||||
timedOut,
|
||||
yieldDetected,
|
||||
}),
|
||||
hasDeliveredSourceReply,
|
||||
markSourceReplyDelivered,
|
||||
onBlockReply,
|
||||
onBlockReplyFlush,
|
||||
sandboxSessionKey,
|
||||
builtinToolNames,
|
||||
replaySafeToolNames,
|
||||
stream: {
|
||||
runtimeChannel,
|
||||
hookRunner,
|
||||
hookAgentId,
|
||||
diagnosticTrace,
|
||||
clientToolCallSlots,
|
||||
toolSearchTargetTranscriptProjections,
|
||||
isReplaySafeTool: (tool) => replaySafeTools.has(tool as never),
|
||||
hasDeliveredSourceReply,
|
||||
markSourceReplyDelivered,
|
||||
sandboxSessionKey,
|
||||
builtinToolNames,
|
||||
replaySafeToolNames,
|
||||
},
|
||||
lifecycle: {
|
||||
isYieldDetected: () => yieldDetected,
|
||||
markRejectedThinkingReplayRepaired: () => {
|
||||
repairedRejectedThinkingReplay = true;
|
||||
},
|
||||
markStreamReady: () => {
|
||||
prepStages.mark("stream-setup");
|
||||
emitPrepStageSummary("stream-ready");
|
||||
},
|
||||
markIdleTimedOut: () => {
|
||||
idleTimedOut = true;
|
||||
},
|
||||
markExternalAbort: () => {
|
||||
externalAbort = true;
|
||||
},
|
||||
markTimedOutDuringCompaction: () => {
|
||||
timedOutDuringCompaction = true;
|
||||
},
|
||||
markTimedOutByRunBudget: () => {
|
||||
timedOutByRunBudget = true;
|
||||
},
|
||||
readRunState: () => ({ aborted, promptError, timedOut, yieldDetected }),
|
||||
setToolSearchCatalogExecutor: (executor) => {
|
||||
toolSearchCatalogExecutor = executor;
|
||||
},
|
||||
},
|
||||
});
|
||||
const {
|
||||
abortable,
|
||||
cache: {
|
||||
observabilityEnabled: cacheObservabilityEnabled,
|
||||
promptToolNames: promptCacheToolNames,
|
||||
},
|
||||
history: {
|
||||
contextEnginePromptAuthority,
|
||||
contextEngineAssemblySucceeded,
|
||||
unwindowedContextEngineMessagesForPrecheck,
|
||||
},
|
||||
isProbeSession,
|
||||
onBlockReplyFlush,
|
||||
promptActiveSession,
|
||||
stream: preparedStream,
|
||||
timeout: attemptTimeout,
|
||||
} = preparedStreamRuntime;
|
||||
let promptCacheChangesForTurn: PromptCacheChange[] | null = null;
|
||||
const {
|
||||
subscription,
|
||||
queueHandle,
|
||||
@@ -662,36 +622,12 @@ export async function runEmbeddedAttempt(
|
||||
getBeforeAgentFinalizeRevisionReason,
|
||||
} = preparedStream;
|
||||
const { unsubscribe, waitForPendingEvents } = subscription;
|
||||
toolSearchCatalogExecutor = preparedStream.toolSearchCatalogExecutor;
|
||||
externalAbortController.setCompactionState({
|
||||
isPendingOrRetrying: subscription.isCompacting,
|
||||
isInFlight: () => activeSession.isCompacting,
|
||||
});
|
||||
let lastAssistant: AssistantMessage | undefined;
|
||||
let currentAttemptAssistant: EmbeddedRunAttemptResult["currentAttemptAssistant"];
|
||||
let attemptUsage: NormalizedUsage | undefined;
|
||||
let cacheBreak: PromptCacheBreak | null = null;
|
||||
let contextBudgetStatus: EmbeddedRunAttemptResult["contextBudgetStatus"];
|
||||
let finalPromptText: string | undefined;
|
||||
queueHandleRef.current = queueHandle;
|
||||
|
||||
const attemptTimeout = prepareEmbeddedAttemptTimeout({
|
||||
attempt: params,
|
||||
activeSession,
|
||||
compactionState: subscription,
|
||||
compactionTimeoutMs,
|
||||
isProbeSession,
|
||||
abortRun,
|
||||
markExternalAbort: () => {
|
||||
externalAbort = true;
|
||||
},
|
||||
markTimedOutDuringCompaction: () => {
|
||||
timedOutDuringCompaction = true;
|
||||
},
|
||||
markTimedOutByRunBudget: () => {
|
||||
timedOutByRunBudget = true;
|
||||
},
|
||||
});
|
||||
const {
|
||||
getRunAbortDeadlineAtMs,
|
||||
clearTimers: clearAttemptTimeoutTimers,
|
||||
|
||||
Reference in New Issue
Block a user