fix(agent): preserve run and custom message ordering

packages/agent/src/agent.ts

@@ -336,33 +336,42 @@ export class Agent {
 
 	/** Continue from the current transcript. The last LLM-context message must be a user or tool-result message. */
 	async continue(): Promise<void> {
-		if (this.activeRun) {
-			throw new Error("Agent is already processing. Wait for completion before continuing.");
-		}
+		const activeRun = this.reserveRun("Agent is already processing. Wait for completion before continuing.");
+		let executor: ((signal: AbortSignal) => Promise<void>) | undefined;
 
-		const llmMessages = await this.convertToLlm(this._state.messages);
-		const lastMessage = llmMessages[llmMessages.length - 1];
-		if (!lastMessage) {
-			throw new Error("No messages to continue from");
-		}
-
-		if (lastMessage.role === "assistant") {
-			const queuedSteering = this.steeringQueue.drain();
-			if (queuedSteering.length > 0) {
-				await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true });
-				return;
+		try {
+			const llmMessages = await this.convertToLlm(this._state.messages);
+			const lastMessage = llmMessages[llmMessages.length - 1];
+			if (!lastMessage) {
+				throw new Error("No messages to continue from");
 			}
 
-			const queuedFollowUps = this.followUpQueue.drain();
-			if (queuedFollowUps.length > 0) {
-				await this.runPromptMessages(queuedFollowUps);
-				return;
+			if (lastMessage.role === "assistant") {
+				const queuedSteering = this.steeringQueue.drain();
+				if (queuedSteering.length > 0) {
+					executor = (signal) =>
+						this.executePromptMessages(queuedSteering, { skipInitialSteeringPoll: true }, signal);
+				} else {
+					const queuedFollowUps = this.followUpQueue.drain();
+					if (queuedFollowUps.length > 0) {
+						executor = (signal) => this.executePromptMessages(queuedFollowUps, {}, signal);
+					} else {
+						throw new Error("Cannot continue from message role: assistant");
+					}
+				}
+			} else {
+				executor = (signal) => this.executeContinuation(signal);
 			}
-
-			throw new Error("Cannot continue from message role: assistant");
+		} catch (error) {
+			this.finishRun();
+			throw error;
 		}
 
-		await this.runContinuation();
+		if (!executor) {
+			this.finishRun();
+			throw new Error("No continuation executor prepared");
+		}
+		await this.runWithLifecycle(executor, activeRun);
 	}
 
 	private normalizePromptInput(
@@ -388,28 +397,32 @@ export class Agent {
 		messages: AgentMessage[],
 		options: { skipInitialSteeringPoll?: boolean } = {},
 	): Promise<void> {
-		await this.runWithLifecycle(async (signal) => {
-			await runAgentLoop(
-				messages,
-				this.createContextSnapshot(),
-				this.createLoopConfig(options),
-				(event) => this.processEvents(event),
-				signal,
-				this.streamFn,
-			);
-		});
+		await this.runWithLifecycle((signal) => this.executePromptMessages(messages, options, signal));
 	}
 
-	private async runContinuation(): Promise<void> {
-		await this.runWithLifecycle(async (signal) => {
-			await runAgentLoopContinue(
-				this.createContextSnapshot(),
-				this.createLoopConfig(),
-				(event) => this.processEvents(event),
-				signal,
-				this.streamFn,
-			);
-		});
+	private async executePromptMessages(
+		messages: AgentMessage[],
+		options: { skipInitialSteeringPoll?: boolean },
+		signal: AbortSignal,
+	): Promise<void> {
+		await runAgentLoop(
+			messages,
+			this.createContextSnapshot(),
+			this.createLoopConfig(options),
+			(event) => this.processEvents(event),
+			signal,
+			this.streamFn,
+		);
+	}
+
+	private async executeContinuation(signal: AbortSignal): Promise<void> {
+		await runAgentLoopContinue(
+			this.createContextSnapshot(),
+			this.createLoopConfig(),
+			(event) => this.processEvents(event),
+			signal,
+			this.streamFn,
+		);
 	}
 
 	private createContextSnapshot(): AgentContext {
@@ -449,9 +462,9 @@ export class Agent {
 		};
 	}
 
-	private async runWithLifecycle(executor: (signal: AbortSignal) => Promise<void>): Promise<void> {
+	private reserveRun(errorMessage: string): ActiveRun {
 		if (this.activeRun) {
-			throw new Error("Agent is already processing.");
+			throw new Error(errorMessage);
 		}
 
 		const abortController = new AbortController();
@@ -459,16 +472,29 @@ export class Agent {
 		const promise = new Promise<void>((resolve) => {
 			resolvePromise = resolve;
 		});
-		this.activeRun = { promise, resolve: resolvePromise, abortController };
+		const activeRun = { promise, resolve: resolvePromise, abortController };
+		this.activeRun = activeRun;
 
 		this._state.isStreaming = true;
 		this._state.streamingMessage = undefined;
 		this._state.errorMessage = undefined;
 
+		return activeRun;
+	}
+
+	private async runWithLifecycle(
+		executor: (signal: AbortSignal) => Promise<void>,
+		activeRun?: ActiveRun,
+	): Promise<void> {
+		const run = activeRun ?? this.reserveRun("Agent is already processing.");
+		if (this.activeRun !== run) {
+			throw new Error("Agent run reservation was lost.");
+		}
+
 		try {
-			await executor(abortController.signal);
+			await executor(run.abortController.signal);
 		} catch (error) {
-			await this.handleRunFailure(error, abortController.signal.aborted);
+			await this.handleRunFailure(error, run.abortController.signal.aborted);
 		} finally {
 			this.finishRun();
 		}

packages/agent/test/agent.test.ts

@@ -560,6 +560,49 @@ describe("Agent", () => {
 		await firstPrompt.catch(() => {});
 	});
 
+	it("continue() should reserve the active run while awaiting async LLM conversion", async () => {
+		const convertStarted = createDeferred();
+		const releaseConvert = createDeferred();
+		let convertCallCount = 0;
+		const agent = new Agent({
+			convertToLlm: async (messages) => {
+				convertCallCount++;
+				if (convertCallCount === 1) {
+					convertStarted.resolve();
+					await releaseConvert.promise;
+				}
+				return messages.filter(
+					(message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult",
+				) as Message[];
+			},
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") });
+				});
+				return stream;
+			},
+		});
+		agent.state.messages = [
+			{
+				role: "user",
+				content: [{ type: "text", text: "Initial" }],
+				timestamp: Date.now(),
+			},
+		];
+
+		const continuePromise = agent.continue();
+		await convertStarted.promise;
+
+		expect(agent.state.isStreaming).toBe(true);
+		await expect(agent.prompt("Second message")).rejects.toThrow("Agent is already processing a prompt");
+
+		releaseConvert.resolve();
+		await continuePromise;
+
+		expect(convertCallCount).toBe(2);
+	});
+
 	it("continue() should process queued follow-up messages after an assistant turn", async () => {
 		const agent = new Agent({
 			streamFn: () => {

packages/coding-agent/src/core/agent-session.ts

@@ -263,6 +263,7 @@ export class AgentSession {
 	// Event subscription state
 	private _unsubscribeAgent?: () => void;
 	private _eventListeners: AgentSessionEventListener[] = [];
+	private _deferredCustomMessagePersistenceStack: CustomMessage[][] = [];
 
 	/** Tracks pending steering messages for UI display. Removed when delivered. */
 	private _steeringMessages: string[] = [];
@@ -496,53 +497,59 @@ export class AgentSession {
 			}
 		}
 
-		// Emit to extensions first
-		await this._emitExtensionEvent(event);
+		const deferredCustomMessages: CustomMessage[] | undefined = event.type === "message_end" ? [] : undefined;
+		if (deferredCustomMessages) {
+			this._deferredCustomMessagePersistenceStack.push(deferredCustomMessages);
+		}
 
-		// Notify all listeners
-		this._emit(event.type === "agent_end" ? { ...event, willRetry: this._willRetryAfterAgentEnd(event) } : event);
+		try {
+			// Emit to extensions first
+			await this._emitExtensionEvent(event);
 
-		// Handle session persistence
-		if (event.type === "message_end") {
-			// Check if this is a custom message from extensions
-			if (event.message.role === "custom") {
-				// Persist as CustomMessageEntry
-				this.sessionManager.appendCustomMessageEntry(
-					event.message.customType,
-					event.message.content,
-					event.message.display,
-					event.message.details,
-					event.message.excludeFromContext,
-				);
-			} else if (
-				event.message.role === "user" ||
-				event.message.role === "assistant" ||
-				event.message.role === "toolResult"
-			) {
-				// Regular LLM message - persist as SessionMessageEntry
-				this.sessionManager.appendMessage(event.message);
+			// Notify all listeners
+			this._emit(event.type === "agent_end" ? { ...event, willRetry: this._willRetryAfterAgentEnd(event) } : event);
+
+			// Handle session persistence
+			if (event.type === "message_end") {
+				// Check if this is a custom message from extensions
+				if (event.message.role === "custom") {
+					// Persist as CustomMessageEntry
+					this._persistCustomMessage(event.message);
+				} else if (
+					event.message.role === "user" ||
+					event.message.role === "assistant" ||
+					event.message.role === "toolResult"
+				) {
+					// Regular LLM message - persist as SessionMessageEntry
+					this.sessionManager.appendMessage(event.message);
+				}
+				// Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
+				this._persistDeferredCustomMessages(deferredCustomMessages ?? []);
+
+				// Track assistant message for auto-compaction (checked on agent_end)
+				if (event.message.role === "assistant") {
+					this._lastAssistantMessage = event.message;
+
+					const assistantMsg = event.message as AssistantMessage;
+					if (assistantMsg.stopReason !== "error") {
+						this._overflowRecoveryAttempted = false;
+					}
+
+					// Reset retry counter immediately on successful assistant response
+					// This prevents accumulation across multiple LLM calls within a turn
+					if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
+						this._emit({
+							type: "auto_retry_end",
+							success: true,
+							attempt: this._retryAttempt,
+						});
+						this._retryAttempt = 0;
+					}
+				}
 			}
-			// Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
-
-			// Track assistant message for auto-compaction (checked on agent_end)
-			if (event.message.role === "assistant") {
-				this._lastAssistantMessage = event.message;
-
-				const assistantMsg = event.message as AssistantMessage;
-				if (assistantMsg.stopReason !== "error") {
-					this._overflowRecoveryAttempted = false;
-				}
-
-				// Reset retry counter immediately on successful assistant response
-				// This prevents accumulation across multiple LLM calls within a turn
-				if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
-					this._emit({
-						type: "auto_retry_end",
-						success: true,
-						attempt: this._retryAttempt,
-					});
-					this._retryAttempt = 0;
-				}
+		} finally {
+			if (deferredCustomMessages) {
+				this._deferredCustomMessagePersistenceStack.pop();
 			}
 		}
 	};
@@ -1350,8 +1357,7 @@ export class AgentSession {
 		}
 	}
 
-	private _recordCustomMessage(message: CustomMessage): void {
-		this.agent.state.messages.push(message);
+	private _persistCustomMessage(message: CustomMessage): void {
 		this.sessionManager.appendCustomMessageEntry(
 			message.customType,
 			message.content,
@@ -1359,6 +1365,23 @@ export class AgentSession {
 			message.details,
 			message.excludeFromContext,
 		);
+	}
+
+	private _persistDeferredCustomMessages(messages: CustomMessage[]): void {
+		for (const message of messages) {
+			this._persistCustomMessage(message);
+		}
+	}
+
+	private _recordCustomMessage(message: CustomMessage): void {
+		this.agent.state.messages.push(message);
+		const deferredMessages =
+			this._deferredCustomMessagePersistenceStack[this._deferredCustomMessagePersistenceStack.length - 1];
+		if (deferredMessages) {
+			deferredMessages.push(message);
+		} else {
+			this._persistCustomMessage(message);
+		}
 		this._emit({ type: "message_start", message });
 		this._emit({ type: "message_end", message });
 	}

packages/coding-agent/test/suite/agent-session-queue.test.ts

@@ -350,6 +350,44 @@ describe("AgentSession queue characterization", () => {
 		expect(harness.getPendingResponseCount()).toBe(1);
 	});
 
+	it("persists excluded custom messages from message_end hooks after the triggering message", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("message_end", (event) => {
+						if (event.message.role !== "assistant") return;
+						pi.sendMessage({
+							customType: "status",
+							content: "display only",
+							display: true,
+							details: {},
+							excludeFromContext: true,
+						});
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("reply")]);
+
+		await harness.session.prompt("hello");
+
+		const stateOrder = harness.session.messages.map((message) =>
+			message.role === "custom" ? `custom:${message.customType}` : message.role,
+		);
+		const branchOrder = harness.sessionManager.getBranch().map((entry) => {
+			if (entry.type === "message") {
+				return entry.message.role;
+			}
+			if (entry.type === "custom_message") {
+				return `custom:${entry.customType}`;
+			}
+			return entry.type;
+		});
+		expect(stateOrder).toEqual(["user", "assistant", "custom:status"]);
+		expect(branchOrder).toEqual(["user", "assistant", "custom:status"]);
+	});
+
 	it("queues custom messages with deliverAs steer while streaming", async () => {
 		const waiting = await createWaitingHarness();
 		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;