diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index bf1747cf5..1052ddbe3 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -183,7 +183,7 @@ unsubscribe = session.subscribe(() => {}); ### Prompting and Message Queueing -`PromptOptions` controls prompt expansion, queueing behavior while streaming, and prompt preflight notifications: +`PromptOptions` controls prompt expansion, queueing behavior while streaming, cancellation, and prompt preflight notifications: ```typescript interface PromptOptions { @@ -192,6 +192,7 @@ interface PromptOptions { streamingBehavior?: "steer" | "followUp"; source?: InputSource; preflightResult?: (success: boolean) => void; + signal?: AbortSignal; } ``` @@ -204,6 +205,8 @@ It fires before `prompt()` resolves. `prompt()` still resolves only after the full accepted run finishes, including retries and coordinated child-agent batches. Failures after acceptance are reported through the normal event and message stream, not through `preflightResult(false)`. +A prompt with `signal` reserves the complete prompt lifecycle synchronously, including asynchronous preflight and the active run. It rejects instead of entering active-run queueing when another prompt is in preflight or an agent run is active, and later prompts reject while its reservation is held. Cancellation before `agent_start` rejects with the exact `signal.reason`, cancels pre-prompt compaction and authentication where active, and starts no user or provider turn. Cancellation after `agent_start` uses the existing session abort behavior, so `prompt()` settles through the normal aborted-response event and persistence path rather than rejecting again with `signal.reason`. Prompts without `signal` retain the existing `streamingBehavior` queue behavior during an active run. + The `prompt()` method handles prompt templates, extension commands, and message sending: ```typescript @@ -218,6 +221,12 @@ await session.prompt("What's in this image?", { // During streaming: must specify how to queue the message await session.prompt("Stop and do this instead", { streamingBehavior: "steer" }); await session.prompt("After you're done, also check X", { streamingBehavior: "followUp" }); + +// Reserve and cancel the complete prompt lifecycle +const controller = new AbortController(); +const pending = session.prompt("Inspect the repository", { signal: controller.signal }); +controller.abort(new Error("Request cancelled")); +await pending.catch((error) => console.error(error)); ``` **Behavior:** diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index f26aade82..eeb625ed3 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -109,8 +109,14 @@ import { ModelRegistry } from "./model-registry.ts"; import type { ModelRuntime } from "./model-runtime.ts"; import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.ts"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.ts"; -import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "./session-manager.ts"; +import type { + BranchSummaryEntry, + CompactionAppendTransaction, + SessionEntry, + SessionManager, +} from "./session-manager.ts"; import { + appendCompactionTransaction, appendCustomMessageWithRollback, CURRENT_SESSION_VERSION, getLatestCompactionEntry, @@ -334,6 +340,8 @@ export interface PromptOptions { source?: InputSource; /** Internal hook used by RPC mode to observe prompt preflight acceptance or rejection. */ preflightResult?: (success: boolean) => void; + /** Cancel prompt preflight or abort the active run after it starts. */ + signal?: AbortSignal; } /** Result from cycleModel() */ @@ -389,6 +397,11 @@ interface SubagentCoordination { waitStarted: boolean; } +interface PromptReservation { + readonly signaled: boolean; + agentStarted: boolean; +} + function estimateMessagesTokens(messages: AgentMessage[]): number { let tokens = 0; for (const message of messages) { @@ -437,6 +450,7 @@ export class AgentSession { private _unsubscribeAgent?: () => void; private _eventListeners: AgentSessionEventListener[] = []; private _isAgentRunActive = false; + private readonly _promptReservations = new Set(); private _idleWaitPromise: Promise | undefined; private _resolveIdleWait: (() => void) | undefined; @@ -568,7 +582,10 @@ export class AgentSession { return this._harnessVersionBundle; } - private async _getRequiredRequestAuth(model: Model): Promise<{ + private async _getRequiredRequestAuth( + model: Model, + signal?: AbortSignal, + ): Promise<{ model: Model; apiKey?: string; headers?: Record; @@ -576,7 +593,7 @@ export class AgentSession { }> { let result: AuthResult | undefined; try { - result = await this._modelRuntime.getAuth(model); + result = await this._modelRuntime.getAuth(model, { signal }); } catch (error) { const cause = error instanceof Error ? error.cause : undefined; if (cause instanceof Error && cause.message === "authHeader requires a resolved API key") { @@ -605,18 +622,22 @@ export class AgentSession { throw new Error(formatNoApiKeyFoundMessage(model.provider)); } - private async _getSummarizationRequestAuth(model: Model): Promise<{ + private async _getSummarizationRequestAuth( + model: Model, + signal?: AbortSignal, + ): Promise<{ model: Model; apiKey?: string; headers?: Record; env?: Record; }> { if (this.agent.streamFunction === streamSimple) { - return this._getRequiredRequestAuth(model); + return this._getRequiredRequestAuth(model, signal); } try { - const result = await this._modelRuntime.getAuth(model); + const result = await this._modelRuntime.getAuth(model, { signal }); + signal?.throwIfAborted(); if (!result) return { model }; const requestModel = result.auth.baseUrl ? { ...model, baseUrl: result.auth.baseUrl } : model; return { @@ -626,6 +647,7 @@ export class AgentSession { env: result.env, }; } catch { + if (signal?.aborted) throw signal.reason; return { model }; } } @@ -828,6 +850,13 @@ export class AgentSession { /** Internal handler for agent events - shared by subscribe and reconnect */ private _handleAgentEvent = async (event: AgentEvent): Promise => { + if (event.type === "agent_start") { + for (const reservation of this._promptReservations) { + if (!reservation.signaled) continue; + reservation.agentStarted = true; + break; + } + } if (event.type === "message_start" && this._pendingSubagentEscalationEnvelope?.message === event.message) { // The exact queued object is now entering parent context. Seal it so later reports form a new envelope. this._pendingSubagentEscalationEnvelope = undefined; @@ -1836,13 +1865,64 @@ export class AgentSession { * @throws Error if streaming and no streamingBehavior specified * @throws Error if no model selected or no API key available (when not streaming) */ - async prompt(text: string, options?: PromptOptions): Promise { - return this._traceOperation("prompt", (operationSpan) => this._prompt(text, options, operationSpan)); + prompt(text: string, options?: PromptOptions): Promise { + const signal = options?.signal; + const preflightResult = options?.preflightResult; + const rejectPreflight = (error: unknown): Promise => { + try { + preflightResult?.(false); + } catch { + // Synchronous rejection gates remain authoritative unless the callback cancels the prompt. + } + return Promise.reject(signal?.aborted ? signal.reason : error); + }; + if (signal?.aborted) { + return rejectPreflight(signal.reason); + } + if ([...this._promptReservations].some((reservation) => reservation.signaled)) { + return rejectPreflight(new Error("Cannot submit a prompt while another prompt owns prompt preflight.")); + } + if (signal && this._promptReservations.size > 0) { + return rejectPreflight(new Error("Cannot submit a prompt while another prompt owns prompt preflight.")); + } + if (signal && this._isAgentRunActive) { + return rejectPreflight(new Error("Cannot submit a signaled prompt while an agent run is active.")); + } + + const reservation: PromptReservation = { signaled: signal !== undefined, agentStarted: false }; + this._promptReservations.add(reservation); + const releaseUnsignaledPreflight = (): void => { + if (!reservation.signaled) this._promptReservations.delete(reservation); + }; + const onAbort = signal + ? (): void => { + if (reservation.agentStarted) { + void this.abort().catch(() => {}); + } else { + this.abortCompaction(); + } + } + : undefined; + if (onAbort && signal) signal.addEventListener("abort", onAbort, { once: true }); + + const cleanup = (): void => { + if (onAbort && signal) signal.removeEventListener("abort", onAbort); + this._promptReservations.delete(reservation); + }; + try { + return this._traceOperation("prompt", (operationSpan) => + this._prompt(text, options, releaseUnsignaledPreflight, operationSpan), + ).finally(cleanup); + } catch (error) { + cleanup(); + return Promise.reject(error); + } } private async _prompt( text: string, options: PromptOptions | undefined, + releaseUnsignaledPreflight: () => void, operationSpan: CodingAgentTelemetrySpan<"diffex.coding_agent.operation">, ): Promise { if (this._harnessSwitchInProgress) { @@ -1850,16 +1930,35 @@ export class AgentSession { } const expandPromptTemplates = options?.expandPromptTemplates ?? true; const preflightResult = options?.preflightResult; + const signal = options?.signal; let messages: AgentMessage[] | undefined; + let consumedNextTurnMessages: CustomMessage[] | undefined; + let preflightAccepted = false; + let pendingCompaction: CompactionAppendTransaction | undefined; + const commitPendingCompaction = (): void => { + if (!pendingCompaction) return; + pendingCompaction.commit(); + pendingCompaction = undefined; + this._armCurrentPlanRestorationIfActive(); + }; + const rollbackPendingCompaction = (): void => { + if (!pendingCompaction) return; + pendingCompaction.rollback(); + pendingCompaction = undefined; + this.agent.state.messages = this.sessionManager.buildSessionContext().messages; + }; try { // Handle extension commands first (execute immediately, even during streaming) // Extension commands manage their own LLM interaction via diffex.sendMessage() if (expandPromptTemplates && text.startsWith("/")) { const handled = await this._tryExecuteExtensionCommand(text); + signal?.throwIfAborted(); if (handled) { // Extension command executed, no prompt to send + preflightAccepted = true; preflightResult?.(true); + signal?.throwIfAborted(); return; } } @@ -1880,8 +1979,11 @@ export class AgentSession { options?.source ?? "interactive", this.isStreaming ? options?.streamingBehavior : undefined, ); + signal?.throwIfAborted(); if (inputResult.action === "handled") { + preflightAccepted = true; preflightResult?.(true); + signal?.throwIfAborted(); return; } if (inputResult.action === "transform") { @@ -1913,14 +2015,13 @@ export class AgentSession { } else { await this._queueSteer(expandedText, currentImages, referencedSkills); } + preflightAccepted = true; preflightResult?.(true); return; } await this._refreshSelectedHarnessAtPromptBoundary(); - - // Flush any pending bash messages before the new prompt - this._flushPendingBashMessages(); + signal?.throwIfAborted(); // Validate model if (!this.model) { @@ -1929,7 +2030,8 @@ export class AgentSession { const hasConfiguredAuth = this._modelRuntime.hasConfiguredAuth(this.model.provider) || - (await this._modelRuntime.checkAuth(this.model.provider)) !== undefined; + (await this._modelRuntime.checkAuth(this.model.provider, { signal })) !== undefined; + signal?.throwIfAborted(); if (!hasConfiguredAuth) { const isOAuth = this._modelRuntime.isUsingOAuth(this.model.provider); if (isOAuth) { @@ -1946,7 +2048,10 @@ export class AgentSession { // The user's new prompt is sent below, so do not call agent.continue() here. const lastAssistant = this._findLastAssistantMessage(); if (lastAssistant) { - await this._checkCompaction(lastAssistant, false); + await this._checkCompaction(lastAssistant, false, signal, (transaction) => { + pendingCompaction = transaction; + }); + signal?.throwIfAborted(); } // Build messages array (custom message if any, then user message) @@ -1964,7 +2069,8 @@ export class AgentSession { }); // Inject any pending "nextTurn" messages as context alongside the user message - for (const msg of this._pendingNextTurnMessages) { + consumedNextTurnMessages = this._pendingNextTurnMessages; + for (const msg of consumedNextTurnMessages) { messages.push(msg); } this._pendingNextTurnMessages = []; @@ -1976,6 +2082,9 @@ export class AgentSession { this._baseSystemPrompt, this._baseSystemPromptOptions, ); + signal?.throwIfAborted(); + // Flush pending bash history only after cancellable preflight has completed. + this._flushPendingBashMessages(); // Add all custom messages from extensions if (result?.messages) { for (const msg of result.messages) { @@ -2000,15 +2109,50 @@ export class AgentSession { this.agent.state.systemPrompt = this._appendSubagentDelegatedMessageTrustGuidance(this._baseSystemPrompt); } } catch (error) { - preflightResult?.(false); - throw error; + let failure = error; + if (!preflightAccepted) { + try { + preflightResult?.(false); + } catch (callbackError) { + failure = callbackError; + } + } + if (signal?.aborted) { + rollbackPendingCompaction(); + } else { + commitPendingCompaction(); + } + if (signal?.aborted && consumedNextTurnMessages) { + this._pendingNextTurnMessages = [...consumedNextTurnMessages, ...this._pendingNextTurnMessages]; + } + throw signal?.aborted ? signal.reason : failure; } if (!messages) { return; } - preflightResult?.(true); + preflightAccepted = true; + try { + preflightResult?.(true); + } catch (error) { + if (signal?.aborted) { + rollbackPendingCompaction(); + this._pendingNextTurnMessages = [...(consumedNextTurnMessages ?? []), ...this._pendingNextTurnMessages]; + throw signal.reason; + } + commitPendingCompaction(); + throw error; + } + if (signal?.aborted) { + rollbackPendingCompaction(); + this._pendingNextTurnMessages = [...(consumedNextTurnMessages ?? []), ...this._pendingNextTurnMessages]; + consumedNextTurnMessages = undefined; + throw signal.reason; + } + commitPendingCompaction(); + consumedNextTurnMessages = undefined; + releaseUnsignaledPreflight(); this._releaseIdleSubagentEscalationsForSelectedBranch(); await this._runAgentPrompt(messages); } @@ -2812,8 +2956,11 @@ export class AgentSession { fromExtension: boolean; usage: Usage | undefined; willRetry: boolean; + signal?: AbortSignal; + deferCommit?: (transaction: CompactionAppendTransaction) => void; }): Promise { - this.sessionManager.appendCompaction( + const transaction = appendCompactionTransaction( + this.sessionManager, options.summary, options.firstKeptEntryId, options.tokensBefore, @@ -2821,25 +2968,36 @@ export class AgentSession { options.fromExtension, options.usage, ); - const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.state.messages = sessionContext.messages; const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages); - // Get the saved compaction entry for the extension event - const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === options.summary) as - | CompactionEntry - | undefined; - - if (this._extensionRunner && savedCompactionEntry) { - await this._extensionRunner.emit({ - type: "session_compact", - compactionEntry: savedCompactionEntry, - fromExtension: options.fromExtension, - reason: options.reason, - willRetry: options.willRetry, - }); + if (this._extensionRunner) { + try { + await this._extensionRunner.emit({ + type: "session_compact", + compactionEntry: transaction.entry, + fromExtension: options.fromExtension, + reason: options.reason, + willRetry: options.willRetry, + }); + } catch (error) { + if (options.signal?.aborted) { + transaction.rollback(); + this.agent.state.messages = this.sessionManager.buildSessionContext().messages; + throw options.signal.reason; + } + transaction.commit(); + throw error; + } } + if (options.signal?.aborted) { + transaction.rollback(); + this.agent.state.messages = this.sessionManager.buildSessionContext().messages; + throw options.signal.reason; + } + if (options.deferCommit) options.deferCommit(transaction); + else transaction.commit(); const compactionResult: CompactionResult = { summary: options.summary, @@ -2849,7 +3007,7 @@ export class AgentSession { usage: options.usage, details: options.details, }; - this._armCurrentPlanRestorationIfActive(); + if (!options.deferCommit) this._armCurrentPlanRestorationIfActive(); return compactionResult; } @@ -2865,7 +3023,12 @@ export class AgentSession { * @param assistantMessage The assistant message to check * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true */ - private async _checkCompaction(assistantMessage: AssistantMessage, skipAbortedCheck = true): Promise { + private async _checkCompaction( + assistantMessage: AssistantMessage, + skipAbortedCheck = true, + signal?: AbortSignal, + deferCommit?: (transaction: CompactionAppendTransaction) => void, + ): Promise { const settings = this.settingsManager.getCompactionSettings(); if (!settings.enabled) return false; @@ -2901,7 +3064,7 @@ export class AgentSession { const willRetry = assistantMessage.stopReason !== "stop"; if (!willRetry) { - return await this._runAutoCompaction("overflow", false); + return await this._runAutoCompaction("overflow", false, signal, deferCommit); } if (this._overflowRecoveryAttempted) { @@ -2924,7 +3087,7 @@ export class AgentSession { if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { this.agent.state.messages = messages.slice(0, -1); } - return await this._runAutoCompaction("overflow", willRetry); + return await this._runAutoCompaction("overflow", willRetry, signal, deferCommit); } // Case 2: Threshold - context is getting large @@ -2953,7 +3116,7 @@ export class AgentSession { contextTokens = directContextTokens; } if (shouldCompact(contextTokens, contextWindow, settings)) { - return await this._runAutoCompaction("threshold", false); + return await this._runAutoCompaction("threshold", false, signal, deferCommit); } return false; } @@ -2961,11 +3124,16 @@ export class AgentSession { /** * Internal: Run auto-compaction with events. */ - private _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise { + private _runAutoCompaction( + reason: "overflow" | "threshold", + willRetry: boolean, + signal?: AbortSignal, + deferCommit?: (transaction: CompactionAppendTransaction) => void, + ): Promise { if (this._disposePromise) return Promise.resolve(false); return this._runOwnedOperation(() => this._traceOperation("compaction", (telemetryContext) => - this._runAutoCompactionOperation(reason, willRetry, telemetryContext), + this._runAutoCompactionOperation(reason, willRetry, telemetryContext, signal, deferCommit), ), ); } @@ -2974,6 +3142,8 @@ export class AgentSession { reason: "overflow" | "threshold", willRetry: boolean, telemetryContext: TelemetryContext = this._telemetryContext, + signal?: AbortSignal, + deferCommit?: (transaction: CompactionAppendTransaction) => void, ): Promise { const settings = this.settingsManager.getCompactionSettings(); let started = false; @@ -2983,7 +3153,13 @@ export class AgentSession { return false; } - const { model: requestModel, apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model); + const { + model: requestModel, + apiKey, + headers, + env, + } = await this._getSummarizationRequestAuth(this.model, signal); + signal?.throwIfAborted(); if (this._disposePromise) return false; const pathEntries = this.sessionManager.getBranch(); @@ -2993,8 +3169,8 @@ export class AgentSession { return false; } - this._emit({ type: "compaction_start", reason }); this._autoCompactionAbortController = new AbortController(); + this._emit({ type: "compaction_start", reason }); started = true; let extensionCompaction: CompactionResult | undefined; @@ -3084,6 +3260,8 @@ export class AgentSession { fromExtension, usage, willRetry, + signal, + deferCommit, }); this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry }); @@ -3105,17 +3283,22 @@ export class AgentSession { return this.agent.hasQueuedMessages(); } catch (error) { const errorMessage = error instanceof Error ? error.message : "compaction failed"; + const aborted = signal?.aborted || this._autoCompactionAbortController?.signal.aborted || false; if (started) { this._emit({ type: "compaction_end", reason, result: undefined, - aborted: false, + aborted, willRetry: false, - errorMessage: - reason === "overflow" - ? `Context overflow recovery failed: ${errorMessage}` - : `Auto-compaction failed: ${errorMessage}`, + ...(aborted + ? {} + : { + errorMessage: + reason === "overflow" + ? `Context overflow recovery failed: ${errorMessage}` + : `Auto-compaction failed: ${errorMessage}`, + }), }); } return false; diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 5c9d0983a..cb24d3a39 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -82,6 +82,13 @@ export interface CompactionEntry extends SessionEntryBase { fromHook?: boolean; } +/** @internal Transaction for one exact compaction entry while post-save observers run. */ +export interface CompactionAppendTransaction { + readonly entry: CompactionEntry; + commit(): void; + rollback(): void; +} + export interface BranchSummaryEntry extends SessionEntryBase { type: "branch_summary"; fromId: string; @@ -231,6 +238,7 @@ function createSessionId(): string { const PERSIST_SESSION_HEADER = Symbol("persistSessionHeader"); const PERSIST_SESSION_STATE = Symbol("persistSessionState"); const APPEND_CUSTOM_MESSAGE_WITH_ROLLBACK = Symbol("appendCustomMessageWithRollback"); +const APPEND_COMPACTION_TRANSACTION = Symbol("appendCompactionTransaction"); export function assertValidSessionId(id: string): void { if (!/^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$/.test(id)) { @@ -1192,6 +1200,84 @@ export class SessionManager { return entry.id; } + [APPEND_COMPACTION_TRANSACTION]( + summary: string, + firstKeptEntryId: string, + tokensBefore: number, + details?: T, + fromHook?: boolean, + usage?: Usage, + ): CompactionAppendTransaction { + const previousLeafId = this.leafId; + const previousFlushed = this.flushed; + const entry: CompactionEntry = { + type: "compaction", + id: generateId(this.byId), + parentId: previousLeafId, + timestamp: new Date().toISOString(), + summary, + firstKeptEntryId, + tokensBefore, + details, + usage, + fromHook, + }; + try { + this._appendEntry(entry); + } catch (error) { + if (this.fileEntries.at(-1) === entry) this.fileEntries.pop(); + this.byId.delete(entry.id); + this.leafId = previousLeafId; + this.flushed = previousFlushed; + throw error; + } + + let active = true; + return { + entry, + commit: () => { + active = false; + }, + rollback: () => { + if (!active) return; + if (this.byId.get(entry.id) !== entry) { + throw new Error(`Compaction entry ${entry.id} is no longer available for rollback`); + } + const entryIndex = this.fileEntries.indexOf(entry); + if (entryIndex < 0) { + throw new Error(`Compaction entry ${entry.id} is missing from session state`); + } + if (this.fileEntries.slice(entryIndex + 1).some((candidate) => candidate.type === "compaction")) { + throw new Error(`Compaction entry ${entry.id} is not the latest compaction`); + } + + const reparentedChildren: SessionEntry[] = []; + for (const candidate of this.fileEntries) { + if (candidate.type !== "session" && candidate.parentId === entry.id) { + candidate.parentId = previousLeafId; + reparentedChildren.push(candidate); + } + } + const rollbackLeafId = this.leafId; + this.fileEntries.splice(entryIndex, 1); + this.byId.delete(entry.id); + if (this.leafId === entry.id) this.leafId = previousLeafId; + try { + if (this.persist && this.sessionFile && this.flushed) this._rewriteFile(); + this.flushed = previousFlushed; + active = false; + } catch (error) { + this.fileEntries.splice(entryIndex, 0, entry); + this.byId.set(entry.id, entry); + this.leafId = rollbackLeafId; + this.flushed = previousFlushed; + for (const child of reparentedChildren) child.parentId = entry.id; + throw error; + } + }, + }; + } + /** Append a custom entry (for extensions) as child of current leaf, then advance leaf. Returns entry id. */ appendCustomEntry(customType: string, data?: unknown): string { const entry: CustomEntry = { @@ -1890,3 +1976,23 @@ export function appendCustomMessageWithRollback( ): string { return sessionManager[APPEND_CUSTOM_MESSAGE_WITH_ROLLBACK](customType, content, display, details); } + +/** Internal exact-entry compaction transaction for post-save extension notification. */ +export function appendCompactionTransaction( + sessionManager: SessionManager, + summary: string, + firstKeptEntryId: string, + tokensBefore: number, + details?: T, + fromHook?: boolean, + usage?: Usage, +): CompactionAppendTransaction { + return sessionManager[APPEND_COMPACTION_TRANSACTION]( + summary, + firstKeptEntryId, + tokensBefore, + details, + fromHook, + usage, + ); +} diff --git a/packages/coding-agent/test/session-manager/compaction-rollback.test.ts b/packages/coding-agent/test/session-manager/compaction-rollback.test.ts new file mode 100644 index 000000000..6ec208c97 --- /dev/null +++ b/packages/coding-agent/test/session-manager/compaction-rollback.test.ts @@ -0,0 +1,82 @@ +import { mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { appendCompactionTransaction, SessionManager } from "../../src/core/session-manager.ts"; + +function appendAssistant(sessionManager: SessionManager, text: string, timestamp: number): string { + return sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text }], + api: "anthropic-messages", + provider: "test", + model: "test", + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp, + }); +} + +describe("SessionManager compaction rollback", () => { + const tempDirs: string[] = []; + + afterEach(() => { + while (tempDirs.length > 0) { + const tempDir = tempDirs.pop(); + if (tempDir) rmSync(tempDir, { recursive: true, force: true }); + } + }); + + it.each(["in-memory", "persisted"] as const)( + "removes only the exact latest compaction and preserves %s branch state", + (storage) => { + const tempDir = join( + tmpdir(), + `diffex-compaction-rollback-${Date.now()}-${Math.random().toString(36).slice(2)}`, + ); + mkdirSync(tempDir, { recursive: true }); + tempDirs.push(tempDir); + const sessionManager = + storage === "persisted" ? SessionManager.create(tempDir, tempDir) : SessionManager.inMemory(tempDir); + const firstUserId = sessionManager.appendMessage({ role: "user", content: "first", timestamp: 1 }); + appendAssistant(sessionManager, "first response", 2); + const priorCompactionId = sessionManager.appendCompaction("prior", firstUserId, 2); + const currentUserId = sessionManager.appendMessage({ role: "user", content: "current", timestamp: 3 }); + const currentAssistantId = appendAssistant(sessionManager, "current response", 4); + const entriesBefore = structuredClone(sessionManager.getEntries()); + + const transaction = appendCompactionTransaction(sessionManager, "cancelled", currentUserId, 4); + const extensionEntryId = sessionManager.appendCustomEntry("extension-side-effect", { preserved: true }); + transaction.rollback(); + const appendedAfterRollbackId = sessionManager.appendCustomEntry("after-rollback", { persisted: true }); + + expect(sessionManager.getEntry(transaction.entry.id)).toBeUndefined(); + expect(sessionManager.getEntry(priorCompactionId)).toMatchObject({ type: "compaction", summary: "prior" }); + expect(sessionManager.getEntries().slice(0, entriesBefore.length)).toEqual(entriesBefore); + expect(sessionManager.getEntry(extensionEntryId)).toMatchObject({ + type: "custom", + parentId: currentAssistantId, + data: { preserved: true }, + }); + expect(sessionManager.getLeafId()).toBe(appendedAfterRollbackId); + expect(sessionManager.getBranch().some((entry) => entry.id === transaction.entry.id)).toBe(false); + + if (storage === "persisted") { + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile) throw new Error("Expected persisted session file"); + const reopened = SessionManager.open(sessionFile, tempDir); + expect(reopened.getEntry(transaction.entry.id)).toBeUndefined(); + expect(reopened.getEntry(priorCompactionId)).toMatchObject({ type: "compaction", summary: "prior" }); + expect(reopened.getEntry(extensionEntryId)).toMatchObject({ parentId: currentAssistantId }); + expect(reopened.getLeafId()).toBe(appendedAfterRollbackId); + } + }, + ); +}); 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 ff40b812b..7250b0842 100644 --- a/packages/coding-agent/test/suite/agent-session-prompt.test.ts +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -2,21 +2,43 @@ import { mkdirSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { type AgentTool, InMemoryTelemetryContext } from "@earendil-works/diffex-agent-core"; -import { fauxAssistantMessage, fauxToolCall, type Model } from "@earendil-works/diffex-ai"; +import { type AssistantMessage, fauxAssistantMessage, fauxToolCall, type Model } from "@earendil-works/diffex-ai"; +import { HarnessStateStore } from "@earendil-works/diffex-harness-state"; import { Type } from "typebox"; -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import type { InputEvent } from "../../src/core/extensions/index.ts"; import type { PromptTemplate } from "../../src/core/prompt-templates.ts"; +import { SessionManager } from "../../src/core/session-manager.ts"; import { createSkill } from "../../src/core/skills.ts"; import { createSyntheticSourceInfo } from "../../src/core/source-info.ts"; import { createTestResourceLoader } from "../utilities.ts"; import { createHarness, getMessageText, type Harness } from "./harness.ts"; +function deferred(): { promise: Promise; resolve: () => void } { + let resolve = (): void => {}; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + +function createUsage(totalTokens: number) { + return { + input: totalTokens, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + describe("AgentSession prompt characterization", () => { const harnesses: Harness[] = []; const tempDirs: string[] = []; afterEach(async () => { + vi.restoreAllMocks(); while (harnesses.length > 0) { await harnesses.pop()?.cleanup(); } @@ -340,7 +362,7 @@ describe("AgentSession prompt characterization", () => { expect(inputEvents[0]?.streamingBehavior).toBeUndefined(); }); - it("reports streamingBehavior to input handlers while streaming", async () => { + it("reports streamingBehavior while preserving concurrent unsignaled queue submissions", async () => { let releaseToolExecution: (() => void) | undefined; const toolRelease = new Promise((resolve) => { releaseToolExecution = resolve; @@ -386,9 +408,14 @@ describe("AgentSession prompt characterization", () => { const promptPromise = harness.session.prompt("start"); await sawToolStart; - await harness.session.prompt("queued", { streamingBehavior: "followUp" }); + const queuedPrompts = Promise.all([ + harness.session.prompt("queued one", { streamingBehavior: "followUp" }), + harness.session.prompt("queued two", { streamingBehavior: "followUp" }), + ]); + await queuedPrompts; - expect(inputEvents.map((event) => event.streamingBehavior)).toEqual([undefined, "followUp"]); + expect(harness.session.getFollowUpMessages()).toEqual(["queued one", "queued two"]); + expect(inputEvents.map((event) => event.streamingBehavior)).toEqual([undefined, "followUp", "followUp"]); releaseToolExecution?.(); await promptPromise; @@ -439,6 +466,634 @@ describe("AgentSession prompt characterization", () => { await promptPromise; }); + it("rejects already-aborted signals before prompt state changes and preserves abort reasons", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("unused")]); + + const defaultAbort = new AbortController(); + defaultAbort.abort(); + await expect(harness.session.prompt("default", { signal: defaultAbort.signal })).rejects.toMatchObject({ + name: "AbortError", + }); + + const customReason = new Error("custom cancellation"); + const customAbort = new AbortController(); + customAbort.abort(customReason); + await expect( + harness.session.prompt("custom", { + signal: customAbort.signal, + preflightResult: () => { + throw new Error("ignored callback failure"); + }, + }), + ).rejects.toBe(customReason); + + expect(harness.session.messages).toEqual([]); + expect(harness.sessionManager.getEntries()).toEqual([]); + expect(harness.faux.state.callCount).toBe(0); + }); + + it.each([ + { firstSignaled: false, label: "signaled prompt behind unsignaled preflight" }, + { firstSignaled: true, label: "unsignaled prompt behind signaled preflight" }, + ])("rejects a $label", async ({ firstSignaled }) => { + const inputStarted = deferred(); + const releaseInput = deferred(); + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("input", async () => { + inputStarted.resolve(); + await releaseInput.promise; + return { action: "continue" }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("first completed")]); + const firstController = new AbortController(); + const firstPrompt = harness.session.prompt( + "first", + firstSignaled ? { signal: firstController.signal } : undefined, + ); + await inputStarted.promise; + + const secondController = new AbortController(); + const secondReason = new Error("cancel rejected overlap"); + const secondOptions = { + ...(firstSignaled ? {} : { signal: secondController.signal }), + preflightResult: () => { + if (!firstSignaled) secondController.abort(secondReason); + throw new Error("ignored callback failure"); + }, + }; + const rejection = expect(harness.session.prompt("second", secondOptions)).rejects; + if (firstSignaled) { + await rejection.toThrow("Cannot submit a prompt while another prompt owns prompt preflight."); + } else { + await rejection.toBe(secondReason); + } + + releaseInput.resolve(); + await firstPrompt; + expect(harness.faux.state.callCount).toBe(1); + }); + + it("rejects a signaled prompt instead of queueing behind an active unsignaled run", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("x".repeat(20_000))]); + const responseStarted = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "message_update") { + unsubscribe(); + resolve(); + } + }); + }); + const activePrompt = harness.session.prompt("active"); + await responseStarted; + + const controller = new AbortController(); + await expect( + harness.session.prompt("must not queue", { signal: controller.signal, streamingBehavior: "followUp" }), + ).rejects.toThrow("Cannot submit a signaled prompt while an agent run is active."); + const cancelledController = new AbortController(); + const cancelledReason = new Error("cancel rejected active overlap"); + await expect( + harness.session.prompt("cancel overlap", { + signal: cancelledController.signal, + streamingBehavior: "followUp", + preflightResult: () => { + cancelledController.abort(cancelledReason); + throw new Error("ignored callback failure"); + }, + }), + ).rejects.toBe(cancelledReason); + + await harness.session.abort(); + await activePrompt; + expect(harness.session.getFollowUpMessages()).toEqual([]); + }); + + it("observes cancellation after an awaited extension command boundary", async () => { + const commandStarted = deferred(); + const releaseCommand = deferred(); + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.registerCommand("wait", { + description: "Wait", + handler: async () => { + commandStarted.resolve(); + await releaseCommand.promise; + }, + }); + }, + ], + }); + harnesses.push(harness); + const controller = new AbortController(); + const reason = new Error("cancel command preflight"); + const prompt = harness.session.prompt("/wait", { signal: controller.signal }); + await commandStarted.promise; + controller.abort(reason); + releaseCommand.resolve(); + + await expect(prompt).rejects.toBe(reason); + expect(harness.session.messages).toEqual([]); + expect(harness.faux.state.callCount).toBe(0); + }); + + it("observes cancellation after awaited input and before-agent-start boundaries", async () => { + const inputStarted = deferred(); + const releaseInput = deferred(); + const beforeStartEntered = deferred(); + const releaseBeforeStart = deferred(); + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("input", async () => { + inputStarted.resolve(); + await releaseInput.promise; + return { action: "continue" }; + }); + pi.on("before_agent_start", async () => { + beforeStartEntered.resolve(); + await releaseBeforeStart.promise; + return undefined; + }); + }, + ], + }); + harnesses.push(harness); + + const inputController = new AbortController(); + const inputReason = new Error("cancel input"); + const inputPrompt = harness.session.prompt("input", { signal: inputController.signal }); + await inputStarted.promise; + inputController.abort(inputReason); + releaseInput.resolve(); + await expect(inputPrompt).rejects.toBe(inputReason); + + const beforeStartController = new AbortController(); + const beforeStartReason = new Error("cancel before start"); + const beforeStartPrompt = harness.session.prompt("before start", { signal: beforeStartController.signal }); + await beforeStartEntered.promise; + beforeStartController.abort(beforeStartReason); + releaseBeforeStart.resolve(); + await expect(beforeStartPrompt).rejects.toBe(beforeStartReason); + + expect(harness.session.messages).toEqual([]); + expect(harness.faux.state.callCount).toBe(0); + }); + + it("observes cancellation after the awaited harness refresh boundary", async () => { + const agentDir = join(tmpdir(), `diffex-prompt-signal-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(agentDir, { recursive: true }); + tempDirs.push(agentDir); + const store = new HarnessStateStore(agentDir); + const root = await store.initialize(); + const harness = await createHarness({ harnessStateStore: store, harnessVersionBundle: root }); + harnesses.push(harness); + const loadStarted = deferred(); + const releaseLoad = deferred(); + const loadCurrent = store.loadCurrent.bind(store); + vi.spyOn(store, "loadCurrent").mockImplementation(async () => { + loadStarted.resolve(); + await releaseLoad.promise; + return loadCurrent(); + }); + const controller = new AbortController(); + const reason = new Error("cancel harness refresh"); + const prompt = harness.session.prompt("refresh", { signal: controller.signal }); + await loadStarted.promise; + controller.abort(reason); + releaseLoad.resolve(); + + await expect(prompt).rejects.toBe(reason); + expect(harness.session.messages).toEqual([]); + expect(harness.faux.state.callCount).toBe(0); + }); + + it("propagates cancellation through authentication", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + const authStarted = deferred(); + vi.spyOn(harness.session.modelRuntime, "checkAuth").mockImplementation(async (_providerId, options) => { + authStarted.resolve(); + const signal = options?.signal; + if (!signal) throw new Error("Expected prompt signal in authentication"); + await new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => reject(signal.reason), { once: true }); + }); + return undefined; + }); + const controller = new AbortController(); + const reason = new Error("cancel authentication"); + const prompt = harness.session.prompt("authenticate", { signal: controller.signal }); + await authStarted.promise; + controller.abort(reason); + + await expect(prompt).rejects.toBe(reason); + expect(harness.session.messages).toEqual([]); + expect(harness.faux.state.callCount).toBe(0); + }); + + it("cancels pre-prompt compaction before recording or starting the turn", async () => { + const compactionStarted = deferred(); + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 100, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", async (event) => { + compactionStarted.resolve(); + await new Promise((resolve) => { + if (event.signal.aborted) resolve(); + else event.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + return { + compaction: { + summary: "must not be recorded", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: {}, + }, + }; + }); + }, + ], + }); + harnesses.push(harness); + const now = Date.now(); + const model = harness.getModel(); + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "previous" }], + timestamp: now - 1000, + }); + const previousAssistant: AssistantMessage = { + ...fauxAssistantMessage("previous response", { stopReason: "length", timestamp: now - 500 }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(100), + }; + harness.sessionManager.appendMessage(previousAssistant); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + + const controller = new AbortController(); + const reason = new Error("cancel compaction"); + const prompt = harness.session.prompt("new prompt", { signal: controller.signal }); + await compactionStarted.promise; + controller.abort(reason); + + await expect(prompt).rejects.toBe(reason); + expect(harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction")).toEqual([]); + expect(harness.session.messages.filter((message) => message.role === "user")).toHaveLength(1); + expect(harness.faux.state.callCount).toBe(0); + }); + + it.each(["in-memory", "persisted"] as const)( + "rolls back a %s compaction when cancellation arrives during session_compact", + async (storage) => { + const sessionDir = join( + tmpdir(), + `diffex-prompt-compact-${Date.now()}-${Math.random().toString(36).slice(2)}`, + ); + mkdirSync(sessionDir, { recursive: true }); + tempDirs.push(sessionDir); + const sessionManager = + storage === "persisted" + ? SessionManager.create(sessionDir, sessionDir) + : SessionManager.inMemory(sessionDir); + const compactHookStarted = deferred(); + const releaseCompactHook = deferred(); + let hookCompleted = false; + const harness = await createHarness({ + sessionManager, + models: [{ id: "faux-1", contextWindow: 100, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "temporarily recorded", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: {}, + }, + })); + pi.on("session_compact", async () => { + compactHookStarted.resolve(); + await releaseCompactHook.promise; + hookCompleted = true; + }); + }, + ], + }); + harnesses.push(harness); + const now = Date.now(); + const model = harness.getModel(); + const previousUserId = sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "previous" }], + timestamp: now - 1000, + }); + const previousAssistantId = sessionManager.appendMessage({ + ...fauxAssistantMessage("previous response", { stopReason: "length", timestamp: now - 500 }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(100), + }); + harness.session.agent.state.messages = sessionManager.buildSessionContext().messages; + const entriesBeforePrompt = structuredClone(sessionManager.getEntries()); + + const controller = new AbortController(); + const reason = new Error("cancel saved compaction"); + const prompt = harness.session.prompt("new prompt", { signal: controller.signal }); + await compactHookStarted.promise; + expect(sessionManager.getLeafEntry()).toMatchObject({ type: "compaction", summary: "temporarily recorded" }); + controller.abort(reason); + releaseCompactHook.resolve(); + + await expect(prompt).rejects.toBe(reason); + expect(hookCompleted).toBe(true); + expect(sessionManager.getEntries()).toEqual(entriesBeforePrompt); + expect(sessionManager.getBranch().map((entry) => entry.id)).toEqual([previousUserId, previousAssistantId]); + expect(sessionManager.getLeafId()).toBe(previousAssistantId); + expect(harness.session.messages.filter((message) => message.role === "user")).toHaveLength(1); + expect(harness.faux.state.callCount).toBe(0); + + if (storage === "persisted") { + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile) throw new Error("Expected persisted session file"); + const reopened = SessionManager.open(sessionFile, sessionDir); + expect(reopened.getEntries()).toEqual(entriesBeforePrompt); + expect(reopened.getEntries().some((entry) => entry.type === "compaction")).toBe(false); + } + }, + ); + + it.each(["success", "preflight failure"])("removes its single abort listener after %s", async (outcome) => { + const harness = await createHarness(); + harnesses.push(harness); + if (outcome === "success") harness.setResponses([fauxAssistantMessage("done")]); + else harness.session.agent.state.model = undefined as unknown as NonNullable>; + const controller = new AbortController(); + const addListener = vi.spyOn(controller.signal, "addEventListener"); + const removeListener = vi.spyOn(controller.signal, "removeEventListener"); + + const prompt = harness.session.prompt("listener", { signal: controller.signal }); + if (outcome === "success") await expect(prompt).resolves.toBeUndefined(); + else await expect(prompt).rejects.toThrow("No model selected."); + + expect(addListener).toHaveBeenCalledOnce(); + expect(addListener).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }); + expect(removeListener).toHaveBeenCalledOnce(); + expect(removeListener).toHaveBeenCalledWith("abort", addListener.mock.calls[0]?.[1]); + }); + + it("preserves pending next-turn messages when cancellation occurs at preflight acceptance", async () => { + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 100, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "completed before acceptance", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: {}, + }, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("done")]); + const now = Date.now(); + const model = harness.getModel(); + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "previous" }], + timestamp: now - 1000, + }); + harness.sessionManager.appendMessage({ + ...fauxAssistantMessage("previous response", { stopReason: "length", timestamp: now - 500 }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(100), + }); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + const entriesBeforePrompt = structuredClone(harness.sessionManager.getEntries()); + await harness.session.sendCustomMessage( + { customType: "pending-context", content: "keep me", display: false }, + { deliverAs: "nextTurn" }, + ); + const controller = new AbortController(); + const reason = new Error("cancel at acceptance"); + const preflightResults: boolean[] = []; + + const prompt = harness.session.prompt("acceptance boundary", { + signal: controller.signal, + preflightResult: (success) => { + preflightResults.push(success); + if (success) { + controller.abort(reason); + throw new Error("ignored callback failure"); + } + }, + }); + + await expect(prompt).rejects.toBe(reason); + expect(preflightResults).toEqual([true]); + expect(harness.sessionManager.getEntries()).toEqual(entriesBeforePrompt); + expect(harness.sessionManager.getEntries().some((entry) => entry.type === "compaction")).toBe(false); + expect(harness.eventsOfType("compaction_end")).toEqual([ + expect.objectContaining({ + aborted: false, + result: expect.objectContaining({ summary: "completed before acceptance" }), + }), + ]); + expect(harness.faux.state.callCount).toBe(0); + + harness.session.setAutoCompactionEnabled(false); + await harness.session.prompt("next prompt"); + expect(harness.session.messages.filter((message) => message.role === "custom")).toEqual([ + expect.objectContaining({ customType: "pending-context", content: "keep me" }), + ]); + }); + + it("commits compaction bookkeeping when accepted preflight notification throws", async () => { + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 100, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "committed before notification failure", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: {}, + }, + })); + }, + ], + }); + harnesses.push(harness); + const now = Date.now(); + const model = harness.getModel(); + harness.sessionManager.appendMessage({ role: "user", content: "previous", timestamp: now - 1000 }); + harness.sessionManager.appendCustomEntry("diffex.current-plan", { + plan: [{ step: "continue", status: "in_progress" }], + }); + harness.sessionManager.appendMessage({ + ...fauxAssistantMessage("previous response", { stopReason: "length", timestamp: now - 500 }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(100), + }); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + const callbackError = new Error("accepted callback failure"); + + await expect( + harness.session.prompt("failed notification", { + preflightResult: (success) => { + if (success) throw callbackError; + }, + }), + ).rejects.toBe(callbackError); + expect(harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction")).toEqual([ + expect.objectContaining({ summary: "committed before notification failure" }), + ]); + + let restoredPlan = false; + harness.session.setAutoCompactionEnabled(false); + harness.setResponses([ + (context) => { + restoredPlan = context.messages.some((message) => + getMessageText(message).includes("Restored unfinished update_plan checklist"), + ); + return fauxAssistantMessage("done"); + }, + ]); + await harness.session.prompt("continue"); + expect(restoredPlan).toBe(true); + }); + + it("restores preflight state when failure notification aborts the prompt", async () => { + let failBeforeStart = true; + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 100, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "completed before failed preflight", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: {}, + }, + })); + pi.on("before_agent_start", () => { + if (!failBeforeStart) return; + failBeforeStart = false; + return { + message: { + customType: "failing-preflight-message", + get content(): string { + throw new Error("before-agent-start result failure"); + }, + display: false, + }, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("done")]); + const now = Date.now(); + const model = harness.getModel(); + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "previous" }], + timestamp: now - 1000, + }); + harness.sessionManager.appendMessage({ + ...fauxAssistantMessage("previous response", { stopReason: "length", timestamp: now - 500 }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(100), + }); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + const entriesBeforePrompt = structuredClone(harness.sessionManager.getEntries()); + await harness.session.sendCustomMessage( + { customType: "pending-context", content: "keep after failure", display: false }, + { deliverAs: "nextTurn" }, + ); + const controller = new AbortController(); + const reason = new Error("cancel during failed preflight notification"); + const preflightResults: boolean[] = []; + + const prompt = harness.session.prompt("failed preflight", { + signal: controller.signal, + preflightResult: (success) => { + preflightResults.push(success); + if (!success) controller.abort(reason); + }, + }); + + await expect(prompt).rejects.toBe(reason); + expect(preflightResults).toEqual([false]); + expect(harness.sessionManager.getEntries()).toEqual(entriesBeforePrompt); + expect(harness.sessionManager.getEntries().some((entry) => entry.type === "compaction")).toBe(false); + expect(harness.session.messages.filter((message) => message.role === "user")).toHaveLength(1); + expect(harness.faux.state.callCount).toBe(0); + + harness.session.setAutoCompactionEnabled(false); + await harness.session.prompt("next prompt"); + expect(harness.session.messages.filter((message) => message.role === "custom")).toEqual([ + expect.objectContaining({ customType: "pending-context", content: "keep after failure" }), + ]); + }); + + it("uses existing aborted-response settlement after agent_start", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("x".repeat(20_000))]); + const responseStarted = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "message_update") { + unsubscribe(); + resolve(); + } + }); + }); + const controller = new AbortController(); + const prompt = harness.session.prompt("start", { signal: controller.signal }); + await responseStarted; + controller.abort(new Error("post-start reason must not reject prompt")); + + await expect(prompt).resolves.toBeUndefined(); + const assistant = harness.session.messages + .slice() + .reverse() + .find((message) => message.role === "assistant"); + expect(assistant).toMatchObject({ role: "assistant", stopReason: "aborted" }); + }); + it("throws when prompted during manual compaction", async () => { let markCompactionStarted = () => {}; const compactionStarted = new Promise((resolve) => {