diff --git a/packages/coding-agent/src/README.md b/packages/coding-agent/src/README.md index c8e290f5b5..a4f8c70273 100644 --- a/packages/coding-agent/src/README.md +++ b/packages/coding-agent/src/README.md @@ -29,10 +29,71 @@ The shared goal types and message formatting remain in `core/goals.ts` during th Use the same ownership rule for the next extraction: move a responsibility's fields, transitions, and cleanup together. Keep request parsing and storage adapters separate when they have independent dependencies. Avoid generic helper folders, modules that receive the entire session, and duplicate copies of feature state. -The next design review should cover input admission and turn lifecycle. The existing `SessionActionStore` already owns action transitions and tickets; build around that ownership when extracting queue dispatch, pause/cancel, and continuation decisions. Do not migrate all callers in the same change as the goal extraction. +## Session input scheduling + +`session/input-scheduler.ts` owns the serialized pump, its preparation epoch, pause leases, and abort/restart suspension. It receives two callbacks: whether the session has work eligible for scheduling, and the operation that runs that work. The scheduler exposes read-only state and named operations; callers cannot change its pause sets or scheduling flags. + +The existing `ActionStore` in `core/session-action-store.ts` owns queued actions, their transitions, and delivery/completion tickets. `session/input-dispatcher.ts` selects and batches those actions, reconciles durable delivery after dispatch, rolls undelivered work back, and settles completion or failure. `AgentSession` supplies turn execution and session-command operations and coordinates goals, child agents, and compaction. The dispatcher shares the existing `ActionStore`; it does not create a second queue or copy the transcript. + +Preserve these distinctions when extending the scheduler: + +- An admission pause blocks new input. A queued-work pause blocks dispatch of already admitted input. They have separate leases and release behavior. +- Starting either pause invalidates asynchronous preparation. Releasing an admission pause also advances the epoch; releasing a queued-work pause retains it. A runner must check its captured epoch after asynchronous work. +- Abort and update restart suspend future scheduling until explicitly resumed. Resume does not release outstanding pause leases. Update restart additionally prevents a custom trigger from implicitly resuming input. +- Pause-release callbacks run once, after the lease is removed. The session retains notification, deferred-message, goal-resumption, and scheduling order. +- Waiting for the pump to settle differs from waiting for the entire session to be idle. Session idle also includes the agent run, event queue, and unfinished actions. + +## Session commit coordination + +`session/commit-fence.ts` owns the FIFO commit queue, its current owner and waiters, asynchronous reentrancy context, and disposal signal. Prompt dispatch, session commands, and branch navigation acquire a lease and run their critical section within its owner context. The session still decides when admission is allowed and when to release the lease. + +- Reentrant work shares the current owner's lease; releasing that nested lease does not release the outer operation. An asynchronous callback from an old owner must queue behind the current owner. +- Cancelling a waiter rejects it promptly but retains its place in the promise chain until its predecessor releases. Later operations cannot overtake that predecessor. +- Disposal rejects waiting and future acquisitions. An already held lease remains owned until its caller releases it. Admission checks still reject disposed sessions before a direct prompt can re-enter. +- Pending work includes queued waiters during the gap between two owners, so daemon passivation cannot mistake a commit handoff for idle. + +The abortable promise helper moved unchanged to `utils/wait-for-abort.ts`, shared by commit acquisition and existing session checkpoint waits. + +The input dispatcher preserves selection and settlement ordering. Batches include only adjacent compatible turns, and preselected turns remain separate. A changed preparation epoch rolls undelivered input back without replaying durable prefix messages. Cancelled actions capturing late messages remain owned until event processing releases them. Checkpoint notifications and queue events retain their distinct positions in these transitions. + +## Session shell commands + +`session/bash.ts` owns shell-command execution, abort controllers, the user-command slot, abort requests during extension dispatch, and deferred transcript output. Its host supplies current shell settings and working directory, extension interception, event delivery, transcript append, and session scheduling notifications. These callbacks read the current runtime so rebuilding extensions or changing settings does not retain stale dependencies. + +`AgentSession` keeps its public shell methods and the cross-feature decision about when to flush deferred output. It also appends messages to agent state before persistence and schedules queued input after the agent becomes idle. The shell owner does not receive the session, kernel, agent loop, or storage manager. + +Execution and recording callbacks preserve dispatch through the public session methods, including wrappers installed by callers. They delegate to the shell owner's corresponding operations; they do not duplicate shell state. + +- User commands reserve the slot before awaiting extensions. Direct executions may overlap and each remains independently abortable. +- Release the user slot and notify waiters before publishing `bash_end`; queued work resumes afterward. +- Extension-provided results take precedence over an abort received during interception. Otherwise, that abort prevents process execution. +- Transient commands publish their lifecycle events but never enter pending output, transcript storage, or model context. Context-excluded commands remain persisted. +- Output produced during streaming waits for the same existing prompt-preparation flush points, preserving tool-call/result ordering. +- Shell event shapes, command options, error behavior, and persisted `bashExecution` messages stay unchanged. + +## Session retry handling + +`session/retry.ts` owns retry attempts, backoff cancellation, retry completion, and authentication-failure tracking. The session reports assistant and agent completion at their existing points in event processing. The retry owner receives current settings, model authentication operations, context inspection, and named operations for continuing or ending a turn. + +- Reserve retry completion synchronously when receiving `agent_end`, before asynchronous event processing. Callers waiting for retry must observe the same pending work. +- Resolve completion before notifying waiters and scheduling queued input. Generation checks keep a rejected continuation from terminating a later retry. +- Preserve provider error classification, retry limits, delay calculation, captured credential identity, and authentication invalidation. Context overflow still belongs to compaction. +- Public retry events and session methods retain their existing shapes and ordering. + +## Turn preparation and action records + +`session/turn-preparation.ts` contains the execution policies for direct, queued, injected, and custom-triggered turns and the ordered preparation pipeline. `TurnPreparer` receives six operations for validation, pending shell output, model selection, compaction, and refinement. The session supplies their implementations and retains transcript dispatch and context rollback. + +Preserve the policy differences: direct prompts flush shell output before validation and compact after model selection; queued turns validate before flushing and compact before model selection. Conditional refinement barriers are checked when reached, so a refinement started during preparation is still awaited. Withdrawing prepared work skips the final barrier and commit. The exported `TurnExecutionPolicy` shape remains available from the session facade. + +`session/prepared-actions.ts` contains prepared action types, delivery records, recovery contracts, input copying, action factories, and queue projections. It has no session dependency. Primary messages retain their identity for durable-delivery checks; separately stored input blocks and prefix messages retain their existing copy behavior. Recovery format version 1 and the public exports from `AgentSession` stay unchanged. ## Validation Controller tests live in `test/goals/`. Session integration coverage remains in `test/suite/agent-session-goal.test.ts`, `test/suite/agent-session-compaction-continuation.test.ts`, and `test/goal-continuation-quiescence.test.ts`. +Scheduler and commit-fence tests live in `test/session/`. Existing queue, action-contract, action-race, and compaction suites cover the integration with `AgentSession`, including pause, cancellation, restart, branch navigation, and goal continuation. + +Shell-owner tests also live in `test/session/`. Session bash/persistence, prompt, queue, and side-question regression suites retain end-to-end coverage of scheduling and transcript behavior using the faux provider and controlled shell operations. + Run the focused files from the coding-agent package root with the repository's prescribed Vitest command, then run `npm run check` from the repository root. Use faux providers for session tests. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 674fe40064..919a636670 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -38,11 +38,44 @@ import { parseGoalSlashCommand } from "../goals/commands.js"; import { GoalController } from "../goals/controller.js"; import { createGoalPersistence } from "../goals/persistence.js"; import { theme } from "../modes/interactive/theme/theme.js"; +import { + type ExecuteBashOptions, + type RunUserBashOptions, + SessionBash, + type SessionBashEvent, +} from "../session/bash.js"; +import { SessionCommitFence, type SessionCommitLease } from "../session/commit-fence.js"; +import { SessionInputDispatcher } from "../session/input-dispatcher.js"; +import { SessionInputScheduler } from "../session/input-scheduler.js"; +import { + buildPromptContent, + cloneCustomMessage, + cloneQueuedAgentMessage, + createDeliveryRecord, + createPreparedTurnAction, + createSessionCommandAction, + DeferredSessionInputError, + normalizeMessageContent, + type PreparedCommandPayload, + type PreparedPromptPreparation, + type PreparedTurnPayload, + primaryDeliveryRecord, + type QueuedAgentMessage, + type QueuedSessionAction, + queuedAgentMessagePreview, + type RestoredPromptInput, + SESSION_ACTION_RECOVERY_FORMAT_VERSION, + type SessionActionRecoverySnapshot, + SessionInputAdmissionPausedError, + type SessionInputSchedule, + visibleSessionActionProjection, +} from "../session/prepared-actions.js"; +import { SessionRetry, type SessionRetryEvent } from "../session/retry.js"; +import { createTurnExecutionPolicy, type TurnExecutionPolicy, TurnPreparer } from "../session/turn-preparation.js"; import { stripFrontmatter } from "../utils/frontmatter.js"; -import { sleep } from "../utils/sleep.js"; +import { waitForPromiseOrAbort } from "../utils/wait-for-abort.js"; import { AGENT_MESSAGE_CUSTOM_TYPE, - AGENT_MESSAGE_RECEIVED_PREVIEW_LABEL, AGENT_MESSAGE_SKILL_NAME, type AgentFamilyCatalogEntry, type AgentSessionMessage, @@ -81,7 +114,6 @@ import { formatNoModelSelectedMessage, isLikelyAuthenticationError, } from "./auth-guidance.js"; -import type { AuthSourceToken } from "./auth-storage.js"; import { type AgentAutonomousConfig, type AgentAutonomousStatus, @@ -97,7 +129,7 @@ import { setAutonomousLimits, UNLIMITED_AUTONOMOUS_LIMIT, } from "./autonomous.js"; -import { type BashResult, executeBashWithOperations } from "./bash-executor.js"; +import type { BashResult } from "./bash-executor.js"; import { COMPACT_SKILL_NAME, type CompactionResult, @@ -167,8 +199,6 @@ import type { McpManager } from "./mcp/mcp-manager.js"; import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE, ASYNC_BASH_COMPLETION_PREVIEW_LABEL, - type AsyncBashCompletionDetails, - type BashExecutionMessage, type CompactionOutcome, type CompactionOutcomeReason, type CustomMessage, @@ -196,15 +226,7 @@ import { import type { ModelRegistry } from "./model-registry.js"; import { throwIfPromptAdmissionCancelled } from "./prompt-admission.js"; import { expandPromptTemplate, type PromptTemplate, parseCommandArgs } from "./prompt-templates.js"; -import { - isAgentLifecycleFailure, - isFauxProviderQueueExhausted, - isPermanentProviderFailureKind, - providerRetryDelay, - providerRetryPolicy, - providerStreamFailureKind, - providerStreamFailureRetryAfterMs, -} from "./provider-retry.js"; +import { providerRetryPolicy } from "./provider-retry.js"; import { type AutoRefineReason, type AutoRefineReview, @@ -272,10 +294,7 @@ import { type RuntimeActivity, type SessionAction, type SessionActionSnapshot, - type SessionCommandPayload, - type SessionTurnPayload, transitionSessionAction, - type WakePolicy, } from "./session-action-store.js"; import type { BranchSummaryEntry, @@ -305,7 +324,6 @@ import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.js"; import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-prompt.js"; import { THINKING_LEVELS } from "./thinking-levels.js"; import { acpMcpToolNames, createAcpMcpToolDefinitions } from "./tools/acp-mcp.js"; -import { type BashOperations, createLocalBashOperations } from "./tools/bash.js"; import { createAllToolDefinitions } from "./tools/index.js"; import { IpythonKernelProvisioner } from "./tools/ipython.js"; import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js"; @@ -377,57 +395,25 @@ export type AgentSessionEvent = errorSeverity?: "warning" | "error"; customInstructions?: string; } - | { - type: "auto_retry_start"; - attempt: number; - maxAttempts: number; - delayMs: number; - errorMessage: string; - } - | { - type: "auto_retry_end"; - success: boolean; - attempt: number; - finalError?: string; - } - | { - type: "auth_stale"; - provider: string; - sourceTokens?: readonly AuthSourceToken[]; - } + | SessionRetryEvent | { type: "rlm_child_update"; child: RlmChildAgentSnapshot } | { type: "recap_update"; recap: string | undefined } | { type: "goal_update"; goal: GoalState } - | { - type: "bash_start"; - command: string; - excludeFromContext: boolean; - transient?: boolean; - runId?: string; - } - | { type: "bash_output"; chunk: string } - | { - type: "bash_end"; - exitCode: number | undefined; - cancelled: boolean; - truncated: boolean; - fullOutputPath?: string; - errorMessage?: string; - transient?: boolean; - runId?: string; - } + | SessionBashEvent | { type: "refine_complete"; result: RefinementResult } | { type: "refine_failed"; error: string }; -export type AgentSessionEventListener = (event: AgentSessionEvent) => void; +export { + SESSION_ACTION_RECOVERY_FORMAT_VERSION, + type SessionActionRecoveryAction, + type SessionActionRecoveryPayload, + type SessionActionRecoveryRecord, + type SessionActionRecoverySnapshot, +} from "../session/prepared-actions.js"; + +export type { TurnExecutionPolicy } from "../session/turn-preparation.js"; -type UserBashEndDetails = { - exitCode: number | undefined; - cancelled: boolean; - truncated: boolean; - fullOutputPath?: string; - errorMessage?: string; -}; +export type AgentSessionEventListener = (event: AgentSessionEvent) => void; export class CompactionSkippedError extends Error {} @@ -590,80 +576,6 @@ type NormalizedSubmission = | { kind: "extensionCommand"; completion: Promise } | { kind: "handled" }; -type PreTurnCompactionTiming = "beforeModelSelection" | "afterModelSelection" | "skip"; -type RefineBarrierPolicy = "always" | "ifInFlight" | "skip"; - -interface CommitPreparationPolicy { - initialRefineBarrier: RefineBarrierPolicy; - flushPendingBashBeforeValidation: boolean; - validateModelAndAuth: boolean; - awaitPendingModelSelection: boolean; - preTurnCompaction: PreTurnCompactionTiming; - finalRefineBarrier: RefineBarrierPolicy; -} - -interface CommitPreparationSteps { - afterValidation?: () => void; - prepare: () => Promise; - shouldCommit?: (prepared: TPrepared) => boolean; - beforeFinalRefineBarrier?: (prepared: TPrepared) => void; - commit: (prepared: TPrepared, passedFinalRefineBarrier: boolean) => TCommitted; -} - -type QueuedAgentMessage = UserMessage | CustomMessage; -type SessionInputSchedule = "steer" | "followUp"; - -export interface TurnExecutionPolicy { - preparation: CommitPreparationPolicy; - runBeforeAgentStart: boolean; - nextTurnContextTiming: "preparation" | "commit" | "skip"; - preserveEmptyExtensionPrompt: boolean; - completionIncludesRetryChain: boolean; -} - -function turnExecutionPoliciesEqual(left: TurnExecutionPolicy, right: TurnExecutionPolicy): boolean { - return ( - left.preparation.initialRefineBarrier === right.preparation.initialRefineBarrier && - left.preparation.flushPendingBashBeforeValidation === right.preparation.flushPendingBashBeforeValidation && - left.preparation.validateModelAndAuth === right.preparation.validateModelAndAuth && - left.preparation.awaitPendingModelSelection === right.preparation.awaitPendingModelSelection && - left.preparation.preTurnCompaction === right.preparation.preTurnCompaction && - left.preparation.finalRefineBarrier === right.preparation.finalRefineBarrier && - left.runBeforeAgentStart === right.runBeforeAgentStart && - left.nextTurnContextTiming === right.nextTurnContextTiming && - left.preserveEmptyExtensionPrompt === right.preserveEmptyExtensionPrompt && - left.completionIncludesRetryChain === right.completionIncludesRetryChain - ); -} - -interface PreparedTurnPayload extends SessionTurnPayload { - images?: ImageContent[]; - content?: (TextContent | ImageContent)[]; - customMessage?: CustomMessage; - prepared?: PreparedPromptPreparation; - executionPolicy: TurnExecutionPolicy; - queueVisible: boolean; - acceptedAgentMessage: boolean; - acceptedBeforeCompletion: boolean; - captureRunMessages?: Set; - cancelledDispatchEnded?: boolean; -} - -interface PreparedCommandPayload extends SessionCommandPayload { - images?: ImageContent[]; -} - -type QueuedSessionAction = SessionAction; - -interface PreparedPromptPreparation { - result: Awaited>; - basePromptSnapshot: string; -} - -class DeferredSessionInputError extends Error {} - -class SessionInputAdmissionPausedError extends Error {} - function oncePreflight( preflightResult: ((success: boolean, queued?: boolean) => void) | undefined, ): (success: boolean, queued?: boolean) => void { @@ -676,121 +588,6 @@ function oncePreflight( }; } -interface RestoredPromptInput { - text: string; - content?: (TextContent | ImageContent)[]; - images?: ImageContent[]; - queueKey?: string; - agentMessageId?: string; - customMessage?: CustomMessage; - prefixMessages?: CustomMessage[]; -} - -export const SESSION_ACTION_RECOVERY_FORMAT_VERSION = 1; - -export interface SessionActionRecoveryRecord { - id: string; - role: DeliveryRecord["role"]; - message: QueuedAgentMessage; - ownerActionId: string; -} - -export type SessionActionRecoveryPayload = - | { - kind: "turn"; - text: string; - preview?: string; - records: SessionActionRecoveryRecord[]; - images?: ImageContent[]; - content?: (TextContent | ImageContent)[]; - customMessage?: CustomMessage; - executionPolicy: TurnExecutionPolicy; - queueVisible: boolean; - acceptedAgentMessage: boolean; - acceptedBeforeCompletion: boolean; - } - | { - kind: "session_command"; - text: string; - command: SessionSlashCommand; - images?: ImageContent[]; - }; - -export interface SessionActionRecoveryAction { - id: string; - source: InputSource | "internal"; - delivery: DeliveryPolicy; - wake: WakePolicy; - payload: SessionActionRecoveryPayload; - queueKey?: string; - agentMessageId?: string; - suppressAutonomousContinuation?: boolean; -} - -export interface SessionActionRecoverySnapshot { - formatVersion: typeof SESSION_ACTION_RECOVERY_FORMAT_VERSION; - actions: SessionActionRecoveryAction[]; -} - -function cloneCustomMessage(message: CustomMessage): CustomMessage { - return { - ...message, - content: Array.isArray(message.content) ? message.content.map((block) => ({ ...block })) : message.content, - }; -} - -function cloneQueuedAgentMessage(message: QueuedAgentMessage): QueuedAgentMessage { - if (message.role === "custom") return cloneCustomMessage(message); - return { - ...message, - content: Array.isArray(message.content) ? message.content.map((block) => ({ ...block })) : message.content, - }; -} - -function primaryDeliveryRecord(action: QueuedSessionAction): DeliveryRecord { - if (action.payload.kind !== "turn") throw new Error(`Session action ${action.id} is not a turn`); - const record = action.payload.records.find((candidate) => candidate.role === "primary"); - if (!record) throw new Error(`Turn action ${action.id} has no primary delivery record`); - return record; -} - -function normalizeMessageContent(content: string | (TextContent | ImageContent)[]): { - text: string; - images?: ImageContent[]; -} { - if (typeof content === "string") return { text: content }; - const text = content - .filter((part): part is TextContent => part.type === "text") - .map((part) => part.text) - .join("\n"); - const images = content.filter((part): part is ImageContent => part.type === "image"); - return { text, ...(images.length > 0 ? { images } : {}) }; -} - -function queuedAgentMessagePreview(action: QueuedSessionAction): string { - const payload = action.payload; - if (payload.kind === "session_command") return payload.text; - if (payload.customMessage && isAgentSessionMessage(payload.customMessage)) { - return `${AGENT_MESSAGE_RECEIVED_PREVIEW_LABEL}: ${payload.customMessage.details.message}`; - } - if (payload.customMessage?.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE) { - const details = payload.customMessage.details as AsyncBashCompletionDetails | undefined; - return details - ? `${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}: pid ${details.pid}, exit ${details.exitCode}` - : ASYNC_BASH_COMPLETION_PREVIEW_LABEL; - } - return payload.preview ?? payload.text; -} - -function visibleSessionActionProjection(actions: readonly QueuedSessionAction[]): readonly QueuedSessionAction[] { - return actions.filter( - (action) => - action.payload.kind === "session_command" || - action.payload.queueVisible || - action.payload.acceptedAgentMessage, - ); -} - const IPYTHON_SENT_AGENT_MESSAGE_CUSTOM_ENTRY = "ipython_sent_agent_message"; interface PersistedIpythonSentAgentMessage { @@ -1143,35 +940,6 @@ function readAssistantText(message: AssistantMessage): string { .join(""); } -function waitForPromiseOrAbort( - promise: Promise, - signal: AbortSignal | undefined, - abortMessage: string, -): Promise { - if (!signal) return promise; - if (signal.aborted) return Promise.reject(new Error(abortMessage)); - return new Promise((resolve, reject) => { - const onAbort = () => { - cleanup(); - reject(new Error(abortMessage)); - }; - const cleanup = () => signal.removeEventListener("abort", onAbort); - signal.addEventListener("abort", onAbort, { once: true }); - // Close the listener-registration race before observing the awaited work. - if (signal.aborted) return onAbort(); - promise.then( - (value) => { - cleanup(); - resolve(value); - }, - (error: unknown) => { - cleanup(); - reject(error); - }, - ); - }); -} - // Bounds how much accumulated child usage a parent process crash can lose. const RLM_CHILD_USAGE_FLUSH_MAX_PENDING_MS = 60_000; @@ -1225,23 +993,44 @@ export class AgentSession { /** Session-owned actions. Items are never fed into Agent.steer/followUp. */ private readonly _actionStore = new ActionStore(); - private _sessionInputPump: Promise = Promise.resolve(); - private _sessionInputPumpRequested = false; - // Invalidates preparation when a branch pause starts and finishes before its next await resumes. - private _sessionInputPumpEpoch = 0; + private readonly _inputScheduler = new SessionInputScheduler({ + canSchedule: () => !this._disposed && !this._disposing && this._hasSelectableSessionInput(), + run: (epoch) => this._inputDispatcher.run(epoch), + }); + private readonly _inputDispatcher = new SessionInputDispatcher(this._actionStore, { + isDisposed: () => this._disposed || this._disposing, + getEpoch: () => this._inputScheduler.epoch, + getActivity: () => this._runtimeActivity(), + isBusy: () => this._isBusyForSessionInput("pump"), + isHandoffDeferred: (epoch) => this._isSessionInputHandoffDeferred(epoch), + getDeliveryMode: (delivery) => (delivery === "next_turn_boundary" ? this.steeringMode : this.followUpMode), + waitForAgentIdle: () => this.agent.waitForIdle(), + hasCancelledDispatchCapture: () => this._hasCancelledDispatchCapture(), + getEventQueue: () => this._agentEventQueue, + waitForRefinement: () => this._waitForRefineIdle(), + getTranscript: () => this.agent.state.messages, + startTurns: (actions, epoch) => this._startPreparedTurnActions(actions, epoch), + executeCommand: (action, epoch) => this._executeSelectedSessionCommand(action, epoch), + settleAgentMessage: (id, leg, error) => this._settleAgentMessage(id, leg, error), + releaseTurn: (id) => { + this._durableRlmTerminalNoticeActionIds.delete(id); + }, + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + emitQueueUpdate: () => this._emitQueueUpdate(), + surfaceError: (error) => this._surfaceSessionInputError(error), + schedule: () => this._scheduleSessionInputPump(), + }); private _sessionInputArrivalEpoch = 0; - // Persists abort/restart suspension after the initiating call returns. - private _sessionInputPumpSuspended = false; - private _sessionInputSuspendedForUpdateRestart = false; - // Branch mutation pause leases can overlap and must all release before dispatch resumes. - private readonly _queuedWorkPauses = new Set(); - private readonly _sessionInputAdmissionPauses = new Set(); private readonly _durableRlmTerminalNoticeActionIds = new Set(); - private _sessionActionCommitTail: Promise = Promise.resolve(); - private _sessionActionCommitOwner: symbol | undefined; - private _pendingSessionActionFenceWaiters = 0; - private readonly _sessionActionCommitContext = new AsyncLocalStorage(); - private readonly _sessionActionCommitDisposeAbortController = new AbortController(); + private readonly _commitFence = new SessionCommitFence(); + private readonly _turnPreparer = new TurnPreparer({ + hasRefinement: () => this._refineInFlight !== undefined, + waitForRefinement: () => this._waitForRefineIdle(), + flushPendingBash: () => this._flushPendingBashMessages(), + validate: () => this._validateCanStartAgentRun(), + compact: () => this._runPreTurnCompaction(), + pendingModelSelection: () => this._pendingModelSelectEmit(), + }); // Checkpoint, handoff, and activity waiters share lifecycle-edge notifications to avoid polling. private readonly _sessionInputCheckpointWaiters = new Set<() => void>(); private _pendingNextTurnMessages: CustomMessage[] = []; @@ -1265,13 +1054,34 @@ export class AgentSession { private _branchSummaryAbortController: AbortController | undefined = undefined; private _branchSummaryOperation: Promise | undefined = undefined; - private _retryAbortController: AbortController | undefined = undefined; - private _retryAttempt = 0; - /** Bumped by every retry resolution; stale scheduled-continue callbacks check it before touching retry state. */ - private _retryGeneration = 0; - private _retryPromise: Promise | undefined = undefined; - private _retryResolve: (() => void) | undefined = undefined; - private _retryAuthFailureSources: AuthSourceToken[] = []; + private readonly _retry = new SessionRetry({ + getRetrySettings: () => this.settingsManager.getRetrySettings(), + getMaxRetryDelayMs: () => this.settingsManager.getProviderRetrySettings().maxRetryDelayMs, + getContextWindow: () => this.model?.contextWindow ?? 0, + getAuthSource: (provider) => this._modelRegistry.getCurrentProviderAuthSourceToken(provider), + markAuthSourceStale: (token) => this._modelRegistry.markProviderAuthSourceStale(token), + markAuthStale: (provider) => this._modelRegistry.markProviderAuthStale(provider), + hasPayloadHooks: () => this._extensionRunner.hasHandlers("before_provider_request"), + prepareTurnRetry: () => this._semanticEdges.prepareTurnRetry(), + clearTurnRetry: () => this._semanticEdges.clearTurnRetry(), + removeLastAssistant: () => { + const messages = this.agent.state.messages; + if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { + this.agent.state.messages = messages.slice(0, -1); + } + }, + continue: () => this.agent.continue(), + waitForIdle: () => this.agent.waitForIdle(), + cancelCompaction: () => { + this._autoCompactionAbortController?.abort(); + this._cancelPostCompactionContinue(); + }, + emit: (event) => this._emit(event), + onResolved: () => { + this._notifySessionInputCheckpointChange(); + this._scheduleSessionInputPump(); + }, + }); private _agentMessageClearEpoch = 0; private _agentMessageOutcomes = new Map(); private _lateIpythonSentAgentMessages = new Map(); @@ -1280,10 +1090,22 @@ export class AgentSession { /** Fresh/empty contexts defer digest injection to the first committed turn so untouched sessions stay empty. */ private _harnessDigestPending = false; - private _bashAbortControllers = new Set(); - private _userBashRunning = false; - private _userBashAbortRequested = false; - private _pendingBashMessages: BashExecutionMessage[] = []; + private readonly _bash = new SessionBash({ + getCwd: () => this.sessionManager.getCwd(), + getShellCommandPrefix: () => this.settingsManager.getShellCommandPrefix(), + getShellPath: () => this.settingsManager.getShellPath(), + isStreaming: () => this.isStreaming, + intercept: (event) => this._extensionRunner.emitUserBash(event), + emit: (event) => this._emit(event), + appendMessage: (message) => { + this.agent.state.messages.push(message); + this.sessionManager.appendMessage(message); + }, + onStateChange: () => this._notifySessionInputCheckpointChange(), + onUserBashEnd: () => this._drainQueuedMessagesAfterBash(), + executeBash: (command, onChunk, options) => this.executeBash(command, onChunk, options), + recordBashResult: (command, result, options) => this.recordBashResult(command, result, options), + }); private _extensionRunner!: ExtensionRunner; private _execEnvProvider?: () => Record | undefined; @@ -2103,7 +1925,7 @@ export class AgentSession { } // Keep the deferral while admission is paused or the pump is suspended // (post-abort); the pause release and resumeQueuedWork retry. - if (this._sessionInputAdmissionPauses.size > 0 || this._sessionInputPumpSuspended) return; + if (this._inputScheduler.admissionPaused || this._inputScheduler.suspended) return; const goalBeforeResume = this._goals.checkpoint(); try { this._ensureGoalRuntimeActive(); @@ -3385,7 +3207,7 @@ export class AgentSession { } private _handleAgentEvent = (event: AgentEvent): void => { - this._createRetryPromiseForAgentEnd(event); + this._retry.observeAgentEnd(event); if (event.type === "message_start" || event.type === "message_end") { for (const action of this._actionStore.ownedActions()) { if ( @@ -3448,30 +3270,6 @@ export class AgentSession { this._agentEventQueue.catch(() => {}); }; - private _createRetryPromiseForAgentEnd(event: AgentEvent): void { - if (event.type !== "agent_end" || this._retryPromise) { - return; - } - - const settings = this.settingsManager.getRetrySettings(); - if (!settings.enabled) { - return; - } - - const lastAssistant = this._findLastAssistantInMessages(event.messages); - const concreteAuthFailure = lastAssistant ? this._isConcreteProviderAuthFailure(lastAssistant) : false; - if (!lastAssistant || (!this._isRetryableError(lastAssistant) && !concreteAuthFailure)) { - return; - } - if (concreteAuthFailure) { - this._captureRetryAuthFailureSource(lastAssistant); - } - - this._retryPromise = new Promise((resolve) => { - this._retryResolve = resolve; - }); - } - private _findLastAssistantInMessages(messages: AgentMessage[]): AssistantMessage | undefined { for (let i = messages.length - 1; i >= 0; i--) { const message = messages[i]; @@ -3532,7 +3330,7 @@ export class AgentSession { this._lastAssistantMessage = undefined; for (const action of cleared) this._actionStore.releaseTerminal(action); this._notifySessionInputCheckpointChange(); - this._resolveRetry(); + this._retry.resolve(); } } @@ -3589,21 +3387,7 @@ export class AgentSession { if (assistantMsg.stopReason !== "error") { this._overflowRecovery = "idle"; } - if (this._isConcreteProviderAuthFailure(assistantMsg)) { - this._captureRetryAuthFailureSource(assistantMsg); - } - - // Reset retry counter immediately on successful assistant response - // This prevents accumulation across multiple LLM calls within a turn - if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) { - this._emit({ - type: "auto_retry_end", - success: true, - attempt: this._retryAttempt, - }); - this._retryAttempt = 0; - this._retryAuthFailureSources = []; - } + this._retry.observeAssistantEnd(assistantMsg); if (this._goals.accountAssistantMessage(assistantMsg)) { const message = createGoalContextMessage(this._goals.state, "budget_limit"); const normalized = normalizeMessageContent(message.content); @@ -3622,33 +3406,22 @@ export class AgentSession { if (event.type === "agent_end") { const msg = this._lastAssistantMessage ?? - (this._retryPromise ? this._findLastAssistantInMessages(event.messages) : undefined); + (this._retry.isRetrying ? this._findLastAssistantInMessages(event.messages) : undefined); this._lastAssistantMessage = undefined; if (!msg) { - this._resolveRetry(); + this._retry.resolve(); return; } - const concreteAuthFailure = this._isConcreteProviderAuthFailure(msg); - const retryConcreteAuthFailure = - concreteAuthFailure && !this._isStructuredPermanentProviderRetryExhausted(msg); - if (this._isRetryableError(msg) || retryConcreteAuthFailure) { - if (retryConcreteAuthFailure) { - this._captureRetryAuthFailureSource(msg); - } - const didRetry = await this._handleRetryableError(msg, { - markAuthStaleOnFailure: retryConcreteAuthFailure, - authSourceTokens: retryConcreteAuthFailure ? this._retryAuthFailureSources : undefined, - }); - if (didRetry) return; // Retry was initiated, don't proceed to compaction - } + const retry = this._retry.retryError(msg); + if (retry && (await retry)) return; const compactionWillRetry = await this._checkCompaction(msg); - if (compactionWillRetry && this._retryAttempt > 0) { + if (compactionWillRetry && this._retry.attempt > 0) { return; } - this._finishActiveRetryWithFailure(msg); - this._resolveRetry(); + this._retry.finishActiveRetryWithFailure(msg); + this._retry.resolve(); if (!compactionWillRetry) { this._finishGoalForTerminalAssistantMessage(msg); // In serialized mode, agent-callable refine.run is serviced @@ -3663,18 +3436,6 @@ export class AgentSession { } } - private _resolveRetry(): void { - this._retryGeneration += 1; - this._semanticEdges.clearTurnRetry(); - if (this._retryResolve) { - this._retryResolve(); - this._retryResolve = undefined; - this._retryPromise = undefined; - this._notifySessionInputCheckpointChange(); - this._scheduleSessionInputPump(); - } - } - private _findLastAssistantMessage(): AssistantMessage | undefined { const messages = this.agent.state.messages; for (let i = messages.length - 1; i >= 0; i--) { @@ -3845,7 +3606,7 @@ export class AgentSession { return this._disposeCallbacksPromise; } this._disposing = true; - this._sessionActionCommitDisposeAbortController.abort(); + this._commitFence.dispose(); await this._disposeAsyncOnce(kernelSnapshot); })(); return this._disposeAsyncPromise; @@ -4050,7 +3811,7 @@ export class AgentSession { this._disposed = true; for (const run of this._unsettledRlmChildRuns) run.suppressTerminalNotice = true; for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); - this._sessionActionCommitDisposeAbortController.abort(); + this._commitFence.dispose(); try { // Invalidate scheduled timers and abort any in-flight review so a late // resolution cannot write harness state or re-subscribe handlers. @@ -4135,7 +3896,7 @@ export class AgentSession { } get retryAttempt(): number { - return this._retryAttempt; + return this._retry.attempt; } getActiveToolNames(): string[] { @@ -4410,42 +4171,6 @@ export class AgentSession { if (lastAssistant) await this._checkCompaction(lastAssistant, false, false); } - private async _prepareForCommit( - policy: CommitPreparationPolicy, - steps: CommitPreparationSteps, - ): Promise { - if ( - policy.initialRefineBarrier === "always" || - (policy.initialRefineBarrier === "ifInFlight" && this._refineInFlight) - ) { - await this._waitForRefineIdle(); - } - if (policy.flushPendingBashBeforeValidation) this._flushPendingBashMessages(); - if (policy.validateModelAndAuth) await this._validateCanStartAgentRun(); - steps.afterValidation?.(); - if (!policy.flushPendingBashBeforeValidation) this._flushPendingBashMessages(); - - if (policy.preTurnCompaction === "beforeModelSelection") await this._runPreTurnCompaction(); - if (policy.awaitPendingModelSelection) { - const pendingModelSelectEmit = this._pendingModelSelectEmit(); - if (pendingModelSelectEmit) await pendingModelSelectEmit; - } - if (policy.preTurnCompaction === "afterModelSelection") await this._runPreTurnCompaction(); - - const prepared = await steps.prepare(); - if (steps.shouldCommit && !steps.shouldCommit(prepared)) return undefined; - steps.beforeFinalRefineBarrier?.(prepared); - let passedFinalRefineBarrier = false; - if ( - policy.finalRefineBarrier === "always" || - (policy.finalRefineBarrier === "ifInFlight" && this._refineInFlight) - ) { - await this._waitForRefineIdle(); - passedFinalRefineBarrier = true; - } - return steps.commit(prepared, passedFinalRefineBarrier); - } - private _applyPreparedSystemPrompt( preparation: PreparedPromptPreparation | undefined, preserveEmptyExtensionPrompt: boolean, @@ -4466,8 +4191,8 @@ export class AgentSession { !this.isCompacting && !this.isRetrying && !this.isBashRunning && - !this._sessionInputPumpSuspended && - this._queuedWorkPauses.size === 0 && + !this._inputScheduler.suspended && + this._inputScheduler.queuedWorkPauseCount === 0 && !this._disposed && !this._disposing ); @@ -4538,7 +4263,7 @@ export class AgentSession { } }; if ( - this._sessionInputPumpSuspended && + this._inputScheduler.suspended && this._isBusyForSessionInput("preflight") && options?.queueIfBusy === true && options.streamingBehavior @@ -4623,7 +4348,7 @@ export class AgentSession { suppressAutonomousContinuation: true, resumeIfIdle: false, source: "internal", - executionPolicy: this._turnExecutionPolicy("injected"), + executionPolicy: createTurnExecutionPolicy("injected"), queueVisible: false, }); this._durableRlmTerminalNoticeActionIds.add(action.id); @@ -4638,9 +4363,9 @@ export class AgentSession { private _flushDeferredRlmTerminalNotices(): void { if ( - this._sessionInputAdmissionPauses.size > 0 || - this._sessionInputPumpSuspended || - this._queuedWorkPauses.size > 0 || + this._inputScheduler.admissionPaused || + this._inputScheduler.suspended || + this._inputScheduler.queuedWorkPauseCount > 0 || this._disposed || this._disposing ) { @@ -4661,9 +4386,9 @@ export class AgentSession { } private async _acquireRlmTerminalNoticeRetentionFence(): Promise<{ owner: symbol; release(): void } | undefined> { - const disposeSignal = this._sessionActionCommitDisposeAbortController.signal; + const disposeSignal = this._commitFence.disposeSignal; while (!this._disposed && !this._disposing && !disposeSignal.aborted) { - if (this._queuedWorkPauses.size > 0) { + if (this._inputScheduler.queuedWorkPauseCount > 0) { let wake = () => {}; const pauseReleased = new Promise((resolve) => { wake = resolve; @@ -4684,7 +4409,7 @@ export class AgentSession { } catch { return undefined; } - if (this._queuedWorkPauses.size === 0 && !this._disposed && !this._disposing) return fence; + if (this._inputScheduler.queuedWorkPauseCount === 0 && !this._disposed && !this._disposing) return fence; fence.release(); } return undefined; @@ -4728,7 +4453,7 @@ export class AgentSession { options?: InternalPromptOptions & { executionPolicy?: TurnExecutionPolicy }, ): Promise { if (!this.isStreaming && options?.resumeIfIdle) this._resumeSessionInputAdmission(); - const admissionEpoch = this._sessionInputPumpEpoch; + const admissionEpoch = this._inputScheduler.epoch; const admissionFence = await this._acquireDirectTurnAdmissionFence(options?.signal).catch((error: unknown) => { throwIfPromptAdmissionCancelled(options?.signal); throw error; @@ -4736,7 +4461,7 @@ export class AgentSession { const reportPreflight = oncePreflight(options?.preflightResult); try { throwIfPromptAdmissionCancelled(options?.signal); - if (admissionEpoch !== this._sessionInputPumpEpoch) { + if (admissionEpoch !== this._inputScheduler.epoch) { throw new Error("Injected session input was invalidated before admission"); } options?.admissionCommitted?.(); @@ -4764,7 +4489,7 @@ export class AgentSession { source: options?.source ?? "internal", executionPolicy: options?.executionPolicy ?? - (visibleQueued ? this._turnExecutionPolicy("queued") : this._turnExecutionPolicy("injected")), + (visibleQueued ? createTurnExecutionPolicy("queued") : createTurnExecutionPolicy("injected")), queueVisible: visibleQueued, }); const result = this._admitSessionInput(action, { @@ -4804,7 +4529,7 @@ export class AgentSession { if (resumeSuspendedInput) this._resumeSessionInputAdmission(); this._assertSessionActionAdmissionAvailable(); } - const admissionEpoch = this._sessionInputPumpEpoch; + const admissionEpoch = this._inputScheduler.epoch; const commitFence = this.isStreaming ? undefined : await this._acquireDirectTurnAdmissionFence(options?.signal).catch((error: unknown) => { @@ -4815,7 +4540,7 @@ export class AgentSession { const run = async () => { try { throwIfPromptAdmissionCancelled(options?.signal); - if (!resumeSuspendedInput && admissionEpoch !== this._sessionInputPumpEpoch) { + if (!resumeSuspendedInput && admissionEpoch !== this._inputScheduler.epoch) { throw new Error("Session input was invalidated before admission"); } options?.admissionCommitted?.(); @@ -4857,7 +4582,7 @@ export class AgentSession { const wasBusy = wasRuntimeBusy || pendingOwnedWork; if (normalized.kind === "sessionCommand") { const schedule = options?.streamingBehavior ?? (this.isStreaming ? "steer" : "followUp"); - const action = this._createSessionCommandAction( + const action = createSessionCommandAction( normalized.text, normalized.command, normalized.images, @@ -4895,7 +4620,7 @@ export class AgentSession { const prefixMessages = visibleQueued ? this._takePendingNextTurnMessages() : undefined; const content = options?.content ? options.content.map((block) => ({ ...block })) - : this._buildPromptContent(normalized.text, normalized.images); + : buildPromptContent(normalized.text, normalized.images); const suppliedMessage = options?.customMessage; const primaryMessage = suppliedMessage ? visibleQueued @@ -4920,8 +4645,8 @@ export class AgentSession { (options?.queueIfBusy === true && canSelectSessionAction(this._runtimeActivity())), source: isInternalPrompt ? "internal" : (options?.source ?? "interactive"), executionPolicy: visibleQueued - ? this._turnExecutionPolicy("queued") - : this._turnExecutionPolicy("directPrompt", { + ? createTurnExecutionPolicy("queued") + : createTurnExecutionPolicy("directPrompt", { returnAfterAccepted: options?.returnAfterAccepted, skipPrePromptWork: options?.skipPrePromptWork, }), @@ -4994,7 +4719,7 @@ export class AgentSession { commitFence?.release(); } }; - return commitFence ? this._sessionActionCommitContext.run(commitFence.owner, run) : run(); + return commitFence ? this._commitFence.run(commitFence, run) : run(); } private _executeExtensionCommand(text: string): Promise | undefined { @@ -5223,7 +4948,7 @@ export class AgentSession { return undefined; } return this._admitSessionInput( - this._createSessionCommandAction(text, customMessage.details.command, images, schedule, { + createSessionCommandAction(text, customMessage.details.command, images, schedule, { agentMessageId, source: "internal", }), @@ -5300,193 +5025,14 @@ export class AgentSession { }); } - private _buildPromptContent(text: string, images?: ImageContent[]): (TextContent | ImageContent)[] { - const content: (TextContent | ImageContent)[] = []; - content.push({ type: "text", text }); - if (images) content.push(...images); - return content; - } - private _takePendingNextTurnMessages(): CustomMessage[] { const messages = this._pendingNextTurnMessages; this._pendingNextTurnMessages = []; return messages; } - private _deliveryPolicy(schedule: SessionInputSchedule): DeliveryPolicy { - return schedule === "steer" ? "next_turn_boundary" : "when_run_idle"; - } - - private _createDeliveryRecord( - actionId: string, - role: DeliveryRecord["role"], - message: QueuedAgentMessage, - ): DeliveryRecord { - return { - id: randomUUID(), - role, - message, - started: false, - durable: false, - ownerActionId: actionId, - }; - } - - private _turnExecutionPolicy( - kind: "queued" | "directPrompt" | "injected" | "customTrigger", - options: { - returnAfterAccepted?: boolean; - skipPrePromptWork?: boolean; - } = {}, - ): TurnExecutionPolicy { - if (kind === "queued") { - return { - preparation: { - initialRefineBarrier: "skip", - flushPendingBashBeforeValidation: false, - validateModelAndAuth: true, - awaitPendingModelSelection: true, - preTurnCompaction: "beforeModelSelection", - finalRefineBarrier: "always", - }, - runBeforeAgentStart: true, - nextTurnContextTiming: "commit", - preserveEmptyExtensionPrompt: true, - completionIncludesRetryChain: true, - }; - } - if (kind === "directPrompt") { - return { - preparation: { - initialRefineBarrier: options.returnAfterAccepted ? "skip" : "always", - flushPendingBashBeforeValidation: true, - validateModelAndAuth: true, - awaitPendingModelSelection: true, - preTurnCompaction: options.skipPrePromptWork ? "skip" : "afterModelSelection", - finalRefineBarrier: "ifInFlight", - }, - runBeforeAgentStart: !options.skipPrePromptWork, - nextTurnContextTiming: "preparation", - preserveEmptyExtensionPrompt: false, - completionIncludesRetryChain: true, - }; - } - if (kind === "injected") { - return { - preparation: { - initialRefineBarrier: "always", - flushPendingBashBeforeValidation: true, - validateModelAndAuth: true, - awaitPendingModelSelection: true, - preTurnCompaction: "beforeModelSelection", - finalRefineBarrier: "ifInFlight", - }, - runBeforeAgentStart: true, - nextTurnContextTiming: "preparation", - preserveEmptyExtensionPrompt: true, - completionIncludesRetryChain: true, - }; - } - return { - preparation: { - initialRefineBarrier: "always", - flushPendingBashBeforeValidation: false, - validateModelAndAuth: false, - awaitPendingModelSelection: false, - preTurnCompaction: "skip", - finalRefineBarrier: "skip", - }, - runBeforeAgentStart: false, - nextTurnContextTiming: "skip", - preserveEmptyExtensionPrompt: false, - completionIncludesRetryChain: false, - }; - } - - private _createPreparedTurnAction( - schedule: SessionInputSchedule, - text: string, - images: ImageContent[] | undefined, - options: { - agentMessageId?: string; - queueKey?: string; - content?: (TextContent | ImageContent)[]; - message?: QueuedAgentMessage; - prefixMessages?: CustomMessage[]; - previewLabel?: string; - suppressAutonomousContinuation?: boolean; - resumeIfIdle?: boolean; - source?: InputSource | "internal"; - executionPolicy?: TurnExecutionPolicy; - queueVisible?: boolean; - acceptedAgentMessage?: boolean; - acceptedBeforeCompletion?: boolean; - }, - ): QueuedSessionAction { - const id = randomUUID(); - const content = options.content ?? this._buildPromptContent(text, images); - const message = - options.message ?? - ({ - role: "user", - content: content.map((block) => ({ ...block })), - timestamp: Date.now(), - } satisfies UserMessage); - const prefixMessages = options.prefixMessages?.map((prefix) => cloneCustomMessage(prefix)) ?? []; - const preview = options.previewLabel ? `${options.previewLabel}: ${text}` : undefined; - const payload: PreparedTurnPayload = { - kind: "turn", - text, - records: [ - ...prefixMessages.map((prefix) => this._createDeliveryRecord(id, "prefix", prefix)), - this._createDeliveryRecord(id, "primary", message), - ], - preview, - images: images?.map((image) => ({ ...image })), - content: content.map((block) => ({ ...block })), - customMessage: options.message?.role === "custom" ? cloneCustomMessage(options.message) : undefined, - executionPolicy: options.executionPolicy ?? this._turnExecutionPolicy("queued"), - queueVisible: options.queueVisible ?? true, - acceptedAgentMessage: options.acceptedAgentMessage ?? false, - acceptedBeforeCompletion: options.acceptedBeforeCompletion ?? false, - }; - return { - id, - source: options.source ?? "internal", - delivery: this._deliveryPolicy(schedule), - wake: - options.resumeIfIdle === true - ? "immediate" - : schedule === "steer" - ? "on_lower_boundary" - : "external_resume", - payload, - lifecycle: { state: "queued" }, - queueKey: options.queueKey, - agentMessageId: options.agentMessageId, - suppressAutonomousContinuation: options.suppressAutonomousContinuation, - }; - } - - private _createSessionCommandAction( - text: string, - command: SessionSlashCommand, - images: ImageContent[] | undefined, - schedule: SessionInputSchedule, - options: { - agentMessageId?: string; - source?: InputSource | "internal"; - } = {}, - ): QueuedSessionAction { - return { - id: randomUUID(), - source: options.source ?? "internal", - delivery: this._deliveryPolicy(schedule), - wake: "immediate", - payload: { kind: "session_command", text, command, images }, - lifecycle: { state: "queued" }, - agentMessageId: options.agentMessageId, - }; + private _createPreparedTurnAction(...args: Parameters): QueuedSessionAction { + return createPreparedTurnAction(...args); } private _coalescedFollowUpOwner(action: QueuedSessionAction): QueuedSessionAction | undefined { @@ -5506,12 +5052,12 @@ export class AgentSession { if (this._disposed || this._disposing) { throw new Error("Cannot admit a session action because the session is disposing or disposed."); } - if (this._sessionInputAdmissionPauses.size > 0) { + if (this._inputScheduler.admissionPaused) { throw new SessionInputAdmissionPausedError( "Cannot admit a session action while session input admission is paused.", ); } - if (this._sessionInputPumpSuspended) { + if (this._inputScheduler.suspended) { throw new Error("Cannot admit a session action while queued session input is suspended."); } } @@ -5532,7 +5078,7 @@ export class AgentSession { if (this._disposed || this._disposing) { throw new Error("Cannot admit a session action because the session is disposing or disposed."); } - if (this._sessionInputAdmissionPauses.size > 0) { + if (this._inputScheduler.admissionPaused) { throw new SessionInputAdmissionPausedError( "Cannot admit a session action while session input admission is paused.", ); @@ -5619,16 +5165,13 @@ export class AgentSession { bash: this.isBashRunning, refinementApply: this._refineInFlight !== undefined, branchMutation: this._branchSummaryOperation !== undefined, - schedulerPauseCount: this._queuedWorkPauses.size + (this._sessionInputPumpSuspended ? 1 : 0), + schedulerPauseCount: this._inputScheduler.queuedWorkPauseCount + (this._inputScheduler.suspended ? 1 : 0), disposing: this._disposed || this._disposing, }; } private _hasSelectableSessionInput(): boolean { - return ( - this._actionStore.queuedActions().length > 0 || - this._actionStore.activeActions().some((action) => action.lifecycle.state === "selected") - ); + return this._inputDispatcher.hasSelectableInput(); } get hasPendingSessionWork(): boolean { @@ -5644,180 +5187,11 @@ export class AgentSession { } get hasPendingAdmissionWaiters(): boolean { - return ( - this._sessionActionCommitOwner !== undefined || - this._pendingSessionActionFenceWaiters > 0 || - this._sessionInputCheckpointWaiters.size > 0 - ); + return this._commitFence.hasPendingWork || this._sessionInputCheckpointWaiters.size > 0; } private _scheduleSessionInputPump(): void { - if (this._sessionInputPumpSuspended || this._queuedWorkPauses.size > 0) return; - if (this._disposed || this._disposing || this._sessionInputPumpRequested || !this._hasSelectableSessionInput()) { - return; - } - this._sessionInputPumpRequested = true; - const epoch = this._sessionInputPumpEpoch; - const pump = async () => { - this._sessionInputPumpRequested = false; - await this._pumpSessionInputs(epoch); - }; - this._sessionInputPump = this._sessionInputPump.then(pump, pump); - this._sessionInputPump.catch(() => {}); - } - - private async _pumpSessionInputs(epoch: number): Promise { - let blocked = false; - try { - while (!this._disposed && !this._disposing && this._hasSelectableSessionInput()) { - await this.agent.waitForIdle(); - const preselected = this._actionStore - .activeActions() - .find((action) => action.lifecycle.state === "selected"); - if (epoch !== this._sessionInputPumpEpoch) { - if (preselected) { - this._actionStore.rollback(preselected); - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - } - return; - } - if (!this._hasCancelledDispatchCapture()) await this._agentEventQueue; - if (!preselected || preselected.payload.kind === "session_command") await this._waitForRefineIdle(); - const activity = this._runtimeActivity(); - const canSelectPreselectedTurn = - preselected?.payload.kind === "turn" && canSelectSessionAction({ ...activity, refinementApply: false }); - if ( - this._isSessionInputHandoffDeferred(epoch) || - (!canSelectPreselectedTurn && !canSelectSessionAction(activity)) - ) { - blocked = true; - this._notifySessionInputCheckpointChange(); - return; - } - const first = preselected ?? this._actionStore.selectFirst(); - if (!first) return; - if (first.payload.kind === "session_command") { - await this._executeSelectedSessionCommand(first, epoch); - return; - } - - const mode = first.delivery === "next_turn_boundary" ? this.steeringMode : this.followUpMode; - const actions: QueuedSessionAction[] = [first]; - while (!preselected && mode === "all") { - const next = this._actionStore.queuedActions(first.delivery)[0]; - if ( - !next || - next.payload.kind !== "turn" || - !turnExecutionPoliciesEqual(first.payload.executionPolicy, next.payload.executionPolicy) - ) { - break; - } - this._actionStore.selectFirst(); - actions.push(next); - } - if (epoch !== this._sessionInputPumpEpoch) { - for (const action of actions) this._actionStore.rollback(action); - return; - } - for (const action of actions) transitionSessionAction(action, { state: "preparing" }); - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - try { - await this._startPreparedTurnActions(actions, epoch); - for (const action of actions) { - if (action.lifecycle.state === "committing") { - const primary = primaryDeliveryRecord(action); - if (this.agent.state.messages.includes(primary.message)) { - primary.durable = true; - transitionSessionAction(action, { - state: "running", - execution: "agent_turn", - }); - } - } - if (action.lifecycle.state === "running") { - transitionSessionAction(action, { state: "completed" }); - this._actionStore.ticketFor(action).settleCompleted(); - this._settleAgentMessage(action.agentMessageId, "completion"); - } - } - } catch (error) { - const transcript = this.agent.state.messages; - const delivered = new Set(transcript); - const undelivered: QueuedSessionAction[] = []; - for (const action of actions) { - if (action.payload.kind !== "turn" || action.lifecycle.state === "cancelled") continue; - for (const record of action.payload.records) record.durable ||= delivered.has(record.message); - action.payload.records = action.payload.records.filter((record) => { - if (record.role === "prefix") return !record.durable; - if (record.role === "next_turn") return record.durable; - return true; - }); - if (!primaryDeliveryRecord(action).durable) undelivered.push(action); - } - if (this._isDeferredSessionInputError(error, epoch)) { - for (const action of undelivered) { - if (action.lifecycle.state === "committing") { - this._actionStore.rollback(action, { - dispatchSettled: true, - transcript, - }); - } else if (action.lifecycle.state === "preparing" || action.lifecycle.state === "selected") { - this._actionStore.rollback(action); - } - } - if (undelivered.length > 0) this._emitQueueUpdate(); - blocked = epoch !== this._sessionInputPumpEpoch || this._isBusyForSessionInput("pump"); - if (blocked) return; - continue; - } - const terminalError = this._asError(error); - for (const action of actions) { - if (action.lifecycle.state === "cancelled") continue; - if (action.lifecycle.state !== "completed" && action.lifecycle.state !== "failed") { - transitionSessionAction(action, { - state: "failed", - error: terminalError, - }); - } - const ticket = this._actionStore.ticketFor(action); - if (undelivered.includes(action)) { - ticket.rejectDelivered(terminalError); - this._settleAgentMessage(action.agentMessageId, "delivery", terminalError); - } - this._settleAgentMessage(action.agentMessageId, "completion", terminalError); - ticket.settleCompleted(terminalError); - } - if (actions.some((action) => action.payload.kind !== "turn" || action.payload.queueVisible)) { - this._surfaceSessionInputError(error); - } - } finally { - for (const action of actions) { - const retainedCancelledDispatch = - action.lifecycle.state === "cancelled" && - action.payload.kind === "turn" && - action.payload.captureRunMessages !== undefined; - if ( - !retainedCancelledDispatch && - (action.lifecycle.state === "completed" || - action.lifecycle.state === "failed" || - action.lifecycle.state === "cancelled") - ) { - this._durableRlmTerminalNoticeActionIds.delete(action.id); - this._actionStore.releaseTerminal(action); - } - } - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - } - if (epoch !== this._sessionInputPumpEpoch || blocked) return; - } - } finally { - if (!blocked && epoch === this._sessionInputPumpEpoch && this._hasSelectableSessionInput()) { - this._scheduleSessionInputPump(); - } - } + this._inputScheduler.schedule(); } private async _executeSelectedSessionCommand(action: QueuedSessionAction, epoch: number): Promise { @@ -5825,7 +5199,7 @@ export class AgentSession { const input = action.payload; const commitFence = await this._acquireSessionActionCommitFence(); try { - await this._sessionActionCommitContext.run(commitFence.owner, async () => { + await this._commitFence.run(commitFence, async () => { const isCancelled = () => action.lifecycle.state === "cancelled"; if (isCancelled()) return; await this._waitForRefineIdle(); @@ -5878,8 +5252,8 @@ export class AgentSession { externalBusy || this._disposed || this._disposing || - this._sessionInputPumpSuspended || - this._queuedWorkPauses.size > 0 || + this._inputScheduler.suspended || + this._inputScheduler.queuedWorkPauseCount > 0 || this._branchSummaryOperation !== undefined ); } @@ -5887,23 +5261,13 @@ export class AgentSession { } private _isSessionInputHandoffDeferred(epoch: number): boolean { - return epoch !== this._sessionInputPumpEpoch || this._isBusyForSessionInput("pump"); + return epoch !== this._inputScheduler.epoch || this._isBusyForSessionInput("pump"); } private _asError(error: unknown): Error { return error instanceof Error ? error : new Error(String(error)); } - private _isDeferredSessionInputError(error: unknown, epoch: number): boolean { - if (error instanceof DeferredSessionInputError) return true; - if (epoch !== this._sessionInputPumpEpoch) return true; - if (this._isBusyForSessionInput("pump")) { - this._surfaceSessionInputError(error); - return true; - } - return false; - } - private _surfaceSessionInputError(error: unknown): void { const normalized = this._asError(error); try { @@ -5939,7 +5303,7 @@ export class AgentSession { nextTurnMessages = []; }; try { - const preparedTurn = await this._prepareForCommit(executionPolicy.preparation, { + const preparedTurn = await this._turnPreparer.prepare(executionPolicy.preparation, { afterValidation: () => { if (this._isSessionInputHandoffDeferred(epoch)) { throw new DeferredSessionInputError("Session input paused before preflight"); @@ -5990,7 +5354,7 @@ export class AgentSession { const commitFence = await this._acquireSessionActionCommitFence(); let promptPromise: Promise; try { - promptPromise = this._sessionActionCommitContext.run(commitFence.owner, () => { + promptPromise = this._commitFence.run(commitFence, () => { if ( this._isSessionInputHandoffDeferred(epoch) || this.isStreaming || @@ -6011,7 +5375,7 @@ export class AgentSession { } } const contextRecords = nextTurnMessages.map((message) => - this._createDeliveryRecord(turns[0].id, "next_turn", message), + createDeliveryRecord(turns[0].id, "next_turn", message), ); const firstPrimaryIndex = turns[0].payload.records.indexOf(primaryDeliveryRecord(turns[0])); turns[0].payload.records.splice(firstPrimaryIndex, 0, ...contextRecords); @@ -6217,7 +5581,7 @@ export class AgentSession { }); } } else if (options?.triggerTurn) { - if (!this._sessionInputSuspendedForUpdateRestart) this._resumeSessionInputAdmission(); + if (!this._inputScheduler.suspendedForUpdateRestart) this._resumeSessionInputAdmission(); const admissionFence = await this._acquireDirectTurnAdmissionFence(); try { const normalized = normalizeMessageContent(message.content); @@ -6225,7 +5589,7 @@ export class AgentSession { const action = this._createPreparedTurnAction("followUp", normalized.text, normalized.images, { message: appMessage, resumeIfIdle: true, - executionPolicy: this._turnExecutionPolicy("customTrigger"), + executionPolicy: createTurnExecutionPolicy("customTrigger"), queueVisible: false, }); const result = this._admitSessionInput(action, { immediatelyEligible }); @@ -6292,7 +5656,7 @@ export class AgentSession { .clearableActions() .filter((action) => action.payload.kind === "session_command" || action.payload.queueVisible); if (clearable.some((action) => action.payload.kind === "turn" && action.lifecycle.state === "preparing")) { - this._sessionInputPumpEpoch++; + this._inputScheduler.invalidatePreparation(); } const steering = clearable .filter((action) => action.delivery === "next_turn_boundary") @@ -6478,7 +5842,7 @@ export class AgentSession { } get isQueuedWorkSuspended(): boolean { - return this._sessionInputPumpSuspended; + return this._inputScheduler.suspended; } get isSessionActive(): boolean { @@ -6704,54 +6068,32 @@ export class AgentSession { } acquireSessionInputPause(): { release(): void } { - const token = Symbol("session-input-admission-pause"); - this._sessionInputAdmissionPauses.add(token); - this._sessionInputPumpRequested = false; - this._sessionInputPumpEpoch++; - let released = false; - return { - release: () => { - if (released) return; - released = true; - this._sessionInputAdmissionPauses.delete(token); - this._sessionInputPumpEpoch++; - this._notifySessionInputCheckpointChange(); - this._flushDeferredRlmTerminalNotices(); - this._maybeResumeGoalContinuationAfterRlmWork(); - this._scheduleSessionInputPump(); - }, - }; + return this._inputScheduler.acquireAdmissionPause(() => { + this._notifySessionInputCheckpointChange(); + this._flushDeferredRlmTerminalNotices(); + this._maybeResumeGoalContinuationAfterRlmWork(); + this._scheduleSessionInputPump(); + }); } acquireQueuedWorkPause(): { release(): void } { - const token = Symbol("queued-work-pause"); - this._queuedWorkPauses.add(token); - this._sessionInputPumpRequested = false; - this._sessionInputPumpEpoch++; - let released = false; - return { - release: () => { - if (released) return; - released = true; - this._queuedWorkPauses.delete(token); - this._notifySessionInputCheckpointChange(); - this._flushDeferredRlmTerminalNotices(); - this._scheduleSessionInputPump(); - }, - }; + return this._inputScheduler.acquireQueuedWorkPause(() => { + this._notifySessionInputCheckpointChange(); + this._flushDeferredRlmTerminalNotices(); + this._scheduleSessionInputPump(); + }); } private async _acquireDirectTurnAdmissionFence(signal?: AbortSignal): Promise<{ owner: symbol; release(): void }> { - const inheritedOwner = this._sessionActionCommitContext.getStore(); - if (inheritedOwner !== undefined && inheritedOwner === this._sessionActionCommitOwner) { + if (this._commitFence.isHeldByCurrentContext) { this._assertSessionActionAdmissionAvailable(); return this._acquireSessionActionCommitFence(signal); } - const disposeSignal = this._sessionActionCommitDisposeAbortController.signal; + const disposeSignal = this._commitFence.disposeSignal; const waitSignal = signal ? AbortSignal.any([signal, disposeSignal]) : disposeSignal; while (true) { this._assertSessionActionAdmissionAvailable(); - if (this._queuedWorkPauses.size > 0) { + if (this._inputScheduler.queuedWorkPauseCount > 0) { let wake = () => {}; const pauseReleased = new Promise((resolve) => { wake = resolve; @@ -6771,7 +6113,7 @@ export class AgentSession { } const fence = await this._acquireSessionActionCommitFence(signal); try { - if (this._queuedWorkPauses.size === 0) { + if (this._inputScheduler.queuedWorkPauseCount === 0) { this._assertSessionActionAdmissionAvailable(); return fence; } @@ -6783,50 +6125,12 @@ export class AgentSession { } } - private async _acquireSessionActionCommitFence(signal?: AbortSignal): Promise<{ owner: symbol; release(): void }> { - const inheritedOwner = this._sessionActionCommitContext.getStore(); - if (inheritedOwner !== undefined && inheritedOwner === this._sessionActionCommitOwner) { - return { owner: inheritedOwner, release: () => {} }; - } - const previous = this._sessionActionCommitTail; - let resolve = () => {}; - this._sessionActionCommitTail = new Promise((release) => { - resolve = release; - }); - const disposeSignal = this._sessionActionCommitDisposeAbortController.signal; - const waitSignal = signal ? AbortSignal.any([signal, disposeSignal]) : disposeSignal; - this._pendingSessionActionFenceWaiters++; - try { - await waitForPromiseOrAbort(previous, waitSignal, "Update restart preparation cancelled"); - } catch (error) { - this._pendingSessionActionFenceWaiters--; - // A cancelled waiter remains in the FIFO chain until its predecessor releases. - void previous.then(resolve, resolve); - if (disposeSignal.aborted) { - throw new Error("Cannot admit a session action because the session is disposing or disposed."); - } - throw error; - } - const owner = Symbol("session-action-commit"); - this._sessionActionCommitOwner = owner; - this._pendingSessionActionFenceWaiters--; - let released = false; - return { - owner, - release: () => { - if (released) return; - released = true; - if (this._sessionActionCommitOwner === owner) this._sessionActionCommitOwner = undefined; - resolve(); - }, - }; + private _acquireSessionActionCommitFence(signal?: AbortSignal): Promise { + return this._commitFence.acquire(signal); } private _resumeSessionInputAdmission(): void { - if (!this._sessionInputPumpSuspended) return; - this._sessionInputPumpSuspended = false; - this._sessionInputSuspendedForUpdateRestart = false; - this._sessionInputPumpEpoch++; + if (!this._inputScheduler.resume()) return; this._notifySessionInputCheckpointChange(); this._flushDeferredRlmTerminalNotices(); } @@ -6839,12 +6143,8 @@ export class AgentSession { return this._hasSelectableSessionInput(); } - async waitForSessionInputIdle(): Promise { - while (true) { - const pump = this._sessionInputPump; - await pump; - if (pump === this._sessionInputPump && !this._sessionInputPumpRequested) return; - } + waitForSessionInputIdle(): Promise { + return this._inputScheduler.waitForIdle(); } async waitForIdle(): Promise { @@ -6860,7 +6160,7 @@ export class AgentSession { private async _waitForIdleOrSettlement(settlement?: PostCompactionContinuationSettlement): Promise { while (settlement === undefined || this._postCompactionContinuationSettlement === settlement) { if (this._actionStore.queuedActions().length > 0) { - if (this._sessionInputPumpSuspended || this._queuedWorkPauses.size > 0) { + if (this._inputScheduler.suspended || this._inputScheduler.queuedWorkPauseCount > 0) { let wake = () => {}; const changed = new Promise((resolve) => { wake = resolve; @@ -6875,15 +6175,15 @@ export class AgentSession { } this._scheduleSessionInputPump(); } - const pump = this._sessionInputPump; + const pump = this._inputScheduler.pendingPump; await pump; await this.agent.waitForIdle(); const agentEventQueue = this._agentEventQueue; await agentEventQueue; if ( - pump === this._sessionInputPump && + pump === this._inputScheduler.pendingPump && agentEventQueue === this._agentEventQueue && - !this._sessionInputPumpRequested && + !this._inputScheduler.requested && !this.agent.state.isStreaming && this.unfinishedActionCount === 0 ) { @@ -6953,10 +6253,7 @@ export class AgentSession { if (run.status === "cancelled") this._abandonRlmRunForQuiescence(run); } for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); - this._sessionInputPumpRequested = false; - this._sessionInputPumpEpoch++; - this._sessionInputPumpSuspended = true; - this._sessionInputSuspendedForUpdateRestart = false; + this._inputScheduler.suspend("abort"); this._demoteRlmTerminalNoticeActions(); this._cancelSessionActions( (action) => @@ -6998,10 +6295,7 @@ export class AgentSession { abortForUpdateRestart(): void { // Cancel scheduled pumps and suspend new ones: queued inputs must survive // into the restart manifest instead of starting a turn during teardown. - this._sessionInputPumpRequested = false; - this._sessionInputPumpEpoch++; - this._sessionInputPumpSuspended = true; - this._sessionInputSuspendedForUpdateRestart = true; + this._inputScheduler.suspend("update-restart"); this._cancelPostCompactionContinue(); this.abortRetry(); for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); @@ -7796,7 +7090,10 @@ export class AgentSession { } private async _waitForQueuedWorkResume(settlement: PostCompactionContinuationSettlement): Promise { - while (this._queuedWorkPauses.size > 0 && this._postCompactionContinuationSettlement === settlement) { + while ( + this._inputScheduler.queuedWorkPauseCount > 0 && + this._postCompactionContinuationSettlement === settlement + ) { let resume = () => {}; const resumed = new Promise((resolve) => { resume = resolve; @@ -7835,7 +7132,7 @@ export class AgentSession { return; } - if (this._queuedWorkPauses.size > 0 || this._compactionOperation || this._refineInFlight) { + if (this._inputScheduler.queuedWorkPauseCount > 0 || this._compactionOperation || this._refineInFlight) { continue; } @@ -7845,7 +7142,7 @@ export class AgentSession { this._scheduleAutoRefineAfterAgentEnd(); return; } - if (this.unfinishedActionCount > 0 || this._sessionInputPumpRequested) { + if (this.unfinishedActionCount > 0 || this._inputScheduler.requested) { this._scheduleSessionInputPump(); waitForSessionInput = true; } else { @@ -9384,7 +8681,7 @@ export class AgentSession { })), "bash.completed": createAsyncBashCompletionHostHandler(async (details) => { const message = createAsyncBashCompletionMessage(details); - const disposeSignal = this._sessionActionCommitDisposeAbortController.signal; + const disposeSignal = this._commitFence.disposeSignal; while (true) { let admissionCommitted = false; try { @@ -9401,7 +8698,7 @@ export class AgentSession { return; } catch (error) { if (admissionCommitted || !(error instanceof SessionInputAdmissionPausedError)) throw error; - while (this._sessionInputAdmissionPauses.size > 0 && !disposeSignal.aborted) { + while (this._inputScheduler.admissionPaused && !disposeSignal.aborted) { await this._waitForSessionActivityChange(disposeSignal); } } @@ -9785,7 +9082,7 @@ export class AgentSession { return false; } run.status = "cancelled"; - if (this._sessionInputPumpSuspended) this._abandonRlmRunForQuiescence(run); + if (this._inputScheduler.suspended) this._abandonRlmRunForQuiescence(run); run.error = reason; run.publication.reject(new Error(reason)); run.abort(); @@ -10478,7 +9775,7 @@ export class AgentSession { const run = session._activeRlmChildRuns.get(childId); if (run) { if (run.status !== "running" && run.status !== "queued" && !run.settled) { - if (session._sessionInputPumpSuspended) session._abandonRlmRunForQuiescence(run); + if (session._inputScheduler.suspended) session._abandonRlmRunForQuiescence(run); else run.suppressTerminalNotice = true; return true; } @@ -11157,275 +10454,16 @@ export class AgentSession { return this._startRlmChildRun(prompt, kwargs, spawnCode); } - private _isRetryableError(message: AssistantMessage): boolean { - if (message.stopReason !== "error" || !message.errorMessage) return false; - - const contextWindow = this.model?.contextWindow ?? 0; - if (isContextOverflow(message, contextWindow)) return false; - - if (this._isFauxProviderQueueExhausted(message)) { - return false; - } - - if (this._isAgentLifecycleFailure(message)) { - return false; - } - - if (this._isStructuredPermanentProviderRetryExhausted(message)) { - return false; - } - - return true; - } - - private _isFauxProviderQueueExhausted(message: AssistantMessage): boolean { - return isFauxProviderQueueExhausted(message); - } - - private _isAgentLifecycleFailure(message: AssistantMessage): boolean { - return isAgentLifecycleFailure(message); - } - - private _getProviderStreamFailureKind(message: AssistantMessage): string | undefined { - return providerStreamFailureKind(message); - } - - private _isStructuredPermanentProviderRetryExhausted(message: AssistantMessage): boolean { - return isPermanentProviderFailureKind(this._getProviderStreamFailureKind(message), this._retryAttempt); - } - - private _isConcreteProviderAuthFailure(message: AssistantMessage): boolean { - if (message.stopReason !== "error" || !message.errorMessage) return false; - // Only the provider's structured classification counts as an auth failure. - return this._getProviderStreamFailureKind(message) === "auth"; - } - - private _captureRetryAuthFailureSource(message: AssistantMessage): AuthSourceToken | undefined { - const token = this._modelRegistry.getCurrentProviderAuthSourceToken(message.provider); - if (!token) { - return undefined; - } - if ( - !this._retryAuthFailureSources.some( - (existing) => - existing.provider === token.provider && - existing.source === token.source && - existing.identityFingerprint === token.identityFingerprint && - existing.valueFingerprint === token.valueFingerprint, - ) - ) { - this._retryAuthFailureSources.push(token); - } - return token; - } - - private _markProviderAuthStale(message: AssistantMessage, authSourceTokens?: readonly AuthSourceToken[]): boolean { - if (authSourceTokens && authSourceTokens.length > 0) { - let marked = false; - for (const token of authSourceTokens) { - marked = this._modelRegistry.markProviderAuthSourceStale(token) || marked; - } - if (marked) { - this._emit({ - type: "auth_stale", - provider: message.provider, - sourceTokens: authSourceTokens, - }); - } - return marked; - } - const marked = this._modelRegistry.markProviderAuthStale(message.provider); - if (marked) { - this._emit({ type: "auth_stale", provider: message.provider }); - } - return marked; - } - - private _markProviderAuthStaleForRetryFailure( - message: AssistantMessage, - options?: { - markAuthStaleOnFailure?: boolean; - authSourceTokens?: readonly AuthSourceToken[]; - }, - ): boolean { - const authSourceTokens = - this._retryAuthFailureSources.length > 0 ? this._retryAuthFailureSources : options?.authSourceTokens; - if ((authSourceTokens?.length ?? 0) > 0 || options?.markAuthStaleOnFailure) { - const marked = this._markProviderAuthStale(message, authSourceTokens); - if (marked && message.errorMessage) { - message.errorMessage = addLoginGuidanceToAuthError(message.errorMessage); - } - return marked; - } - return false; - } - - private _finishActiveRetryWithFailure(message: AssistantMessage): void { - if (this._retryAttempt === 0) { - return; - } - this._markProviderAuthStaleForRetryFailure(message); - this._emit({ - type: "auto_retry_end", - success: false, - attempt: this._retryAttempt, - finalError: message.errorMessage, - }); - this._retryAttempt = 0; - this._retryAuthFailureSources = []; - } - - private async _handleRetryableError( - message: AssistantMessage, - options?: { - markAuthStaleOnFailure?: boolean; - authSourceTokens?: readonly AuthSourceToken[]; - }, - ): Promise { - const settings = this.settingsManager.getRetrySettings(); - if (!settings.enabled) { - this._markProviderAuthStaleForRetryFailure(message, options); - this._retryAuthFailureSources = []; - this._resolveRetry(); - return false; - } - - if (!this._retryPromise) { - this._retryPromise = new Promise((resolve) => { - this._retryResolve = resolve; - }); - } - - this._retryAttempt++; - - if (this._retryAttempt > settings.maxRetries) { - this._markProviderAuthStaleForRetryFailure(message, options); - this._emit({ - type: "auto_retry_end", - success: false, - attempt: this._retryAttempt - 1, - finalError: message.errorMessage, - }); - this._retryAttempt = 0; - this._retryAuthFailureSources = []; - this._resolveRetry(); // Resolve so waitForRetry() completes - return false; - } - - // Server-requested waits are honored, capped by retry.provider.maxRetryDelayMs (0 disables). - const maxRetryDelayMs = this.settingsManager.getProviderRetrySettings().maxRetryDelayMs; - const delay = providerRetryDelay(this._retryAttempt, providerStreamFailureRetryAfterMs(message), { - baseDelayMs: settings.baseDelayMs, - maxRetryDelayMs, - }); - if (delay.kind === "exceeds-cap") { - this._markProviderAuthStaleForRetryFailure(message, options); - this._emit({ - type: "auto_retry_end", - success: false, - attempt: this._retryAttempt - 1, - finalError: `Provider requested a ${Math.ceil(delay.retryAfterMs / 1000)}s wait before retrying (above retry.provider.maxRetryDelayMs=${maxRetryDelayMs}ms): ${message.errorMessage || "unknown error"}`, - }); - this._retryAttempt = 0; - this._retryAuthFailureSources = []; - this._resolveRetry(); - return false; - } - - const delayMs = delay.delayMs; - // Park now: the retry re-issues the failed call and must reuse its Idempotency-Key. - // Payload hooks mutate the wire body after the hash point, so reuse is forfeited. - if (!this._extensionRunner.hasHandlers("before_provider_request")) { - this._semanticEdges.prepareTurnRetry(); - } - - this._emit({ - type: "auto_retry_start", - attempt: this._retryAttempt, - maxAttempts: settings.maxRetries, - delayMs, - errorMessage: message.errorMessage || "Unknown error", - }); - - const messages = this.agent.state.messages; - if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { - this.agent.state.messages = messages.slice(0, -1); - } - - this._retryAbortController = new AbortController(); - try { - await sleep(delayMs, this._retryAbortController.signal); - } catch { - const attempt = this._retryAttempt; - this._markProviderAuthStaleForRetryFailure(message, options); - this._retryAttempt = 0; - this._retryAbortController = undefined; - this._emit({ - type: "auto_retry_end", - success: false, - attempt, - finalError: "Retry cancelled", - }); - this._resolveRetry(); - this._retryAuthFailureSources = []; - return false; - } - this._retryAbortController = undefined; - - const retryGeneration = this._retryGeneration; - setTimeout(() => { - this.agent.continue().catch((error: unknown) => { - // A continue that never starts must still resolve the retry (else isRetrying - // sticks forever) — unless a newer retry owns the state by now. - if (this._retryGeneration !== retryGeneration || !this.isRetrying) return; - this._markProviderAuthStaleForRetryFailure(message, options); - const attempt = this._retryAttempt; - this._retryAttempt = 0; - this._retryAuthFailureSources = []; - this._emit({ - type: "auto_retry_end", - success: false, - attempt, - finalError: error instanceof Error ? error.message : String(error), - }); - this._resolveRetry(); - }); - }, 0); - - return true; - } - abortRetry(): void { - if (this._retryAbortController) { - this._retryAbortController.abort(); - return; - } - if (this._retryAttempt > 0) { - this._autoCompactionAbortController?.abort(); - this._cancelPostCompactionContinue(); - this._emit({ - type: "auto_retry_end", - success: false, - attempt: this._retryAttempt, - finalError: "Retry cancelled", - }); - this._retryAttempt = 0; - } - this._retryAuthFailureSources = []; - this._resolveRetry(); + this._retry.abortRetry(); } - private async waitForRetry(): Promise { - if (!this._retryPromise) { - return; - } - - await this._retryPromise; - await this.agent.waitForIdle(); + private waitForRetry(): Promise { + return this._retry.waitForRetry(); } get isRetrying(): boolean { - return this._retryPromise !== undefined; + return this._retry.isRetrying; } get hasAcceptedPromptInFlight(): boolean { @@ -11447,99 +10485,14 @@ export class AgentSession { this.settingsManager.setRetryEnabled(enabled); } - /** - * Execute a bash command. - * Adds result to agent context and session. - * @param command The bash command to execute - * @param onChunk Optional streaming callback for output - * @param options.excludeFromContext If true, command output won't be sent to LLM (!! prefix) - * @param options.operations Custom BashOperations for remote execution - */ - async executeBash( - command: string, - onChunk?: (chunk: string) => void, - options?: { - excludeFromContext?: boolean; - operations?: BashOperations; - transient?: boolean; - }, - ): Promise { - // Each invocation owns its controller so abortBash reaches every in-flight command. - const abortController = new AbortController(); - this._bashAbortControllers.add(abortController); - - const prefix = this.settingsManager.getShellCommandPrefix(); - const shellPath = this.settingsManager.getShellPath(); - const resolvedCommand = prefix ? `${prefix}\n${command}` : command; - - try { - const result = await executeBashWithOperations( - resolvedCommand, - this.sessionManager.getCwd(), - options?.operations ?? createLocalBashOperations({ shellPath }), - { - onChunk, - signal: abortController.signal, - }, - ); - - if (!options?.transient) { - this.recordBashResult(command, result, options); - } - return result; - } finally { - this._bashAbortControllers.delete(abortController); - this._notifySessionInputCheckpointChange(); - } + /** Execute a shell command and record its result unless transient. */ + executeBash(command: string, onChunk?: (chunk: string) => void, options?: ExecuteBashOptions): Promise { + return this._bash.executeBash(command, onChunk, options); } - /** - * Run a user-initiated bash command (! / !! prefix), emitting bash_start, - * bash_output, and bash_end session events so any attached client can render - * streaming output. Extensions can intercept execution via the user_bash event. - * Execution failures are reported through bash_end rather than a rejected promise; - * only the already-running guard and extension dispatch errors reject. - * @param command The bash command to execute - * @param options.excludeFromContext If true, command output won't be sent to LLM (!! prefix) - */ - async runUserBash( - command: string, - options?: { - excludeFromContext?: boolean; - transient?: boolean; - runId?: string; - }, - ): Promise { - if (this.isBashRunning) { - throw new Error("A bash command is already running"); - } - // Claim the bash slot synchronously: isBashRunning is otherwise false until - // executeBash installs its abort controller, which would let a second command - // slip through during the user_bash extension dispatch below. - this._userBashRunning = true; - this._userBashAbortRequested = false; - // Echoed on bash_start/bash_end so the requesting client can tell its own - // run apart from other clients' runs broadcast on the same session. - const identity = { - ...(options?.transient ? { transient: true } : {}), - ...(options?.runId !== undefined ? { runId: options.runId } : {}), - }; - let end: UserBashEndDetails; - try { - end = await this.runUserBashLocked( - command, - options?.excludeFromContext ?? false, - options?.transient ?? false, - identity, - ); - } finally { - this._userBashRunning = false; - this._notifySessionInputCheckpointChange(); - } - // Emitted after the slot is released so clients never observe a bash_end - // while the session still rejects new commands as already running. - this._emit({ type: "bash_end", ...end, ...identity }); - void this._drainQueuedMessagesAfterBash().catch(() => undefined); + /** Run ! / !! input with extension interception and bash lifecycle events. */ + runUserBash(command: string, options?: RunUserBashOptions): Promise { + return this._bash.runUserBash(command, options); } private async _drainQueuedMessagesAfterBash(): Promise { @@ -11547,150 +10500,25 @@ export class AgentSession { this._scheduleSessionInputPump(); } - private async runUserBashLocked( - command: string, - excludeFromContext: boolean, - transient: boolean, - identity: { transient?: boolean; runId?: string }, - ): Promise { - const eventResult = await this._extensionRunner.emitUserBash({ - type: "user_bash", - command, - excludeFromContext, - cwd: this.sessionManager.getCwd(), - }); - - // Transient runs (side-conversation bash) live only in their pane: they - // are never recorded, so reloads and rebuilds cannot resurface them. - const record = transient - ? () => {} - : (result: BashResult) => this.recordBashResult(command, result, { excludeFromContext }); - - this._emit({ - type: "bash_start", - command, - excludeFromContext, - ...identity, - }); - try { - // If an extension returned a full result, surface it without executing - if (eventResult?.result) { - const result = eventResult.result; - if (result.output) { - this._emit({ type: "bash_output", chunk: result.output }); - } - record(result); - return { - exitCode: result.exitCode, - cancelled: result.cancelled, - truncated: result.truncated, - fullOutputPath: result.fullOutputPath, - }; - } - - // An abort that arrived before the process spawned (during extension - // dispatch) has no abort controller to act on; honor it here instead. - if (this._userBashAbortRequested) { - record({ - output: "", - exitCode: undefined, - cancelled: true, - truncated: false, - }); - return { exitCode: undefined, cancelled: true, truncated: false }; - } - - const result = await this.executeBash(command, (chunk) => this._emit({ type: "bash_output", chunk }), { - excludeFromContext, - operations: eventResult?.operations, - transient, - }); - return { - exitCode: result.exitCode, - cancelled: result.cancelled, - truncated: result.truncated, - fullOutputPath: result.fullOutputPath, - }; - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error); - // Persist the failure like every other outcome so replayed transcripts - // and the LLM context reflect that the command did not run. - record({ - output: `bash failed: ${errorMessage}`, - exitCode: undefined, - cancelled: false, - truncated: false, - }); - return { - exitCode: undefined, - cancelled: false, - truncated: false, - errorMessage, - }; - } - } - recordBashResult(command: string, result: BashResult, options?: { excludeFromContext?: boolean }): void { - const bashMessage: BashExecutionMessage = { - role: "bashExecution", - command, - output: result.output, - exitCode: result.exitCode, - cancelled: result.cancelled, - truncated: result.truncated, - fullOutputPath: result.fullOutputPath, - timestamp: Date.now(), - excludeFromContext: options?.excludeFromContext, - }; - - // If agent is streaming, defer adding to avoid breaking tool_use/tool_result ordering - if (this.isStreaming) { - this._pendingBashMessages.push(bashMessage); - } else { - this.agent.state.messages.push(bashMessage); - - this.sessionManager.appendMessage(bashMessage); - } + this._bash.recordBashResult(command, result, options); } - /** - * Cancel running bash command. - */ + /** Cancel every in-flight shell command, including pending extension dispatch. */ abortBash(): void { - // A user bash command may not have spawned yet (extension dispatch in - // progress); flag the request so runUserBash cancels before executing. - // runUserBash clears the flag at each start, so a stale flag is harmless. - if (this._userBashRunning) { - this._userBashAbortRequested = true; - } - for (const controller of this._bashAbortControllers) { - controller.abort(); - } + this._bash.abortBash(); } get isBashRunning(): boolean { - return this._bashAbortControllers.size > 0 || this._userBashRunning; + return this._bash.isBashRunning; } - /** Whether there are pending bash messages waiting to be flushed */ get hasPendingBashMessages(): boolean { - return this._pendingBashMessages.length > 0; + return this._bash.hasPendingBashMessages; } - /** - * Flush pending bash messages to agent state and session. - * Called after agent turn completes to maintain proper message ordering. - */ private _flushPendingBashMessages(): void { - if (this._pendingBashMessages.length === 0) return; - - for (const bashMessage of this._pendingBashMessages) { - this.agent.state.messages.push(bashMessage); - - this.sessionManager.appendMessage(bashMessage); - } - - this._pendingBashMessages = []; + this._bash.flushPendingMessages(); } getRlmMaxDepthStatus(): RlmMaxDepthStatus { @@ -11805,7 +10633,7 @@ export class AgentSession { try { // Branch navigation and turn dispatch mutate the same transcript leaf. commitFence = await this._acquireSessionActionCommitFence(); - return await this._sessionActionCommitContext.run(commitFence.owner, async () => { + return await this._commitFence.run(commitFence, async () => { await this.agent.waitForIdle(); await this._agentEventQueue; return this._navigateTreeUnderPause(targetId, targetEntry, options); @@ -12289,10 +11117,8 @@ export class AgentSession { return text.trim() || undefined; } - // ========================================================================= - // Extension System - // ========================================================================= - + // ================================================================== // Extension System + // ================================================================== createReplacedSessionContext(): ReplacedSessionContext { const context = Object.defineProperties( {}, diff --git a/packages/coding-agent/src/session/bash.ts b/packages/coding-agent/src/session/bash.ts new file mode 100644 index 0000000000..ee998cbc60 --- /dev/null +++ b/packages/coding-agent/src/session/bash.ts @@ -0,0 +1,288 @@ +import { type BashResult, executeBashWithOperations } from "../core/bash-executor.js"; +import type { UserBashEvent, UserBashEventResult } from "../core/extensions/types.js"; +import type { BashExecutionMessage } from "../core/messages.js"; +import { type BashOperations, createLocalBashOperations } from "../core/tools/bash.js"; + +export interface ExecuteBashOptions { + excludeFromContext?: boolean; + operations?: BashOperations; + transient?: boolean; +} + +export interface RunUserBashOptions { + excludeFromContext?: boolean; + transient?: boolean; + runId?: string; +} + +type UserBashEndDetails = { + exitCode: number | undefined; + cancelled: boolean; + truncated: boolean; + fullOutputPath?: string; + errorMessage?: string; +}; + +export type SessionBashEvent = + | { + type: "bash_start"; + command: string; + excludeFromContext: boolean; + transient?: boolean; + runId?: string; + } + | { type: "bash_output"; chunk: string } + | ({ type: "bash_end"; transient?: boolean; runId?: string } & UserBashEndDetails); + +export interface SessionBashHost { + getCwd(): string; + getShellCommandPrefix(): string | undefined; + getShellPath(): string | undefined; + isStreaming(): boolean; + intercept(event: UserBashEvent): Promise; + emit(event: SessionBashEvent): void; + appendMessage(message: BashExecutionMessage): void; + onStateChange(): void; + onUserBashEnd(): Promise; + executeBash(command: string, onChunk?: (chunk: string) => void, options?: ExecuteBashOptions): Promise; + recordBashResult(command: string, result: BashResult, options?: { excludeFromContext?: boolean }): void; +} + +export class SessionBash { + private _bashAbortControllers = new Set(); + private _userBashRunning = false; + private _userBashAbortRequested = false; + private _pendingBashMessages: BashExecutionMessage[] = []; + + constructor(private readonly _host: SessionBashHost) {} + + /** + * Execute a bash command. + * Adds result to agent context and session. + * @param command The bash command to execute + * @param onChunk Optional streaming callback for output + * @param options.excludeFromContext If true, command output won't be sent to LLM (!! prefix) + * @param options.operations Custom BashOperations for remote execution + */ + async executeBash( + command: string, + onChunk?: (chunk: string) => void, + options?: ExecuteBashOptions, + ): Promise { + // Each invocation owns its controller so abortBash reaches every in-flight command. + const abortController = new AbortController(); + this._bashAbortControllers.add(abortController); + + const prefix = this._host.getShellCommandPrefix(); + const shellPath = this._host.getShellPath(); + const resolvedCommand = prefix ? `${prefix}\n${command}` : command; + + try { + const result = await executeBashWithOperations( + resolvedCommand, + this._host.getCwd(), + options?.operations ?? createLocalBashOperations({ shellPath }), + { + onChunk, + signal: abortController.signal, + }, + ); + + if (!options?.transient) { + this._host.recordBashResult(command, result, options); + } + return result; + } finally { + this._bashAbortControllers.delete(abortController); + this._host.onStateChange(); + } + } + + /** + * Run a user-initiated bash command (! / !! prefix), emitting bash_start, + * bash_output, and bash_end session events so any attached client can render + * streaming output. Extensions can intercept execution via the user_bash event. + * Execution failures are reported through bash_end rather than a rejected promise; + * only the already-running guard and extension dispatch errors reject. + * @param command The bash command to execute + * @param options.excludeFromContext If true, command output won't be sent to LLM (!! prefix) + */ + async runUserBash(command: string, options?: RunUserBashOptions): Promise { + if (this.isBashRunning) { + throw new Error("A bash command is already running"); + } + // Claim the bash slot synchronously: isBashRunning is otherwise false until + // executeBash installs its abort controller, which would let a second command + // slip through during the user_bash extension dispatch below. + this._userBashRunning = true; + this._userBashAbortRequested = false; + // Echoed on bash_start/bash_end so the requesting client can tell its own + // run apart from other clients' runs broadcast on the same session. + const identity = { + ...(options?.transient ? { transient: true } : {}), + ...(options?.runId !== undefined ? { runId: options.runId } : {}), + }; + let end: UserBashEndDetails; + try { + end = await this.runUserBashLocked( + command, + options?.excludeFromContext ?? false, + options?.transient ?? false, + identity, + ); + } finally { + this._userBashRunning = false; + this._host.onStateChange(); + } + // Emitted after the slot is released so clients never observe a bash_end + // while the session still rejects new commands as already running. + this._host.emit({ type: "bash_end", ...end, ...identity }); + void this._host.onUserBashEnd().catch(() => undefined); + } + + private async runUserBashLocked( + command: string, + excludeFromContext: boolean, + transient: boolean, + identity: { transient?: boolean; runId?: string }, + ): Promise { + const eventResult = await this._host.intercept({ + type: "user_bash", + command, + excludeFromContext, + cwd: this._host.getCwd(), + }); + + // Transient runs (side-conversation bash) live only in their pane: they + // are never recorded, so reloads and rebuilds cannot resurface them. + const record = transient + ? () => {} + : (result: BashResult) => this._host.recordBashResult(command, result, { excludeFromContext }); + + this._host.emit({ + type: "bash_start", + command, + excludeFromContext, + ...identity, + }); + try { + // If an extension returned a full result, surface it without executing + if (eventResult?.result) { + const result = eventResult.result; + if (result.output) { + this._host.emit({ type: "bash_output", chunk: result.output }); + } + record(result); + return { + exitCode: result.exitCode, + cancelled: result.cancelled, + truncated: result.truncated, + fullOutputPath: result.fullOutputPath, + }; + } + + // An abort that arrived before the process spawned (during extension + // dispatch) has no abort controller to act on; honor it here instead. + if (this._userBashAbortRequested) { + record({ + output: "", + exitCode: undefined, + cancelled: true, + truncated: false, + }); + return { exitCode: undefined, cancelled: true, truncated: false }; + } + + const result = await this._host.executeBash( + command, + (chunk) => this._host.emit({ type: "bash_output", chunk }), + { + excludeFromContext, + operations: eventResult?.operations, + transient, + }, + ); + return { + exitCode: result.exitCode, + cancelled: result.cancelled, + truncated: result.truncated, + fullOutputPath: result.fullOutputPath, + }; + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + // Persist the failure like every other outcome so replayed transcripts + // and the LLM context reflect that the command did not run. + record({ + output: `bash failed: ${errorMessage}`, + exitCode: undefined, + cancelled: false, + truncated: false, + }); + return { + exitCode: undefined, + cancelled: false, + truncated: false, + errorMessage, + }; + } + } + + recordBashResult(command: string, result: BashResult, options?: { excludeFromContext?: boolean }): void { + const bashMessage: BashExecutionMessage = { + role: "bashExecution", + command, + output: result.output, + exitCode: result.exitCode, + cancelled: result.cancelled, + truncated: result.truncated, + fullOutputPath: result.fullOutputPath, + timestamp: Date.now(), + excludeFromContext: options?.excludeFromContext, + }; + + // If agent is streaming, defer adding to avoid breaking tool_use/tool_result ordering + if (this._host.isStreaming()) { + this._pendingBashMessages.push(bashMessage); + } else { + this._host.appendMessage(bashMessage); + } + } + + /** + * Cancel running bash command. + */ + abortBash(): void { + // A user bash command may not have spawned yet (extension dispatch in + // progress); flag the request so runUserBash cancels before executing. + // runUserBash clears the flag at each start, so a stale flag is harmless. + if (this._userBashRunning) { + this._userBashAbortRequested = true; + } + for (const controller of this._bashAbortControllers) { + controller.abort(); + } + } + + get isBashRunning(): boolean { + return this._bashAbortControllers.size > 0 || this._userBashRunning; + } + + /** Whether there are pending bash messages waiting to be flushed */ + get hasPendingBashMessages(): boolean { + return this._pendingBashMessages.length > 0; + } + + /** + * Flush pending bash messages to agent state and session. + * Called after agent turn completes to maintain proper message ordering. + */ + flushPendingMessages(): void { + if (this._pendingBashMessages.length === 0) return; + + for (const bashMessage of this._pendingBashMessages) { + this._host.appendMessage(bashMessage); + } + + this._pendingBashMessages = []; + } +} diff --git a/packages/coding-agent/src/session/commit-fence.ts b/packages/coding-agent/src/session/commit-fence.ts new file mode 100644 index 0000000000..a28f5c49d6 --- /dev/null +++ b/packages/coding-agent/src/session/commit-fence.ts @@ -0,0 +1,75 @@ +import { AsyncLocalStorage } from "node:async_hooks"; +import { waitForPromiseOrAbort } from "../utils/wait-for-abort.js"; + +export interface SessionCommitLease { + readonly owner: symbol; + release(): void; +} + +export class SessionCommitFence { + private _tail: Promise = Promise.resolve(); + private _owner: symbol | undefined; + private _pendingWaiters = 0; + private readonly _context = new AsyncLocalStorage(); + private readonly _disposeAbortController = new AbortController(); + + get disposeSignal(): AbortSignal { + return this._disposeAbortController.signal; + } + + get hasPendingWork(): boolean { + return this._owner !== undefined || this._pendingWaiters > 0; + } + + get isHeldByCurrentContext(): boolean { + const inheritedOwner = this._context.getStore(); + return inheritedOwner !== undefined && inheritedOwner === this._owner; + } + + run(lease: SessionCommitLease, callback: () => T): T { + return this._context.run(lease.owner, callback); + } + + async acquire(signal?: AbortSignal): Promise { + const inheritedOwner = this._context.getStore(); + if (inheritedOwner !== undefined && inheritedOwner === this._owner) { + return { owner: inheritedOwner, release: () => {} }; + } + const previous = this._tail; + let resolve = () => {}; + this._tail = new Promise((release) => { + resolve = release; + }); + const disposeSignal = this._disposeAbortController.signal; + const waitSignal = signal ? AbortSignal.any([signal, disposeSignal]) : disposeSignal; + this._pendingWaiters++; + try { + await waitForPromiseOrAbort(previous, waitSignal, "Update restart preparation cancelled"); + } catch (error) { + this._pendingWaiters--; + // A cancelled waiter remains in the FIFO chain until its predecessor releases. + void previous.then(resolve, resolve); + if (disposeSignal.aborted) { + throw new Error("Cannot admit a session action because the session is disposing or disposed."); + } + throw error; + } + const owner = Symbol("session-action-commit"); + this._owner = owner; + this._pendingWaiters--; + let released = false; + return { + owner, + release: () => { + if (released) return; + released = true; + if (this._owner === owner) this._owner = undefined; + resolve(); + }, + }; + } + + dispose(): void { + this._disposeAbortController.abort(); + } +} diff --git a/packages/coding-agent/src/session/input-dispatcher.ts b/packages/coding-agent/src/session/input-dispatcher.ts new file mode 100644 index 0000000000..baef151b13 --- /dev/null +++ b/packages/coding-agent/src/session/input-dispatcher.ts @@ -0,0 +1,208 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { + type ActionStore, + canSelectSessionAction, + type DeliveryPolicy, + type RuntimeActivity, + transitionSessionAction, +} from "../core/session-action-store.js"; +import { DeferredSessionInputError, primaryDeliveryRecord, type QueuedSessionAction } from "./prepared-actions.js"; +import { turnExecutionPoliciesEqual } from "./turn-preparation.js"; + +export interface SessionInputDispatcherHost { + isDisposed(): boolean; + getEpoch(): number; + getActivity(): RuntimeActivity; + isBusy(): boolean; + isHandoffDeferred(epoch: number): boolean; + getDeliveryMode(delivery: DeliveryPolicy): "all" | "one-at-a-time"; + waitForAgentIdle(): Promise; + hasCancelledDispatchCapture(): boolean; + getEventQueue(): Promise; + waitForRefinement(): Promise; + getTranscript(): readonly AgentMessage[]; + startTurns(actions: QueuedSessionAction[], epoch: number): Promise; + executeCommand(action: QueuedSessionAction, epoch: number): Promise; + settleAgentMessage(id: string | undefined, leg: "delivery" | "completion", error?: Error): void; + releaseTurn(id: string): void; + notifyCheckpoints(): void; + emitQueueUpdate(): void; + surfaceError(error: unknown): void; + schedule(): void; +} + +export class SessionInputDispatcher { + constructor( + private readonly actions: ActionStore, + private readonly host: SessionInputDispatcherHost, + ) {} + + hasSelectableInput(): boolean { + return ( + this.actions.queuedActions().length > 0 || + this.actions.activeActions().some((action) => action.lifecycle.state === "selected") + ); + } + + async run(epoch: number): Promise { + let blocked = false; + try { + while (!this.host.isDisposed() && this.hasSelectableInput()) { + await this.host.waitForAgentIdle(); + const preselected = this.actions.activeActions().find((action) => action.lifecycle.state === "selected"); + if (epoch !== this.host.getEpoch()) { + if (preselected) { + this.actions.rollback(preselected); + this.host.notifyCheckpoints(); + this.host.emitQueueUpdate(); + } + return; + } + if (!this.host.hasCancelledDispatchCapture()) await this.host.getEventQueue(); + if (!preselected || preselected.payload.kind === "session_command") await this.host.waitForRefinement(); + const activity = this.host.getActivity(); + const canSelectPreselectedTurn = + preselected?.payload.kind === "turn" && canSelectSessionAction({ ...activity, refinementApply: false }); + if ( + this.host.isHandoffDeferred(epoch) || + (!canSelectPreselectedTurn && !canSelectSessionAction(activity)) + ) { + blocked = true; + this.host.notifyCheckpoints(); + return; + } + const first = preselected ?? this.actions.selectFirst(); + if (!first) return; + if (first.payload.kind === "session_command") { + await this.host.executeCommand(first, epoch); + return; + } + + const mode = this.host.getDeliveryMode(first.delivery); + const actions: QueuedSessionAction[] = [first]; + while (!preselected && mode === "all") { + const next = this.actions.queuedActions(first.delivery)[0]; + if ( + !next || + next.payload.kind !== "turn" || + !turnExecutionPoliciesEqual(first.payload.executionPolicy, next.payload.executionPolicy) + ) { + break; + } + this.actions.selectFirst(); + actions.push(next); + } + if (epoch !== this.host.getEpoch()) { + for (const action of actions) this.actions.rollback(action); + return; + } + for (const action of actions) transitionSessionAction(action, { state: "preparing" }); + this.host.notifyCheckpoints(); + this.host.emitQueueUpdate(); + try { + await this.host.startTurns(actions, epoch); + for (const action of actions) { + if (action.lifecycle.state === "committing") { + const primary = primaryDeliveryRecord(action); + if (this.host.getTranscript().includes(primary.message)) { + primary.durable = true; + transitionSessionAction(action, { + state: "running", + execution: "agent_turn", + }); + } + } + if (action.lifecycle.state === "running") { + transitionSessionAction(action, { state: "completed" }); + this.actions.ticketFor(action).settleCompleted(); + this.host.settleAgentMessage(action.agentMessageId, "completion"); + } + } + } catch (error) { + const transcript = this.host.getTranscript(); + const delivered = new Set(transcript); + const undelivered: QueuedSessionAction[] = []; + for (const action of actions) { + if (action.payload.kind !== "turn" || action.lifecycle.state === "cancelled") continue; + for (const record of action.payload.records) record.durable ||= delivered.has(record.message); + action.payload.records = action.payload.records.filter((record) => { + if (record.role === "prefix") return !record.durable; + if (record.role === "next_turn") return record.durable; + return true; + }); + if (!primaryDeliveryRecord(action).durable) undelivered.push(action); + } + if (this.isDeferredError(error, epoch)) { + for (const action of undelivered) { + if (action.lifecycle.state === "committing") { + this.actions.rollback(action, { + dispatchSettled: true, + transcript, + }); + } else if (action.lifecycle.state === "preparing" || action.lifecycle.state === "selected") { + this.actions.rollback(action); + } + } + if (undelivered.length > 0) this.host.emitQueueUpdate(); + blocked = epoch !== this.host.getEpoch() || this.host.isBusy(); + if (blocked) return; + continue; + } + const terminalError = error instanceof Error ? error : new Error(String(error)); + for (const action of actions) { + if (action.lifecycle.state === "cancelled") continue; + if (action.lifecycle.state !== "completed" && action.lifecycle.state !== "failed") { + transitionSessionAction(action, { + state: "failed", + error: terminalError, + }); + } + const ticket = this.actions.ticketFor(action); + if (undelivered.includes(action)) { + ticket.rejectDelivered(terminalError); + this.host.settleAgentMessage(action.agentMessageId, "delivery", terminalError); + } + this.host.settleAgentMessage(action.agentMessageId, "completion", terminalError); + ticket.settleCompleted(terminalError); + } + if (actions.some((action) => action.payload.kind !== "turn" || action.payload.queueVisible)) { + this.host.surfaceError(error); + } + } finally { + for (const action of actions) { + const retainedCancelledDispatch = + action.lifecycle.state === "cancelled" && + action.payload.kind === "turn" && + action.payload.captureRunMessages !== undefined; + if ( + !retainedCancelledDispatch && + (action.lifecycle.state === "completed" || + action.lifecycle.state === "failed" || + action.lifecycle.state === "cancelled") + ) { + this.host.releaseTurn(action.id); + this.actions.releaseTerminal(action); + } + } + this.host.notifyCheckpoints(); + this.host.emitQueueUpdate(); + } + if (epoch !== this.host.getEpoch() || blocked) return; + } + } finally { + if (!blocked && epoch === this.host.getEpoch() && this.hasSelectableInput()) { + this.host.schedule(); + } + } + } + + private isDeferredError(error: unknown, epoch: number): boolean { + if (error instanceof DeferredSessionInputError) return true; + if (epoch !== this.host.getEpoch()) return true; + if (this.host.isBusy()) { + this.host.surfaceError(error); + return true; + } + return false; + } +} diff --git a/packages/coding-agent/src/session/input-scheduler.ts b/packages/coding-agent/src/session/input-scheduler.ts new file mode 100644 index 0000000000..4bcf2f5014 --- /dev/null +++ b/packages/coding-agent/src/session/input-scheduler.ts @@ -0,0 +1,119 @@ +interface SessionInputSchedulerDependencies { + canSchedule(): boolean; + run(epoch: number): Promise; +} + +export class SessionInputScheduler { + private _pump: Promise = Promise.resolve(); + private _requested = false; + private _epoch = 0; + private _suspended = false; + private _suspendedForUpdateRestart = false; + private readonly _queuedWorkPauses = new Set(); + private readonly _admissionPauses = new Set(); + + constructor(private readonly dependencies: SessionInputSchedulerDependencies) {} + + get pendingPump(): Promise { + return this._pump; + } + + get requested(): boolean { + return this._requested; + } + + get epoch(): number { + return this._epoch; + } + + get suspended(): boolean { + return this._suspended; + } + + get suspendedForUpdateRestart(): boolean { + return this._suspendedForUpdateRestart; + } + + get queuedWorkPauseCount(): number { + return this._queuedWorkPauses.size; + } + + get admissionPaused(): boolean { + return this._admissionPauses.size > 0; + } + + schedule(): void { + if (this._suspended || this._queuedWorkPauses.size > 0) return; + if (this._requested || !this.dependencies.canSchedule()) return; + this._requested = true; + const epoch = this._epoch; + const pump = async () => { + this._requested = false; + await this.dependencies.run(epoch); + }; + this._pump = this._pump.then(pump, pump); + this._pump.catch(() => {}); + } + + // The runner checks this epoch after asynchronous preparation, including pauses + // that have already been released by the time preparation finishes. + invalidatePreparation(): void { + this._epoch++; + } + + acquireAdmissionPause(onRelease: () => void): { release(): void } { + const token = Symbol("session-input-admission-pause"); + this._admissionPauses.add(token); + this._requested = false; + this._epoch++; + let released = false; + return { + release: () => { + if (released) return; + released = true; + this._admissionPauses.delete(token); + this._epoch++; + onRelease(); + }, + }; + } + + acquireQueuedWorkPause(onRelease: () => void): { release(): void } { + const token = Symbol("queued-work-pause"); + this._queuedWorkPauses.add(token); + this._requested = false; + this._epoch++; + let released = false; + return { + release: () => { + if (released) return; + released = true; + this._queuedWorkPauses.delete(token); + onRelease(); + }, + }; + } + + suspend(reason: "abort" | "update-restart"): void { + this._requested = false; + this._epoch++; + this._suspended = true; + this._suspendedForUpdateRestart = reason === "update-restart"; + } + + resume(): boolean { + if (!this._suspended) return false; + this._suspended = false; + this._suspendedForUpdateRestart = false; + this._epoch++; + return true; + } + + async waitForIdle(): Promise { + while (true) { + const pump = this._pump; + await pump; + if (pump === this._pump && !this._requested) return; + } + } +} diff --git a/packages/coding-agent/src/session/prepared-actions.ts b/packages/coding-agent/src/session/prepared-actions.ts new file mode 100644 index 0000000000..67a0088b56 --- /dev/null +++ b/packages/coding-agent/src/session/prepared-actions.ts @@ -0,0 +1,277 @@ +import { randomUUID } from "node:crypto"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { ImageContent, TextContent, UserMessage } from "@earendil-works/pi-ai"; +import { AGENT_MESSAGE_RECEIVED_PREVIEW_LABEL, isAgentSessionMessage } from "../core/agent-messages.js"; +import type { ExtensionRunner, InputSource } from "../core/extensions/index.js"; +import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ASYNC_BASH_COMPLETION_PREVIEW_LABEL, + type AsyncBashCompletionDetails, + type CustomMessage, +} from "../core/messages.js"; +import type { + DeliveryPolicy, + DeliveryRecord, + SessionAction, + SessionCommandPayload, + SessionTurnPayload, + WakePolicy, +} from "../core/session-action-store.js"; +import type { SessionSlashCommand } from "../core/slash-commands.js"; +import { createTurnExecutionPolicy, type TurnExecutionPolicy } from "./turn-preparation.js"; + +export type QueuedAgentMessage = UserMessage | CustomMessage; +export type SessionInputSchedule = "steer" | "followUp"; + +export interface PreparedTurnPayload extends SessionTurnPayload { + images?: ImageContent[]; + content?: (TextContent | ImageContent)[]; + customMessage?: CustomMessage; + prepared?: PreparedPromptPreparation; + executionPolicy: TurnExecutionPolicy; + queueVisible: boolean; + acceptedAgentMessage: boolean; + acceptedBeforeCompletion: boolean; + captureRunMessages?: Set; + cancelledDispatchEnded?: boolean; +} + +export interface PreparedCommandPayload extends SessionCommandPayload { + images?: ImageContent[]; +} + +export type QueuedSessionAction = SessionAction; + +export interface PreparedPromptPreparation { + result: Awaited>; + basePromptSnapshot: string; +} + +export class DeferredSessionInputError extends Error {} + +export class SessionInputAdmissionPausedError extends Error {} + +export interface RestoredPromptInput { + text: string; + content?: (TextContent | ImageContent)[]; + images?: ImageContent[]; + queueKey?: string; + agentMessageId?: string; + customMessage?: CustomMessage; + prefixMessages?: CustomMessage[]; +} + +export const SESSION_ACTION_RECOVERY_FORMAT_VERSION = 1; + +export interface SessionActionRecoveryRecord { + id: string; + role: DeliveryRecord["role"]; + message: QueuedAgentMessage; + ownerActionId: string; +} + +export type SessionActionRecoveryPayload = + | { + kind: "turn"; + text: string; + preview?: string; + records: SessionActionRecoveryRecord[]; + images?: ImageContent[]; + content?: (TextContent | ImageContent)[]; + customMessage?: CustomMessage; + executionPolicy: TurnExecutionPolicy; + queueVisible: boolean; + acceptedAgentMessage: boolean; + acceptedBeforeCompletion: boolean; + } + | { + kind: "session_command"; + text: string; + command: SessionSlashCommand; + images?: ImageContent[]; + }; + +export interface SessionActionRecoveryAction { + id: string; + source: InputSource | "internal"; + delivery: DeliveryPolicy; + wake: WakePolicy; + payload: SessionActionRecoveryPayload; + queueKey?: string; + agentMessageId?: string; + suppressAutonomousContinuation?: boolean; +} + +export interface SessionActionRecoverySnapshot { + formatVersion: typeof SESSION_ACTION_RECOVERY_FORMAT_VERSION; + actions: SessionActionRecoveryAction[]; +} + +export function cloneCustomMessage(message: CustomMessage): CustomMessage { + return { + ...message, + content: Array.isArray(message.content) ? message.content.map((block) => ({ ...block })) : message.content, + }; +} + +export function cloneQueuedAgentMessage(message: QueuedAgentMessage): QueuedAgentMessage { + if (message.role === "custom") return cloneCustomMessage(message); + return { + ...message, + content: Array.isArray(message.content) ? message.content.map((block) => ({ ...block })) : message.content, + }; +} + +export function primaryDeliveryRecord(action: QueuedSessionAction): DeliveryRecord { + if (action.payload.kind !== "turn") throw new Error(`Session action ${action.id} is not a turn`); + const record = action.payload.records.find((candidate) => candidate.role === "primary"); + if (!record) throw new Error(`Turn action ${action.id} has no primary delivery record`); + return record; +} + +export function normalizeMessageContent(content: string | (TextContent | ImageContent)[]): { + text: string; + images?: ImageContent[]; +} { + if (typeof content === "string") return { text: content }; + const text = content + .filter((part): part is TextContent => part.type === "text") + .map((part) => part.text) + .join("\n"); + const images = content.filter((part): part is ImageContent => part.type === "image"); + return { text, ...(images.length > 0 ? { images } : {}) }; +} + +export function queuedAgentMessagePreview(action: QueuedSessionAction): string { + const payload = action.payload; + if (payload.kind === "session_command") return payload.text; + if (payload.customMessage && isAgentSessionMessage(payload.customMessage)) { + return `${AGENT_MESSAGE_RECEIVED_PREVIEW_LABEL}: ${payload.customMessage.details.message}`; + } + if (payload.customMessage?.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE) { + const details = payload.customMessage.details as AsyncBashCompletionDetails | undefined; + return details + ? `${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}: pid ${details.pid}, exit ${details.exitCode}` + : ASYNC_BASH_COMPLETION_PREVIEW_LABEL; + } + return payload.preview ?? payload.text; +} + +export function visibleSessionActionProjection( + actions: readonly QueuedSessionAction[], +): readonly QueuedSessionAction[] { + return actions.filter( + (action) => + action.payload.kind === "session_command" || + action.payload.queueVisible || + action.payload.acceptedAgentMessage, + ); +} + +export function buildPromptContent(text: string, images?: ImageContent[]): (TextContent | ImageContent)[] { + const content: (TextContent | ImageContent)[] = []; + content.push({ type: "text", text }); + if (images) content.push(...images); + return content; +} + +function deliveryPolicy(schedule: SessionInputSchedule): DeliveryPolicy { + return schedule === "steer" ? "next_turn_boundary" : "when_run_idle"; +} + +export function createDeliveryRecord( + actionId: string, + role: DeliveryRecord["role"], + message: QueuedAgentMessage, +): DeliveryRecord { + return { + id: randomUUID(), + role, + message, + started: false, + durable: false, + ownerActionId: actionId, + }; +} + +export function createPreparedTurnAction( + schedule: SessionInputSchedule, + text: string, + images: ImageContent[] | undefined, + options: { + agentMessageId?: string; + queueKey?: string; + content?: (TextContent | ImageContent)[]; + message?: QueuedAgentMessage; + prefixMessages?: CustomMessage[]; + previewLabel?: string; + suppressAutonomousContinuation?: boolean; + resumeIfIdle?: boolean; + source?: InputSource | "internal"; + executionPolicy?: TurnExecutionPolicy; + queueVisible?: boolean; + acceptedAgentMessage?: boolean; + acceptedBeforeCompletion?: boolean; + }, +): QueuedSessionAction { + const id = randomUUID(); + const content = options.content ?? buildPromptContent(text, images); + const message = + options.message ?? + ({ + role: "user", + content: content.map((block) => ({ ...block })), + timestamp: Date.now(), + } satisfies UserMessage); + const prefixMessages = options.prefixMessages?.map((prefix) => cloneCustomMessage(prefix)) ?? []; + const preview = options.previewLabel ? `${options.previewLabel}: ${text}` : undefined; + const payload: PreparedTurnPayload = { + kind: "turn", + text, + records: [ + ...prefixMessages.map((prefix) => createDeliveryRecord(id, "prefix", prefix)), + createDeliveryRecord(id, "primary", message), + ], + preview, + images: images?.map((image) => ({ ...image })), + content: content.map((block) => ({ ...block })), + customMessage: options.message?.role === "custom" ? cloneCustomMessage(options.message) : undefined, + executionPolicy: options.executionPolicy ?? createTurnExecutionPolicy("queued"), + queueVisible: options.queueVisible ?? true, + acceptedAgentMessage: options.acceptedAgentMessage ?? false, + acceptedBeforeCompletion: options.acceptedBeforeCompletion ?? false, + }; + return { + id, + source: options.source ?? "internal", + delivery: deliveryPolicy(schedule), + wake: + options.resumeIfIdle === true ? "immediate" : schedule === "steer" ? "on_lower_boundary" : "external_resume", + payload, + lifecycle: { state: "queued" }, + queueKey: options.queueKey, + agentMessageId: options.agentMessageId, + suppressAutonomousContinuation: options.suppressAutonomousContinuation, + }; +} + +export function createSessionCommandAction( + text: string, + command: SessionSlashCommand, + images: ImageContent[] | undefined, + schedule: SessionInputSchedule, + options: { + agentMessageId?: string; + source?: InputSource | "internal"; + } = {}, +): QueuedSessionAction { + return { + id: randomUUID(), + source: options.source ?? "internal", + delivery: deliveryPolicy(schedule), + wake: "immediate", + payload: { kind: "session_command", text, command, images }, + lifecycle: { state: "queued" }, + agentMessageId: options.agentMessageId, + }; +} diff --git a/packages/coding-agent/src/session/retry.ts b/packages/coding-agent/src/session/retry.ts new file mode 100644 index 0000000000..cfa92ae06a --- /dev/null +++ b/packages/coding-agent/src/session/retry.ts @@ -0,0 +1,411 @@ +import type { AgentEvent } from "@earendil-works/pi-agent-core"; +import { type AssistantMessage, isContextOverflow } from "@earendil-works/pi-ai"; +import { addLoginGuidanceToAuthError } from "../core/auth-guidance.js"; +import type { AuthSourceToken } from "../core/auth-storage.js"; +import { + isAgentLifecycleFailure, + isFauxProviderQueueExhausted, + isPermanentProviderFailureKind, + providerRetryDelay, + providerStreamFailureKind, + providerStreamFailureRetryAfterMs, +} from "../core/provider-retry.js"; +import type { SettingsManager } from "../core/settings-manager.js"; +import { sleep } from "../utils/sleep.js"; + +export type SessionRetryEvent = + | { + type: "auto_retry_start"; + attempt: number; + maxAttempts: number; + delayMs: number; + errorMessage: string; + } + | { + type: "auto_retry_end"; + success: boolean; + attempt: number; + finalError?: string; + } + | { + type: "auth_stale"; + provider: string; + sourceTokens?: readonly AuthSourceToken[]; + }; + +export interface SessionRetryHost { + getRetrySettings(): ReturnType; + getMaxRetryDelayMs(): number; + getContextWindow(): number; + getAuthSource(provider: string): AuthSourceToken | undefined; + markAuthSourceStale(token: AuthSourceToken): boolean; + markAuthStale(provider: string): boolean; + hasPayloadHooks(): boolean; + prepareTurnRetry(): void; + clearTurnRetry(): void; + removeLastAssistant(): void; + continue(): Promise; + waitForIdle(): Promise; + cancelCompaction(): void; + emit(event: SessionRetryEvent): void; + onResolved(): void; +} + +/** Owns one session's retry chain; callbacks read live collaborators at each boundary. */ +export class SessionRetry { + private _retryAbortController: AbortController | undefined = undefined; + private _retryAttempt = 0; + /** Bumped by every retry resolution; stale scheduled-continue callbacks check it before touching retry state. */ + private _retryGeneration = 0; + private _retryPromise: Promise | undefined = undefined; + private _retryResolve: (() => void) | undefined = undefined; + private _retryAuthFailureSources: AuthSourceToken[] = []; + + constructor(private readonly host: SessionRetryHost) {} + + observeAgentEnd(event: AgentEvent): void { + if (event.type !== "agent_end" || this._retryPromise) { + return; + } + + const settings = this.host.getRetrySettings(); + if (!settings.enabled) { + return; + } + + let lastAssistant: AssistantMessage | undefined; + for (let i = event.messages.length - 1; i >= 0; i--) { + const message = event.messages[i]; + if (message.role === "assistant") { + lastAssistant = message as AssistantMessage; + break; + } + } + const concreteAuthFailure = lastAssistant ? this._isConcreteProviderAuthFailure(lastAssistant) : false; + if (!lastAssistant || (!this._isRetryableError(lastAssistant) && !concreteAuthFailure)) { + return; + } + if (concreteAuthFailure) { + this._captureRetryAuthFailureSource(lastAssistant); + } + + this._retryPromise = new Promise((resolve) => { + this._retryResolve = resolve; + }); + } + + observeAssistantEnd(assistantMsg: AssistantMessage): void { + if (this._isConcreteProviderAuthFailure(assistantMsg)) { + this._captureRetryAuthFailureSource(assistantMsg); + } + + // Reset retry counter immediately on successful assistant response + // This prevents accumulation across multiple LLM calls within a turn + if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) { + this.host.emit({ + type: "auto_retry_end", + success: true, + attempt: this._retryAttempt, + }); + this._retryAttempt = 0; + this._retryAuthFailureSources = []; + } + } + + /** Undefined preserves the caller's synchronous fallthrough when no retry applies. */ + retryError(msg: AssistantMessage): Promise | undefined { + const concreteAuthFailure = this._isConcreteProviderAuthFailure(msg); + const retryConcreteAuthFailure = concreteAuthFailure && !this._isStructuredPermanentProviderRetryExhausted(msg); + if (this._isRetryableError(msg) || retryConcreteAuthFailure) { + if (retryConcreteAuthFailure) { + this._captureRetryAuthFailureSource(msg); + } + return this._handleRetryableError(msg, { + markAuthStaleOnFailure: retryConcreteAuthFailure, + authSourceTokens: retryConcreteAuthFailure ? this._retryAuthFailureSources : undefined, + }); + } + } + + resolve(): void { + this._retryGeneration += 1; + this.host.clearTurnRetry(); + if (this._retryResolve) { + this._retryResolve(); + this._retryResolve = undefined; + this._retryPromise = undefined; + this.host.onResolved(); + } + } + + private _isRetryableError(message: AssistantMessage): boolean { + if (message.stopReason !== "error" || !message.errorMessage) return false; + + const contextWindow = this.host.getContextWindow(); + if (isContextOverflow(message, contextWindow)) return false; + + if (this._isFauxProviderQueueExhausted(message)) { + return false; + } + + if (this._isAgentLifecycleFailure(message)) { + return false; + } + + if (this._isStructuredPermanentProviderRetryExhausted(message)) { + return false; + } + + return true; + } + + private _isFauxProviderQueueExhausted(message: AssistantMessage): boolean { + return isFauxProviderQueueExhausted(message); + } + + private _isAgentLifecycleFailure(message: AssistantMessage): boolean { + return isAgentLifecycleFailure(message); + } + + private _getProviderStreamFailureKind(message: AssistantMessage): string | undefined { + return providerStreamFailureKind(message); + } + + private _isStructuredPermanentProviderRetryExhausted(message: AssistantMessage): boolean { + return isPermanentProviderFailureKind(this._getProviderStreamFailureKind(message), this._retryAttempt); + } + + private _isConcreteProviderAuthFailure(message: AssistantMessage): boolean { + if (message.stopReason !== "error" || !message.errorMessage) return false; + // Only the provider's structured classification counts as an auth failure. + return this._getProviderStreamFailureKind(message) === "auth"; + } + + private _captureRetryAuthFailureSource(message: AssistantMessage): AuthSourceToken | undefined { + const token = this.host.getAuthSource(message.provider); + if (!token) { + return undefined; + } + if ( + !this._retryAuthFailureSources.some( + (existing) => + existing.provider === token.provider && + existing.source === token.source && + existing.identityFingerprint === token.identityFingerprint && + existing.valueFingerprint === token.valueFingerprint, + ) + ) { + this._retryAuthFailureSources.push(token); + } + return token; + } + + private _markProviderAuthStale(message: AssistantMessage, authSourceTokens?: readonly AuthSourceToken[]): boolean { + if (authSourceTokens && authSourceTokens.length > 0) { + let marked = false; + for (const token of authSourceTokens) { + marked = this.host.markAuthSourceStale(token) || marked; + } + if (marked) { + this.host.emit({ + type: "auth_stale", + provider: message.provider, + sourceTokens: authSourceTokens, + }); + } + return marked; + } + const marked = this.host.markAuthStale(message.provider); + if (marked) { + this.host.emit({ type: "auth_stale", provider: message.provider }); + } + return marked; + } + + private _markProviderAuthStaleForRetryFailure( + message: AssistantMessage, + options?: { + markAuthStaleOnFailure?: boolean; + authSourceTokens?: readonly AuthSourceToken[]; + }, + ): boolean { + const authSourceTokens = + this._retryAuthFailureSources.length > 0 ? this._retryAuthFailureSources : options?.authSourceTokens; + if ((authSourceTokens?.length ?? 0) > 0 || options?.markAuthStaleOnFailure) { + const marked = this._markProviderAuthStale(message, authSourceTokens); + if (marked && message.errorMessage) { + message.errorMessage = addLoginGuidanceToAuthError(message.errorMessage); + } + return marked; + } + return false; + } + + finishActiveRetryWithFailure(message: AssistantMessage): void { + if (this._retryAttempt === 0) { + return; + } + this._markProviderAuthStaleForRetryFailure(message); + this.host.emit({ + type: "auto_retry_end", + success: false, + attempt: this._retryAttempt, + finalError: message.errorMessage, + }); + this._retryAttempt = 0; + this._retryAuthFailureSources = []; + } + + private async _handleRetryableError( + message: AssistantMessage, + options?: { + markAuthStaleOnFailure?: boolean; + authSourceTokens?: readonly AuthSourceToken[]; + }, + ): Promise { + const settings = this.host.getRetrySettings(); + if (!settings.enabled) { + this._markProviderAuthStaleForRetryFailure(message, options); + this._retryAuthFailureSources = []; + this.resolve(); + return false; + } + + if (!this._retryPromise) { + this._retryPromise = new Promise((resolve) => { + this._retryResolve = resolve; + }); + } + + this._retryAttempt++; + + if (this._retryAttempt > settings.maxRetries) { + this._markProviderAuthStaleForRetryFailure(message, options); + this.host.emit({ + type: "auto_retry_end", + success: false, + attempt: this._retryAttempt - 1, + finalError: message.errorMessage, + }); + this._retryAttempt = 0; + this._retryAuthFailureSources = []; + this.resolve(); // Resolve so waitForRetry() completes + return false; + } + + // Server-requested waits are honored, capped by retry.provider.maxRetryDelayMs (0 disables). + const maxRetryDelayMs = this.host.getMaxRetryDelayMs(); + const delay = providerRetryDelay(this._retryAttempt, providerStreamFailureRetryAfterMs(message), { + baseDelayMs: settings.baseDelayMs, + maxRetryDelayMs, + }); + if (delay.kind === "exceeds-cap") { + this._markProviderAuthStaleForRetryFailure(message, options); + this.host.emit({ + type: "auto_retry_end", + success: false, + attempt: this._retryAttempt - 1, + finalError: `Provider requested a ${Math.ceil(delay.retryAfterMs / 1000)}s wait before retrying (above retry.provider.maxRetryDelayMs=${maxRetryDelayMs}ms): ${message.errorMessage || "unknown error"}`, + }); + this._retryAttempt = 0; + this._retryAuthFailureSources = []; + this.resolve(); + return false; + } + + const delayMs = delay.delayMs; + // Park now: the retry re-issues the failed call and must reuse its Idempotency-Key. + // Payload hooks mutate the wire body after the hash point, so reuse is forfeited. + if (!this.host.hasPayloadHooks()) { + this.host.prepareTurnRetry(); + } + + this.host.emit({ + type: "auto_retry_start", + attempt: this._retryAttempt, + maxAttempts: settings.maxRetries, + delayMs, + errorMessage: message.errorMessage || "Unknown error", + }); + + this.host.removeLastAssistant(); + + this._retryAbortController = new AbortController(); + try { + await sleep(delayMs, this._retryAbortController.signal); + } catch { + const attempt = this._retryAttempt; + this._markProviderAuthStaleForRetryFailure(message, options); + this._retryAttempt = 0; + this._retryAbortController = undefined; + this.host.emit({ + type: "auto_retry_end", + success: false, + attempt, + finalError: "Retry cancelled", + }); + this.resolve(); + this._retryAuthFailureSources = []; + return false; + } + this._retryAbortController = undefined; + + const retryGeneration = this._retryGeneration; + setTimeout(() => { + this.host.continue().catch((error: unknown) => { + // A continue that never starts must still resolve the retry (else isRetrying + // sticks forever) — unless a newer retry owns the state by now. + if (this._retryGeneration !== retryGeneration || !this.isRetrying) return; + this._markProviderAuthStaleForRetryFailure(message, options); + const attempt = this._retryAttempt; + this._retryAttempt = 0; + this._retryAuthFailureSources = []; + this.host.emit({ + type: "auto_retry_end", + success: false, + attempt, + finalError: error instanceof Error ? error.message : String(error), + }); + this.resolve(); + }); + }, 0); + + return true; + } + + abortRetry(): void { + if (this._retryAbortController) { + this._retryAbortController.abort(); + return; + } + if (this._retryAttempt > 0) { + this.host.cancelCompaction(); + this.host.emit({ + type: "auto_retry_end", + success: false, + attempt: this._retryAttempt, + finalError: "Retry cancelled", + }); + this._retryAttempt = 0; + } + this._retryAuthFailureSources = []; + this.resolve(); + } + + async waitForRetry(): Promise { + if (!this._retryPromise) { + return; + } + + await this._retryPromise; + await this.host.waitForIdle(); + } + + get isRetrying(): boolean { + return this._retryPromise !== undefined; + } + + get attempt(): number { + return this._retryAttempt; + } +} diff --git a/packages/coding-agent/src/session/turn-preparation.ts b/packages/coding-agent/src/session/turn-preparation.ts new file mode 100644 index 0000000000..d63e19e8b5 --- /dev/null +++ b/packages/coding-agent/src/session/turn-preparation.ts @@ -0,0 +1,162 @@ +type PreTurnCompactionTiming = "beforeModelSelection" | "afterModelSelection" | "skip"; +type RefineBarrierPolicy = "always" | "ifInFlight" | "skip"; + +export interface CommitPreparationPolicy { + initialRefineBarrier: RefineBarrierPolicy; + flushPendingBashBeforeValidation: boolean; + validateModelAndAuth: boolean; + awaitPendingModelSelection: boolean; + preTurnCompaction: PreTurnCompactionTiming; + finalRefineBarrier: RefineBarrierPolicy; +} + +export interface CommitPreparationSteps { + afterValidation?: () => void; + prepare: () => Promise; + shouldCommit?: (prepared: TPrepared) => boolean; + beforeFinalRefineBarrier?: (prepared: TPrepared) => void; + commit: (prepared: TPrepared, passedFinalRefineBarrier: boolean) => TCommitted; +} + +export interface TurnExecutionPolicy { + preparation: CommitPreparationPolicy; + runBeforeAgentStart: boolean; + nextTurnContextTiming: "preparation" | "commit" | "skip"; + preserveEmptyExtensionPrompt: boolean; + completionIncludesRetryChain: boolean; +} + +export function turnExecutionPoliciesEqual(left: TurnExecutionPolicy, right: TurnExecutionPolicy): boolean { + return ( + left.preparation.initialRefineBarrier === right.preparation.initialRefineBarrier && + left.preparation.flushPendingBashBeforeValidation === right.preparation.flushPendingBashBeforeValidation && + left.preparation.validateModelAndAuth === right.preparation.validateModelAndAuth && + left.preparation.awaitPendingModelSelection === right.preparation.awaitPendingModelSelection && + left.preparation.preTurnCompaction === right.preparation.preTurnCompaction && + left.preparation.finalRefineBarrier === right.preparation.finalRefineBarrier && + left.runBeforeAgentStart === right.runBeforeAgentStart && + left.nextTurnContextTiming === right.nextTurnContextTiming && + left.preserveEmptyExtensionPrompt === right.preserveEmptyExtensionPrompt && + left.completionIncludesRetryChain === right.completionIncludesRetryChain + ); +} + +export function createTurnExecutionPolicy( + kind: "queued" | "directPrompt" | "injected" | "customTrigger", + options: { + returnAfterAccepted?: boolean; + skipPrePromptWork?: boolean; + } = {}, +): TurnExecutionPolicy { + if (kind === "queued") { + return { + preparation: { + initialRefineBarrier: "skip", + flushPendingBashBeforeValidation: false, + validateModelAndAuth: true, + awaitPendingModelSelection: true, + preTurnCompaction: "beforeModelSelection", + finalRefineBarrier: "always", + }, + runBeforeAgentStart: true, + nextTurnContextTiming: "commit", + preserveEmptyExtensionPrompt: true, + completionIncludesRetryChain: true, + }; + } + if (kind === "directPrompt") { + return { + preparation: { + initialRefineBarrier: options.returnAfterAccepted ? "skip" : "always", + flushPendingBashBeforeValidation: true, + validateModelAndAuth: true, + awaitPendingModelSelection: true, + preTurnCompaction: options.skipPrePromptWork ? "skip" : "afterModelSelection", + finalRefineBarrier: "ifInFlight", + }, + runBeforeAgentStart: !options.skipPrePromptWork, + nextTurnContextTiming: "preparation", + preserveEmptyExtensionPrompt: false, + completionIncludesRetryChain: true, + }; + } + if (kind === "injected") { + return { + preparation: { + initialRefineBarrier: "always", + flushPendingBashBeforeValidation: true, + validateModelAndAuth: true, + awaitPendingModelSelection: true, + preTurnCompaction: "beforeModelSelection", + finalRefineBarrier: "ifInFlight", + }, + runBeforeAgentStart: true, + nextTurnContextTiming: "preparation", + preserveEmptyExtensionPrompt: true, + completionIncludesRetryChain: true, + }; + } + return { + preparation: { + initialRefineBarrier: "always", + flushPendingBashBeforeValidation: false, + validateModelAndAuth: false, + awaitPendingModelSelection: false, + preTurnCompaction: "skip", + finalRefineBarrier: "skip", + }, + runBeforeAgentStart: false, + nextTurnContextTiming: "skip", + preserveEmptyExtensionPrompt: false, + completionIncludesRetryChain: false, + }; +} + +export interface TurnPreparationHost { + hasRefinement(): boolean; + waitForRefinement(): Promise; + flushPendingBash(): void; + validate(): Promise; + compact(): Promise; + pendingModelSelection(): Promise | undefined; +} + +export class TurnPreparer { + constructor(private readonly _host: TurnPreparationHost) {} + + async prepare( + policy: CommitPreparationPolicy, + steps: CommitPreparationSteps, + ): Promise { + if ( + policy.initialRefineBarrier === "always" || + (policy.initialRefineBarrier === "ifInFlight" && this._host.hasRefinement()) + ) { + await this._host.waitForRefinement(); + } + if (policy.flushPendingBashBeforeValidation) this._host.flushPendingBash(); + if (policy.validateModelAndAuth) await this._host.validate(); + steps.afterValidation?.(); + if (!policy.flushPendingBashBeforeValidation) this._host.flushPendingBash(); + + if (policy.preTurnCompaction === "beforeModelSelection") await this._host.compact(); + if (policy.awaitPendingModelSelection) { + const pendingModelSelectEmit = this._host.pendingModelSelection(); + if (pendingModelSelectEmit) await pendingModelSelectEmit; + } + if (policy.preTurnCompaction === "afterModelSelection") await this._host.compact(); + + const prepared = await steps.prepare(); + if (steps.shouldCommit && !steps.shouldCommit(prepared)) return undefined; + steps.beforeFinalRefineBarrier?.(prepared); + let passedFinalRefineBarrier = false; + if ( + policy.finalRefineBarrier === "always" || + (policy.finalRefineBarrier === "ifInFlight" && this._host.hasRefinement()) + ) { + await this._host.waitForRefinement(); + passedFinalRefineBarrier = true; + } + return steps.commit(prepared, passedFinalRefineBarrier); + } +} diff --git a/packages/coding-agent/src/utils/wait-for-abort.ts b/packages/coding-agent/src/utils/wait-for-abort.ts new file mode 100644 index 0000000000..fc9c5b83ae --- /dev/null +++ b/packages/coding-agent/src/utils/wait-for-abort.ts @@ -0,0 +1,28 @@ +export function waitForPromiseOrAbort( + promise: Promise, + signal: AbortSignal | undefined, + abortMessage: string, +): Promise { + if (!signal) return promise; + if (signal.aborted) return Promise.reject(new Error(abortMessage)); + return new Promise((resolve, reject) => { + const onAbort = () => { + cleanup(); + reject(new Error(abortMessage)); + }; + const cleanup = () => signal.removeEventListener("abort", onAbort); + signal.addEventListener("abort", onAbort, { once: true }); + // Close the listener-registration race before observing the awaited work. + if (signal.aborted) return onAbort(); + promise.then( + (value) => { + cleanup(); + resolve(value); + }, + (error: unknown) => { + cleanup(); + reject(error); + }, + ); + }); +} diff --git a/packages/coding-agent/test/goal-continuation-quiescence.test.ts b/packages/coding-agent/test/goal-continuation-quiescence.test.ts index 1ab9434a8c..091f5684a3 100644 --- a/packages/coding-agent/test/goal-continuation-quiescence.test.ts +++ b/packages/coding-agent/test/goal-continuation-quiescence.test.ts @@ -2,14 +2,14 @@ import { describe, expect, it, vi } from "vitest"; import { AgentSession } from "../src/core/agent-session.js"; import { emptyGoalState } from "../src/core/goals.js"; import { GoalController } from "../src/goals/controller.js"; +import { SessionInputScheduler } from "../src/session/input-scheduler.js"; type Harness = { _goals: GoalController; _goalContinuationAwaitsRlmWork: boolean; _disposed: boolean; _disposing: boolean; - _sessionInputAdmissionPauses: Set; - _sessionInputPumpSuspended: boolean; + _inputScheduler: SessionInputScheduler; _hasUnsettledRlmQuiescenceWork: () => boolean; _stopGoalContinuationForTerminalMessage: () => boolean; _ensureGoalRuntimeActive: () => void; @@ -33,8 +33,7 @@ function harness(overrides: Partial = {}): Harness { _goalContinuationAwaitsRlmWork: false, _disposed: false, _disposing: false, - _sessionInputAdmissionPauses: new Set(), - _sessionInputPumpSuspended: false, + _inputScheduler: new SessionInputScheduler({ canSchedule: () => false, run: async () => {} }), _hasUnsettledRlmQuiescenceWork: () => false, _stopGoalContinuationForTerminalMessage: () => false, _ensureGoalRuntimeActive: () => {}, @@ -79,20 +78,21 @@ describe("goal continuation vs unsettled subagent work", () => { it("keeps the deferral while admission is paused and retries after release", () => { const paused = harness({ _goalContinuationAwaitsRlmWork: true, - _sessionInputAdmissionPauses: new Set([Symbol("pause")]), }); + const pause = paused._inputScheduler.acquireAdmissionPause(() => {}); maybeResume.call(paused); expect(paused._admitSessionInput).not.toHaveBeenCalled(); expect(paused._goalContinuationAwaitsRlmWork).toBe(true); - paused._sessionInputAdmissionPauses.clear(); + pause.release(); maybeResume.call(paused); expect(paused._admitSessionInput).toHaveBeenCalledTimes(1); expect(paused._goalContinuationAwaitsRlmWork).toBe(false); }); it("keeps the deferral while the pump is suspended after an abort", () => { - const mode = harness({ _goalContinuationAwaitsRlmWork: true, _sessionInputPumpSuspended: true }); + const mode = harness({ _goalContinuationAwaitsRlmWork: true }); + mode._inputScheduler.suspend("abort"); maybeResume.call(mode); expect(mode._admitSessionInput).not.toHaveBeenCalled(); expect(mode._goalContinuationAwaitsRlmWork).toBe(true); diff --git a/packages/coding-agent/test/session/bash.test.ts b/packages/coding-agent/test/session/bash.test.ts new file mode 100644 index 0000000000..236d921c29 --- /dev/null +++ b/packages/coding-agent/test/session/bash.test.ts @@ -0,0 +1,206 @@ +import { Buffer } from "node:buffer"; +import { describe, expect, it, vi } from "vitest"; +import type { BashResult } from "../../src/core/bash-executor.js"; +import type { BashExecutionMessage } from "../../src/core/messages.js"; +import { SessionBash, type SessionBashEvent, type SessionBashHost } from "../../src/session/bash.js"; + +function deferred() { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +function result(output = "done"): BashResult { + return { output, exitCode: 0, cancelled: false, truncated: false }; +} + +function createShell(overrides: Partial = {}) { + const events: SessionBashEvent[] = []; + const messages: BashExecutionMessage[] = []; + const onStateChange = vi.fn(); + const onUserBashEnd = vi.fn(async () => {}); + const host: SessionBashHost = { + getCwd: () => "/workspace", + getShellCommandPrefix: () => undefined, + getShellPath: () => undefined, + isStreaming: () => false, + intercept: async () => undefined, + emit: (event) => events.push(event), + appendMessage: (message) => messages.push(message), + onStateChange, + onUserBashEnd, + executeBash: (command, onChunk, options) => shell.executeBash(command, onChunk, options), + recordBashResult: (command, outcome, options) => shell.recordBashResult(command, outcome, options), + ...overrides, + }; + const shell = new SessionBash(host); + return { shell, events, messages, onStateChange, onUserBashEnd }; +} + +describe("SessionBash", () => { + it("reads current shell settings and cwd while recording the original command", async () => { + let cwd = "/first"; + let prefix = "set -e"; + const { shell, messages } = createShell({ getCwd: () => cwd, getShellCommandPrefix: () => prefix }); + const calls: string[][] = []; + const chunks: string[] = []; + const options = { + excludeFromContext: true, + operations: { + exec: vi.fn(async (command, directory, execution) => { + calls.push([command, directory]); + execution.onData(Buffer.from("output")); + return { exitCode: 0 }; + }), + }, + } satisfies Parameters[2]; + await shell.executeBash("first", (chunk) => chunks.push(chunk), options); + cwd = "/second"; + prefix = ""; + await shell.executeBash("second", undefined, options); + expect(calls).toEqual([ + ["set -e\nfirst", "/first"], + ["second", "/second"], + ]); + expect(chunks).toEqual(["output"]); + expect(messages.map((message) => [message.command, message.excludeFromContext])).toEqual([ + ["first", true], + ["second", true], + ]); + }); + + it("holds the user slot during interception and honours abort before execution", async () => { + const gate = deferred(); + const exec = vi.fn(async () => ({ exitCode: 0 })); + const { shell, events, messages } = createShell({ + intercept: async () => { + await gate.promise; + return { operations: { exec } }; + }, + }); + const run = shell.runUserBash("cancel me", { runId: "first" }); + expect(shell.isBashRunning).toBe(true); + await expect(shell.runUserBash("second")).rejects.toThrow("already running"); + shell.abortBash(); + gate.resolve(); + await run; + expect(exec).not.toHaveBeenCalled(); + expect(messages).toMatchObject([{ command: "cancel me", cancelled: true, output: "" }]); + expect(events).toEqual([ + { type: "bash_start", command: "cancel me", excludeFromContext: false, runId: "first" }, + { type: "bash_end", exitCode: undefined, cancelled: true, truncated: false, runId: "first" }, + ]); + expect(shell.isBashRunning).toBe(false); + }); + + it("keeps an extension replacement result authoritative when abort arrives during interception", async () => { + const gate = deferred(); + const exec = vi.fn(async () => ({ exitCode: 1 })); + const replacement = { ...result("replacement"), truncated: true, fullOutputPath: "/output/complete" }; + const { shell, events, messages } = createShell({ + intercept: async () => { + await gate.promise; + return { result: replacement, operations: { exec } }; + }, + }); + const run = shell.runUserBash("intercepted"); + shell.abortBash(); + gate.resolve(); + await run; + expect(exec).not.toHaveBeenCalled(); + expect(messages).toMatchObject([replacement]); + expect(events.at(-1)).toMatchObject({ + type: "bash_end", + cancelled: false, + truncated: true, + fullOutputPath: "/output/complete", + }); + }); + + it("releases the slot on interception failure without fabricating execution events", async () => { + let fail = true; + const { shell, events, messages, onStateChange, onUserBashEnd } = createShell({ + intercept: async () => { + if (fail) throw new Error("dispatch failed"); + return { result: result() }; + }, + }); + await expect(shell.runUserBash("first")).rejects.toThrow("dispatch failed"); + expect(shell.isBashRunning).toBe(false); + expect(onStateChange).toHaveBeenCalledOnce(); + expect(onUserBashEnd).not.toHaveBeenCalled(); + expect(events).toEqual([]); + expect(messages).toEqual([]); + fail = false; + await shell.runUserBash("second"); + expect(messages).toHaveLength(1); + }); + + it.each(["execution", "replacement", "failure"] as const)( + "keeps transient %s output out of pending messages and persistence", + async (mode) => { + const { shell, events, messages } = createShell({ + isStreaming: () => true, + intercept: async () => + mode === "replacement" + ? { result: result() } + : { + operations: { + exec: async (_command, _cwd, options) => { + if (mode === "failure") throw new Error("failed"); + options.onData(Buffer.from("transient")); + return { exitCode: 0 }; + }, + }, + }, + }); + await shell.runUserBash("side command", { transient: true, runId: "side" }); + expect(events[0]).toMatchObject({ type: "bash_start", transient: true, runId: "side" }); + expect(events.at(-1)).toMatchObject({ type: "bash_end", transient: true, runId: "side" }); + expect(messages).toEqual([]); + expect(shell.hasPendingBashMessages).toBe(false); + shell.flushPendingMessages(); + expect(messages).toEqual([]); + }, + ); + + it("flushes deferred output in order only at the explicit flush boundary", () => { + let streaming = true; + const { shell, messages } = createShell({ isStreaming: () => streaming }); + shell.recordBashResult("first", result("one")); + shell.recordBashResult("second", result("two"), { excludeFromContext: true }); + streaming = false; + expect(messages).toEqual([]); + expect(shell.hasPendingBashMessages).toBe(true); + shell.flushPendingMessages(); + shell.flushPendingMessages(); + expect(messages.map((message) => message.command)).toEqual(["first", "second"]); + expect(messages[1]?.excludeFromContext).toBe(true); + expect(shell.hasPendingBashMessages).toBe(false); + }); + + it("allows a completion subscriber to synchronously start another user command", async () => { + const events: SessionBashEvent[] = []; + let secondRun: Promise | undefined; + const { shell, messages } = createShell({ + intercept: async (event) => ({ result: result(event.command) }), + emit: (event) => { + events.push(event); + if (event.type === "bash_end" && event.runId === "first") { + expect(shell.isBashRunning).toBe(false); + secondRun = shell.runUserBash("second", { runId: "second" }); + } + }, + }); + await shell.runUserBash("first", { runId: "first" }); + await secondRun; + expect(messages.map((message) => message.command)).toEqual(["first", "second"]); + expect(events.filter((event) => event.type === "bash_end").map((event) => event.runId)).toEqual([ + "first", + "second", + ]); + expect(shell.isBashRunning).toBe(false); + }); +}); diff --git a/packages/coding-agent/test/session/commit-fence.test.ts b/packages/coding-agent/test/session/commit-fence.test.ts new file mode 100644 index 0000000000..731c29086a --- /dev/null +++ b/packages/coding-agent/test/session/commit-fence.test.ts @@ -0,0 +1,184 @@ +import { describe, expect, it } from "vitest"; +import { SessionCommitFence } from "../../src/session/commit-fence.js"; + +function deferred() { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +function yieldToEventLoop(): Promise { + return new Promise((resolve) => setImmediate(resolve)); +} + +describe("SessionCommitFence", () => { + it("grants leases in arrival order and reports pending work between owners", async () => { + const fence = new SessionCommitFence(); + const first = await fence.acquire(); + const order: string[] = []; + const secondPromise = fence.acquire().then((lease) => { + order.push("second"); + return lease; + }); + const thirdPromise = fence.acquire().then((lease) => { + order.push("third"); + return lease; + }); + await yieldToEventLoop(); + expect(order).toEqual([]); + first.release(); + expect(fence.hasPendingWork).toBe(true); + const second = await secondPromise; + expect(order).toEqual(["second"]); + first.release(); + await yieldToEventLoop(); + expect(order).toEqual(["second"]); + second.release(); + const third = await thirdPromise; + expect(order).toEqual(["second", "third"]); + third.release(); + expect(fence.hasPendingWork).toBe(false); + }); + + it("cancels a middle waiter without allowing later work to overtake its predecessor", async () => { + const fence = new SessionCommitFence(); + const first = await fence.acquire(); + const controller = new AbortController(); + const cancelled = fence.acquire(controller.signal); + const rejection = expect(cancelled).rejects.toThrow("Update restart preparation cancelled"); + let lastAcquired = false; + const lastPromise = fence.acquire().then((lease) => { + lastAcquired = true; + return lease; + }); + controller.abort(); + await rejection; + await yieldToEventLoop(); + expect(lastAcquired).toBe(false); + expect(fence.hasPendingWork).toBe(true); + first.release(); + const last = await lastPromise; + last.release(); + expect(fence.hasPendingWork).toBe(false); + }); + + it("rejects pre-aborted acquisition without blocking subsequent work", async () => { + const fence = new SessionCommitFence(); + const controller = new AbortController(); + controller.abort(); + await expect(fence.acquire(controller.signal)).rejects.toThrow("Update restart preparation cancelled"); + expect(fence.hasPendingWork).toBe(false); + const lease = await fence.acquire(); + lease.release(); + expect(fence.hasPendingWork).toBe(false); + }); + + it("re-enters through an asynchronous hook without releasing the outer lease", async () => { + const fence = new SessionCommitFence(); + const outer = await fence.acquire(); + let unrelatedAcquired = false; + const unrelated = fence.acquire().then((lease) => { + unrelatedAcquired = true; + return lease; + }); + await fence.run(outer, async () => { + await yieldToEventLoop(); + expect(fence.isHeldByCurrentContext).toBe(true); + const nested = await fence.acquire(); + expect(nested.owner).toBe(outer.owner); + nested.release(); + await yieldToEventLoop(); + expect(unrelatedAcquired).toBe(false); + expect(fence.hasPendingWork).toBe(true); + }); + expect(fence.isHeldByCurrentContext).toBe(false); + outer.release(); + const next = await unrelated; + next.release(); + }); + + it("queues a stale asynchronous context behind the current owner", async () => { + const fence = new SessionCommitFence(); + const first = await fence.acquire(); + const delayedHook = deferred(); + let staleAcquired = false; + const stale = fence.run(first, async () => { + await delayedHook.promise; + expect(fence.isHeldByCurrentContext).toBe(false); + const lease = await fence.acquire(); + staleAcquired = true; + return lease; + }); + first.release(); + const second = await fence.acquire(); + delayedHook.resolve(); + await yieldToEventLoop(); + expect(staleAcquired).toBe(false); + second.release(); + const third = await stale; + expect(third.owner).not.toBe(first.owner); + expect(third.owner).not.toBe(second.owner); + third.release(); + }); + + it("retains ownership when the acquiring caller's signal is aborted after acquisition", async () => { + const fence = new SessionCommitFence(); + const controller = new AbortController(); + const first = await fence.acquire(controller.signal); + controller.abort(); + let nextAcquired = false; + const next = fence.acquire().then((lease) => { + nextAcquired = true; + return lease; + }); + await yieldToEventLoop(); + expect(nextAcquired).toBe(false); + first.release(); + (await next).release(); + }); + + it("rejects pending and future acquisitions on disposal without revoking a held lease", async () => { + const fence = new SessionCommitFence(); + const first = await fence.acquire(); + const second = expect(fence.acquire()).rejects.toThrow("session is disposing or disposed"); + const third = expect(fence.acquire()).rejects.toThrow("session is disposing or disposed"); + fence.dispose(); + fence.dispose(); + await Promise.all([second, third]); + expect(fence.disposeSignal.aborted).toBe(true); + expect(fence.hasPendingWork).toBe(true); + await expect(fence.acquire()).rejects.toThrow("session is disposing or disposed"); + first.release(); + expect(fence.hasPendingWork).toBe(false); + }); + + it("preserves reentrant acquisition for an existing owner even when its wait signal is aborted", async () => { + const fence = new SessionCommitFence(); + const outer = await fence.acquire(); + const controller = new AbortController(); + controller.abort(); + await fence.run(outer, async () => { + const nested = await fence.acquire(controller.signal); + expect(nested.owner).toBe(outer.owner); + nested.release(); + }); + outer.release(); + }); + + it("restores the caller's context after a hook throws and allows explicit cleanup", async () => { + const fence = new SessionCommitFence(); + const lease = await fence.acquire(); + expect(() => + fence.run(lease, () => { + throw new Error("hook failed"); + }), + ).toThrow("hook failed"); + expect(fence.isHeldByCurrentContext).toBe(false); + expect(fence.hasPendingWork).toBe(true); + lease.release(); + (await fence.acquire()).release(); + expect(fence.hasPendingWork).toBe(false); + }); +}); diff --git a/packages/coding-agent/test/session/input-dispatcher.test.ts b/packages/coding-agent/test/session/input-dispatcher.test.ts new file mode 100644 index 0000000000..19e1679477 --- /dev/null +++ b/packages/coding-agent/test/session/input-dispatcher.test.ts @@ -0,0 +1,215 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { describe, expect, it, vi } from "vitest"; +import { ActionStore, transitionSessionAction } from "../../src/core/session-action-store.js"; +import { SessionInputDispatcher, type SessionInputDispatcherHost } from "../../src/session/input-dispatcher.js"; +import { + createDeliveryRecord, + createPreparedTurnAction, + DeferredSessionInputError, + primaryDeliveryRecord, + type QueuedSessionAction, +} from "../../src/session/prepared-actions.js"; +import { createTurnExecutionPolicy } from "../../src/session/turn-preparation.js"; + +function createFixture() { + const actions = new ActionStore(); + const state = { epoch: 0, busy: false, transcript: [] as AgentMessage[] }; + const host = { + isDisposed: () => false, + getEpoch: () => state.epoch, + getActivity: () => ({ + lowerAgentRun: false, + compaction: false, + retry: false, + bash: state.busy, + refinementApply: false, + branchMutation: false, + schedulerPauseCount: 0, + disposing: false, + }), + isBusy: () => state.busy, + isHandoffDeferred: (epoch: number) => epoch !== state.epoch || state.busy, + getDeliveryMode: () => "all" as const, + waitForAgentIdle: vi.fn(async () => {}), + hasCancelledDispatchCapture: () => false, + getEventQueue: () => Promise.resolve(), + waitForRefinement: vi.fn(async () => {}), + getTranscript: () => state.transcript, + startTurns: vi.fn(async (batch: QueuedSessionAction[], _epoch: number) => { + for (const action of batch) { + transitionSessionAction(action, { state: "committing" }); + state.transcript.push(primaryDeliveryRecord(action).message); + actions.ticketFor(action).settleDelivered({ status: "delivered" }); + } + }), + executeCommand: vi.fn(async () => {}), + settleAgentMessage: vi.fn(), + releaseTurn: vi.fn(), + notifyCheckpoints: vi.fn(), + emitQueueUpdate: vi.fn(), + surfaceError: vi.fn(), + schedule: vi.fn(), + } satisfies SessionInputDispatcherHost; + const dispatcher = new SessionInputDispatcher(actions, host); + const enqueue = (text: string, options: Parameters[3] = {}) => { + const action = createPreparedTurnAction("followUp", text, undefined, options); + actions.enqueue(action); + return action; + }; + return { actions, state, host, dispatcher, enqueue }; +} + +describe("session input dispatch", () => { + it("batches only adjacent turns with matching execution policies", async () => { + const { enqueue, dispatcher, host, actions } = createFixture(); + const first = enqueue("first"); + const second = enqueue("second"); + const direct = enqueue("direct", { executionPolicy: createTurnExecutionPolicy("directPrompt") }); + const last = enqueue("last"); + const completions = [first, second, direct, last].map((action) => actions.ticketFor(action).ticket.completed); + + await dispatcher.run(0); + await Promise.all(completions); + + expect(host.startTurns.mock.calls.map(([batch]) => batch.map((action) => action.payload.text))).toEqual([ + ["first", "second"], + ["direct"], + ["last"], + ]); + expect(actions.ownedActions()).toEqual([]); + }); + + it("dispatches a preselected turn separately even in all mode", async () => { + const { enqueue, dispatcher, host, actions } = createFixture(); + const first = enqueue("already selected"); + actions.selectFirst(); + const second = enqueue("queued later"); + + await dispatcher.run(0); + + expect(host.startTurns.mock.calls.map(([batch]) => batch)).toEqual([[first], [second]]); + }); + + it("rolls selection back when the epoch changes while waiting for the agent", async () => { + const { enqueue, dispatcher, host, actions, state } = createFixture(); + const action = enqueue("selected before pause"); + actions.selectFirst(); + let release = () => {}; + host.waitForAgentIdle.mockReturnValueOnce( + new Promise((resolve) => { + release = resolve; + }), + ); + + const dispatch = dispatcher.run(0); + state.epoch++; + release(); + await dispatch; + + expect(action.lifecycle.state).toBe("queued"); + expect(host.startTurns).not.toHaveBeenCalled(); + expect(host.notifyCheckpoints).toHaveBeenCalledOnce(); + expect(host.emitQueueUpdate).toHaveBeenCalledOnce(); + expect(host.schedule).not.toHaveBeenCalled(); + }); + + it("requeues an undelivered primary without replaying its durable prefix", async () => { + const { enqueue, dispatcher, host, actions, state } = createFixture(); + const prefix = { + role: "custom" as const, + customType: "prefix", + content: "context", + display: false, + timestamp: 1, + }; + const action = enqueue("retry delivery", { prefixMessages: [prefix] }); + if (action.payload.kind !== "turn") throw new Error("Expected turn"); + const durablePrefix = action.payload.records[0].message; + const nextTurn = createDeliveryRecord(action.id, "next_turn", { ...prefix, customType: "next-turn" }); + action.payload.records.splice(1, 0, nextTurn); + host.startTurns.mockImplementationOnce(async () => { + transitionSessionAction(action, { state: "committing" }); + state.transcript.push(durablePrefix); + state.epoch++; + throw new DeferredSessionInputError("paused at handoff"); + }); + const completed = actions.ticketFor(action).ticket.completed; + + await dispatcher.run(0); + + expect(action.lifecycle.state).toBe("queued"); + expect(action.payload.records.map((record) => record.role)).toEqual(["primary"]); + expect(host.releaseTurn).not.toHaveBeenCalled(); + expect(host.surfaceError).not.toHaveBeenCalled(); + expect(host.schedule).not.toHaveBeenCalled(); + + await dispatcher.run(state.epoch); + await completed; + expect(state.transcript.filter((message) => message === durablePrefix)).toHaveLength(1); + expect(state.transcript.filter((message) => message === primaryDeliveryRecord(action).message)).toHaveLength(1); + }); + + it("preserves completed delivery while rejecting completion after a partial batch failure", async () => { + const { enqueue, dispatcher, host, actions, state } = createFixture(); + const first = enqueue("delivered", { agentMessageId: "first" }); + const second = enqueue("not delivered", { agentMessageId: "second" }); + const firstTicket = actions.ticketFor(first).ticket; + const secondTicket = actions.ticketFor(second).ticket; + const failure = new Error("dispatch failed"); + host.startTurns.mockImplementationOnce(async (batch) => { + for (const action of batch) transitionSessionAction(action, { state: "committing" }); + state.transcript.push(primaryDeliveryRecord(first).message); + actions.ticketFor(first).settleDelivered({ status: "delivered" }); + throw failure; + }); + + await dispatcher.run(0); + + await expect(firstTicket.delivered).resolves.toEqual({ status: "delivered" }); + await expect(secondTicket.delivered).rejects.toBe(failure); + await expect(firstTicket.completed).rejects.toBe(failure); + await expect(secondTicket.completed).rejects.toBe(failure); + expect(host.settleAgentMessage.mock.calls).toEqual([ + ["first", "completion", failure], + ["second", "delivery", failure], + ["second", "completion", failure], + ]); + expect(host.surfaceError).toHaveBeenCalledExactlyOnceWith(failure); + expect(actions.ownedActions()).toEqual([]); + }); + + it.each([false, true])( + "retains cancelled work only when capturing late dispatch messages: %s", + async (capturing) => { + const { enqueue, dispatcher, host, actions } = createFixture(); + const action = enqueue("cancelled"); + host.startTurns.mockImplementationOnce(async () => { + if (action.payload.kind !== "turn") throw new Error("Expected turn"); + if (capturing) action.payload.captureRunMessages = new Set(); + transitionSessionAction(action, { state: "cancelled" }); + }); + + await dispatcher.run(0); + + expect(actions.ownedActions()).toEqual(capturing ? [action] : []); + expect(host.releaseTurn.mock.calls).toEqual(capturing ? [] : [[action.id]]); + }, + ); + + it("surfaces a preparation error once and preserves queued work when runtime becomes busy", async () => { + const { enqueue, dispatcher, host, state } = createFixture(); + const action = enqueue("deferred by shell"); + const failure = new Error("preparation failed while busy"); + host.startTurns.mockImplementationOnce(async () => { + state.busy = true; + throw failure; + }); + + await dispatcher.run(0); + + expect(action.lifecycle.state).toBe("queued"); + expect(host.surfaceError).toHaveBeenCalledExactlyOnceWith(failure); + expect(host.schedule).not.toHaveBeenCalled(); + expect(host.releaseTurn).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/coding-agent/test/session/input-scheduler.test.ts b/packages/coding-agent/test/session/input-scheduler.test.ts new file mode 100644 index 0000000000..05829c3f6a --- /dev/null +++ b/packages/coding-agent/test/session/input-scheduler.test.ts @@ -0,0 +1,166 @@ +import { describe, expect, it, vi } from "vitest"; +import { SessionInputScheduler } from "../../src/session/input-scheduler.js"; + +function deferred() { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describe("SessionInputScheduler", () => { + it("coalesces pending requests and serializes work scheduled during a run", async () => { + const firstRun = deferred(); + const firstStarted = deferred(); + const secondRun = deferred(); + const secondStarted = deferred(); + const run = vi.fn(async () => { + if (run.mock.calls.length === 1) { + firstStarted.resolve(); + await firstRun.promise; + } else { + secondStarted.resolve(); + await secondRun.promise; + } + }); + const scheduler = new SessionInputScheduler({ canSchedule: () => true, run }); + scheduler.schedule(); + scheduler.schedule(); + await firstStarted.promise; + expect(run).toHaveBeenCalledTimes(1); + let idle = false; + const waiting = scheduler.waitForIdle().then(() => { + idle = true; + }); + scheduler.schedule(); + scheduler.schedule(); + await Promise.resolve(); + expect(run).toHaveBeenCalledTimes(1); + firstRun.resolve(); + await secondStarted.promise; + expect(idle).toBe(false); + secondRun.resolve(); + await waiting; + expect(run).toHaveBeenCalledTimes(2); + }); + + it("checks session eligibility before scheduling", async () => { + let eligible = false; + const run = vi.fn(async () => {}); + const scheduler = new SessionInputScheduler({ canSchedule: () => eligible, run }); + scheduler.schedule(); + await scheduler.waitForIdle(); + expect(run).not.toHaveBeenCalled(); + eligible = true; + scheduler.schedule(); + await scheduler.waitForIdle(); + expect(run).toHaveBeenCalledOnce(); + }); + + it("keeps queued work paused until every lease releases, including duplicate releases", async () => { + const run = vi.fn(async () => {}); + const scheduler = new SessionInputScheduler({ canSchedule: () => true, run }); + const onRelease = vi.fn(() => scheduler.schedule()); + const first = scheduler.acquireQueuedWorkPause(onRelease); + const second = scheduler.acquireQueuedWorkPause(onRelease); + first.release(); + first.release(); + await scheduler.waitForIdle(); + expect(scheduler.queuedWorkPauseCount).toBe(1); + expect(run).not.toHaveBeenCalled(); + second.release(); + second.release(); + await scheduler.waitForIdle(); + expect(scheduler.queuedWorkPauseCount).toBe(0); + expect(onRelease).toHaveBeenCalledTimes(2); + expect(run).toHaveBeenCalledOnce(); + }); + + it("keeps overlapping admission pauses independent from already admitted work", async () => { + const run = vi.fn(async () => {}); + const scheduler = new SessionInputScheduler({ canSchedule: () => true, run }); + const first = scheduler.acquireAdmissionPause(() => scheduler.schedule()); + const second = scheduler.acquireAdmissionPause(() => scheduler.schedule()); + first.release(); + first.release(); + expect(scheduler.admissionPaused).toBe(true); + await scheduler.waitForIdle(); + expect(run).toHaveBeenCalledOnce(); + second.release(); + expect(scheduler.admissionPaused).toBe(false); + await scheduler.waitForIdle(); + }); + + it("invalidates preparation even when an admission pause ends before preparation resolves", async () => { + const prepared = deferred(); + const started = deferred(); + const delivered = vi.fn(); + const scheduler = new SessionInputScheduler({ + canSchedule: () => true, + run: async (epoch) => { + started.resolve(); + await prepared.promise; + if (epoch === scheduler.epoch) delivered(); + }, + }); + scheduler.schedule(); + await started.promise; + const pause = scheduler.acquireAdmissionPause(() => {}); + pause.release(); + prepared.resolve(); + await scheduler.waitForIdle(); + expect(delivered).not.toHaveBeenCalled(); + scheduler.schedule(); + await scheduler.waitForIdle(); + expect(delivered).toHaveBeenCalledOnce(); + }); + + it.each(["abort", "update-restart"] as const)( + "keeps %s suspension until resumed and retains outstanding pause leases", + async (reason) => { + const run = vi.fn(async () => {}); + const scheduler = new SessionInputScheduler({ canSchedule: () => true, run }); + const pause = scheduler.acquireQueuedWorkPause(() => scheduler.schedule()); + const admission = scheduler.acquireAdmissionPause(() => scheduler.schedule()); + scheduler.suspend(reason); + admission.release(); + await scheduler.waitForIdle(); + expect(scheduler.suspended).toBe(true); + expect(scheduler.suspendedForUpdateRestart).toBe(reason === "update-restart"); + expect(run).not.toHaveBeenCalled(); + expect(scheduler.resume()).toBe(true); + expect(scheduler.suspendedForUpdateRestart).toBe(false); + expect(scheduler.resume()).toBe(false); + scheduler.schedule(); + await scheduler.waitForIdle(); + expect(run).not.toHaveBeenCalled(); + pause.release(); + await scheduler.waitForIdle(); + expect(run).toHaveBeenCalledOnce(); + }, + ); + + it("preserves the captured epoch so a pending runner can reject work invalidated by abort", async () => { + const run = vi.fn(async () => {}); + const scheduler = new SessionInputScheduler({ canSchedule: () => true, run }); + const epoch = scheduler.epoch; + scheduler.schedule(); + scheduler.suspend("abort"); + scheduler.schedule(); + await scheduler.waitForIdle(); + expect(run).toHaveBeenCalledExactlyOnceWith(epoch); + expect(scheduler.epoch).not.toBe(epoch); + expect(scheduler.suspended).toBe(true); + }); + + it("lets idle callers observe failure and schedules later work after a failed run", async () => { + const run = vi.fn(async () => {}).mockRejectedValueOnce(new Error("preparation failed")); + const scheduler = new SessionInputScheduler({ canSchedule: () => true, run }); + scheduler.schedule(); + await expect(scheduler.waitForIdle()).rejects.toThrow("preparation failed"); + scheduler.schedule(); + await expect(scheduler.waitForIdle()).resolves.toBeUndefined(); + expect(run).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/coding-agent/test/session/retry.test.ts b/packages/coding-agent/test/session/retry.test.ts new file mode 100644 index 0000000000..9d958c52a6 --- /dev/null +++ b/packages/coding-agent/test/session/retry.test.ts @@ -0,0 +1,232 @@ +import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { AuthSourceToken } from "../../src/core/auth-storage.js"; +import { SessionRetry, type SessionRetryEvent, type SessionRetryHost } from "../../src/session/retry.js"; + +function failure(kind?: string): AssistantMessage { + return { + ...fauxAssistantMessage("", { stopReason: "error", errorMessage: "provider failed" }), + diagnostics: kind ? [{ type: "provider_stream_failure", timestamp: 0, details: { kind } }] : [], + }; +} + +function setup() { + const events: SessionRetryEvent[] = []; + const order: string[] = []; + const settings = { enabled: true, maxRetries: 3, baseDelayMs: 10 }; + const host = { + getRetrySettings: () => settings, + getMaxRetryDelayMs: () => 1000, + getContextWindow: () => 10000, + getAuthSource: vi.fn(() => undefined), + markAuthSourceStale: vi.fn((_token: AuthSourceToken) => true), + markAuthStale: vi.fn(() => true), + hasPayloadHooks: vi.fn(() => false), + prepareTurnRetry: vi.fn(() => { + order.push("prepare"); + }), + clearTurnRetry: vi.fn(() => { + order.push("clear"); + }), + removeLastAssistant: vi.fn(() => { + order.push("remove"); + }), + continue: vi.fn(() => Promise.resolve()), + waitForIdle: vi.fn(() => Promise.resolve()), + cancelCompaction: vi.fn(() => { + order.push("cancel-compaction"); + }), + emit: vi.fn((event: SessionRetryEvent) => { + events.push(event); + order.push(event.type); + }), + onResolved: vi.fn(() => { + order.push("resolved"); + }), + } satisfies SessionRetryHost; + return { retry: new SessionRetry(host), host, events, order, settings }; +} + +async function schedule(retry: SessionRetry, message = failure()): Promise { + const pending = retry.retryError(message); + expect(pending).toBeInstanceOf(Promise); + await vi.advanceTimersByTimeAsync(10); + expect(await pending).toBe(true); +} + +function token(value: string): AuthSourceToken { + return { provider: "faux", source: "runtime", identityFingerprint: "identity", valueFingerprint: value }; +} + +describe("SessionRetry ownership", () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(() => { + vi.clearAllTimers(); + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("creates the wait boundary synchronously and waits for lower-agent idle after resolution", async () => { + const { retry, host, order } = setup(); + retry.observeAgentEnd({ type: "agent_end", messages: [failure()] }); + expect(retry.isRetrying).toBe(true); + expect(retry.attempt).toBe(0); + let releaseIdle = () => {}; + host.waitForIdle.mockReturnValue( + new Promise((resolve) => { + releaseIdle = resolve; + }), + ); + let done = false; + const waiting = retry.waitForRetry().then(() => { + done = true; + }); + retry.resolve(); + expect(order).toEqual(["clear", "resolved"]); + expect(host.waitForIdle).not.toHaveBeenCalled(); + await Promise.resolve(); + expect(host.waitForIdle).toHaveBeenCalledOnce(); + expect(done).toBe(false); + releaseIdle(); + await waiting; + expect(done).toBe(true); + retry.resolve(); + expect(host.clearTurnRetry).toHaveBeenCalledTimes(2); + expect(host.onResolved).toHaveBeenCalledOnce(); + }); + + it.each(["invalid_request", "refusal", "permission"])("falls through synchronously for %s", (kind) => { + const { retry, host } = setup(); + retry.observeAgentEnd({ type: "agent_end", messages: [failure(kind)] }); + expect(retry.retryError(failure(kind))).toBeUndefined(); + expect(retry.isRetrying).toBe(false); + expect(host.prepareTurnRetry).not.toHaveBeenCalled(); + }); + + it("retains the promise after an aborted assistant resets the attempt as a success", async () => { + const { retry, events } = setup(); + await schedule(retry); + retry.observeAssistantEnd(fauxAssistantMessage("", { stopReason: "aborted" })); + expect(retry.attempt).toBe(0); + expect(retry.isRetrying).toBe(true); + expect(events.at(-1)).toEqual({ type: "auto_retry_end", success: true, attempt: 1 }); + retry.resolve(); + }); + + it("preserves captured source identity and auth/event order on backoff cancellation", async () => { + const { retry, host, events, order } = setup(); + const old = token("old"); + host.getAuthSource.mockReturnValue(old); + const message = failure("auth"); + retry.observeAgentEnd({ type: "agent_end", messages: [message] }); + retry.observeAssistantEnd(message); + const pending = retry.retryError(message); + host.getAuthSource.mockReturnValue(token("fresh")); + retry.abortRetry(); + expect(retry.attempt).toBe(1); + expect(retry.isRetrying).toBe(true); + expect(await pending).toBe(false); + expect(host.markAuthSourceStale).toHaveBeenCalledExactlyOnceWith(old); + expect(host.markAuthStale).not.toHaveBeenCalled(); + expect(order).toEqual([ + "prepare", + "auto_retry_start", + "remove", + "auth_stale", + "auto_retry_end", + "clear", + "resolved", + ]); + expect(events.at(-1)).toMatchObject({ attempt: 1, finalError: "Retry cancelled" }); + expect(message.errorMessage).toContain("Run /login"); + }); + + it("deduplicates captured tokens while retaining changed values across failures", async () => { + const { retry, host, events } = setup(); + const old = token("old"), + fresh = token("fresh"); + host.getAuthSource.mockReturnValue(old); + await schedule(retry, failure("auth")); + host.getAuthSource.mockReturnValue(fresh); + const message = failure("auth"); + retry.observeAssistantEnd(message); + retry.observeAssistantEnd(message); + expect(retry.retryError(message)).toBeUndefined(); + retry.finishActiveRetryWithFailure(message); + expect(host.markAuthSourceStale.mock.calls).toEqual([[old], [fresh]]); + expect(events.at(-2)).toMatchObject({ type: "auth_stale", sourceTokens: [old, fresh] }); + retry.resolve(); + }); + + it("clears successful auth history before a later unrelated failure", async () => { + const { retry, host } = setup(); + host.getAuthSource.mockReturnValue(token("old")); + await schedule(retry, failure("auth")); + retry.observeAssistantEnd(fauxAssistantMessage("recovered")); + retry.resolve(); + await schedule(retry); + retry.finishActiveRetryWithFailure(failure("permission")); + expect(host.markAuthSourceStale).not.toHaveBeenCalled(); + retry.resolve(); + }); + + it("leaves a newer retry untouched when an old scheduled continuation rejects", async () => { + const { retry, host, events } = setup(); + let rejectOld = (_error: Error) => {}; + host.continue.mockReturnValueOnce( + new Promise((_resolve, reject) => { + rejectOld = reject; + }), + ); + await schedule(retry); + await vi.advanceTimersByTimeAsync(1); + retry.abortRetry(); + await schedule(retry); + rejectOld(new Error("old failure")); + await Promise.resolve(); + expect(retry.isRetrying).toBe(true); + expect(retry.attempt).toBe(1); + expect(events.filter((event) => event.type === "auto_retry_end")).toHaveLength(1); + retry.abortRetry(); + }); + + it("still invokes an already scheduled continuation after abort, ignoring its rejection", async () => { + const { retry, host, events } = setup(); + host.continue.mockRejectedValue(new Error("late failure")); + await schedule(retry); + expect(host.continue).not.toHaveBeenCalled(); + retry.abortRetry(); + await vi.advanceTimersByTimeAsync(1); + expect(host.continue).toHaveBeenCalledOnce(); + expect(events.filter((event) => event.type === "auto_retry_end")).toEqual([ + { type: "auto_retry_end", success: false, attempt: 1, finalError: "Retry cancelled" }, + ]); + }); + + it("uses live payload hooks and preserves exceptions before emitting the start", async () => { + const { retry, host, events } = setup(); + host.prepareTurnRetry.mockImplementationOnce(() => { + throw new Error("ledger failed"); + }); + await expect(retry.retryError(failure())).rejects.toThrow("ledger failed"); + expect(retry.attempt).toBe(1); + expect(retry.isRetrying).toBe(true); + expect(events).toEqual([]); + retry.abortRetry(); + host.hasPayloadHooks.mockReturnValue(true); + await schedule(retry); + expect(host.prepareTurnRetry).toHaveBeenCalledOnce(); + retry.abortRetry(); + }); + + it("reads disabled settings at the next error without resetting the active attempt early", async () => { + const { retry, settings, events } = setup(); + await schedule(retry); + settings.enabled = false; + expect(await retry.retryError(failure())).toBe(false); + expect(retry.isRetrying).toBe(false); + expect(retry.attempt).toBe(1); + retry.finishActiveRetryWithFailure(failure()); + expect(events.at(-1)).toMatchObject({ type: "auto_retry_end", success: false, attempt: 1 }); + }); +}); diff --git a/packages/coding-agent/test/session/turn-preparation.test.ts b/packages/coding-agent/test/session/turn-preparation.test.ts new file mode 100644 index 0000000000..011b5f08ac --- /dev/null +++ b/packages/coding-agent/test/session/turn-preparation.test.ts @@ -0,0 +1,167 @@ +import { describe, expect, it } from "vitest"; +import { + createTurnExecutionPolicy, + type TurnPreparationHost, + TurnPreparer, +} from "../../src/session/turn-preparation.js"; + +function createPreparation(overrides: Partial = {}) { + const order: string[] = []; + const preparer = new TurnPreparer({ + hasRefinement: () => false, + waitForRefinement: async () => { + order.push("refine"); + }, + flushPendingBash: () => { + order.push("flush"); + }, + validate: async () => { + order.push("validate"); + }, + compact: async () => { + order.push("compact"); + }, + pendingModelSelection: () => { + order.push("model"); + return Promise.resolve(); + }, + ...overrides, + }); + const steps = { + afterValidation: () => { + order.push("after validation"); + }, + prepare: async () => { + order.push("prepare"); + return "prepared"; + }, + beforeFinalRefineBarrier: () => { + order.push("before barrier"); + }, + commit: (value: string, passedBarrier: boolean) => { + order.push(`commit:${passedBarrier}`); + return value; + }, + }; + return { preparer, order, steps }; +} + +describe("TurnPreparer", () => { + it.each([ + { + kind: "queued", + order: [ + "validate", + "after validation", + "flush", + "compact", + "model", + "prepare", + "before barrier", + "refine", + "commit:true", + ], + }, + { + kind: "directPrompt", + order: [ + "refine", + "flush", + "validate", + "after validation", + "model", + "compact", + "prepare", + "before barrier", + "commit:false", + ], + }, + { + kind: "injected", + order: [ + "refine", + "flush", + "validate", + "after validation", + "compact", + "model", + "prepare", + "before barrier", + "commit:false", + ], + }, + { + kind: "customTrigger", + order: ["refine", "after validation", "flush", "prepare", "before barrier", "commit:false"], + }, + ] as const)("preserves $kind preparation order", async ({ kind, order: expected }) => { + const { preparer, order, steps } = createPreparation(); + expect(await preparer.prepare(createTurnExecutionPolicy(kind).preparation, steps)).toBe("prepared"); + expect(order).toEqual(expected); + }); + + it("rechecks conditional refinement after preparation and before commit", async () => { + let refining = false; + const { preparer, order, steps } = createPreparation({ hasRefinement: () => refining }); + await preparer.prepare(createTurnExecutionPolicy("directPrompt", { returnAfterAccepted: true }).preparation, { + ...steps, + beforeFinalRefineBarrier: () => { + refining = true; + }, + }); + expect(order).toEqual([ + "flush", + "validate", + "after validation", + "model", + "compact", + "prepare", + "refine", + "commit:true", + ]); + }); + + it("skips the final barrier and commit when preparation was withdrawn", async () => { + const { preparer, order, steps } = createPreparation(); + const committed = await preparer.prepare(createTurnExecutionPolicy("queued").preparation, { + ...steps, + shouldCommit: () => false, + }); + expect(committed).toBeUndefined(); + expect(order).toEqual(["validate", "after validation", "flush", "compact", "model", "prepare"]); + }); + + it.each(["queued", "directPrompt"] as const)("retains the %s flush boundary when validation fails", async (kind) => { + const { preparer, order, steps } = createPreparation({ + validate: async () => { + throw new Error("auth failed"); + }, + }); + await expect(preparer.prepare(createTurnExecutionPolicy(kind).preparation, steps)).rejects.toThrow("auth failed"); + expect(order).toEqual(kind === "queued" ? [] : ["refine", "flush"]); + }); + + it("does not prepare or compact after pending model selection rejects on a direct prompt", async () => { + const { preparer, order, steps } = createPreparation({ + pendingModelSelection: () => Promise.reject(new Error("selection failed")), + }); + await expect(preparer.prepare(createTurnExecutionPolicy("directPrompt").preparation, steps)).rejects.toThrow( + "selection failed", + ); + expect(order).toEqual(["refine", "flush", "validate", "after validation"]); + }); + + it("awaits the final refinement barrier before committing the prepared value", async () => { + let release = () => {}; + const gate = new Promise((resolve) => { + release = resolve; + }); + const { preparer, order, steps } = createPreparation({ waitForRefinement: () => gate }); + const preparation = preparer.prepare(createTurnExecutionPolicy("queued").preparation, steps); + await new Promise((resolve) => setImmediate(resolve)); + expect(order.at(-1)).toBe("before barrier"); + release(); + expect(await preparation).toBe("prepared"); + expect(order.at(-1)).toBe("commit:true"); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-action-races.test.ts b/packages/coding-agent/test/suite/agent-session-action-races.test.ts index 8d4e29ea65..5f3e911157 100644 --- a/packages/coding-agent/test/suite/agent-session-action-races.test.ts +++ b/packages/coding-agent/test/suite/agent-session-action-races.test.ts @@ -9,7 +9,6 @@ type ActionKind = "turn" | "command"; interface CommitFenceInternals { _actionStore: ActionStore; - _pendingSessionActionFenceWaiters: number; _refineInFlight?: Promise; _scheduleSessionInputPump(): void; _acquireDirectTurnAdmissionFence(signal?: AbortSignal): Promise<{ release(): void }>; @@ -132,7 +131,6 @@ describe("AgentSession action commit-fence races", () => { const internals = harness.session as unknown as CommitFenceInternals; const heldFence = await internals._acquireSessionActionCommitFence(); const nextFencePromise = internals._acquireSessionActionCommitFence(); - await vi.waitFor(() => expect(internals._pendingSessionActionFenceWaiters).toBe(1)); heldFence.release(); expect(harness.session.hasPendingAdmissionWaiters).toBe(true); diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index 65afe62088..9eb28a649e 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -507,7 +507,6 @@ describe("AgentSession compaction characterization", () => { _schedulePostCompactionContinue(): void; _cancelPostCompactionContinue(): void; _sessionInputCheckpointWaiters: Set<() => void>; - _sessionInputPumpSuspended: boolean; }; // A queued follow-up held back by a pause, then a pump suspension (the // requestAbort teardown state): the queue stays populated but undispatchable. @@ -516,7 +515,7 @@ describe("AgentSession compaction characterization", () => { expect(session.queuedActionCount).toBe(1); session.requestAbort(); pause.release(); - expect(internals._sessionInputPumpSuspended).toBe(true); + expect(session.isQueuedWorkSuspended).toBe(true); expect(session.queuedActionCount).toBe(1); // The runner passes its pre-dispatch guards (no pauses, agent idle) and diff --git a/packages/coding-agent/test/suite/agent-session-prompt.test.ts b/packages/coding-agent/test/suite/agent-session-prompt.test.ts index 93255ea85b..88dc4550d7 100644 --- a/packages/coding-agent/test/suite/agent-session-prompt.test.ts +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -1067,7 +1067,6 @@ stale post-hook extension instructions`, }); const sessionInternals = harness.session as unknown as { _refineInFlight?: Promise; - _userBashRunning?: boolean; }; sessionInternals._refineInFlight = refineGate; @@ -1076,12 +1075,26 @@ stale post-hook extension instructions`, { expandPromptTemplates: false, queueIfBusy: true }, ); await vi.waitFor(() => expect(harness.session.getPendingNextTurnMessageSnapshots()).toEqual([])); - sessionInternals._userBashRunning = true; + const bashGate = createDeferred(); + const bashRun = harness.session.executeBash("hold handoff", undefined, { + transient: true, + operations: { + exec: async () => { + await bashGate.promise; + return { exitCode: 0 }; + }, + }, + }); + expect(harness.session.isBashRunning).toBe(true); sessionInternals._refineInFlight = undefined; releaseRefine?.(); - await expect(accepted).rejects.toThrow("Agent became busy before prompt delivery"); - sessionInternals._userBashRunning = false; + try { + await expect(accepted).rejects.toThrow("Agent became busy before prompt delivery"); + } finally { + bashGate.resolve(); + await bashRun; + } let sawRestoredContext = false; harness.setResponses([ diff --git a/packages/coding-agent/test/suite/agent-session-queue.test.ts b/packages/coding-agent/test/suite/agent-session-queue.test.ts index c993bd6582..44b70ebe01 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/packages/coding-agent/test/suite/agent-session-queue.test.ts @@ -2883,19 +2883,12 @@ describe("AgentSession queue characterization", () => { const harness = await createHarness(); harnesses.push(harness); harness.setResponses([fauxAssistantMessage("queued input resumed")]); - const internals = harness.session as unknown as { - _sessionInputPumpRequested: boolean; - _scheduleSessionInputPump(): void; - }; - const schedule = vi.spyOn(internals, "_scheduleSessionInputPump").mockImplementation(() => {}); await harness.session.followUp("queued before pause"); expect(harness.session.getFollowUpMessages()).toEqual(["queued before pause"]); - schedule.mockRestore(); - // Model a pump that was requested before the pause invalidated its epoch. - internals._sessionInputPumpRequested = true; + // Schedule a real pump, then invalidate its epoch before it can run. + expect(harness.session.resumeQueuedWork()).toBe(true); const pause = harness.session.acquireSessionInputPause(); - expect(internals._sessionInputPumpRequested).toBe(false); pause.release(); await harness.session.waitForIdle(); expect(getUserTexts(harness)).toEqual(["queued before pause"]); diff --git a/packages/coding-agent/test/suite/agent-session-retry-events.test.ts b/packages/coding-agent/test/suite/agent-session-retry-events.test.ts index acd0a433e6..99776df104 100644 --- a/packages/coding-agent/test/suite/agent-session-retry-events.test.ts +++ b/packages/coding-agent/test/suite/agent-session-retry-events.test.ts @@ -2,6 +2,7 @@ import { AgentContinueError, type AgentEvent, type AgentTool } from "@earendil-w import { type AssistantMessage, fauxAssistantMessage, fauxThinking, fauxToolCall } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; +import type { SessionRetry } from "../../src/session/retry.js"; import { createHarness, type Harness } from "./harness.js"; function normalizeEventOrder(events: Harness["events"]): string[] { @@ -51,9 +52,7 @@ function rateLimitedFailure(retryAfterMs: number): AssistantMessage { } type SessionRetryCompactionInternals = { - _retryAttempt: number; - _retryPromise: Promise | undefined; - _retryResolve: (() => void) | undefined; + _retry: SessionRetry; _autoCompactionAbortController: AbortController | undefined; _postCompactionContinuationScheduled: boolean; _processAgentEvent: (event: AgentEvent) => Promise; @@ -66,6 +65,7 @@ describe("AgentSession retry and event characterization", () => { const harnesses: Harness[] = []; afterEach(() => { + vi.restoreAllMocks(); while (harnesses.length > 0) { harnesses.pop()?.cleanup(); } @@ -100,15 +100,12 @@ describe("AgentSession retry and event characterization", () => { if (event.type === "auto_retry_start") retryEvents.push(`start:${event.attempt}`); if (event.type === "auto_retry_end") retryEvents.push(`end:${event.success}:${event.finalError}`); }); - harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" })]); + harness.setResponses([structuredProviderFailure("auth")]); vi.spyOn(harness.session.agent, "continue").mockRejectedValue( new AgentContinueError("nothing-to-continue", "Nothing to continue"), ); - const markStale = vi.spyOn( - harness.session as unknown as { _markProviderAuthStaleForRetryFailure: () => void }, - "_markProviderAuthStaleForRetryFailure", - ); + const markStale = vi.spyOn(harness.session.modelRegistry, "markProviderAuthSourceStale"); // Pre-fix this hangs: the swallowed rejection leaves the retry unresolved forever. await harness.session.prompt("test"); @@ -409,16 +406,15 @@ describe("AgentSession retry and event characterization", () => { stopReason: "error", errorMessage: "prompt is too long", }); - internals._retryAttempt = 1; - internals._retryPromise = new Promise((resolve) => { - internals._retryResolve = resolve; - }); + const continuation = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); + await internals._retry.retryError(fauxAssistantMessage("", { stopReason: "error", errorMessage: "transient" })); + await vi.waitFor(() => expect(continuation).toHaveBeenCalledOnce()); internals._checkCompaction = async () => true; try { await internals._processAgentEvent({ type: "agent_end", messages: [overflowMessage] } as AgentEvent); - expect(internals._retryAttempt).toBe(1); + expect(harness.session.retryAttempt).toBe(1); expect(harness.session.isRetrying).toBe(true); expect(harness.eventsOfType("auto_retry_end")).toEqual([]); } finally { @@ -432,10 +428,9 @@ describe("AgentSession retry and event characterization", () => { harnesses.push(harness); const internals = harness.session as unknown as SessionRetryCompactionInternals; const compactionAbortController = new AbortController(); - internals._retryAttempt = 1; - internals._retryPromise = new Promise((resolve) => { - internals._retryResolve = resolve; - }); + const continuation = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); + await internals._retry.retryError(fauxAssistantMessage("", { stopReason: "error", errorMessage: "transient" })); + await vi.waitFor(() => expect(continuation).toHaveBeenCalledOnce()); internals._autoCompactionAbortController = compactionAbortController; internals._schedulePostCompactionContinue(); @@ -446,7 +441,7 @@ describe("AgentSession retry and event characterization", () => { expect(compactionAbortController.signal.aborted).toBe(true); expect(internals._postCompactionContinuationScheduled).toBe(false); - expect(internals._retryAttempt).toBe(0); + expect(harness.session.retryAttempt).toBe(0); expect(harness.session.isRetrying).toBe(false); expect(harness.eventsOfType("auto_retry_end").at(-1)).toMatchObject({ success: false, diff --git a/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts b/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts index 04b65316be..55d014667f 100644 --- a/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts +++ b/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts @@ -3,6 +3,7 @@ import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi- import { afterEach, describe, expect, it, vi } from "vitest"; import type { AgentSessionRuntime } from "../../../src/core/agent-session-runtime.js"; import { InProcessAgentConnection } from "../../../src/modes/agent-connection/in-process-agent-connection.js"; +import type { SessionRetry } from "../../../src/session/retry.js"; import { createHarness, type Harness } from "../harness.js"; function structuredFailureMessage(kind: string, status: number, errorMessage: string): AssistantMessage { @@ -160,12 +161,16 @@ describe("issue #4491 provider stale after repeated 401", () => { harnesses.push(harness); const event = { type: "agent_end", messages: [provider401Message()] } as AgentEvent; const session = harness.session as unknown as { - _retryAttempt: number; - _createRetryPromiseForAgentEnd(event: AgentEvent): void; + _retry: SessionRetry; }; - session._retryAttempt = 1; + const continuation = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); + await session._retry.retryError(provider401Message()); + await vi.waitFor(() => expect(continuation).toHaveBeenCalledOnce()); + session._retry.resolve(); + expect(harness.session.retryAttempt).toBe(1); + expect(harness.session.isRetrying).toBe(false); - session._createRetryPromiseForAgentEnd(event); + session._retry.observeAgentEnd(event); expect(harness.session.isRetrying).toBe(true); harness.session.abortRetry(); @@ -329,11 +334,11 @@ describe("issue #4491 provider stale after repeated 401", () => { const message = provider401Message(); const event = { type: "agent_end", messages: [message] } as AgentEvent; const session = harness.session as unknown as { - _createRetryPromiseForAgentEnd(event: AgentEvent): void; + _retry: SessionRetry; _processAgentEvent(event: AgentEvent): Promise; }; - session._createRetryPromiseForAgentEnd(event); + session._retry.observeAgentEnd(event); await session._processAgentEvent(event); expect(harness.session.isRetrying).toBe(false);