diff --git a/packages/coding-agent/src/README.md b/packages/coding-agent/src/README.md index a4f8c70273..ab7a526a48 100644 --- a/packages/coding-agent/src/README.md +++ b/packages/coding-agent/src/README.md @@ -88,6 +88,26 @@ Preserve the policy differences: direct prompts flush shell output before valida `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. +## Session context + +| File | Responsibility | +| --- | --- | +| `session/compaction.ts` | Manual and automatic compaction lifecycle, pending requests, cancellation, thresholds, and overflow recovery. | +| `session/compaction-execution.ts` | Summary generation, extension interception, request accounting, persistence, and context rebuild ordering. | +| `session/refinement.ts` | Refinement admission, planning and application barriers, serialized plan ownership, and disposal drains. | +| `session/auto-refinement.ts` | Review triggers, cooldowns, pending reviews, timers, and automatic operation cleanup. | +| `session/refinement-execution.ts` | Planning against current dependencies, applying harness edits, and persisting outcomes and notices. | +| `session/continuation.ts` | Resuming work after compaction, settlement, cancellation, and ownership of continuation messages. | + +Each owner keeps its mutable state and cleanup together. Typed host operations connect the owners to current model, authentication, extensions, storage, and scheduling. `AgentSession` composes them and retains public methods, events, and decisions that cross features, including goal and autonomous continuation admission. The summary algorithms and harness storage remain in their existing `core/` feature modules. + +Preserve these boundaries when changing context behavior: + +- Compaction releases its operation and reconnects event handling before resuming work. Summarization request accounting completes before the transcript append; failed persistence must preserve the existing live outcome disclosure. +- Refinement planning may overlap active work. Applying a plan waits for the relevant agent, event, compaction, and branch operations to settle, and rechecks their identities before mutating context. Serialized checkpoints claim a background plan once. +- A cancelled continuation settles its own waiters. Late results cannot clear a replacement operation or consume its messages. The commit lease is released before awaiting a continued agent turn, and checkpoint waiters are removed when cancellation wins. +- Public session methods delegate without additional asynchronous wrappers. Optional waits retain their original positions, and dependency callbacks read current runtime state when invoked. + ## 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`. @@ -96,4 +116,6 @@ Scheduler and commit-fence tests live in `test/session/`. Existing queue, action 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. +Compaction and continuation owner tests in `test/session/` exercise lifecycle and cancellation boundaries. Compaction, refinement, serialized refinement, queue, concurrency, and semantic-edge suites cover their integration with session persistence, goals, and disposal. `test/suite/session-refinement-owner.test.ts` checks the refinement owner boundary using the shared faux-provider harness. + 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 919a636670..b926239ecc 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -6,7 +6,6 @@ import { basename, dirname, join, resolve } from "node:path"; import { Agent, type AgentContext, - AgentContinueError, type AgentEvent, type AgentMessage, type AgentState, @@ -29,7 +28,6 @@ import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, - isContextOverflow, modelsAreEqual, resetApiProviders, supportsFastMode, @@ -45,6 +43,14 @@ import { type SessionBashEvent, } from "../session/bash.js"; import { SessionCommitFence, type SessionCommitLease } from "../session/commit-fence.js"; +import { SessionCompaction, type SessionCompactionEvent } from "../session/compaction.js"; +import { + type CompactionExecutionHost, + type CompactionExecutionOptions, + CompactionSkippedError, + performSessionCompaction, +} from "../session/compaction-execution.js"; +import { type ContinuationToken, SessionContinuation } from "../session/continuation.js"; import { SessionInputDispatcher } from "../session/input-dispatcher.js"; import { SessionInputScheduler } from "../session/input-scheduler.js"; import { @@ -70,6 +76,7 @@ import { type SessionInputSchedule, visibleSessionActionProjection, } from "../session/prepared-actions.js"; +import { type AutoRefineReviewer, SessionRefinement } from "../session/refinement.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"; @@ -135,11 +142,9 @@ import { type CompactionResult, calculateContextTokens, collectEntriesForBranchSummary, - compact, estimateContextTokens, generateBranchSummary, prepareCompaction, - serializeConversation, shouldCompact, } from "./compaction/index.js"; import { @@ -165,8 +170,6 @@ import { type MessageStartEvent, type MessageUpdateEvent, type ReplacedSessionContext, - type SessionBeforeCompactResult, - type SessionBeforeRefineResult, type SessionBeforeTreeResult, type SessionStartEvent, type ShutdownHandler, @@ -199,16 +202,10 @@ import type { McpManager } from "./mcp/mcp-manager.js"; import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE, ASYNC_BASH_COMPLETION_PREVIEW_LABEL, - type CompactionOutcome, - type CompactionOutcomeReason, type CustomMessage, - convertToLlm, createAsyncBashCompletionMessage, - createCompactionOutcomeMessage, createHarnessDigestMessage, createHeartbeatPromptMessage, - createRefinementNoticeMessage, - createRefinementOutcomeMessage, createRlmChildFailureMessage, createRlmChildTerminalNoticeMessage, createSessionSlashCommandMessage, @@ -228,28 +225,11 @@ import { throwIfPromptAdmissionCancelled } from "./prompt-admission.js"; import { expandPromptTemplate, type PromptTemplate, parseCommandArgs } from "./prompt-templates.js"; import { providerRetryPolicy } from "./provider-retry.js"; import { - type AutoRefineReason, - type AutoRefineReview, - appendGlobalRefinement, - applyRefinementProposal, formatHarnessStateForPrompt, - generateRefinementId, getGlobalHarnessStateDir, getLocalHarnessStateDir, - getRefinementHistory, - type HarnessState, - inferRefinementResultScope, - loadGlobalRefinementHistory, - loadHarnessState, - mergeHarnessStates, - mergeRefinementHistory, - normalizeRefinementProposal, - planRefinement, REFINE_SKILL_NAME, - type RefinementPlan, type RefinementResult, - reviewAutoRefine, - saveHarnessState, } from "./refinement/index.js"; import { resolveConfigValue } from "./resolve-config-value.js"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.js"; @@ -275,12 +255,7 @@ import { type RlmSubagentRuntime, type SubagentRuntimeHost, } from "./rlm-runtime.js"; -import { - modelRequestHeaders, - SemanticEdgeRecorder, - semanticEdgeLedgerPath, - wrapStreamFnWithSemanticEdges, -} from "./semantic-edges.js"; +import { SemanticEdgeRecorder, semanticEdgeLedgerPath, wrapStreamFnWithSemanticEdges } from "./semantic-edges.js"; import { ActionStore, type ActionTicket, @@ -299,7 +274,6 @@ import { import type { BranchSummaryEntry, ChildUsageAttributionEntry, - CompactionEntry, SessionContext, SessionEntry, SessionMessageEntry, @@ -367,7 +341,7 @@ export interface RlmChildAgentSnapshot { error?: string; } -export type CompactionReason = "manual" | "threshold" | "overflow" | "requested"; +export type { CompactionReason } from "../session/compaction.js"; export type AgentSessionEvent = | AgentEvent @@ -377,24 +351,10 @@ export type AgentSessionEvent = message: KernelSentAgentMessage; } | { type: "session_action_update"; actions: SessionActionSnapshot } - | { - type: "compaction_start"; - reason: CompactionReason; - customInstructions?: string; - } + | SessionCompactionEvent | { type: "session_info_changed"; name: string | undefined } | { type: "thinking_level_changed"; level: ThinkingLevel } | { type: "service_tier_changed"; serviceTier: ServiceTier } - | { - type: "compaction_end"; - reason: CompactionReason; - result: CompactionResult | undefined; - aborted: boolean; - willRetry: boolean; - errorMessage?: string; - errorSeverity?: "warning" | "error"; - customInstructions?: string; - } | SessionRetryEvent | { type: "rlm_child_update"; child: RlmChildAgentSnapshot } | { type: "recap_update"; recap: string | undefined } @@ -415,10 +375,9 @@ export type { TurnExecutionPolicy } from "../session/turn-preparation.js"; export type AgentSessionEventListener = (event: AgentSessionEvent) => void; -export class CompactionSkippedError extends Error {} +export { CompactionSkippedError } from "../session/compaction-execution.js"; -/** Thrown when a session_before_refine extension skips the refinement round. */ -export class RefineSkippedError extends Error {} +export { RefineSkippedError } from "../session/refinement.js"; export interface AgentSessionConfig { agent: Agent; @@ -497,39 +456,11 @@ export interface ExtensionBindings { onError?: ExtensionErrorListener; } -export interface AutoRefineReviewRequest { - reason: AutoRefineReason; - turnsSinceLastReview: number; -} - -/** - * Discriminated result from a serialized-mode background planning pass. - * - "plan": review approved and planning succeeded; carry the exact plan, - * options, and abort controller so the boundary can apply directly - * without a second planning request. - * - "skip": reviewer declined; no refine needed. - * - "failure": review or planning threw; boundary should not retry. - */ -export type SerializedBackgroundPlanResult = - | { - status: "plan"; - plan: RefinementPlan; - options: { instructions?: string; rollbackId?: string; global?: boolean }; - abort: AbortController; - branchVersion: number; - source: Exclude; - } - | { status: "skip"; explicit?: boolean } - | { status: "invalidated"; branchVersion: number } - | { - status: "failure"; - explicit: boolean; - options: { instructions?: string; rollbackId?: string; global?: boolean }; - branchVersion: number; - }; - -export type AutoRefineReviewer = (request: AutoRefineReviewRequest, signal?: AbortSignal) => Promise; - +export type { + AutoRefineReviewer, + AutoRefineReviewRequest, + SerializedBackgroundPlanResult, +} from "../session/refinement.js"; export interface PromptOptions { expandPromptTemplates?: boolean; images?: ImageContent[]; @@ -683,16 +614,6 @@ function createAgentMessageDeferred(): AgentMessageDeferred { return deferred; } -/** One-shot settlement for a scheduled post-compaction continuation; a settled failure is never re-exposed to later waiters. */ -interface PostCompactionContinuationSettlement extends AgentMessageDeferred { - continueAfterSessionInput: boolean; - settled: boolean; -} - -function createPostCompactionContinuationSettlement(): PostCompactionContinuationSettlement { - return { ...createAgentMessageDeferred(), continueAfterSessionInput: false, settled: false }; -} - export interface ModelCycleResult { model: Model; thinkingLevel: ThinkingLevel; @@ -781,14 +702,6 @@ const RLM_MAX_DEPTH_STATE_CUSTOM_TYPE = "rlm_max_depth_state"; function noopRlmChildAbort(): void {} function noopRlmChildEventUnsubscribe(): void {} -function autoRefineInstructions(reason: AutoRefineReason, review: AutoRefineReview): string { - const detail = review.instructions - ? ` -Reviewer instructions: ${review.instructions}` - : ""; - return `Automatic refine review triggered by ${reason}. Only create/update/delete local harness entries if there is clear evidence that should help this session continue. Prefer an empty edits array over speculative or one-off memories. Do not promote anything global unless explicitly requested. Reviewer rationale: ${review.rationale}${detail}`; -} - function isNonNegativeInteger(value: unknown): value is number { return typeof value === "number" && Number.isSafeInteger(value) && value >= 0; } @@ -972,6 +885,7 @@ function attributeChildUsage(parentUsage: Usage, childUsage: Usage): void { } export class AgentSession { + private readonly _refinement: SessionRefinement; readonly agent: Agent; readonly sessionManager: SessionManager; readonly settingsManager: SettingsManager; @@ -1007,7 +921,7 @@ export class AgentSession { waitForAgentIdle: () => this.agent.waitForIdle(), hasCancelledDispatchCapture: () => this._hasCancelledDispatchCapture(), getEventQueue: () => this._agentEventQueue, - waitForRefinement: () => this._waitForRefineIdle(), + waitForRefinement: () => this._refinement._waitForRefineIdle(), getTranscript: () => this.agent.state.messages, startTurns: (actions, epoch) => this._startPreparedTurnActions(actions, epoch), executeCommand: (action, epoch) => this._executeSelectedSessionCommand(action, epoch), @@ -1024,8 +938,8 @@ export class AgentSession { private readonly _durableRlmTerminalNoticeActionIds = new Set(); private readonly _commitFence = new SessionCommitFence(); private readonly _turnPreparer = new TurnPreparer({ - hasRefinement: () => this._refineInFlight !== undefined, - waitForRefinement: () => this._waitForRefineIdle(), + hasRefinement: () => this._refinement.isApplying, + waitForRefinement: () => this._refinement._waitForRefineIdle(), flushPendingBash: () => this._flushPendingBashMessages(), validate: () => this._validateCanStartAgentRun(), compact: () => this._runPreTurnCompaction(), @@ -1042,14 +956,63 @@ export class AgentSession { private _autonomousContinuationSuppressionDepth = 0; private _autonomousContinuationSuppressedMessages = new WeakSet(); - private _compactionAbortController: AbortController | undefined = undefined; - private _autoCompactionAbortController: AbortController | undefined = undefined; - private _compactionOperation: Promise | undefined = undefined; - /** One recovery attempt per overflow; "reported" dedups the failure notice. */ - private _overflowRecovery: "idle" | "attempted" | "reported" = "idle"; - private _continueAfterThresholdCompaction = false; - private _pendingRequestedCompaction: { customInstructions?: string } | undefined; - private _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; + private readonly _compaction = new SessionCompaction({ + getModel: () => this.model, + isStreaming: () => this.isStreaming, + getSettings: () => this.settingsManager.getCompactionSettings(), + runAutomatic: (reason, willRetry) => this._runAutoCompaction(reason, willRetry), + queueGoalContinuation: (message) => this._queueGoalContinuationForThresholdCompaction(message), + queueAutonomousContinuation: (message) => this._queueAutonomousContinuationForThresholdCompaction(message), + beginRefinementAbort: () => this._refinement.beginAbortedTurnCleanup(), + getRequiredAuth: (model) => this._getRequiredRequestAuth(model), + getAuth: (model) => this._modelRegistry.getApiKeyAndHeaders(model), + perform: (options) => this._performCompaction(options), + disconnect: () => this._disconnectFromAgent(), + reconnect: () => this._reconnectToAgent(), + abortSession: () => this.abort(), + getContinuationState: () => ({ + scheduled: this._continuation.isScheduled, + continueAfterSessionInput: this._continuation.current?.continueAfterSessionInput ?? false, + }), + afterManualCompaction: (signal, scheduled, continueAfterInput) => + this._afterManualCompaction(signal, scheduled, continueAfterInput), + getMessages: () => this.agent.state.messages, + replaceMessages: (messages) => { + this.agent.state.messages = messages; + }, + hasAgentQueuedMessages: () => this.agent.hasQueuedMessages(), + hasPendingSessionWork: () => this.hasPendingSessionWork, + scheduleContinuation: (continueAfterInput) => this._schedulePostCompactionContinue(continueAfterInput), + scheduleRefinement: (willContinue) => this._refinement._scheduleAutoRefineAfterCompaction(willContinue), + takeThresholdAutonomousMessages: () => this._pendingThresholdCompactionAutonomousMessages.splice(0), + getThresholdGoalContinuation: () => this._queuedGoalThresholdContinuation, + clearAutonomousContinuations: (shouldContinue, messages) => + this._clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction(shouldContinue, messages), + clearGoalContinuation: (message) => this._clearQueuedGoalContinuationAfterCancelledThresholdCompaction(message), + getSessionStore: () => this.sessionManager, + retainUnpersistedOutcome: (message) => { + this._unpersistedOutcomes.push(message); + }, + emit: (event) => this._emit(event), + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + scheduleInput: () => this._scheduleSessionInputPump(), + }); + private readonly _compactionExecution: CompactionExecutionHost = { + getSessionStore: () => this.sessionManager, + getSettings: () => this.settingsManager.getCompactionSettings(), + getSemanticEdges: () => this._semanticEdges, + getExtensions: () => this._extensionRunner, + getThinkingLevel: () => this.thinkingLevel, + getRetryPolicy: () => providerRetryPolicy(this.settingsManager), + getHarnessDigest: () => this._harnessDigest(), + rebuildContext: () => { + this.agent.state.messages = this.sessionManager.buildSessionContext().messages; + this._mergeUnpersistedOutcomes(this.agent.state.messages); + this._restoreLateIpythonSentAgentMessages(); + }, + syncKernelState: () => this._syncKernelStateAfterCompaction(), + reapDeletedChildren: () => this._reapDeletedRlmSubagentRuntimesAfterCompaction(), + }; private _branchSummaryAbortController: AbortController | undefined = undefined; private _branchSummaryOperation: Promise | undefined = undefined; @@ -1073,7 +1036,7 @@ export class AgentSession { continue: () => this.agent.continue(), waitForIdle: () => this.agent.waitForIdle(), cancelCompaction: () => { - this._autoCompactionAbortController?.abort(); + this._compaction.abortAutomatic(); this._cancelPostCompactionContinue(); }, emit: (event) => this._emit(event), @@ -1198,40 +1161,77 @@ export class AgentSession { private _baseSystemPrompt = ""; private _baseSystemPromptOptions!: BuildSystemPromptOptions; - private _assistantTurnsSinceAutoRefine = 0; - private _lastAutoRefineReviewAt = 0; - private _autoRefineInProgress = false; - private readonly _autoRefineOperations = new Set>(); - private readonly _scheduledAutoRefineTimers = new Set>(); - private _compactAutoRefinePending = false; - private _turnIntervalAutoRefinePending = false; - private _postCompactionContinuationScheduled = false; - private _postCompactionContinuationSettlement: PostCompactionContinuationSettlement | undefined; - private _postCompactionContinuationMessages: AgentMessage[] = []; - private _scheduledPostCompactionContinuationMessages: AgentMessage[] = []; + private readonly _continuation = new SessionContinuation({ + waitForAgentIdle: () => this.agent.waitForIdle(), + waitForRetry: () => this.waitForRetry(), + waitForRefinement: () => this._refinement._waitForRefineIdle(), + queuedWorkPauseCount: () => this._inputScheduler.queuedWorkPauseCount, + addCheckpointWaiter: (waiter) => { + this._sessionInputCheckpointWaiters.add(waiter); + }, + removeCheckpointWaiter: (waiter) => { + this._sessionInputCheckpointWaiters.delete(waiter); + }, + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + compactionOperation: () => this._compaction.operation, + isRefinementApplying: () => this._refinement.isApplying, + acquireCommitFence: () => this._acquireSessionActionCommitFence(), + scheduleRefinement: () => this._refinement._scheduleAutoRefineAfterAgentEnd(), + unfinishedActionCount: () => this.unfinishedActionCount, + isInputRequested: () => this._inputScheduler.requested, + scheduleInput: () => this._scheduleSessionInputPump(), + continue: () => this.agent.continue(), + waitForIdleOrSettlement: (token) => this._waitForIdleOrSettlement(token), + removeQueuedMessages: (predicate) => this.agent.removeQueuedMessages(predicate), + followUp: (message) => this.agent.followUp(message), + onMessageConsumed: (message) => { + this._queuedAutonomousContinuationSnapshots.delete(message); + }, + }); private _queuedAutonomousThresholdContinuations = new WeakMap(); private _queuedAutonomousContinuationSnapshots = new WeakMap(); private _pendingThresholdCompactionAutonomousMessages: AgentMessage[] = []; private _queuedGoalThresholdContinuation: AgentMessage | undefined; - private _pendingAutoRefineReview: { reason: AutoRefineReason; review: AutoRefineReview } | undefined; - private _autoRefineBranchVersion = 0; - private _autoRefineReviewAbort?: AbortController; - private _refineAbortController?: AbortController; - private readonly _autoRefineReviewer?: AutoRefineReviewer; - private readonly _serializedRefine: boolean; - private _refineInFlight?: Promise; - private _refinePlanInFlight?: Promise; - private _serializedPlanInFlight?: Promise; - private _serializedPlanClaim?: Promise; - private _serializedExplicitRefineOptions?: { - instructions?: string; - global?: boolean; - }; constructor(config: AgentSessionConfig) { this.agent = config.agent; this.sessionManager = config.sessionManager; this.settingsManager = config.settingsManager; + this._refinement = new SessionRefinement( + { + sessionManager: this.sessionManager, + settingsManager: this.settingsManager, + getRetryPolicy: () => providerRetryPolicy(this.settingsManager), + isDisposed: () => this._disposed, + isDisposing: () => this._disposing, + isStreaming: () => this.isStreaming, + isCompacting: () => this.isCompacting, + getDepth: () => this._rlmDepth, + getRlmSessionDir: () => this._rlmSessionDir, + getModel: () => this.model, + getThinkingLevel: () => this.thinkingLevel, + getMessages: () => this.agent.state.messages, + getRequiredRequestAuth: (model) => this._getRequiredRequestAuth(model), + getExtensionRunner: () => this._extensionRunner, + getEventQueue: () => this._agentEventQueue, + getCompactionOperation: () => this._compaction.operation, + getBranchSummaryOperation: () => this._branchSummaryOperation, + waitForAgentIdle: () => this.agent.waitForIdle(), + dispatchRefine: (options, internal) => this.refine(options, internal), + disconnect: () => this._disconnectFromAgent(), + reconnect: () => this._reconnectToAgent(), + emit: (event) => this._emit(event), + retainUnpersistedOutcome: (message) => this._unpersistedOutcomes.push(message), + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + scheduleInputPump: () => this._scheduleSessionInputPump(), + isContinuationScheduled: () => this._continuation.isScheduled, + cancelContinuation: () => this._cancelPostCompactionContinue(), + }, + { + serializedRefine: config.serializedRefine, + autoRefineReviewer: config.autoRefineReviewer?.bind(this), + }, + ); this._serviceTierPreference = config.serviceTierPreference ?? config.agent.state.serviceTier; this._scopedModels = config.scopedModels ?? []; this._resourceLoader = config.resourceLoader; @@ -1262,8 +1262,7 @@ export class AgentSession { this._rlmMaxDepth = resolvedRlmMaxDepth.maxDepth; this._rlmMaxDepthSource = resolvedRlmMaxDepth.source; this._prewarmIpythonKernel = (config.prewarmIpythonKernel ?? false) && this._rlmDepth === 0; - this._autoRefineReviewer = config.autoRefineReviewer; - this._serializedRefine = config.serializedRefine ?? false; + this._rlmSessionDir = config.rlmSessionDir; this._rlmParentNodeId = config.rlmParentNodeId; this._rlmParentAgent = config.rlmParentAgent; @@ -2034,11 +2033,11 @@ export class AgentSession { // This MUST run BEFORE threshold compaction to prevent the // compaction model call from overlapping an in-flight refine // plan/apply that was started at message_end. - if (this._serializedRefine) { + if (this._refinement.serialized) { // Ensure the preceding message_end processing (counter increment, // background plan kickoff) has completed before the checkpoint. await this._agentEventQueue; - await this._runSerializedRefineCheckpoint(); + await this._refinement._runSerializedRefineCheckpoint(); } if (await this._shouldStopForThresholdCompaction(context)) { return true; @@ -2049,495 +2048,17 @@ export class AgentSession { } private async _shouldStopForThresholdCompaction(context: ShouldStopAfterTurnContext): Promise { - this._continueAfterThresholdCompaction = false; - if (this._pendingRequestedCompaction === undefined && !(await this._thresholdCompactionNeeded(context))) { + this._compaction.resetContinuation(); + if (!this._compaction.hasPendingRequest && !(await this._thresholdCompactionNeeded(context))) { return false; } const lastMessage = this.agent.state.messages[this.agent.state.messages.length - 1]; // A queued continuation disproves the assistant-last "task finished" heuristic, so preserve a true set above. - this._continueAfterThresholdCompaction ||= lastMessage !== undefined && lastMessage.role !== "assistant"; + if (lastMessage !== undefined && lastMessage.role !== "assistant") this._compaction.requestContinuation(); return true; } - /** - * Serialized-mode auto-refine checkpoint called from _shouldStopAfterTurn. - * Runs the review, planning, and application phases inline between turns - * at the quiescent shouldStopAfterTurn boundary. This path NEVER calls - * _maybeAutoRefine, _runApprovedRefine, public refine(), agent.abort(), - * or agent.waitForIdle — all of which would deadlock or defer because - * the agent loop still owns activeRun at this point. Instead it calls - * _reviewAutoRefine, _planRefine, and _applyRefine directly with proper - * in-flight guards and counter resets. - */ - private async _runSerializedRefineCheckpoint(): Promise { - if (this._disposed || this._disposing) { - return; - } - - // 1. Await any background plan that was started at message_end - // (either for a pending refine.run or for interval-triggered - // auto-refine). This must be checked BEFORE the pending and - // interval checks because background planning may have consumed - // the pending request at message_end. - const branchVersion = this._autoRefineBranchVersion; - const bgConsumption = await this._consumeSerializedBackgroundPlan(async (bgResult) => { - if (this._disposed || this._disposing) { - return true; - } - - if (bgResult?.status === "plan") { - if (bgResult.branchVersion !== this._autoRefineBranchVersion) { - if (!this._pendingRequestedRefine) { - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - return true; - } - } else { - // Apply the EXACT background plan directly via _applyRefine - // (no second _planRefine call). - try { - await this._applySerializedPlan(bgResult); - } catch (error) { - this._emitRefineFailed(error); - } - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - if (!this._pendingRequestedRefine) { - return true; - } - } - } - - if (bgResult?.status === "skip") { - // Reviewer declined or an extension skipped during background planning. - // Reset exactly once. Never retry the interval review; only fall through for a separate pending refine.run. - if (bgResult.explicit) { - this._emitRefineFailed(new RefineSkippedError("Refinement skipped by extension")); - } - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - if (!this._pendingRequestedRefine) { - return true; - } - } - - if (bgResult?.status === "failure") { - // Background review or planning failure stamps cooldown without a synchronous retry. - // A separately queued refine.run may still be serviced below. - if (branchVersion === this._autoRefineBranchVersion) { - this._lastAutoRefineReviewAt = Date.now(); - } - // Re-queue an explicit refine.run whose background plan failed, - // but only when branchVersion is still current and no newer - // pending request has arrived since the background plan consumed - // the original one. A newer request retains priority; interval - // failures keep existing no-retry cooldown semantics. - if ( - bgResult.explicit && - bgResult.branchVersion === this._autoRefineBranchVersion && - !this._pendingRequestedRefine - ) { - this._pendingRequestedRefine = bgResult.options; - } - if (!this._pendingRequestedRefine) { - return true; - } - } - - if (bgResult?.status === "invalidated" && !this._pendingRequestedRefine) { - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - return true; - } - - await this._runSerializedRefineCheckpointAfterBackground(branchVersion); - return true; - }); - if (this._disposed || this._disposing || bgConsumption !== "none") { - return; - } - await this._runSerializedRefineCheckpointAfterBackground(branchVersion); - } - - private async _runSerializedRefineCheckpointAfterBackground(branchVersion: number): Promise { - // No background result, or a refine.run arrived while the background result was - // in flight. Fall through so an explicit pending request is serviced at this boundary. - - // 2. Agent-callable refine.run requests that were NOT consumed by - // background planning (e.g. interval not reached at message_end, - // or cooldown was active). Service them synchronously. - const pending = this._pendingRequestedRefine; - if (pending) { - this._pendingRequestedRefine = undefined; - try { - await this._runSerializedRefine(pending, "self"); - } catch (error) { - this._emitRefineFailed(error); - } - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - return; - } - - // 3. Post-compaction auto-refine. Serialized sessions defer the - // compaction trigger to this boundary instead of entering the interactive - // path, which waits for agent idle and can never run inside a tool loop. - if (!this._autoRefineAllowedForSession()) { - this._compactAutoRefinePending = false; - return; - } - const settings = this.settingsManager.getAutoRefineSettings(); - if (!settings.enabled) { - this._compactAutoRefinePending = false; - return; - } - if (this._compactAutoRefinePending) { - if (!settings.compact) { - this._compactAutoRefinePending = false; - } else { - const nowMs = Date.now(); - const underCooldown = - this._lastAutoRefineReviewAt > 0 && nowMs - this._lastAutoRefineReviewAt < settings.cooldownMs; - if (underCooldown) { - // Preserve the compact trigger for a later boundary, matching the - // interactive path's pending behavior while the cooldown is active. - return; - } - this._compactAutoRefinePending = false; - await this._runSerializedAutoRefineReview("compact", branchVersion); - return; - } - } - - // 4. Interval-triggered auto-refine (no background plan was started). - if (this._assistantTurnsSinceAutoRefine < settings.turnInterval) { - return; - } - const nowMs = Date.now(); - const underCooldown = - this._lastAutoRefineReviewAt > 0 && nowMs - this._lastAutoRefineReviewAt < settings.cooldownMs; - if (underCooldown) { - return; - } - await this._runSerializedAutoRefineReview("turn_interval", branchVersion); - } - - private async _runSerializedAutoRefineReview( - reason: "compact" | "turn_interval", - branchVersion: number, - ): Promise { - const reviewAbort = new AbortController(); - this._autoRefineReviewAbort = reviewAbort; - this._autoRefineInProgress = true; - try { - const review = await this._reviewAutoRefine( - { reason, turnsSinceLastReview: this._assistantTurnsSinceAutoRefine }, - reviewAbort.signal, - ); - if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { - return; - } - if (!review.shouldRefine) { - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - return; - } - await this._runSerializedRefine({ instructions: autoRefineInstructions(reason, review) }, "auto"); - if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { - return; - } - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - } catch (error) { - if (branchVersion === this._autoRefineBranchVersion) { - this._lastAutoRefineReviewAt = Date.now(); - // An extension skip is an intentional non-round, not a failure. - if (error instanceof RefineSkippedError) { - this._assistantTurnsSinceAutoRefine = 0; - } else { - this._emitRefineFailed(error); - } - } - } finally { - if (this._autoRefineReviewAbort === reviewAbort) { - this._autoRefineReviewAbort = undefined; - } - this._autoRefineInProgress = false; - } - } - - /** - * Claim and process the serialized background plan if one is in flight. - * A concurrent caller waits for the claim holder's full processing callback - * instead of resuming as soon as planning settles. - */ - private async _consumeSerializedBackgroundPlan( - consume: (result: SerializedBackgroundPlanResult | undefined) => Promise, - ): Promise<"none" | "waited" | "continue" | "stop"> { - if (this._serializedPlanClaim) { - await this._serializedPlanClaim.catch(() => undefined); - return "waited"; - } - const planInFlight = this._serializedPlanInFlight; - if (!planInFlight) { - return "none"; - } - - let releaseClaim: () => void = () => {}; - const claim = new Promise((resolve) => { - releaseClaim = resolve; - }); - this._serializedPlanClaim = claim; - try { - const result = await planInFlight.catch(() => undefined); - if (this._serializedPlanInFlight === planInFlight) { - this._serializedPlanInFlight = undefined; - this._serializedExplicitRefineOptions = undefined; - } - return (await consume(result)) ? "stop" : "continue"; - } finally { - releaseClaim(); - if (this._serializedPlanClaim === claim) { - this._serializedPlanClaim = undefined; - } - } - } - - /** - * Apply an exact background plan directly via _applyRefine without - * calling _planRefine again. Sets _refineInFlight for safety. - */ - private async _applySerializedPlan( - bgResult: Extract, - ): Promise { - let resolveApplySettled: () => void = () => {}; - const applySettled = new Promise((resolve) => { - resolveApplySettled = resolve; - }); - this._refineInFlight = applySettled; - try { - await this._applyRefine(bgResult.plan, bgResult.options, bgResult.abort, bgResult.source); - } finally { - resolveApplySettled(); - if (this._refineInFlight === applySettled) { - this._refineInFlight = undefined; - } - this._notifySessionInputCheckpointChange(); - this._scheduleSessionInputPump(); - } - } - - /** - * Start background refinement planning at assistant message_end, while - * tools are still executing. The plan (if any) is awaited at the - * shouldStopAfterTurn boundary before applying. Planning overlaps tool - * execution only — never another model request. - */ - private _maybeStartSerializedBackgroundPlan(): void { - if (!this._serializedRefine || this._disposed || this._disposing) { - return; - } - // Don't start if a plan is already in flight. - if (this._serializedPlanInFlight || this._refineInFlight || this._refinePlanInFlight) { - return; - } - - // Start background planning for a pending agent-callable - // refine.run request, so its plan is ready at the shouldStopAfterTurn - // boundary. The pending request is consumed (cleared) here so the - // boundary doesn't re-plan it. Explicit refine.run skips the review gate. - const pending = this._pendingRequestedRefine; - if (pending) { - this._pendingRequestedRefine = undefined; - this._serializedExplicitRefineOptions = pending; - const refineAbort = new AbortController(); - this._refineAbortController = refineAbort; - const branchVersion = this._autoRefineBranchVersion; - this._serializedPlanInFlight = this._runBackgroundPlan(pending, refineAbort, branchVersion, true); - return; - } - - // Interval-triggered auto-refine background planning. - if (!this._autoRefineAllowedForSession()) { - return; - } - const settings = this.settingsManager.getAutoRefineSettings(); - if (!settings.enabled) { - return; - } - if (this._assistantTurnsSinceAutoRefine < settings.turnInterval) { - return; - } - const nowMs = Date.now(); - const underCooldown = - this._lastAutoRefineReviewAt > 0 && nowMs - this._lastAutoRefineReviewAt < settings.cooldownMs; - if (underCooldown) { - return; - } - - const refineAbort = new AbortController(); - this._refineAbortController = refineAbort; - const branchVersion = this._autoRefineBranchVersion; - // Pass empty options — _runBackgroundPlan derives instructions from - // the review result for interval-triggered auto-refine. - this._serializedPlanInFlight = this._runBackgroundPlan({}, refineAbort, branchVersion); - } - - /** - * Shared background planning coroutine. Runs review + planRefine and - * returns a discriminated result so the boundary can distinguish - * reviewer-declined ("skip") from failure ("failure") from a ready - * plan ("plan") and apply that exact plan without re-planning. - */ - private async _runBackgroundPlan( - options: { instructions?: string; rollbackId?: string; global?: boolean }, - refineAbort: AbortController, - branchVersion: number, - skipReview = false, - ): Promise { - try { - let planOptions = options; - if (!skipReview) { - // Interval-triggered: run the review gate first, then derive - // instructions from the review result (not prepopulated). - const review = await this._reviewAutoRefine( - { - reason: "turn_interval", - turnsSinceLastReview: this._assistantTurnsSinceAutoRefine, - }, - refineAbort.signal, - ); - if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { - return { status: "invalidated", branchVersion }; - } - if (!review.shouldRefine) { - return { status: "skip" }; - } - planOptions = { - instructions: autoRefineInstructions("turn_interval", review), - }; - } - // For explicit refine.run (skipReview=true), plan directly with - // the user-provided options — no auto-review gate. - const plan = await this._planRefine(planOptions, refineAbort.signal, skipReview ? "manual" : "auto"); - if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { - return { status: "invalidated", branchVersion }; - } - return { - status: "plan", - plan, - options: planOptions, - abort: refineAbort, - branchVersion, - source: skipReview ? "self" : "auto", - }; - } catch (error) { - if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { - return { status: "invalidated", branchVersion }; - } - if (error instanceof RefineSkippedError) { - return { status: "skip", explicit: skipReview }; - } - return { - status: "failure", - explicit: skipReview, - options, - branchVersion, - }; - } finally { - if (this._refineAbortController === refineAbort) { - this._refineAbortController = undefined; - } - } - } - - /** - * Direct serialized plan+apply. Calls _planRefine and _applyRefine with - * proper in-flight guards but NEVER agent.waitForIdle or agent.abort. - * The caller (shouldStopAfterTurn) is already at the quiescent boundary, - * so the agent is between turns and _applyRefine's disconnect/reconnect - * is safe. - */ - private async _runSerializedRefine( - options: { - instructions?: string; - rollbackId?: string; - global?: boolean; - }, - source: Exclude, - ): Promise { - if (this._disposed || this._disposing) { - return; - } - // Guard: serialize against concurrent _runSerializedRefine calls. - // _serializedPlanInFlight covers background planning; _refineInFlight - // covers the apply phase. Both must be settled before starting a new - // plan+apply cycle. - while (this._serializedPlanInFlight || this._refineInFlight || this._refinePlanInFlight) { - if (this._serializedPlanInFlight) { - await this._consumeSerializedBackgroundPlan(async () => false); - } else if (this._refineInFlight) { - await this._refineInFlight; - } else { - await this._refinePlanInFlight; - } - } - if (this._disposed || this._disposing) { - return; - } - - const refineAbort = new AbortController(); - this._refineAbortController = refineAbort; - - const planRun = this._planRefine(options, refineAbort.signal, source === "auto" ? "auto" : "manual"); - const planSettled = planRun.then( - () => undefined, - () => undefined, - ); - this._refinePlanInFlight = planSettled; - let plan: RefinementPlan; - try { - plan = await planRun; - } catch (error) { - if (this._refineAbortController === refineAbort) { - this._refineAbortController = undefined; - } - this._scheduleSessionInputPump(); - throw error; - } finally { - if (this._refinePlanInFlight === planSettled) { - this._refinePlanInFlight = undefined; - } - } - - if (this._disposed || refineAbort.signal.aborted) { - if (this._refineAbortController === refineAbort) { - this._refineAbortController = undefined; - } - this._scheduleSessionInputPump(); - return; - } - - // Do NOT call agent.waitForIdle() — we are at the quiescent boundary - // already (shouldStopAfterTurn). _applyRefine handles disconnect/reconnect internally. - let resolveApplySettled: () => void = () => {}; - const applySettled = new Promise((resolve) => { - resolveApplySettled = resolve; - }); - this._refineInFlight = applySettled; - try { - await this._applyRefine(plan, options, refineAbort, source); - } finally { - resolveApplySettled(); - if (this._refineInFlight === applySettled) { - this._refineInFlight = undefined; - } - this._notifySessionInputCheckpointChange(); - this._scheduleSessionInputPump(); - } - } - private async _thresholdCompactionNeeded(context: ShouldStopAfterTurnContext): Promise { const settings = this.settingsManager.getCompactionSettings(); if (!settings.enabled) return false; @@ -2549,16 +2070,16 @@ export class AgentSession { return false; } - const contextTokens = this._getThresholdContextTokens(context.message, compactionTimestamp); + const contextTokens = this._compaction.getThresholdContextTokens(context.message, compactionTimestamp); if (contextTokens === undefined || !shouldCompact(contextTokens, contextWindow, settings)) { return false; } // Goal continuation takes exclusive priority over autonomous continuation, matching _getContinuationMessages. if (this._queueGoalContinuationForThresholdCompaction(context.message)) { - this._continueAfterThresholdCompaction = true; + this._compaction.requestContinuation(); } else if (await this._queueAutonomousContinuationForThresholdCompaction(context.message)) { - this._continueAfterThresholdCompaction = true; + this._compaction.requestContinuation(); } return true; } @@ -2589,7 +2110,7 @@ export class AgentSession { message: AssistantMessage, ): Promise { const queuedMessage = this._queuedAutonomousThresholdContinuations.get(message); - if (queuedMessage && this._postCompactionContinuationMessages.includes(queuedMessage)) { + if (queuedMessage && this._continuation.messages.includes(queuedMessage)) { return queuedMessage; } const snapshot = this._snapshotAutonomousRuntimeState(); @@ -2607,7 +2128,7 @@ export class AgentSession { } this._queuedAutonomousThresholdContinuations.set(message, autonomousMessage); this._queuedAutonomousContinuationSnapshots.set(autonomousMessage, snapshot); - this._postCompactionContinuationMessages.push(autonomousMessage); + this._continuation.track(autonomousMessage); this._pendingThresholdCompactionAutonomousMessages.push(autonomousMessage); const text = typeof autonomousMessage.content === "string" @@ -2684,18 +2205,14 @@ export class AgentSession { private _clearQueuedAutonomousContinuations( options: { restoreAutonomousState?: boolean; messages?: AgentMessage[] } = {}, ): void { - const requestedMessages = options.messages ?? [...this._postCompactionContinuationMessages]; + const requestedMessages = options.messages ?? [...this._continuation.messages]; const requestedMessageSet = new Set(requestedMessages); - const queuedMessages = this._postCompactionContinuationMessages.filter((message) => - requestedMessageSet.has(message), - ); + const queuedMessages = this._continuation.messages.filter((message) => requestedMessageSet.has(message)); if (queuedMessages.length === 0) { return; } const queuedMessageSet = new Set(queuedMessages); - this._postCompactionContinuationMessages = this._postCompactionContinuationMessages.filter( - (message) => !queuedMessageSet.has(message), - ); + this._continuation.remove(queuedMessageSet); this.agent.removeQueuedMessages((message) => queuedMessageSet.has(message)); this._cancelSessionActions( (action) => action.payload.kind === "turn" && queuedMessageSet.has(primaryDeliveryRecord(action).message), @@ -2718,7 +2235,7 @@ export class AgentSession { (message) => !queuedMessageSet.has(message), ); if (options.messages === undefined) { - this._continueAfterThresholdCompaction = false; + this._compaction.resetContinuation(); } if (!this.agent.hasQueuedMessages() && this.unfinishedActionCount === 0) { this._cancelPostCompactionContinue(); @@ -2781,7 +2298,7 @@ export class AgentSession { tokens: usage?.tokens ?? null, context_window: usage?.contextWindow ?? null, percent: usage?.percent ?? null, - scheduled: this._pendingRequestedCompaction !== undefined, + scheduled: this._compaction.hasPendingRequest, }; } case "compact.run": { @@ -2806,7 +2323,7 @@ export class AgentSession { reason: lastEntry?.type === "compaction" ? "already compacted" : "session is too short to compact", }; } - this._pendingRequestedCompaction = { customInstructions: instructions }; + this._compaction.request(instructions); return { scheduled: true, note: "Compaction runs when the current turn ends; you resume automatically afterwards. Continue working normally.", @@ -2825,75 +2342,19 @@ export class AgentSession { * if refine() awaited agent idle from within the active tool call. */ handleRefineHostRequest(type: string, payload: Record = {}): Record { - switch (type) { - case "refine.status": { - return { - pending: this._pendingRequestedRefine !== undefined, - in_flight: - this._refineInFlight !== undefined || - this._refinePlanInFlight !== undefined || - this._serializedPlanInFlight !== undefined, - }; - } - case "refine.run": { - const instructions = payload.instructions; - if (instructions !== undefined && typeof instructions !== "string") { - throw new Error("refine.run instructions must be a string when provided"); - } - const globalFlag = payload.global; - if (globalFlag !== undefined && typeof globalFlag !== "boolean") { - throw new Error("refine.run global must be a boolean when provided"); - } - if (!this.isStreaming) { - return { - scheduled: false, - reason: "no active turn; refine can only be requested while a turn is running", - }; - } - const previous = this._pendingRequestedRefine ?? this._serializedExplicitRefineOptions; - this._pendingRequestedRefine = { - instructions: instructions ?? previous?.instructions, - global: globalFlag ?? previous?.global, - }; - // In serialized mode, kick off background planning immediately - // (the primary response ended at message_end, tools are active). - // This lets planning overlap tool execution rather than waiting - // for the shouldStopAfterTurn boundary. - if (this._serializedRefine) { - if (this._serializedPlanInFlight) { - this._autoRefineBranchVersion++; - if (this._refineAbortController) { - this._refineAbortController.abort(); - } else { - this._serializedPlanInFlight = Promise.resolve({ - status: "invalidated", - branchVersion: this._autoRefineBranchVersion, - }); - } - } else { - this._maybeStartSerializedBackgroundPlan(); - } - } - return { - scheduled: true, - note: "Refinement runs when the current turn ends; applied edits are appended to your context as a refinement notice and you resume automatically. Continue working normally.", - }; - } - default: - throw new Error(`unknown refine request type "${type}"`); - } - } - - /** - * Handle an rlm_heartbeat.* request from the bundled rlm-heartbeat skill. - * These heartbeats are internal to this active session and never read or - * mutate the user-level /heartbeat. - */ - handleRlmHeartbeatHostRequest(type: string, payload: Record = {}): Record { - const controller = this._rlmHeartbeatController; - if (!controller) { - throw new Error("RLM heartbeat skill is not available in this session"); - } + return this._refinement.handleRefineHostRequest(type, payload); + } + + /** + * Handle an rlm_heartbeat.* request from the bundled rlm-heartbeat skill. + * These heartbeats are internal to this active session and never read or + * mutate the user-level /heartbeat. + */ + handleRlmHeartbeatHostRequest(type: string, payload: Record = {}): Record { + const controller = this._rlmHeartbeatController; + if (!controller) { + throw new Error("RLM heartbeat skill is not available in this session"); + } switch (type) { case "rlm_heartbeat.list": { const includeInactive = payload.include_inactive === true || payload.includeInactive === true; @@ -3335,7 +2796,7 @@ export class AgentSession { } if (event.type === "message_start" && startsAgentRun(event.message)) { - this._overflowRecovery = "idle"; + this._compaction.resetOverflowRecovery(); } await this._emitExtensionEvent(event); @@ -3376,16 +2837,15 @@ export class AgentSession { addAutonomousUsage(this._autonomousState, assistantMsg.usage); } if (assistantMsg.stopReason !== "error" && assistantMsg.stopReason !== "aborted") { - this._assistantTurnsSinceAutoRefine++; + this._refinement.observeAssistantEnd(); // In serialized mode, kick off background refinement planning // immediately after the primary stream finishes, while tools // are still executing. The plan is awaited at shouldStopAfterTurn // before applying, so planning overlaps tools only — never another // model request. - this._maybeStartSerializedBackgroundPlan(); } if (assistantMsg.stopReason !== "error") { - this._overflowRecovery = "idle"; + this._compaction.resetOverflowRecovery(); } this._retry.observeAssistantEnd(assistantMsg); if (this._goals.accountAssistantMessage(assistantMsg)) { @@ -3426,10 +2886,10 @@ export class AgentSession { this._finishGoalForTerminalAssistantMessage(msg); // In serialized mode, agent-callable refine.run is serviced // at the shouldStopAfterTurn boundary, not here at agent_end. - if (!this._serializedRefine) { - const consumedRequestedRefine = this._consumePendingRequestedRefine(); + if (!this._refinement.serialized) { + const consumedRequestedRefine = this._refinement._consumePendingRequestedRefine(); if (!consumedRequestedRefine) { - this._scheduleAutoRefineAfterAgentEnd(); + this._refinement._scheduleAutoRefineAfterAgentEnd(); } } } @@ -3601,7 +3061,7 @@ export class AgentSession { this._disposeAsyncPromise = (async () => { // Drain before marking _disposing so a refine triggered at the final // agent_end completes instead of being aborted by dispose(). - await this._drainPendingRefinementForDisposal(); + await this._refinement._drainPendingRefinementForDisposal(); if (this._disposed) { return this._disposeCallbacksPromise; } @@ -3612,138 +3072,6 @@ export class AgentSession { return this._disposeAsyncPromise; } - /** - * Await any in-flight refinement (planning or application) and run a - * pending auto-refine that was scheduled but not yet started. Called - * from disposeAsync before _disposing is set so refinement completes - * before disposal. - */ - private async _drainPendingRefinementForDisposal(): Promise { - for (const timer of this._scheduledAutoRefineTimers) { - clearTimeout(timer); - } - this._scheduledAutoRefineTimers.clear(); - await Promise.allSettled([...this._autoRefineOperations]); - for (const timer of this._scheduledAutoRefineTimers) { - clearTimeout(timer); - } - this._scheduledAutoRefineTimers.clear(); - // Wait for in-flight refinement (including serialized background plan) to settle. - while (this._refineInFlight || this._refinePlanInFlight || this._serializedPlanInFlight) { - if (this._refineInFlight) { - await this._refineInFlight; - } else if (this._refinePlanInFlight) { - await this._refinePlanInFlight; - } else if (this._serializedPlanInFlight) { - // Await the background plan and apply a ready "plan" result before teardown. - await this._consumeSerializedBackgroundPlan(async (bgResult) => { - if (bgResult?.status === "plan" && bgResult.branchVersion === this._autoRefineBranchVersion) { - try { - await this._applySerializedPlan(bgResult); - } catch (error) { - this._emitRefineFailed(error); - } - // Stamp cooldown and reset counter so the interval - // check below does not trigger a duplicate refine. - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - } - // Preserve a consumed explicit request when its background plan failed, - // matching the turn-boundary recovery path. The pending drain below - // retries it once before disposal. - if ( - bgResult?.status === "failure" && - bgResult.explicit && - bgResult.branchVersion === this._autoRefineBranchVersion && - !this._pendingRequestedRefine - ) { - this._pendingRequestedRefine = bgResult.options; - } - if (bgResult?.status === "skip" && bgResult.explicit) { - this._emitRefineFailed(new RefineSkippedError("Refinement skipped by extension")); - } - // For "skip" or "failure", stamp cooldown and reset counter - // so the interval check below does not trigger a duplicate - // terminal retry. - if ( - bgResult?.status === "skip" || - bgResult?.status === "failure" || - bgResult?.status === "invalidated" - ) { - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - } - return false; - }); - } else { - await new Promise((resolve) => setTimeout(resolve, 0)); - } - } - // Drain an agent-callable refine.run request that was scheduled but - // not yet consumed. Use the direct serialized path (no waitForIdle) - // since the agent may still own activeRun at the final agent_end. - if (this._pendingRequestedRefine) { - const pending = this._pendingRequestedRefine; - this._pendingRequestedRefine = undefined; - try { - await this._runSerializedRefine(pending, "self"); - } catch { - // Best-effort drain; refinement errors must not block disposal. - } - // Stamp cooldown and reset counter so the interval check below - // does not trigger a duplicate refine after the explicit drain. - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - } - // A serialized compaction can finish without another model turn. Drain its - // pending review here so disposal does not silently lose the trigger. - if (this._serializedRefine && this._compactAutoRefinePending && this._autoRefineAllowedForSession()) { - const compactSettings = this.settingsManager.getAutoRefineSettings(); - if (!compactSettings.enabled || !compactSettings.compact) { - this._compactAutoRefinePending = false; - } else { - const nowMs = Date.now(); - const underCooldown = - this._lastAutoRefineReviewAt > 0 && nowMs - this._lastAutoRefineReviewAt < compactSettings.cooldownMs; - this._compactAutoRefinePending = false; - if (!underCooldown) { - try { - await this._runSerializedAutoRefineReview("compact", this._autoRefineBranchVersion); - } catch { - // Best-effort drain; refinement errors must not block disposal. - } - return; - } - } - } - - // If auto-refine is due but has not started yet, run it now so the - // refinement is persisted before disposal. Use the direct serialized - // path in serialized mode, or _maybeAutoRefine in interactive mode - // (where the agent is idle at this point). - if (this._disposed || !this._autoRefineAllowedForSession()) { - return; - } - const settings = this.settingsManager.getAutoRefineSettings(); - if (!settings.enabled) { - return; - } - if (this._assistantTurnsSinceAutoRefine < settings.turnInterval) { - return; - } - const nowMs = Date.now(); - const underCooldown = - this._lastAutoRefineReviewAt > 0 && nowMs - this._lastAutoRefineReviewAt < settings.cooldownMs; - if (underCooldown) { - return; - } - if (this._serializedRefine) { - await this._runSerializedRefineCheckpoint(); - } else { - await this._maybeAutoRefine("turn_interval"); - } - } - private async _disposeAsyncOnce(kernelSnapshot: boolean): Promise { // Flush kernels/traces for both still-running and retained children; the sync // dispose() below only tears them down synchronously. @@ -3815,17 +3143,7 @@ export class AgentSession { try { // Invalidate scheduled timers and abort any in-flight review so a late // resolution cannot write harness state or re-subscribe handlers. - this._autoRefineReviewAbort?.abort(); - this._refineAbortController?.abort(); - for (const timer of this._scheduledAutoRefineTimers) { - clearTimeout(timer); - } - this._scheduledAutoRefineTimers.clear(); - this._serializedPlanInFlight = undefined; - this._serializedExplicitRefineOptions = undefined; - this._pendingRequestedRefine = undefined; - this._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); - this._autoRefineBranchVersion++; + this._refinement.dispose(); this._cancelActiveRlmChildRuns("Parent session disposed"); for (const unsubscribe of this._rlmChildUnsubscribes.values()) { unsubscribe(); @@ -3938,11 +3256,7 @@ export class AgentSession { } get isCompacting(): boolean { - return ( - this._autoCompactionAbortController !== undefined || - this._compactionAbortController !== undefined || - this._branchSummaryAbortController !== undefined - ); + return this._compaction.isRunning || this._branchSummaryAbortController !== undefined; } get messages(): AgentMessage[] { @@ -5163,7 +4477,7 @@ export class AgentSession { compaction: this.isCompacting, retry: this.isRetrying, bash: this.isBashRunning, - refinementApply: this._refineInFlight !== undefined, + refinementApply: this._refinement.isApplying, branchMutation: this._branchSummaryOperation !== undefined, schedulerPauseCount: this._inputScheduler.queuedWorkPauseCount + (this._inputScheduler.suspended ? 1 : 0), disposing: this._disposed || this._disposing, @@ -5202,7 +4516,7 @@ export class AgentSession { await this._commitFence.run(commitFence, async () => { const isCancelled = () => action.lifecycle.state === "cancelled"; if (isCancelled()) return; - await this._waitForRefineIdle(); + await this._refinement._waitForRefineIdle(); if (isCancelled()) return; if (this._isSessionInputHandoffDeferred(epoch) || !canSelectSessionAction(this._runtimeActivity())) { this._actionStore.rollback(action); @@ -5449,7 +4763,7 @@ export class AgentSession { } catch (error) { // Only a failure of the refinement itself is a refine failure; a later // result-row persist error must not report a completed refinement as failed. - this._emitRefineFailed(this._asError(error)); + this._refinement._emitRefineFailed(this._asError(error)); throw error; } const applied = result.appliedEdits.filter((edit) => edit.applied).length; @@ -5852,9 +5166,9 @@ export class AgentSession { this.isCompacting || this.isRetrying || this.isBashRunning || - this._refineInFlight !== undefined || + this._refinement.isApplying || this._branchSummaryOperation !== undefined || - this._postCompactionContinuationSettlement !== undefined || + this._continuation.current !== undefined || this.unfinishedActionCount > 0 ); } @@ -6157,8 +5471,8 @@ export class AgentSession { * waiter registered (a leaked waiter holds hasPendingAdmissionWaiters true and * blocks daemon passivation). */ - private async _waitForIdleOrSettlement(settlement?: PostCompactionContinuationSettlement): Promise { - while (settlement === undefined || this._postCompactionContinuationSettlement === settlement) { + private async _waitForIdleOrSettlement(settlement?: ContinuationToken): Promise { + while (settlement === undefined || this._continuation.current === settlement) { if (this._actionStore.queuedActions().length > 0) { if (this._inputScheduler.suspended || this._inputScheduler.queuedWorkPauseCount > 0) { let wake = () => {}; @@ -6196,7 +5510,7 @@ export class AgentSession { async waitForHeadlessIdle(): Promise { while (true) { await this.waitForIdle(); - const postCompactionContinuation = this._postCompactionContinuationSettlement?.promise; + const postCompactionContinuation = this._continuation.current?.promise; if (!postCompactionContinuation) return; await postCompactionContinuation; } @@ -6267,15 +5581,12 @@ export class AgentSession { this.abortCompaction(); this.abortBranchSummary(); this.abortBash(); - this._pendingRequestedRefine = undefined; - this._autoRefineBranchVersion++; - this._autoRefineReviewAbort?.abort(); - this._refineAbortController?.abort(); + this._refinement.requestAbort(); this.agent.abort(); } async abort(): Promise { - const compactionOperation = this._compactionOperation; + const compactionOperation = this._compaction.operation; const branchSummaryOperation = this._branchSummaryOperation; this.requestAbort(); this._cancelActiveRlmChildRuns("Parent session aborted"); @@ -6683,99 +5994,32 @@ export class AgentSession { this.settingsManager.setFollowUpMode(mode); } - async compact(customInstructions?: string, options: { skipAbort?: boolean } = {}): Promise { - if (options.skipAbort && this.isStreaming) { - throw new Error("Cannot compact without aborting while the agent is running."); - } - const hadPostCompactionContinue = this._postCompactionContinuationScheduled; - const continueAfterSessionInput = this._postCompactionContinuationSettlement?.continueAfterSessionInput ?? false; - this._disconnectFromAgent(); - if (!options.skipAbort) await this.abort(); - let didCompact = false; - const compactionAbort = new AbortController(); - this._compactionAbortController = compactionAbort; - let resolveCompactionOperation: () => void = () => {}; - const compactionOperation = new Promise((resolve) => { - resolveCompactionOperation = resolve; - }); - this._compactionOperation = compactionOperation; - this._emit({ - type: "compaction_start", - reason: "manual", - customInstructions, - }); - - try { - if (!this.model) { - throw new Error(formatNoModelSelectedMessage()); - } - - const { apiKey, headers } = await this._getRequiredRequestAuth(this.model); - const result = await this._performCompaction({ - model: this.model, - apiKey, - headers, - customInstructions, - signal: compactionAbort.signal, - }); + compact(customInstructions?: string, options: { skipAbort?: boolean } = {}): Promise { + return this._compaction.compact(customInstructions, options); + } - this._emit({ - type: "compaction_end", - reason: "manual", - result, - aborted: false, - willRetry: false, - customInstructions, - }); - didCompact = true; - // A manual compaction satisfies any pending model request; on failure the - // request stays scheduled for the next turn boundary. - this._pendingRequestedCompaction = undefined; - return result; - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - const aborted = message === "Compaction cancelled" || (error instanceof Error && error.name === "AbortError"); - const skipped = error instanceof CompactionSkippedError; - this._emit({ - type: "compaction_end", - reason: "manual", - result: undefined, - aborted, - willRetry: false, - errorMessage: aborted ? undefined : skipped ? message : `Compaction failed: ${message}`, - errorSeverity: skipped ? "warning" : "error", - customInstructions, - }); - throw error; - } finally { - this._compactionAbortController = undefined; - this._reconnectToAgent(); - if (this._compactionOperation === compactionOperation) { - this._compactionOperation = undefined; - } - resolveCompactionOperation(); - this._notifySessionInputCheckpointChange(); - this._scheduleSessionInputPump(); - if (didCompact) { - this._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); - if (this._goals.state.status === "active" && !compactionAbort.signal.aborted) { - this._goalContinuationAwaitsRlmWork ||= !this.agent.hasQueuedMessages(); - this.resumeQueuedWork(); - if (this.agent.hasQueuedMessages()) this._schedulePostCompactionContinue(); - } - if (hadPostCompactionContinue) { - this._schedulePostCompactionContinue(continueAfterSessionInput); - } - // Queued agent or session-owned inputs resume the loop; defer refine - // behind them instead of interleaving it before their turns. - this._scheduleAutoRefineAfterCompaction( - this._goalContinuationAwaitsRlmWork || - hadPostCompactionContinue || - this.agent.hasQueuedMessages() || - this.unfinishedActionCount > 0, - ); - } - } + private _afterManualCompaction( + signal: AbortSignal, + hadPostCompactionContinue: boolean, + continueAfterSessionInput: boolean, + ): void { + this._refinement._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); + if (this._goals.state.status === "active" && !signal.aborted) { + this._goalContinuationAwaitsRlmWork ||= !this.agent.hasQueuedMessages(); + this.resumeQueuedWork(); + if (this.agent.hasQueuedMessages()) this._schedulePostCompactionContinue(); + } + if (hadPostCompactionContinue) { + this._schedulePostCompactionContinue(continueAfterSessionInput); + } + // Queued agent or session-owned inputs resume the loop; defer refine + // behind them instead of interleaving it before their turns. + this._refinement._scheduleAutoRefineAfterCompaction( + this._goalContinuationAwaitsRlmWork || + hadPostCompactionContinue || + this.agent.hasQueuedMessages() || + this.unfinishedActionCount > 0, + ); } /** @@ -6783,156 +6027,8 @@ export class AgentSession { * skill. Throws CompactionSkippedError when there is nothing to compact and * Error("Compaction cancelled") on abort or extension cancel. */ - private async _performCompaction(options: { - model: Model; - apiKey: string; - headers?: Record; - customInstructions?: string; - signal: AbortSignal; - }): Promise { - const { model, apiKey, headers, customInstructions, signal } = options; - const pathEntries = this.sessionManager.getBranch(); - const settings = this.settingsManager.getCompactionSettings(); - - const preparation = prepareCompaction(pathEntries, settings); - if (!preparation) { - const lastEntry = pathEntries[pathEntries.length - 1]; - if (lastEntry?.type === "compaction") { - throw new CompactionSkippedError("Already compacted"); - } - throw new CompactionSkippedError("Session is too short to compact — try again once it grows"); - } - - let extensionCompaction: CompactionResult | undefined; - let fromExtension = false; - - const semanticCompaction = this._semanticEdges.beginCompaction(); - let compactionRecorded = false; - const uncommittedSlices: string[] = []; - let compactionSettled = false; - let summary: string; - let firstKeptEntryId: string; - let tokensBefore: number; - let details: CompactionResult["details"]; - let usage: CompactionResult["usage"]; - try { - if (this._extensionRunner.hasHandlers("session_before_compact")) { - const result = (await this._extensionRunner.emit({ - type: "session_before_compact", - preparation, - branchEntries: pathEntries, - customInstructions, - signal, - })) as SessionBeforeCompactResult | undefined; - - if (result?.cancel) { - throw new Error("Compaction cancelled"); - } - - if (result?.compaction) { - extensionCompaction = result.compaction; - fromExtension = true; - } - } - - if (extensionCompaction) { - ({ summary, firstKeptEntryId, tokensBefore, details, usage } = extensionCompaction); - } else { - // Each summary wire call gets its own request ID: split turns send two - // different bodies, and one Idempotency-Key must never cover both. A slice - // that succeeds on the wire stays uncommitted until the compaction itself - // commits: a racing sibling's failure (or an abort) must leave no committed - // summary request for the next turn's continuation edge to attach to. - const summaryCall = async ( - call: (callHeaders: Record | undefined) => Promise, - ): Promise => { - const requestId = this._semanticEdges.startCompactionRequest(semanticCompaction.compactionId); - if (requestId === undefined) { - return call(headers); - } - try { - const result = await call({ ...headers, ...modelRequestHeaders(requestId) }); - // A slice resolving after a sibling's rejection already settled the - // compaction would push into a drained list and stay in-flight forever. - if (compactionSettled) { - this._semanticEdges.failRequest(requestId); - } else { - uncommittedSlices.push(requestId); - } - return result; - } catch (error) { - this._semanticEdges.failRequest(requestId); - throw error; - } - }; - ({ summary, firstKeptEntryId, tokensBefore, details, usage } = await compact( - preparation, - model, - apiKey, - headers, - customInstructions, - signal, - this.thinkingLevel, - summaryCall, - providerRetryPolicy(this.settingsManager), - )); - } - - if (signal.aborted) { - throw new Error("Compaction cancelled"); - } - - // Ledger-before-effect: the compaction outcome is durable before the transcript - // commits it. Marked first: the ID is consumed even when the write throws, and a - // second finish attempt would mask the original I/O error. - compactionRecorded = true; - compactionSettled = true; - for (const requestId of uncommittedSlices.splice(0)) { - this._semanticEdges.finishRequest(requestId); - } - this._semanticEdges.finishCompaction(semanticCompaction.compactionId, "completed"); - // Attached mechanically; the digest never flows through the summarizer LLM. - this.sessionManager.appendCompaction( - summary, - firstKeptEntryId, - tokensBefore, - details, - fromExtension, - customInstructions, - usage, - this._harnessDigest(), - ); - } catch (error) { - compactionSettled = true; - for (const requestId of uncommittedSlices.splice(0)) { - this._semanticEdges.failRequest(requestId); - } - if (!compactionRecorded) { - const cancelled = - error instanceof Error && (error.name === "AbortError" || error.message === "Compaction cancelled"); - this._semanticEdges.finishCompaction(semanticCompaction.compactionId, cancelled ? "cancelled" : "failed"); - } - throw error; - } - const newEntries = this.sessionManager.getEntries(); - this.agent.state.messages = this.sessionManager.buildSessionContext().messages; - this._mergeUnpersistedOutcomes(this.agent.state.messages); - this._restoreLateIpythonSentAgentMessages(); - - const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as - | CompactionEntry - | undefined; - if (savedCompactionEntry) { - await this._extensionRunner.emit({ - type: "session_compact", - compactionEntry: savedCompactionEntry, - fromExtension, - }); - } - await this._syncKernelStateAfterCompaction(); - await this._reapDeletedRlmSubagentRuntimesAfterCompaction(); - - return { summary, firstKeptEntryId, tokensBefore, details }; + private _performCompaction(options: CompactionExecutionOptions): Promise { + return performSessionCompaction(this._compactionExecution, options); } private async _reapDeletedRlmSubagentRuntimesAfterCompaction(): Promise { @@ -6943,1260 +6039,125 @@ export class AgentSession { } abortCompaction(): void { - this._compactionAbortController?.abort(); - this._autoCompactionAbortController?.abort(); + this._compaction.abort(); } - private _localHarnessStateDir(): string | undefined { - return ( - getLocalHarnessStateDir(this.sessionManager.getSessionArtifactDir()) ?? - (this._rlmSessionDir ? getLocalHarnessStateDir(this._rlmSessionDir) : undefined) - ); + private _cancelPostCompactionContinue(): void { + this._continuation.cancel(); } - private _autoRefineAllowedForSession(): boolean { - return this._rlmDepth === 0 && this._localHarnessStateDir() !== undefined; + private _schedulePostCompactionContinue(continueAfterSessionInput = false): void { + this._continuation.schedule(continueAfterSessionInput); } - private _settlePostCompactionContinue(error?: Error): void { - if (!error && this._postCompactionContinuationScheduled) return; - const settlement = this._postCompactionContinuationSettlement; - if (!settlement || settlement.settled) return; - settlement.settled = true; - this._postCompactionContinuationSettlement = undefined; - if (error) settlement.reject(error); - else settlement.resolve(); - this._notifySessionInputCheckpointChange(); + private _forgetConsumedPostCompactionContinuations(messages: AgentMessage[]): void { + this._continuation.forgetConsumed(messages); } - private _cancelPostCompactionContinue(): void { - this._postCompactionContinuationScheduled = false; - this._scheduledPostCompactionContinuationMessages = []; - this._settlePostCompactionContinue(); + /** The compact harness digest delivered at cold context boundaries (session start, resume, compaction head). */ + private _harnessDigest(): string { + const tools = this.getActiveToolNames(); + const hasIpython = tools.includes("ipython"); + const visibleSkills = this._modelVisibleSkills().filter((skill) => !skill.disableModelInvocation); + const hasRefineSkill = visibleSkills.some((skill) => skill.name === REFINE_SKILL_NAME); + return formatHarnessStateForPrompt(this._refinement._loadMergedHarnessState(), { + includeIpythonExamples: hasIpython, + includeShellExamples: tools.includes("bash"), + includeRefineExamples: hasIpython && hasRefineSkill, + }); } - private _discardPendingAutoRefine(options: { cancelPostCompactionContinue?: boolean } = {}): void { - this._compactAutoRefinePending = false; - this._turnIntervalAutoRefinePending = false; - this._pendingAutoRefineReview = undefined; - if (options.cancelPostCompactionContinue) { - this._cancelPostCompactionContinue(); + /** Cold-boundary digest delivery: empty contexts defer to the first committed turn (untouched sessions must stay empty); non-empty contexts append only when the newest in-context digest mismatches disk. */ + private _ensureHarnessDigestContext(): void { + if (this.agent.state.messages.length === 0) { + this._harnessDigestPending = true; + return; } + this._harnessDigestPending = false; + this._appendHarnessDigestIfStale(); } - private async _invalidatePendingAutoRefineForBranchChange(): Promise { - this._autoRefineReviewAbort?.abort(); - this._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); - this._assistantTurnsSinceAutoRefine = 0; - // Increment branch version BEFORE aborting/awaiting the serialized plan. - // This invalidates the plan's branchVersion check at the boundary - // so even if the plan completes, the boundary will reject it - // (bgResult.branchVersion !== this._autoRefineBranchVersion). - this._autoRefineBranchVersion++; - // Abort the in-flight refine/bplan controller so any pending - // _planRefine or _reviewAutoRefine call settles via signal abort - // rather than hanging forever. - this._refineAbortController?.abort(); - if (this._serializedPlanInFlight) { - await this._consumeSerializedBackgroundPlan(async () => false); + private _appendHarnessDigestIfStale(): void { + const digest = this._harnessDigest(); + if (this._latestContextHarnessDigest() === digest) return; + const message = createHarnessDigestMessage(digest); + try { + this.sessionManager.appendCustomMessageEntryWithRollback( + message.customType, + message.content, + message.display, + message.details, + ); + } catch { + // Unpersisted session: context-only injection. } - while (this._refinePlanInFlight) { - await this._refinePlanInFlight; + this.agent.state.messages.push(message); + } + + private _latestContextHarnessDigest(): string | undefined { + // Retained pre-compaction messages follow the compaction head, so recency is by timestamp, not position. + let latest: { timestamp: number; digest: string } | undefined; + for (const message of this.agent.state.messages) { + let digest: string | undefined; + if (message.role === "custom" && message.customType === HARNESS_DIGEST_CUSTOM_TYPE) { + digest = (message.details as HarnessDigestDetails | undefined)?.digest; + } else if (message.role === "compactionSummary") { + digest = message.harnessDigest; + } else { + continue; + } + if (digest !== undefined && (!latest || message.timestamp >= latest.timestamp)) { + latest = { timestamp: message.timestamp, digest }; + } } - await this._waitForRefineIdle(); + return latest?.digest; } /** - * Consume a refine request that was scheduled by the agent-callable refine - * skill (refine.run). Fire-and-forget: the refine() method handles its own - * background planning, idle wait, application, and error recovery. Called - * at the turn boundary after compaction checks and before auto-refine - * scheduling so the manual request takes priority. + * Refine editable continual harness state: prompt notes, memory, skills, and subagent specs. + * The base system prompt is intentionally not editable through this path. + * + * Planning runs in the background and does NOT block turn entry points + * (`_waitForRefineIdle` only waits for `_refineInFlight`). Only the fast + * application phase (disk I/O + in-memory mutation) blocks turn entry points. */ - private _emitRefineFailed(error: unknown): void { - this._emit({ - type: "refine_failed", - error: error instanceof Error ? error.message : String(error), - }); + refine( + options: { instructions?: string; rollbackId?: string; global?: boolean } = {}, + internal: { skipAbort?: boolean; trigger?: "manual" | "auto"; source?: RefinementSource } = {}, + ): Promise { + return this._refinement.refine(options, internal); } - private _consumePendingRequestedRefine(): boolean { - const pending = this._pendingRequestedRefine; - if (!pending) return false; - this._pendingRequestedRefine = undefined; - void this.refine(pending, { source: "self" }).catch((error) => this._emitRefineFailed(error)); - return true; + abortBranchSummary(): void { + this._branchSummaryAbortController?.abort(); } - private _scheduleAutoRefineAfterAgentEnd(): void { - if (!this._autoRefineAllowedForSession()) { - return; - } - if (this._pendingAutoRefineReview) { - this._scheduleAutoRefine(this._pendingAutoRefineReview.reason); - return; - } - if (this._compactAutoRefinePending) { - if (this._postCompactionContinuationScheduled) { - return; - } - this._scheduleAutoRefine("compact"); - return; - } - - this._scheduleAutoRefine("turn_interval"); - } - - private _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void { - if (!this._autoRefineAllowedForSession()) { - return; - } - if (this._serializedRefine) { - // Serialized sessions must service compaction-triggered refinement at - // shouldStopAfterTurn (or disposal), never through the interactive path. - this._compactAutoRefinePending = true; - return; - } - if (willContinueAfterCompaction) { - this._compactAutoRefinePending = true; - return; - } - - this._scheduleAutoRefine("compact"); - } - - private _schedulePostCompactionContinue(continueAfterSessionInput = false): void { - if (!this._postCompactionContinuationSettlement || this._postCompactionContinuationSettlement.settled) { - this._postCompactionContinuationSettlement = createPostCompactionContinuationSettlement(); - } - const settlement = this._postCompactionContinuationSettlement; - settlement.continueAfterSessionInput ||= continueAfterSessionInput; - if (this._postCompactionContinuationScheduled) { - return; - } - this._postCompactionContinuationScheduled = true; - this._scheduledPostCompactionContinuationMessages = [...this._postCompactionContinuationMessages]; - void this._runScheduledPostCompactionContinue(settlement) - .catch(() => undefined) - .finally(() => { - if (this._postCompactionContinuationSettlement === settlement) { - this._settlePostCompactionContinue(); - } - }); - } - - private _sessionOwnsScheduledContinuations(continuationMessages: AgentMessage[]): boolean { - return continuationMessages.some((message) => this._postCompactionContinuationMessages.includes(message)); - } - - private async _waitForQueuedWorkResume(settlement: PostCompactionContinuationSettlement): Promise { - while ( - this._inputScheduler.queuedWorkPauseCount > 0 && - this._postCompactionContinuationSettlement === settlement - ) { - let resume = () => {}; - const resumed = new Promise((resolve) => { - resume = resolve; - this._sessionInputCheckpointWaiters.add(resolve); - }); - try { - await Promise.race([resumed, settlement.promise]); - } finally { - this._sessionInputCheckpointWaiters.delete(resume); - } - } - } - - private async _runScheduledPostCompactionContinue(settlement: PostCompactionContinuationSettlement): Promise { - while (this._postCompactionContinuationScheduled && this._postCompactionContinuationSettlement === settlement) { - await this.agent.waitForIdle(); - await this.waitForRetry(); - await this._waitForRefineIdle(); - await this._waitForQueuedWorkResume(settlement); - const compactionOperation = this._compactionOperation; - if (compactionOperation) { - await Promise.race([compactionOperation, settlement.promise]); - continue; - } - - const commitFence = await this._acquireSessionActionCommitFence(); - let continuation: Promise | undefined; - let continuationMessages: AgentMessage[] = []; - let waitForSessionInput = false; - try { - await this.agent.waitForIdle(); - if ( - !this._postCompactionContinuationScheduled || - this._postCompactionContinuationSettlement !== settlement - ) { - return; - } - - if (this._inputScheduler.queuedWorkPauseCount > 0 || this._compactionOperation || this._refineInFlight) { - continue; - } - - continuationMessages = [...this._scheduledPostCompactionContinuationMessages]; - if (continuationMessages.length > 0 && !this._sessionOwnsScheduledContinuations(continuationMessages)) { - this._cancelPostCompactionContinue(); - this._scheduleAutoRefineAfterAgentEnd(); - return; - } - if (this.unfinishedActionCount > 0 || this._inputScheduler.requested) { - this._scheduleSessionInputPump(); - waitForSessionInput = true; - } else { - this._postCompactionContinuationScheduled = false; - continuation = this.agent.continue(); - } - } finally { - commitFence.release(); - } - - if (waitForSessionInput) { - await this._waitForIdleOrSettlement(settlement); - if (this._postCompactionContinuationSettlement !== settlement) return; - const shouldContinue = - (settlement.continueAfterSessionInput && continuationMessages.length === 0) || - this._sessionOwnsScheduledContinuations(continuationMessages); - if (shouldContinue) { - this._scheduledPostCompactionContinuationMessages = [...this._postCompactionContinuationMessages]; - continue; - } - this._postCompactionContinuationScheduled = false; - this._scheduledPostCompactionContinuationMessages = []; - this._scheduleAutoRefineAfterAgentEnd(); - return; - } - - try { - await continuation; - if (this._postCompactionContinuationSettlement === settlement) { - this._forgetConsumedPostCompactionContinuations(continuationMessages); - } - return; - } catch (error) { - const code = error instanceof AgentContinueError ? error.code : undefined; - if (code === "busy") { - if (this._postCompactionContinuationSettlement === settlement) { - this._postCompactionContinuationScheduled = true; - this._scheduledPostCompactionContinuationMessages = [...this._postCompactionContinuationMessages]; - } - continue; - } - if (code !== "nothing-to-continue" && this._postCompactionContinuationSettlement === settlement) { - this._settlePostCompactionContinue(this._asError(error)); - } - return; - } - } - } - - private _forgetConsumedPostCompactionContinuations(continuationMessages: AgentMessage[]): void { - if (continuationMessages.length === 0) { - return; - } - const continuationMessageSet = new Set(continuationMessages); - const stillQueued = new Set(this.agent.removeQueuedMessages((message) => continuationMessageSet.has(message))); - for (const message of stillQueued) { - this.agent.followUp(message); - } - for (const message of continuationMessages) { - if (!stillQueued.has(message)) { - this._queuedAutonomousContinuationSnapshots.delete(message); - } - } - this._postCompactionContinuationMessages = this._postCompactionContinuationMessages.filter( - (message) => !continuationMessageSet.has(message) || stillQueued.has(message), - ); - } - - private _shouldSkipAutoRefineForActiveAgent(): boolean { - return this.isStreaming || this.isCompacting; - } - - private _scheduleDeferredAutoRefineIfIdle(): void { - if (this._autoRefineInProgress || this._shouldSkipAutoRefineForActiveAgent() || this._pendingAutoRefineReview) { - return; - } - if (this._turnIntervalAutoRefinePending) { - this._turnIntervalAutoRefinePending = false; - this._scheduleAutoRefine("turn_interval"); - } - } - - private _scheduleAutoRefine(reason: AutoRefineReason, branchVersion = this._autoRefineBranchVersion): void { - const timer = setTimeout(() => { - this._scheduledAutoRefineTimers.delete(timer); - if (branchVersion !== this._autoRefineBranchVersion) { - return; - } - const operation = this._maybeAutoRefine(reason); - this._autoRefineOperations.add(operation); - void operation.finally(() => this._autoRefineOperations.delete(operation)).catch(() => undefined); - }, 0); - this._scheduledAutoRefineTimers.add(timer); - } - - private async _maybeAutoRefine(reason: AutoRefineReason): Promise { - if (this._disposed || this._disposing) { - this._discardPendingAutoRefine(); - return; - } - if (!this._autoRefineAllowedForSession()) { - this._discardPendingAutoRefine(); - return; - } - - const settings = this.settingsManager.getAutoRefineSettings(); - if (!settings.enabled) { - this._discardPendingAutoRefine(); - return; - } - if (this._autoRefineInProgress || this._shouldSkipAutoRefineForActiveAgent()) { - if (reason === "compact") { - this._compactAutoRefinePending = true; - } else { - this._turnIntervalAutoRefinePending = true; - } - return; - } - - const nowMs = Date.now(); - const underCooldown = - this._lastAutoRefineReviewAt > 0 && nowMs - this._lastAutoRefineReviewAt < settings.cooldownMs; - - const pendingReview = this._pendingAutoRefineReview; - if (pendingReview) { - // A failed refine stamps the cooldown; keep the pending review for later. - if (underCooldown) { - return; - } - await this._runApprovedRefine(pendingReview.reason, pendingReview.review); - return; - } - - if (reason === "compact" && !settings.compact) { - this._compactAutoRefinePending = false; - reason = "turn_interval"; - } - if (reason === "turn_interval" && this._assistantTurnsSinceAutoRefine < settings.turnInterval) { - return; - } - if (underCooldown) { - if (reason === "compact") { - this._compactAutoRefinePending = true; - } else { - this._turnIntervalAutoRefinePending = true; - } - return; - } - if (reason === "turn_interval") { - this._turnIntervalAutoRefinePending = false; - } - if (!this.model) { - if (reason === "compact") { - this._compactAutoRefinePending = true; - } - return; - } - this._autoRefineInProgress = true; - const turnsSinceLastReview = this._assistantTurnsSinceAutoRefine; - const branchVersion = this._autoRefineBranchVersion; - const reviewAbort = new AbortController(); - this._autoRefineReviewAbort = reviewAbort; - let approvedReview: AutoRefineReview | undefined; - try { - const review = await this._reviewAutoRefine({ reason, turnsSinceLastReview }, reviewAbort.signal); - if (this._disposed || this._disposing || branchVersion !== this._autoRefineBranchVersion) { - return; - } - if (!review.shouldRefine) { - const preserveTurnIntervalReview = - reason === "compact" && this._assistantTurnsSinceAutoRefine >= settings.turnInterval; - if (preserveTurnIntervalReview) { - this._turnIntervalAutoRefinePending = true; - } else { - this._lastAutoRefineReviewAt = nowMs; - this._assistantTurnsSinceAutoRefine = 0; - } - if (reason === "compact") { - this._compactAutoRefinePending = false; - } - return; - } - if (this._shouldSkipAutoRefineForActiveAgent()) { - this._pendingAutoRefineReview = { reason, review }; - return; - } - approvedReview = review; - } catch { - // Failed review: stamp the cooldown so a persistent failure (bad auth, - // unparseable output) doesn't retry a full review on every agent end. - if (branchVersion === this._autoRefineBranchVersion) { - this._lastAutoRefineReviewAt = Date.now(); - } - } finally { - if (this._autoRefineReviewAbort === reviewAbort) { - this._autoRefineReviewAbort = undefined; - } - this._autoRefineInProgress = false; - // When a refine follows, _runApprovedRefine schedules the deferred pass. - if (!approvedReview) { - this._scheduleDeferredAutoRefineIfIdle(); - } - } - if (approvedReview) { - await this._runApprovedRefine(reason, approvedReview); - } - } - - private async _runApprovedRefine(reason: AutoRefineReason, review: AutoRefineReview): Promise { - this._autoRefineInProgress = true; - try { - await this.refine({ instructions: autoRefineInstructions(reason, review) }, { trigger: "auto" }); - this._pendingAutoRefineReview = undefined; - this._turnIntervalAutoRefinePending = false; - this._lastAutoRefineReviewAt = Date.now(); - this._assistantTurnsSinceAutoRefine = 0; - if (reason === "compact") { - this._compactAutoRefinePending = false; - } - } catch (error) { - // Auto-refine is opportunistic; manual /refine remains available. - // Stamp the cooldown so a persistently failing refine doesn't retry - // (via a retained pending review) on every agent end. - this._lastAutoRefineReviewAt = Date.now(); - if (error instanceof RefineSkippedError) { - // A skipped round is consumed like a reviewer decline, not retained for retry. - this._pendingAutoRefineReview = undefined; - this._turnIntervalAutoRefinePending = false; - this._assistantTurnsSinceAutoRefine = 0; - if (reason === "compact") this._compactAutoRefinePending = false; - } - } finally { - this._autoRefineInProgress = false; - this._scheduleDeferredAutoRefineIfIdle(); - } - } - - private async _reviewAutoRefine(context: AutoRefineReviewRequest, signal?: AbortSignal): Promise { - if (this._autoRefineReviewer) { - return this._autoRefineReviewer(context, signal); - } - const model = this.model; - if (!model) { - return { shouldRefine: false, rationale: "No model selected." }; - } - const { apiKey, headers } = await this._getRequiredRequestAuth(model); - return reviewAutoRefine( - this.agent.state.messages, - this._loadMergedHarnessState(), - this._loadRefinementHistory(), - model, - apiKey, - context, - headers, - signal, - this.thinkingLevel, - providerRetryPolicy(this.settingsManager), - ); - } - - /** The compact harness digest delivered at cold context boundaries (session start, resume, compaction head). */ - private _harnessDigest(): string { - const tools = this.getActiveToolNames(); - const hasIpython = tools.includes("ipython"); - const visibleSkills = this._modelVisibleSkills().filter((skill) => !skill.disableModelInvocation); - const hasRefineSkill = visibleSkills.some((skill) => skill.name === REFINE_SKILL_NAME); - return formatHarnessStateForPrompt(this._loadMergedHarnessState(), { - includeIpythonExamples: hasIpython, - includeShellExamples: tools.includes("bash"), - includeRefineExamples: hasIpython && hasRefineSkill, - }); - } - - /** Cold-boundary digest delivery: empty contexts defer to the first committed turn (untouched sessions must stay empty); non-empty contexts append only when the newest in-context digest mismatches disk. */ - private _ensureHarnessDigestContext(): void { - if (this.agent.state.messages.length === 0) { - this._harnessDigestPending = true; - return; - } - this._harnessDigestPending = false; - this._appendHarnessDigestIfStale(); - } - - private _appendHarnessDigestIfStale(): void { - const digest = this._harnessDigest(); - if (this._latestContextHarnessDigest() === digest) return; - const message = createHarnessDigestMessage(digest); - try { - this.sessionManager.appendCustomMessageEntryWithRollback( - message.customType, - message.content, - message.display, - message.details, - ); - } catch { - // Unpersisted session: context-only injection. - } - this.agent.state.messages.push(message); - } - - private _latestContextHarnessDigest(): string | undefined { - // Retained pre-compaction messages follow the compaction head, so recency is by timestamp, not position. - let latest: { timestamp: number; digest: string } | undefined; - for (const message of this.agent.state.messages) { - let digest: string | undefined; - if (message.role === "custom" && message.customType === HARNESS_DIGEST_CUSTOM_TYPE) { - digest = (message.details as HarnessDigestDetails | undefined)?.digest; - } else if (message.role === "compactionSummary") { - digest = message.harnessDigest; - } else { - continue; - } - if (digest !== undefined && (!latest || message.timestamp >= latest.timestamp)) { - latest = { timestamp: message.timestamp, digest }; - } - } - return latest?.digest; - } - - /** Global harness state overlaid with this session's local state, when persisted. */ - private _loadMergedHarnessState(): HarnessState { - const localHarnessStateDir = this._localHarnessStateDir(); - return mergeHarnessStates( - loadHarnessState(getGlobalHarnessStateDir(), "global"), - localHarnessStateDir ? loadHarnessState(localHarnessStateDir, "local") : undefined, - ); - } - - private _loadRefinementHistory(): RefinementResult[] { - return mergeRefinementHistory( - loadGlobalRefinementHistory(getGlobalHarnessStateDir()), - getRefinementHistory(this.sessionManager.getEntries().filter((entry) => entry.type === "custom")), - ); - } - - /** - * Refine editable continual harness state: prompt notes, memory, skills, and subagent specs. - * The base system prompt is intentionally not editable through this path. - * - * Planning runs in the background and does NOT block turn entry points - * (`_waitForRefineIdle` only waits for `_refineInFlight`). Only the fast - * application phase (disk I/O + in-memory mutation) blocks turn entry points. - */ - async refine( - options: { - instructions?: string; - rollbackId?: string; - global?: boolean; - } = {}, - internal: { skipAbort?: boolean; trigger?: "manual" | "auto"; source?: RefinementSource } = {}, - ): Promise { - // Queued /refine executes from the session-input pump between turns; - // refine never aborts the agent (planning is backgrounded and the apply - // phase waits for quiescence), so skipAbort only asserts the pump's - // idle invariant instead of changing abort behavior. - if (internal.skipAbort && this.isStreaming) { - throw new Error("Cannot refine without aborting while the agent is running."); - } - // Wait for any existing refine (both planning and application) before - // starting a new run. This serializes concurrent /refine calls so two - // planning phases cannot race into concurrent _applyRefine calls that - // overwrite harness state. - while (this._refineInFlight || this._refinePlanInFlight || this._serializedPlanInFlight) { - if (this._refineInFlight) { - await this._refineInFlight; - } else if (this._refinePlanInFlight) { - await this._refinePlanInFlight; - } else { - // A serialized background plan is in flight (started during an - // active turn at message_end). Wait for planning and for the active - // turn to settle so its normal checkpoint can consume the plan. - const serializedPlanInFlight = this._serializedPlanInFlight; - await serializedPlanInFlight; - if (this._refineInFlight || this._refinePlanInFlight) { - continue; - } - await this.agent.waitForIdle(); - // Aborted turns skip shouldStopAfterTurn. Drop their settled plan - // after idle so a later public refine cannot spin on it forever. - if (this._serializedPlanInFlight === serializedPlanInFlight) { - this._serializedPlanInFlight = undefined; - this._serializedExplicitRefineOptions = undefined; - } - } - } - - const refineAbort = new AbortController(); - this._refineAbortController = refineAbort; - - const planRun = this._planRefine(options, refineAbort.signal, internal.trigger ?? "manual"); - const planSettled = planRun.then( - () => undefined, - () => undefined, - ); - this._refinePlanInFlight = planSettled; - let plan: RefinementPlan; - try { - plan = await planRun; - } catch (e) { - if (this._refineAbortController === refineAbort) { - this._refineAbortController = undefined; - } - this._scheduleSessionInputPump(); - throw e; - } finally { - if (this._refinePlanInFlight === planSettled) { - this._refinePlanInFlight = undefined; - } - } - - // Block new turns before waiting for the current turn to finish. One shared - // settled promise covers the full transition and apply critical section. - let resolveApplySettled: () => void = () => {}; - const applySettled = new Promise((resolve) => { - resolveApplySettled = resolve; - }); - this._refineInFlight = applySettled; - try { - // Wait for the session to become quiescent before applying. Planning is - // allowed to overlap active user work, but application must not disconnect - // event handling until that work and its queued events have completed. - await this.agent.waitForIdle(); - while (true) { - const eventQueue = this._agentEventQueue; - const compactionOp = this._compactionOperation; - const branchSummaryOp = this._branchSummaryOperation; - await Promise.allSettled([ - eventQueue, - ...(compactionOp ? [compactionOp] : []), - ...(branchSummaryOp ? [branchSummaryOp] : []), - ]); - if ( - eventQueue === this._agentEventQueue && - compactionOp === this._compactionOperation && - branchSummaryOp === this._branchSummaryOperation - ) { - break; - } - } - if (this._disposed || refineAbort.signal.aborted) { - throw new Error("Refinement cancelled because the session was disposed."); - } - return await this._applyRefine( - plan, - options, - refineAbort, - internal.source ?? (internal.trigger === "auto" ? "auto" : "user"), - ); - } finally { - resolveApplySettled(); - if (this._refineInFlight === applySettled) { - this._refineInFlight = undefined; - } - this._notifySessionInputCheckpointChange(); - this._scheduleSessionInputPump(); - } - } - - /** - * Block a new agent turn until any in-flight refine application phase has - * reattached event handling; otherwise the turn's messages are never - * persisted or rendered. - * - * The idle-wait and application phase (`_refineInFlight`) block here. The - * background planning phase (`_refinePlanInFlight`) does NOT block turns. - * Refine failures surface to the refine caller, not here. - */ - private async _waitForRefineIdle(): Promise { - while (this._refineInFlight) { - await this._refineInFlight; - } - } - - /** - * Background planning phase: runs the LLM planning call via `planRefinement`. - * Does not disconnect from or abort the agent. Returns the plan without - * applying anything. - */ - private async _planRefine( - options: { instructions?: string; rollbackId?: string; global?: boolean }, - signal: AbortSignal, - trigger: "manual" | "auto" = "manual", - ): Promise { - if (this._disposed) { - throw new Error("Cannot refine a disposed session."); - } - - if (!this.model) { - throw new Error(formatNoModelSelectedMessage()); - } - - const model = this.model; - const { apiKey, headers } = await this._getRequiredRequestAuth(model); - const globalHarnessStateDir = getGlobalHarnessStateDir(); - const localHarnessStateDir = this._localHarnessStateDir(); - const requestedScope = options.global ? "global" : "local"; - if (!options.rollbackId && requestedScope === "local" && !localHarnessStateDir) { - throw new Error("Local harness refinement requires a persisted session; use global refinement instead."); - } - const globalPlanningState = loadHarnessState(globalHarnessStateDir, "global"); - const localPlanningState = localHarnessStateDir ? loadHarnessState(localHarnessStateDir, "local") : undefined; - const planningState = - requestedScope === "global" - ? globalPlanningState - : mergeHarnessStates(globalPlanningState, localPlanningState); - const history = this._loadRefinementHistory(); - const rollbackTarget = options.rollbackId ? history.find((item) => item.id === options.rollbackId) : undefined; - let baselineScope = rollbackTarget - ? (inferRefinementResultScope(rollbackTarget) ?? requestedScope) - : requestedScope; - let baselineHarnessStateDir = baselineScope === "global" ? globalHarnessStateDir : localHarnessStateDir; - if (rollbackTarget?.harnessStatePath) { - baselineHarnessStateDir = dirname(rollbackTarget.harnessStatePath); - baselineScope = resolve(baselineHarnessStateDir) === resolve(globalHarnessStateDir) ? "global" : "local"; - } - if (!baselineHarnessStateDir) { - throw new Error("Local harness refinement requires a persisted session; use global refinement instead."); - } - const baselineState = rollbackTarget - ? loadHarnessState(baselineHarnessStateDir, baselineScope) - : baselineScope === "global" - ? globalPlanningState - : localPlanningState!; - if (!options.rollbackId && this._extensionRunner.hasHandlers("session_before_refine")) { - const result = (await this._extensionRunner.emit({ - type: "session_before_refine", - preparation: { - trigger, - instructions: options.instructions, - scope: requestedScope, - planningState, - history, - conversationText: serializeConversation(convertToLlm(this.agent.state.messages)).slice(-80_000), - }, - signal, - })) as SessionBeforeRefineResult | undefined; - if (this._disposed || signal.aborted) { - throw new Error("Refinement cancelled because the session was disposed."); - } - if (result?.skip) { - throw new RefineSkippedError("Refinement skipped by extension"); - } - if (result?.proposal !== undefined) { - return { - proposal: normalizeRefinementProposal(result.proposal), - id: generateRefinementId(), - baselineState, - }; - } - } - const plan = await planRefinement( - this.agent.state.messages, - planningState, - history, - model, - apiKey, - { ...options, retry: providerRetryPolicy(this.settingsManager) }, - headers, - signal, - this.thinkingLevel, - ); - if (this._disposed || signal.aborted) { - throw new Error("Refinement cancelled because the session was disposed."); - } - return { ...plan, baselineState }; - } - - private _recordRefinementOutcome(result: RefinementResult): void { - this._appendDurableRefineMessage(createRefinementOutcomeMessage(result)); - } - - /** In-context notice for an applied refinement. Refinements with zero applied edits emit nothing. */ - private _recordRefinementNotice(result: RefinementResult, source: RefinementSource): void { - if (!result.appliedEdits.some((edit) => edit.applied)) return; - this._appendDurableRefineMessage(createRefinementNoticeMessage(result, source)); - } - - private _appendDurableRefineMessage(message: CustomMessage): void { - try { - this.sessionManager.appendCustomMessageEntryWithRollback( - message.customType, - message.content, - message.display, - message.details, - ); - } catch { - // Not in the session file, so context rebuilds would drop the outcome. - this._unpersistedOutcomes.push(message); - } - this.agent.state.messages.push(message); - this._emit({ type: "message_start", message }); - this._emit({ type: "message_end", message }); - } - - /** - * Synchronous application phase: disconnects from the agent, aborts any - * in-flight agent run, applies the refinement plan to disk and memory, then - * reconnects. This is the only phase that blocks turn entry points. - */ - private async _applyRefine( - plan: RefinementPlan, - options: { instructions?: string; rollbackId?: string; global?: boolean }, - refineAbort: AbortController, - source: RefinementSource, - ): Promise { - if (this._disposed) { - throw new Error("Cannot refine a disposed session."); - } - // The caller has already set _refineInFlight and waited for agent idle. - // Disconnect only for the brief apply + save + reconnect critical section. - this._disconnectFromAgent(); - - try { - const globalHarnessStateDir = getGlobalHarnessStateDir(); - const localHarnessStateDir = this._localHarnessStateDir(); - const requestedScope = options.global ? "global" : "local"; - const history = this._loadRefinementHistory(); - const rollbackTarget = options.rollbackId ? history.find((item) => item.id === options.rollbackId) : undefined; - let targetScope = plan.rollbackScope ?? requestedScope; - let targetHarnessStateDir = targetScope === "global" ? globalHarnessStateDir : localHarnessStateDir; - if (targetScope === "local" && rollbackTarget?.harnessStatePath) { - if (!existsSync(rollbackTarget.harnessStatePath)) { - throw new Error( - `Local refinement ${rollbackTarget.id} state file not found: ${rollbackTarget.harnessStatePath}`, - ); - } - targetHarnessStateDir = dirname(rollbackTarget.harnessStatePath); - // Legacy records predate scope fields and default to "local" but may point - // at the global store; honor the recorded path so its entries stay global. - if (resolve(targetHarnessStateDir) === resolve(globalHarnessStateDir)) { - targetScope = "global"; - } - } - if (!targetHarnessStateDir) { - throw new Error("Local harness refinement requires a persisted session; use global refinement instead."); - } - // Re-read the target state immediately before applying so concurrent kernel - // (`rlm.harness`) writes during the LLM pass are not clobbered. - const state = loadHarnessState(targetHarnessStateDir, targetScope); - const proposal = { - ...plan.proposal, - edits: plan.proposal.edits.map((edit) => { - const localPrefix = "local:"; - const globalPrefix = "global:"; - return { - ...edit, - id: edit.id?.startsWith(localPrefix) - ? edit.id.slice(localPrefix.length) - : edit.id?.startsWith(globalPrefix) - ? edit.id.slice(globalPrefix.length) - : edit.id, - }; - }), - }; - if (this._disposed || refineAbort.signal.aborted) { - throw new Error("Refinement cancelled because the session was disposed."); - } - const result = applyRefinementProposal(state, proposal, { - id: plan.id, - rollbackOf: plan.rollbackOf, - scope: targetScope, - baselineState: plan.baselineState, - }); - result.harnessStatePath = saveHarnessState(targetHarnessStateDir, state); - if (targetScope === "global") { - appendGlobalRefinement(globalHarnessStateDir, result); - } - let refinementAuditAppendError: { error: unknown } | undefined; - try { - this.sessionManager.appendCustomEntry("prime-agent.refinement", result); - } catch (error) { - refinementAuditAppendError = { error }; - } - try { - this._recordRefinementOutcome(result); - } catch (error) { - if (!refinementAuditAppendError) throw error; - } - if (refinementAuditAppendError) throw refinementAuditAppendError.error; - // The prompt stays byte-identical so the provider prefix cache survives; the notice carries the change. - this._recordRefinementNotice(result, source); - try { - this._emit({ type: "refine_complete", result }); - } catch { - // Listener failures must not flip a successful refinement into - // a reported failure — the refinement is already persisted. - } - try { - await this._extensionRunner.emit({ - type: "refine_complete", - id: result.id, - summary: result.summary, - appliedEdits: result.appliedEdits.filter((edit) => edit.applied).length, - scope: result.scope ?? "local", - }); - } catch { - // Extension emit failures must not flip a successful refinement - // into a reported failure — the refinement is already persisted. - } - return result; - } finally { - if (this._refineAbortController === refineAbort) { - this._refineAbortController = undefined; - } - if (!this._disposed) { - this._reconnectToAgent(); - } - } - } - - abortBranchSummary(): void { - this._branchSummaryAbortController?.abort(); - } - - /** - * Check if compaction is needed and run it. - * Called after agent_end and before prompt submission. - * - * Two cases: - * 1. Overflow: LLM returned context overflow error, remove error message from agent state, compact, auto-retry - * 2. Threshold: Context over threshold, compact, and continue only for stopped in-progress loops or queued messages - * - * @param assistantMessage The assistant message to check - * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true - */ - private _getThresholdContextTokens( - assistantMessage: AssistantMessage, - compactionTimestamp: number | undefined, - ): number | undefined { - const messages = this.agent.state.messages; - const estimate = estimateContextTokens(messages); - if (estimate.lastUsageIndex !== null) { - // Verify the usage source is post-compaction. Kept pre-compaction messages - // have stale usage reflecting the old (larger) context and would falsely - // trigger compaction right after one just finished. - const usageMsg = messages[estimate.lastUsageIndex]; - if ( - compactionTimestamp !== undefined && - usageMsg.role === "assistant" && - (usageMsg as AssistantMessage).timestamp <= compactionTimestamp - ) { - return undefined; - } - return estimate.tokens; - } - if (assistantMessage.stopReason === "error") return undefined; - return calculateContextTokens(assistantMessage.usage); - } - - private async _checkCompaction( - assistantMessage: AssistantMessage, - skipAbortedCheck = true, - queueAutonomousContinuation = true, - ): Promise { - // An abort drops any compaction the model requested this turn, even on the - // pre-prompt path (skipAbortedCheck=false) which continues to threshold checks. - if (assistantMessage.stopReason === "aborted") { - this._pendingRequestedCompaction = undefined; - // An abort also drops any pending explicit refine.run request: the - // turn that would service it (non-serialized: _consumePendingRequestedRefine - // at agent_end; serialized: the shouldStopAfterTurn checkpoint) never - // runs for an aborted turn, so a stale request would leak into the - // next turn or checkpoint. - this._pendingRequestedRefine = undefined; - if (this._serializedPlanInFlight) { - const serializedPlanInFlight = this._serializedPlanInFlight; - this._autoRefineBranchVersion++; - this._refineAbortController?.abort(); - await serializedPlanInFlight.catch(() => undefined); - if (this._serializedPlanInFlight === serializedPlanInFlight) { - this._serializedPlanInFlight = undefined; - this._serializedExplicitRefineOptions = undefined; - } - } - if (skipAbortedCheck) return false; - } - - const settings = this.settingsManager.getCompactionSettings(); - const contextWindow = this.model?.contextWindow ?? 0; - - // Skip overflow check if the message came from a different model. - // This handles the case where user switched from a smaller-context model (e.g. opus) - // to a larger-context model (e.g. codex) - the overflow error from the old model - // shouldn't trigger compaction for the new model. - const sameModel = - this.model && assistantMessage.provider === this.model.provider && assistantMessage.model === this.model.id; - - // Skip overflow/threshold checks if this assistant message is older than the - // latest compaction boundary. This prevents a stale pre-compaction usage/error - // from retriggering compaction on the first prompt after compaction. - const compactionEntry = getLatestCompactionEntry(this.sessionManager.getBranch()); - const compactionTimestamp = compactionEntry ? new Date(compactionEntry.timestamp).getTime() : undefined; - const assistantIsFromBeforeCompaction = - compactionTimestamp !== undefined && assistantMessage.timestamp <= compactionTimestamp; - - // Case 1: Overflow - takes priority over a pending model request so the error - // strip + retry still happen; the compaction it runs consumes the request. - if ( - !assistantIsFromBeforeCompaction && - (settings.enabled || this._pendingRequestedCompaction !== undefined) && - sameModel && - isContextOverflow(assistantMessage, contextWindow) - ) { - if (this._overflowRecovery !== "idle") { - if (this._overflowRecovery === "attempted") { - this._overflowRecovery = "reported"; - this._endCompactionUnsuccessfully( - "overflow", - "failed", - "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", - ); - } - return false; - } - - this._overflowRecovery = "attempted"; - // Remove the error message from agent state (it IS saved to session for history, - // but we don't want it in context for the retry) - const messages = this.agent.state.messages; - if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { - this.agent.state.messages = messages.slice(0, -1); - } - return await this._runAutoCompaction("overflow", true); - } - - if (this._pendingRequestedCompaction !== undefined) { - return await this._runAutoCompaction("requested", false); - } - - if (!settings.enabled || assistantIsFromBeforeCompaction) return false; - - // Case 3: Threshold - context is getting large. - // Use the full-session estimate so messages appended after the last successful - // assistant usage are included, matching the /usage context display. - const contextTokens = this._getThresholdContextTokens(assistantMessage, compactionTimestamp); - if (contextTokens === undefined) return false; - if (shouldCompact(contextTokens, contextWindow, settings)) { - if (queueAutonomousContinuation && this._queueGoalContinuationForThresholdCompaction(assistantMessage)) { - this._continueAfterThresholdCompaction = true; - } else if ( - queueAutonomousContinuation && - (await this._queueAutonomousContinuationForThresholdCompaction(assistantMessage)) - ) { - this._continueAfterThresholdCompaction = true; - } - return await this._runAutoCompaction("threshold", false); - } - return false; + /** + * Check if compaction is needed and run it. + * Called after agent_end and before prompt submission. + * + * Two cases: + * 1. Overflow: LLM returned context overflow error, remove error message from agent state, compact, auto-retry + * 2. Threshold: Context over threshold, compact, and continue only for stopped in-progress loops or queued messages + * + * @param assistantMessage The assistant message to check + * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true + */ + private _checkCompaction( + assistantMessage: AssistantMessage, + skipAbortedCheck = true, + queueAutonomousContinuation = true, + ): Promise { + return this._compaction.check(assistantMessage, skipAbortedCheck, queueAutonomousContinuation); } /** * Internal: Run automatic (threshold/overflow) or model-requested compaction * with events. */ - private _endCompactionUnsuccessfully( - reason: CompactionOutcomeReason, - outcome: CompactionOutcome, - message: string, - options: { - aborted?: boolean; - errorSeverity?: "warning" | "error"; - customInstructions?: string; - } = {}, - ): void { - this._persistCompactionOutcome(reason, outcome, message); - this._emit({ - type: "compaction_end", - reason, - result: undefined, - aborted: options.aborted ?? false, - willRetry: false, - // Aborts are user-initiated; they carry no error message on the event. - errorMessage: options.aborted ? undefined : message, - errorSeverity: options.errorSeverity, - customInstructions: options.customInstructions, - }); - } - - private _persistCompactionOutcome( - reason: CompactionOutcomeReason, - outcome: CompactionOutcome, - message: string, - ): void { - let outcomeMessage = createCompactionOutcomeMessage(message, { - reason, - outcome, - }); - try { - this.sessionManager.appendCustomMessageEntryWithRollback( - outcomeMessage.customType, - outcomeMessage.content, - outcomeMessage.display, - outcomeMessage.details, - ); - } catch (error) { - const persistenceError = error instanceof Error ? error.message : String(error); - outcomeMessage = createCompactionOutcomeMessage( - `${message}\n\nThis compaction outcome could not be saved to session history: ${persistenceError}`, - { reason, outcome }, - ); - // Not in the session file, so context rebuilds would drop the disclosure. - this._unpersistedOutcomes.push(outcomeMessage); - } - this.agent.state.messages.push(outcomeMessage); - this._emit({ type: "message_start", message: outcomeMessage }); - this._emit({ type: "message_end", message: outcomeMessage }); - } - - private async _runAutoCompaction( - reason: "overflow" | "threshold" | "requested", - willRetry: boolean, - ): Promise { - // Any compaction consumes a pending model request and honors its instructions - // (overflow recovery can fire first and take the request with it). - const pending = this._pendingRequestedCompaction; - this._pendingRequestedCompaction = undefined; - const customInstructions = pending?.customInstructions; - const shouldContinueAfterCompaction = - (reason === "threshold" || reason === "requested") && this._continueAfterThresholdCompaction; - const queuedAutonomousContinuationsForThisCompaction = - reason === "threshold" && shouldContinueAfterCompaction - ? this._pendingThresholdCompactionAutonomousMessages.splice(0) - : []; - const queuedGoalContinuationForThisCompaction = - reason === "threshold" && shouldContinueAfterCompaction ? this._queuedGoalThresholdContinuation : undefined; - this._continueAfterThresholdCompaction = false; - - // Requested/threshold stop the loop on purpose, so a failed or skipped compaction must not stall it. - // Overflow stays excluded: a failed overflow recovery must not re-issue the overflowing request. - const resumeAfterFailure = () => { - if ( - (reason === "requested" || reason === "threshold") && - (shouldContinueAfterCompaction || this.agent.hasQueuedMessages() || this.hasPendingSessionWork) - ) { - this._schedulePostCompactionContinue(shouldContinueAfterCompaction); - } - }; - - this._emit({ type: "compaction_start", reason, customInstructions }); - this._autoCompactionAbortController = new AbortController(); - let resolveCompactionOperation: () => void = () => {}; - const compactionOperation = new Promise((resolve) => { - resolveCompactionOperation = resolve; - }); - this._compactionOperation = compactionOperation; - - try { - const authResult = this.model ? await this._modelRegistry.getApiKeyAndHeaders(this.model) : undefined; - if (!this.model || !authResult || !authResult.ok || !authResult.apiKey) { - const detail = - !this.model || !authResult - ? "no model is selected" - : authResult.ok - ? "no API key is available" - : authResult.error; - this._endCompactionUnsuccessfully(reason, "failed", `Compaction failed: ${detail}`); - this._clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction( - reason === "threshold" && shouldContinueAfterCompaction, - queuedAutonomousContinuationsForThisCompaction, - ); - resumeAfterFailure(); - return false; - } - const result = await this._performCompaction({ - model: this.model, - apiKey: authResult.apiKey, - headers: authResult.headers, - customInstructions, - signal: this._autoCompactionAbortController.signal, - }); - - this._emit({ - type: "compaction_end", - reason, - result, - aborted: false, - willRetry, - customInstructions, - }); - // Queued work lives in both the agent queues and the session-owned queues. - const hasQueuedMessages = this.agent.hasQueuedMessages() || this.hasPendingSessionWork; - const willContinueAfterCompaction = willRetry || shouldContinueAfterCompaction || hasQueuedMessages; - - if (willRetry) { - const messages = this.agent.state.messages; - const lastMsg = messages[messages.length - 1]; - if (lastMsg?.role === "assistant" && (lastMsg as AssistantMessage).stopReason === "error") { - this.agent.state.messages = messages.slice(0, -1); - } - - this._schedulePostCompactionContinue(true); - this._scheduleAutoRefineAfterCompaction(willContinueAfterCompaction); - return true; - } else if (shouldContinueAfterCompaction || hasQueuedMessages) { - // Compaction can intentionally stop a tool loop between turns. - // Queued follow-up/steering/custom messages can also be waiting. - this._schedulePostCompactionContinue(shouldContinueAfterCompaction); - this._scheduleAutoRefineAfterCompaction(willContinueAfterCompaction); - } else { - this._scheduleAutoRefineAfterCompaction(willContinueAfterCompaction); - } - return false; - } catch (error) { - this._clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction( - reason === "threshold" && shouldContinueAfterCompaction, - queuedAutonomousContinuationsForThisCompaction, - ); - const errorMessage = error instanceof Error ? error.message : "compaction failed"; - const aborted = - errorMessage === "Compaction cancelled" || (error instanceof Error && error.name === "AbortError"); - if (aborted) { - this._clearQueuedGoalContinuationAfterCancelledThresholdCompaction(queuedGoalContinuationForThisCompaction); - this._endCompactionUnsuccessfully( - reason, - "cancelled", - `${reason === "requested" ? "Requested c" : "C"}ompaction cancelled`, - { aborted: true, customInstructions }, - ); - return false; - } - if (error instanceof CompactionSkippedError) { - this._endCompactionUnsuccessfully( - reason, - "skipped", - reason === "requested" - ? `Requested compaction skipped: ${errorMessage}` - : `Auto-compaction skipped: ${errorMessage}`, - { errorSeverity: "warning", customInstructions }, - ); - resumeAfterFailure(); - return false; - } - this._endCompactionUnsuccessfully( - reason, - "failed", - reason === "overflow" - ? `Context overflow recovery failed: ${errorMessage}` - : reason === "requested" - ? `Requested compaction failed: ${errorMessage}` - : `Auto-compaction failed: ${errorMessage}`, - { customInstructions }, - ); - resumeAfterFailure(); - return false; - } finally { - this._autoCompactionAbortController = undefined; - if (this._compactionOperation === compactionOperation) { - this._compactionOperation = undefined; - } - resolveCompactionOperation(); - this._notifySessionInputCheckpointChange(); - this._scheduleSessionInputPump(); - } + private _runAutoCompaction(reason: "overflow" | "threshold" | "requested", willRetry: boolean): Promise { + return this._compaction.runAutomatic(reason, willRetry); } setAutoCompactionEnabled(enabled: boolean): void { @@ -8656,7 +6617,7 @@ export class AgentSession { if (!this._includeCompactSkill) { skills = skills.filter((skill) => skill.name !== COMPACT_SKILL_NAME); } - if (!this._autoRefineAllowedForSession()) { + if (!this._refinement._autoRefineAllowedForSession()) { skills = skills.filter((skill) => skill.name !== REFINE_SKILL_NAME); } if (!this._agentMessageController) { @@ -8723,7 +6684,7 @@ export class AgentSession { handlers[type] = async (payload) => this.handleCompactHostRequest(type, payload); } } - if (this._autoRefineAllowedForSession()) { + if (this._refinement._autoRefineAllowedForSession()) { for (const type of ["refine.run", "refine.status"]) { handlers[type] = async (payload) => this.handleRefineHostRequest(type, payload); } @@ -8854,7 +6815,8 @@ export class AgentSession { // Keep kernel writes and host reads (system prompt, review, /refine) on // the same local harness path. Subagents prefer their own artifact dir; // ephemeral sessions fall back to the RLM session dir once it exists. - env.RLM_HARNESS_STATE_DIR = this._localHarnessStateDir() ?? getLocalHarnessStateDir(rlmSessionDir)!; + env.RLM_HARNESS_STATE_DIR = + this._refinement._localHarnessStateDir() ?? getLocalHarnessStateDir(rlmSessionDir)!; } this._addWebsearchKeyEnv(env); return env; @@ -10668,7 +8630,7 @@ export class AgentSession { // Do not switch branches while /refine has detached event handling and is // about to persist harness/session entries for the current branch. - await this._invalidatePendingAutoRefineForBranchChange(); + await this._refinement._invalidatePendingAutoRefineForBranchChange(); const { entries: entriesToSummarize, commonAncestorId } = collectEntriesForBranchSummary( this.sessionManager, diff --git a/packages/coding-agent/src/session/auto-refinement.ts b/packages/coding-agent/src/session/auto-refinement.ts new file mode 100644 index 0000000000..bebc5b5860 --- /dev/null +++ b/packages/coding-agent/src/session/auto-refinement.ts @@ -0,0 +1,365 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { AutoRefineReason, AutoRefineReview, RefinementResult } from "../core/refinement/index.js"; +import type { SettingsManager } from "../core/settings-manager.js"; +import { RefineSkippedError } from "./refinement-execution.js"; +export interface AutoRefineReviewRequest { + reason: AutoRefineReason; + turnsSinceLastReview: number; +} +export type AutoRefineReviewer = (request: AutoRefineReviewRequest, signal?: AbortSignal) => Promise; + +export function autoRefineInstructions(reason: AutoRefineReason, review: AutoRefineReview): string { + const detail = review.instructions + ? ` +Reviewer instructions: ${review.instructions}` + : ""; + return `Automatic refine review triggered by ${reason}. Only create/update/delete local harness entries if there is clear evidence that should help this session continue. Prefer an empty edits array over speculative or one-off memories. Do not promote anything global unless explicitly requested. Reviewer rationale: ${review.rationale}${detail}`; +} + +export interface AutoRefinementHost { + settingsManager: Pick; + isDisposed(): boolean; + isDisposing(): boolean; + isStreaming(): boolean; + isCompacting(): boolean; + isAllowed(): boolean; + getModel(): Model | undefined; + isContinuationScheduled(): boolean; + cancelContinuation(): void; + refine(options: { instructions?: string }, internal: { trigger: "auto" }): Promise; + runSerialized(options: { instructions?: string }, source: "auto"): Promise; + emitFailure(error: unknown): void; + review(context: AutoRefineReviewRequest, signal?: AbortSignal): Promise; +} + +/** Owns automatic review triggers, cooldowns, cancellation, and scheduled work. */ +export class AutoRefinement { + private _assistantTurnsSinceAutoRefine = 0; + private _lastAutoRefineReviewAt = 0; + private _autoRefineInProgress = false; + private readonly _autoRefineOperations = new Set>(); + private readonly _scheduledAutoRefineTimers = new Set>(); + private _compactAutoRefinePending = false; + private _turnIntervalAutoRefinePending = false; + private _pendingAutoRefineReview: { reason: AutoRefineReason; review: AutoRefineReview } | undefined; + private _autoRefineBranchVersion = 0; + private _autoRefineReviewAbort?: AbortController; + private readonly _autoRefineReviewer?: AutoRefineReviewer; + + constructor( + private readonly _host: AutoRefinementHost, + private readonly _serialized: boolean, + reviewer?: AutoRefineReviewer, + ) { + this._autoRefineReviewer = reviewer; + } + + get branchVersion(): number { + return this._autoRefineBranchVersion; + } + get turnsSinceReview(): number { + return this._assistantTurnsSinceAutoRefine; + } + get lastReviewAt(): number { + return this._lastAutoRefineReviewAt; + } + get hasPendingCompact(): boolean { + return this._compactAutoRefinePending; + } + observeAssistantEnd(): void { + this._assistantTurnsSinceAutoRefine++; + } + resetTurns(): void { + this._assistantTurnsSinceAutoRefine = 0; + } + stampCooldown(): void { + this._lastAutoRefineReviewAt = Date.now(); + } + invalidatePlans(): void { + this._autoRefineBranchVersion++; + } + abortReview(): void { + this._autoRefineReviewAbort?.abort(); + } + discardCompact(): void { + this._compactAutoRefinePending = false; + } + cancelScheduled(): void { + for (const timer of this._scheduledAutoRefineTimers) clearTimeout(timer); + this._scheduledAutoRefineTimers.clear(); + } + pendingOperations(): Promise[] { + return [...this._autoRefineOperations]; + } + + async _runSerializedAutoRefineReview(reason: "compact" | "turn_interval", branchVersion: number): Promise { + const reviewAbort = new AbortController(); + this._autoRefineReviewAbort = reviewAbort; + this._autoRefineInProgress = true; + try { + const review = await this._reviewAutoRefine( + { reason, turnsSinceLastReview: this._assistantTurnsSinceAutoRefine }, + reviewAbort.signal, + ); + if (this._host.isDisposed() || this._host.isDisposing() || branchVersion !== this._autoRefineBranchVersion) { + return; + } + if (!review.shouldRefine) { + this._lastAutoRefineReviewAt = Date.now(); + this._assistantTurnsSinceAutoRefine = 0; + return; + } + await this._host.runSerialized({ instructions: autoRefineInstructions(reason, review) }, "auto"); + if (this._host.isDisposed() || this._host.isDisposing() || branchVersion !== this._autoRefineBranchVersion) { + return; + } + this._lastAutoRefineReviewAt = Date.now(); + this._assistantTurnsSinceAutoRefine = 0; + } catch (error) { + if (branchVersion === this._autoRefineBranchVersion) { + this._lastAutoRefineReviewAt = Date.now(); + // An extension skip is an intentional non-round, not a failure. + if (error instanceof RefineSkippedError) { + this._assistantTurnsSinceAutoRefine = 0; + } else { + this._host.emitFailure(error); + } + } + } finally { + if (this._autoRefineReviewAbort === reviewAbort) { + this._autoRefineReviewAbort = undefined; + } + this._autoRefineInProgress = false; + } + } + + _discardPendingAutoRefine(options: { cancelPostCompactionContinue?: boolean } = {}): void { + this._compactAutoRefinePending = false; + this._turnIntervalAutoRefinePending = false; + this._pendingAutoRefineReview = undefined; + if (options.cancelPostCompactionContinue) { + this._host.cancelContinuation(); + } + } + + _scheduleAutoRefineAfterAgentEnd(): void { + if (!this._host.isAllowed()) { + return; + } + if (this._pendingAutoRefineReview) { + this._scheduleAutoRefine(this._pendingAutoRefineReview.reason); + return; + } + if (this._compactAutoRefinePending) { + if (this._host.isContinuationScheduled()) { + return; + } + this._scheduleAutoRefine("compact"); + return; + } + + this._scheduleAutoRefine("turn_interval"); + } + + _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void { + if (!this._host.isAllowed()) { + return; + } + if (this._serialized) { + // Serialized sessions must service compaction-triggered refinement at + // shouldStopAfterTurn (or disposal), never through the interactive path. + this._compactAutoRefinePending = true; + return; + } + if (willContinueAfterCompaction) { + this._compactAutoRefinePending = true; + return; + } + + this._scheduleAutoRefine("compact"); + } + + private _shouldSkipAutoRefineForActiveAgent(): boolean { + return this._host.isStreaming() || this._host.isCompacting(); + } + + private _scheduleDeferredAutoRefineIfIdle(): void { + if (this._autoRefineInProgress || this._shouldSkipAutoRefineForActiveAgent() || this._pendingAutoRefineReview) { + return; + } + if (this._turnIntervalAutoRefinePending) { + this._turnIntervalAutoRefinePending = false; + this._scheduleAutoRefine("turn_interval"); + } + } + + _scheduleAutoRefine(reason: AutoRefineReason, branchVersion = this._autoRefineBranchVersion): void { + const timer = setTimeout(() => { + this._scheduledAutoRefineTimers.delete(timer); + if (branchVersion !== this._autoRefineBranchVersion) { + return; + } + const operation = this._maybeAutoRefine(reason); + this._autoRefineOperations.add(operation); + void operation.finally(() => this._autoRefineOperations.delete(operation)).catch(() => undefined); + }, 0); + this._scheduledAutoRefineTimers.add(timer); + } + + async _maybeAutoRefine(reason: AutoRefineReason): Promise { + if (this._host.isDisposed() || this._host.isDisposing()) { + this._discardPendingAutoRefine(); + return; + } + if (!this._host.isAllowed()) { + this._discardPendingAutoRefine(); + return; + } + + const settings = this._host.settingsManager.getAutoRefineSettings(); + if (!settings.enabled) { + this._discardPendingAutoRefine(); + return; + } + if (this._autoRefineInProgress || this._shouldSkipAutoRefineForActiveAgent()) { + if (reason === "compact") { + this._compactAutoRefinePending = true; + } else { + this._turnIntervalAutoRefinePending = true; + } + return; + } + + const nowMs = Date.now(); + const underCooldown = + this._lastAutoRefineReviewAt > 0 && nowMs - this._lastAutoRefineReviewAt < settings.cooldownMs; + + const pendingReview = this._pendingAutoRefineReview; + if (pendingReview) { + // A failed refine stamps the cooldown; keep the pending review for later. + if (underCooldown) { + return; + } + await this._runApprovedRefine(pendingReview.reason, pendingReview.review); + return; + } + + if (reason === "compact" && !settings.compact) { + this._compactAutoRefinePending = false; + reason = "turn_interval"; + } + if (reason === "turn_interval" && this._assistantTurnsSinceAutoRefine < settings.turnInterval) { + return; + } + if (underCooldown) { + if (reason === "compact") { + this._compactAutoRefinePending = true; + } else { + this._turnIntervalAutoRefinePending = true; + } + return; + } + if (reason === "turn_interval") { + this._turnIntervalAutoRefinePending = false; + } + if (!this._host.getModel()) { + if (reason === "compact") { + this._compactAutoRefinePending = true; + } + return; + } + this._autoRefineInProgress = true; + const turnsSinceLastReview = this._assistantTurnsSinceAutoRefine; + const branchVersion = this._autoRefineBranchVersion; + const reviewAbort = new AbortController(); + this._autoRefineReviewAbort = reviewAbort; + let approvedReview: AutoRefineReview | undefined; + try { + const review = await this._reviewAutoRefine({ reason, turnsSinceLastReview }, reviewAbort.signal); + if (this._host.isDisposed() || this._host.isDisposing() || branchVersion !== this._autoRefineBranchVersion) { + return; + } + if (!review.shouldRefine) { + const preserveTurnIntervalReview = + reason === "compact" && this._assistantTurnsSinceAutoRefine >= settings.turnInterval; + if (preserveTurnIntervalReview) { + this._turnIntervalAutoRefinePending = true; + } else { + this._lastAutoRefineReviewAt = nowMs; + this._assistantTurnsSinceAutoRefine = 0; + } + if (reason === "compact") { + this._compactAutoRefinePending = false; + } + return; + } + if (this._shouldSkipAutoRefineForActiveAgent()) { + this._pendingAutoRefineReview = { reason, review }; + return; + } + approvedReview = review; + } catch { + // Failed review: stamp the cooldown so a persistent failure (bad auth, + // unparseable output) doesn't retry a full review on every agent end. + if (branchVersion === this._autoRefineBranchVersion) { + this._lastAutoRefineReviewAt = Date.now(); + } + } finally { + if (this._autoRefineReviewAbort === reviewAbort) { + this._autoRefineReviewAbort = undefined; + } + this._autoRefineInProgress = false; + // When a refine follows, _runApprovedRefine schedules the deferred pass. + if (!approvedReview) { + this._scheduleDeferredAutoRefineIfIdle(); + } + } + if (approvedReview) { + await this._runApprovedRefine(reason, approvedReview); + } + } + + private async _runApprovedRefine(reason: AutoRefineReason, review: AutoRefineReview): Promise { + this._autoRefineInProgress = true; + try { + await this._host.refine({ instructions: autoRefineInstructions(reason, review) }, { trigger: "auto" }); + this._pendingAutoRefineReview = undefined; + this._turnIntervalAutoRefinePending = false; + this._lastAutoRefineReviewAt = Date.now(); + this._assistantTurnsSinceAutoRefine = 0; + if (reason === "compact") { + this._compactAutoRefinePending = false; + } + } catch (error) { + // Auto-refine is opportunistic; manual /refine remains available. + // Stamp the cooldown so a persistently failing refine doesn't retry + // (via a retained pending review) on every agent end. + this._lastAutoRefineReviewAt = Date.now(); + if (error instanceof RefineSkippedError) { + // A skipped round is consumed like a reviewer decline, not retained for retry. + this._pendingAutoRefineReview = undefined; + this._turnIntervalAutoRefinePending = false; + this._assistantTurnsSinceAutoRefine = 0; + if (reason === "compact") this._compactAutoRefinePending = false; + } + } finally { + this._autoRefineInProgress = false; + this._scheduleDeferredAutoRefineIfIdle(); + } + } + + _reviewAutoRefine(context: AutoRefineReviewRequest, signal?: AbortSignal): Promise { + if (this._autoRefineReviewer) { + return this._reviewWithCustomReviewer(this._autoRefineReviewer, context, signal); + } + return this._host.review(context, signal); + } + + private async _reviewWithCustomReviewer( + reviewer: AutoRefineReviewer, + context: AutoRefineReviewRequest, + signal?: AbortSignal, + ): Promise { + return reviewer.call(this, context, signal); + } +} diff --git a/packages/coding-agent/src/session/compaction-execution.ts b/packages/coding-agent/src/session/compaction-execution.ts new file mode 100644 index 0000000000..db7e6c7e71 --- /dev/null +++ b/packages/coding-agent/src/session/compaction-execution.ts @@ -0,0 +1,187 @@ +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { Api, Model } from "@earendil-works/pi-ai"; +import { + type CompactionResult, + type CompactionSettings, + compact, + prepareCompaction, +} from "../core/compaction/index.js"; +import type { ExtensionRunner, SessionBeforeCompactResult } from "../core/extensions/index.js"; +import type { ProviderRetryPolicy } from "../core/provider-retry.js"; +import { modelRequestHeaders, type SemanticEdgeRecorder } from "../core/semantic-edges.js"; +import type { CompactionEntry, SessionManager } from "../core/session-manager.js"; + +export class CompactionSkippedError extends Error {} + +export interface CompactionExecutionOptions { + model: Model; + apiKey: string; + headers?: Record; + customInstructions?: string; + signal: AbortSignal; +} + +export interface CompactionExecutionHost { + getSessionStore(): Pick; + getSettings(): CompactionSettings; + getSemanticEdges(): Pick< + SemanticEdgeRecorder, + "beginCompaction" | "startCompactionRequest" | "failRequest" | "finishRequest" | "finishCompaction" + >; + getExtensions(): Pick; + getThinkingLevel(): ThinkingLevel; + getRetryPolicy(): ProviderRetryPolicy; + getHarnessDigest(): string; + rebuildContext(): void; + syncKernelState(): Promise; + reapDeletedChildren(): Promise; +} + +export async function performSessionCompaction( + host: CompactionExecutionHost, + options: CompactionExecutionOptions, +): Promise { + const { model, apiKey, headers, customInstructions, signal } = options; + const pathEntries = host.getSessionStore().getBranch(); + const settings = host.getSettings(); + + const preparation = prepareCompaction(pathEntries, settings); + if (!preparation) { + const lastEntry = pathEntries[pathEntries.length - 1]; + if (lastEntry?.type === "compaction") { + throw new CompactionSkippedError("Already compacted"); + } + throw new CompactionSkippedError("Session is too short to compact — try again once it grows"); + } + + let extensionCompaction: CompactionResult | undefined; + let fromExtension = false; + + const semanticCompaction = host.getSemanticEdges().beginCompaction(); + let compactionRecorded = false; + const uncommittedSlices: string[] = []; + let compactionSettled = false; + let summary: string; + let firstKeptEntryId: string; + let tokensBefore: number; + let details: CompactionResult["details"]; + let usage: CompactionResult["usage"]; + try { + if (host.getExtensions().hasHandlers("session_before_compact")) { + const result = (await host.getExtensions().emit({ + type: "session_before_compact", + preparation, + branchEntries: pathEntries, + customInstructions, + signal, + })) as SessionBeforeCompactResult | undefined; + + if (result?.cancel) { + throw new Error("Compaction cancelled"); + } + + if (result?.compaction) { + extensionCompaction = result.compaction; + fromExtension = true; + } + } + + if (extensionCompaction) { + ({ summary, firstKeptEntryId, tokensBefore, details, usage } = extensionCompaction); + } else { + // Each summary wire call gets its own request ID: split turns send two + // different bodies, and one Idempotency-Key must never cover both. A slice + // that succeeds on the wire stays uncommitted until the compaction itself + // commits: a racing sibling's failure (or an abort) must leave no committed + // summary request for the next turn's continuation edge to attach to. + const summaryCall = async ( + call: (callHeaders: Record | undefined) => Promise, + ): Promise => { + const requestId = host.getSemanticEdges().startCompactionRequest(semanticCompaction.compactionId); + if (requestId === undefined) { + return call(headers); + } + try { + const result = await call({ ...headers, ...modelRequestHeaders(requestId) }); + // A slice resolving after a sibling's rejection already settled the + // compaction would push into a drained list and stay in-flight forever. + if (compactionSettled) { + host.getSemanticEdges().failRequest(requestId); + } else { + uncommittedSlices.push(requestId); + } + return result; + } catch (error) { + host.getSemanticEdges().failRequest(requestId); + throw error; + } + }; + ({ summary, firstKeptEntryId, tokensBefore, details, usage } = await compact( + preparation, + model, + apiKey, + headers, + customInstructions, + signal, + host.getThinkingLevel(), + summaryCall, + host.getRetryPolicy(), + )); + } + + if (signal.aborted) { + throw new Error("Compaction cancelled"); + } + + // Ledger-before-effect: the compaction outcome is durable before the transcript + // commits it. Marked first: the ID is consumed even when the write throws, and a + // second finish attempt would mask the original I/O error. + compactionRecorded = true; + compactionSettled = true; + for (const requestId of uncommittedSlices.splice(0)) { + host.getSemanticEdges().finishRequest(requestId); + } + host.getSemanticEdges().finishCompaction(semanticCompaction.compactionId, "completed"); + // Attached mechanically; the digest never flows through the summarizer LLM. + host + .getSessionStore() + .appendCompaction( + summary, + firstKeptEntryId, + tokensBefore, + details, + fromExtension, + customInstructions, + usage, + host.getHarnessDigest(), + ); + } catch (error) { + compactionSettled = true; + for (const requestId of uncommittedSlices.splice(0)) { + host.getSemanticEdges().failRequest(requestId); + } + if (!compactionRecorded) { + const cancelled = + error instanceof Error && (error.name === "AbortError" || error.message === "Compaction cancelled"); + host.getSemanticEdges().finishCompaction(semanticCompaction.compactionId, cancelled ? "cancelled" : "failed"); + } + throw error; + } + const newEntries = host.getSessionStore().getEntries(); + host.rebuildContext(); + + const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as + | CompactionEntry + | undefined; + if (savedCompactionEntry) { + await host.getExtensions().emit({ + type: "session_compact", + compactionEntry: savedCompactionEntry, + fromExtension, + }); + } + await host.syncKernelState(); + await host.reapDeletedChildren(); + + return { summary, firstKeptEntryId, tokensBefore, details }; +} diff --git a/packages/coding-agent/src/session/compaction.ts b/packages/coding-agent/src/session/compaction.ts new file mode 100644 index 0000000000..d2ebebc5fe --- /dev/null +++ b/packages/coding-agent/src/session/compaction.ts @@ -0,0 +1,510 @@ +import type { AgentEvent, AgentMessage } from "@earendil-works/pi-agent-core"; +import { type Api, type AssistantMessage, isContextOverflow, type Model } from "@earendil-works/pi-ai"; +import { formatNoModelSelectedMessage } from "../core/auth-guidance.js"; +import { + type CompactionResult, + type CompactionSettings, + calculateContextTokens, + estimateContextTokens, + shouldCompact, +} from "../core/compaction/index.js"; +import { + type CompactionOutcome, + type CompactionOutcomeReason, + type CustomMessage, + createCompactionOutcomeMessage, +} from "../core/messages.js"; +import type { ModelRegistry } from "../core/model-registry.js"; +import { getLatestCompactionEntry, type SessionManager } from "../core/session-manager.js"; +import { type CompactionExecutionOptions, CompactionSkippedError } from "./compaction-execution.js"; + +export type CompactionReason = "manual" | "threshold" | "overflow" | "requested"; +export type SessionCompactionEvent = + | { type: "compaction_start"; reason: CompactionReason; customInstructions?: string } + | { + type: "compaction_end"; + reason: CompactionReason; + result: CompactionResult | undefined; + aborted: boolean; + willRetry: boolean; + errorMessage?: string; + errorSeverity?: "warning" | "error"; + customInstructions?: string; + }; + +export interface SessionCompactionHost { + getSettings(): CompactionSettings; + runAutomatic(reason: "overflow" | "threshold" | "requested", willRetry: boolean): Promise; + queueGoalContinuation(message: AssistantMessage): boolean; + queueAutonomousContinuation(message: AssistantMessage): Promise; + beginRefinementAbort(): { promise: Promise; finish(): void } | undefined; + getModel(): Model | undefined; + isStreaming(): boolean; + getRequiredAuth(model: Model): Promise<{ apiKey: string; headers?: Record }>; + getAuth(model: Model): ReturnType; + perform(options: CompactionExecutionOptions): Promise; + disconnect(): void; + reconnect(): void; + abortSession(): Promise; + getContinuationState(): { scheduled: boolean; continueAfterSessionInput: boolean }; + afterManualCompaction(signal: AbortSignal, wasScheduled: boolean, continueAfterSessionInput: boolean): void; + getMessages(): AgentMessage[]; + replaceMessages(messages: AgentMessage[]): void; + hasAgentQueuedMessages(): boolean; + hasPendingSessionWork(): boolean; + scheduleContinuation(continueAfterSessionInput?: boolean): void; + scheduleRefinement(willContinue: boolean): void; + takeThresholdAutonomousMessages(): AgentMessage[]; + getThresholdGoalContinuation(): AgentMessage | undefined; + clearAutonomousContinuations(shouldContinue: boolean, messages: AgentMessage[]): void; + clearGoalContinuation(message: AgentMessage | undefined): void; + getSessionStore(): Pick; + retainUnpersistedOutcome(message: CustomMessage): void; + emit(event: SessionCompactionEvent | Extract): void; + notifyCheckpoints(): void; + scheduleInput(): void; +} + +export class SessionCompaction { + private manualAbort: AbortController | undefined; + private automaticAbort: AbortController | undefined; + private activeOperation: Promise | undefined; + private overflowStage: "idle" | "attempted" | "reported" = "idle"; + private pendingRequest: { customInstructions?: string } | undefined; + private continueAfterThreshold = false; + + constructor(private readonly host: SessionCompactionHost) {} + + private get model(): Model | undefined { + return this.host.getModel(); + } + get operation(): Promise | undefined { + return this.activeOperation; + } + get isRunning(): boolean { + return this.automaticAbort !== undefined || this.manualAbort !== undefined; + } + get hasPendingRequest(): boolean { + return this.pendingRequest !== undefined; + } + get overflowRecovery(): "idle" | "attempted" | "reported" { + return this.overflowStage; + } + + request(customInstructions?: string): void { + this.pendingRequest = { customInstructions }; + } + clearRequest(): void { + this.pendingRequest = undefined; + } + requestContinuation(): void { + this.continueAfterThreshold = true; + } + resetContinuation(): void { + this.continueAfterThreshold = false; + } + resetOverflowRecovery(): void { + this.overflowStage = "idle"; + } + markOverflowAttempted(): void { + this.overflowStage = "attempted"; + } + markOverflowReported(): void { + this.overflowStage = "reported"; + } + abort(): void { + this.manualAbort?.abort(); + this.automaticAbort?.abort(); + } + abortAutomatic(): void { + this.automaticAbort?.abort(); + } + + getThresholdContextTokens( + assistantMessage: AssistantMessage, + compactionTimestamp: number | undefined, + ): number | undefined { + const messages = this.host.getMessages(); + const estimate = estimateContextTokens(messages); + if (estimate.lastUsageIndex !== null) { + // Verify the usage source is post-compaction. Kept pre-compaction messages + // have stale usage reflecting the old (larger) context and would falsely + // trigger compaction right after one just finished. + const usageMsg = messages[estimate.lastUsageIndex]; + if ( + compactionTimestamp !== undefined && + usageMsg.role === "assistant" && + (usageMsg as AssistantMessage).timestamp <= compactionTimestamp + ) { + return undefined; + } + return estimate.tokens; + } + if (assistantMessage.stopReason === "error") return undefined; + return calculateContextTokens(assistantMessage.usage); + } + + async check( + assistantMessage: AssistantMessage, + skipAbortedCheck = true, + queueAutonomousContinuation = true, + ): Promise { + // An abort drops any compaction the model requested this turn, even on the + // pre-prompt path (skipAbortedCheck=false) which continues to threshold checks. + if (assistantMessage.stopReason === "aborted") { + this.clearRequest(); + const refinementAbort = this.host.beginRefinementAbort(); + if (refinementAbort) { + await refinementAbort.promise.catch(() => undefined); + refinementAbort.finish(); + } + if (skipAbortedCheck) return false; + } + + const settings = this.host.getSettings(); + const contextWindow = this.model?.contextWindow ?? 0; + + // Skip overflow check if the message came from a different model. + // This handles the case where user switched from a smaller-context model (e.g. opus) + // to a larger-context model (e.g. codex) - the overflow error from the old model + // shouldn't trigger compaction for the new model. + const sameModel = + this.model && assistantMessage.provider === this.model.provider && assistantMessage.model === this.model.id; + + // Skip overflow/threshold checks if this assistant message is older than the + // latest compaction boundary. This prevents a stale pre-compaction usage/error + // from retriggering compaction on the first prompt after compaction. + const compactionEntry = getLatestCompactionEntry(this.host.getSessionStore().getBranch()); + const compactionTimestamp = compactionEntry ? new Date(compactionEntry.timestamp).getTime() : undefined; + const assistantIsFromBeforeCompaction = + compactionTimestamp !== undefined && assistantMessage.timestamp <= compactionTimestamp; + + // Case 1: Overflow - takes priority over a pending model request so the error + // strip + retry still happen; the compaction it runs consumes the request. + if ( + !assistantIsFromBeforeCompaction && + (settings.enabled || this.hasPendingRequest) && + sameModel && + isContextOverflow(assistantMessage, contextWindow) + ) { + if (this.overflowRecovery !== "idle") { + if (this.overflowRecovery === "attempted") { + this.markOverflowReported(); + this.endUnsuccessfully( + "overflow", + "failed", + "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", + ); + } + return false; + } + + this.markOverflowAttempted(); + // Remove the error message from agent state (it IS saved to session for history, + // but we don't want it in context for the retry) + const messages = this.host.getMessages(); + if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { + this.host.replaceMessages(messages.slice(0, -1)); + } + return await this.host.runAutomatic("overflow", true); + } + + if (this.hasPendingRequest) { + return await this.host.runAutomatic("requested", false); + } + + if (!settings.enabled || assistantIsFromBeforeCompaction) return false; + + // Case 3: Threshold - context is getting large. + // Use the full-session estimate so messages appended after the last successful + // assistant usage are included, matching the /usage context display. + const contextTokens = this.getThresholdContextTokens(assistantMessage, compactionTimestamp); + if (contextTokens === undefined) return false; + if (shouldCompact(contextTokens, contextWindow, settings)) { + if (queueAutonomousContinuation && this.host.queueGoalContinuation(assistantMessage)) { + this.requestContinuation(); + } else if (queueAutonomousContinuation && (await this.host.queueAutonomousContinuation(assistantMessage))) { + this.requestContinuation(); + } + return await this.host.runAutomatic("threshold", false); + } + return false; + } + + async compact(customInstructions?: string, options: { skipAbort?: boolean } = {}): Promise { + if (options.skipAbort && this.host.isStreaming()) { + throw new Error("Cannot compact without aborting while the agent is running."); + } + const { scheduled: hadPostCompactionContinue, continueAfterSessionInput } = this.host.getContinuationState(); + this.host.disconnect(); + if (!options.skipAbort) await this.host.abortSession(); + let didCompact = false; + const compactionAbort = new AbortController(); + this.manualAbort = compactionAbort; + let resolveCompactionOperation: () => void = () => {}; + const compactionOperation = new Promise((resolve) => { + resolveCompactionOperation = resolve; + }); + this.activeOperation = compactionOperation; + this.host.emit({ + type: "compaction_start", + reason: "manual", + customInstructions, + }); + + try { + if (!this.model) { + throw new Error(formatNoModelSelectedMessage()); + } + + const { apiKey, headers } = await this.host.getRequiredAuth(this.model); + const result = await this.host.perform({ + model: this.model, + apiKey, + headers, + customInstructions, + signal: compactionAbort.signal, + }); + + this.host.emit({ + type: "compaction_end", + reason: "manual", + result, + aborted: false, + willRetry: false, + customInstructions, + }); + didCompact = true; + // A manual compaction satisfies any pending model request; on failure the + // request stays scheduled for the next turn boundary. + this.pendingRequest = undefined; + return result; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + const aborted = message === "Compaction cancelled" || (error instanceof Error && error.name === "AbortError"); + const skipped = error instanceof CompactionSkippedError; + this.host.emit({ + type: "compaction_end", + reason: "manual", + result: undefined, + aborted, + willRetry: false, + errorMessage: aborted ? undefined : skipped ? message : `Compaction failed: ${message}`, + errorSeverity: skipped ? "warning" : "error", + customInstructions, + }); + throw error; + } finally { + this.manualAbort = undefined; + this.host.reconnect(); + if (this.activeOperation === compactionOperation) { + this.activeOperation = undefined; + } + resolveCompactionOperation(); + this.host.notifyCheckpoints(); + this.host.scheduleInput(); + if (didCompact) { + this.host.afterManualCompaction( + compactionAbort.signal, + hadPostCompactionContinue, + continueAfterSessionInput, + ); + } + } + } + + endUnsuccessfully( + reason: CompactionOutcomeReason, + outcome: CompactionOutcome, + message: string, + options: { + aborted?: boolean; + errorSeverity?: "warning" | "error"; + customInstructions?: string; + } = {}, + ): void { + this.persistOutcome(reason, outcome, message); + this.host.emit({ + type: "compaction_end", + reason, + result: undefined, + aborted: options.aborted ?? false, + willRetry: false, + // Aborts are user-initiated; they carry no error message on the event. + errorMessage: options.aborted ? undefined : message, + errorSeverity: options.errorSeverity, + customInstructions: options.customInstructions, + }); + } + + private persistOutcome(reason: CompactionOutcomeReason, outcome: CompactionOutcome, message: string): void { + let outcomeMessage = createCompactionOutcomeMessage(message, { + reason, + outcome, + }); + try { + this.host + .getSessionStore() + .appendCustomMessageEntryWithRollback( + outcomeMessage.customType, + outcomeMessage.content, + outcomeMessage.display, + outcomeMessage.details, + ); + } catch (error) { + const persistenceError = error instanceof Error ? error.message : String(error); + outcomeMessage = createCompactionOutcomeMessage( + `${message}\n\nThis compaction outcome could not be saved to session history: ${persistenceError}`, + { reason, outcome }, + ); + // Not in the session file, so context rebuilds would drop the disclosure. + this.host.retainUnpersistedOutcome(outcomeMessage); + } + this.host.getMessages().push(outcomeMessage); + this.host.emit({ type: "message_start", message: outcomeMessage }); + this.host.emit({ type: "message_end", message: outcomeMessage }); + } + + async runAutomatic(reason: "overflow" | "threshold" | "requested", willRetry: boolean): Promise { + // Any compaction consumes a pending model request and honors its instructions + // (overflow recovery can fire first and take the request with it). + const pending = this.pendingRequest; + this.pendingRequest = undefined; + const customInstructions = pending?.customInstructions; + const shouldContinueAfterCompaction = + (reason === "threshold" || reason === "requested") && this.continueAfterThreshold; + const queuedAutonomousContinuationsForThisCompaction = + reason === "threshold" && shouldContinueAfterCompaction ? this.host.takeThresholdAutonomousMessages() : []; + const queuedGoalContinuationForThisCompaction = + reason === "threshold" && shouldContinueAfterCompaction ? this.host.getThresholdGoalContinuation() : undefined; + this.continueAfterThreshold = false; + + // Requested/threshold stop the loop on purpose, so a failed or skipped compaction must not stall it. + // Overflow stays excluded: a failed overflow recovery must not re-issue the overflowing request. + const resumeAfterFailure = () => { + if ( + (reason === "requested" || reason === "threshold") && + (shouldContinueAfterCompaction || this.host.hasAgentQueuedMessages() || this.host.hasPendingSessionWork()) + ) { + this.host.scheduleContinuation(shouldContinueAfterCompaction); + } + }; + + this.host.emit({ type: "compaction_start", reason, customInstructions }); + this.automaticAbort = new AbortController(); + let resolveCompactionOperation: () => void = () => {}; + const compactionOperation = new Promise((resolve) => { + resolveCompactionOperation = resolve; + }); + this.activeOperation = compactionOperation; + + try { + const authResult = this.model ? await this.host.getAuth(this.model) : undefined; + if (!this.model || !authResult || !authResult.ok || !authResult.apiKey) { + const detail = + !this.model || !authResult + ? "no model is selected" + : authResult.ok + ? "no API key is available" + : authResult.error; + this.endUnsuccessfully(reason, "failed", `Compaction failed: ${detail}`); + this.host.clearAutonomousContinuations( + reason === "threshold" && shouldContinueAfterCompaction, + queuedAutonomousContinuationsForThisCompaction, + ); + resumeAfterFailure(); + return false; + } + + const result = await this.host.perform({ + model: this.model, + apiKey: authResult.apiKey, + headers: authResult.headers, + customInstructions, + signal: this.automaticAbort.signal, + }); + + this.host.emit({ + type: "compaction_end", + reason, + result, + aborted: false, + willRetry, + customInstructions, + }); + // Queued work lives in both the agent queues and the session-owned queues. + const hasQueuedMessages = this.host.hasAgentQueuedMessages() || this.host.hasPendingSessionWork(); + const willContinueAfterCompaction = willRetry || shouldContinueAfterCompaction || hasQueuedMessages; + + if (willRetry) { + const messages = this.host.getMessages(); + const lastMsg = messages[messages.length - 1]; + if (lastMsg?.role === "assistant" && (lastMsg as AssistantMessage).stopReason === "error") { + this.host.replaceMessages(messages.slice(0, -1)); + } + + this.host.scheduleContinuation(true); + this.host.scheduleRefinement(willContinueAfterCompaction); + return true; + } else if (shouldContinueAfterCompaction || hasQueuedMessages) { + // Compaction can intentionally stop a tool loop between turns. + // Queued follow-up/steering/custom messages can also be waiting. + this.host.scheduleContinuation(shouldContinueAfterCompaction); + this.host.scheduleRefinement(willContinueAfterCompaction); + } else { + this.host.scheduleRefinement(willContinueAfterCompaction); + } + return false; + } catch (error) { + this.host.clearAutonomousContinuations( + reason === "threshold" && shouldContinueAfterCompaction, + queuedAutonomousContinuationsForThisCompaction, + ); + const errorMessage = error instanceof Error ? error.message : "compaction failed"; + const aborted = + errorMessage === "Compaction cancelled" || (error instanceof Error && error.name === "AbortError"); + if (aborted) { + this.host.clearGoalContinuation(queuedGoalContinuationForThisCompaction); + this.endUnsuccessfully( + reason, + "cancelled", + `${reason === "requested" ? "Requested c" : "C"}ompaction cancelled`, + { aborted: true, customInstructions }, + ); + return false; + } + if (error instanceof CompactionSkippedError) { + this.endUnsuccessfully( + reason, + "skipped", + reason === "requested" + ? `Requested compaction skipped: ${errorMessage}` + : `Auto-compaction skipped: ${errorMessage}`, + { errorSeverity: "warning", customInstructions }, + ); + resumeAfterFailure(); + return false; + } + this.endUnsuccessfully( + reason, + "failed", + reason === "overflow" + ? `Context overflow recovery failed: ${errorMessage}` + : reason === "requested" + ? `Requested compaction failed: ${errorMessage}` + : `Auto-compaction failed: ${errorMessage}`, + { customInstructions }, + ); + resumeAfterFailure(); + return false; + } finally { + this.automaticAbort = undefined; + if (this.activeOperation === compactionOperation) { + this.activeOperation = undefined; + } + resolveCompactionOperation(); + this.host.notifyCheckpoints(); + this.host.scheduleInput(); + } + } +} diff --git a/packages/coding-agent/src/session/continuation.ts b/packages/coding-agent/src/session/continuation.ts new file mode 100644 index 0000000000..a68e2e0d6a --- /dev/null +++ b/packages/coding-agent/src/session/continuation.ts @@ -0,0 +1,234 @@ +import { AgentContinueError, type AgentMessage } from "@earendil-works/pi-agent-core"; +import type { SessionCommitLease } from "./commit-fence.js"; + +export interface ContinuationToken { + readonly promise: Promise; + readonly continueAfterSessionInput: boolean; +} + +interface ContinuationSettlement extends ContinuationToken { + continueAfterSessionInput: boolean; + settled: boolean; + resolve(): void; + reject(error: Error): void; +} + +function createSettlement(): ContinuationSettlement { + let resolve!: () => void; + let reject!: (error: Error) => void; + const promise = new Promise((done, fail) => { + resolve = done; + reject = fail; + }); + promise.catch(() => undefined); + return { promise, resolve, reject, continueAfterSessionInput: false, settled: false }; +} + +export interface SessionContinuationHost { + waitForAgentIdle(): Promise; + waitForRetry(): Promise; + waitForRefinement(): Promise; + queuedWorkPauseCount(): number; + addCheckpointWaiter(waiter: () => void): void; + removeCheckpointWaiter(waiter: () => void): void; + notifyCheckpoints(): void; + compactionOperation(): Promise | undefined; + isRefinementApplying(): boolean; + acquireCommitFence(): Promise; + scheduleRefinement(): void; + unfinishedActionCount(): number; + isInputRequested(): boolean; + scheduleInput(): void; + continue(): Promise; + waitForIdleOrSettlement(token: ContinuationToken): Promise; + removeQueuedMessages(predicate: (message: AgentMessage) => boolean): AgentMessage[]; + followUp(message: AgentMessage): void; + onMessageConsumed(message: AgentMessage): void; +} + +/** Owns continuation settlement and message identity across cancellation and replacement runs. */ +export class SessionContinuation { + private scheduled = false; + private settlement: ContinuationSettlement | undefined; + private trackedMessages: AgentMessage[] = []; + private scheduledMessages: AgentMessage[] = []; + + constructor(private readonly host: SessionContinuationHost) {} + + get isScheduled(): boolean { + return this.scheduled; + } + get current(): ContinuationToken | undefined { + return this.settlement; + } + get messages(): readonly AgentMessage[] { + return this.trackedMessages; + } + track(message: AgentMessage): void { + this.trackedMessages.push(message); + } + remove(messages: ReadonlySet): void { + this.trackedMessages = this.trackedMessages.filter((message) => !messages.has(message)); + } + + private settle(error?: Error): void { + if (!error && this.scheduled) return; + const settlement = this.settlement; + if (!settlement || settlement.settled) return; + settlement.settled = true; + this.settlement = undefined; + if (error) settlement.reject(error); + else settlement.resolve(); + this.host.notifyCheckpoints(); + } + + cancel(): void { + this.scheduled = false; + this.scheduledMessages = []; + this.settle(); + } + + schedule(continueAfterSessionInput = false): void { + if (!this.settlement || this.settlement.settled) { + this.settlement = createSettlement(); + } + const settlement = this.settlement; + settlement.continueAfterSessionInput ||= continueAfterSessionInput; + if (this.scheduled) { + return; + } + this.scheduled = true; + this.scheduledMessages = [...this.trackedMessages]; + void this.runScheduled(settlement) + .catch(() => undefined) + .finally(() => { + if (this.settlement === settlement) { + this.settle(); + } + }); + } + + private ownsScheduledMessages(continuationMessages: AgentMessage[]): boolean { + return continuationMessages.some((message) => this.trackedMessages.includes(message)); + } + + private async waitForQueuedWorkResume(settlement: ContinuationSettlement): Promise { + while (this.host.queuedWorkPauseCount() > 0 && this.settlement === settlement) { + let resume = () => {}; + const resumed = new Promise((resolve) => { + resume = resolve; + this.host.addCheckpointWaiter(resolve); + }); + try { + await Promise.race([resumed, settlement.promise]); + } finally { + this.host.removeCheckpointWaiter(resume); + } + } + } + + private async runScheduled(settlement: ContinuationSettlement): Promise { + while (this.scheduled && this.settlement === settlement) { + await this.host.waitForAgentIdle(); + await this.host.waitForRetry(); + await this.host.waitForRefinement(); + await this.waitForQueuedWorkResume(settlement); + const compactionOperation = this.host.compactionOperation(); + if (compactionOperation) { + await Promise.race([compactionOperation, settlement.promise]); + continue; + } + + const commitFence = await this.host.acquireCommitFence(); + let continuation: Promise | undefined; + let continuationMessages: AgentMessage[] = []; + let waitForSessionInput = false; + try { + await this.host.waitForAgentIdle(); + if (!this.scheduled || this.settlement !== settlement) { + return; + } + + if ( + this.host.queuedWorkPauseCount() > 0 || + this.host.compactionOperation() || + this.host.isRefinementApplying() + ) { + continue; + } + + continuationMessages = [...this.scheduledMessages]; + if (continuationMessages.length > 0 && !this.ownsScheduledMessages(continuationMessages)) { + this.cancel(); + this.host.scheduleRefinement(); + return; + } + if (this.host.unfinishedActionCount() > 0 || this.host.isInputRequested()) { + this.host.scheduleInput(); + waitForSessionInput = true; + } else { + this.scheduled = false; + continuation = this.host.continue(); + } + } finally { + commitFence.release(); + } + + if (waitForSessionInput) { + await this.host.waitForIdleOrSettlement(settlement); + if (this.settlement !== settlement) return; + const shouldContinue = + (settlement.continueAfterSessionInput && continuationMessages.length === 0) || + this.ownsScheduledMessages(continuationMessages); + if (shouldContinue) { + this.scheduledMessages = [...this.trackedMessages]; + continue; + } + this.scheduled = false; + this.scheduledMessages = []; + this.host.scheduleRefinement(); + return; + } + + try { + await continuation; + if (this.settlement === settlement) { + this.forgetConsumed(continuationMessages); + } + return; + } catch (error) { + const code = error instanceof AgentContinueError ? error.code : undefined; + if (code === "busy") { + if (this.settlement === settlement) { + this.scheduled = true; + this.scheduledMessages = [...this.trackedMessages]; + } + continue; + } + if (code !== "nothing-to-continue" && this.settlement === settlement) { + this.settle(error instanceof Error ? error : new Error(String(error))); + } + return; + } + } + } + + forgetConsumed(continuationMessages: AgentMessage[]): void { + if (continuationMessages.length === 0) { + return; + } + const continuationMessageSet = new Set(continuationMessages); + const stillQueued = new Set(this.host.removeQueuedMessages((message) => continuationMessageSet.has(message))); + for (const message of stillQueued) { + this.host.followUp(message); + } + for (const message of continuationMessages) { + if (!stillQueued.has(message)) { + this.host.onMessageConsumed(message); + } + } + this.trackedMessages = this.trackedMessages.filter( + (message) => !continuationMessageSet.has(message) || stillQueued.has(message), + ); + } +} diff --git a/packages/coding-agent/src/session/refinement-execution.ts b/packages/coding-agent/src/session/refinement-execution.ts new file mode 100644 index 0000000000..3889cec492 --- /dev/null +++ b/packages/coding-agent/src/session/refinement-execution.ts @@ -0,0 +1,336 @@ +import { existsSync } from "node:fs"; +import { dirname, resolve } from "node:path"; +import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { Api, Model } from "@earendil-works/pi-ai"; +import { formatNoModelSelectedMessage } from "../core/auth-guidance.js"; +import { serializeConversation } from "../core/compaction/index.js"; +import type { ExtensionRunner, SessionBeforeRefineResult } from "../core/extensions/index.js"; +import { + type CustomMessage, + convertToLlm, + createRefinementNoticeMessage, + createRefinementOutcomeMessage, + type RefinementSource, +} from "../core/messages.js"; +import type { ProviderRetryPolicy } from "../core/provider-retry.js"; +import { + type AutoRefineReview, + appendGlobalRefinement, + applyRefinementProposal, + generateRefinementId, + getGlobalHarnessStateDir, + getLocalHarnessStateDir, + getRefinementHistory, + type HarnessState, + inferRefinementResultScope, + loadGlobalRefinementHistory, + loadHarnessState, + mergeHarnessStates, + mergeRefinementHistory, + normalizeRefinementProposal, + planRefinement, + type RefinementPlan, + type RefinementResult, + reviewAutoRefine, + saveHarnessState, +} from "../core/refinement/index.js"; +import type { SessionManager } from "../core/session-manager.js"; +import type { AutoRefineReviewRequest } from "./auto-refinement.js"; +export type SessionRefinementEvent = + | { type: "refine_complete"; result: RefinementResult } + | { type: "refine_failed"; error: string } + | { type: "message_start" | "message_end"; message: CustomMessage }; +export interface RefinementExecutionHost { + sessionManager: Pick< + SessionManager, + "getSessionArtifactDir" | "getEntries" | "appendCustomMessageEntryWithRollback" | "appendCustomEntry" + >; + isDisposed(): boolean; + getRlmSessionDir(): string | undefined; + getModel(): Model | undefined; + getThinkingLevel(): ThinkingLevel; + getMessages(): AgentMessage[]; + getRequiredRequestAuth(model: Model): Promise<{ apiKey: string; headers?: Record }>; + getRetryPolicy(): ProviderRetryPolicy; + getExtensionRunner(): Pick; + disconnect(): void; + reconnect(): void; + emit(event: SessionRefinementEvent): void; + retainUnpersistedOutcome(message: CustomMessage): void; +} + +/** Thrown when a session_before_refine extension skips the refinement round. */ +export class RefineSkippedError extends Error {} + +/** Plans against current session dependencies and persists results during the apply barrier. */ +export class RefinementExecution { + constructor( + private readonly _host: RefinementExecutionHost, + private readonly _releaseAbort: (abort: AbortController) => void, + ) {} + _localHarnessStateDir(): string | undefined { + return ( + getLocalHarnessStateDir(this._host.sessionManager.getSessionArtifactDir()) ?? + (this._host.getRlmSessionDir() ? getLocalHarnessStateDir(this._host.getRlmSessionDir()) : undefined) + ); + } + + _loadMergedHarnessState(): HarnessState { + const localHarnessStateDir = this._localHarnessStateDir(); + return mergeHarnessStates( + loadHarnessState(getGlobalHarnessStateDir(), "global"), + localHarnessStateDir ? loadHarnessState(localHarnessStateDir, "local") : undefined, + ); + } + + private _loadRefinementHistory(): RefinementResult[] { + return mergeRefinementHistory( + loadGlobalRefinementHistory(getGlobalHarnessStateDir()), + getRefinementHistory(this._host.sessionManager.getEntries().filter((entry) => entry.type === "custom")), + ); + } + + async _planRefine( + options: { instructions?: string; rollbackId?: string; global?: boolean }, + signal: AbortSignal, + trigger: "manual" | "auto" = "manual", + ): Promise { + if (this._host.isDisposed()) { + throw new Error("Cannot refine a disposed session."); + } + + if (!this._host.getModel()) { + throw new Error(formatNoModelSelectedMessage()); + } + + const model = this._host.getModel()!; + const { apiKey, headers } = await this._host.getRequiredRequestAuth(model); + const globalHarnessStateDir = getGlobalHarnessStateDir(); + const localHarnessStateDir = this._localHarnessStateDir(); + const requestedScope = options.global ? "global" : "local"; + if (!options.rollbackId && requestedScope === "local" && !localHarnessStateDir) { + throw new Error("Local harness refinement requires a persisted session; use global refinement instead."); + } + const globalPlanningState = loadHarnessState(globalHarnessStateDir, "global"); + const localPlanningState = localHarnessStateDir ? loadHarnessState(localHarnessStateDir, "local") : undefined; + const planningState = + requestedScope === "global" + ? globalPlanningState + : mergeHarnessStates(globalPlanningState, localPlanningState); + const history = this._loadRefinementHistory(); + const rollbackTarget = options.rollbackId ? history.find((item) => item.id === options.rollbackId) : undefined; + let baselineScope = rollbackTarget + ? (inferRefinementResultScope(rollbackTarget) ?? requestedScope) + : requestedScope; + let baselineHarnessStateDir = baselineScope === "global" ? globalHarnessStateDir : localHarnessStateDir; + if (rollbackTarget?.harnessStatePath) { + baselineHarnessStateDir = dirname(rollbackTarget.harnessStatePath); + baselineScope = resolve(baselineHarnessStateDir) === resolve(globalHarnessStateDir) ? "global" : "local"; + } + if (!baselineHarnessStateDir) { + throw new Error("Local harness refinement requires a persisted session; use global refinement instead."); + } + const baselineState = rollbackTarget + ? loadHarnessState(baselineHarnessStateDir, baselineScope) + : baselineScope === "global" + ? globalPlanningState + : localPlanningState!; + if (!options.rollbackId && this._host.getExtensionRunner().hasHandlers("session_before_refine")) { + const result = (await this._host.getExtensionRunner().emit({ + type: "session_before_refine", + preparation: { + trigger, + instructions: options.instructions, + scope: requestedScope, + planningState, + history, + conversationText: serializeConversation(convertToLlm(this._host.getMessages())).slice(-80_000), + }, + signal, + })) as SessionBeforeRefineResult | undefined; + if (this._host.isDisposed() || signal.aborted) { + throw new Error("Refinement cancelled because the session was disposed."); + } + if (result?.skip) { + throw new RefineSkippedError("Refinement skipped by extension"); + } + if (result?.proposal !== undefined) { + return { + proposal: normalizeRefinementProposal(result.proposal), + id: generateRefinementId(), + baselineState, + }; + } + } + const plan = await planRefinement( + this._host.getMessages(), + planningState, + history, + model, + apiKey, + { ...options, retry: this._host.getRetryPolicy() }, + headers, + signal, + this._host.getThinkingLevel(), + ); + if (this._host.isDisposed() || signal.aborted) { + throw new Error("Refinement cancelled because the session was disposed."); + } + return { ...plan, baselineState }; + } + + async _applyRefine( + plan: RefinementPlan, + options: { instructions?: string; rollbackId?: string; global?: boolean }, + refineAbort: AbortController, + source: RefinementSource, + ): Promise { + if (this._host.isDisposed()) { + throw new Error("Cannot refine a disposed session."); + } + // The caller has already set _refineInFlight and waited for agent idle. + // Disconnect only for the brief apply + save + reconnect critical section. + this._host.disconnect(); + + try { + const globalHarnessStateDir = getGlobalHarnessStateDir(); + const localHarnessStateDir = this._localHarnessStateDir(); + const requestedScope = options.global ? "global" : "local"; + const history = this._loadRefinementHistory(); + const rollbackTarget = options.rollbackId ? history.find((item) => item.id === options.rollbackId) : undefined; + let targetScope = plan.rollbackScope ?? requestedScope; + let targetHarnessStateDir = targetScope === "global" ? globalHarnessStateDir : localHarnessStateDir; + if (targetScope === "local" && rollbackTarget?.harnessStatePath) { + if (!existsSync(rollbackTarget.harnessStatePath)) { + throw new Error( + `Local refinement ${rollbackTarget.id} state file not found: ${rollbackTarget.harnessStatePath}`, + ); + } + targetHarnessStateDir = dirname(rollbackTarget.harnessStatePath); + // Legacy records predate scope fields and default to "local" but may point + // at the global store; honor the recorded path so its entries stay global. + if (resolve(targetHarnessStateDir) === resolve(globalHarnessStateDir)) { + targetScope = "global"; + } + } + if (!targetHarnessStateDir) { + throw new Error("Local harness refinement requires a persisted session; use global refinement instead."); + } + // Re-read the target state immediately before applying so concurrent kernel + // (`rlm.harness`) writes during the LLM pass are not clobbered. + const state = loadHarnessState(targetHarnessStateDir, targetScope); + const proposal = { + ...plan.proposal, + edits: plan.proposal.edits.map((edit) => { + const localPrefix = "local:"; + const globalPrefix = "global:"; + return { + ...edit, + id: edit.id?.startsWith(localPrefix) + ? edit.id.slice(localPrefix.length) + : edit.id?.startsWith(globalPrefix) + ? edit.id.slice(globalPrefix.length) + : edit.id, + }; + }), + }; + if (this._host.isDisposed() || refineAbort.signal.aborted) { + throw new Error("Refinement cancelled because the session was disposed."); + } + const result = applyRefinementProposal(state, proposal, { + id: plan.id, + rollbackOf: plan.rollbackOf, + scope: targetScope, + baselineState: plan.baselineState, + }); + result.harnessStatePath = saveHarnessState(targetHarnessStateDir, state); + if (targetScope === "global") { + appendGlobalRefinement(globalHarnessStateDir, result); + } + let refinementAuditAppendError: { error: unknown } | undefined; + try { + this._host.sessionManager.appendCustomEntry("prime-agent.refinement", result); + } catch (error) { + refinementAuditAppendError = { error }; + } + try { + this._recordRefinementOutcome(result); + } catch (error) { + if (!refinementAuditAppendError) throw error; + } + if (refinementAuditAppendError) throw refinementAuditAppendError.error; + // The prompt stays byte-identical so the provider prefix cache survives; the notice carries the change. + this._recordRefinementNotice(result, source); + try { + this._host.emit({ type: "refine_complete", result }); + } catch { + // Listener failures must not flip a successful refinement into + // a reported failure — the refinement is already persisted. + } + try { + await this._host.getExtensionRunner().emit({ + type: "refine_complete", + id: result.id, + summary: result.summary, + appliedEdits: result.appliedEdits.filter((edit) => edit.applied).length, + scope: result.scope ?? "local", + }); + } catch { + // Extension emit failures must not flip a successful refinement + // into a reported failure — the refinement is already persisted. + } + return result; + } finally { + this._releaseAbort(refineAbort); + if (!this._host.isDisposed()) { + this._host.reconnect(); + } + } + } + + private _recordRefinementOutcome(result: RefinementResult): void { + this._appendDurableRefineMessage(createRefinementOutcomeMessage(result)); + } + + private _recordRefinementNotice(result: RefinementResult, source: RefinementSource): void { + if (!result.appliedEdits.some((edit) => edit.applied)) return; + this._appendDurableRefineMessage(createRefinementNoticeMessage(result, source)); + } + + private _appendDurableRefineMessage(message: CustomMessage): void { + try { + this._host.sessionManager.appendCustomMessageEntryWithRollback( + message.customType, + message.content, + message.display, + message.details, + ); + } catch { + // Not in the session file, so context rebuilds would drop the outcome. + this._host.retainUnpersistedOutcome(message); + } + this._host.getMessages().push(message); + this._host.emit({ type: "message_start", message }); + this._host.emit({ type: "message_end", message }); + } + + async review(context: AutoRefineReviewRequest, signal?: AbortSignal): Promise { + const model = this._host.getModel(); + if (!model) { + return { shouldRefine: false, rationale: "No model selected." }; + } + const { apiKey, headers } = await this._host.getRequiredRequestAuth(model); + return reviewAutoRefine( + this._host.getMessages(), + this._loadMergedHarnessState(), + this._loadRefinementHistory(), + model, + apiKey, + context, + headers, + signal, + this._host.getThinkingLevel(), + this._host.getRetryPolicy(), + ); + } +} diff --git a/packages/coding-agent/src/session/refinement.ts b/packages/coding-agent/src/session/refinement.ts new file mode 100644 index 0000000000..b218d93a4d --- /dev/null +++ b/packages/coding-agent/src/session/refinement.ts @@ -0,0 +1,925 @@ +import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { ExtensionRunner } from "../core/extensions/index.js"; +import type { CustomMessage, RefinementSource } from "../core/messages.js"; +import type { ProviderRetryPolicy } from "../core/provider-retry.js"; +import type { HarnessState, RefinementPlan, RefinementResult } from "../core/refinement/index.js"; +import type { SessionManager } from "../core/session-manager.js"; +import type { SettingsManager } from "../core/settings-manager.js"; +import { AutoRefinement, type AutoRefineReviewer, autoRefineInstructions } from "./auto-refinement.js"; +import { RefinementExecution, RefineSkippedError, type SessionRefinementEvent } from "./refinement-execution.js"; + +export type { AutoRefineReviewer, AutoRefineReviewRequest } from "./auto-refinement.js"; + +export interface SessionRefinementHost { + sessionManager: Pick< + SessionManager, + "getSessionArtifactDir" | "getEntries" | "appendCustomMessageEntryWithRollback" | "appendCustomEntry" + >; + settingsManager: Pick; + getRetryPolicy(): ProviderRetryPolicy; + isDisposed(): boolean; + isDisposing(): boolean; + isStreaming(): boolean; + isCompacting(): boolean; + getDepth(): number; + getRlmSessionDir(): string | undefined; + getModel(): Model | undefined; + getThinkingLevel(): ThinkingLevel; + getMessages(): AgentMessage[]; + getRequiredRequestAuth(model: Model): Promise<{ apiKey: string; headers?: Record }>; + getExtensionRunner(): Pick; + getEventQueue(): Promise; + getCompactionOperation(): Promise | undefined; + getBranchSummaryOperation(): Promise | undefined; + waitForAgentIdle(): Promise; + dispatchRefine( + options: { instructions?: string; global?: boolean }, + internal: { source: "self" } | { trigger: "auto" }, + ): Promise; + disconnect(): void; + reconnect(): void; + emit(event: SessionRefinementEvent): void; + retainUnpersistedOutcome(message: CustomMessage): void; + notifyCheckpoints(): void; + scheduleInputPump(): void; + isContinuationScheduled(): boolean; + cancelContinuation(): void; +} + +/** Thrown when a session_before_refine extension skips the refinement round. */ +export { RefineSkippedError } from "./refinement-execution.js"; + +/** + * Discriminated result from a serialized-mode background planning pass. + * - "plan": review approved and planning succeeded; carry the exact plan, + * options, and abort controller so the boundary can apply directly + * without a second planning request. + * - "skip": reviewer declined; no refine needed. + * - "failure": review or planning threw; boundary should not retry. + */ +export type SerializedBackgroundPlanResult = + | { + status: "plan"; + plan: RefinementPlan; + options: { instructions?: string; rollbackId?: string; global?: boolean }; + abort: AbortController; + branchVersion: number; + source: Exclude; + } + | { status: "skip"; explicit?: boolean } + | { status: "invalidated"; branchVersion: number } + | { + status: "failure"; + explicit: boolean; + options: { instructions?: string; rollbackId?: string; global?: boolean }; + branchVersion: number; + }; + +/** Owns refinement admission, planning/apply barriers, and serialized plan claims. */ +export class SessionRefinement { + private readonly _auto: AutoRefinement; + private readonly _execution: RefinementExecution; + private _refineAbortController?: AbortController; + private readonly _serializedRefine: boolean; + private _refineInFlight?: Promise; + private _refinePlanInFlight?: Promise; + private _serializedPlanInFlight?: Promise; + private _serializedPlanClaim?: Promise; + private _serializedExplicitRefineOptions?: { + instructions?: string; + global?: boolean; + }; + private _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; + + constructor( + private readonly _host: SessionRefinementHost, + config: { autoRefineReviewer?: AutoRefineReviewer; serializedRefine?: boolean }, + ) { + this._execution = new RefinementExecution(_host, (abort) => { + if (this._refineAbortController === abort) this._refineAbortController = undefined; + }); + this._auto = new AutoRefinement( + { + settingsManager: _host.settingsManager, + isDisposed: () => _host.isDisposed(), + isDisposing: () => _host.isDisposing(), + isStreaming: () => _host.isStreaming(), + isCompacting: () => _host.isCompacting(), + isAllowed: () => this._autoRefineAllowedForSession(), + getModel: () => _host.getModel(), + isContinuationScheduled: () => _host.isContinuationScheduled(), + cancelContinuation: () => _host.cancelContinuation(), + refine: (options, internal) => _host.dispatchRefine(options, internal), + runSerialized: (options, source) => this._runSerializedRefine(options, source), + emitFailure: (error) => this._emitRefineFailed(error), + review: (context, signal) => this._execution.review(context, signal), + }, + config.serializedRefine ?? false, + config.autoRefineReviewer, + ); + this._serializedRefine = config.serializedRefine ?? false; + } + + requestAbort(): void { + this._pendingRequestedRefine = undefined; + this._auto.invalidatePlans(); + this._auto.abortReview(); + this._refineAbortController?.abort(); + } + observeAssistantEnd(): void { + this._auto.observeAssistantEnd(); + this._maybeStartSerializedBackgroundPlan(); + } + get isApplying(): boolean { + return this._refineInFlight !== undefined; + } + get serialized(): boolean { + return this._serializedRefine; + } + /** Split settlement preserves the caller's existing await boundary. */ + beginAbortedTurnCleanup(): { promise: Promise; finish(): void } | undefined { + this._pendingRequestedRefine = undefined; + const plan = this._serializedPlanInFlight; + if (!plan) return undefined; + this._auto.invalidatePlans(); + this._refineAbortController?.abort(); + return { + promise: plan, + finish: () => { + if (this._serializedPlanInFlight === plan) { + this._serializedPlanInFlight = undefined; + this._serializedExplicitRefineOptions = undefined; + } + }, + }; + } + dispose(): void { + this._auto.abortReview(); + this._refineAbortController?.abort(); + this._auto.cancelScheduled(); + this._serializedPlanInFlight = undefined; + this._serializedExplicitRefineOptions = undefined; + this._pendingRequestedRefine = undefined; + this._auto._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); + this._auto.invalidatePlans(); + } + + async _runSerializedRefineCheckpoint(): Promise { + if (this._host.isDisposed() || this._host.isDisposing()) { + return; + } + + // 1. Await any background plan that was started at message_end + // (either for a pending refine.run or for interval-triggered + // auto-refine). This must be checked BEFORE the pending and + // interval checks because background planning may have consumed + // the pending request at message_end. + const branchVersion = this._auto.branchVersion; + const bgConsumption = await this._consumeSerializedBackgroundPlan(async (bgResult) => { + if (this._host.isDisposed() || this._host.isDisposing()) { + return true; + } + + if (bgResult?.status === "plan") { + if (bgResult.branchVersion !== this._auto.branchVersion) { + if (!this._pendingRequestedRefine) { + this._auto.stampCooldown(); + this._auto.resetTurns(); + return true; + } + } else { + // Apply the EXACT background plan directly via _applyRefine + // (no second _planRefine call). + try { + await this._applySerializedPlan(bgResult); + } catch (error) { + this._emitRefineFailed(error); + } + this._auto.stampCooldown(); + this._auto.resetTurns(); + if (!this._pendingRequestedRefine) { + return true; + } + } + } + + if (bgResult?.status === "skip") { + // Reviewer declined or an extension skipped during background planning. + // Reset exactly once. Never retry the interval review; only fall through for a separate pending refine.run. + if (bgResult.explicit) { + this._emitRefineFailed(new RefineSkippedError("Refinement skipped by extension")); + } + this._auto.stampCooldown(); + this._auto.resetTurns(); + if (!this._pendingRequestedRefine) { + return true; + } + } + + if (bgResult?.status === "failure") { + // Background review or planning failure stamps cooldown without a synchronous retry. + // A separately queued refine.run may still be serviced below. + if (branchVersion === this._auto.branchVersion) { + this._auto.stampCooldown(); + } + // Re-queue an explicit refine.run whose background plan failed, + // but only when branchVersion is still current and no newer + // pending request has arrived since the background plan consumed + // the original one. A newer request retains priority; interval + // failures keep existing no-retry cooldown semantics. + if ( + bgResult.explicit && + bgResult.branchVersion === this._auto.branchVersion && + !this._pendingRequestedRefine + ) { + this._pendingRequestedRefine = bgResult.options; + } + if (!this._pendingRequestedRefine) { + return true; + } + } + + if (bgResult?.status === "invalidated" && !this._pendingRequestedRefine) { + this._auto.stampCooldown(); + this._auto.resetTurns(); + return true; + } + + await this._runSerializedRefineCheckpointAfterBackground(branchVersion); + return true; + }); + if (this._host.isDisposed() || this._host.isDisposing() || bgConsumption !== "none") { + return; + } + await this._runSerializedRefineCheckpointAfterBackground(branchVersion); + } + + private async _runSerializedRefineCheckpointAfterBackground(branchVersion: number): Promise { + // No background result, or a refine.run arrived while the background result was + // in flight. Fall through so an explicit pending request is serviced at this boundary. + + // 2. Agent-callable refine.run requests that were NOT consumed by + // background planning (e.g. interval not reached at message_end, + // or cooldown was active). Service them synchronously. + const pending = this._pendingRequestedRefine; + if (pending) { + this._pendingRequestedRefine = undefined; + try { + await this._runSerializedRefine(pending, "self"); + } catch (error) { + this._emitRefineFailed(error); + } + this._auto.stampCooldown(); + this._auto.resetTurns(); + return; + } + + // 3. Post-compaction auto-refine. Serialized sessions defer the + // compaction trigger to this boundary instead of entering the interactive + // path, which waits for agent idle and can never run inside a tool loop. + if (!this._autoRefineAllowedForSession()) { + this._auto.discardCompact(); + return; + } + const settings = this._host.settingsManager.getAutoRefineSettings(); + if (!settings.enabled) { + this._auto.discardCompact(); + return; + } + if (this._auto.hasPendingCompact) { + if (!settings.compact) { + this._auto.discardCompact(); + } else { + const nowMs = Date.now(); + const underCooldown = this._auto.lastReviewAt > 0 && nowMs - this._auto.lastReviewAt < settings.cooldownMs; + if (underCooldown) { + // Preserve the compact trigger for a later boundary, matching the + // interactive path's pending behavior while the cooldown is active. + return; + } + this._auto.discardCompact(); + await this._auto._runSerializedAutoRefineReview("compact", branchVersion); + return; + } + } + + // 4. Interval-triggered auto-refine (no background plan was started). + if (this._auto.turnsSinceReview < settings.turnInterval) { + return; + } + const nowMs = Date.now(); + const underCooldown = this._auto.lastReviewAt > 0 && nowMs - this._auto.lastReviewAt < settings.cooldownMs; + if (underCooldown) { + return; + } + await this._auto._runSerializedAutoRefineReview("turn_interval", branchVersion); + } + + private async _consumeSerializedBackgroundPlan( + consume: (result: SerializedBackgroundPlanResult | undefined) => Promise, + ): Promise<"none" | "waited" | "continue" | "stop"> { + if (this._serializedPlanClaim) { + await this._serializedPlanClaim.catch(() => undefined); + return "waited"; + } + const planInFlight = this._serializedPlanInFlight; + if (!planInFlight) { + return "none"; + } + + let releaseClaim: () => void = () => {}; + const claim = new Promise((resolve) => { + releaseClaim = resolve; + }); + this._serializedPlanClaim = claim; + try { + const result = await planInFlight.catch(() => undefined); + if (this._serializedPlanInFlight === planInFlight) { + this._serializedPlanInFlight = undefined; + this._serializedExplicitRefineOptions = undefined; + } + return (await consume(result)) ? "stop" : "continue"; + } finally { + releaseClaim(); + if (this._serializedPlanClaim === claim) { + this._serializedPlanClaim = undefined; + } + } + } + + private async _applySerializedPlan( + bgResult: Extract, + ): Promise { + let resolveApplySettled: () => void = () => {}; + const applySettled = new Promise((resolve) => { + resolveApplySettled = resolve; + }); + this._refineInFlight = applySettled; + try { + await this._execution._applyRefine(bgResult.plan, bgResult.options, bgResult.abort, bgResult.source); + } finally { + resolveApplySettled(); + if (this._refineInFlight === applySettled) { + this._refineInFlight = undefined; + } + this._host.notifyCheckpoints(); + this._host.scheduleInputPump(); + } + } + + private _maybeStartSerializedBackgroundPlan(): void { + if (!this._serializedRefine || this._host.isDisposed() || this._host.isDisposing()) { + return; + } + // Don't start if a plan is already in flight. + if (this._serializedPlanInFlight || this._refineInFlight || this._refinePlanInFlight) { + return; + } + + // Start background planning for a pending agent-callable + // refine.run request, so its plan is ready at the shouldStopAfterTurn + // boundary. The pending request is consumed (cleared) here so the + // boundary doesn't re-plan it. Explicit refine.run skips the review gate. + const pending = this._pendingRequestedRefine; + if (pending) { + this._pendingRequestedRefine = undefined; + this._serializedExplicitRefineOptions = pending; + const refineAbort = new AbortController(); + this._refineAbortController = refineAbort; + const branchVersion = this._auto.branchVersion; + this._serializedPlanInFlight = this._runBackgroundPlan(pending, refineAbort, branchVersion, true); + return; + } + + // Interval-triggered auto-refine background planning. + if (!this._autoRefineAllowedForSession()) { + return; + } + const settings = this._host.settingsManager.getAutoRefineSettings(); + if (!settings.enabled) { + return; + } + if (this._auto.turnsSinceReview < settings.turnInterval) { + return; + } + const nowMs = Date.now(); + const underCooldown = this._auto.lastReviewAt > 0 && nowMs - this._auto.lastReviewAt < settings.cooldownMs; + if (underCooldown) { + return; + } + + const refineAbort = new AbortController(); + this._refineAbortController = refineAbort; + const branchVersion = this._auto.branchVersion; + // Pass empty options — _runBackgroundPlan derives instructions from + // the review result for interval-triggered auto-refine. + this._serializedPlanInFlight = this._runBackgroundPlan({}, refineAbort, branchVersion); + } + + private async _runBackgroundPlan( + options: { instructions?: string; rollbackId?: string; global?: boolean }, + refineAbort: AbortController, + branchVersion: number, + skipReview = false, + ): Promise { + try { + let planOptions = options; + if (!skipReview) { + // Interval-triggered: run the review gate first, then derive + // instructions from the review result (not prepopulated). + const review = await this._auto._reviewAutoRefine( + { + reason: "turn_interval", + turnsSinceLastReview: this._auto.turnsSinceReview, + }, + refineAbort.signal, + ); + if (this._host.isDisposed() || this._host.isDisposing() || branchVersion !== this._auto.branchVersion) { + return { status: "invalidated", branchVersion }; + } + if (!review.shouldRefine) { + return { status: "skip" }; + } + planOptions = { + instructions: autoRefineInstructions("turn_interval", review), + }; + } + // For explicit refine.run (skipReview=true), plan directly with + // the user-provided options — no auto-review gate. + const plan = await this._execution._planRefine( + planOptions, + refineAbort.signal, + skipReview ? "manual" : "auto", + ); + if (this._host.isDisposed() || this._host.isDisposing() || branchVersion !== this._auto.branchVersion) { + return { status: "invalidated", branchVersion }; + } + return { + status: "plan", + plan, + options: planOptions, + abort: refineAbort, + branchVersion, + source: skipReview ? "self" : "auto", + }; + } catch (error) { + if (this._host.isDisposed() || this._host.isDisposing() || branchVersion !== this._auto.branchVersion) { + return { status: "invalidated", branchVersion }; + } + if (error instanceof RefineSkippedError) { + return { status: "skip", explicit: skipReview }; + } + return { + status: "failure", + explicit: skipReview, + options, + branchVersion, + }; + } finally { + if (this._refineAbortController === refineAbort) { + this._refineAbortController = undefined; + } + } + } + + private async _runSerializedRefine( + options: { + instructions?: string; + rollbackId?: string; + global?: boolean; + }, + source: Exclude, + ): Promise { + if (this._host.isDisposed() || this._host.isDisposing()) { + return; + } + // Guard: serialize against concurrent _runSerializedRefine calls. + // _serializedPlanInFlight covers background planning; _refineInFlight + // covers the apply phase. Both must be settled before starting a new + // plan+apply cycle. + while (this._serializedPlanInFlight || this._refineInFlight || this._refinePlanInFlight) { + if (this._serializedPlanInFlight) { + await this._consumeSerializedBackgroundPlan(async () => false); + } else if (this._refineInFlight) { + await this._refineInFlight; + } else { + await this._refinePlanInFlight; + } + } + if (this._host.isDisposed() || this._host.isDisposing()) { + return; + } + + const refineAbort = new AbortController(); + this._refineAbortController = refineAbort; + + const planRun = this._execution._planRefine(options, refineAbort.signal, source === "auto" ? "auto" : "manual"); + const planSettled = planRun.then( + () => undefined, + () => undefined, + ); + this._refinePlanInFlight = planSettled; + let plan: RefinementPlan; + try { + plan = await planRun; + } catch (error) { + if (this._refineAbortController === refineAbort) { + this._refineAbortController = undefined; + } + this._host.scheduleInputPump(); + throw error; + } finally { + if (this._refinePlanInFlight === planSettled) { + this._refinePlanInFlight = undefined; + } + } + + if (this._host.isDisposed() || refineAbort.signal.aborted) { + if (this._refineAbortController === refineAbort) { + this._refineAbortController = undefined; + } + this._host.scheduleInputPump(); + return; + } + + // Do NOT call agent.waitForIdle() — we are at the quiescent boundary + // already (shouldStopAfterTurn). _applyRefine handles disconnect/reconnect internally. + let resolveApplySettled: () => void = () => {}; + const applySettled = new Promise((resolve) => { + resolveApplySettled = resolve; + }); + this._refineInFlight = applySettled; + try { + await this._execution._applyRefine(plan, options, refineAbort, source); + } finally { + resolveApplySettled(); + if (this._refineInFlight === applySettled) { + this._refineInFlight = undefined; + } + this._host.notifyCheckpoints(); + this._host.scheduleInputPump(); + } + } + + handleRefineHostRequest(type: string, payload: Record = {}): Record { + switch (type) { + case "refine.status": { + return { + pending: this._pendingRequestedRefine !== undefined, + in_flight: + this._refineInFlight !== undefined || + this._refinePlanInFlight !== undefined || + this._serializedPlanInFlight !== undefined, + }; + } + case "refine.run": { + const instructions = payload.instructions; + if (instructions !== undefined && typeof instructions !== "string") { + throw new Error("refine.run instructions must be a string when provided"); + } + const globalFlag = payload.global; + if (globalFlag !== undefined && typeof globalFlag !== "boolean") { + throw new Error("refine.run global must be a boolean when provided"); + } + if (!this._host.isStreaming()) { + return { + scheduled: false, + reason: "no active turn; refine can only be requested while a turn is running", + }; + } + const previous = this._pendingRequestedRefine ?? this._serializedExplicitRefineOptions; + this._pendingRequestedRefine = { + instructions: instructions ?? previous?.instructions, + global: globalFlag ?? previous?.global, + }; + // In serialized mode, kick off background planning immediately + // (the primary response ended at message_end, tools are active). + // This lets planning overlap tool execution rather than waiting + // for the shouldStopAfterTurn boundary. + if (this._serializedRefine) { + if (this._serializedPlanInFlight) { + this._auto.invalidatePlans(); + if (this._refineAbortController) { + this._refineAbortController.abort(); + } else { + this._serializedPlanInFlight = Promise.resolve({ + status: "invalidated", + branchVersion: this._auto.branchVersion, + }); + } + } else { + this._maybeStartSerializedBackgroundPlan(); + } + } + return { + scheduled: true, + note: "Refinement runs when the current turn ends; applied edits are appended to your context as a refinement notice and you resume automatically. Continue working normally.", + }; + } + default: + throw new Error(`unknown refine request type "${type}"`); + } + } + + async _drainPendingRefinementForDisposal(): Promise { + this._auto.cancelScheduled(); + await Promise.allSettled(this._auto.pendingOperations()); + this._auto.cancelScheduled(); + // Wait for in-flight refinement (including serialized background plan) to settle. + while (this._refineInFlight || this._refinePlanInFlight || this._serializedPlanInFlight) { + if (this._refineInFlight) { + await this._refineInFlight; + } else if (this._refinePlanInFlight) { + await this._refinePlanInFlight; + } else if (this._serializedPlanInFlight) { + // Await the background plan and apply a ready "plan" result before teardown. + await this._consumeSerializedBackgroundPlan(async (bgResult) => { + if (bgResult?.status === "plan" && bgResult.branchVersion === this._auto.branchVersion) { + try { + await this._applySerializedPlan(bgResult); + } catch (error) { + this._emitRefineFailed(error); + } + // Stamp cooldown and reset counter so the interval + // check below does not trigger a duplicate refine. + this._auto.stampCooldown(); + this._auto.resetTurns(); + } + // Preserve a consumed explicit request when its background plan failed, + // matching the turn-boundary recovery path. The pending drain below + // retries it once before disposal. + if ( + bgResult?.status === "failure" && + bgResult.explicit && + bgResult.branchVersion === this._auto.branchVersion && + !this._pendingRequestedRefine + ) { + this._pendingRequestedRefine = bgResult.options; + } + if (bgResult?.status === "skip" && bgResult.explicit) { + this._emitRefineFailed(new RefineSkippedError("Refinement skipped by extension")); + } + // For "skip" or "failure", stamp cooldown and reset counter + // so the interval check below does not trigger a duplicate + // terminal retry. + if ( + bgResult?.status === "skip" || + bgResult?.status === "failure" || + bgResult?.status === "invalidated" + ) { + this._auto.stampCooldown(); + this._auto.resetTurns(); + } + return false; + }); + } else { + await new Promise((resolve) => setTimeout(resolve, 0)); + } + } + // Drain an agent-callable refine.run request that was scheduled but + // not yet consumed. Use the direct serialized path (no waitForIdle) + // since the agent may still own activeRun at the final agent_end. + if (this._pendingRequestedRefine) { + const pending = this._pendingRequestedRefine; + this._pendingRequestedRefine = undefined; + try { + await this._runSerializedRefine(pending, "self"); + } catch { + // Best-effort drain; refinement errors must not block disposal. + } + // Stamp cooldown and reset counter so the interval check below + // does not trigger a duplicate refine after the explicit drain. + this._auto.stampCooldown(); + this._auto.resetTurns(); + } + // A serialized compaction can finish without another model turn. Drain its + // pending review here so disposal does not silently lose the trigger. + if (this._serializedRefine && this._auto.hasPendingCompact && this._autoRefineAllowedForSession()) { + const compactSettings = this._host.settingsManager.getAutoRefineSettings(); + if (!compactSettings.enabled || !compactSettings.compact) { + this._auto.discardCompact(); + } else { + const nowMs = Date.now(); + const underCooldown = + this._auto.lastReviewAt > 0 && nowMs - this._auto.lastReviewAt < compactSettings.cooldownMs; + this._auto.discardCompact(); + if (!underCooldown) { + try { + await this._auto._runSerializedAutoRefineReview("compact", this._auto.branchVersion); + } catch { + // Best-effort drain; refinement errors must not block disposal. + } + return; + } + } + } + + // If auto-refine is due but has not started yet, run it now so the + // refinement is persisted before disposal. Use the direct serialized + // path in serialized mode, or _maybeAutoRefine in interactive mode + // (where the agent is idle at this point). + if (this._host.isDisposed() || !this._autoRefineAllowedForSession()) { + return; + } + const settings = this._host.settingsManager.getAutoRefineSettings(); + if (!settings.enabled) { + return; + } + if (this._auto.turnsSinceReview < settings.turnInterval) { + return; + } + const nowMs = Date.now(); + const underCooldown = this._auto.lastReviewAt > 0 && nowMs - this._auto.lastReviewAt < settings.cooldownMs; + if (underCooldown) { + return; + } + if (this._serializedRefine) { + await this._runSerializedRefineCheckpoint(); + } else { + await this._auto._maybeAutoRefine("turn_interval"); + } + } + + _autoRefineAllowedForSession(): boolean { + return this._host.getDepth() === 0 && this._execution._localHarnessStateDir() !== undefined; + } + + async _invalidatePendingAutoRefineForBranchChange(): Promise { + this._auto.abortReview(); + this._auto._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); + this._auto.resetTurns(); + // Increment branch version BEFORE aborting/awaiting the serialized plan. + // This invalidates the plan's branchVersion check at the boundary + // so even if the plan completes, the boundary will reject it + // (bgResult.branchVersion !== this._auto.branchVersion). + this._auto.invalidatePlans(); + // Abort the in-flight refine/bplan controller so any pending + // _planRefine or _reviewAutoRefine call settles via signal abort + // rather than hanging forever. + this._refineAbortController?.abort(); + if (this._serializedPlanInFlight) { + await this._consumeSerializedBackgroundPlan(async () => false); + } + while (this._refinePlanInFlight) { + await this._refinePlanInFlight; + } + await this._waitForRefineIdle(); + } + + _emitRefineFailed(error: unknown): void { + this._host.emit({ + type: "refine_failed", + error: error instanceof Error ? error.message : String(error), + }); + } + + _consumePendingRequestedRefine(): boolean { + const pending = this._pendingRequestedRefine; + if (!pending) return false; + this._pendingRequestedRefine = undefined; + void this._host.dispatchRefine(pending, { source: "self" }).catch((error) => this._emitRefineFailed(error)); + return true; + } + + async refine( + options: { + instructions?: string; + rollbackId?: string; + global?: boolean; + } = {}, + internal: { skipAbort?: boolean; trigger?: "manual" | "auto"; source?: RefinementSource } = {}, + ): Promise { + // Queued /refine executes from the session-input pump between turns; + // refine never aborts the agent (planning is backgrounded and the apply + // phase waits for quiescence), so skipAbort only asserts the pump's + // idle invariant instead of changing abort behavior. + if (internal.skipAbort && this._host.isStreaming()) { + throw new Error("Cannot refine without aborting while the agent is running."); + } + // Wait for any existing refine (both planning and application) before + // starting a new run. This serializes concurrent /refine calls so two + // planning phases cannot race into concurrent _applyRefine calls that + // overwrite harness state. + while (this._refineInFlight || this._refinePlanInFlight || this._serializedPlanInFlight) { + if (this._refineInFlight) { + await this._refineInFlight; + } else if (this._refinePlanInFlight) { + await this._refinePlanInFlight; + } else { + // A serialized background plan is in flight (started during an + // active turn at message_end). Wait for planning and for the active + // turn to settle so its normal checkpoint can consume the plan. + const serializedPlanInFlight = this._serializedPlanInFlight; + await serializedPlanInFlight; + if (this._refineInFlight || this._refinePlanInFlight) { + continue; + } + await this._host.waitForAgentIdle(); + // Aborted turns skip shouldStopAfterTurn. Drop their settled plan + // after idle so a later public refine cannot spin on it forever. + if (this._serializedPlanInFlight === serializedPlanInFlight) { + this._serializedPlanInFlight = undefined; + this._serializedExplicitRefineOptions = undefined; + } + } + } + + const refineAbort = new AbortController(); + this._refineAbortController = refineAbort; + + const planRun = this._execution._planRefine(options, refineAbort.signal, internal.trigger ?? "manual"); + const planSettled = planRun.then( + () => undefined, + () => undefined, + ); + this._refinePlanInFlight = planSettled; + let plan: RefinementPlan; + try { + plan = await planRun; + } catch (e) { + if (this._refineAbortController === refineAbort) { + this._refineAbortController = undefined; + } + this._host.scheduleInputPump(); + throw e; + } finally { + if (this._refinePlanInFlight === planSettled) { + this._refinePlanInFlight = undefined; + } + } + + // Block new turns before waiting for the current turn to finish. One shared + // settled promise covers the full transition and apply critical section. + let resolveApplySettled: () => void = () => {}; + const applySettled = new Promise((resolve) => { + resolveApplySettled = resolve; + }); + this._refineInFlight = applySettled; + try { + // Wait for the session to become quiescent before applying. Planning is + // allowed to overlap active user work, but application must not disconnect + // event handling until that work and its queued events have completed. + await this._host.waitForAgentIdle(); + while (true) { + const eventQueue = this._host.getEventQueue(); + const compactionOp = this._host.getCompactionOperation(); + const branchSummaryOp = this._host.getBranchSummaryOperation(); + await Promise.allSettled([ + eventQueue, + ...(compactionOp ? [compactionOp] : []), + ...(branchSummaryOp ? [branchSummaryOp] : []), + ]); + if ( + eventQueue === this._host.getEventQueue() && + compactionOp === this._host.getCompactionOperation() && + branchSummaryOp === this._host.getBranchSummaryOperation() + ) { + break; + } + } + if (this._host.isDisposed() || refineAbort.signal.aborted) { + throw new Error("Refinement cancelled because the session was disposed."); + } + return await this._execution._applyRefine( + plan, + options, + refineAbort, + internal.source ?? (internal.trigger === "auto" ? "auto" : "user"), + ); + } finally { + resolveApplySettled(); + if (this._refineInFlight === applySettled) { + this._refineInFlight = undefined; + } + this._host.notifyCheckpoints(); + this._host.scheduleInputPump(); + } + } + + async _waitForRefineIdle(): Promise { + while (this._refineInFlight) { + await this._refineInFlight; + } + } + + _discardPendingAutoRefine(options: { cancelPostCompactionContinue?: boolean } = {}): void { + this._auto._discardPendingAutoRefine(options); + } + + _scheduleAutoRefineAfterAgentEnd(): void { + this._auto._scheduleAutoRefineAfterAgentEnd(); + } + + _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void { + this._auto._scheduleAutoRefineAfterCompaction(willContinueAfterCompaction); + } + + _localHarnessStateDir(): string | undefined { + return this._execution._localHarnessStateDir(); + } + + _loadMergedHarnessState(): HarnessState { + return this._execution._loadMergedHarnessState(); + } +} diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index 1c8a6f950f..c4bc2bd5e6 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -21,6 +21,7 @@ import { ModelRegistry } from "../src/core/model-registry.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; import type { BuildSystemPromptOptions } from "../src/core/system-prompt.js"; +import type { SessionRefinement } from "../src/session/refinement.js"; import { createTestExtensionsResult, createTestResourceLoader } from "./utilities.js"; class MockAssistantStream extends EventStream { @@ -155,9 +156,9 @@ describe("AgentSession concurrent prompt guard", () => { }); let drainStarted = false; const internals = session as unknown as { - _drainPendingRefinementForDisposal: () => Promise; + _refinement: SessionRefinement; }; - vi.spyOn(internals, "_drainPendingRefinementForDisposal").mockImplementation(async () => { + vi.spyOn(internals._refinement, "_drainPendingRefinementForDisposal").mockImplementation(async () => { drainStarted = true; await drainGate; }); diff --git a/packages/coding-agent/test/session/compaction.test.ts b/packages/coding-agent/test/session/compaction.test.ts new file mode 100644 index 0000000000..5d58a5404e --- /dev/null +++ b/packages/coding-agent/test/session/compaction.test.ts @@ -0,0 +1,243 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage, registerFauxProvider } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { SessionManager } from "../../src/core/session-manager.js"; +import { SessionCompaction, type SessionCompactionHost } from "../../src/session/compaction.js"; +import { CompactionSkippedError } from "../../src/session/compaction-execution.js"; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: Error) => void; + const promise = new Promise((done, fail) => { + resolve = done; + reject = fail; + }); + return { promise, resolve, reject }; +} + +const providers: ReturnType[] = []; + +function setup() { + const provider = registerFauxProvider(); + providers.push(provider); + const store = SessionManager.inMemory(); + const messages: AgentMessage[] = []; + const order: string[] = []; + const result = { summary: "summary", firstKeptEntryId: "kept", tokensBefore: 1000 }; + const host = { + getSettings: vi.fn(() => ({ enabled: false, reserveTokens: 100, keepRecentTokens: 10 })), + runAutomatic: vi.fn(async () => false), + queueGoalContinuation: vi.fn(() => false), + queueAutonomousContinuation: vi.fn(async () => undefined), + beginRefinementAbort: vi.fn(() => undefined), + getModel: vi.fn(() => provider.getModel()), + isStreaming: vi.fn(() => false), + getRequiredAuth: vi.fn(async () => ({ apiKey: "faux" })), + getAuth: vi.fn(async () => ({ ok: true, apiKey: "faux" })), + perform: vi.fn(async () => result), + disconnect: vi.fn(() => { + order.push("disconnect"); + }), + reconnect: vi.fn(() => { + order.push("reconnect"); + }), + abortSession: vi.fn(async () => { + order.push("abort"); + }), + getContinuationState: vi.fn(() => ({ scheduled: true, continueAfterSessionInput: true })), + afterManualCompaction: vi.fn(() => { + order.push("resume"); + }), + getMessages: () => messages, + replaceMessages: vi.fn(), + hasAgentQueuedMessages: vi.fn(() => false), + hasPendingSessionWork: vi.fn(() => false), + scheduleContinuation: vi.fn(), + scheduleRefinement: vi.fn(), + takeThresholdAutonomousMessages: vi.fn(() => []), + getThresholdGoalContinuation: vi.fn(() => undefined), + clearAutonomousContinuations: vi.fn(), + clearGoalContinuation: vi.fn(), + getSessionStore: () => store, + retainUnpersistedOutcome: vi.fn(), + emit: vi.fn((event) => { + order.push(event.type); + }), + notifyCheckpoints: vi.fn(() => { + order.push("checkpoint"); + }), + scheduleInput: vi.fn(() => { + order.push("input"); + }), + } satisfies SessionCompactionHost; + return { compaction: new SessionCompaction(host), host, order, result, provider }; +} + +afterEach(() => { + while (providers.length) providers.pop()?.unregister(); + vi.restoreAllMocks(); +}); + +describe("SessionCompaction boundaries", () => { + it("waits for abort before starting, then reconnects and releases waiters before resuming", async () => { + const { compaction, host, order, result } = setup(); + const aborted = deferred(); + const performed = deferred(); + host.abortSession.mockReturnValue(aborted.promise); + host.perform.mockReturnValue(performed.promise); + compaction.request("pending instructions"); + const pending = compaction.compact("manual instructions"); + expect(compaction.operation).toBeUndefined(); + expect(host.emit).not.toHaveBeenCalled(); + aborted.resolve(); + await Promise.resolve(); + const operation = compaction.operation; + expect(operation).toBeInstanceOf(Promise); + expect(compaction.isRunning).toBe(true); + await Promise.resolve(); + host.afterManualCompaction.mockImplementation((signal, scheduled, continueAfterInput) => { + expect(compaction.operation).toBeUndefined(); + expect(compaction.isRunning).toBe(false); + expect([signal.aborted, scheduled, continueAfterInput]).toEqual([false, true, true]); + order.push("resume"); + }); + performed.resolve(result); + expect(await pending).toBe(result); + await operation; + expect(compaction.hasPendingRequest).toBe(false); + expect(order).toEqual([ + "disconnect", + "compaction_start", + "compaction_end", + "reconnect", + "checkpoint", + "input", + "resume", + ]); + }); + + it.each([new CompactionSkippedError("too short"), new Error("write failed")])( + "retains a pending request and releases manual waiters after %s", + async (error) => { + const { compaction, host } = setup(); + compaction.request(); + host.perform.mockRejectedValue(error); + const pending = compaction.compact(undefined, { skipAbort: true }); + const operation = compaction.operation; + await expect(pending).rejects.toBe(error); + await operation; + expect(compaction.operation).toBeUndefined(); + expect(compaction.hasPendingRequest).toBe(true); + expect(host.reconnect).toHaveBeenCalledOnce(); + expect(host.afterManualCompaction).not.toHaveBeenCalled(); + }, + ); + + it("consumes requests before automatic start and reads the current model after authentication", async () => { + const { compaction, host, provider } = setup(); + const authentication = deferred>>(); + host.getAuth.mockReturnValue(authentication.promise); + host.emit.mockImplementation((event) => { + if (event.type !== "compaction_start") return; + expect(compaction.hasPendingRequest).toBe(false); + expect(compaction.isRunning).toBe(false); + expect(event.customInstructions).toBe("keep decisions"); + }); + compaction.request("keep decisions"); + const pending = compaction.runAutomatic("requested", false); + const nextModel = { ...provider.getModel(), id: "next-model" }; + host.getModel.mockReturnValue(nextModel); + authentication.resolve({ ok: true, apiKey: "faux", headers: { "x-test": "auth" } }); + await pending; + expect(host.perform).toHaveBeenCalledWith( + expect.objectContaining({ + model: nextModel, + apiKey: "faux", + headers: { "x-test": "auth" }, + customInstructions: "keep decisions", + }), + ); + }); + + it.each(["threshold", "requested", "overflow"] as const)( + "handles queued work after %s failure without repeating an overflowing request", + async (reason) => { + const { compaction, host } = setup(); + host.hasPendingSessionWork.mockReturnValue(true); + host.perform.mockRejectedValue(new Error("summary failed")); + await compaction.runAutomatic(reason, reason === "overflow"); + expect(host.scheduleContinuation).toHaveBeenCalledTimes(reason === "overflow" ? 0 : 1); + expect(host.emit).toHaveBeenCalledWith( + expect.objectContaining({ type: "compaction_end", reason, willRetry: false }), + ); + expect(compaction.operation).toBeUndefined(); + }, + ); + + it("cancels active automatic work and withdraws its goal continuation without resuming", async () => { + const { compaction, host } = setup(); + const entered = deferred(); + const goalMessage = fauxAssistantMessage("queued goal"); + host.getThresholdGoalContinuation.mockReturnValue(goalMessage); + host.perform.mockImplementation( + ({ signal }) => + new Promise((_resolve, reject) => { + entered.resolve(signal); + signal.addEventListener("abort", () => reject(new Error("Compaction cancelled")), { once: true }); + }), + ); + compaction.requestContinuation(); + const pending = compaction.runAutomatic("threshold", false); + const signal = await entered.promise; + compaction.abortAutomatic(); + expect(signal.aborted).toBe(true); + await pending; + expect(host.clearGoalContinuation).toHaveBeenCalledExactlyOnceWith(goalMessage); + expect(host.scheduleContinuation).not.toHaveBeenCalled(); + expect(host.emit).toHaveBeenCalledWith( + expect.objectContaining({ type: "compaction_end", aborted: true, errorMessage: undefined }), + ); + expect(compaction.isRunning).toBe(false); + }); + + it("keeps the newer operation visible when an older manual compaction settles", async () => { + const { compaction, host, result } = setup(); + const first = deferred(); + const second = deferred(); + host.perform.mockReturnValueOnce(first.promise).mockReturnValueOnce(second.promise); + const firstRun = compaction.compact(undefined, { skipAbort: true }); + await Promise.resolve(); + const secondRun = compaction.compact(undefined, { skipAbort: true }); + await Promise.resolve(); + const operation = compaction.operation; + first.resolve(result); + await firstRun; + expect(compaction.operation).toBe(operation); + second.resolve(result); + await secondRun; + await operation; + expect(compaction.operation).toBeUndefined(); + }); + + it("does not add an await to aborted-turn checks when no refinement plan needs cleanup", async () => { + const { compaction, host } = setup(); + compaction.request(); + const pending = compaction.check(fauxAssistantMessage("", { stopReason: "aborted" }), false); + expect(compaction.hasPendingRequest).toBe(false); + expect(host.getSettings).toHaveBeenCalledOnce(); + await pending; + }); + + it("waits for an aborted refinement plan before proceeding with pre-prompt checks", async () => { + const { compaction, host } = setup(); + const plan = deferred(); + const finish = vi.fn(); + host.beginRefinementAbort.mockReturnValue({ promise: plan.promise, finish }); + const pending = compaction.check(fauxAssistantMessage("", { stopReason: "aborted" }), false); + expect(host.getSettings).not.toHaveBeenCalled(); + plan.reject(new Error("plan aborted")); + await pending; + expect(finish).toHaveBeenCalledOnce(); + expect(host.getSettings).toHaveBeenCalledOnce(); + }); +}); diff --git a/packages/coding-agent/test/session/continuation.test.ts b/packages/coding-agent/test/session/continuation.test.ts new file mode 100644 index 0000000000..f698bd90e6 --- /dev/null +++ b/packages/coding-agent/test/session/continuation.test.ts @@ -0,0 +1,159 @@ +import { AgentContinueError, type AgentMessage } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it, vi } from "vitest"; +import { SessionCommitFence } from "../../src/session/commit-fence.js"; +import { SessionContinuation, type SessionContinuationHost } from "../../src/session/continuation.js"; + +function deferred() { + let resolve!: () => void; + let reject!: (error: Error) => void; + const promise = new Promise((done, fail) => { + resolve = done; + reject = fail; + }); + return { promise, resolve, reject }; +} + +function turn(): Promise { + return new Promise((resolve) => setImmediate(resolve)); +} + +function setup() { + const fence = new SessionCommitFence(); + const waiters = new Set<() => void>(); + const host = { + waitForAgentIdle: vi.fn(async () => {}), + waitForRetry: vi.fn(async () => {}), + waitForRefinement: vi.fn(async () => {}), + queuedWorkPauseCount: vi.fn(() => 0), + addCheckpointWaiter: (waiter: () => void) => { + waiters.add(waiter); + }, + removeCheckpointWaiter: (waiter: () => void) => { + waiters.delete(waiter); + }, + notifyCheckpoints: vi.fn(() => { + for (const waiter of waiters) waiter(); + }), + compactionOperation: vi.fn(() => undefined), + isRefinementApplying: vi.fn(() => false), + acquireCommitFence: () => fence.acquire(), + scheduleRefinement: vi.fn(), + unfinishedActionCount: vi.fn(() => 0), + isInputRequested: vi.fn(() => false), + scheduleInput: vi.fn(), + continue: vi.fn(async () => {}), + waitForIdleOrSettlement: vi.fn(async () => {}), + removeQueuedMessages: vi.fn(() => []), + followUp: vi.fn(), + onMessageConsumed: vi.fn(), + } satisfies SessionContinuationHost; + return { continuation: new SessionContinuation(host), host, fence, waiters }; +} + +describe("SessionContinuation boundaries", () => { + it("settles cancelled paused work and removes its checkpoint waiter", async () => { + const { continuation, host, waiters } = setup(); + host.queuedWorkPauseCount.mockReturnValue(1); + continuation.schedule(); + const token = continuation.current!; + await turn(); + expect(waiters.size).toBe(1); + continuation.cancel(); + await token.promise; + await turn(); + expect(waiters.size).toBe(0); + expect(continuation.current).toBeUndefined(); + expect(host.continue).not.toHaveBeenCalled(); + }); + + it("releases the commit fence before waiting for the continued turn", async () => { + const { continuation, host, fence } = setup(); + const running = deferred(); + host.continue.mockReturnValue(running.promise); + continuation.schedule(); + const token = continuation.current!; + await turn(); + expect(host.continue).toHaveBeenCalledOnce(); + expect(fence.hasPendingWork).toBe(false); + const lease = await fence.acquire(); + lease.release(); + expect(continuation.current).toBe(token); + running.resolve(); + await token.promise; + expect(continuation.current).toBeUndefined(); + }); + + it("does not expose an old rejection or consume replacement messages after cancellation", async () => { + const { continuation, host } = setup(); + const oldRun = deferred(); + const replacementRun = deferred(); + const message = fauxAssistantMessage("owned continuation"); + continuation.track(message); + host.continue.mockReturnValueOnce(oldRun.promise).mockReturnValueOnce(replacementRun.promise); + continuation.schedule(); + await turn(); + const oldToken = continuation.current!; + continuation.cancel(); + continuation.schedule(); + const replacement = continuation.current!; + await turn(); + oldRun.reject(new Error("old turn failed")); + await oldToken.promise; + await turn(); + expect(continuation.current).toBe(replacement); + expect(continuation.messages).toEqual([message]); + expect(host.onMessageConsumed).not.toHaveBeenCalled(); + replacementRun.resolve(); + await replacement.promise; + expect(continuation.messages).toEqual([]); + expect(host.onMessageConsumed).toHaveBeenCalledExactlyOnceWith(message); + }); + + it("rechecks a refinement that starts while waiting for the commit fence", async () => { + const { continuation, host, fence } = setup(); + const held = await fence.acquire(); + const refinement = deferred(); + continuation.schedule(); + const token = continuation.current!; + await turn(); + host.isRefinementApplying.mockReturnValue(true); + host.waitForRefinement.mockReturnValue(refinement.promise); + held.release(); + await turn(); + expect(host.continue).not.toHaveBeenCalled(); + expect(fence.hasPendingWork).toBe(false); + host.isRefinementApplying.mockReturnValue(false); + refinement.resolve(); + await token.promise; + expect(host.continue).toHaveBeenCalledOnce(); + }); + + it.each([false, true])("honors continue-after-input=%s after the input pump completes", async (shouldContinue) => { + const { continuation, host } = setup(); + host.unfinishedActionCount.mockReturnValue(1); + host.waitForIdleOrSettlement.mockImplementation(async () => { + host.unfinishedActionCount.mockReturnValue(0); + }); + continuation.schedule(shouldContinue); + await continuation.current!.promise; + expect(host.scheduleInput).toHaveBeenCalledOnce(); + expect(host.continue).toHaveBeenCalledTimes(shouldContinue ? 1 : 0); + expect(continuation.isScheduled).toBe(false); + }); + + it("retains messages still in the agent queue after a busy continuation retries", async () => { + const { continuation, host } = setup(); + const message = fauxAssistantMessage("queued continuation"); + const stillQueued: AgentMessage[] = [message]; + continuation.track(message); + host.continue.mockRejectedValueOnce(new AgentContinueError("busy", "running")); + host.removeQueuedMessages.mockImplementation((predicate) => stillQueued.filter(predicate)); + continuation.schedule(); + await continuation.current!.promise; + expect(host.continue).toHaveBeenCalledTimes(2); + expect(host.followUp).toHaveBeenCalledExactlyOnceWith(message); + expect(host.onMessageConsumed).not.toHaveBeenCalled(); + expect(continuation.messages).toEqual([message]); + }); +}); diff --git a/packages/coding-agent/test/suite/acp-features.test.ts b/packages/coding-agent/test/suite/acp-features.test.ts index 5e92bc424d..4ca41fc56c 100644 --- a/packages/coding-agent/test/suite/acp-features.test.ts +++ b/packages/coding-agent/test/suite/acp-features.test.ts @@ -426,9 +426,13 @@ describe("ACP mode preserves prime-agent features", () => { // then runs for real: it applies the proposal, persists harness state, and // emits the genuine refine_complete event this test is about. const internals = harness.session as unknown as { - _planRefine: (...args: unknown[]) => Promise; + _refinement: { + _execution: { + _planRefine: (...args: unknown[]) => Promise; + }; + }; }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ id: "acp_refine_plan", proposal: { summary: "refined for ACP", 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 5f3e911157..2182654a45 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 @@ -8,8 +8,12 @@ import { createDeferred } from "./scheduling.js"; type ActionKind = "turn" | "command"; interface CommitFenceInternals { + _refinement: { + _refineInFlight?: Promise; + }; + _actionStore: ActionStore; - _refineInFlight?: Promise; + _scheduleSessionInputPump(): void; _acquireDirectTurnAdmissionFence(signal?: AbortSignal): Promise<{ release(): void }>; _acquireSessionActionCommitFence(signal?: AbortSignal): Promise<{ release(): void }>; @@ -246,9 +250,10 @@ describe("AgentSession action commit-fence races", () => { const refineGate = createDeferred(); const refineInFlight = refineGate.promise.finally(() => { - if (internals._refineInFlight === refineInFlight) internals._refineInFlight = undefined; + if (internals._refinement._refineInFlight === refineInFlight) + internals._refinement._refineInFlight = undefined; }); - internals._refineInFlight = refineInFlight; + internals._refinement._refineInFlight = refineInFlight; heldFence.release(); await yieldToEventLoop(); diff --git a/packages/coding-agent/test/suite/agent-session-autonomous.test.ts b/packages/coding-agent/test/suite/agent-session-autonomous.test.ts index 5d110d607b..4143af04a0 100644 --- a/packages/coding-agent/test/suite/agent-session-autonomous.test.ts +++ b/packages/coding-agent/test/suite/agent-session-autonomous.test.ts @@ -495,10 +495,7 @@ describe("AgentSession autonomous mode", () => { }); harnesses.push(harness); harness.setResponses([fauxAssistantMessage("Still failing.")]); - const sessionInternals = harness.session as unknown as { - _compactionAbortController?: AbortController; - }; - sessionInternals._compactionAbortController = new AbortController(); + const compacting = vi.spyOn(harness.session, "isCompacting", "get").mockReturnValue(true); const heartbeatJob = { id: "heartbeat-test", status: "active", @@ -518,7 +515,7 @@ describe("AgentSession autonomous mode", () => { streamingBehavior: "followUp", suppressAutonomousContinuation: true, }); - sessionInternals._compactionAbortController = undefined; + compacting.mockRestore(); expect(harness.session.resumeQueuedWork()).toBe(true); await vi.waitFor(() => expect(harness.session.queuedActionCount).toBe(0)); await harness.session.waitForSessionInputIdle(); diff --git a/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts b/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts index b1b1425bc1..92af80fc73 100644 --- a/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts @@ -1,6 +1,7 @@ import type { ShouldStopAfterTurnContext } from "@earendil-works/pi-agent-core"; import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; +import type { SessionCompaction } from "../../src/session/compaction.js"; import { createHarness, type Harness } from "./harness.js"; type SessionInternals = { @@ -212,7 +213,7 @@ describe("AgentSession compact skill host requests", () => { harnesses.push(harness); await harness.session.prompt("one"); - (harness.session as unknown as { _pendingRequestedCompaction?: object })._pendingRequestedCompaction = {}; + (harness.session as unknown as { _compaction: SessionCompaction })._compaction.request(); harness.session.agent.followUp({ role: "custom", customType: "test", diff --git a/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts b/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts index 06c9a215e4..a65438017c 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts @@ -15,6 +15,7 @@ import { import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { AgentSession } from "../../src/core/agent-session.js"; +import type { SessionCompaction } from "../../src/session/compaction.js"; import { createHarness, type Harness } from "./harness.js"; type SessionInternals = { @@ -27,7 +28,7 @@ type SessionInternals = { customInstructions?: string; signal: AbortSignal; }) => Promise; - _continueAfterThresholdCompaction: boolean; + _compaction: Pick & { readonly continueAfterThreshold: boolean }; }; function createUsage(totalTokens: number): Usage { @@ -128,7 +129,7 @@ describe("compaction continuation", () => { // toolResult-last makes the session stop the loop for compaction AND continue afterwards. const shouldStop = await internals._shouldStopAfterTurn(context); expect(shouldStop).toBe(true); - expect(internals._continueAfterThresholdCompaction).toBe(true); + expect(internals._compaction.continueAfterThreshold).toBe(true); const continueSpy = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); @@ -152,7 +153,7 @@ describe("compaction continuation", () => { harnesses.push(harness); const internals = harness.session as unknown as SessionInternals; midToolLoopContext(harness); - internals._continueAfterThresholdCompaction = true; + internals._compaction.requestContinuation(); const continueSpy = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); @@ -335,7 +336,7 @@ describe("compaction continuation", () => { const shouldStop = await internals._shouldStopAfterTurn(context); expect(shouldStop).toBe(true); - expect(internals._continueAfterThresholdCompaction).toBe(true); + expect(internals._compaction.continueAfterThreshold).toBe(true); expect(harness.session.queuedActionCount).toBe(1); expect(harness.session.goalState.continuationsUsed).toBe(1); 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 9eb28a649e..3d6d6cef83 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -11,6 +11,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { convertToLlm } from "../../src/core/messages.js"; import { getLocalHarnessStateDir, loadHarnessState, saveHarnessState } from "../../src/core/refinement/index.js"; import { SessionManager } from "../../src/core/session-manager.js"; +import type { SessionCompaction } from "../../src/session/compaction.js"; +import type { SessionContinuation } from "../../src/session/continuation.js"; import { createHarness, getMessageText, type Harness } from "./harness.js"; import { createDeferred } from "./scheduling.js"; @@ -22,11 +24,7 @@ type SessionWithCompactionInternals = { ) => Promise; _runAutoCompaction: (reason: "overflow" | "threshold" | "requested", willRetry: boolean) => Promise; _shouldStopAfterTurn: (context: ShouldStopAfterTurnContext) => boolean | Promise; - _persistCompactionOutcome: ( - reason: "overflow" | "threshold" | "requested", - outcome: "skipped" | "cancelled" | "failed", - message: string, - ) => void; + _compaction: SessionCompaction; }; function createUsage(totalTokens: number) { @@ -351,7 +349,7 @@ describe("AgentSession compaction characterization", () => { const internals = harness.session as unknown as { _schedulePostCompactionContinue(): void; _cancelPostCompactionContinue(): void; - _postCompactionContinuationScheduled: boolean; + _continuation: SessionContinuation; }; try { await harness.session.prompt("one"); @@ -360,7 +358,7 @@ describe("AgentSession compaction characterization", () => { await harness.session.compact(); - expect(internals._postCompactionContinuationScheduled).toBe(true); + expect(internals._continuation.isScheduled).toBe(true); } finally { internals._cancelPostCompactionContinue(); } @@ -474,10 +472,15 @@ describe("AgentSession compaction characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as { + _refinement: { + _auto: { + _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; + }; + }; + _cancelPostCompactionContinue(): void; - _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; }; - const scheduleAutoRefineSpy = vi.spyOn(internals, "_scheduleAutoRefineAfterCompaction"); + const scheduleAutoRefineSpy = vi.spyOn(internals._refinement._auto, "_scheduleAutoRefineAfterCompaction"); try { await harness.session.prompt("one"); await harness.session.prompt("two"); @@ -578,10 +581,15 @@ describe("AgentSession compaction characterization", () => { await harness.session.followUp("preparing across compaction", undefined, { resumeIfIdle: true }); await vi.waitFor(() => expect(preparationReached).toHaveBeenCalledOnce()); const internals = harness.session as unknown as SessionWithCompactionInternals & { + _refinement: { + _auto: { + _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; + }; + }; + _cancelPostCompactionContinue(): void; - _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; }; - const scheduleAutoRefineSpy = vi.spyOn(internals, "_scheduleAutoRefineAfterCompaction"); + const scheduleAutoRefineSpy = vi.spyOn(internals._refinement._auto, "_scheduleAutoRefineAfterCompaction"); try { await internals._runAutoCompaction("requested", false); expect(scheduleAutoRefineSpy).toHaveBeenCalledWith(true); @@ -689,7 +697,7 @@ describe("AgentSession compaction characterization", () => { shouldContinueAfterThreshold: boolean, queuedMessages: AgentMessage[], ): void; - _postCompactionContinuationMessages: AgentMessage[]; + _continuation: SessionContinuation; }; const firstAssistant = createAssistant(harness, { stopReason: "toolUse", totalTokens: 10_000 }); const secondAssistant = createAssistant(harness, { stopReason: "toolUse", totalTokens: 10_000 }); @@ -703,7 +711,7 @@ describe("AgentSession compaction characterization", () => { sessionInternals._clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction(true, [secondQueued!]); expect(harness.session.getAutonomousStatus().continuationsUsed).toBe(1); - expect(sessionInternals._postCompactionContinuationMessages).toEqual([firstQueued]); + expect(sessionInternals._continuation.messages).toEqual([firstQueued]); expect(harness.session.getFollowUpMessages()).toHaveLength(1); }); @@ -1066,7 +1074,7 @@ describe("AgentSession compaction characterization", () => { harnesses.push(harness); const sessionInternals = harness.session as unknown as { _schedulePostCompactionContinue(): void; - _postCompactionContinuationMessages: AgentMessage[]; + _continuation: SessionContinuation; }; const steeringMessage = { role: "user", @@ -1078,7 +1086,7 @@ describe("AgentSession compaction characterization", () => { content: [{ type: "text", text: "autonomous follow-up" }], timestamp: Date.now(), } satisfies AgentMessage; - sessionInternals._postCompactionContinuationMessages = [autonomousMessage]; + sessionInternals._continuation.track(autonomousMessage); harness.session.agent.state.messages = [{ ...fauxAssistantMessage("done"), timestamp: Date.now() - 1000 }]; harness.session.agent.steer(steeringMessage); harness.session.agent.followUp(autonomousMessage); @@ -1089,7 +1097,7 @@ describe("AgentSession compaction characterization", () => { await vi.advanceTimersByTimeAsync(100); expect(continueSpy).toHaveBeenCalledTimes(1); - expect(sessionInternals._postCompactionContinuationMessages).toEqual([autonomousMessage]); + expect(sessionInternals._continuation.messages).toEqual([autonomousMessage]); expect(followUpSpy).toHaveBeenCalledWith(autonomousMessage); }); @@ -1119,8 +1127,7 @@ describe("AgentSession compaction characterization", () => { harnesses.push(harness); const sessionInternals = harness.session as unknown as { _schedulePostCompactionContinue(continueAfterSessionInput?: boolean): void; - _postCompactionContinuationMessages: AgentMessage[]; - _postCompactionContinuationScheduled: boolean; + _continuation: SessionContinuation; _createPreparedTurnAction( schedule: "followUp", text: string, @@ -1134,7 +1141,7 @@ describe("AgentSession compaction characterization", () => { content: [{ type: "text", text }], timestamp: Date.now(), } satisfies AgentMessage; - if (tracked) sessionInternals._postCompactionContinuationMessages = [continuation]; + if (tracked) sessionInternals._continuation.track(continuation); harness.setResponses([fauxAssistantMessage(response)]); sessionInternals._admitSessionInput( sessionInternals._createPreparedTurnAction("followUp", text, undefined, { @@ -1148,8 +1155,8 @@ describe("AgentSession compaction characterization", () => { await vi.advanceTimersByTimeAsync(200); expect(continueSpy).toHaveBeenCalledTimes(continueAfterSessionInput ? 1 : 0); - expect(sessionInternals._postCompactionContinuationScheduled).toBe(false); - expect(sessionInternals._postCompactionContinuationMessages).toEqual([]); + expect(sessionInternals._continuation.isScheduled).toBe(false); + expect(sessionInternals._continuation.messages).toEqual([]); expect(harness.session.messages.at(-1)).toMatchObject({ role: "assistant", content: [{ type: "text", text: response }], @@ -1171,15 +1178,14 @@ describe("AgentSession compaction characterization", () => { const sessionInternals = harness.session as unknown as { _schedulePostCompactionContinue(): void; _cancelPostCompactionContinue(): void; - _postCompactionContinuationMessages: AgentMessage[]; - _postCompactionContinuationScheduled: boolean; + _continuation: SessionContinuation; }; const queuedMessage = { role: "user", content: [{ type: "text", text: "autonomous follow-up" }], timestamp: Date.now(), } satisfies AgentMessage; - sessionInternals._postCompactionContinuationMessages = [queuedMessage]; + sessionInternals._continuation.track(queuedMessage); harness.session.agent.state.messages = [ { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 }, ]; @@ -1194,8 +1200,8 @@ describe("AgentSession compaction characterization", () => { sessionInternals._schedulePostCompactionContinue(); await vi.waitFor(() => expect(continueSpy).toHaveBeenCalledTimes(1)); - expect(sessionInternals._postCompactionContinuationMessages).toEqual([queuedMessage]); - expect(sessionInternals._postCompactionContinuationScheduled).toBe(true); + expect(sessionInternals._continuation.messages).toEqual([queuedMessage]); + expect(sessionInternals._continuation.isScheduled).toBe(true); sessionInternals._cancelPostCompactionContinue(); activeRunSettled.resolve(); }); @@ -1206,14 +1212,14 @@ describe("AgentSession compaction characterization", () => { const sessionInternals = harness.session as unknown as { _schedulePostCompactionContinue(): void; _cancelPostCompactionContinue(): void; - _postCompactionContinuationMessages: AgentMessage[]; + _continuation: SessionContinuation; }; const queuedMessage = { role: "user", content: [{ type: "text", text: "autonomous follow-up" }], timestamp: Date.now(), } satisfies AgentMessage; - sessionInternals._postCompactionContinuationMessages = [queuedMessage]; + sessionInternals._continuation.track(queuedMessage); const staleRun = createDeferred(); const replacementRun = createDeferred(); const continueSpy = vi @@ -1229,19 +1235,22 @@ describe("AgentSession compaction characterization", () => { staleRun.resolve(); await new Promise(setImmediate); - expect(sessionInternals._postCompactionContinuationMessages).toEqual([queuedMessage]); + expect(sessionInternals._continuation.messages).toEqual([queuedMessage]); replacementRun.resolve(); await harness.session.waitForHeadlessIdle(); - expect(sessionInternals._postCompactionContinuationMessages).toEqual([]); + expect(sessionInternals._continuation.messages).toEqual([]); }); it("waits for an in-flight refine application before continuing", async () => { const harness = await createHarness(); harnesses.push(harness); const sessionInternals = harness.session as unknown as { + _refinement: { + _refineInFlight: Promise | undefined; + }; + _schedulePostCompactionContinue(): void; - _refineInFlight: Promise | undefined; }; const continueSpy = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); const pause = harness.session.acquireQueuedWorkPause(); @@ -1250,12 +1259,12 @@ describe("AgentSession compaction characterization", () => { // Refine enters its apply phase while the runner waits out the pause. const refineApply = createDeferred(); - sessionInternals._refineInFlight = refineApply.promise; + sessionInternals._refinement._refineInFlight = refineApply.promise; pause.release(); await new Promise(setImmediate); expect(continueSpy).not.toHaveBeenCalled(); - sessionInternals._refineInFlight = undefined; + sessionInternals._refinement._refineInFlight = undefined; refineApply.resolve(); await harness.session.waitForHeadlessIdle(); expect(continueSpy).toHaveBeenCalledTimes(1); @@ -1569,7 +1578,7 @@ describe("AgentSession compaction characterization", () => { }); expect(() => - internals._persistCompactionOutcome("requested", "failed", "Requested compaction failed"), + internals._compaction.endUnsuccessfully("requested", "failed", "Requested compaction failed"), ).not.toThrow(); // The live outcome message discloses that it was not saved. expect(harness.session.messages.at(-1)).toMatchObject({ @@ -1641,7 +1650,7 @@ describe("AgentSession compaction characterization", () => { vi.spyOn(harness.sessionManager, "_persist").mockImplementationOnce(() => { throw new Error("disk full"); }); - internals._persistCompactionOutcome("requested", "failed", "Requested compaction failed"); + internals._compaction.endUnsuccessfully("requested", "failed", "Requested compaction failed"); // Compaction reloads agent.state.messages from the session file; the // memory-only disclosure must survive. 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 88dc4550d7..bfb656c3fd 100644 --- a/packages/coding-agent/test/suite/agent-session-prompt.test.ts +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -575,17 +575,25 @@ stale extension instructions`, }); harnesses.push(harness); const internals = harness.session as unknown as { + _refinement: { + _refineInFlight?: Promise; + _refineAbortController?: AbortController; + _execution: { + _planRefine(options: unknown, signal: AbortSignal): Promise; + _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; + }; + }; + _baseSystemPrompt: string; - _refineInFlight?: Promise; - _refineAbortController?: AbortController; - _planRefine(options: unknown, signal: AbortSignal): Promise; - _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "race-plan", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "race-plan", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockImplementation(async () => { internals._baseSystemPrompt = "refined $& base"; harness.session.agent.state.systemPrompt = "refined $& base"; - internals._refineAbortController = undefined; + internals._refinement._refineAbortController = undefined; return { id: "refine_race", summary: "refined", @@ -607,7 +615,7 @@ stale extension instructions`, const promptPromise = harness.session.prompt("normal prompt"); await hookStarted; await harness.session.refine({ instructions: "complete while the extension hook is suspended" }); - expect(internals._refineInFlight).toBeUndefined(); + expect(internals._refinement._refineInFlight).toBeUndefined(); releaseHook(); await promptPromise; @@ -640,16 +648,24 @@ stale extension instructions`, }); harnesses.push(harness); const internals = harness.session as unknown as { + _refinement: { + _refineAbortController?: AbortController; + _execution: { + _planRefine(options: unknown, signal: AbortSignal): Promise; + _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; + }; + }; + _baseSystemPrompt: string; - _refineAbortController?: AbortController; - _planRefine(options: unknown, signal: AbortSignal): Promise; - _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "race-plan", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "race-plan", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockImplementation(async () => { internals._baseSystemPrompt = "refined base"; harness.session.agent.state.systemPrompt = "refined base"; - internals._refineAbortController = undefined; + internals._refinement._refineAbortController = undefined; return { id: "refine_independent", summary: "refined", @@ -709,11 +725,17 @@ stale injected extension instructions`, }); harnesses.push(harness); const internals = harness.session as unknown as { + _refinement: { + _refineInFlight?: Promise; + _refineAbortController?: AbortController; + _execution: { + _planRefine(options: unknown, signal: AbortSignal): Promise; + _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; + }; + }; + _baseSystemPrompt: string; - _refineInFlight?: Promise; - _refineAbortController?: AbortController; - _planRefine(options: unknown, signal: AbortSignal): Promise; - _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; + _promptInjectedMessage( text: string, message: { @@ -726,11 +748,14 @@ stale injected extension instructions`, }, ): Promise; }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "race-plan", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "race-plan", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockImplementation(async () => { internals._baseSystemPrompt = "refined injected base"; harness.session.agent.state.systemPrompt = "refined injected base"; - internals._refineAbortController = undefined; + internals._refinement._refineAbortController = undefined; return { id: "refine_race", summary: "refined", @@ -759,7 +784,7 @@ stale injected extension instructions`, }); await hookStarted; await harness.session.refine({ instructions: "complete while the injected hook is suspended" }); - expect(internals._refineInFlight).toBeUndefined(); + expect(internals._refinement._refineInFlight).toBeUndefined(); releaseHook(); await injectedPrompt; @@ -809,7 +834,10 @@ stale injected extension instructions`, it("preserves an empty extension system prompt across an injected refine handoff wait", async () => { let sessionInternals: { - _refineInFlight?: Promise; + _refinement: { + _refineInFlight?: Promise; + }; + _promptInjectedMessage( text: string, message: { @@ -828,11 +856,11 @@ stale injected extension instructions`, (pi) => { pi.on("before_agent_start", async () => { let releaseRefine: (() => void) | undefined; - sessionInternals._refineInFlight = new Promise((resolve) => { + sessionInternals._refinement._refineInFlight = new Promise((resolve) => { releaseRefine = resolve; }); setTimeout(() => { - sessionInternals._refineInFlight = undefined; + sessionInternals._refinement._refineInFlight = undefined; releaseRefine?.(); }, 0); return { systemPrompt: "" }; @@ -863,18 +891,18 @@ stale injected extension instructions`, }); it("preserves an extension system prompt across a refine handoff wait", async () => { - let sessionInternals: { _refineInFlight?: Promise }; + let sessionInternals: { _refinement: { _refineInFlight?: Promise } }; const harness = await createHarness({ systemPrompt: "base prompt", extensionFactories: [ (pi) => { pi.on("before_agent_start", async (event) => { let releaseRefine: (() => void) | undefined; - sessionInternals._refineInFlight = new Promise((resolve) => { + sessionInternals._refinement._refineInFlight = new Promise((resolve) => { releaseRefine = resolve; }); setTimeout(() => { - sessionInternals._refineInFlight = undefined; + sessionInternals._refinement._refineInFlight = undefined; releaseRefine?.(); }, 0); return { @@ -887,7 +915,7 @@ extension instructions`, ], }); harnesses.push(harness); - sessionInternals = harness.session as unknown as { _refineInFlight?: Promise }; + sessionInternals = harness.session as unknown as { _refinement: { _refineInFlight?: Promise } }; let providerSystemPrompt = ""; harness.setResponses([ (context) => { @@ -908,11 +936,16 @@ extension instructions`, // The post-wait guard must detect the base change and discard the // stale extension prompt, using the refined base instead. let sessionInternals: { + _refinement: { + _refineInFlight?: Promise; + _refineAbortController?: AbortController; + _execution: { + _planRefine(options: unknown, signal: AbortSignal): Promise; + _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; + }; + }; + _baseSystemPrompt: string; - _refineInFlight?: Promise; - _refineAbortController?: AbortController; - _planRefine(options: unknown, signal: AbortSignal): Promise; - _applyRefine(plan: unknown, options: unknown, abort: AbortController): Promise; }; let releaseRefine: (() => void) | undefined; const harness = await createHarness({ @@ -922,7 +955,7 @@ extension instructions`, (pi) => { pi.on("before_agent_start", async (event) => { // Set up _refineInFlight so the post-hook wait triggers. - sessionInternals._refineInFlight = new Promise((resolve) => { + sessionInternals._refinement._refineInFlight = new Promise((resolve) => { releaseRefine = resolve; }); return { @@ -936,11 +969,14 @@ stale post-hook extension instructions`, }); harnesses.push(harness); sessionInternals = harness.session as unknown as typeof sessionInternals; - vi.spyOn(sessionInternals, "_planRefine").mockResolvedValue({ id: "race-plan", proposal: { edits: [] } }); - vi.spyOn(sessionInternals, "_applyRefine").mockImplementation(async () => { + vi.spyOn(sessionInternals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "race-plan", + proposal: { edits: [] }, + }); + vi.spyOn(sessionInternals._refinement._execution, "_applyRefine").mockImplementation(async () => { sessionInternals._baseSystemPrompt = "refined post-hook base"; harness.session.agent.state.systemPrompt = "refined post-hook base"; - sessionInternals._refineAbortController = undefined; + sessionInternals._refinement._refineAbortController = undefined; return { id: "refine_post_hook", summary: "refined", @@ -962,12 +998,12 @@ stale post-hook extension instructions`, // extension prompt. The prompt path enters _waitForRefineIdle. // Wait for the hook to fire and set _refineInFlight. await vi.waitFor(() => { - expect(sessionInternals._refineInFlight).toBeDefined(); + expect(sessionInternals._refinement._refineInFlight).toBeDefined(); }); // Complete the refine during the wait: change the base and resolve. sessionInternals._baseSystemPrompt = "refined post-hook base"; releaseRefine?.(); - sessionInternals._refineInFlight = undefined; + sessionInternals._refinement._refineInFlight = undefined; await promptPromise; expect(providerSystemPrompt).toContain("refined post-hook base"); @@ -983,10 +1019,7 @@ stale post-hook extension instructions`, { customType: "next-turn", content: "queued context", display: true, details: {} }, { deliverAs: "nextTurn" }, ); - const sessionInternals = harness.session as unknown as { - _compactionAbortController?: AbortController; - }; - sessionInternals._compactionAbortController = new AbortController(); + const compacting = vi.spyOn(harness.session, "isCompacting", "get").mockReturnValue(true); await harness.session.acceptAgentMessagePrompt(agentPrompt, { expandPromptTemplates: false, @@ -995,7 +1028,7 @@ stale post-hook extension instructions`, }); expect(harness.session.getFollowUpMessages()).toEqual([agentPrompt]); - sessionInternals._compactionAbortController = undefined; + compacting.mockRestore(); let queuedTurnSawSeparateNextTurnContext = false; harness.setResponses([ fauxAssistantMessage("first turn"), @@ -1066,9 +1099,11 @@ stale post-hook extension instructions`, releaseRefine = resolve; }); const sessionInternals = harness.session as unknown as { - _refineInFlight?: Promise; + _refinement: { + _refineInFlight?: Promise; + }; }; - sessionInternals._refineInFlight = refineGate; + sessionInternals._refinement._refineInFlight = refineGate; const accepted = harness.session.acceptAgentMessagePrompt( "Agent-to-agent message received.\nSource: agent_message\nTo: Target, active target, session session-target\nMessage id: agentmsg_handoff_reject\n\nagent text", @@ -1086,7 +1121,7 @@ stale post-hook extension instructions`, }, }); expect(harness.session.isBashRunning).toBe(true); - sessionInternals._refineInFlight = undefined; + sessionInternals._refinement._refineInFlight = undefined; releaseRefine?.(); try { @@ -1461,10 +1496,7 @@ stale post-hook extension instructions`, it("does not run built-in slash commands immediately while queueIfBusy backpressure is active", async () => { const harness = await createHarness(); harnesses.push(harness); - const sessionInternals = harness.session as unknown as { - _compactionAbortController?: AbortController; - }; - sessionInternals._compactionAbortController = new AbortController(); + const compacting = vi.spyOn(harness.session, "isCompacting", "get").mockReturnValue(true); await harness.session.prompt("/autonomous on", { queueIfBusy: true, @@ -1473,7 +1505,7 @@ stale post-hook extension instructions`, expect(harness.session.getAutonomousStatus().enabled).toBe(false); expect(harness.session.getFollowUpMessages()).toEqual(["/autonomous on"]); - sessionInternals._compactionAbortController = undefined; + compacting.mockRestore(); expect(harness.session.resumeQueuedWork()).toBe(true); await harness.session.waitForSessionInputIdle(); expect(harness.session.getAutonomousStatus().enabled).toBe(true); 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 44b70ebe01..b50305a380 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/packages/coding-agent/test/suite/agent-session-queue.test.ts @@ -29,6 +29,7 @@ import { saveHarnessState, } from "../../src/core/refinement/index.js"; import { parseSessionSlashCommand } from "../../src/core/slash-commands.js"; +import type { SessionContinuation } from "../../src/session/continuation.js"; import { conversationMessages, createHarness, @@ -42,21 +43,28 @@ import { createDeferred, createWaitingHarness, gatedHook, withStreaming } from " type AutoRefineReason = "turn_interval" | "compact"; type AutoRefineInternals = { - _maybeAutoRefine(reason: AutoRefineReason): Promise; - _scheduleAutoRefine(reason: AutoRefineReason): void; - _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; - _scheduleAutoRefineAfterAgentEnd(): void; + _refinement: { + _invalidatePendingAutoRefineForBranchChange(): Promise; + _auto: { + _maybeAutoRefine(reason: AutoRefineReason): Promise; + _scheduleAutoRefine(reason: AutoRefineReason): void; + _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; + _scheduleAutoRefineAfterAgentEnd(): void; + _assistantTurnsSinceAutoRefine: number; + _lastAutoRefineReviewAt: number; + _compactAutoRefinePending: boolean; + _turnIntervalAutoRefinePending: boolean; + _pendingAutoRefineReview?: unknown; + _autoRefineInProgress: boolean; + _autoRefineBranchVersion: number; + }; + }; + _schedulePostCompactionContinue(continueAfterSessionInput?: boolean): void; - _invalidatePendingAutoRefineForBranchChange(): Promise; + _cancelPostCompactionContinue(): void; - _assistantTurnsSinceAutoRefine: number; - _lastAutoRefineReviewAt: number; - _compactAutoRefinePending: boolean; - _turnIntervalAutoRefinePending: boolean; - _postCompactionContinuationScheduled: boolean; - _pendingAutoRefineReview?: unknown; - _autoRefineInProgress: boolean; - _autoRefineBranchVersion: number; + + _continuation: Pick; }; type SteeringStopInternals = { @@ -131,7 +139,7 @@ describe("AgentSession queue characterization", () => { await harness.session.prompt("fail once"); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); it.each([ @@ -225,13 +233,20 @@ describe("AgentSession queue characterization", () => { const reviewer = vi.fn(async () => review); const harness = await createAutoRefineHarness({ settings, autoRefineReviewer: reviewer }); harnesses.push(harness); - const refine = vi.spyOn(harness.session, "refine").mockResolvedValue(emptyRefinementResult()); + const refine = vi + .spyOn( + (harness.session as unknown as { _refinement: Pick })._refinement, + "refine", + ) + .mockResolvedValue(emptyRefinementResult()); const internals = harness.session as unknown as AutoRefineInternals; - internals._assistantTurnsSinceAutoRefine = turns; - const scheduleAutoRefine = vi.spyOn(internals, "_scheduleAutoRefine").mockImplementation(() => {}); + internals._refinement._auto._assistantTurnsSinceAutoRefine = turns; + const scheduleAutoRefine = vi + .spyOn(internals._refinement._auto, "_scheduleAutoRefine") + .mockImplementation(() => {}); if (queuedMessages) vi.spyOn(harness.session.agent, "hasQueuedMessages").mockReturnValue(true); - await internals._maybeAutoRefine(reason); + await internals._refinement._auto._maybeAutoRefine(reason); if (expectedReviewContext !== undefined) { expect(reviewer).toHaveBeenCalledWith(expectedReviewContext, expect.any(AbortSignal)); @@ -247,8 +262,10 @@ describe("AgentSession queue characterization", () => { ); } } - if (turnsAfter !== undefined) expect(internals._assistantTurnsSinceAutoRefine).toBe(turnsAfter); - if (compactPendingAfter !== undefined) expect(internals._compactAutoRefinePending).toBe(compactPendingAfter); + if (turnsAfter !== undefined) + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(turnsAfter); + if (compactPendingAfter !== undefined) + expect(internals._refinement._auto._compactAutoRefinePending).toBe(compactPendingAfter); if (scheduleCalledWith !== undefined) expect(scheduleAutoRefine).toHaveBeenCalledWith(scheduleCalledWith); }, ); @@ -257,29 +274,29 @@ describe("AgentSession queue characterization", () => { { name: "waits for planned post-compaction continuation", act: (internals: AutoRefineInternals, expectSchedule: (called: boolean) => void) => { - internals._scheduleAutoRefineAfterCompaction(true); - expect(internals._compactAutoRefinePending).toBe(true); + internals._refinement._auto._scheduleAutoRefineAfterCompaction(true); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(true); expectSchedule(false); - internals._scheduleAutoRefineAfterAgentEnd(); - expect(internals._compactAutoRefinePending).toBe(true); + internals._refinement._auto._scheduleAutoRefineAfterAgentEnd(); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(true); }, }, { name: "waits until the scheduled post-compaction continuation starts", act: (internals: AutoRefineInternals, expectSchedule: (called: boolean) => void) => { - internals._compactAutoRefinePending = true; - internals._postCompactionContinuationScheduled = true; - internals._scheduleAutoRefineAfterAgentEnd(); + internals._refinement._auto._compactAutoRefinePending = true; + const scheduled = vi.spyOn(internals._continuation, "isScheduled", "get").mockReturnValue(true); + internals._refinement._auto._scheduleAutoRefineAfterAgentEnd(); expectSchedule(false); - internals._postCompactionContinuationScheduled = false; - internals._scheduleAutoRefineAfterAgentEnd(); + scheduled.mockRestore(); + internals._refinement._auto._scheduleAutoRefineAfterAgentEnd(); }, }, { name: "runs immediately when no post-compaction continuation is planned", act: (internals: AutoRefineInternals) => { - internals._scheduleAutoRefineAfterCompaction(false); - expect(internals._compactAutoRefinePending).toBe(false); + internals._refinement._auto._scheduleAutoRefineAfterCompaction(false); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); }, }, ])("auto-refine compact hook $name", async ({ act }) => { @@ -288,7 +305,9 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - const scheduleAutoRefine = vi.spyOn(internals, "_scheduleAutoRefine").mockImplementation(() => {}); + const scheduleAutoRefine = vi + .spyOn(internals._refinement._auto, "_scheduleAutoRefine") + .mockImplementation(() => {}); act(internals, (called) => called ? expect(scheduleAutoRefine).toHaveBeenCalled() : expect(scheduleAutoRefine).not.toHaveBeenCalled(), @@ -313,22 +332,22 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._assistantTurnsSinceAutoRefine = 2; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 2; try { - const compactReview = internals._maybeAutoRefine("compact"); + const compactReview = internals._refinement._auto._maybeAutoRefine("compact"); await Promise.resolve(); - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); - expect(internals._turnIntervalAutoRefinePending).toBe(true); + expect(internals._refinement._auto._turnIntervalAutoRefinePending).toBe(true); compactReviewGate.resolve(); await compactReview; await vi.runOnlyPendingTimersAsync(); expect(reviewer.mock.calls.map(([context]) => context.reason)).toEqual(["compact", "turn_interval"]); - expect(internals._turnIntervalAutoRefinePending).toBe(false); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._turnIntervalAutoRefinePending).toBe(false); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); } finally { vi.useRealTimers(); } @@ -353,11 +372,11 @@ describe("AgentSession queue characterization", () => { internals._schedulePostCompactionContinue(); await vi.waitFor(() => expect(continueAgent).toHaveBeenCalledTimes(1)); - expect(internals._postCompactionContinuationScheduled).toBe(true); + expect(internals._continuation.isScheduled).toBe(true); activeRunSettled.resolve(); await vi.waitFor(() => expect(continueAgent).toHaveBeenCalledTimes(2)); - expect(internals._postCompactionContinuationScheduled).toBe(false); + expect(internals._continuation.isScheduled).toBe(false); }); it("does not let a failed cancelled continuation reject its replacement", async () => { @@ -412,11 +431,11 @@ describe("AgentSession queue characterization", () => { try { internals._schedulePostCompactionContinue(); - await internals._invalidatePendingAutoRefineForBranchChange(); + await internals._refinement._invalidatePendingAutoRefineForBranchChange(); await vi.advanceTimersByTimeAsync(100); expect(continueAgent).not.toHaveBeenCalled(); - expect(internals._postCompactionContinuationScheduled).toBe(false); + expect(internals._continuation.isScheduled).toBe(false); } finally { vi.useRealTimers(); } @@ -443,7 +462,7 @@ describe("AgentSession queue characterization", () => { await vi.advanceTimersByTimeAsync(100); expect(continueAgent).not.toHaveBeenCalled(); - expect(internals._postCompactionContinuationScheduled).toBe(false); + expect(internals._continuation.isScheduled).toBe(false); expect(harness.session.getFollowUpMessages()).toEqual(["queued across abort"]); } finally { vi.useRealTimers(); @@ -464,7 +483,7 @@ describe("AgentSession queue characterization", () => { "Session is too short to compact", ); - expect(internals._postCompactionContinuationScheduled).toBe(true); + expect(internals._continuation.isScheduled).toBe(true); internals._cancelPostCompactionContinue(); idle.resolve(); }); @@ -475,39 +494,44 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._pendingAutoRefineReview = { + internals._refinement._auto._pendingAutoRefineReview = { reason: "turn_interval", review: { shouldRefine: true, rationale: "durable lesson" }, }; let guardWasSetDuringRefine = false; - const refine = vi.spyOn(harness.session, "refine").mockImplementation(async () => { - guardWasSetDuringRefine = internals._autoRefineInProgress; - throw new Error("refine failed"); - }); + const refine = vi + .spyOn( + (harness.session as unknown as { _refinement: Pick })._refinement, + "refine", + ) + .mockImplementation(async () => { + guardWasSetDuringRefine = internals._refinement._auto._autoRefineInProgress; + throw new Error("refine failed"); + }); - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(refine).toHaveBeenCalledWith( expect.objectContaining({ instructions: expect.stringContaining("durable lesson") }), { trigger: "auto" }, ); expect(guardWasSetDuringRefine).toBe(true); - expect(internals._autoRefineInProgress).toBe(false); - expect(internals._pendingAutoRefineReview).toBeDefined(); + expect(internals._refinement._auto._autoRefineInProgress).toBe(false); + expect(internals._refinement._auto._pendingAutoRefineReview).toBeDefined(); // The failure stamps the cooldown so the retained pending review does not // retry on every agent end. - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); refine.mockResolvedValueOnce(emptyRefinementResult()); - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(refine).toHaveBeenCalledTimes(1); - expect(internals._pendingAutoRefineReview).toBeDefined(); + expect(internals._refinement._auto._pendingAutoRefineReview).toBeDefined(); - internals._lastAutoRefineReviewAt = 0; - await internals._maybeAutoRefine("turn_interval"); + internals._refinement._auto._lastAutoRefineReviewAt = 0; + await internals._refinement._auto._maybeAutoRefine("turn_interval"); - expect(internals._pendingAutoRefineReview).toBeUndefined(); + expect(internals._refinement._auto._pendingAutoRefineReview).toBeUndefined(); }); it("keeps the turn counter and stamps the cooldown when an approved immediate refine fails", async () => { @@ -518,17 +542,20 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._assistantTurnsSinceAutoRefine = 2; - vi.spyOn(harness.session, "refine").mockRejectedValueOnce(new Error("refine failed")); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 2; + vi.spyOn( + (harness.session as unknown as { _refinement: Pick })._refinement, + "refine", + ).mockRejectedValueOnce(new Error("refine failed")); - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(reviewer).toHaveBeenCalledWith( { reason: "turn_interval", turnsSinceLastReview: 2 }, expect.any(AbortSignal), ); - expect(internals._assistantTurnsSinceAutoRefine).toBe(2); - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(2); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); }); it("does not refine when a review resolves after the session is disposed", async () => { @@ -547,10 +574,15 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._assistantTurnsSinceAutoRefine = 1; - const refine = vi.spyOn(harness.session, "refine").mockResolvedValue(emptyRefinementResult()); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + const refine = vi + .spyOn( + (harness.session as unknown as { _refinement: Pick })._refinement, + "refine", + ) + .mockResolvedValue(emptyRefinementResult()); - const autoRefinePromise = internals._maybeAutoRefine("turn_interval"); + const autoRefinePromise = internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(reviewer).toHaveBeenCalledTimes(1); const entriesBeforeDispose = harness.sessionManager.getEntries().length; harness.session.dispose(); @@ -559,11 +591,11 @@ describe("AgentSession queue characterization", () => { await autoRefinePromise; expect(refine).not.toHaveBeenCalled(); - expect(internals._pendingAutoRefineReview).toBeUndefined(); + expect(internals._refinement._auto._pendingAutoRefineReview).toBeUndefined(); expect(harness.sessionManager.getEntries().length).toBe(entriesBeforeDispose); // Disposal also invalidates any newly scheduled auto-refine. - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(reviewer).toHaveBeenCalledTimes(1); }); @@ -577,14 +609,14 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(reviewer).toHaveBeenCalledTimes(1); - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(reviewer).toHaveBeenCalledTimes(1); }); @@ -595,17 +627,22 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._pendingAutoRefineReview = { + internals._refinement._auto._pendingAutoRefineReview = { reason: "turn_interval", review: { shouldRefine: true, rationale: "durable lesson" }, }; - internals._lastAutoRefineReviewAt = Date.now(); - const refine = vi.spyOn(harness.session, "refine").mockResolvedValue(emptyRefinementResult()); + internals._refinement._auto._lastAutoRefineReviewAt = Date.now(); + const refine = vi + .spyOn( + (harness.session as unknown as { _refinement: Pick })._refinement, + "refine", + ) + .mockResolvedValue(emptyRefinementResult()); - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(refine).not.toHaveBeenCalled(); - expect(internals._pendingAutoRefineReview).toBeDefined(); + expect(internals._refinement._auto._pendingAutoRefineReview).toBeDefined(); }); it("serializes concurrent refine calls", async () => { @@ -838,13 +875,20 @@ describe("AgentSession queue characterization", () => { const harness = await makeHarness(); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._assistantTurnsSinceAutoRefine = 1; - const refine = vi.spyOn(harness.session, "refine").mockResolvedValue(emptyRefinementResult()); - const scheduleAutoRefine = vi.spyOn(internals, "_scheduleAutoRefine").mockImplementation(() => {}); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + const refine = vi + .spyOn( + (harness.session as unknown as { _refinement: Pick })._refinement, + "refine", + ) + .mockResolvedValue(emptyRefinementResult()); + const scheduleAutoRefine = vi + .spyOn(internals._refinement._auto, "_scheduleAutoRefine") + .mockImplementation(() => {}); - await internals._maybeAutoRefine("turn_interval"); - internals._scheduleAutoRefineAfterCompaction(false); - internals._scheduleAutoRefineAfterAgentEnd(); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); + internals._refinement._auto._scheduleAutoRefineAfterCompaction(false); + internals._refinement._auto._scheduleAutoRefineAfterAgentEnd(); expect(skipReviewer).not.toHaveBeenCalled(); if (expectRefineChecked) expect(refine).not.toHaveBeenCalled(); @@ -860,9 +904,9 @@ describe("AgentSession queue characterization", () => { state.model = undefined; const internals = harness.session as unknown as AutoRefineInternals; - await internals._maybeAutoRefine("compact"); + await internals._refinement._auto._maybeAutoRefine("compact"); - expect(internals._compactAutoRefinePending).toBe(true); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(true); }); it.each([ @@ -878,13 +922,13 @@ describe("AgentSession queue characterization", () => { }); harnesses.push(harness); const internals = harness.session as unknown as AutoRefineInternals; - internals._assistantTurnsSinceAutoRefine = turns; - internals._lastAutoRefineReviewAt = Date.now(); + internals._refinement._auto._assistantTurnsSinceAutoRefine = turns; + internals._refinement._auto._lastAutoRefineReviewAt = Date.now(); - await internals._maybeAutoRefine(reason); + await internals._refinement._auto._maybeAutoRefine(reason); expect(reviewer).not.toHaveBeenCalled(); - expect(internals[pendingFlag]).toBe(true); + expect(internals._refinement._auto[pendingFlag]).toBe(true); }, ); @@ -1812,7 +1856,7 @@ describe("AgentSession queue characterization", () => { }); it("keeps cleared prompts out of the handoff snapshot during the refine wait", async () => { - let sessionInternals: { _refineInFlight?: Promise }; + let sessionInternals: { _refinement: { _refineInFlight?: Promise } }; let clearDuringRefineWait: (() => void) | undefined; const harness = await createHarness({ extensionFactories: [ @@ -1822,12 +1866,12 @@ describe("AgentSession queue characterization", () => { // message inside that window; the handoff snapshot must not // deliver it. let releaseRefine: (() => void) | undefined; - sessionInternals._refineInFlight = new Promise((resolve) => { + sessionInternals._refinement._refineInFlight = new Promise((resolve) => { releaseRefine = resolve; }); setTimeout(() => { clearDuringRefineWait?.(); - sessionInternals._refineInFlight = undefined; + sessionInternals._refinement._refineInFlight = undefined; releaseRefine?.(); }, 0); return {}; @@ -1836,7 +1880,7 @@ describe("AgentSession queue characterization", () => { ], }); harnesses.push(harness); - sessionInternals = harness.session as unknown as { _refineInFlight?: Promise }; + sessionInternals = harness.session as unknown as { _refinement: { _refineInFlight?: Promise } }; harness.setResponses([fauxAssistantMessage("kept response")]); const clearedAgentMessage = agentPromptText("agentmsg_cleared", "cleared"); diff --git a/packages/coding-agent/test/suite/agent-session-refine-extension.test.ts b/packages/coding-agent/test/suite/agent-session-refine-extension.test.ts index 6a94ed8045..89971eda67 100644 --- a/packages/coding-agent/test/suite/agent-session-refine-extension.test.ts +++ b/packages/coding-agent/test/suite/agent-session-refine-extension.test.ts @@ -217,15 +217,19 @@ describe("AgentSession session_before_refine extension hook", () => { await harness.session.prompt("hello").catch(() => {}); const internals = harness.session as unknown as { - _planRefine(options: unknown, signal: AbortSignal): Promise; + _refinement: { + _execution: { + _planRefine(options: unknown, signal: AbortSignal): Promise; + }; + }; }; // The handler runs but does not short-circuit: planning proceeds to the // built-in planner LLM call, which fails here (no faux response queued) // rather than being skipped. const refineAbort = new AbortController(); - await expect(internals._planRefine({ instructions: "x" }, refineAbort.signal)).rejects.not.toThrow( - RefineSkippedError, - ); + await expect( + internals._refinement._execution._planRefine({ instructions: "x" }, refineAbort.signal), + ).rejects.not.toThrow(RefineSkippedError); expect(handlerCalls).toBe(1); }); @@ -246,22 +250,26 @@ describe("AgentSession session_before_refine extension hook", () => { }); harnesses.push(harness); const internals = harness.session as unknown as { - _maybeAutoRefine(reason: "turn_interval"): Promise; - _assistantTurnsSinceAutoRefine: number; - _turnIntervalAutoRefinePending: boolean; - _pendingAutoRefineReview?: unknown; + _refinement: { + _auto: { + _maybeAutoRefine(reason: "turn_interval"): Promise; + _assistantTurnsSinceAutoRefine: number; + _turnIntervalAutoRefinePending: boolean; + _pendingAutoRefineReview?: unknown; + }; + }; }; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(handlerCalls).toBe(1); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); - expect(internals._turnIntervalAutoRefinePending).toBe(false); - expect(internals._pendingAutoRefineReview).toBeUndefined(); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._turnIntervalAutoRefinePending).toBe(false); + expect(internals._refinement._auto._pendingAutoRefineReview).toBeUndefined(); expect(harness.eventsOfType("refine_failed")).toHaveLength(0); - await internals._maybeAutoRefine("turn_interval"); + await internals._refinement._auto._maybeAutoRefine("turn_interval"); expect(handlerCalls).toBe(1); }); @@ -289,10 +297,20 @@ describe("AgentSession session_before_refine extension hook", () => { await harness.session.prompt("hello").catch(() => {}); const internals = harness.session as unknown as { - _runSerializedAutoRefineReview(reason: "compact" | "turn_interval", branchVersion: number): Promise; - _autoRefineBranchVersion: number; + _refinement: { + _auto: { + _runSerializedAutoRefineReview( + reason: "compact" | "turn_interval", + branchVersion: number, + ): Promise; + _autoRefineBranchVersion: number; + }; + }; }; - await internals._runSerializedAutoRefineReview("turn_interval", internals._autoRefineBranchVersion); + await internals._refinement._auto._runSerializedAutoRefineReview( + "turn_interval", + internals._refinement._auto._autoRefineBranchVersion, + ); expect(events).toHaveLength(1); expect(events[0]?.preparation.trigger).toBe("auto"); diff --git a/packages/coding-agent/test/suite/agent-session-refine-skill.test.ts b/packages/coding-agent/test/suite/agent-session-refine-skill.test.ts index fa10a7c4d6..79f3b8992d 100644 --- a/packages/coding-agent/test/suite/agent-session-refine-skill.test.ts +++ b/packages/coding-agent/test/suite/agent-session-refine-skill.test.ts @@ -2,14 +2,17 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { createHarness, type Harness } from "./harness.js"; type SessionInternals = { - _consumePendingRequestedRefine: () => boolean; - _emitRefineFailed: (error: unknown) => void; - _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; - _serializedPlanInFlight?: Promise; - _serializedExplicitRefineOptions?: { instructions?: string; global?: boolean }; - _refineAbortController?: AbortController; + _refinement: { + _consumePendingRequestedRefine: () => boolean; + _emitRefineFailed: (error: unknown) => void; + _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; + _serializedPlanInFlight?: Promise; + _serializedExplicitRefineOptions?: { instructions?: string; global?: boolean }; + _refineAbortController?: AbortController; + refine: (options: { instructions?: string; global?: boolean }) => Promise; + }; + _createKernelHostHandlers: () => Record; - refine: (options: { instructions?: string; global?: boolean }) => Promise; }; function setStreaming(harness: Harness, streaming: boolean) { @@ -53,7 +56,7 @@ describe("AgentSession refine skill host requests", () => { setStreaming(harness, false); const internals = harness.session as unknown as SessionInternals; - expect(internals._pendingRequestedRefine?.global).toBe(true); + expect(internals._refinement._pendingRequestedRefine?.global).toBe(true); }); it("defaults to local scope when global is not provided", async () => { @@ -67,8 +70,8 @@ describe("AgentSession refine skill host requests", () => { setStreaming(harness, false); const internals = harness.session as unknown as SessionInternals; - expect(internals._pendingRequestedRefine?.global).toBeUndefined(); - expect(internals._pendingRequestedRefine?.instructions).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine?.global).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine?.instructions).toBeUndefined(); }); it("updates pending request when called again", async () => { @@ -83,8 +86,8 @@ describe("AgentSession refine skill host requests", () => { setStreaming(harness, false); const internals = harness.session as unknown as SessionInternals; - expect(internals._pendingRequestedRefine?.instructions).toBe("second"); - expect(internals._pendingRequestedRefine?.global).toBe(true); + expect(internals._refinement._pendingRequestedRefine?.instructions).toBe("second"); + expect(internals._refinement._pendingRequestedRefine?.global).toBe(true); }); it("replaces an in-flight serialized plan instead of applying both requests", async () => { @@ -92,34 +95,34 @@ describe("AgentSession refine skill host requests", () => { harnesses.push(harness); const internals = harness.session as unknown as SessionInternals; const abort = new AbortController(); - internals._serializedPlanInFlight = new Promise(() => {}); - internals._serializedExplicitRefineOptions = { instructions: "first", global: true }; - internals._refineAbortController = abort; + internals._refinement._serializedPlanInFlight = new Promise(() => {}); + internals._refinement._serializedExplicitRefineOptions = { instructions: "first", global: true }; + internals._refinement._refineAbortController = abort; setStreaming(harness, true); harness.session.handleRefineHostRequest("refine.run", { instructions: "replacement" }); setStreaming(harness, false); expect(abort.signal.aborted).toBe(true); - expect(internals._pendingRequestedRefine).toEqual({ instructions: "replacement", global: true }); + expect(internals._refinement._pendingRequestedRefine).toEqual({ instructions: "replacement", global: true }); }); it("discards a settled serialized plan when a replacement request arrives", async () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SessionInternals; - internals._serializedPlanInFlight = Promise.resolve({ status: "plan" }); - internals._serializedExplicitRefineOptions = { instructions: "first", global: true }; + internals._refinement._serializedPlanInFlight = Promise.resolve({ status: "plan" }); + internals._refinement._serializedExplicitRefineOptions = { instructions: "first", global: true }; setStreaming(harness, true); harness.session.handleRefineHostRequest("refine.run", { instructions: "replacement" }); setStreaming(harness, false); - await expect(internals._serializedPlanInFlight).resolves.toEqual({ + await expect(internals._refinement._serializedPlanInFlight).resolves.toEqual({ status: "invalidated", branchVersion: expect.any(Number), }); - expect(internals._pendingRequestedRefine).toEqual({ instructions: "replacement", global: true }); + expect(internals._refinement._pendingRequestedRefine).toEqual({ instructions: "replacement", global: true }); }); it("rejects refine.run while no turn is active", async () => { @@ -179,10 +182,10 @@ describe("AgentSession refine skill host requests", () => { setStreaming(harness, false); const internals = harness.session as unknown as SessionInternals; - const refineSpy = vi.spyOn(internals, "refine").mockResolvedValue({}); - internals._consumePendingRequestedRefine(); + const refineSpy = vi.spyOn(internals._refinement, "refine").mockResolvedValue({}); + internals._refinement._consumePendingRequestedRefine(); expect(refineSpy).toHaveBeenCalledWith({ instructions: "test", global: undefined }, { source: "self" }); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("does nothing when no pending refine at turn boundary", async () => { @@ -191,8 +194,8 @@ describe("AgentSession refine skill host requests", () => { await harness.session.prompt("one"); const internals = harness.session as unknown as SessionInternals; - const refineSpy = vi.spyOn(internals, "refine").mockResolvedValue({}); - expect(internals._consumePendingRequestedRefine()).toBe(false); + const refineSpy = vi.spyOn(internals._refinement, "refine").mockResolvedValue({}); + expect(internals._refinement._consumePendingRequestedRefine()).toBe(false); expect(refineSpy).not.toHaveBeenCalled(); }); @@ -207,7 +210,7 @@ describe("AgentSession refine skill host requests", () => { setStreaming(harness, false); const internals = harness.session as unknown as SessionInternals; - vi.spyOn(internals, "refine").mockRejectedValue(new Error("refine failed")); + vi.spyOn(internals._refinement, "refine").mockRejectedValue(new Error("refine failed")); const failed = new Promise((resolve) => { const unsubscribe = harness.session.subscribe((event) => { if (event.type === "refine_failed") { @@ -217,9 +220,9 @@ describe("AgentSession refine skill host requests", () => { }); }); - expect(internals._consumePendingRequestedRefine()).toBe(true); + expect(internals._refinement._consumePendingRequestedRefine()).toBe(true); expect(await failed).toBe("refine failed"); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("continues notifying refine listeners after one throws", async () => { @@ -236,7 +239,7 @@ describe("AgentSession refine skill host requests", () => { } }); - expect(() => internals._emitRefineFailed(new Error("planning failed"))).not.toThrow(); + expect(() => internals._refinement._emitRefineFailed(new Error("planning failed"))).not.toThrow(); expect(observed).toEqual(["planning failed"]); }); 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 99776df104..ea88d3ce77 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,10 @@ 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 { CompactionResult } from "../../src/core/compaction/index.js"; +import type { SessionCompaction } from "../../src/session/compaction.js"; +import type { CompactionExecutionOptions } from "../../src/session/compaction-execution.js"; +import type { SessionContinuation } from "../../src/session/continuation.js"; import type { SessionRetry } from "../../src/session/retry.js"; import { createHarness, type Harness } from "./harness.js"; @@ -53,8 +57,9 @@ function rateLimitedFailure(retryAfterMs: number): AssistantMessage { type SessionRetryCompactionInternals = { _retry: SessionRetry; - _autoCompactionAbortController: AbortController | undefined; - _postCompactionContinuationScheduled: boolean; + _compaction: SessionCompaction; + _performCompaction(options: CompactionExecutionOptions): Promise; + _continuation: SessionContinuation; _processAgentEvent: (event: AgentEvent) => Promise; _checkCompaction: (message: AssistantMessage) => Promise; _schedulePostCompactionContinue: () => void; @@ -427,20 +432,27 @@ describe("AgentSession retry and event characterization", () => { const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } } }); harnesses.push(harness); const internals = harness.session as unknown as SessionRetryCompactionInternals; - const compactionAbortController = new AbortController(); + let compactionSignal: AbortSignal | undefined; + const performCompaction = vi.spyOn(internals, "_performCompaction").mockImplementation(({ signal }) => { + compactionSignal = signal; + return new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => reject(new Error("Compaction cancelled")), { once: true }); + }); + }); 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; + const compaction = internals._compaction.runAutomatic("overflow", true); + await vi.waitFor(() => expect(compactionSignal).toBeDefined()); internals._schedulePostCompactionContinue(); try { - expect(internals._postCompactionContinuationScheduled).toBe(true); + expect(internals._continuation.isScheduled).toBe(true); harness.session.abortRetry(); - expect(compactionAbortController.signal.aborted).toBe(true); - expect(internals._postCompactionContinuationScheduled).toBe(false); + expect(compactionSignal?.aborted).toBe(true); + expect(internals._continuation.isScheduled).toBe(false); expect(harness.session.retryAttempt).toBe(0); expect(harness.session.isRetrying).toBe(false); expect(harness.eventsOfType("auto_retry_end").at(-1)).toMatchObject({ @@ -449,7 +461,9 @@ describe("AgentSession retry and event characterization", () => { finalError: "Retry cancelled", }); } finally { - internals._autoCompactionAbortController = undefined; + internals._compaction.abortAutomatic(); + await compaction; + performCompaction.mockRestore(); internals._cancelPostCompactionContinue(); } }); diff --git a/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts b/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts index 948dba16b8..7ba5cedba7 100644 --- a/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts +++ b/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts @@ -2,29 +2,62 @@ import type { AgentTool } from "@earendil-works/pi-agent-core"; import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; +import { getLocalHarnessStateDir, loadHarnessState, saveHarnessState } from "../../src/core/refinement/index.js"; import { createHarness, getMessageText, type Harness } from "./harness.js"; type SerializedInternals = { + _refinement: { + _host: { getCompactionOperation(): Promise | undefined }; + _runSerializedRefineCheckpoint(): Promise; + _runSerializedRefine( + options: { instructions?: string; global?: boolean }, + source: "auto" | "self", + ): Promise; + _consumeSerializedBackgroundPlan(consume: (result: unknown) => Promise): Promise; + _consumePendingRequestedRefine(): boolean; + _serializedRefine: boolean; + _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; + _refineInFlight?: Promise; + _refinePlanInFlight?: Promise; + _serializedPlanInFlight?: Promise; + _refineAbortController?: AbortController | undefined; + _maybeStartSerializedBackgroundPlan: () => void; + _drainPendingRefinementForDisposal(): Promise; + _auto: { + _runSerializedAutoRefineReview(reason: "turn_interval" | "compact", branchVersion: number): Promise; + _maybeAutoRefine(reason: string): Promise; + _reviewAutoRefine( + context: { reason: string; turnsSinceLastReview: number }, + signal?: AbortSignal, + ): Promise; + _assistantTurnsSinceAutoRefine: number; + _lastAutoRefineReviewAt: number; + _autoRefineInProgress: boolean; + _autoRefineBranchVersion: number; + _scheduleAutoRefineAfterAgentEnd(): void; + _autoRefineReviewAbort?: AbortController | undefined; + _compactAutoRefinePending: boolean; + _scheduleAutoRefine(reason: "turn_interval" | "compact"): void; + _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; + _autoRefineOperations: Set>; + }; + _execution: { + _planRefine(options: { instructions?: string; global?: boolean }, signal: AbortSignal): Promise; + _applyRefine( + plan: unknown, + options: { instructions?: string; global?: boolean }, + abort: AbortController, + ): Promise; + }; + }; + _shouldStopAfterTurn(context: { message: { stopReason?: string; content: unknown[]; role: string; usage?: unknown; timestamp?: number }; toolResults: unknown[]; context: unknown; newMessages: unknown[]; }): Promise; - _runSerializedRefineCheckpoint(): Promise; - _runSerializedRefine(options: { instructions?: string; global?: boolean }, source: "auto" | "self"): Promise; - _consumeSerializedBackgroundPlan(consume: (result: unknown) => Promise): Promise; - _runSerializedAutoRefineReview(reason: "turn_interval" | "compact", branchVersion: number): Promise; - _consumePendingRequestedRefine(): boolean; - _maybeAutoRefine(reason: string): Promise; - _reviewAutoRefine(context: { reason: string; turnsSinceLastReview: number }, signal?: AbortSignal): Promise; - _planRefine(options: { instructions?: string; global?: boolean }, signal: AbortSignal): Promise; - _applyRefine( - plan: unknown, - options: { instructions?: string; global?: boolean }, - abort: AbortController, - ): Promise; - _serializedRefine: boolean; + _createPreparedTurnAction( schedule: "steer", text: string, @@ -32,24 +65,15 @@ type SerializedInternals = { options: Record, ): unknown; _admitSessionInput(action: unknown, options?: { wake?: boolean }): { accepted: boolean }; - _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; - _assistantTurnsSinceAutoRefine: number; - _lastAutoRefineReviewAt: number; - _autoRefineInProgress: boolean; - _autoRefineBranchVersion: number; - _refineInFlight?: Promise; - _refinePlanInFlight?: Promise; - _serializedPlanInFlight?: Promise; + _disposing: boolean; _disposed: boolean; - _scheduleAutoRefineAfterAgentEnd(): void; + _checkCompaction(message: unknown): Promise; _lastAssistantMessage: unknown; _handleAgentEvent(event: { type: string; messages?: unknown[] }): void; _agentEventQueue: Promise; - _autoRefineReviewAbort?: AbortController | undefined; - _refineAbortController?: AbortController | undefined; - _compactionOperation?: Promise | undefined; + _branchSummaryOperation?: Promise | undefined; requestAbort(): void; abortCompaction(): void; @@ -58,18 +82,13 @@ type SerializedInternals = { _pendingMessageResumeEpoch: number; _pendingMessageResumeQueue: Promise; _schedulePendingMessageResume(request?: boolean): void; - _maybeStartSerializedBackgroundPlan: () => void; - _compactAutoRefinePending: boolean; - _scheduleAutoRefine(reason: "turn_interval" | "compact"): void; - _scheduleAutoRefineAfterCompaction(willContinueAfterCompaction: boolean): void; + _getRequiredRequestAuth(model: unknown): Promise<{ apiKey: string; headers?: Record }>; _performCompaction(options: unknown): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; }>; - _drainPendingRefinementForDisposal(): Promise; - _autoRefineOperations: Set>; }; function emptyRefinementResult() { @@ -102,9 +121,9 @@ function makeCtx(text: string) { function mockSerializedRefine(harness: Harness) { const internals = harness.session as unknown as SerializedInternals; const plan = { id: "test-plan", proposal: { edits: [] } }; - vi.spyOn(internals, "_planRefine").mockResolvedValue(plan); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); - return { plan, applyRefine: internals._applyRefine }; + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue(plan); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + return { plan, applyRefine: internals._refinement._execution._applyRefine }; } describe("Serialized auto-refine checkpoint", () => { @@ -134,7 +153,7 @@ describe("Serialized auto-refine checkpoint", () => { const internals = harness.session as unknown as SerializedInternals; for (let turn = 1; turn <= 26; turn++) { - internals._assistantTurnsSinceAutoRefine++; // simulate message_end increment + internals._refinement._auto._assistantTurnsSinceAutoRefine++; // simulate message_end increment await internals._shouldStopAfterTurn(makeCtx(`turn ${turn}`)); } @@ -156,15 +175,15 @@ describe("Serialized auto-refine checkpoint", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; let applyInFlight = false; let resolveApply: () => void = () => {}; const applyPromise = new Promise((resolve) => { resolveApply = resolve; }); - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockImplementation(async () => { applyInFlight = true; await applyPromise; applyInFlight = false; @@ -202,13 +221,13 @@ describe("Serialized auto-refine checkpoint", () => { const { applyRefine } = mockSerializedRefine(harness); const internals = harness.session as unknown as SerializedInternals; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // _applyRefine was called (which rebuilds system prompt). expect(applyRefine).toHaveBeenCalledTimes(1); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); it("final agent_end pending refine completes before dispose", async () => { @@ -226,7 +245,7 @@ describe("Serialized auto-refine checkpoint", () => { const { applyRefine } = mockSerializedRefine(harness); const internals = harness.session as unknown as SerializedInternals; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; await harness.session.disposeAsync(); @@ -249,14 +268,14 @@ describe("Serialized auto-refine checkpoint", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - expect(internals._serializedRefine).toBe(false); + expect(internals._refinement._serializedRefine).toBe(false); let checkpointCalled = false; - vi.spyOn(internals, "_runSerializedRefineCheckpoint").mockImplementation(async () => { + vi.spyOn(internals._refinement, "_runSerializedRefineCheckpoint").mockImplementation(async () => { checkpointCalled = true; }); - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; await internals._shouldStopAfterTurn(makeCtx("test")); expect(checkpointCalled).toBe(false); @@ -280,7 +299,7 @@ describe("Serialized auto-refine checkpoint", () => { (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; const internals = harness.session as unknown as SerializedInternals; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; // Must resolve within 5 seconds — if _maybeAutoRefine were called, // _shouldSkipAutoRefineForActiveAgent would defer, and if waitForIdle @@ -289,10 +308,10 @@ describe("Serialized auto-refine checkpoint", () => { setTimeout(() => reject(new Error("Serialized checkpoint deadlocked")), 5000), ); - await Promise.race([internals._runSerializedRefineCheckpoint(), timeout]); + await Promise.race([internals._refinement._runSerializedRefineCheckpoint(), timeout]); expect(applyRefine).toHaveBeenCalledTimes(1); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); }); @@ -315,7 +334,7 @@ describe("Serialized agent-callable refine", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; const { applyRefine } = mockSerializedRefine(harness); - internals._pendingRequestedRefine = { instructions: "capture a lesson" }; + internals._refinement._pendingRequestedRefine = { instructions: "capture a lesson" }; internals._admitSessionInput(internals._createPreparedTurnAction("steer", "steer", undefined, {})); const compactionSpy = vi .spyOn( @@ -348,14 +367,14 @@ describe("Serialized agent-callable refine", () => { // In serialized mode, handleRefineHostRequest immediately kicks off // background planning, consuming the pending request. - expect(internals._pendingRequestedRefine).toBeUndefined(); - expect(internals._serializedPlanInFlight).toBeDefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._serializedPlanInFlight).toBeDefined(); // Wait for background planning to complete await new Promise((resolve) => setTimeout(resolve, 50)); // At the boundary, the plan should be applied - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(applyRefine).toHaveBeenCalledTimes(1); }); @@ -375,19 +394,19 @@ describe("Serialized agent-callable refine", () => { const { applyRefine } = mockSerializedRefine(harness); const internals = harness.session as unknown as SerializedInternals; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; harness.session.handleRefineHostRequest("refine.run", { instructions: "callable lesson" }); (harness.session.agent.state as { isStreaming: boolean }).isStreaming = false; - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // _applyRefine called once (for the callable request only). // The explicit refine satisfied the interval and reset the counter, // so no interval-triggered auto-refine follows. expect(applyRefine).toHaveBeenCalledTimes(1); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); // The reviewer was NOT called because the explicit refine // satisfied the interval check. expect(reviewer).not.toHaveBeenCalled(); @@ -409,8 +428,8 @@ describe("Serialized agent-callable refine", () => { // In serialized mode, the pending request is consumed immediately // by background planning, NOT left for fire-and-forget at agent_end. - expect(internals._pendingRequestedRefine).toBeUndefined(); - expect(internals._serializedPlanInFlight).toBeDefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._serializedPlanInFlight).toBeDefined(); }); it("pending agent-callable refine drained before disposal", async () => { @@ -465,7 +484,7 @@ describe("Serialized autonomous continuation", () => { // Simulate shouldStopAfterTurn calls for turns 1-5. const turnResults: boolean[] = []; for (let turn = 1; turn <= 5; turn++) { - internals._assistantTurnsSinceAutoRefine++; // simulate message_end increment + internals._refinement._auto._assistantTurnsSinceAutoRefine++; // simulate message_end increment const stopResult = await Promise.race([ internals._shouldStopAfterTurn(makeCtx(`autonomous turn ${turn}`)), new Promise((_, reject) => setTimeout(() => reject(new Error(`Deadlock at turn ${turn}`)), 5000)), @@ -482,7 +501,7 @@ describe("Serialized autonomous continuation", () => { // After the refine at turn 3, counter was reset to 0. // Turns 4 and 5 set it to 1 and 2, both < 3. - expect(internals._assistantTurnsSinceAutoRefine).toBe(2); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(2); }); }); @@ -517,20 +536,18 @@ describe("Serialized background planning during tools", () => { let planStarted = false; let planFinished = false; - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { planStarted = true; await planPromise; planFinished = true; return { id: "bg-plan", proposal: { edits: [] } }; }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Step 1: simulate message_end — this should kick off background planning. - internals._assistantTurnsSinceAutoRefine++; + internals._refinement._auto._assistantTurnsSinceAutoRefine++; // Call the private method that starts background planning - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._maybeStartSerializedBackgroundPlan(); // Wait a tick for the async plan to start await new Promise((resolve) => setTimeout(resolve, 10)); @@ -543,16 +560,16 @@ describe("Serialized background planning during tools", () => { // The checkpoint should be waiting for the plan (not yet resolved). await new Promise((resolve) => setTimeout(resolve, 10)); // _applyRefine has NOT been called yet because planning is still in flight. - expect(internals._applyRefine).not.toHaveBeenCalled(); + expect(internals._refinement._execution._applyRefine).not.toHaveBeenCalled(); // Release the plan — the checkpoint should proceed to apply. resolvePlan(); await checkpointPromise; // Now apply has been called. - expect(internals._applyRefine).toHaveBeenCalledTimes(1); + expect(internals._refinement._execution._applyRefine).toHaveBeenCalledTimes(1); expect(planFinished).toBe(true); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); it("lets refine.run supersede an interval plan that ignores abort", async () => { @@ -575,37 +592,37 @@ describe("Serialized background planning during tools", () => { resolveIntervalPlan = resolve; }); let planCalls = 0; - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { planCalls++; if (planCalls === 1) { await intervalPlanReady; } return { id: `plan-${planCalls}`, proposal: { edits: [] } }; }); - const apply = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const apply = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); - internals._assistantTurnsSinceAutoRefine = 1; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + internals._refinement._maybeStartSerializedBackgroundPlan(); await new Promise((resolve) => setTimeout(resolve, 10)); expect(planCalls).toBe(1); (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; harness.session.handleRefineHostRequest("refine.run", { instructions: "explicit plan" }); (harness.session.agent.state as { isStreaming: boolean }).isStreaming = false; - expect(internals._pendingRequestedRefine?.instructions).toBe("explicit plan"); + expect(internals._refinement._pendingRequestedRefine?.instructions).toBe("explicit plan"); - const checkpoint = internals._runSerializedRefineCheckpoint(); + const checkpoint = internals._refinement._runSerializedRefineCheckpoint(); await new Promise((resolve) => setTimeout(resolve, 10)); resolveIntervalPlan(); await checkpoint; expect(reviewer).toHaveBeenCalledTimes(1); - expect(internals._planRefine).toHaveBeenCalledTimes(2); + expect(internals._refinement._execution._planRefine).toHaveBeenCalledTimes(2); expect(apply).toHaveBeenCalledTimes(1); expect(apply.mock.calls[0]?.[0]).toMatchObject({ id: "plan-2" }); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("max concurrent model calls is one: planning does not start another model request", async () => { @@ -628,14 +645,12 @@ describe("Serialized background planning during tools", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Simulate message_end - internals._assistantTurnsSinceAutoRefine++; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine++; + internals._refinement._maybeStartSerializedBackgroundPlan(); // Wait for planning to complete (reviewer is async) await new Promise((resolve) => setTimeout(resolve, 50)); @@ -647,8 +662,8 @@ describe("Serialized background planning during tools", () => { // and apply it without starting a new model request. await internals._shouldStopAfterTurn(makeCtx("boundary")); - expect(internals._applyRefine).toHaveBeenCalledTimes(1); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._execution._applyRefine).toHaveBeenCalledTimes(1); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); }); @@ -697,8 +712,8 @@ describe("PR #503 model persistence regression", () => { const internals = harness.session as unknown as SerializedInternals; const modelBefore = harness.session.model; - internals._assistantTurnsSinceAutoRefine = 1; - await internals._runSerializedRefineCheckpoint(); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + await internals._refinement._runSerializedRefineCheckpoint(); // Model is preserved after the checkpoint. expect(applyRefine).toHaveBeenCalledTimes(1); @@ -731,8 +746,8 @@ describe("Serialized refine review-fix regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Queue a refine.run request — in serialized mode, this immediately // kicks off background planning (skipping the review gate). @@ -741,7 +756,7 @@ describe("Serialized refine review-fix regressions", () => { (harness.session.agent.state as { isStreaming: boolean }).isStreaming = false; // Background planning was started automatically by handleRefineHostRequest. - expect(internals._serializedPlanInFlight).toBeDefined(); + expect(internals._refinement._serializedPlanInFlight).toBeDefined(); // Wait for background planning to complete await new Promise((resolve) => setTimeout(resolve, 50)); @@ -750,9 +765,9 @@ describe("Serialized refine review-fix regressions", () => { expect(reviewer).not.toHaveBeenCalled(); // At the boundary, the plan should be applied - await internals._runSerializedRefineCheckpoint(); - expect(internals._applyRefine).toHaveBeenCalledTimes(1); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + await internals._refinement._runSerializedRefineCheckpoint(); + expect(internals._refinement._execution._applyRefine).toHaveBeenCalledTimes(1); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); it("failure result stamps cooldown and does NOT retry synchronously", async () => { @@ -771,40 +786,40 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; // Make background planning fail - vi.spyOn(internals, "_planRefine").mockRejectedValue(new Error("plan failed")); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockRejectedValue(new Error("plan failed")); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Start background planning - internals._assistantTurnsSinceAutoRefine++; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine++; + internals._refinement._maybeStartSerializedBackgroundPlan(); // Wait for background planning to fail await new Promise((resolve) => setTimeout(resolve, 50)); // At the boundary, the failure result should stamp cooldown and return. // It should NOT fall through to synchronous review+plan (no duplicate model call). - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // Reviewer was called exactly once (during background planning). // It should NOT be called again at the boundary (failure -> no retry). expect(reviewer).toHaveBeenCalledTimes(1); // _applyRefine was NOT called (planning failed). - expect(internals._applyRefine).not.toHaveBeenCalled(); + expect(internals._refinement._execution._applyRefine).not.toHaveBeenCalled(); // Cooldown was stamped. - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); }); it("non-serialized explicit refine suppresses interval auto-refine for that turn", async () => { const harness = await createHarness({ persistSession: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - refine(options: { instructions?: string }): Promise; + _refinement: { + refine(options: { instructions?: string }): Promise; + }; }; - internals._pendingRequestedRefine = { instructions: "explicit" }; - const refine = vi.spyOn(internals, "refine").mockResolvedValue(emptyRefinementResult()); - const schedule = vi.spyOn(internals, "_scheduleAutoRefineAfterAgentEnd"); + internals._refinement._pendingRequestedRefine = { instructions: "explicit" }; + const refine = vi.spyOn(internals._refinement, "refine").mockResolvedValue(emptyRefinementResult()); + const schedule = vi.spyOn(internals._refinement._auto, "_scheduleAutoRefineAfterAgentEnd"); const assistant = fauxAssistantMessage("done"); internals._lastAssistantMessage = assistant; @@ -831,25 +846,23 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; // Make background planning fail - vi.spyOn(internals, "_planRefine").mockRejectedValue(new Error("plan failed")); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockRejectedValue(new Error("plan failed")); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Start background planning and let it fail - internals._assistantTurnsSinceAutoRefine = 1; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + internals._refinement._maybeStartSerializedBackgroundPlan(); await new Promise((resolve) => setTimeout(resolve, 50)); // At the boundary, the failure result should stamp cooldown and return. - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // Reviewer was called exactly once (during background planning). expect(reviewer).toHaveBeenCalledTimes(1); // _applyRefine was NOT called (planning failed). - expect(internals._applyRefine).not.toHaveBeenCalled(); + expect(internals._refinement._execution._applyRefine).not.toHaveBeenCalled(); // Cooldown was stamped. - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); // Now simulate the actual agent_end event path. // Set _lastAssistantMessage so agent_end has a non-error assistant to process. @@ -857,7 +870,7 @@ describe("Serialized refine review-fix regressions", () => { internals._lastAssistantMessage = fauxAssistant; // Spy on _scheduleAutoRefineAfterAgentEnd — it should NOT be called in serialized mode. - const scheduleSpy = vi.spyOn(internals, "_scheduleAutoRefineAfterAgentEnd"); + const scheduleSpy = vi.spyOn(internals._refinement._auto, "_scheduleAutoRefineAfterAgentEnd"); // Mock _checkCompaction so agent_end reaches the scheduling guard. vi.spyOn(internals, "_checkCompaction").mockResolvedValue(false); @@ -887,12 +900,12 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; let planCalls = 0; - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { planCalls++; if (planCalls === 1) throw new Error("bg plan failed"); return { id: "replan", proposal: { edits: [] } }; }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Queue an explicit refine.run — this starts background planning at message_end. (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; @@ -901,15 +914,15 @@ describe("Serialized refine review-fix regressions", () => { // Let the background plan fail. await new Promise((resolve) => setTimeout(resolve, 50)); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); // At the boundary, the failure re-queues the explicit options and the // synchronous pending path replans + applies. - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(planCalls).toBe(2); // bg plan (failed) + synchronous replan - expect(internals._applyRefine).toHaveBeenCalledTimes(1); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._execution._applyRefine).toHaveBeenCalledTimes(1); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("interval bg plan failure does not retry synchronously", async () => { @@ -928,43 +941,48 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; // Interval background planning fails (not explicit refine.run). - vi.spyOn(internals, "_planRefine").mockRejectedValue(new Error("plan failed")); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockRejectedValue(new Error("plan failed")); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); - internals._assistantTurnsSinceAutoRefine = 1; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + internals._refinement._maybeStartSerializedBackgroundPlan(); await new Promise((resolve) => setTimeout(resolve, 50)); - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // Reviewer called once (interval bg plan), NOT retried at boundary. expect(reviewer).toHaveBeenCalledTimes(1); - expect(internals._applyRefine).not.toHaveBeenCalled(); - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._execution._applyRefine).not.toHaveBeenCalled(); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("classifies an aborted stale background plan as invalidated", async () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - _runBackgroundPlan( - options: { instructions?: string }, - abort: AbortController, - branchVersion: number, - skipReview?: boolean, - ): Promise; + _refinement: { + _runBackgroundPlan( + options: { instructions?: string }, + abort: AbortController, + branchVersion: number, + skipReview?: boolean, + ): Promise; + }; }; - const branchVersion = internals._autoRefineBranchVersion; - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { - internals._autoRefineBranchVersion++; + const branchVersion = internals._refinement._auto._autoRefineBranchVersion; + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { + internals._refinement._auto._autoRefineBranchVersion++; throw new Error("branch changed"); }); await expect( - internals._runBackgroundPlan({ instructions: "stale" }, new AbortController(), branchVersion, true), + internals._refinement._runBackgroundPlan( + { instructions: "stale" }, + new AbortController(), + branchVersion, + true, + ), ).resolves.toEqual({ status: "invalidated", branchVersion }); }); @@ -983,8 +1001,8 @@ describe("Serialized refine review-fix regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - vi.spyOn(internals, "_planRefine").mockRejectedValue(new Error("plan failed")); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockRejectedValue(new Error("plan failed")); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Queue explicit refine.run — starts background planning. (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; @@ -993,13 +1011,13 @@ describe("Serialized refine review-fix regressions", () => { await new Promise((resolve) => setTimeout(resolve, 50)); // Invalidate the branch (simulates a newer turn or branch reset). - internals._autoRefineBranchVersion++; + internals._refinement._auto._autoRefineBranchVersion++; - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // Stale branch: no re-queue, no apply. - expect(internals._applyRefine).not.toHaveBeenCalled(); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._execution._applyRefine).not.toHaveBeenCalled(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("newer pending supersedes failed older explicit refine.run", async () => { @@ -1017,11 +1035,11 @@ describe("Serialized refine review-fix regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - vi.spyOn(internals, "_planRefine").mockImplementation(async (opts) => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async (opts) => { if (opts.instructions === "older") throw new Error("bg plan failed"); return { id: "newer-plan", proposal: { edits: [] } }; }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Queue older request — starts background planning. (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; @@ -1030,13 +1048,13 @@ describe("Serialized refine review-fix regressions", () => { await new Promise((resolve) => setTimeout(resolve, 50)); // Queue newer request before the boundary. - internals._pendingRequestedRefine = { instructions: "newer" }; + internals._refinement._pendingRequestedRefine = { instructions: "newer" }; - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // Newer request is serviced (plan with "newer", not "older"). - expect(internals._applyRefine).toHaveBeenCalledTimes(1); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._execution._applyRefine).toHaveBeenCalledTimes(1); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("disposal does not double-refine after explicit drain", async () => { @@ -1055,29 +1073,27 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; let planCalls = 0; - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { planCalls++; return { id: `plan-${planCalls}`, proposal: { edits: [] } }; }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Set up: assistant turns met, explicit refine.run pending. - internals._assistantTurnsSinceAutoRefine = 1; - internals._pendingRequestedRefine = { instructions: "explicit" }; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + internals._refinement._pendingRequestedRefine = { instructions: "explicit" }; // Drive the disposal drain path. - ( - internals as unknown as { _drainPendingRefinementForDisposal: () => Promise } - )._drainPendingRefinementForDisposal(); + internals._refinement._drainPendingRefinementForDisposal(); await new Promise((resolve) => setTimeout(resolve, 50)); // The explicit drain should have planned + applied exactly once. expect(planCalls).toBe(1); - expect(internals._applyRefine).toHaveBeenCalledTimes(1); + expect(internals._refinement._execution._applyRefine).toHaveBeenCalledTimes(1); // Counter reset prevents the interval check from triggering a second refine. - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("public refine waits for serialized bg plan: max concurrency 1, bg then public apply", async () => { @@ -1127,7 +1143,7 @@ describe("Serialized refine review-fix regressions", () => { resolveBgPlan = resolve; }); const applyCalls: { plan: unknown; options: unknown }[] = []; - vi.spyOn(internals, "_planRefine").mockImplementation(async (_opts) => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async (_opts) => { activePlans++; maxConcurrentPlans = Math.max(maxConcurrentPlans, activePlans); planCalls++; @@ -1138,7 +1154,7 @@ describe("Serialized refine review-fix regressions", () => { activePlans--; return { id: `plan-${planCalls}`, proposal: { edits: [] } }; }); - vi.spyOn(internals, "_applyRefine").mockImplementation(async (plan, options) => { + vi.spyOn(internals._refinement._execution, "_applyRefine").mockImplementation(async (plan, options) => { applyCalls.push({ plan, options }); return emptyRefinementResult(); }); @@ -1161,7 +1177,7 @@ describe("Serialized refine review-fix regressions", () => { harness.session.handleRefineHostRequest("refine.run", { instructions: "bg" }); await new Promise((resolve) => setTimeout(resolve, 50)); expect(planCalls).toBe(1); - expect(internals._serializedPlanInFlight).toBeDefined(); + expect(internals._refinement._serializedPlanInFlight).toBeDefined(); // Call public refine concurrently while bg plan is in flight. const publicRefinePromise = harness.session.refine({ instructions: "public" }); @@ -1202,23 +1218,30 @@ describe("Serialized refine review-fix regressions", () => { const compactionOperation = new Promise((resolve) => { releaseCompaction = resolve; }); - internals._compactionOperation = compactionOperation; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "public-plan", proposal: { edits: [] } }); - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const getCompactionOperation = vi + .spyOn(internals._refinement._host, "getCompactionOperation") + .mockReturnValue(compactionOperation); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "public-plan", + proposal: { edits: [] }, + }); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); const abortCompaction = vi.spyOn(internals, "abortCompaction"); const refinePromise = harness.session.refine({ instructions: "public" }); - await vi.waitFor(() => expect(internals._refineInFlight).toBeDefined()); + await vi.waitFor(() => expect(internals._refinement._refineInFlight).toBeDefined()); expect(abortCompaction).not.toHaveBeenCalled(); expect(applyRefine).not.toHaveBeenCalled(); - expect(internals._refineAbortController?.signal.aborted).toBe(false); + expect(internals._refinement._refineAbortController?.signal.aborted).toBe(false); releaseCompaction(); await refinePromise; expect(applyRefine).toHaveBeenCalledOnce(); - internals._compactionOperation = undefined; - internals._refineAbortController = undefined; + getCompactionOperation.mockRestore(); + internals._refinement._refineAbortController = undefined; }); it("public refine waits for active branch summary without aborting it", async () => { @@ -1230,22 +1253,27 @@ describe("Serialized refine review-fix regressions", () => { releaseBranchSummary = resolve; }); internals._branchSummaryOperation = branchSummaryOperation; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "public-plan", proposal: { edits: [] } }); - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "public-plan", + proposal: { edits: [] }, + }); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); const abortBranchSummary = vi.spyOn(internals, "abortBranchSummary"); const refinePromise = harness.session.refine({ instructions: "public" }); - await vi.waitFor(() => expect(internals._refineInFlight).toBeDefined()); + await vi.waitFor(() => expect(internals._refinement._refineInFlight).toBeDefined()); expect(abortBranchSummary).not.toHaveBeenCalled(); expect(applyRefine).not.toHaveBeenCalled(); - expect(internals._refineAbortController?.signal.aborted).toBe(false); + expect(internals._refinement._refineAbortController?.signal.aborted).toBe(false); releaseBranchSummary(); await refinePromise; expect(applyRefine).toHaveBeenCalledOnce(); internals._branchSummaryOperation = undefined; - internals._refineAbortController = undefined; + internals._refinement._refineAbortController = undefined; }); it("public refine snapshots final events and operations after the agent becomes idle", async () => { @@ -1257,8 +1285,13 @@ describe("Serialized refine review-fix regressions", () => { releaseIdle = resolve; }); const waitForIdle = vi.spyOn(harness.session.agent, "waitForIdle").mockReturnValue(idleOperation); - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "public-plan", proposal: { edits: [] } }); - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "public-plan", + proposal: { edits: [] }, + }); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); const refinePromise = harness.session.refine({ instructions: "public" }); await vi.waitFor(() => expect(waitForIdle).toHaveBeenCalledOnce()); @@ -1268,9 +1301,12 @@ describe("Serialized refine review-fix regressions", () => { releaseEventQueue = resolve; }); let releaseCompaction: () => void = () => {}; - internals._compactionOperation = new Promise((resolve) => { + const compactionOperation = new Promise((resolve) => { releaseCompaction = resolve; }); + const getCompactionOperation = vi + .spyOn(internals._refinement._host, "getCompactionOperation") + .mockReturnValue(compactionOperation); releaseIdle(); await new Promise((resolve) => setTimeout(resolve, 20)); @@ -1284,8 +1320,8 @@ describe("Serialized refine review-fix regressions", () => { await refinePromise; expect(applyRefine).toHaveBeenCalledOnce(); - internals._compactionOperation = undefined; - internals._refineAbortController = undefined; + getCompactionOperation.mockRestore(); + internals._refinement._refineAbortController = undefined; }); it("interval background plan derives instructions from review, not prepopulated", async () => { @@ -1304,17 +1340,17 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; let capturedPlanOptions: { instructions?: string } | undefined; - vi.spyOn(internals, "_planRefine").mockImplementation(async (opts: { instructions?: string }) => { - capturedPlanOptions = opts; - return { id: "p", proposal: { edits: [] } }; - }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation( + async (opts: { instructions?: string }) => { + capturedPlanOptions = opts; + return { id: "p", proposal: { edits: [] } }; + }, + ); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Start background planning (interval-triggered) - internals._assistantTurnsSinceAutoRefine++; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine++; + internals._refinement._maybeStartSerializedBackgroundPlan(); // Wait for background planning await new Promise((resolve) => setTimeout(resolve, 50)); @@ -1330,14 +1366,14 @@ describe("Serialized refine review-fix regressions", () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._serializedPlanInFlight = Promise.resolve({ + internals._refinement._serializedPlanInFlight = Promise.resolve({ status: "plan", plan: { id: "failed-plan", proposal: { edits: [] } }, options: { instructions: "explicit" }, abort: new AbortController(), - branchVersion: internals._autoRefineBranchVersion, + branchVersion: internals._refinement._auto._autoRefineBranchVersion, }); - vi.spyOn(internals, "_applyRefine").mockRejectedValue(new Error("harness write failed")); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockRejectedValue(new Error("harness write failed")); const failed = new Promise((resolve) => { const unsubscribe = harness.session.subscribe((event) => { if (event.type === "refine_failed") { @@ -1347,7 +1383,7 @@ describe("Serialized refine review-fix regressions", () => { }); }); - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(await failed).toBe("harness write failed"); }); @@ -1367,14 +1403,14 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; const exactPlan = { id: "disposal-plan", proposal: { edits: [] } }; - vi.spyOn(internals, "_planRefine").mockResolvedValue(exactPlan); - const applySpy = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue(exactPlan); + const applySpy = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); // Start background planning - internals._assistantTurnsSinceAutoRefine++; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine++; + internals._refinement._maybeStartSerializedBackgroundPlan(); // Wait for the background plan to settle (plan ready) await new Promise((resolve) => setTimeout(resolve, 50)); @@ -1407,13 +1443,13 @@ describe("Serialized refine review-fix regressions", () => { const internals = harness.session as unknown as SerializedInternals; const { applyRefine } = mockSerializedRefine(harness); - internals._scheduleAutoRefineAfterCompaction(true); - expect(internals._compactAutoRefinePending).toBe(true); - await internals._runSerializedRefineCheckpoint(); + internals._refinement._auto._scheduleAutoRefineAfterCompaction(true); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(true); + await internals._refinement._runSerializedRefineCheckpoint(); expect(reviewer).toHaveBeenCalledWith(expect.objectContaining({ reason: "compact" }), expect.any(AbortSignal)); expect(applyRefine).toHaveBeenCalledTimes(1); - expect(internals._compactAutoRefinePending).toBe(false); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); }); it("defers serialized compaction refinement even when no continuation was scheduled", async () => { @@ -1430,17 +1466,17 @@ describe("Serialized refine review-fix regressions", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - const interactiveSpy = vi.spyOn(internals, "_maybeAutoRefine"); + const interactiveSpy = vi.spyOn(internals._refinement._auto, "_maybeAutoRefine"); const { applyRefine } = mockSerializedRefine(harness); - internals._scheduleAutoRefineAfterCompaction(false); - expect(internals._compactAutoRefinePending).toBe(true); + internals._refinement._auto._scheduleAutoRefineAfterCompaction(false); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(true); expect(interactiveSpy).not.toHaveBeenCalled(); - await internals._drainPendingRefinementForDisposal(); + await internals._refinement._drainPendingRefinementForDisposal(); expect(reviewer).toHaveBeenCalledWith(expect.objectContaining({ reason: "compact" }), expect.any(AbortSignal)); expect(applyRefine).toHaveBeenCalledTimes(1); - expect(internals._compactAutoRefinePending).toBe(false); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); }); it("defers manual compaction refinement to the serialized checkpoint", async () => { @@ -1457,12 +1493,12 @@ describe("Serialized refine review-fix regressions", () => { firstKeptEntryId: "entry-1", tokensBefore: 100, }); - const interactiveSpy = vi.spyOn(internals, "_scheduleAutoRefine"); + const interactiveSpy = vi.spyOn(internals._refinement._auto, "_scheduleAutoRefine"); await harness.session.compact(); expect(interactiveSpy).not.toHaveBeenCalled(); - expect(internals._compactAutoRefinePending).toBe(true); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(true); }); it("preserves a serialized compaction trigger while cooldown is active", async () => { @@ -1475,13 +1511,13 @@ describe("Serialized refine review-fix regressions", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._lastAutoRefineReviewAt = Date.now(); - internals._scheduleAutoRefineAfterCompaction(true); + internals._refinement._auto._lastAutoRefineReviewAt = Date.now(); + internals._refinement._auto._scheduleAutoRefineAfterCompaction(true); - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(reviewer).not.toHaveBeenCalled(); - expect(internals._compactAutoRefinePending).toBe(true); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(true); }); it("clears a serialized compaction trigger when disposal occurs during cooldown", async () => { @@ -1494,13 +1530,13 @@ describe("Serialized refine review-fix regressions", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._lastAutoRefineReviewAt = Date.now(); - internals._compactAutoRefinePending = true; + internals._refinement._auto._lastAutoRefineReviewAt = Date.now(); + internals._refinement._auto._compactAutoRefinePending = true; - await internals._drainPendingRefinementForDisposal(); + await internals._refinement._drainPendingRefinementForDisposal(); expect(reviewer).not.toHaveBeenCalled(); - expect(internals._compactAutoRefinePending).toBe(false); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); }); it("continues to the interval drain after a compact trigger hits cooldown", async () => { @@ -1515,14 +1551,14 @@ describe("Serialized refine review-fix regressions", () => { vi.spyOn(harness.session.settingsManager, "getAutoRefineSettings") .mockReturnValueOnce(settings) .mockReturnValue({ ...settings, cooldownMs: 0 }); - const checkpoint = vi.spyOn(internals, "_runSerializedRefineCheckpoint").mockResolvedValue(); - internals._lastAutoRefineReviewAt = Date.now(); - internals._assistantTurnsSinceAutoRefine = 1; - internals._compactAutoRefinePending = true; + const checkpoint = vi.spyOn(internals._refinement, "_runSerializedRefineCheckpoint").mockResolvedValue(); + internals._refinement._auto._lastAutoRefineReviewAt = Date.now(); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._compactAutoRefinePending = true; - await internals._drainPendingRefinementForDisposal(); + await internals._refinement._drainPendingRefinementForDisposal(); - expect(internals._compactAutoRefinePending).toBe(false); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); expect(checkpoint).toHaveBeenCalledOnce(); }); @@ -1541,17 +1577,17 @@ describe("Serialized refine review-fix regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; const { applyRefine } = mockSerializedRefine(harness); - internals._assistantTurnsSinceAutoRefine = 1; - internals._scheduleAutoRefineAfterCompaction(true); + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._scheduleAutoRefineAfterCompaction(true); - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(reviewer).toHaveBeenCalledWith( expect.objectContaining({ reason: "turn_interval" }), expect.any(AbortSignal), ); expect(applyRefine).toHaveBeenCalledTimes(1); - expect(internals._compactAutoRefinePending).toBe(false); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); }); it("clears a compact trigger after a review decides no refinement is needed", async () => { @@ -1565,14 +1601,14 @@ describe("Serialized refine review-fix regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; const { applyRefine } = mockSerializedRefine(harness); - internals._scheduleAutoRefineAfterCompaction(true); + internals._refinement._auto._scheduleAutoRefineAfterCompaction(true); - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(reviewer).toHaveBeenCalledWith(expect.objectContaining({ reason: "compact" }), expect.any(AbortSignal)); expect(applyRefine).not.toHaveBeenCalled(); - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); - expect(internals._compactAutoRefinePending).toBe(false); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); }); it("does not let a compact review failure block disposal", async () => { @@ -1583,13 +1619,13 @@ describe("Serialized refine review-fix regressions", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._compactAutoRefinePending = true; - vi.spyOn(internals, "_runSerializedAutoRefineReview").mockRejectedValue( + internals._refinement._auto._compactAutoRefinePending = true; + vi.spyOn(internals._refinement._auto, "_runSerializedAutoRefineReview").mockRejectedValue( new Error("unexpected compact review failure"), ); - await expect(internals._drainPendingRefinementForDisposal()).resolves.toBeUndefined(); - expect(internals._compactAutoRefinePending).toBe(false); + await expect(internals._refinement._drainPendingRefinementForDisposal()).resolves.toBeUndefined(); + expect(internals._refinement._auto._compactAutoRefinePending).toBe(false); }); it("retries a failed explicit background plan during disposal", async () => { @@ -1601,19 +1637,19 @@ describe("Serialized refine review-fix regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; const options = { instructions: "explicit recovery" }; - internals._serializedPlanInFlight = Promise.resolve({ + internals._refinement._serializedPlanInFlight = Promise.resolve({ status: "failure", explicit: true, options, - branchVersion: internals._autoRefineBranchVersion, + branchVersion: internals._refinement._auto._autoRefineBranchVersion, }); - const runSpy = vi.spyOn(internals, "_runSerializedRefine").mockResolvedValue(); + const runSpy = vi.spyOn(internals._refinement, "_runSerializedRefine").mockResolvedValue(); - await internals._drainPendingRefinementForDisposal(); + await internals._refinement._drainPendingRefinementForDisposal(); expect(runSpy).toHaveBeenCalledTimes(1); expect(runSpy).toHaveBeenCalledWith(options, "self"); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); }); @@ -1647,9 +1683,11 @@ describe("Serialized refine event-ordering integration", () => { const internals = harness.session as unknown as SerializedInternals; const planSpy = vi - .spyOn(internals, "_planRefine") + .spyOn(internals._refinement._execution, "_planRefine") .mockResolvedValue({ id: "p", proposal: { edits: [] } } as never); - const applySpy = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const applySpy = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); // Send a prompt that produces a text response (no tool calls). // This goes through the real agent loop: agent_start -> turn_start -> @@ -1658,12 +1696,12 @@ describe("Serialized refine event-ordering integration", () => { harness.setResponses([fauxAssistantMessage("response 1")]); // Track counter before and after prompt - const counterBefore = internals._assistantTurnsSinceAutoRefine; + const counterBefore = internals._refinement._auto._assistantTurnsSinceAutoRefine; await harness.session.prompt("test prompt"); // After the first turn, the counter should be incremented by the // real message_end handler (not by direct mutation). - expect(internals._assistantTurnsSinceAutoRefine).toBe(counterBefore + 1); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(counterBefore + 1); // Now send a second prompt to reach the interval (turnInterval=2). harness.setResponses([fauxAssistantMessage("response 2")]); @@ -1674,7 +1712,7 @@ describe("Serialized refine event-ordering integration", () => { // _shouldStopAfterTurn ensured the counter was incremented before // the checkpoint checked the threshold. expect(applySpy).toHaveBeenCalled(); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); // The reviewer was called exactly once (during background planning // or synchronous boundary review, not duplicated). @@ -1716,12 +1754,12 @@ describe("P0 concurrency regressions", () => { let planResolved = false; let applyFinished = false; - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { await new Promise((resolve) => setTimeout(resolve, 20)); planResolved = true; return { id: "p", proposal: { edits: [] } }; }); - vi.spyOn(internals, "_applyRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_applyRefine").mockImplementation(async () => { expect(planResolved).toBe(true); await new Promise((resolve) => setTimeout(resolve, 10)); applyFinished = true; @@ -1741,10 +1779,8 @@ describe("P0 concurrency regressions", () => { }); // Start background planning at message_end - internals._assistantTurnsSinceAutoRefine++; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine++; + internals._refinement._maybeStartSerializedBackgroundPlan(); // Wait for plan to start but not finish await new Promise((resolve) => setTimeout(resolve, 10)); @@ -1755,7 +1791,7 @@ describe("P0 concurrency regressions", () => { expect(compactionSpy).toHaveBeenCalledTimes(1); expect(result).toBe(true); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); it("branch navigation with in-flight serialized plan: aborts signal, no apply", async () => { const harness = await createHarness({ @@ -1771,42 +1807,46 @@ describe("P0 concurrency regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - const applySpy = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const applySpy = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); // Make _planRefine block on its AbortSignal so we can test abort behavior. let planSignal: AbortSignal | undefined; - vi.spyOn(internals, "_planRefine").mockImplementation(async (_opts: unknown, signal: AbortSignal) => { - planSignal = signal; - // Block until the signal aborts, then throw. - if (signal.aborted) throw new Error("Plan aborted by branch change"); - await new Promise((_, reject) => { - signal.addEventListener("abort", () => reject(new Error("aborted"))); - }).catch(() => undefined); - if (signal.aborted) throw new Error("Plan aborted by branch change"); - return { id: "p", proposal: { edits: [] } }; - }); + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation( + async (_opts: unknown, signal: AbortSignal) => { + planSignal = signal; + // Block until the signal aborts, then throw. + if (signal.aborted) throw new Error("Plan aborted by branch change"); + await new Promise((_, reject) => { + signal.addEventListener("abort", () => reject(new Error("aborted"))); + }).catch(() => undefined); + if (signal.aborted) throw new Error("Plan aborted by branch change"); + return { id: "p", proposal: { edits: [] } }; + }, + ); // Start background planning - internals._assistantTurnsSinceAutoRefine++; - ( - internals as unknown as { _maybeStartSerializedBackgroundPlan: () => void } - )._maybeStartSerializedBackgroundPlan(); + internals._refinement._auto._assistantTurnsSinceAutoRefine++; + internals._refinement._maybeStartSerializedBackgroundPlan(); // Wait for the plan to start blocking await new Promise((resolve) => setTimeout(resolve, 20)); - expect(internals._serializedPlanInFlight).toBeDefined(); + expect(internals._refinement._serializedPlanInFlight).toBeDefined(); expect(planSignal).toBeDefined(); // Call _invalidatePendingAutoRefineForBranchChange while plan is pending. // This increments _autoRefineBranchVersion and awaits _serializedPlanInFlight. const invalidatePromise = ( harness.session as unknown as { - _invalidatePendingAutoRefineForBranchChange: () => Promise; + _refinement: { + _invalidatePendingAutoRefineForBranchChange: () => Promise; + }; } - )._invalidatePendingAutoRefineForBranchChange(); + )._refinement._invalidatePendingAutoRefineForBranchChange(); // The branchVersion was incremented, invalidating the plan. - const branchVersionAfter = internals._autoRefineBranchVersion; + const branchVersionAfter = internals._refinement._auto._autoRefineBranchVersion; expect(branchVersionAfter).toBeGreaterThan(0); // Wait for invalidation to complete (the plan promise settles as "failure" @@ -1814,10 +1854,10 @@ describe("P0 concurrency regressions", () => { await invalidatePromise; // _serializedPlanInFlight is cleared. - expect(internals._serializedPlanInFlight).toBeUndefined(); + expect(internals._refinement._serializedPlanInFlight).toBeUndefined(); // Run the checkpoint — no background plan to apply (cleared by invalidation). - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // _applyRefine was NOT called. expect(applySpy).not.toHaveBeenCalled(); @@ -1866,7 +1906,7 @@ describe("P0 concurrency regressions", () => { ], }, }; - vi.spyOn(internals, "_planRefine").mockResolvedValue(fauxPlan as never); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue(fauxPlan as never); // Spy on _rebuildSystemPrompt (call-through) to assert it was invoked. const rebuildSpy = vi.spyOn( @@ -1885,7 +1925,7 @@ describe("P0 concurrency regressions", () => { // Run the serialized refine (real _applyRefine runs). const promptBefore = harness.session.agent.state.systemPrompt; - await internals._runSerializedRefine({ instructions: "add a memory" }, "self"); + await internals._refinement._runSerializedRefine({ instructions: "add a memory" }, "self"); // The cache pin: applying a refinement never rebuilds or swaps the prompt. expect(rebuildSpy).not.toHaveBeenCalled(); @@ -1896,12 +1936,10 @@ describe("P0 concurrency regressions", () => { expect(getMessageText(notice)).toMatch(/^\[self-refinement\]\n\n/); // Harness state persisted to disk. - const localDir = (await import("../../src/core/refinement/index.js")).getLocalHarnessStateDir( - harness.sessionManager.getSessionArtifactDir(), - ); + const localDir = getLocalHarnessStateDir(harness.sessionManager.getSessionArtifactDir()); expect(localDir).toBeDefined(); if (localDir) { - const state = (await import("../../src/core/refinement/index.js")).loadHarnessState(localDir, "local"); + const state = loadHarnessState(localDir, "local"); const memoryEntries = Object.values(state.entries.memory ?? {}); const memoryEntry = memoryEntries.find((m) => m.title === "P0 concurrency test memory"); expect(memoryEntry).toBeDefined(); @@ -1923,8 +1961,10 @@ describe("P0 concurrency regressions", () => { // Make planning fail for the explicit refine.run request. const planError = new Error("planning network failure"); - vi.spyOn(internals, "_planRefine").mockRejectedValue(planError); - const applySpy = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockRejectedValue(planError); + const applySpy = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); // Queue a refine.run request — this starts background planning. (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; @@ -1935,16 +1975,16 @@ describe("P0 concurrency regressions", () => { await new Promise((resolve) => setTimeout(resolve, 50)); // Run the checkpoint — the background plan failed ("failure" result). - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); // _applyRefine was NOT called (planning failed). expect(applySpy).not.toHaveBeenCalled(); // Cooldown was stamped (failure -> stamp cooldown, no retry). - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(0); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(0); // _serializedPlanInFlight is cleared. - expect(internals._serializedPlanInFlight).toBeUndefined(); + expect(internals._refinement._serializedPlanInFlight).toBeUndefined(); }); it("serialized same-entry Python harness-write: concurrent kernel write rejected via baselineState", async () => { @@ -1961,9 +2001,6 @@ describe("P0 concurrency regressions", () => { const internals = harness.session as unknown as SerializedInternals; // Seed a memory entry on disk so the plan has something to update. - const { getLocalHarnessStateDir, loadHarnessState, saveHarnessState } = await import( - "../../src/core/refinement/index.js" - ); const localDir = getLocalHarnessStateDir(harness.sessionManager.getSessionArtifactDir()); expect(localDir).toBeDefined(); if (!localDir) return; @@ -2000,7 +2037,7 @@ describe("P0 concurrency regressions", () => { // Capture the baseline state at planning time (before the LLM call would run) // so the real _applyRefine can compare it against the current on-disk state. const baselineState = loadHarnessState(localDir, "local"); - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { planStarted?.(); await planGate; return { @@ -2024,7 +2061,10 @@ describe("P0 concurrency regressions", () => { }); // Start the serialized refine — begins background planning. - const refinePromise = internals._runSerializedRefine({ instructions: "update shared memory" }, "self"); + const refinePromise = internals._refinement._runSerializedRefine( + { instructions: "update shared memory" }, + "self", + ); // Wait for planning to start. await planStartedPromise; @@ -2055,14 +2095,19 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - refine: (options: { instructions?: string }) => Promise; + _refinement: { + refine: (options: { instructions?: string }) => Promise; + }; }; // Seed a pending explicit refine.run request. - internals._pendingRequestedRefine = { instructions: "leaked-explicit" }; - const refineSpy = vi.spyOn(internals, "refine").mockResolvedValue(undefined); - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "plan", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + internals._refinement._pendingRequestedRefine = { instructions: "leaked-explicit" }; + const refineSpy = vi.spyOn(internals._refinement, "refine").mockResolvedValue(undefined); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "plan", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); // Simulate an aborted assistant message arriving at agent_end. // Do NOT mock _checkCompaction — the real aborted path must clear the pending refine. @@ -2073,7 +2118,7 @@ describe("P0 concurrency regressions", () => { await new Promise((resolve) => setTimeout(resolve, 20)); // _checkCompaction's aborted block cleared _pendingRequestedRefine. - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); // The non-serialized agent_end path did NOT call refine (pending was already cleared). expect(refineSpy).not.toHaveBeenCalled(); }); @@ -2082,9 +2127,11 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - refine: (options: { instructions?: string }) => Promise; + _refinement: { + refine: (options: { instructions?: string }) => Promise; + }; }; - const refineSpy = vi.spyOn(internals, "refine").mockResolvedValue(undefined); + const refineSpy = vi.spyOn(internals._refinement, "refine").mockResolvedValue(undefined); (harness.session.agent.state as { isStreaming: boolean }).isStreaming = true; expect( @@ -2092,10 +2139,10 @@ describe("P0 concurrency regressions", () => { instructions: "must not survive cancellation", }), ).toMatchObject({ scheduled: true }); - expect(internals._pendingRequestedRefine).toBeDefined(); + expect(internals._refinement._pendingRequestedRefine).toBeDefined(); harness.session.requestAbort(); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); (harness.session.agent.state as { isStreaming: boolean }).isStreaming = false; const toolUseAssistant = fauxAssistantMessage([fauxToolCall("ipython", { code: "await refine.run()" })], { @@ -2113,9 +2160,11 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); let planSignal: AbortSignal | undefined; - vi.spyOn(internals, "_planRefine").mockImplementation( + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation( (_options: { instructions?: string }, signal: AbortSignal) => { planSignal = signal; return new Promise((_, reject) => { @@ -2129,17 +2178,17 @@ describe("P0 concurrency regressions", () => { instructions: "must not survive cancellation", }); await vi.waitFor(() => expect(planSignal).toBeDefined()); - expect(internals._serializedPlanInFlight).toBeDefined(); - const branchVersion = internals._autoRefineBranchVersion; + expect(internals._refinement._serializedPlanInFlight).toBeDefined(); + const branchVersion = internals._refinement._auto._autoRefineBranchVersion; harness.session.requestAbort(); expect(planSignal?.aborted).toBe(true); - expect(internals._autoRefineBranchVersion).toBeGreaterThan(branchVersion); + expect(internals._refinement._auto._autoRefineBranchVersion).toBeGreaterThan(branchVersion); (harness.session.agent.state as { isStreaming: boolean }).isStreaming = false; - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(applyRefine).not.toHaveBeenCalled(); - expect(internals._serializedPlanInFlight).toBeUndefined(); + expect(internals._refinement._serializedPlanInFlight).toBeUndefined(); }); it("explicit abort cancels an in-flight public refinement plan", async () => { @@ -2147,7 +2196,7 @@ describe("P0 concurrency regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; let planSignal: AbortSignal | undefined; - vi.spyOn(internals, "_planRefine").mockImplementation( + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation( (_options: { instructions?: string }, signal: AbortSignal) => { planSignal = signal; return new Promise((_, reject) => { @@ -2169,9 +2218,11 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); let planSignal: AbortSignal | undefined; - vi.spyOn(internals, "_planRefine").mockImplementation( + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation( (_options: { instructions?: string }, signal: AbortSignal) => { planSignal = signal; return new Promise((_, reject) => { @@ -2180,17 +2231,17 @@ describe("P0 concurrency regressions", () => { }, ); - const refine = internals._runSerializedRefine({ instructions: "cancel direct plan" }, "self"); + const refine = internals._refinement._runSerializedRefine({ instructions: "cancel direct plan" }, "self"); await vi.waitFor(() => expect(planSignal).toBeDefined()); - expect(internals._serializedPlanInFlight).toBeUndefined(); - expect(internals._refinePlanInFlight).toBeDefined(); - expect(internals._refineInFlight).toBeUndefined(); - const branchVersion = internals._autoRefineBranchVersion; + expect(internals._refinement._serializedPlanInFlight).toBeUndefined(); + expect(internals._refinement._refinePlanInFlight).toBeDefined(); + expect(internals._refinement._refineInFlight).toBeUndefined(); + const branchVersion = internals._refinement._auto._autoRefineBranchVersion; harness.session.requestAbort(); expect(planSignal?.aborted).toBe(true); - expect(internals._autoRefineBranchVersion).toBeGreaterThan(branchVersion); + expect(internals._refinement._auto._autoRefineBranchVersion).toBeGreaterThan(branchVersion); await expect(refine).rejects.toThrow("cancelled"); expect(applyRefine).not.toHaveBeenCalled(); }); @@ -2206,7 +2257,7 @@ describe("P0 concurrency regressions", () => { let planCalls = 0; let activePlans = 0; let maxActivePlans = 0; - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { planCalls++; activePlans++; maxActivePlans = Math.max(maxActivePlans, activePlans); @@ -2216,10 +2267,12 @@ describe("P0 concurrency regressions", () => { activePlans--; return { id: `plan-${planCalls}`, proposal: { edits: [] } }; }); - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); - const directRefine = internals._runSerializedRefine({ instructions: "direct" }, "self"); - await vi.waitFor(() => expect(internals._refinePlanInFlight).toBeDefined()); + const directRefine = internals._refinement._runSerializedRefine({ instructions: "direct" }, "self"); + await vi.waitFor(() => expect(internals._refinement._refinePlanInFlight).toBeDefined()); const publicRefine = harness.session.refine({ instructions: "public" }); await new Promise((resolve) => setTimeout(resolve, 20)); @@ -2239,84 +2292,88 @@ describe("P0 concurrency regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; let reviewSignal: AbortSignal | undefined; - vi.spyOn(internals, "_reviewAutoRefine").mockImplementation((_context, signal) => { + vi.spyOn(internals._refinement._auto, "_reviewAutoRefine").mockImplementation((_context, signal) => { reviewSignal = signal; return new Promise((_, reject) => { signal?.addEventListener("abort", () => reject(new Error("cancelled")), { once: true }); }); }); - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); - const branchVersion = internals._autoRefineBranchVersion; - const review = internals._runSerializedAutoRefineReview("turn_interval", branchVersion); + const branchVersion = internals._refinement._auto._autoRefineBranchVersion; + const review = internals._refinement._auto._runSerializedAutoRefineReview("turn_interval", branchVersion); await vi.waitFor(() => expect(reviewSignal).toBeDefined()); - expect(internals._autoRefineReviewAbort?.signal).toBe(reviewSignal); + expect(internals._refinement._auto._autoRefineReviewAbort?.signal).toBe(reviewSignal); harness.session.requestAbort(); expect(reviewSignal?.aborted).toBe(true); - expect(internals._autoRefineBranchVersion).toBeGreaterThan(branchVersion); + expect(internals._refinement._auto._autoRefineBranchVersion).toBeGreaterThan(branchVersion); await expect(review).resolves.toBeUndefined(); expect(applyRefine).not.toHaveBeenCalled(); - expect(internals._autoRefineReviewAbort).toBeUndefined(); + expect(internals._refinement._auto._autoRefineReviewAbort).toBeUndefined(); }); it("does not stamp cooldown when the branch changes during serialized refinement", async () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._assistantTurnsSinceAutoRefine = 4; - internals._lastAutoRefineReviewAt = 0; - vi.spyOn(internals, "_reviewAutoRefine").mockResolvedValue({ + internals._refinement._auto._assistantTurnsSinceAutoRefine = 4; + internals._refinement._auto._lastAutoRefineReviewAt = 0; + vi.spyOn(internals._refinement._auto, "_reviewAutoRefine").mockResolvedValue({ shouldRefine: true, rationale: "capture lesson", instructions: "capture it", }); - vi.spyOn(internals, "_runSerializedRefine").mockImplementation(async () => { - internals._autoRefineBranchVersion++; + vi.spyOn(internals._refinement, "_runSerializedRefine").mockImplementation(async () => { + internals._refinement._auto._autoRefineBranchVersion++; }); - const branchVersion = internals._autoRefineBranchVersion; + const branchVersion = internals._refinement._auto._autoRefineBranchVersion; - await internals._runSerializedAutoRefineReview("turn_interval", branchVersion); + await internals._refinement._auto._runSerializedAutoRefineReview("turn_interval", branchVersion); - expect(internals._lastAutoRefineReviewAt).toBe(0); - expect(internals._assistantTurnsSinceAutoRefine).toBe(4); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(4); }); it("services a pending request after an invalidated background plan", async () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._serializedPlanInFlight = Promise.resolve({ status: "invalidated" }); - internals._pendingRequestedRefine = { instructions: "latest" }; - const run = vi.spyOn(internals, "_runSerializedRefine").mockResolvedValue(); + internals._refinement._serializedPlanInFlight = Promise.resolve({ status: "invalidated" }); + internals._refinement._pendingRequestedRefine = { instructions: "latest" }; + const run = vi.spyOn(internals._refinement, "_runSerializedRefine").mockResolvedValue(); - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(run).toHaveBeenCalledWith({ instructions: "latest" }, "self"); - expect(internals._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); }); it("records cooldown after dropping a stale ready background plan", async () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - const apply = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); - const reviewAtBeforeCheckpoint = internals._lastAutoRefineReviewAt; - internals._assistantTurnsSinceAutoRefine = 1; - internals._serializedPlanInFlight = Promise.resolve({ + const apply = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); + const reviewAtBeforeCheckpoint = internals._refinement._auto._lastAutoRefineReviewAt; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; + internals._refinement._serializedPlanInFlight = Promise.resolve({ status: "plan", plan: { id: "stale-plan", proposal: { edits: [] } }, options: {}, abort: new AbortController(), - branchVersion: internals._autoRefineBranchVersion - 1, + branchVersion: internals._refinement._auto._autoRefineBranchVersion - 1, }); - await internals._runSerializedRefineCheckpoint(); + await internals._refinement._runSerializedRefineCheckpoint(); expect(apply).not.toHaveBeenCalled(); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); - expect(internals._lastAutoRefineReviewAt).toBeGreaterThan(reviewAtBeforeCheckpoint); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._lastAutoRefineReviewAt).toBeGreaterThan(reviewAtBeforeCheckpoint); }); it("aborted serialized turn clears _pendingRequestedRefine so it does not leak to next checkpoint", async () => { @@ -2332,20 +2389,31 @@ describe("P0 concurrency regressions", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - refine: (options: { instructions?: string }) => Promise; + _refinement: { + refine: (options: { instructions?: string }) => Promise; + }; }; // Seed a pending explicit refine.run request. - internals._pendingRequestedRefine = { instructions: "leaked-explicit" }; - const refineSpy = vi.spyOn(internals, "refine").mockResolvedValue(undefined); - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "plan", proposal: { edits: [] } }); - const applyRefine = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + internals._refinement._pendingRequestedRefine = { instructions: "leaked-explicit" }; + const refineSpy = vi.spyOn(internals._refinement, "refine").mockResolvedValue(undefined); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "plan", + proposal: { edits: [] }, + }); + const applyRefine = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); const stalePlanAbort = new AbortController(); - internals._refineAbortController = stalePlanAbort; - internals._serializedPlanInFlight = new Promise((resolve) => { + internals._refinement._refineAbortController = stalePlanAbort; + internals._refinement._serializedPlanInFlight = new Promise((resolve) => { stalePlanAbort.signal.addEventListener( "abort", - () => resolve({ status: "invalidated", branchVersion: internals._autoRefineBranchVersion - 1 }), + () => + resolve({ + status: "invalidated", + branchVersion: internals._refinement._auto._autoRefineBranchVersion - 1, + }), { once: true }, ); }); @@ -2359,14 +2427,14 @@ describe("P0 concurrency regressions", () => { await new Promise((resolve) => setTimeout(resolve, 20)); // _checkCompaction's aborted block cleared _pendingRequestedRefine. - expect(internals._pendingRequestedRefine).toBeUndefined(); - expect(internals._serializedPlanInFlight).toBeUndefined(); + expect(internals._refinement._pendingRequestedRefine).toBeUndefined(); + expect(internals._refinement._serializedPlanInFlight).toBeUndefined(); expect(stalePlanAbort.signal.aborted).toBe(true); expect(refineSpy).not.toHaveBeenCalled(); // Run the next serialized checkpoint; no plan or apply should fire. - await internals._runSerializedRefineCheckpoint(); - expect(internals._planRefine).not.toHaveBeenCalled(); + await internals._refinement._runSerializedRefineCheckpoint(); + expect(internals._refinement._execution._planRefine).not.toHaveBeenCalled(); expect(applyRefine).not.toHaveBeenCalled(); }); @@ -2383,8 +2451,11 @@ describe("P0 concurrency regressions", () => { await harness.session.queueAgentMessagePrompt("queued follow-up", "followUp"); expect(harness.session.queuedActionCount).toBe(1); - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "plan", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "plan", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); await harness.session.refine({ instructions: "public" }); @@ -2399,8 +2470,8 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._pendingRequestedRefine = { instructions: "fail planning" }; - vi.spyOn(internals, "_planRefine").mockRejectedValue(new Error("planner unavailable")); + internals._refinement._pendingRequestedRefine = { instructions: "fail planning" }; + vi.spyOn(internals._refinement._execution, "_planRefine").mockRejectedValue(new Error("planner unavailable")); const failed = new Promise((resolve) => { const unsubscribe = harness.session.subscribe((event) => { if (event.type === "refine_failed") { @@ -2410,7 +2481,7 @@ describe("P0 concurrency regressions", () => { }); }); - await expect(internals._runSerializedRefineCheckpoint()).resolves.toBeUndefined(); + await expect(internals._refinement._runSerializedRefineCheckpoint()).resolves.toBeUndefined(); expect(await failed).toBe("planner unavailable"); }); @@ -2418,9 +2489,12 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._pendingRequestedRefine = { instructions: "fail at apply" }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "plan", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockRejectedValue(new Error("disk full")); + internals._refinement._pendingRequestedRefine = { instructions: "fail at apply" }; + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "plan", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockRejectedValue(new Error("disk full")); const failed = new Promise((resolve) => { const unsubscribe = harness.session.subscribe((event) => { if (event.type === "refine_failed") { @@ -2430,7 +2504,7 @@ describe("P0 concurrency regressions", () => { }); }); - await expect(internals._runSerializedRefineCheckpoint()).resolves.toBeUndefined(); + await expect(internals._refinement._runSerializedRefineCheckpoint()).resolves.toBeUndefined(); expect(await failed).toBe("disk full"); }); @@ -2439,7 +2513,7 @@ describe("P0 concurrency regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; let resolvePlan: (value: unknown) => void = () => {}; - internals._serializedPlanInFlight = new Promise((resolve) => { + internals._refinement._serializedPlanInFlight = new Promise((resolve) => { resolvePlan = resolve; }); let releaseProcessing: () => void = () => {}; @@ -2453,9 +2527,9 @@ describe("P0 concurrency regressions", () => { }); const secondConsumer = vi.fn(async () => false); - const first = internals._consumeSerializedBackgroundPlan(firstConsumer); - const second = internals._consumeSerializedBackgroundPlan(secondConsumer); - expect(internals._serializedPlanInFlight).toBeDefined(); + const first = internals._refinement._consumeSerializedBackgroundPlan(firstConsumer); + const second = internals._refinement._consumeSerializedBackgroundPlan(secondConsumer); + expect(internals._refinement._serializedPlanInFlight).toBeDefined(); resolvePlan({ status: "skip" }); await vi.waitFor(() => expect(firstConsumer).toHaveBeenCalledOnce()); @@ -2477,27 +2551,27 @@ describe("P0 concurrency regressions", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; let resolvePlan: (value: unknown) => void = () => {}; - internals._serializedPlanInFlight = new Promise((resolve) => { + internals._refinement._serializedPlanInFlight = new Promise((resolve) => { resolvePlan = resolve; }); - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; let releaseApply: () => void = () => {}; const applyBlocked = new Promise((resolve) => { releaseApply = resolve; }); - const apply = vi.spyOn(internals, "_applyRefine").mockImplementation(async () => { + const apply = vi.spyOn(internals._refinement._execution, "_applyRefine").mockImplementation(async () => { await applyBlocked; return emptyRefinementResult(); }); - const checkpoint = internals._runSerializedRefineCheckpoint(); - const drain = internals._drainPendingRefinementForDisposal(); + const checkpoint = internals._refinement._runSerializedRefineCheckpoint(); + const drain = internals._refinement._drainPendingRefinementForDisposal(); resolvePlan({ status: "plan", plan: { id: "claimed-plan", proposal: { edits: [] } }, options: {}, abort: new AbortController(), - branchVersion: internals._autoRefineBranchVersion, + branchVersion: internals._refinement._auto._autoRefineBranchVersion, }); await vi.waitFor(() => expect(apply).toHaveBeenCalledOnce()); @@ -2511,7 +2585,7 @@ describe("P0 concurrency regressions", () => { releaseApply(); await Promise.all([checkpoint, drain]); expect(apply).toHaveBeenCalledOnce(); - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); }); it("drains a due interactive auto-refine without waiting for agent idle", async () => { @@ -2521,11 +2595,11 @@ describe("P0 concurrency regressions", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; const waitForIdle = vi.spyOn(harness.session.agent, "waitForIdle").mockResolvedValue(); - const maybeAutoRefine = vi.spyOn(internals, "_maybeAutoRefine").mockResolvedValue(); + const maybeAutoRefine = vi.spyOn(internals._refinement._auto, "_maybeAutoRefine").mockResolvedValue(); - await internals._drainPendingRefinementForDisposal(); + await internals._refinement._drainPendingRefinementForDisposal(); expect(waitForIdle).not.toHaveBeenCalled(); expect(maybeAutoRefine).toHaveBeenCalledWith("turn_interval"); @@ -2539,9 +2613,9 @@ describe("P0 concurrency regressions", () => { const operation = new Promise((resolve) => { releaseOperation = resolve; }); - internals._autoRefineOperations.add(operation); + internals._refinement._auto._autoRefineOperations.add(operation); let settled = false; - const drain = internals._drainPendingRefinementForDisposal().then(() => { + const drain = internals._refinement._drainPendingRefinementForDisposal().then(() => { settled = true; }); @@ -2560,11 +2634,11 @@ describe("P0 concurrency regressions", () => { const drain = new Promise((resolve) => { releaseDrain = resolve; }); - vi.spyOn(internals, "_drainPendingRefinementForDisposal").mockReturnValue(drain); + vi.spyOn(internals._refinement, "_drainPendingRefinementForDisposal").mockReturnValue(drain); const first = harness.session.disposeAsync(); const second = harness.session.disposeAsync(); - expect(internals._drainPendingRefinementForDisposal).toHaveBeenCalledTimes(1); + expect(internals._refinement._drainPendingRefinementForDisposal).toHaveBeenCalledTimes(1); releaseDrain(); await Promise.all([first, second]); }); @@ -2573,19 +2647,26 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true, serializedRefine: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - _serializedExplicitRefineOptions?: { instructions?: string; global?: boolean }; + _refinement: { + _serializedExplicitRefineOptions?: { instructions?: string; global?: boolean }; + }; }; - internals._serializedPlanInFlight = Promise.resolve({ status: "skip" }); - internals._serializedExplicitRefineOptions = { instructions: "stale", global: true }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "public-plan", proposal: { edits: [] } }); - const apply = vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + internals._refinement._serializedPlanInFlight = Promise.resolve({ status: "skip" }); + internals._refinement._serializedExplicitRefineOptions = { instructions: "stale", global: true }; + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "public-plan", + proposal: { edits: [] }, + }); + const apply = vi + .spyOn(internals._refinement._execution, "_applyRefine") + .mockResolvedValue(emptyRefinementResult()); await expect(harness.session.refine({ instructions: "public after abort" })).resolves.toEqual( emptyRefinementResult(), ); - expect(internals._serializedPlanInFlight).toBeUndefined(); - expect(internals._serializedExplicitRefineOptions).toBeUndefined(); + expect(internals._refinement._serializedPlanInFlight).toBeUndefined(); + expect(internals._refinement._serializedExplicitRefineOptions).toBeUndefined(); expect(apply).toHaveBeenCalledOnce(); }); @@ -2597,8 +2678,11 @@ describe("P0 concurrency regressions", () => { _cancelActiveRlmChildRuns: () => void; }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "plan", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue(emptyRefinementResult()); + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "plan", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue(emptyRefinementResult()); const cancelChildRuns = vi.spyOn(internals, "_cancelActiveRlmChildRuns"); const requestAbort = vi.spyOn(harness.session, "requestAbort"); diff --git a/packages/coding-agent/test/suite/daemon-serialized-refine.test.ts b/packages/coding-agent/test/suite/daemon-serialized-refine.test.ts index 2cd468f9ab..2738217369 100644 --- a/packages/coding-agent/test/suite/daemon-serialized-refine.test.ts +++ b/packages/coding-agent/test/suite/daemon-serialized-refine.test.ts @@ -20,15 +20,25 @@ import type { DaemonCommand } from "../../src/modes/daemon/daemon-protocol.js"; import { createHarness, type Harness } from "./harness.js"; type SerializedInternals = { - _serializedRefine: boolean; + _refinement: { + _serializedRefine: boolean; + _auto: { + _assistantTurnsSinceAutoRefine: number; + _lastAutoRefineReviewAt: number; + _reviewAutoRefine: ( + ctx: { reason: string; turnsSinceLastReview: number }, + signal?: AbortSignal, + ) => Promise; + }; + _execution: { + _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; + _applyRefine: (plan: unknown, opts: unknown, abort: AbortController) => Promise; + }; + }; + _rlmHeartbeatController?: unknown; _agentMessageController?: unknown; _agentObserveController?: unknown; - _assistantTurnsSinceAutoRefine: number; - _lastAutoRefineReviewAt: number; - _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; - _applyRefine: (plan: unknown, opts: unknown, abort: AbortController) => Promise; - _reviewAutoRefine: (ctx: { reason: string; turnsSinceLastReview: number }, signal?: AbortSignal) => Promise; }; describe("Daemon-backed serializedRefine propagation", () => { @@ -81,7 +91,7 @@ describe("Daemon-backed serializedRefine propagation", () => { // The session created by the daemon has _serializedRefine=true. const sessionInternals = state.runtime.session as unknown as SerializedInternals; - expect(sessionInternals._serializedRefine).toBe(true); + expect(sessionInternals._refinement._serializedRefine).toBe(true); state.runtime.session.dispose(); }); @@ -120,7 +130,7 @@ describe("Daemon-backed serializedRefine propagation", () => { const state = await internals.createRuntime({ type: "create", sessionPath: sessionFile }); const sessionInternals = state.runtime.session as unknown as SerializedInternals; - expect(sessionInternals._serializedRefine).toBe(false); + expect(sessionInternals._refinement._serializedRefine).toBe(false); state.runtime.session.dispose(); }); @@ -148,8 +158,11 @@ describe("Daemon-backed serializedRefine propagation", () => { stashedHarness = harness; const session = harness.session; const internals = session as unknown as SerializedInternals; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue({ + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ + id: "p", + proposal: { edits: [] }, + }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue({ id: "refine_test", summary: "test", rationale: "test", @@ -186,14 +199,14 @@ describe("Daemon-backed serializedRefine propagation", () => { // Turn 1: counter goes to 1 (< threshold 2) harness.setResponses([fauxAssistantMessage("response 1")]); await session.prompt("prompt 1"); - expect(sessionInternals._assistantTurnsSinceAutoRefine).toBe(1); + expect(sessionInternals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(1); // Turn 2: counter goes to 2 (>= threshold), checkpoint fires harness.setResponses([fauxAssistantMessage("response 2")]); await session.prompt("prompt 2"); // After checkpoint, counter reset to 0 - expect(sessionInternals._assistantTurnsSinceAutoRefine).toBe(0); + expect(sessionInternals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); // Reviewer called exactly once expect(reviewer).toHaveBeenCalledTimes(1); diff --git a/packages/coding-agent/test/suite/regressions/2190-refinement-dispatch.test.ts b/packages/coding-agent/test/suite/regressions/2190-refinement-dispatch.test.ts new file mode 100644 index 0000000000..c23425119f --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/2190-refinement-dispatch.test.ts @@ -0,0 +1,157 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { AutoRefinement } from "../../../src/session/auto-refinement.js"; +import type { SessionRefinement, SessionRefinementHost } from "../../../src/session/refinement.js"; +import { createHarness, type Harness } from "../harness.js"; +import { withStreaming } from "../scheduling.js"; + +type RefinementInternals = Pick & { + _auto: AutoRefinement; + _host: SessionRefinementHost; +}; +function owner(harness: Harness): RefinementInternals { + return (harness.session as unknown as { _refinement: RefinementInternals })._refinement; +} + +describe("refinement public dispatch preservation", () => { + const harnesses: Harness[] = []; + afterEach(() => { + vi.restoreAllMocks(); + for (const harness of harnesses.splice(0)) harness.cleanup(); + }); + + for (const source of ["self", "auto"] as const) { + it(`dispatches ${source} refinement through a live public wrapper before planning`, async () => { + const order: string[] = []; + const harness = await createHarness({ + persistSession: true, + autoRefineReviewer: async () => { + order.push("review"); + return { shouldRefine: true, rationale: "test" }; + }, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", () => { + order.push("plan"); + return { proposal: { summary: "test", rationale: "test", expectedOutcome: "test", edits: [] } }; + }); + }, + ], + }); + harnesses.push(harness); + const refinement = owner(harness); + vi.spyOn(harness.settingsManager, "getAutoRefineSettings").mockReturnValue({ + ...harness.settingsManager.getAutoRefineSettings(), + enabled: true, + turnInterval: 1, + }); + const original = harness.session.refine.bind(harness.session); + const wrapper = vi.spyOn(harness.session, "refine").mockImplementation((options, internal) => { + order.push("wrapper"); + if (source === "self") expect(refinement._consumePendingRequestedRefine()).toBe(false); + return original(options, internal); + }); + if (source === "self") { + withStreaming(harness, true); + harness.session.handleRefineHostRequest("refine.run", { instructions: "test" }); + withStreaming(harness, false); + expect(refinement._consumePendingRequestedRefine()).toBe(true); + } else { + refinement._auto.observeAssistantEnd(); + await refinement._auto._maybeAutoRefine("turn_interval"); + } + expect(wrapper).toHaveBeenCalledOnce(); + await wrapper.mock.results[0].value; + expect(wrapper.mock.calls[0][1]).toEqual(source === "self" ? { source: "self" } : { trigger: "auto" }); + expect(order).toEqual(source === "self" ? ["wrapper", "plan"] : ["review", "wrapper", "plan"]); + }); + + it(`preserves ${source} public wrapper rejection handling`, async () => { + const harness = await createHarness({ + persistSession: true, + autoRefineReviewer: async () => ({ shouldRefine: true, rationale: "test" }), + }); + harnesses.push(harness); + const refinement = owner(harness); + vi.spyOn(harness.settingsManager, "getAutoRefineSettings").mockReturnValue({ + ...harness.settingsManager.getAutoRefineSettings(), + enabled: true, + turnInterval: 1, + }); + const failure = new Error("public wrapper rejected"); + const wrapper = vi.spyOn(harness.session, "refine").mockRejectedValue(failure); + const emit = vi.spyOn(refinement._host, "emit"); + if (source === "self") { + withStreaming(harness, true); + harness.session.handleRefineHostRequest("refine.run", {}); + withStreaming(harness, false); + expect(refinement._consumePendingRequestedRefine()).toBe(true); + await Promise.resolve(); + expect(emit).toHaveBeenCalledWith({ type: "refine_failed", error: failure.message }); + expect(refinement._consumePendingRequestedRefine()).toBe(false); + } else { + refinement._auto.observeAssistantEnd(); + await expect(refinement._auto._maybeAutoRefine("turn_interval")).resolves.toBeUndefined(); + expect(refinement._auto.lastReviewAt).toBeGreaterThan(0); + expect(refinement._auto.turnsSinceReview).toBe(1); + expect(emit).not.toHaveBeenCalled(); + } + expect(wrapper).toHaveBeenCalledOnce(); + }); + } + + it("settles the default no-model review before the next queued microtask", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const refinement = owner(harness); + vi.spyOn(refinement._host, "getModel").mockReturnValue(undefined); + const order: string[] = []; + const review = refinement._auto._reviewAutoRefine({ reason: "turn_interval", turnsSinceLastReview: 1 }); + void review.then(() => order.push("review")); + queueMicrotask(() => order.push("microtask")); + await expect(review).resolves.toEqual({ shouldRefine: false, rationale: "No model selected." }); + expect(order).toEqual(["review", "microtask"]); + }); + + it("preserves the custom reviewer's async promise adoption", async () => { + const result = { shouldRefine: false, rationale: "custom decline" }; + const inner = Promise.resolve(result); + const reviewer = vi.fn(() => inner); + const harness = await createHarness({ autoRefineReviewer: reviewer }); + harnesses.push(harness); + const order: string[] = []; + const review = owner(harness)._auto._reviewAutoRefine({ reason: "turn_interval", turnsSinceLastReview: 1 }); + expect(reviewer).toHaveBeenCalledOnce(); + expect(review).not.toBe(inner); + void review.then(() => order.push("review")); + queueMicrotask(() => order.push("microtask")); + await expect(review).resolves.toEqual(result); + expect(order).toEqual(["microtask", "review"]); + }); + + it("calls a non-arrow custom reviewer with the public session as its receiver", async () => { + let receiver: unknown; + const reviewer = vi.fn(function (this: unknown) { + receiver = this; + return Promise.resolve({ shouldRefine: false, rationale: "custom decline" }); + }); + const harness = await createHarness({ autoRefineReviewer: reviewer }); + harnesses.push(harness); + const context = { reason: "turn_interval" as const, turnsSinceLastReview: 1 }; + const signal = new AbortController().signal; + await owner(harness)._auto._reviewAutoRefine(context, signal); + expect(receiver === harness.session).toBe(true); + expect(reviewer).toHaveBeenCalledExactlyOnceWith(context, signal); + }); + + it("turns synchronous custom reviewer throws into promise rejections", async () => { + const failure = new Error("review failed"); + const harness = await createHarness({ + autoRefineReviewer: () => { + throw failure; + }, + }); + harnesses.push(harness); + const review = owner(harness)._auto._reviewAutoRefine({ reason: "turn_interval", turnsSinceLastReview: 1 }); + await expect(review).rejects.toBe(failure); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts b/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts index 4aa258a4f8..2064c7c52b 100644 --- a/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts +++ b/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts @@ -17,6 +17,7 @@ import { } from "../../../src/modes/interactive/components/injected-prompt-message.js"; import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.js"; import { getMarkdownTheme, initTheme } from "../../../src/modes/interactive/theme/theme.js"; +import type { SessionCompaction } from "../../../src/session/compaction.js"; import { conversationMessages, createHarness, getMessageText, getUserTexts, type Harness } from "../harness.js"; type AddMessageToChatHost = { @@ -188,14 +189,14 @@ describe("ENG-4482 heartbeat injected prompt UI", () => { it("resets overflow recovery state when heartbeat prompt turns start", async () => { const harness = await createHarness(); harnesses.push(harness); - const sessionInternals = harness.session as unknown as { _overflowRecovery: string }; - sessionInternals._overflowRecovery = "attempted"; + const sessionInternals = harness.session as unknown as { _compaction: SessionCompaction }; + sessionInternals._compaction.markOverflowAttempted(); harness.setResponses([fauxAssistantMessage("heartbeat handled")]); await harness.session.promptHeartbeat(createHeartbeat()); await harness.session.agent.waitForIdle(); - expect(sessionInternals._overflowRecovery).toBe("idle"); + expect(sessionInternals._compaction.overflowRecovery).toBe("idle"); }); it("keeps pending nextTurn context separate from queued heartbeat prompts", async () => { diff --git a/packages/coding-agent/test/suite/serialized-refine-config-integration.test.ts b/packages/coding-agent/test/suite/serialized-refine-config-integration.test.ts index 17d0d8d781..c9801f6c7c 100644 --- a/packages/coding-agent/test/suite/serialized-refine-config-integration.test.ts +++ b/packages/coding-agent/test/suite/serialized-refine-config-integration.test.ts @@ -12,11 +12,16 @@ import type { AgentRlmHeartbeatController } from "../../src/core/cron-jobs.js"; import { createHarness, type Harness } from "./harness.js"; type SerializedInternals = { - _serializedRefine: boolean; - _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; - _assistantTurnsSinceAutoRefine: number; - _lastAutoRefineReviewAt: number; - _autoRefineInProgress: boolean; + _refinement: { + _serializedRefine: boolean; + _pendingRequestedRefine: { instructions?: string; global?: boolean } | undefined; + _auto: { + _assistantTurnsSinceAutoRefine: number; + _lastAutoRefineReviewAt: number; + _autoRefineInProgress: boolean; + }; + }; + _rlmHeartbeatController?: unknown; _agentMessageController?: unknown; _agentObserveController?: unknown; @@ -64,7 +69,7 @@ describe("Serialized refine config integration (unit)", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - expect(internals._serializedRefine).toBe(true); + expect(internals._refinement._serializedRefine).toBe(true); }); it("serializedRefine=false (default) produces a session with _serializedRefine=false", async () => { @@ -74,7 +79,7 @@ describe("Serialized refine config integration (unit)", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - expect(internals._serializedRefine).toBe(false); + expect(internals._refinement._serializedRefine).toBe(false); }); it("real autonomous loop crosses threshold and resumes with serialized refine", async () => { @@ -94,12 +99,16 @@ describe("Serialized refine config integration (unit)", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; - _applyRefine: (plan: unknown, opts: unknown, abort: AbortController) => Promise; + _refinement: { + _execution: { + _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; + _applyRefine: (plan: unknown, opts: unknown, abort: AbortController) => Promise; + }; + }; }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue({ + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue({ id: "refine_test", summary: "test", rationale: "test", @@ -111,19 +120,19 @@ describe("Serialized refine config integration (unit)", () => { // Turn 1: counter goes to 1 (< threshold 2, no refine) harness.setResponses([fauxAssistantMessage("response 1")]); await harness.session.prompt("prompt 1"); - expect(internals._assistantTurnsSinceAutoRefine).toBe(1); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(1); // Turn 2: counter goes to 2 (>= threshold 2, serialized checkpoint fires) harness.setResponses([fauxAssistantMessage("response 2")]); await harness.session.prompt("prompt 2"); // After the checkpoint, counter is reset to 0. - expect(internals._assistantTurnsSinceAutoRefine).toBe(0); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(0); // Turn 3: counter goes to 1 again (< threshold, loop continues) harness.setResponses([fauxAssistantMessage("response 3")]); await harness.session.prompt("prompt 3"); - expect(internals._assistantTurnsSinceAutoRefine).toBe(1); + expect(internals._refinement._auto._assistantTurnsSinceAutoRefine).toBe(1); // Reviewer called exactly once (at turn 2). expect(reviewer).toHaveBeenCalledTimes(1); @@ -183,7 +192,11 @@ describe("Serialized refine controller availability (unit)", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; + _refinement: { + _execution: { + _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; + }; + }; }; // Initially not in flight. @@ -195,7 +208,7 @@ describe("Serialized refine controller availability (unit)", () => { const planPromise = new Promise((resolve) => { resolvePlan = resolve; }); - vi.spyOn(internals, "_planRefine").mockImplementation(async () => { + vi.spyOn(internals._refinement._execution, "_planRefine").mockImplementation(async () => { await planPromise; return { id: "p", proposal: { edits: [] } }; }); @@ -216,8 +229,8 @@ describe("Serialized refine controller availability (unit)", () => { // Run the checkpoint to consume the background plan. await ( - harness.session as unknown as { _runSerializedRefineCheckpoint: () => Promise } - )._runSerializedRefineCheckpoint(); + harness.session as unknown as { _refinement: { _runSerializedRefineCheckpoint: () => Promise } } + )._refinement._runSerializedRefineCheckpoint(); // After the checkpoint, nothing should be in flight. const statusAfter = harness.session.handleRefineHostRequest("refine.status"); @@ -262,12 +275,16 @@ describe("PR #503 model preservation (unit)", () => { harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; - _applyRefine: (plan: unknown, opts: unknown, abort: AbortController) => Promise; + _refinement: { + _execution: { + _planRefine: (opts: { instructions?: string }, signal: AbortSignal) => Promise; + _applyRefine: (plan: unknown, opts: unknown, abort: AbortController) => Promise; + }; + }; }; - vi.spyOn(internals, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); - vi.spyOn(internals, "_applyRefine").mockResolvedValue({ + vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ id: "p", proposal: { edits: [] } }); + vi.spyOn(internals._refinement._execution, "_applyRefine").mockResolvedValue({ id: "refine_test", summary: "test", rationale: "test", @@ -277,7 +294,7 @@ describe("PR #503 model preservation (unit)", () => { }); const modelBefore = harness.session.model; - internals._assistantTurnsSinceAutoRefine = 1; + internals._refinement._auto._assistantTurnsSinceAutoRefine = 1; await harness.session.prompt("test"); // Wait for the checkpoint to complete. diff --git a/packages/coding-agent/test/suite/session-refinement-owner.test.ts b/packages/coding-agent/test/suite/session-refinement-owner.test.ts new file mode 100644 index 0000000000..4d628a76e0 --- /dev/null +++ b/packages/coding-agent/test/suite/session-refinement-owner.test.ts @@ -0,0 +1,179 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { RefinementPlan, RefinementResult } from "../../src/core/refinement/index.js"; +import type { AutoRefinement } from "../../src/session/auto-refinement.js"; +import type { + SerializedBackgroundPlanResult, + SessionRefinement, + SessionRefinementHost, +} from "../../src/session/refinement.js"; +import type { RefinementExecution } from "../../src/session/refinement-execution.js"; +import { createHarness, type Harness } from "./harness.js"; +import { createDeferred, withStreaming } from "./scheduling.js"; + +type RefinementInternals = Pick & { + _host: SessionRefinementHost; + _auto: AutoRefinement; + _execution: RefinementExecution; + _serializedPlanInFlight?: Promise; +}; + +function owner(harness: Harness): RefinementInternals { + return (harness.session as unknown as { _refinement: RefinementInternals })._refinement; +} + +function plan(): RefinementPlan { + return { + id: "owner-plan", + proposal: { summary: "owner test", rationale: "test", expectedOutcome: "test", edits: [] }, + }; +} + +function result(): RefinementResult { + return { + id: "owner-result", + summary: "test", + rationale: "test", + expectedOutcome: "test", + appliedEdits: [], + harnessStatePath: "/tmp/owner-state.json", + }; +} + +describe("SessionRefinement ownership boundary", () => { + const harnesses: Harness[] = []; + afterEach(() => { + vi.restoreAllMocks(); + while (harnesses.length) harnesses.pop()?.cleanup(); + }); + + it("preserves the owner's public promise and synchronous host-request errors", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const refinement = owner(harness); + const gate = createDeferred(); + const refine = vi.spyOn(refinement, "refine").mockReturnValue(gate.promise); + expect(harness.session.refine({ instructions: "capture" })).toBe(gate.promise); + expect(refine).toHaveBeenCalledWith({ instructions: "capture" }, {}); + expect(() => harness.session.handleRefineHostRequest("refine.run", { instructions: 1 })).toThrow( + "refine.run instructions must be a string", + ); + gate.resolve(result()); + await gate.promise; + }); + + it("invalidates before abort and returns the exact plan for caller-owned cleanup settlement", async () => { + const harness = await createHarness({ persistSession: true, serializedRefine: true }); + harnesses.push(harness); + const refinement = owner(harness); + expect(refinement.beginAbortedTurnCleanup()).toBeUndefined(); + const planning = createDeferred(); + let versionAtAbort: number | undefined; + const initialVersion = refinement._auto.branchVersion; + vi.spyOn(refinement._execution, "_planRefine").mockImplementation((_options, signal) => { + signal.addEventListener("abort", () => { + versionAtAbort = refinement._auto.branchVersion; + }); + return planning.promise; + }); + withStreaming(harness, true); + harness.session.handleRefineHostRequest("refine.run", { instructions: "original" }); + withStreaming(harness, false); + const original = refinement._serializedPlanInFlight; + const cleanup = refinement.beginAbortedTurnCleanup(); + expect(cleanup?.promise).toBe(original); + expect(versionAtAbort).toBe(initialVersion + 1); + expect(refinement._serializedPlanInFlight).toBe(original); + planning.resolve(plan()); + await expect(cleanup?.promise).resolves.toEqual({ status: "invalidated", branchVersion: initialVersion }); + const replacement = Promise.resolve({ status: "skip" }); + refinement._serializedPlanInFlight = replacement; + cleanup?.finish(); + expect(refinement._serializedPlanInFlight).toBe(replacement); + refinement._serializedPlanInFlight = undefined; + }); + + it("keeps planning nonblocking and rechecks operations added while waiting to apply", async () => { + const harness = await createHarness({ persistSession: true }); + harnesses.push(harness); + const refinement = owner(harness); + const planning = createDeferred(); + const firstOperation = createDeferred(); + const secondOperation = createDeferred(); + let operation = firstOperation.promise; + vi.spyOn(refinement._host, "getCompactionOperation").mockImplementation(() => operation); + vi.spyOn(refinement._execution, "_planRefine").mockReturnValue(planning.promise); + const apply = vi.spyOn(refinement._execution, "_applyRefine").mockResolvedValue(result()); + const run = harness.session.refine(); + expect(refinement.isApplying).toBe(false); + expect(harness.session.handleRefineHostRequest("refine.status").in_flight).toBe(true); + planning.resolve(plan()); + await vi.waitFor(() => expect(refinement.isApplying).toBe(true)); + operation = secondOperation.promise; + firstOperation.resolve(); + await new Promise(setImmediate); + expect(apply).not.toHaveBeenCalled(); + secondOperation.resolve(); + await run; + expect(apply).toHaveBeenCalledOnce(); + expect(refinement.isApplying).toBe(false); + }); + + it("reads the current model, auth resolver, and extension runner for each plan", async () => { + const firstHook = vi.fn(() => ({ proposal: plan().proposal })); + const secondHook = vi.fn(() => ({ proposal: { ...plan().proposal, summary: "replacement runner" } })); + const first = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", firstHook); + }, + ], + }); + const second = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_refine", secondHook); + }, + ], + }); + harnesses.push(first, second); + const refinement = owner(first); + const firstModel = first.getModel(); + const secondModel = { ...firstModel, id: "replacement-model" }; + const getModel = vi.spyOn(refinement._host, "getModel").mockReturnValue(firstModel); + const getAuth = vi.spyOn(refinement._host, "getRequiredRequestAuth").mockResolvedValue({ apiKey: "test-first" }); + await first.session.refine(); + expect(firstHook).toHaveBeenCalledOnce(); + expect(getAuth).toHaveBeenLastCalledWith(firstModel); + getModel.mockReturnValue(secondModel); + getAuth.mockResolvedValue({ apiKey: "test-replacement" }); + vi.spyOn(refinement._host, "getExtensionRunner").mockReturnValue(owner(second)._host.getExtensionRunner()); + const applied = await first.session.refine(); + expect(applied.summary).toBe("replacement runner"); + expect(getAuth).toHaveBeenLastCalledWith(secondModel); + expect(firstHook).toHaveBeenCalledOnce(); + expect(secondHook).toHaveBeenCalledOnce(); + }); + + it("preserves audit failure precedence over outcome publication and still reconnects", async () => { + const harness = await createHarness({ persistSession: true }); + harnesses.push(harness); + const refinement = owner(harness); + vi.spyOn(refinement._execution, "_planRefine").mockResolvedValue(plan()); + const auditError = new Error("audit append failed"); + vi.spyOn(harness.sessionManager, "appendCustomEntry").mockImplementationOnce(() => { + throw auditError; + }); + vi.spyOn(refinement._host, "emit").mockImplementation((event) => { + if (event.type === "message_start") throw new Error("outcome listener failed"); + }); + const reconnect = vi.spyOn(refinement._host, "reconnect"); + await expect(harness.session.refine()).rejects.toBe(auditError); + expect(harness.session.messages).toContainEqual( + expect.objectContaining({ role: "custom", customType: "refinement_outcome" }), + ); + expect(reconnect).toHaveBeenCalledOnce(); + expect(refinement.isApplying).toBe(false); + }); +});