From 1edcc76f89505ebf719ad6faf73681ed70a75d53 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 14 Jun 2026 17:58:42 +0200 Subject: [PATCH] fix(agent): preserve run and custom message ordering --- packages/agent/src/agent.ts | 118 +++++++++++------- packages/agent/test/agent.test.ts | 43 +++++++ .../coding-agent/src/core/agent-session.ts | 115 ++++++++++------- .../test/suite/agent-session-queue.test.ts | 38 ++++++ 4 files changed, 222 insertions(+), 92 deletions(-) diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 56283849b..3ab2b2be5 100644 --- a/packages/agent/src/agent.ts +++ b/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 { - 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) | 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 { - 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 { - 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 { + await runAgentLoop( + messages, + this.createContextSnapshot(), + this.createLoopConfig(options), + (event) => this.processEvents(event), + signal, + this.streamFn, + ); + } + + private async executeContinuation(signal: AbortSignal): Promise { + 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): Promise { + 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((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, + activeRun?: ActiveRun, + ): Promise { + 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(); } diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 8839aa393..16675a38f 100644 --- a/packages/agent/test/agent.test.ts +++ b/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: () => { diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 4d82185a6..e722bfc7f 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/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 }); } diff --git a/packages/coding-agent/test/suite/agent-session-queue.test.ts b/packages/coding-agent/test/suite/agent-session-queue.test.ts index 7da0b70f0..119b688fb 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/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;