From 1ee0d28d7b9701e8a5d10f491370b29afe56787f Mon Sep 17 00:00:00 2001 From: Duncan Ogilvie Date: Sun, 29 Mar 2026 19:23:40 +0200 Subject: [PATCH] fix(coding-agent): properly queue extension messages (#2674) --- .../coding-agent/src/core/agent-session.ts | 20 ++- .../src/modes/interactive/interactive-mode.ts | 5 + .../test/agent-session-concurrent.test.ts | 123 +++++++++++++++++- 3 files changed, 145 insertions(+), 3 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 821a6182..42124f41 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -112,6 +112,11 @@ export function parseSkillBlock(text: string): ParsedSkillBlock | null { /** Session-specific events that extend the core AgentEvent */ export type AgentSessionEvent = | AgentEvent + | { + type: "queue_update"; + steering: readonly string[]; + followUp: readonly string[]; + } | { type: "compaction_start"; reason: "manual" | "threshold" | "overflow" } | { type: "compaction_end"; @@ -414,6 +419,14 @@ export class AgentSession { } } + private _emitQueueUpdate(): void { + this._emit({ + type: "queue_update", + steering: [...this._steeringMessages], + followUp: [...this._followUpMessages], + }); + } + // Track last assistant message for auto-compaction check private _lastAssistantMessage: AssistantMessage | undefined = undefined; @@ -476,11 +489,13 @@ export class AgentSession { const steeringIndex = this._steeringMessages.indexOf(messageText); if (steeringIndex !== -1) { this._steeringMessages.splice(steeringIndex, 1); + this._emitQueueUpdate(); } else { // Check follow-up queue const followUpIndex = this._followUpMessages.indexOf(messageText); if (followUpIndex !== -1) { this._followUpMessages.splice(followUpIndex, 1); + this._emitQueueUpdate(); } } } @@ -1155,6 +1170,7 @@ export class AgentSession { */ private async _queueSteer(text: string, images?: ImageContent[]): Promise { this._steeringMessages.push(text); + this._emitQueueUpdate(); const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; if (images) { content.push(...images); @@ -1171,6 +1187,7 @@ export class AgentSession { */ private async _queueFollowUp(text: string, images?: ImageContent[]): Promise { this._followUpMessages.push(text); + this._emitQueueUpdate(); const content: (TextContent | ImageContent)[] = [{ type: "text", text }]; if (images) { content.push(...images); @@ -1282,7 +1299,7 @@ export class AgentSession { expandPromptTemplates: false, streamingBehavior: options?.deliverAs, images, - source: "extension", + source: "interactive", }); } @@ -1297,6 +1314,7 @@ export class AgentSession { this._steeringMessages = []; this._followUpMessages = []; this.agent.clearAllQueues(); + this._emitQueueUpdate(); return { steering, followUp }; } diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 64876740..7c933e67 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -2244,6 +2244,11 @@ export class InteractiveMode { this.ui.requestRender(); break; + case "queue_update": + this.updatePendingMessagesDisplay(); + this.ui.requestRender(); + break; + case "message_start": if (event.message.role === "custom") { this.addMessageToChat(event.message); diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index 4b7e2f94..f3281e97 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -6,7 +6,14 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent } from "@mariozechner/pi-agent-core"; -import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@mariozechner/pi-ai"; +import { + type AssistantMessage, + type AssistantMessageEvent, + EventStream, + getModel, + type ImageContent, + type TextContent, +} from "@mariozechner/pi-ai"; import { Type } from "@sinclair/typebox"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.js"; @@ -14,7 +21,7 @@ import { AuthStorage } from "../src/core/auth-storage.js"; import { ModelRegistry } from "../src/core/model-registry.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; -import { createTestResourceLoader } from "./utilities.js"; +import { createTestExtensionsResult, createTestResourceLoader } from "./utilities.js"; // Mock stream that mimics AssistantMessageEventStream class MockAssistantStream extends EventStream { @@ -60,6 +67,8 @@ describe("AgentSession concurrent prompt guard", () => { }); afterEach(async () => { + delete (globalThis as typeof globalThis & { testExtensionApi?: unknown }).testExtensionApi; + delete (globalThis as typeof globalThis & { testCommandRuns?: unknown }).testCommandRuns; if (session) { session.dispose(); } @@ -171,6 +180,116 @@ describe("AgentSession concurrent prompt guard", () => { await firstPrompt.catch(() => {}); }); + it("should queue extension-origin steering messages while streaming", async () => { + const model = getModel("anthropic", "claude-sonnet-4-5")!; + let abortSignal: AbortSignal | undefined; + let sawSteeringMessage = false; + let lastInputSource: string | undefined; + const queueEvents: Array<{ steering: readonly string[]; followUp: readonly string[] }> = []; + + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { + model, + systemPrompt: "Test", + tools: [], + }, + streamFn: (_model, context, options) => { + abortSignal = options?.signal; + const stream = new MockAssistantStream(); + queueMicrotask(() => { + const userTexts = context.messages + .filter((message) => message.role === "user") + .map((message) => { + if (typeof message.content === "string") { + return message.content; + } + return message.content + .filter((part): part is TextContent | ImageContent => typeof part === "object" && part !== null) + .filter((part): part is TextContent => part.type === "text") + .map((part) => part.text) + .join("\n"); + }); + + if (userTexts.includes("Steer from extension")) { + sawSteeringMessage = true; + stream.push({ type: "start", partial: createAssistantMessage("") }); + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Steered") }); + return; + } + + stream.push({ type: "start", partial: createAssistantMessage("") }); + const checkAbort = () => { + if (abortSignal?.aborted) { + stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }); + } else { + setTimeout(checkAbort, 5); + } + }; + checkAbort(); + }); + return stream; + }, + }); + + const sessionManager = SessionManager.inMemory(); + const settingsManager = SettingsManager.create(tempDir, tempDir); + const authStorage = AuthStorage.create(join(tempDir, "auth.json")); + const modelRegistry = new ModelRegistry(authStorage, tempDir); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + + const extensionsResult = await createTestExtensionsResult([ + (pi) => { + (globalThis as typeof globalThis & { testExtensionApi?: unknown }).testExtensionApi = pi; + }, + (pi) => { + pi.on("input", async (event) => { + lastInputSource = event.source; + }); + }, + ]); + + session = new AgentSession({ + agent, + sessionManager, + settingsManager, + cwd: tempDir, + modelRegistry, + resourceLoader: createTestResourceLoader({ extensionsResult }), + }); + session.subscribe((event) => { + if (event.type === "queue_update") { + queueEvents.push({ steering: event.steering, followUp: event.followUp }); + } + }); + + const firstPrompt = session.prompt("First message"); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(session.isStreaming).toBe(true); + + const pi = ( + globalThis as typeof globalThis & { + testExtensionApi?: { + sendUserMessage: (content: string, options?: { deliverAs?: "steer" | "followUp" }) => void; + }; + } + ).testExtensionApi; + expect(pi).toBeDefined(); + + pi!.sendUserMessage("Steer from extension", { deliverAs: "steer" }); + await new Promise((resolve) => setTimeout(resolve, 25)); + + expect(session.pendingMessageCount).toBe(1); + expect(session.getSteeringMessages()).toContain("Steer from extension"); + expect(lastInputSource).toBe("interactive"); + expect(queueEvents.some((event) => event.steering.includes("Steer from extension"))).toBe(true); + + await session.abort(); + await firstPrompt.catch(() => {}); + + expect(sawSteeringMessage).toBe(true); + }); + it("should allow prompt() after previous completes", async () => { // Create session with a stream that completes immediately const model = getModel("anthropic", "claude-sonnet-4-5")!;