diff --git a/packages/pi-agent-core/src/agent.ts b/packages/pi-agent-core/src/agent.ts index 14a0a33ac..112573650 100644 --- a/packages/pi-agent-core/src/agent.ts +++ b/packages/pi-agent-core/src/agent.ts @@ -103,6 +103,15 @@ export interface AgentOptions { maxRetryDelayMs?: number; } +/** + * Internal wrapper that tracks message origin for origin-aware queue clearing. + * "user" = typed by human in TUI; "system" = generated by extensions/background jobs. + */ +interface QueueEntry { + message: AgentMessage; + origin: "user" | "system"; +} + export class Agent { private _state: AgentState = { systemPrompt: "", @@ -120,8 +129,8 @@ export class Agent { private abortController?: AbortController; private convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; private transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; - private steeringQueue: AgentMessage[] = []; - private followUpQueue: AgentMessage[] = []; + private steeringQueue: QueueEntry[] = []; + private followUpQueue: QueueEntry[] = []; private steeringMode: "all" | "one-at-a-time"; private followUpMode: "all" | "one-at-a-time"; public streamFn: StreamFn; @@ -279,16 +288,16 @@ export class Agent { * Queue a steering message to interrupt the agent mid-run. * Delivered after current tool execution, skips remaining tools. */ - steer(m: AgentMessage) { - this.steeringQueue.push(m); + steer(m: AgentMessage, origin: "user" | "system" = "system") { + this.steeringQueue.push({ message: m, origin }); } /** * Queue a follow-up message to be processed after the agent finishes. * Delivered only when agent has no more tool calls or steering messages. */ - followUp(m: AgentMessage) { - this.followUpQueue.push(m); + followUp(m: AgentMessage, origin: "user" | "system" = "system") { + this.followUpQueue.push({ message: m, origin }); } clearSteeringQueue() { @@ -304,6 +313,18 @@ export class Agent { this.followUpQueue = []; } + /** + * Drain user-origin messages from queues, leaving system messages in place. + * Used during abort to preserve messages the user explicitly typed. + */ + drainUserMessages(): { steering: AgentMessage[]; followUp: AgentMessage[] } { + const userSteering = this.steeringQueue.filter((e) => e.origin === "user").map((e) => e.message); + const userFollowUp = this.followUpQueue.filter((e) => e.origin === "user").map((e) => e.message); + this.steeringQueue = this.steeringQueue.filter((e) => e.origin !== "user"); + this.followUpQueue = this.followUpQueue.filter((e) => e.origin !== "user"); + return { steering: userSteering, followUp: userFollowUp }; + } + hasQueuedMessages(): boolean { return this.steeringQueue.length > 0 || this.followUpQueue.length > 0; } @@ -313,12 +334,12 @@ export class Agent { if (this.steeringQueue.length > 0) { const first = this.steeringQueue[0]; this.steeringQueue = this.steeringQueue.slice(1); - return [first]; + return [first.message]; } return []; } - const steering = this.steeringQueue.slice(); + const steering = this.steeringQueue.map((e) => e.message); this.steeringQueue = []; return steering; } @@ -328,12 +349,12 @@ export class Agent { if (this.followUpQueue.length > 0) { const first = this.followUpQueue[0]; this.followUpQueue = this.followUpQueue.slice(1); - return [first]; + return [first.message]; } return []; } - const followUp = this.followUpQueue.slice(); + const followUp = this.followUpQueue.map((e) => e.message); this.followUpQueue = []; return followUp; } diff --git a/packages/pi-coding-agent/src/core/agent-session.ts b/packages/pi-coding-agent/src/core/agent-session.ts index d09946b39..acd234702 100644 --- a/packages/pi-coding-agent/src/core/agent-session.ts +++ b/packages/pi-coding-agent/src/core/agent-session.ts @@ -1171,11 +1171,14 @@ export class AgentSession { if (images) { content.push(...images); } - this.agent.steer({ - role: "user", - content, - timestamp: Date.now(), - }); + this.agent.steer( + { + role: "user", + content, + timestamp: Date.now(), + }, + "user", + ); } /** @@ -1187,11 +1190,14 @@ export class AgentSession { if (images) { content.push(...images); } - this.agent.followUp({ - role: "user", - content, - timestamp: Date.now(), - }); + this.agent.followUp( + { + role: "user", + content, + timestamp: Date.now(), + }, + "user", + ); } /** @@ -1304,10 +1310,28 @@ export class AgentSession { * @returns Object with steering and followUp arrays */ clearQueue(): { steering: string[]; followUp: string[] } { - const steering = [...this._steeringMessages]; - const followUp = [...this._followUpMessages]; + // Drain user-origin messages from agent queues before clearing. + // This preserves messages the user explicitly typed during streaming, + // while system-generated messages (extension notifications, etc.) are discarded. + const userMessages = this.agent.drainUserMessages(); + + // Extract text content from preserved user messages + const extractText = (m: AgentMessage): string => { + if (!("content" in m) || !Array.isArray(m.content)) return ""; + const textPart = m.content.find((c: { type: string }) => c.type === "text"); + return textPart && "text" in textPart ? (textPart as { text: string }).text : ""; + }; + const preservedSteering = userMessages.steering.map(extractText).filter((t) => t.length > 0); + const preservedFollowUp = userMessages.followUp.map(extractText).filter((t) => t.length > 0); + + // Session-level string arrays track what was queued for display purposes. + // Return the full set (session-tracked + any agent-only user messages). + const steering = [...this._steeringMessages, ...preservedSteering]; + const followUp = [...this._followUpMessages, ...preservedFollowUp]; this._steeringMessages = []; this._followUpMessages = []; + + // Clear remaining system messages from agent queues this.agent.clearAllQueues(); return { steering, followUp }; }