Merge PR #7539: fix(coding-agent): complete usage-aware fallback integration (@eggpeat)

packages/agent/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- Preserved queued steering and follow-up messages when a continuation is cancelled before or during pre-dequeue hooks, and propagated the caller's cancellation signal through every continuation model-call loop.
+
 ## [17.2.6] - 2026-08-03
 
 ### Fixed

packages/agent/src/agent-loop.ts

@@ -1016,7 +1016,7 @@ async function runLoopBody(
 		// Skip when the run is already externally aborted — dequeuing would strand
 		// the messages in a run that is about to die.
 		try {
-			pendingMessages = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
+			pendingMessages = signal?.aborted ? [] : (await config.getSteeringMessages?.(signal)) || [];
 		} catch (error) {
 			stream.push({ type: "turn_start" });
 			emitInputMessages(stream, messagesToEmit);
@@ -1075,7 +1075,7 @@ async function runLoopBody(
 				let gateResult: AgentPreModelCallResult;
 				try {
 					if (config.syncContextBeforeModelCall) {
-						await config.syncContextBeforeModelCall(currentContext);
+						await config.syncContextBeforeModelCall(currentContext, signal);
 					}
 
 					if (!directiveResolvedForTurn) {
@@ -1421,7 +1421,7 @@ async function runLoopBody(
 				// instantly aborts — message lands in history, agent never responds. The
 				// mid-batch interrupt poll only peeks (hasSteeringMessages), so the queue
 				// still owns every message until this dequeue.
-				const steering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
+				const steering = signal?.aborted ? [] : (await config.getSteeringMessages?.(signal)) || [];
 				if (hasMoreToolCalls) {
 					// Mid-work: fold any non-interrupting asides into the next turn alongside steering.
 					const asides = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.());
@@ -1450,9 +1450,9 @@ async function runLoopBody(
 			// Re-poll steering too: a steer can land between the stop-boundary dequeue
 			// above and this yield point (e.g. queued while onBeforeYield ran). Without
 			// this poll it would strand in the queue until the next manual prompt.
-			const lateSteering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
+			const lateSteering = signal?.aborted ? [] : (await config.getSteeringMessages?.(signal)) || [];
 			const asideMessages = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.());
-			const followUpMessages = signal?.aborted ? [] : (await config.getFollowUpMessages?.()) || [];
+			const followUpMessages = signal?.aborted ? [] : (await config.getFollowUpMessages?.(signal)) || [];
 			if (lateSteering.length > 0 || asideMessages.length > 0 || followUpMessages.length > 0) {
 				// Set as pending so the inner loop processes them before stopping.
 				pendingMessages = [...lateSteering, ...asideMessages, ...followUpMessages];

packages/agent/src/agent.ts

@@ -426,6 +426,8 @@ export class Agent {
 	#asideMessageProvider?: () => AsideMessage[] | Promise<AsideMessage[]>;
 	#telemetry?: AgentLoopConfig["telemetry"];
 	#appendOnlyContext?: AppendOnlyContextManager;
+	#beforeQueuedMessageDequeueHooks = new Set<(signal?: AbortSignal) => Promise<void> | void>();
+	#beforeModelCallHooks = new Set<(signal?: AbortSignal) => Promise<void> | void>();
 
 	/** Buffered Cursor tool results with text length at time of call (for correct ordering) */
 	#cursorToolResultBuffer: CursorToolResultEntry[] = [];
@@ -784,6 +786,40 @@ export class Agent {
 		return () => this.#listeners.delete(fn);
 	}
 
+	/** Register an independently removable hook that runs before queued messages are consumed. */
+	addBeforeQueuedMessageDequeueHook(hook: (signal?: AbortSignal) => Promise<void> | void): () => void {
+		const registration = (signal?: AbortSignal) => hook(signal);
+		this.#beforeQueuedMessageDequeueHooks.add(registration);
+		return () => this.#beforeQueuedMessageDequeueHooks.delete(registration);
+	}
+
+	/** Register an independently removable hook that runs immediately before each model call. */
+	addBeforeModelCallHook(hook: (signal?: AbortSignal) => Promise<void> | void): () => void {
+		const registration = (signal?: AbortSignal) => hook(signal);
+		this.#beforeModelCallHooks.add(registration);
+		return () => this.#beforeModelCallHooks.delete(registration);
+	}
+
+	async #runBeforeModelCallHooks(signal?: AbortSignal): Promise<void> {
+		for (const hook of this.#beforeModelCallHooks) await hook(signal);
+	}
+
+	async #runBeforeQueuedMessageDequeueHooks(signal?: AbortSignal): Promise<void> {
+		for (const hook of this.#beforeQueuedMessageDequeueHooks) await hook(signal);
+	}
+
+	async #dequeueSteeringMessagesAfterHooks(signal?: AbortSignal): Promise<AgentMessage[]> {
+		if (signal?.aborted || this.#steeringQueue.length === 0) return [];
+		await this.#runBeforeQueuedMessageDequeueHooks(signal);
+		return signal?.aborted ? [] : this.#dequeueSteeringMessages();
+	}
+
+	async #dequeueFollowUpMessagesAfterHooks(signal?: AbortSignal): Promise<AgentMessage[]> {
+		if (signal?.aborted || this.#followUpQueue.length === 0) return [];
+		await this.#runBeforeQueuedMessageDequeueHooks(signal);
+		return signal?.aborted ? [] : this.#dequeueFollowUpMessages();
+	}
+
 	setProviderResponseInterceptor(fn: SimpleStreamOptions["onResponse"] | undefined): void {
 		this.#onResponse = fn;
 	}
@@ -1137,48 +1173,90 @@ export class Agent {
 	/**
 	 * Continue from current context (used for retries and resuming queued messages).
 	 */
-	async continue() {
+	#continuationDequeueSignal(signal?: AbortSignal): AbortSignal | undefined {
+		const signals: AbortSignal[] = [];
+		if (this.#abortController) signals.push(this.#abortController.signal);
+		if (signal) signals.push(signal);
+		if (this.#deadline !== undefined) {
+			const delay = this.#deadline - Date.now();
+			if (delay <= 0) {
+				const controller = new AbortController();
+				controller.abort(new DOMException("Deadline exceeded", "TimeoutError"));
+				signals.push(controller.signal);
+			} else {
+				signals.push(AbortSignal.timeout(delay));
+			}
+		}
+		if (signals.length === 0) return undefined;
+		return signals.length === 1 ? signals[0] : AbortSignal.any(signals);
+	}
+
+	async continue(signal?: AbortSignal) {
 		if (this.#state.isStreaming) {
 			throw new AgentBusyError();
 		}
 
-		const messages = this.#state.messages;
-		if (messages.length === 0) {
-			// An empty transcript has nothing to resume, but a queued steer/follow-up
-			// must still be delivered as the opening turn — mirroring the assistant-tail
-			// branch below. Throwing here leaves the message undeliverable, and idle-drain
-			// callers (AgentSession#scheduleQueuedMessageDrain) re-arm continue() on every
-			// microtask because hasQueuedMessages() never clears, spinning an unbounded
-			// allocation loop until OOM (issue #6344).
-			const queuedSteering = this.#dequeueSteeringMessages();
-			if (queuedSteering.length > 0) {
-				await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true });
-				return;
+		const { promise, resolve } = Promise.withResolvers<void>();
+		this.#runningPrompt = promise;
+		this.#resolveRunningPrompt = resolve;
+		const continuationAbortController = new AbortController();
+		this.#abortController = continuationAbortController;
+		this.#state.isStreaming = true;
+		this.#state.streamMessage = null;
+		this.#state.error = undefined;
+
+		try {
+			const dequeueSignal = this.#continuationDequeueSignal(signal);
+			const messages = this.#state.messages;
+			if (messages.length === 0) {
+				// An empty transcript has nothing to resume, but a queued steer/follow-up
+				// must still be delivered as the opening turn — mirroring the assistant-tail
+				// branch below. Throwing here leaves the message undeliverable, and idle-drain
+				// callers (AgentSession#scheduleQueuedMessageDrain) re-arm continue() on every
+				// microtask because hasQueuedMessages() never clears, spinning an unbounded
+				// allocation loop until OOM (issue #6344).
+				const queuedSteering = await this.#dequeueSteeringMessagesAfterHooks(dequeueSignal);
+				if (queuedSteering.length > 0) {
+					await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }, signal, true);
+					return;
+				}
+				const queuedFollowUp = await this.#dequeueFollowUpMessagesAfterHooks(dequeueSignal);
+				if (queuedFollowUp.length > 0) {
+					await this.#runLoop(queuedFollowUp, undefined, signal, true);
+					return;
+				}
+				throw new Error("No messages to continue from");
 			}
-			const queuedFollowUp = this.#dequeueFollowUpMessages();
-			if (queuedFollowUp.length > 0) {
-				await this.#runLoop(queuedFollowUp);
-				return;
+			if (messages[messages.length - 1].role === "assistant") {
+				const queuedSteering = await this.#dequeueSteeringMessagesAfterHooks(dequeueSignal);
+				if (queuedSteering.length > 0) {
+					await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }, signal, true);
+					return;
+				}
+
+				const queuedFollowUp = await this.#dequeueFollowUpMessagesAfterHooks(dequeueSignal);
+				if (queuedFollowUp.length > 0) {
+					await this.#runLoop(queuedFollowUp, undefined, signal, true);
+					return;
+				}
+
+				throw new Error("Cannot continue from message role: assistant");
+			}
+
+			await this.#runLoop(undefined, undefined, signal, true);
+		} finally {
+			resolve();
+			if (this.#abortController === continuationAbortController) {
+				this.#state.isStreaming = false;
+				this.#state.streamMessage = null;
+				this.#state.pendingToolCalls.clear();
+				this.#abortController = undefined;
+				if (this.#runningPrompt === promise) {
+					this.#runningPrompt = undefined;
+					this.#resolveRunningPrompt = undefined;
+				}
 			}
-			throw new Error("No messages to continue from");
 		}
-		if (messages[messages.length - 1].role === "assistant") {
-			const queuedSteering = this.#dequeueSteeringMessages();
-			if (queuedSteering.length > 0) {
-				await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true });
-				return;
-			}
-
-			const queuedFollowUp = this.#dequeueFollowUpMessages();
-			if (queuedFollowUp.length > 0) {
-				await this.#runLoop(queuedFollowUp);
-				return;
-			}
-
-			throw new Error("Cannot continue from message role: assistant");
-		}
-
-		await this.#runLoop(undefined);
 	}
 
 	/**
@@ -1186,17 +1264,29 @@ export class Agent {
 	 * If messages are provided, starts a new conversation turn with those messages.
 	 * Otherwise, continues from existing context.
 	 */
-	async #runLoop(messages?: AgentMessage[], options?: AgentPromptOptions & { skipInitialSteeringPoll?: boolean }) {
+	async #runLoop(
+		messages?: AgentMessage[],
+		options?: AgentPromptOptions & { skipInitialSteeringPoll?: boolean },
+		continuationSignal?: AbortSignal,
+		runStateClaimed = false,
+	) {
 		const model = this.#state.model;
 		if (!model) throw new Error("No model configured");
 
 		let skipInitialSteeringPoll = options?.skipInitialSteeringPoll === true;
 		using _ = new EventLoopKeepalive();
-		const { promise, resolve } = Promise.withResolvers<void>();
-		this.#runningPrompt = promise;
-		this.#resolveRunningPrompt = resolve;
-
-		this.#abortController = new AbortController();
+		if (!runStateClaimed) {
+			const { promise, resolve } = Promise.withResolvers<void>();
+			this.#runningPrompt = promise;
+			this.#resolveRunningPrompt = resolve;
+			this.#abortController = new AbortController();
+		}
+		const resolveRun = this.#resolveRunningPrompt;
+		const loopAbortController = this.#abortController;
+		if (!loopAbortController) throw new Error("Agent run state was not initialized");
+		const loopSignal = continuationSignal
+			? AbortSignal.any([loopAbortController.signal, continuationSignal])
+			: loopAbortController.signal;
 		this.#state.isStreaming = true;
 		this.#state.streamMessage = null;
 		this.#state.error = undefined;
@@ -1315,7 +1405,8 @@ export class Agent {
 			onSseEvent: this.#onSseEvent,
 			getApiKey: this.getApiKey,
 			getToolContext: this.#getToolContext,
-			syncContextBeforeModelCall: async context => {
+			syncContextBeforeModelCall: async (context, signal) => {
+				await this.#runBeforeModelCallHooks(signal);
 				if (this.#listeners.size > 0) {
 					await Bun.sleep(0);
 				}
@@ -1362,12 +1453,12 @@ export class Agent {
 			getReasoning: () => this.#state.thinkingLevel,
 			getDisableReasoning: () => this.#state.disableReasoning,
 			getServiceTier: this.#serviceTierResolver,
-			getSteeringMessages: async () => {
+			getSteeringMessages: async signal => {
 				if (skipInitialSteeringPoll) {
 					skipInitialSteeringPoll = false;
 					return [];
 				}
-				return this.#dequeueSteeringMessages();
+				return this.#dequeueSteeringMessagesAfterHooks(signal);
 			},
 			hasSteeringMessages: () => {
 				if (this.#steeringQueue.length === 0) {
@@ -1392,7 +1483,7 @@ export class Agent {
 			},
 			waitForSteeringMessages: signal => this.#waitForSteeringMessages(signal),
 			hasIrcInterrupts: this.hasIrcInterrupts,
-			getFollowUpMessages: async () => this.#dequeueFollowUpMessages(),
+			getFollowUpMessages: signal => this.#dequeueFollowUpMessagesAfterHooks(signal),
 			getAsideMessages: async () => (await this.#asideMessageProvider?.()) ?? [],
 			onBeforeYield: () => this.#onBeforeYield?.(),
 			telemetry: this.#telemetry,
@@ -1404,8 +1495,8 @@ export class Agent {
 
 		try {
 			const stream = messages
-				? agentLoop(messages, context, config, this.#abortController.signal, this.streamFn)
-				: agentLoopContinue(context, config, this.#abortController.signal, this.streamFn);
+				? agentLoop(messages, context, config, loopSignal, this.streamFn)
+				: agentLoopContinue(context, config, loopSignal, this.streamFn);
 
 			for await (const event of stream) {
 				if (event.type === "turn_start") turnOpen = true;
@@ -1472,15 +1563,15 @@ export class Agent {
 				if (!onlyEmpty) {
 					this.appendMessage(partial);
 				} else {
-					if (this.#abortController?.signal.aborted) {
+					if (loopSignal.aborted) {
 						throw new Error("Request was aborted");
 					}
 				}
 			}
 		} catch (err) {
-			const stoppedForAbort = this.#abortController?.signal.aborted === true;
+			const stoppedForAbort = loopSignal.aborted;
 			const errorMessage = stoppedForAbort
-				? abortReasonText(this.#abortController?.signal)
+				? abortReasonText(loopSignal)
 				: err instanceof Error
 					? err.message
 					: String(err);
@@ -1582,13 +1673,15 @@ export class Agent {
 				this.#emit({ type: "agent_end", messages: [errorMsg] });
 			}
 		} finally {
-			this.#state.isStreaming = false;
-			this.#state.streamMessage = null;
-			this.#state.pendingToolCalls.clear();
-			this.#abortController = undefined;
-			this.#resolveRunningPrompt?.();
-			this.#runningPrompt = undefined;
-			this.#resolveRunningPrompt = undefined;
+			resolveRun?.();
+			if (this.#abortController === loopAbortController) {
+				this.#state.isStreaming = false;
+				this.#state.streamMessage = null;
+				this.#state.pendingToolCalls.clear();
+				this.#abortController = undefined;
+				this.#runningPrompt = undefined;
+				this.#resolveRunningPrompt = undefined;
+			}
 		}
 	}
 

packages/agent/src/types.ts

@@ -240,7 +240,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
 	 * mid-batch interrupt poll uses {@link hasSteeringMessages} instead and
 	 * never consumes the queue.
 	 */
-	getSteeringMessages?: () => Promise<AgentMessage[]>;
+	getSteeringMessages?: (signal?: AbortSignal) => Promise<AgentMessage[]>;
 
 	/**
 	 * Peeks whether steering messages are queued, without consuming them.
@@ -285,7 +285,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
 	 * If messages are returned, they're added to the context and the agent
 	 * continues with another turn.
 	 */
-	getFollowUpMessages?: () => Promise<AgentMessage[]>;
+	getFollowUpMessages?: (signal?: AbortSignal) => Promise<AgentMessage[]>;
 	/**
 	 * Returns non-interrupting "aside" messages to inject at a step boundary.
 	 *
@@ -319,7 +319,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
 	 * Mutate the agent context here; use `beforeModelCall` to inspect the
 	 * provider-bound context.
 	 */
-	syncContextBeforeModelCall?: (context: AgentContext) => void | Promise<void>;
+	syncContextBeforeModelCall?: (context: AgentContext, signal?: AbortSignal) => void | Promise<void>;
 
 	/**
 	 * Asked after the complete provider context has been built, including

packages/agent/test/agent.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "bun:test";
-import { Agent, type AgentEvent, type AgentTool, ThinkingLevel } from "@oh-my-pi/pi-agent-core";
+import { Agent, AgentBusyError, type AgentEvent, type AgentTool, ThinkingLevel } from "@oh-my-pi/pi-agent-core";
 import { type SimpleStreamOptions, type ToolResultMessage, z } from "@oh-my-pi/pi-ai";
 import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
 import { kCursorExecResolved } from "@oh-my-pi/pi-ai/utils/block-symbols";
@@ -240,6 +240,213 @@ describe("Agent", () => {
 		}
 	});
 
+	it("removes duplicate queued-message hooks independently", async () => {
+		const mock = createMockModel({ responses: [{ content: ["first"] }, { content: ["second"] }] });
+		const agent = new Agent({ streamFn: mock.stream });
+		agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]);
+		let calls = 0;
+		const signals: Array<AbortSignal | undefined> = [];
+		const hook = (signal?: AbortSignal) => {
+			calls++;
+			signals.push(signal);
+		};
+		const removeFirst = agent.addBeforeQueuedMessageDequeueHook(hook);
+		const removeSecond = agent.addBeforeQueuedMessageDequeueHook(hook);
+
+		const controller = new AbortController();
+		removeFirst();
+		agent.followUp({ role: "user", content: "first turn", timestamp: Date.now() });
+		await agent.continue(controller.signal);
+		expect(calls).toBe(1);
+		expect(signals).toEqual([controller.signal]);
+
+		removeSecond();
+		agent.followUp({ role: "user", content: "second turn", timestamp: Date.now() });
+		await agent.continue();
+		expect(calls).toBe(1);
+	});
+
+	it("continue() leaves queued messages owned when its signal is already aborted", async () => {
+		const agent = new Agent();
+		agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]);
+		agent.followUp({ role: "user", content: "stay queued", timestamp: Date.now() });
+		const controller = new AbortController();
+		controller.abort();
+
+		await expect(agent.continue(controller.signal)).rejects.toThrow("Cannot continue from message role: assistant");
+		expect(agent.peekFollowUpQueue()).toHaveLength(1);
+	});
+	it("keeps follow-up ownership when the deadline expires during a dequeue hook", async () => {
+		const mock = createMockModel({ responses: [{ content: ["done"] }] });
+		const agent = new Agent({ streamFn: mock.stream, deadline: Date.now() + 25 });
+		let hookSignal: AbortSignal | undefined;
+		agent.addBeforeQueuedMessageDequeueHook(async signal => {
+			if (!signal) throw new Error("Expected the active loop signal");
+			hookSignal = signal;
+			if (signal.aborted) return;
+			const { promise, resolve } = Promise.withResolvers<void>();
+			signal.addEventListener("abort", () => resolve(), { once: true });
+			await promise;
+		});
+		agent.followUp({ role: "user", content: "stay queued after deadline", timestamp: Date.now() });
+
+		await agent.prompt("start");
+
+		expect(hookSignal?.aborted).toBe(true);
+		expect(agent.peekFollowUpQueue()).toHaveLength(1);
+	});
+	it("keeps queued work when continue() reaches its deadline inside a dequeue hook", async () => {
+		const agent = new Agent({ deadline: Date.now() + 25 });
+		agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]);
+		agent.addBeforeQueuedMessageDequeueHook(async signal => {
+			if (!signal) throw new Error("Expected the deadline-aware dequeue signal");
+			if (signal.aborted) return;
+			const { promise, resolve } = Promise.withResolvers<void>();
+			signal.addEventListener("abort", () => resolve(), { once: true });
+			await promise;
+		});
+		agent.followUp({ role: "user", content: "stay queued before run loop", timestamp: Date.now() });
+
+		await expect(agent.continue()).rejects.toThrow("Cannot continue from message role: assistant");
+
+		expect(agent.peekFollowUpQueue()).toHaveLength(1);
+	});
+
+	it("claims an abortable busy state while continue() awaits dequeue hooks", async () => {
+		const agent = new Agent();
+		agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]);
+		agent.followUp({ role: "user", content: "stay queued", timestamp: Date.now() });
+		const hookStarted = Promise.withResolvers<void>();
+		agent.addBeforeQueuedMessageDequeueHook(async signal => {
+			if (!signal) throw new Error("Expected continuation dequeue signal");
+			hookStarted.resolve();
+			if (signal.aborted) return;
+			const { promise, resolve } = Promise.withResolvers<void>();
+			signal.addEventListener("abort", () => resolve(), { once: true });
+			await promise;
+		});
+
+		const continuing = agent.continue();
+		await hookStarted.promise;
+		let idleResolved = false;
+		const idle = agent.waitForIdle().then(() => {
+			idleResolved = true;
+		});
+		await Promise.resolve();
+
+		expect(agent.state.isStreaming).toBe(true);
+		expect(idleResolved).toBe(false);
+		await expect(agent.prompt("must not overlap")).rejects.toBeInstanceOf(AgentBusyError);
+
+		agent.abort("cancel dequeue");
+		await expect(continuing).rejects.toThrow("Cannot continue from message role: assistant");
+		await idle;
+		expect(idleResolved).toBe(true);
+		expect(agent.state.isStreaming).toBe(false);
+		expect(agent.peekFollowUpQueue()).toHaveLength(1);
+	});
+
+	it("does not clear a successor prompt after continue() releases idle waiters", async () => {
+		const firstStarted = Promise.withResolvers<void>();
+		const releaseFirst = Promise.withResolvers<void>();
+		const secondStarted = Promise.withResolvers<void>();
+		const releaseSecond = Promise.withResolvers<void>();
+		const mock = createMockModel({
+			responses: [
+				async () => {
+					firstStarted.resolve();
+					await releaseFirst.promise;
+					return { content: ["continued"] };
+				},
+				async () => {
+					secondStarted.resolve();
+					await releaseSecond.promise;
+					return { content: ["successor"] };
+				},
+			],
+		});
+		const agent = new Agent({ streamFn: mock.stream });
+		agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]);
+		agent.followUp({ role: "user", content: "continue", timestamp: Date.now() });
+
+		const continuing = agent.continue();
+		await firstStarted.promise;
+		const successor = agent.waitForIdle().then(() => agent.prompt("next prompt"));
+		releaseFirst.resolve();
+		await secondStarted.promise;
+		await continuing;
+
+		expect(agent.state.isStreaming).toBe(true);
+		releaseSecond.resolve();
+		await successor;
+		expect(agent.state.isStreaming).toBe(false);
+	});
+
+	it("resolves a predecessor idle waiter when agent_end starts a successor", async () => {
+		const secondStarted = Promise.withResolvers<void>();
+		const releaseSecond = Promise.withResolvers<void>();
+		const mock = createMockModel({
+			responses: [
+				{ content: ["first"] },
+				async () => {
+					secondStarted.resolve();
+					await releaseSecond.promise;
+					return { content: ["second"] };
+				},
+			],
+		});
+		const agent = new Agent({ streamFn: mock.stream });
+		let successor: Promise<void> | undefined;
+		agent.subscribe(event => {
+			if (event.type === "agent_end" && !successor) {
+				successor = agent.prompt("successor");
+			}
+		});
+
+		const predecessor = agent.prompt("predecessor");
+		let predecessorIdleResolved = false;
+		void agent.waitForIdle().then(() => {
+			predecessorIdleResolved = true;
+		});
+		await secondStarted.promise;
+		await predecessor;
+		expect(agent.state.isStreaming).toBe(true);
+
+		releaseSecond.resolve();
+		await successor;
+		await Promise.resolve();
+		expect(predecessorIdleResolved).toBe(true);
+		expect(agent.state.isStreaming).toBe(false);
+	});
+
+	it("classifies an in-flight continuation cancellation as aborted", async () => {
+		const providerStarted = Promise.withResolvers<AbortSignal>();
+		const agent = new Agent({
+			streamFn: (_model, _context, options) => {
+				const signal = options?.signal;
+				if (!signal) throw new Error("Expected provider abort signal");
+				providerStarted.resolve(signal);
+				const stream = new AssistantMessageEventStream();
+				signal.addEventListener("abort", () => stream.fail(new Error("provider aborted")), { once: true });
+				return stream;
+			},
+		});
+		agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]);
+		agent.followUp({ role: "user", content: "cancel this continuation", timestamp: Date.now() });
+		const controller = new AbortController();
+
+		const running = agent.continue(controller.signal);
+		await providerStarted.promise;
+		controller.abort("caller cancelled");
+		await running;
+
+		const finalMessage = agent.state.messages.at(-1);
+		expect(finalMessage?.role).toBe("assistant");
+		if (finalMessage?.role !== "assistant") throw new Error("Expected aborted assistant message");
+		expect(finalMessage.stopReason).toBe("aborted");
+		expect(finalMessage.errorMessage).toBe("caller cancelled");
+	});
+
 	it("continue() should process queued follow-up messages after an assistant turn", async () => {
 		const mock = createMockModel({ responses: [{ content: ["Processed"] }] });
 		const agent = new Agent({ streamFn: mock.stream });
@@ -276,6 +483,12 @@ describe("Agent", () => {
 			responses: [{ content: ["Processed 1"] }, { content: ["Processed 2"] }],
 		});
 		const agent = new Agent({ streamFn: mock.stream });
+		let dequeueHooks = 0;
+		const dequeueSignals: Array<AbortSignal | undefined> = [];
+		agent.addBeforeQueuedMessageDequeueHook(signal => {
+			dequeueHooks++;
+			dequeueSignals.push(signal);
+		});
 
 		agent.replaceMessages([
 			{
@@ -297,11 +510,16 @@ describe("Agent", () => {
 			timestamp: Date.now() + 1,
 		});
 
-		await expect(agent.continue()).resolves.toBeUndefined();
+		const controller = new AbortController();
+		await expect(agent.continue(controller.signal)).resolves.toBeUndefined();
 
 		const recentMessages = agent.state.messages.slice(-4);
 		expect(recentMessages.map(m => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
 		expect(mock.calls.length).toBe(2);
+		expect(dequeueHooks).toBe(2);
+		expect(dequeueSignals).toHaveLength(2);
+		controller.abort();
+		expect(dequeueSignals.every(signal => signal?.aborted === true)).toBe(true);
 	});
 
 	it("delivers a steer that lands at the yield boundary instead of stranding it", async () => {
@@ -856,6 +1074,10 @@ describe("Agent", () => {
 			},
 			streamFn: mock.stream,
 		});
+		let beforeModelCalls = 0;
+		agent.addBeforeModelCallHook(() => {
+			beforeModelCalls++;
+		});
 
 		const unsubscribe = agent.subscribe(event => {
 			if (event.type === "message_end" && event.message.role === "toolResult") {
@@ -875,6 +1097,7 @@ describe("Agent", () => {
 			{ systemPrompt: "prompt-one", toolNames: ["alpha"] },
 			{ systemPrompt: "prompt-two", toolNames: ["alpha", "beta"] },
 		]);
+		expect(beforeModelCalls).toBe(2);
 	});
 
 	it("prompt() drops stale forced toolChoice after same-turn tool refresh", async () => {