diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 36ef5fa84..c1410182d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -32,6 +32,7 @@ - Fixed extension tools without parameter schemas to be rejected during registration instead of breaking provider requests ([#9300](https://github.com/earendil-works/pi/issues/9300)). - Fixed `before_agent_start` handlers returning `systemPrompt` (and `forceSystemPrompt`) on models with mid-conversation system messages: the forced prompt is now sent as the provider's leading system prompt instead of being appended as a section patch after the original prompt. - Fixed loaded llama.cpp models with `enable_thinking` chat templates ignoring Pi's thinking level ([#9528](https://github.com/earendil-works/pi/issues/9528)). +- Fixed cancellation races that could start automatic compaction, leave stale retry state, or miss cancellation while waiting for summarization authentication ([#9340](https://github.com/earendil-works/pi/issues/9340), [#9777](https://github.com/earendil-works/pi/issues/9777)). ## [0.85.1] - 2026-09-05 diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 5c7fc7ec6..7448916b3 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -322,6 +322,7 @@ export class AgentSession { private _unsubscribeAgent?: () => void; private _eventListeners: AgentSessionEventListener[] = []; private _isAgentRunActive = false; + private _agentRunAbortRequested = false; private _idleWaitPromise: Promise | undefined; private _resolveIdleWait: (() => void) | undefined; @@ -418,7 +419,10 @@ export class AgentSession { return this._modelRuntime; } - private async _getRequiredRequestAuth(model: Model): Promise<{ + private async _getRequiredRequestAuth( + model: Model, + signal?: AbortSignal, + ): Promise<{ model: Model; apiKey?: string; headers?: Record; @@ -426,7 +430,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") { @@ -455,18 +459,21 @@ 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 }); if (!result) return { model }; const requestModel = result.auth.baseUrl ? { ...model, baseUrl: result.auth.baseUrl } : model; return { @@ -475,7 +482,8 @@ export class AgentSession { headers: withoutDeletedHeaders(result.auth.headers), env: result.env, }; - } catch { + } catch (error) { + if (signal?.aborted) throw error; return { model }; } } @@ -741,6 +749,7 @@ export class AgentSession { }; private _willRetryAfterAgentEnd(event: Extract): boolean { + if (this._agentRunAbortRequested) return false; const settings = this.settingsManager.getRetrySettings(); if (!settings.enabled || this._retryAttempt >= settings.maxRetries) { return false; @@ -1175,13 +1184,16 @@ export class AgentSession { // ========================================================================= private async _runAgentPrompt(messages: AgentMessage | AgentMessage[]): Promise { + this._agentRunAbortRequested = false; this._isAgentRunActive = true; try { await this.agent.prompt(messages); while (await this._handlePostAgentRun()) { + if (this._agentRunAbortRequested) break; await this.agent.continue(); } } finally { + if (this._agentRunAbortRequested) this._finishCancelledRetry(); this._runSystemPromptOptions = undefined; this._flushPendingBashMessages(); this._flushPendingCustomMessages(); @@ -1192,12 +1204,21 @@ export class AgentSession { private async _handlePostAgentRun(): Promise { const msg = this._lastAssistantMessage; this._lastAssistantMessage = undefined; + if (this._agentRunAbortRequested) { + this._finishCancelledRetry(); + return false; + } if (!msg) { return false; } if (this._isRetryableError(msg) && (await this._prepareRetry(msg))) { - return true; + if (this._agentRunAbortRequested) this._finishCancelledRetry(); + return !this._agentRunAbortRequested; + } + if (this._agentRunAbortRequested) { + this._finishCancelledRetry(); + return false; } if (msg.stopReason === "error" && this._retryAttempt > 0) { @@ -1211,12 +1232,12 @@ export class AgentSession { } if (await this._checkCompaction(msg)) { - return true; + return !this._agentRunAbortRequested; } // The agent loop drains both queues before emitting agent_end. Any messages // here were queued by agent_end extension handlers and need a continuation. - return this.agent.hasQueuedMessages(); + return !this._agentRunAbortRequested && this.agent.hasQueuedMessages(); } private async _runInputHandlers( @@ -1712,6 +1733,9 @@ export class AgentSession { * Abort current operation and wait for agent to become idle. */ async abort(): Promise { + if (this._isAgentRunActive) { + this._agentRunAbortRequested = true; + } this.abortRetry(); this.abortCompaction(); this.abortBranchSummary(); @@ -2043,6 +2067,7 @@ export class AgentSession { this._compactionAbortController = new AbortController(); this._emit({ type: "compaction_start", reason: "manual" }); let fromExtension = false; + let cancelledByExtension = false; try { const model = this.model; @@ -2051,7 +2076,12 @@ export class AgentSession { } const settings = this.settingsManager.getCompactionSettings(model); - const { model: requestModel, apiKey, headers, env } = await this._getSummarizationRequestAuth(model); + const { + model: requestModel, + apiKey, + headers, + env, + } = await this._getSummarizationRequestAuth(model, this._compactionAbortController.signal); const pathEntries = this.sessionManager.getBranch(); @@ -2079,6 +2109,7 @@ export class AgentSession { })) as SessionBeforeCompactResult | undefined; if (result?.cancel) { + cancelledByExtension = true; throw new Error("Compaction cancelled"); } @@ -2165,7 +2196,7 @@ export class AgentSession { return compactionResult; } catch (error) { const message = error instanceof Error ? error.message : String(error); - const aborted = message === "Compaction cancelled" || (error instanceof Error && error.name === "AbortError"); + const aborted = this._compactionAbortController.signal.aborted || cancelledByExtension; const errorMessage = aborted ? undefined : `Compaction failed: ${message}`; this._clearManualCompactionState(); this._emit({ @@ -2344,26 +2375,35 @@ export class AgentSession { private async _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise { const model = this.model; const settings = this.settingsManager.getCompactionSettings(model); + let abortController: AbortController | undefined; let started = false; let fromExtension = false; + let cancelledByExtension = false; try { if (!model) { return false; } - const { model: requestModel, apiKey, headers, env } = await this._getSummarizationRequestAuth(model); - const pathEntries = this.sessionManager.getBranch(); - const preparation = prepareCompaction(pathEntries, settings); if (!preparation) { return false; } - this._emit({ type: "compaction_start", reason }); - this._autoCompactionAbortController = new AbortController(); + abortController = new AbortController(); + this._autoCompactionAbortController = abortController; started = true; + this._emit({ type: "compaction_start", reason }); + abortController.signal.throwIfAborted(); + + const { + model: requestModel, + apiKey, + headers, + env, + } = await this._getSummarizationRequestAuth(model, abortController.signal); + abortController.signal.throwIfAborted(); let extensionCompaction: CompactionResult | undefined; @@ -2375,24 +2415,12 @@ export class AgentSession { customInstructions: undefined, reason, willRetry, - signal: this._autoCompactionAbortController.signal, + signal: abortController.signal, })) as SessionBeforeCompactResult | undefined; if (extensionResult?.cancel) { - this._emit({ - type: "compaction_end", - reason, - result: undefined, - aborted: true, - willRetry: false, - }); - await this._emitSessionCompactFailed({ - reason, - aborted: true, - willRetry: false, - fromExtension: false, - }); - return false; + cancelledByExtension = true; + throw new Error("Compaction cancelled"); } if (extensionResult?.compaction) { @@ -2400,6 +2428,7 @@ export class AgentSession { fromExtension = true; } } + abortController.signal.throwIfAborted(); let summary: string; let firstKeptEntryId: string; @@ -2422,7 +2451,7 @@ export class AgentSession { apiKey, headers, undefined, - this._autoCompactionAbortController.signal, + abortController.signal, env, reason, ); @@ -2432,23 +2461,7 @@ export class AgentSession { usage = compactResult.usage; details = compactResult.details; } - - if (this._autoCompactionAbortController.signal.aborted) { - this._emit({ - type: "compaction_end", - reason, - result: undefined, - aborted: true, - willRetry: false, - }); - await this._emitSessionCompactFailed({ - reason, - aborted: true, - willRetry: false, - fromExtension, - }); - return false; - } + abortController.signal.throwIfAborted(); this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage); const newEntries = this.sessionManager.getEntries(); @@ -2498,31 +2511,35 @@ export class AgentSession { // Continue once so queued messages are delivered. return this.agent.hasQueuedMessages(); } catch (error) { - const errorMessage = error instanceof Error ? error.message : "compaction failed"; + const message = error instanceof Error ? error.message : "compaction failed"; + const aborted = abortController?.signal.aborted === true || cancelledByExtension; if (started) { - const formattedErrorMessage = - reason === "overflow" - ? `Context overflow recovery failed: ${errorMessage}` - : `Auto-compaction failed: ${errorMessage}`; + const errorMessage = aborted + ? undefined + : reason === "overflow" + ? `Context overflow recovery failed: ${message}` + : `Auto-compaction failed: ${message}`; this._emit({ type: "compaction_end", reason, result: undefined, - aborted: false, + aborted, willRetry: false, - errorMessage: formattedErrorMessage, + errorMessage, }); await this._emitSessionCompactFailed({ reason, - errorMessage: formattedErrorMessage, - aborted: false, + errorMessage, + aborted, willRetry: false, fromExtension, }); } return false; } finally { - this._autoCompactionAbortController = undefined; + if (this._autoCompactionAbortController === abortController) { + this._autoCompactionAbortController = undefined; + } this._resolveIdleWaitIfIdle(); } } @@ -2983,6 +3000,18 @@ export class AgentSession { }; } + private _finishCancelledRetry(): void { + if (this._retryAttempt === 0) return; + const attempt = this._retryAttempt; + this._retryAttempt = 0; + this._emit({ + type: "auto_retry_end", + success: false, + attempt, + finalError: "Retry cancelled", + }); + } + /** * Prepare a retryable error for continuation with exponential backoff. * @returns true if the caller should continue the agent, false otherwise @@ -3023,14 +3052,7 @@ export class AgentSession { await sleep(delayMs, this._retryAbortController.signal); } catch { // Aborted during sleep - emit end event so UI can clean up - const attempt = this._retryAttempt; - this._retryAttempt = 0; - this._emit({ - type: "auto_retry_end", - success: false, - attempt, - finalError: "Retry cancelled", - }); + this._finishCancelledRetry(); return false; } finally { this._retryAbortController = undefined; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 4b514a77a..dd41a664f 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -4473,7 +4473,7 @@ export class InteractiveMode { if (allQueued.length === 0) { this.updatePendingMessagesDisplay(); if (options?.abort) { - this.agent.abort(); + void this.session.abort(); } return 0; } @@ -4483,7 +4483,7 @@ export class InteractiveMode { this.editor.setText(combinedText); this.updatePendingMessagesDisplay(); if (options?.abort) { - this.agent.abort(); + void this.session.abort(); } return allQueued.length; } diff --git a/packages/coding-agent/test/interactive-mode-compaction.test.ts b/packages/coding-agent/test/interactive-mode-compaction.test.ts index 378be9303..5b0980ba7 100644 --- a/packages/coding-agent/test/interactive-mode-compaction.test.ts +++ b/packages/coding-agent/test/interactive-mode-compaction.test.ts @@ -228,6 +228,24 @@ describe("InteractiveMode compaction events", () => { expect(fakeThis.ui.requestRender).toHaveBeenCalledTimes(2); }); + // Regression test for #9340. + test("routes interactive response aborts through AgentSession", () => { + const abort = vi.fn(async () => {}); + const ui = { + clearAllQueues: () => ({ steering: [], followUp: [] }), + updatePendingMessagesDisplay: vi.fn(), + session: { abort }, + }; + const restoreQueuedMessagesToEditor = Reflect.get(InteractiveMode.prototype, "restoreQueuedMessagesToEditor") as ( + this: typeof ui, + options?: { abort?: boolean }, + ) => number; + + restoreQueuedMessagesToEditor.call(ui, { abort: true }); + + expect(abort).toHaveBeenCalledOnce(); + }); + test("preserves steering behavior when flushing into an active agent run", async () => { const fakeThis = { compactionQueuedMessages: [{ text: "change direction", mode: "steer" as const }], 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 17648e8b9..597acc498 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 @@ -73,6 +73,32 @@ describe("AgentSession retry and event characterization", () => { expect(harness.faux.state.callCount).toBe(3); }); + // Regression test for #9340. + it("finalizes retry state when abort is requested after a retry attempt fails", async () => { + const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 0 } } }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + ]); + + let errorCount = 0; + harness.session.subscribe((event) => { + if (event.type !== "message_end" || event.message.role !== "assistant") return; + if (event.message.stopReason === "error" && ++errorCount === 2) void harness.session.abort(); + }); + + await harness.session.prompt("test"); + + expect(harness.session.retryAttempt).toBe(0); + expect(harness.eventsOfType("agent_end").at(-1)?.willRetry).toBe(false); + expect(harness.eventsOfType("auto_retry_end").at(-1)).toMatchObject({ + success: false, + attempt: 1, + finalError: "Retry cancelled", + }); + }); + it("exhausts max retries and emits a failure event", async () => { const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 2, baseDelayMs: 1 } } }); harnesses.push(harness); diff --git a/packages/coding-agent/test/suite/regressions/9340-9777-auto-compaction-cancellation.test.ts b/packages/coding-agent/test/suite/regressions/9340-9777-auto-compaction-cancellation.test.ts new file mode 100644 index 000000000..8f9792562 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/9340-9777-auto-compaction-cancellation.test.ts @@ -0,0 +1,158 @@ +import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createHarness, type Harness } from "../harness.ts"; + +type SessionWithCompactionInternals = { + _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; +}; + +function seedCompactableSession(harness: Harness): void { + const model = harness.getModel(); + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "x".repeat(500) }], + timestamp: 1, + }); + const assistant: AssistantMessage = { + ...fauxAssistantMessage("y".repeat(200), { timestamp: 2 }), + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 100, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 100, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }; + harness.sessionManager.appendMessage(assistant); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; +} + +function runAutoCompaction(harness: Harness): Promise { + return (harness.session as unknown as SessionWithCompactionInternals)._runAutoCompaction("threshold", false); +} + +describe("automatic compaction cancellation regressions", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + vi.restoreAllMocks(); + while (harnesses.length > 0) harnesses.pop()?.cleanup(); + }); + + // Regression test for #9340. + it("does not start post-run auto-compaction after abort", async () => { + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 200, maxTokens: 50 }], + settings: { + compaction: { enabled: true, reserveTokens: 50, keepRecentTokens: 1 }, + retry: { enabled: false }, + }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", () => ({ cancel: true })); + }, + ], + }); + harnesses.push(harness); + seedCompactableSession(harness); + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "Synthetic network failure" }), + ]); + harness.session.subscribe((event) => { + if (event.type === "message_end" && event.message.role === "assistant") { + harness.session.abortCompaction(); + void harness.session.abort(); + } + }); + + await harness.session.prompt("z".repeat(1000)); + + expect(harness.eventsOfType("compaction_start")).toHaveLength(0); + }); + + // Regression test for #9777. + it("cancels summarization authentication", async () => { + const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } } }); + harnesses.push(harness); + seedCompactableSession(harness); + let markAuthStarted = () => {}; + const authStarted = new Promise((resolve) => { + markAuthStarted = resolve; + }); + let authSignal: AbortSignal | undefined; + vi.spyOn(harness.session.modelRuntime, "getAuth").mockImplementation(async (_model, options) => { + authSignal = options?.signal; + markAuthStarted(); + if (!authSignal) throw new Error("Missing auth abort signal"); + return await new Promise((_resolve, reject) => { + authSignal?.addEventListener("abort", () => reject(authSignal?.reason), { once: true }); + }); + }); + + const compaction = runAutoCompaction(harness); + await authStarted; + const started = harness.eventsOfType("compaction_start").length; + const wasCompacting = harness.session.isCompacting; + await Promise.all([compaction, harness.session.abort()]); + + expect({ started, wasCompacting, authAborted: authSignal?.aborted }).toEqual({ + started: 1, + wasCompacting: true, + authAborted: true, + }); + expect(harness.eventsOfType("compaction_end").at(-1)?.aborted).toBe(true); + }); + + // Regression test for #9777. + it("cancels synchronously from compaction_start", async () => { + const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } } }); + harnesses.push(harness); + seedCompactableSession(harness); + harness.session.subscribe((event) => { + if (event.type === "compaction_start") harness.session.abortCompaction(); + }); + + await runAutoCompaction(harness); + + expect(harness.faux.state.callCount).toBe(0); + expect(harness.eventsOfType("compaction_end").at(-1)?.aborted).toBe(true); + }); + + // Regression test for #9777. + it.each([ + ["matching error text", () => new Error("Compaction cancelled")], + ["an unrelated AbortError", () => Object.assign(new Error("auth failed"), { name: "AbortError" })], + ] as const)("reports %s as a failure", async (_label, createError) => { + const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } } }); + harnesses.push(harness); + seedCompactableSession(harness); + vi.spyOn(harness.session.modelRuntime, "getAuth").mockRejectedValue(createError()); + + await runAutoCompaction(harness); + + const event = harness.eventsOfType("compaction_end").at(-1); + expect(event?.aborted).toBe(false); + expect(event?.errorMessage).toContain(createError().message); + }); + + it("reports extension cancellation as aborted", async () => { + const harness = await createHarness({ + settings: { compaction: { keepRecentTokens: 1 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", () => ({ cancel: true })); + }, + ], + }); + harnesses.push(harness); + seedCompactableSession(harness); + + await runAutoCompaction(harness); + + expect(harness.eventsOfType("compaction_end").at(-1)?.aborted).toBe(true); + }); +});