From 73838101df3bf65947aec6e753f664d174490a3b Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 14 Jul 2026 05:50:58 +0100 Subject: [PATCH] refactor(agents): extract attempt stream runtime setup --- .../attempt-stream-runtime-prepare.test.ts | 249 ++++++++++++++++++ .../run/attempt-stream-runtime-prepare.ts | 192 ++++++++++++++ .../embedded-agent-runner/run/attempt.ts | 218 ++++++--------- 3 files changed, 518 insertions(+), 141 deletions(-) create mode 100644 src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.test.ts create mode 100644 src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.ts diff --git a/src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.test.ts b/src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.test.ts new file mode 100644 index 000000000000..80eefcd2c332 --- /dev/null +++ b/src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.test.ts @@ -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[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) => 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(); + }); +}); diff --git a/src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.ts b/src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.ts new file mode 100644 index 000000000000..43b4c9fa6475 --- /dev/null +++ b/src/agents/embedded-agent-runner/run/attempt-stream-runtime-prepare.ts @@ -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[0]; +type HistoryInput = Parameters[0]; +type StreamInput = Parameters[0]; +type ToolResultFlushInput = Parameters[0]; +type ExternalAbortController = Pick< + ReturnType, + "setCompactionState" | "setRunAbort" +>; +type StreamGuardPhaseInput = Omit< + StreamGuardInput, + | "abortSignal" + | "attempt" + | "isYieldDetected" + | "onIdleTimeout" + | "onRejectedThinkingReplayRepaired" + | "session" + | "sessionLockController" + | "sessionManager" +>; +type HistoryPhaseInput = Omit; +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; + sessionLockController: StreamGuardInput["sessionLockController"]; + ownedTranscriptWriteContext: Parameters[0]; + runAbortController: AbortController; + externalAbortController: ExternalAbortController; + abortActiveSession: Parameters[0]["abortActiveSession"]; + abortState: Parameters[0]["state"]; + trackPromptSettlePromise: (promise: Promise) => Promise; + 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["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>; + 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 = (promise: Promise): Promise => + abortableWithSignal(input.runAbortController.signal, promise); + const promptActiveSession = ( + prompt: string, + options?: Parameters[1], + ): Promise => + 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, + }; +} diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts index b8cc5fabfbb7..b87b8ec2be82 100644 --- a/src/agents/embedded-agent-runner/run/attempt.ts +++ b/src/agents/embedded-agent-runner/run/attempt.ts @@ -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>; - 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 = (promise: Promise): Promise => - abortableWithSignal(runAbortController.signal, promise); - const promptActiveSession = ( - prompt: string, - options?: Parameters[1], - ): Promise => - 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,