From 2abd1d9e199583f88ac84006c5ef9330bff37b67 Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Thu, 28 May 2026 14:15:52 +0200 Subject: [PATCH] fix: queue prompts during compaction --- .changeset/queue-prompts-during-compaction.md | 5 + src/server/sessions/piSessionService.test.ts | 83 +++++++++- src/server/sessions/piSessionService.ts | 156 +++++++++++++++--- 3 files changed, 219 insertions(+), 25 deletions(-) create mode 100644 .changeset/queue-prompts-during-compaction.md diff --git a/.changeset/queue-prompts-during-compaction.md b/.changeset/queue-prompts-during-compaction.md new file mode 100644 index 0000000..45ebbef --- /dev/null +++ b/.changeset/queue-prompts-during-compaction.md @@ -0,0 +1,5 @@ +--- +"@jmfederico/pi-web": patch +--- + +Queue prompts submitted during session compaction in pi-web and deliver them only after compaction finishes. diff --git a/src/server/sessions/piSessionService.test.ts b/src/server/sessions/piSessionService.test.ts index 74e3e69..eb23fc1 100644 --- a/src/server/sessions/piSessionService.test.ts +++ b/src/server/sessions/piSessionService.test.ts @@ -45,6 +45,7 @@ function sessionRecord(id: string, cwd = "/workspace") { function fakeRuntime(sessionId = "session-1", patch: Partial = {}) { const promptCalls: { text: string; options: unknown }[] = []; + const listeners: ((event: unknown) => void)[] = []; const calls = { abort: 0, clearQueue: 0, dispose: 0, prompt: promptCalls }; const session: TestSession = { sessionId, @@ -63,7 +64,13 @@ function fakeRuntime(sessionId = "session-1", patch: Partial = {}) extensionRunner: { getRegisteredCommands: () => [] }, promptTemplates: [], resourceLoader: { getSkills: () => ({ skills: [] }) }, - subscribe: () => () => undefined, + subscribe: (listener: (event: unknown) => void) => { + listeners.push(listener); + return () => { + const index = listeners.indexOf(listener); + if (index !== -1) listeners.splice(index, 1); + }; + }, getSessionStats: () => ({ sessionId, totalMessages: 0, userMessages: 0, assistantMessages: 0, toolCalls: 0, tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, cost: 0 }), getContextUsage: () => undefined, prompt: (text: string, options: unknown) => { @@ -101,7 +108,7 @@ function fakeRuntime(sessionId = "session-1", patch: Partial = {}) return Promise.resolve(); }, }; - return { runtime, session, calls }; + return { runtime, session, calls, emit: (event: unknown) => { for (const listener of [...listeners]) listener(event); } }; } function runtimeCreator(runtime: PiSessionRuntime): RuntimeCreator { @@ -414,6 +421,60 @@ describe("PiSessionService", () => { await service.dispose(); }); + it("holds prompts sent during compaction until compaction finishes", async () => { + const hub = new CapturingSessionEventHub(); + const fake = fakeRuntime("compacting-session", { isCompacting: true }); + let resolveFirstPrompt: (() => void) | undefined; + fake.session.prompt = (text: string, options?: { streamingBehavior?: "steer" | "followUp" }) => { + fake.calls.prompt.push({ text, options }); + if (options === undefined) { + fake.session.isStreaming = true; + return new Promise((resolve) => { resolveFirstPrompt = resolve; }); + } + return Promise.resolve(); + }; + const service = new PiSessionService(hub, { + createAgentRuntime: runtimeCreator(fake.runtime), + sessionManager: sessionGateway([sessionRecord("compacting-session")]), + heartbeatIntervalMs: 60_000, + }); + + await service.prompt("compacting-session", "Start task 1", "followUp"); + await service.prompt("compacting-session", "Then task 2", "followUp"); + + expect(fake.calls.prompt).toEqual([]); + expect(hub.sessionEvents.some(({ event }) => event.type === "message.append")).toBe(false); + await expect(service.status("compacting-session")).resolves.toMatchObject({ + pendingMessageCount: 2, + queuedMessages: [{ kind: "followUp", text: "Start task 1" }, { kind: "followUp", text: "Then task 2" }], + }); + + fake.session.isCompacting = false; + fake.emit({ type: "compaction_end" }); + await new Promise((resolve) => setTimeout(resolve, 5)); + + expect(fake.calls.prompt).toEqual([{ text: "Start task 1", options: undefined }]); + expect(hub.sessionEvents.some(({ event }) => event.type === "message.append" && JSON.stringify(event.message).includes("Start task 1"))).toBe(true); + await expect(service.status("compacting-session")).resolves.toMatchObject({ + pendingMessageCount: 1, + queuedMessages: [{ kind: "followUp", text: "Then task 2" }], + }); + + fake.emit({ type: "agent_start" }); + await new Promise((resolve) => setTimeout(resolve, 5)); + + expect(fake.calls.prompt).toEqual([ + { text: "Start task 1", options: undefined }, + { text: "Then task 2", options: { streamingBehavior: "followUp" } }, + ]); + await expect(service.status("compacting-session")).resolves.toMatchObject({ + pendingMessageCount: 0, + queuedMessages: [], + }); + resolveFirstPrompt?.(); + await service.dispose(); + }); + it("clears queued messages when aborting active work", async () => { const fake = fakeRuntime("abort-session"); const service = new PiSessionService(new CapturingSessionEventHub(), { @@ -430,6 +491,24 @@ describe("PiSessionService", () => { await service.dispose(); }); + it("clears prompts queued during compaction when aborting active work", async () => { + const fake = fakeRuntime("abort-compaction-session", { isCompacting: true }); + const service = new PiSessionService(new CapturingSessionEventHub(), { + createAgentRuntime: runtimeCreator(fake.runtime), + sessionManager: sessionGateway([sessionRecord("abort-compaction-session")]), + heartbeatIntervalMs: 60_000, + }); + + await service.prompt("abort-compaction-session", "Do not deliver after abort", "followUp"); + await expect(service.status("abort-compaction-session")).resolves.toMatchObject({ pendingMessageCount: 1 }); + await service.abort("abort-compaction-session"); + + expect(fake.calls.clearQueue).toBe(1); + expect(fake.calls.prompt).toEqual([]); + await expect(service.status("abort-compaction-session")).resolves.toMatchObject({ pendingMessageCount: 0, queuedMessages: [] }); + await service.dispose(); + }); + it("refreshes auth state and dedupes warnings when logout removes the current model's credentials", async () => { const hub = new CapturingSessionEventHub(); const authStorage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "sk-test" } }); diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index 93e9d64..cdb29ed 100644 --- a/src/server/sessions/piSessionService.ts +++ b/src/server/sessions/piSessionService.ts @@ -34,6 +34,13 @@ function authLossWarningKey(sessionId: string, provider: string, modelId: string return `${sessionId}:${provider}/${modelId}`; } +type QueuedPromptKind = "steer" | "followUp"; + +interface QueuedPrompt { + kind: QueuedPromptKind; + text: string; +} + type SessionArchiveRepository = Pick; interface PiSessionListEntry { id: string; @@ -180,6 +187,8 @@ export class PiSessionService { private readonly activities = new Map(); private readonly heartbeat: NodeJS.Timeout; private readonly commandService: SessionCommandService; + private readonly compactionPromptQueues = new Map(); + private readonly compactionDrainTimers = new Map(); private readonly authLossWarnings = new Set(); private readonly archiveStore: SessionArchiveRepository; private readonly agentDir: string; @@ -222,9 +231,11 @@ export class PiSessionService { async dispose(): Promise { clearInterval(this.heartbeat); + this.clearCompactionDrainTimers(); const activeSessions = Array.from(new Set(this.active.values())); this.active.clear(); this.activities.clear(); + this.compactionPromptQueues.clear(); this.authLossWarnings.clear(); await Promise.all(activeSessions.map(async (active) => { active.unsubscribe(); @@ -355,18 +366,36 @@ export class PiSessionService { this.maybeGenerateSessionName(session, text); const isQueued = session.isStreaming || session.isCompacting; const behavior = isQueued ? streamingBehavior ?? "followUp" : undefined; - if (isQueued && hasQueuedMessageText(session, text)) { + if (isQueued && this.hasQueuedMessageText(session, text)) { this.publishActivity(session, "duplicate queued message ignored", "active"); this.publishStatus(session); return; } - this.publishActivity(session, session.isCompacting ? "message queued during compaction" : behavior === "steer" ? "steering queued" : behavior === "followUp" ? "message queued" : "prompt accepted", "active"); - if (!isQueued) this.events.publish(sessionId, { type: "message.append", message: userTextMessage(text) }); - void session.prompt(text, behavior === undefined ? undefined : { streamingBehavior: behavior }).catch((error: unknown) => { + if (session.isCompacting) { + this.enqueuePromptDuringCompaction(session, text, behavior ?? "followUp"); + return; + } + void this.submitPrompt(session, text, behavior); + } + + private submitPrompt(session: PiAgentSession, text: string, behavior: QueuedPromptKind | undefined): Promise { + this.publishActivity(session, behavior === "steer" ? "steering queued" : behavior === "followUp" ? "message queued" : "prompt accepted", "active"); + if (behavior === undefined) this.events.publish(session.sessionId, { type: "message.append", message: userTextMessage(text) }); + const promptPromise = session.prompt(text, behavior === undefined ? undefined : { streamingBehavior: behavior }).catch((error: unknown) => { const message = error instanceof Error ? error.message : String(error); this.publishActivity(session, "error", "error", message); - this.events.publish(sessionId, { type: "session.error", message }); + this.events.publish(session.sessionId, { type: "session.error", message }); }); + void promptPromise; + return promptPromise; + } + + private enqueuePromptDuringCompaction(session: PiAgentSession, text: string, kind: QueuedPromptKind): void { + const queue = this.compactionPromptQueues.get(session.sessionId) ?? []; + queue.push({ kind, text }); + this.compactionPromptQueues.set(session.sessionId, queue); + this.publishActivity(session, "message queued during compaction", "active"); + this.publishStatus(session); } async shell(sessionId: string, text: string): Promise { @@ -416,7 +445,7 @@ export class PiSessionService { async archive(sessionId: string): Promise { const session = await this.getOrOpen(sessionId); - if (sessionHasActiveWork(session)) throw new Error("Stop current session activity before archiving"); + if (this.hasActiveWork(session)) throw new Error("Stop current session activity before archiving"); const archiveInput = await this.archiveInputForSession(session); await this.closeActive(session.sessionId); await this.archiveStore.archive(archiveInput); @@ -427,7 +456,7 @@ export class PiSessionService { const catalog = await this.workspaceArchiveCandidates(session.sessionManager.getCwd()); const root = findArchiveCandidateByIdOrPrefix(catalog, session.sessionId) ?? archiveCandidateFromActiveSession(session, false); const plan = planSessionArchiveTree(root, catalog); - const busy = plan.targets.map((target) => target.activeSession).find((target) => target !== undefined && sessionHasActiveWork(target)); + const busy = plan.targets.map((target) => target.activeSession).find((target) => target !== undefined && this.hasActiveWork(target)); if (busy !== undefined) throw new Error(`Stop current session activity before archiving ${sessionDisplayName(busy)}`); const archiveInputs = plan.unarchivedTargets.map((target) => archiveInputFromCandidate(target)); @@ -457,6 +486,7 @@ export class PiSessionService { async abort(sessionId: string): Promise { const active = this.active.get(sessionId); if (!active) return; + this.clearCompactionPromptQueue(sessionId); clearSessionQueue(active.runtime.session); await active.runtime.session.abort(); this.publishActivity(active.runtime.session, "stopped", "idle"); @@ -551,6 +581,7 @@ export class PiSessionService { this.activities.delete(sessionId); this.workspaceActivity?.removeSession(sessionId, active.runtime.session.sessionManager.getCwd()); this.clearAuthLossWarningsForSession(sessionId); + this.clearCompactionPromptQueue(sessionId); clearSessionQueue(active.runtime.session); active.unsubscribe(); try { @@ -595,18 +626,84 @@ export class PiSessionService { private bindRuntime(active: ActiveSession): void { active.unsubscribe(); - for (const [sessionId, candidate] of this.active.entries()) { - if (candidate === active) this.active.delete(sessionId); - } const { session } = active.runtime; + for (const [sessionId, candidate] of this.active.entries()) { + if (candidate === active) { + this.active.delete(sessionId); + if (sessionId !== session.sessionId) this.clearCompactionPromptQueue(sessionId); + } + } active.unsubscribe = session.subscribe((event) => { this.events.publish(session.sessionId, toClientEvent(event)); this.publishActivityForEvent(session, event); + const eventType = getString(event, "type"); + if (eventType === "compaction_end") this.scheduleCompactionQueueDrain(session.sessionId); + if (eventType === "agent_start" || eventType === "agent_end") this.scheduleCompactionQueueDrain(session.sessionId); this.publishStatus(session); }); this.active.set(session.sessionId, active); } + private scheduleCompactionQueueDrain(sessionId: string, delayMs = 0): void { + if (!this.compactionPromptQueues.has(sessionId) || this.compactionDrainTimers.has(sessionId)) return; + const timer = setTimeout(() => { + this.compactionDrainTimers.delete(sessionId); + this.drainCompactionPromptQueue(sessionId); + }, delayMs); + this.compactionDrainTimers.set(sessionId, timer); + } + + private drainCompactionPromptQueue(sessionId: string): void { + const active = this.active.get(sessionId); + if (active === undefined) return; + const { session } = active.runtime; + if (session.isCompacting) { + this.scheduleCompactionQueueDrain(sessionId, 100); + return; + } + + if (session.isStreaming) { + const queued = this.takeCompactionPromptQueue(sessionId); + if (queued.length === 0) return; + this.publishStatus(session); + for (const prompt of queued) void this.submitPrompt(session, prompt.text, prompt.kind); + return; + } + + const prompt = this.shiftCompactionPrompt(sessionId); + if (prompt === undefined) return; + this.publishStatus(session); + const submitted = this.submitPrompt(session, prompt.text, undefined); + void submitted.finally(() => { this.scheduleCompactionQueueDrain(sessionId); }); + } + + private takeCompactionPromptQueue(sessionId: string): QueuedPrompt[] { + const queued = this.compactionPromptQueues.get(sessionId) ?? []; + this.compactionPromptQueues.delete(sessionId); + return queued; + } + + private shiftCompactionPrompt(sessionId: string): QueuedPrompt | undefined { + const queue = this.compactionPromptQueues.get(sessionId); + const prompt = queue?.shift(); + if (queue === undefined || queue.length === 0) this.compactionPromptQueues.delete(sessionId); + return prompt; + } + + private clearCompactionPromptQueue(sessionId: string): void { + this.compactionPromptQueues.delete(sessionId); + const timer = this.compactionDrainTimers.get(sessionId); + if (timer !== undefined) { + clearTimeout(timer); + this.compactionDrainTimers.delete(sessionId); + } + } + + private clearCompactionDrainTimers(): void { + for (const timer of this.compactionDrainTimers.values()) clearTimeout(timer); + this.compactionDrainTimers.clear(); + } + private maybeGenerateSessionName(session: PiAgentSession, firstMessage: string): void { if (session.sessionName !== undefined || session.messages.length !== 0 || session.isStreaming || session.isCompacting) return; const model = session.model; @@ -674,7 +771,7 @@ export class PiSessionService { for (const active of this.active.values()) { const { session } = active.runtime; const activity = this.activities.get(session.sessionId); - if (!sessionHasActiveWork(session)) { + if (!this.hasActiveWork(session)) { if (activity?.phase === "active") this.publishStatus(session); continue; } @@ -688,10 +785,14 @@ export class PiSessionService { if (session.isCompacting) return "compacting"; if (session.isBashRunning) return "running bash"; if (session.isStreaming) return "agent running"; - if (session.pendingMessageCount) return "queued"; + if (this.pendingMessageCount(session) > 0) return "queued"; return "active"; } + private hasActiveWork(session: PiAgentSession): boolean { + return sessionHasActiveWork(session, this.compactionQueuedMessages(session.sessionId).length); + } + private publishActivityForEvent(session: PiAgentSession, event: unknown): void { const eventType = getString(event, "type"); if (eventType === undefined) return; @@ -716,7 +817,7 @@ export class PiSessionService { } if (eventType === "bash_execution_start") { this.publishActivity(session, "running bash", "active"); return; } if (eventType === "bash_execution_end") { this.publishActivity(session, "bash complete", "idle"); return; } - if (sessionHasActiveWork(session)) this.publishActivity(session, eventType.replaceAll("_", " "), "active"); + if (this.hasActiveWork(session)) this.publishActivity(session, eventType.replaceAll("_", " "), "active"); } private publishActivity(session: PiAgentSession, label: string, phase: "active" | "idle" | "error", detail?: string): void { @@ -739,7 +840,7 @@ export class PiSessionService { private clearStaleActiveActivity(session: PiAgentSession): void { const current = this.activities.get(session.sessionId); - if (current?.phase !== "active" || sessionHasActiveWork(session)) return; + if (current?.phase !== "active" || this.hasActiveWork(session)) return; const at = new Date().toISOString(); const stored = { phase: "idle" as const, label: "idle", at }; this.activities.set(session.sessionId, stored); @@ -759,14 +860,26 @@ export class PiSessionService { isStreaming: session.isStreaming, isCompacting: session.isCompacting, isBashRunning: session.isBashRunning, - pendingMessageCount: session.pendingMessageCount, - queuedMessages: queuedMessagesFromSession(session), + pendingMessageCount: this.pendingMessageCount(session), + queuedMessages: queuedMessagesFromSession(session, this.compactionQueuedMessages(session.sessionId)), messageCount: session.messages.length, tokens: stats.tokens, cost: stats.cost, ...(contextUsage === undefined ? {} : { contextUsage }), }; } + + private pendingMessageCount(session: PiAgentSession): number { + return session.pendingMessageCount + this.compactionQueuedMessages(session.sessionId).length; + } + + private compactionQueuedMessages(sessionId: string): readonly QueuedPrompt[] { + return this.compactionPromptQueues.get(sessionId) ?? []; + } + + private hasQueuedMessageText(session: PiAgentSession, text: string): boolean { + return queuedMessagesFromSession(session, this.compactionQueuedMessages(session.sessionId)).some((message) => message.text === text); + } } function modelToClientModel(model: PiAgentSession["model"]): ClientSessionModel { @@ -872,8 +985,8 @@ function archiveInputFromCandidate(candidate: WorkspaceArchiveCandidate): Archiv throw new Error(`Session is not available for archiving: ${candidate.id}`); } -function sessionHasActiveWork(session: PiAgentSession): boolean { - return session.isStreaming || session.isCompacting || session.isBashRunning || session.pendingMessageCount > 0; +function sessionHasActiveWork(session: PiAgentSession, extraQueuedMessageCount = 0): boolean { + return session.isStreaming || session.isCompacting || session.isBashRunning || session.pendingMessageCount + extraQueuedMessageCount > 0; } function sessionDisplayName(session: PiAgentSession): string { @@ -938,14 +1051,11 @@ function clearSessionQueue(session: PiAgentSession): void { session.clearQueue(); } -function hasQueuedMessageText(session: PiAgentSession, text: string): boolean { - return queuedMessagesFromSession(session).some((message) => message.text === text); -} - -function queuedMessagesFromSession(session: PiAgentSession): { kind: "steer" | "followUp"; text: string }[] { +function queuedMessagesFromSession(session: PiAgentSession, extraQueuedMessages: readonly QueuedPrompt[] = []): { kind: "steer" | "followUp"; text: string }[] { return [ ...session.getSteeringMessages().map((text) => ({ kind: "steer" as const, text })), ...session.getFollowUpMessages().map((text) => ({ kind: "followUp" as const, text })), + ...extraQueuedMessages, ]; }