diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index f8f8561b..0ed48705 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -1890,13 +1890,15 @@ export class CodexAcpServer { return pendingTurnStart; }; const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); + let eventHandler: CodexEventHandler | null = null; try { - const eventHandler = new CodexEventHandler( + const promptEventHandler = new CodexEventHandler( this.connection, sessionState, clientSupportsPlanUpdates(this.clientCapabilities), ); + eventHandler = promptEventHandler; const approvalHandler = new CodexApprovalHandler(this.connection, sessionState, activePrompt.signal); const elicitationHandler = new CodexElicitationHandler( this.connection, @@ -1907,7 +1909,7 @@ export class CodexAcpServer { await this.codexAcpClient.subscribeToSessionEvents(params.sessionId, async (event) => { await elicitationHandler.handleNotification(event); - return eventHandler.handleNotification(event); + return promptEventHandler.handleNotification(event); }, approvalHandler, elicitationHandler); @@ -2041,6 +2043,7 @@ export class CodexAcpServer { await this.codexAcpClient.waitForSessionNotifications(params.sessionId); if (turnCompleted.turn.status === "interrupted") { + await eventHandler.flushPendingPlanUpdates(); await this.notifyConversationInterrupted(params.sessionId); return this.cancelledPromptResponse(sessionState); } @@ -2051,6 +2054,7 @@ export class CodexAcpServer { throw error; } + await eventHandler.flushPendingPlanUpdates(); const completedPlan = eventHandler.takeCompletedPlan(); if ( completedPlan !== null @@ -2116,6 +2120,7 @@ export class CodexAcpServer { await this.codexAcpClient.waitForSessionNotifications(params.sessionId); if (turnCompleted.turn.status === "interrupted") { + await eventHandler.flushPendingPlanUpdates(); await this.notifyConversationInterrupted(params.sessionId); return this.cancelledPromptResponse(sessionState); } @@ -2142,6 +2147,7 @@ export class CodexAcpServer { throw err; } finally { logger.log("Prompt completed", {sessionId: params.sessionId}); + await eventHandler?.dispose(); disposePromptRequestCancellation(); sessionState.currentTurnId = null; const registeredPendingTurnStart = this.pendingTurnStarts.get(params.sessionId); diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index 2a06ddff..895c1c59 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -65,6 +65,7 @@ import { createAgentTextThoughtChunk, } from "./ContentChunks"; import {sameThreadGoalSnapshot, toThreadGoalSnapshot} from "./ThreadGoalSnapshot"; +import {logger} from "./Logger"; export { stripShellPrefix }; @@ -75,7 +76,8 @@ export type CompletedPlan = { export class CodexEventHandler { - private readonly connection: AcpClientConnection; + private static readonly PLAN_UPDATE_INTERVAL_MS = 150; + private readonly sessionState: SessionState; private readonly supportsPlanUpdates: boolean; private failure: RequestError | null = null; @@ -85,6 +87,12 @@ export class CodexEventHandler { private readonly activeImageGenerationItems = new Set(); private readonly emittedImageViewItems = new Set(); private readonly planDeltaTextByItemId = new Map(); + private readonly pendingPlanItemIds = new Set(); + private readonly lastEmittedPlanTextByItemId = new Map(); + private readonly session: ACPSessionConnection; + private planUpdateTimer: ReturnType | null = null; + private planUpdateChain: Promise = Promise.resolve(); + private disposed = false; private readonly seenReasoningDeltaItemIds = new Set(); private readonly terminalCommandIds = new Set(); private readonly terminalCommandOutputIds = new Set(); @@ -96,9 +104,9 @@ export class CodexEventHandler { sessionState: SessionState, supportsPlanUpdates = false, ) { - this.connection = connection; this.sessionState = sessionState; this.supportsPlanUpdates = supportsPlanUpdates; + this.session = new ACPSessionConnection(connection, sessionState.sessionId); } getFailure(): RequestError | null { @@ -112,13 +120,37 @@ export class CodexEventHandler { } async handleNotification(notification: ServerNotification) { - const session = new ACPSessionConnection(this.connection, this.sessionState.sessionId); const updateEvent = await this.createUpdateEvent(notification); if (updateEvent) { - await session.update(updateEvent); + await this.session.update(updateEvent); } } + async flushPendingPlanUpdates(): Promise { + this.cancelPlanUpdateTimer(); + do { + const itemIds = [...this.pendingPlanItemIds]; + this.pendingPlanItemIds.clear(); + await Promise.all(itemIds.map(itemId => { + const text = this.planDeltaTextByItemId.get(itemId) ?? ""; + return text.length > 0 + ? this.enqueuePlanSnapshot(itemId, text) + : Promise.resolve(); + })); + await this.planUpdateChain; + } while (this.pendingPlanItemIds.size > 0); + } + + async dispose(): Promise { + if (this.disposed) return; + await this.flushPendingPlanUpdates(); + this.disposed = true; + this.cancelPlanUpdateTimer(); + this.pendingPlanItemIds.clear(); + this.planDeltaTextByItemId.clear(); + this.lastEmittedPlanTextByItemId.clear(); + } + private async createUpdateEvent(notification: ServerNotification): Promise { /* TODO split UpdateSessionEvent to improve completion @@ -144,6 +176,8 @@ export class CodexEventHandler { this.sessionState.currentTurnId = notification.params.turn.id; return null; case "turn/completed": + await this.flushPendingPlanUpdates(); + this.clearPlanTurnState(); this.sessionState.currentTurnId = null; return null; case "thread/tokenUsage/updated": @@ -314,16 +348,18 @@ export class CodexEventHandler { return this.createAgentThoughtEvent(event.delta, event.itemId); } - private createPlanDeltaEvent(event: PlanDeltaNotification): UpdateSessionEvent | null { + private createPlanDeltaEvent(event: PlanDeltaNotification): null { if (event.delta.length === 0) { return null; } const text = this.planDeltaTextByItemId.get(event.itemId) ?? ""; const updatedText = text + event.delta; this.planDeltaTextByItemId.set(event.itemId, updatedText); - return this.supportsPlanUpdates - ? this.createPlanUpdateEvent(updatedText, event.itemId) - : null; + if (this.supportsPlanUpdates) { + this.pendingPlanItemIds.add(event.itemId); + this.schedulePlanUpdate(); + } + return null; } private createReasoningSectionBreakEvent(event: ReasoningSummaryPartAddedNotification): UpdateSessionEvent { @@ -424,8 +460,7 @@ export class CodexEventHandler { return null; case "plan": { const deltaText = this.planDeltaTextByItemId.get(event.item.id) ?? ""; - this.planDeltaTextByItemId.delete(event.item.id); - return this.createCompletedPlanEvent(event.item, deltaText); + return await this.createCompletedPlanEvent(event.item, deltaText); } case "exitedReviewMode": return this.createExitedReviewModeEvent(event.item); @@ -460,18 +495,59 @@ export class CodexEventHandler { return this.createAgentThoughtEvent(text, item.id); } - private createCompletedPlanEvent( + private async createCompletedPlanEvent( item: ThreadItem & { type: "plan" }, deltaText: string, - ): UpdateSessionEvent | null { + ): Promise { const text = item.text.length > 0 ? item.text : deltaText; + this.pendingPlanItemIds.delete(item.id); + if (this.pendingPlanItemIds.size === 0) { + this.cancelPlanUpdateTimer(); + } + this.planDeltaTextByItemId.delete(item.id); if (text.length === 0) { return null; } this.completedPlan = {itemId: item.id, text}; - return this.supportsPlanUpdates - ? this.createPlanUpdateEvent(text, item.id) - : this.createPlanTextEvent(text, item.id); + if (this.supportsPlanUpdates) { + await this.enqueuePlanSnapshot(item.id, text); + return null; + } + return this.createPlanTextEvent(text, item.id); + } + + private schedulePlanUpdate(): void { + if (this.disposed || this.planUpdateTimer !== null) return; + this.planUpdateTimer = setTimeout(() => { + this.planUpdateTimer = null; + void this.flushPendingPlanUpdates().catch(error => { + logger.error("Failed to flush throttled plan updates", error); + }); + }, CodexEventHandler.PLAN_UPDATE_INTERVAL_MS); + } + + private cancelPlanUpdateTimer(): void { + if (this.planUpdateTimer === null) return; + clearTimeout(this.planUpdateTimer); + this.planUpdateTimer = null; + } + + private enqueuePlanSnapshot(itemId: string, text: string): Promise { + const send = async () => { + if (this.lastEmittedPlanTextByItemId.get(itemId) === text) return; + await this.session.update(this.createPlanUpdateEvent(text, itemId)); + this.lastEmittedPlanTextByItemId.set(itemId, text); + }; + const result = this.planUpdateChain.then(send); + this.planUpdateChain = result.catch(() => {}); + return result; + } + + private clearPlanTurnState(): void { + this.cancelPlanUpdateTimer(); + this.pendingPlanItemIds.clear(); + this.planDeltaTextByItemId.clear(); + this.lastEmittedPlanTextByItemId.clear(); } private createPlanUpdateEvent(text: string, planId: string): UpdateSessionEvent { diff --git a/src/__tests__/CodexACPAgent/plan-events.test.ts b/src/__tests__/CodexACPAgent/plan-events.test.ts index 67ca0a92..aa22e89a 100644 --- a/src/__tests__/CodexACPAgent/plan-events.test.ts +++ b/src/__tests__/CodexACPAgent/plan-events.test.ts @@ -1,7 +1,9 @@ -import {beforeEach, describe, expect, it, vi} from "vitest"; +import {afterEach, beforeEach, describe, expect, it, vi} from "vitest"; import type {ServerNotification} from "../../app-server"; import {AgentMode} from "../../AgentMode"; import type {SessionState} from "../../CodexAcpServer"; +import {CodexEventHandler} from "../../CodexEventHandler"; +import type {AcpClientConnection} from "../../ACPSessionConnection"; import { createCodexMockTestFixture, createTestSessionState, @@ -18,6 +20,10 @@ describe("CodexEventHandler - plan events", () => { vi.clearAllMocks(); }); + afterEach(() => { + vi.useRealTimers(); + }); + const sessionState: SessionState = createTestSessionState({ sessionId, currentModelId: "model-id[effort]", @@ -173,4 +179,165 @@ describe("CodexEventHandler - plan events", () => { "data/plan-checklist-update.json", ); }); + + describe("plan update coalescing", () => { + function createHandler( + notify = vi.fn(async (_method: unknown, _params: unknown) => {}), + ) { + const connection = { + notify, + request: vi.fn(), + } as unknown as AcpClientConnection; + const handler = new CodexEventHandler(connection, sessionState, true); + const planUpdates = () => notify.mock.calls + .map(call => call[1] as {update?: {sessionUpdate?: string, plan?: {planId: string, content: string}}}) + .filter(params => params.update?.sessionUpdate === "plan_update") + .map(params => params.update!.plan!); + return {handler, planUpdates}; + } + + function planDelta(itemId: string, delta: string): ServerNotification { + return { + method: "item/plan/delta", + params: {threadId: sessionId, turnId: "turn-1", itemId, delta}, + }; + } + + function completedPlan(itemId: string, text: string): ServerNotification { + return { + method: "item/completed", + params: { + threadId: sessionId, + turnId: "turn-1", + completedAtMs: 0, + item: {type: "plan", id: itemId, text}, + }, + }; + } + + function completedTurn(status: "completed" | "interrupted"): ServerNotification { + return { + method: "turn/completed", + params: { + threadId: sessionId, + turn: { + id: "turn-1", + items: [], + itemsView: "notLoaded", + status, + error: null, + startedAt: null, + completedAt: null, + durationMs: null, + }, + }, + }; + } + + it("coalesces many small deltas and emits the complete final snapshot", async () => { + vi.useFakeTimers(); + const {handler, planUpdates} = createHandler(); + let fullText = ""; + + for (let index = 0; index < 200; index += 1) { + const delta = `${index % 10}`; + fullText += delta; + await handler.handleNotification(planDelta("plan-many", delta)); + if (index % 10 === 9) { + await vi.advanceTimersByTimeAsync(25); + } + } + await handler.handleNotification(completedPlan("plan-many", fullText)); + + expect(planUpdates().length).toBeLessThan(20); + expect(planUpdates().length).toBeGreaterThan(1); + expect(planUpdates().at(-1)).toEqual({type: "markdown", planId: "plan-many", content: fullText}); + await handler.dispose(); + }); + + it.each(["completed", "interrupted"] as const)("flushes a pending snapshot when the turn is %s", async status => { + vi.useFakeTimers(); + const {handler, planUpdates} = createHandler(); + await handler.handleNotification(planDelta("plan-boundary", "full pending plan")); + + await handler.handleNotification(completedTurn(status)); + + expect(planUpdates()).toEqual([{type: "markdown", planId: "plan-boundary", content: "full pending plan"}]); + await vi.advanceTimersByTimeAsync(1_000); + expect(planUpdates()).toHaveLength(1); + await handler.dispose(); + }); + + it("does not duplicate an identical completed snapshot", async () => { + vi.useFakeTimers(); + const {handler, planUpdates} = createHandler(); + await handler.handleNotification(planDelta("plan-same", "same text")); + await vi.advanceTimersByTimeAsync(150); + + await handler.handleNotification(completedPlan("plan-same", "same text")); + + expect(planUpdates()).toEqual([{type: "markdown", planId: "plan-same", content: "same text"}]); + await handler.dispose(); + }); + + it("serializes an in-flight throttled snapshot before the completed snapshot", async () => { + vi.useFakeTimers(); + let releaseFirstSend!: () => void; + let markFirstSendStarted!: () => void; + const firstSendStarted = new Promise(resolve => { + markFirstSendStarted = resolve; + }); + const firstSendReleased = new Promise(resolve => { + releaseFirstSend = resolve; + }); + let firstSend = true; + const notify = vi.fn(async (_method: unknown, _params: unknown) => { + if (!firstSend) return; + firstSend = false; + markFirstSendStarted(); + await firstSendReleased; + }); + const {handler, planUpdates} = createHandler(notify); + await handler.handleNotification(planDelta("plan-race", "partial")); + + await vi.advanceTimersByTimeAsync(150); + await firstSendStarted; + const completion = handler.handleNotification(completedPlan("plan-race", "partial and final")); + releaseFirstSend(); + await completion; + + expect(planUpdates().map(plan => plan.content)).toEqual(["partial", "partial and final"]); + await handler.dispose(); + }); + + it("flushes and cancels pending work when disposed", async () => { + vi.useFakeTimers(); + const {handler, planUpdates} = createHandler(); + await handler.handleNotification(planDelta("plan-dispose", "last session snapshot")); + + await handler.dispose(); + await vi.advanceTimersByTimeAsync(1_000); + + expect(planUpdates()).toEqual([ + {type: "markdown", planId: "plan-dispose", content: "last session snapshot"}, + ]); + }); + + it("keeps independently streamed plans separate", async () => { + vi.useFakeTimers(); + const {handler, planUpdates} = createHandler(); + await handler.handleNotification(planDelta("plan-a", "A1")); + await handler.handleNotification(planDelta("plan-b", "B1")); + await handler.handleNotification(planDelta("plan-a", "A2")); + await handler.handleNotification(planDelta("plan-b", "B2")); + + await handler.handleNotification(completedTurn("completed")); + + expect(planUpdates()).toEqual([ + {type: "markdown", planId: "plan-a", content: "A1A2"}, + {type: "markdown", planId: "plan-b", content: "B1B2"}, + ]); + await handler.dispose(); + }); + }); }); diff --git a/src/__tests__/CodexACPAgent/plan-review-events.test.ts b/src/__tests__/CodexACPAgent/plan-review-events.test.ts index 088869cb..c31d2d85 100644 --- a/src/__tests__/CodexACPAgent/plan-review-events.test.ts +++ b/src/__tests__/CodexACPAgent/plan-review-events.test.ts @@ -170,6 +170,15 @@ describe("CodexACPAgent - plan review", () => { }, }], }); + const finalPlanUpdateIndex = events.reduce((lastIndex, event, index) => + event.method === "sessionUpdate" + && (event.args[0] as {update?: {sessionUpdate?: string}}).update?.sessionUpdate === "plan_update" + ? index + : lastIndex, + -1); + const permissionIndex = events.findIndex(event => event.method === "requestPermission"); + expect(finalPlanUpdateIndex).toBeGreaterThanOrEqual(0); + expect(permissionIndex).toBeGreaterThan(finalPlanUpdateIndex); expect(sessionState.collaborationMode).toBe("default"); implementationTurn.resolve({