From ece7d0945c427ab4a01d87f20b4c1baf4649a63d Mon Sep 17 00:00:00 2001 From: Vincent Koc <25068+vincentkoc@users.noreply.github.com> Date: Fri, 19 Jun 2026 14:17:31 +0800 Subject: [PATCH] fix(copilot): cancel deferred compaction on abort --- extensions/copilot/harness.test.ts | 8 ++-- extensions/copilot/harness.ts | 54 +++++++++++++------------- extensions/copilot/src/attempt.test.ts | 17 +++++--- extensions/copilot/src/attempt.ts | 37 +++++++++++------- 4 files changed, 66 insertions(+), 50 deletions(-) diff --git a/extensions/copilot/harness.test.ts b/extensions/copilot/harness.test.ts index 75cadea7d286..042fa38c69d6 100644 --- a/extensions/copilot/harness.test.ts +++ b/extensions/copilot/harness.test.ts @@ -537,7 +537,7 @@ describe("createCopilotAgentHarness", () => { pooledClient: { key: {} as any, client: {} as any }, sessionConfig: TEST_SESSION_CONFIG, }); - deps.onTimedOutCompaction?.({ + deps.onDeferredCompaction?.({ abort, cleanup: cleanup.promise, sdkSessionId: "sdk-sess-pending-cleanup", @@ -561,7 +561,7 @@ describe("createCopilotAgentHarness", () => { pooledClient: { key: {} as any, client: {} as any }, sessionConfig: TEST_SESSION_CONFIG, }); - deps.onTimedOutCompaction?.({ + deps.onDeferredCompaction?.({ abort, cleanup: cleanup.promise, sdkSessionId: "sdk-sess-reset-cleanup", @@ -636,7 +636,7 @@ describe("createCopilotAgentHarness", () => { pooledClient: { key: {} as any, client: {} as any }, sessionConfig: TEST_SESSION_CONFIG, }); - deps.onTimedOutCompaction?.({ + deps.onDeferredCompaction?.({ abort: () => undefined, cleanup: cleanup.promise, sdkSessionId: "sdk-sess-compacting", @@ -678,7 +678,7 @@ describe("createCopilotAgentHarness", () => { pooledClient: { key: {} as any, client: {} as any }, sessionConfig: TEST_SESSION_CONFIG, }); - deps.onTimedOutCompaction?.({ + deps.onDeferredCompaction?.({ abort: () => undefined, cleanup: cleanup.promise, sdkSessionId: "sdk-sess-cancelled", diff --git a/extensions/copilot/harness.ts b/extensions/copilot/harness.ts index eb03d444687d..a125d84cd625 100644 --- a/extensions/copilot/harness.ts +++ b/extensions/copilot/harness.ts @@ -94,7 +94,7 @@ type LegacyCopilotSessionBinding = { }; type CopilotAttemptSessionBinding = Pick; -type TimedOutCompactionCleanupOutcome = "aborted" | "completed" | "deadline"; +type DeferredCompactionCleanupOutcome = "aborted" | "completed" | "deadline"; type CopilotSessionBindingStore = Pick< PluginStateSyncKeyedStore, @@ -425,9 +425,9 @@ export function createCopilotAgentHarness( let disposed = false; let disposePromise: Promise | undefined; const inFlight = new Set>(); - const timedOutCompactionCleanups = new Map< + const deferredCompactionCleanups = new Map< string, - Map, () => void> + Map, () => void> >(); // Maps OpenClaw session id (from AgentHarnessAttemptParams.sessionId) to // the SDK session id + client that owns it. Populated by @@ -450,38 +450,38 @@ export function createCopilotAgentHarness( return poolPromise; } - function trackTimedOutCompactionCleanup(params: { + function trackDeferredCompactionCleanup(params: { abort: () => void; - cleanup: Promise; + cleanup: Promise; sessionId: string; }): void { const cleanups = - timedOutCompactionCleanups.get(params.sessionId) ?? - new Map, () => void>(); + deferredCompactionCleanups.get(params.sessionId) ?? + new Map, () => void>(); cleanups.set(params.cleanup, params.abort); - timedOutCompactionCleanups.set(params.sessionId, cleanups); + deferredCompactionCleanups.set(params.sessionId, cleanups); void params.cleanup.then( - () => removeTimedOutCompactionCleanup(params.sessionId, params.cleanup), - () => removeTimedOutCompactionCleanup(params.sessionId, params.cleanup), + () => removeDeferredCompactionCleanup(params.sessionId, params.cleanup), + () => removeDeferredCompactionCleanup(params.sessionId, params.cleanup), ); } - function removeTimedOutCompactionCleanup( + function removeDeferredCompactionCleanup( sessionId: string, - cleanup: Promise, + cleanup: Promise, ): void { - const cleanups = timedOutCompactionCleanups.get(sessionId); + const cleanups = deferredCompactionCleanups.get(sessionId); if (!cleanups) { return; } cleanups.delete(cleanup); if (cleanups.size === 0) { - timedOutCompactionCleanups.delete(sessionId); + deferredCompactionCleanups.delete(sessionId); } } - async function abortTimedOutCompactionCleanups(sessionId: string): Promise { - const cleanups = timedOutCompactionCleanups.get(sessionId); + async function abortDeferredCompactionCleanups(sessionId: string): Promise { + const cleanups = deferredCompactionCleanups.get(sessionId); if (!cleanups) { return; } @@ -553,7 +553,7 @@ export function createCopilotAgentHarness( const currentCompatKey = computeSessionCompatKey(params); const currentCompactKey = computeSessionCompactKey(params); const compactionCleanupPending = - openclawSessionId !== undefined && timedOutCompactionCleanups.has(openclawSessionId); + openclawSessionId !== undefined && deferredCompactionCleanups.has(openclawSessionId); const tracked = openclawSessionId && !compactionCleanupPending ? trackedSessions.get(openclawSessionId) @@ -614,21 +614,21 @@ export function createCopilotAgentHarness( } } : undefined, - onTimedOutCompaction: openclawSessionId + onDeferredCompaction: openclawSessionId ? ({ abort, cleanup, sdkSessionId, }: { abort: () => void; - cleanup: Promise; + cleanup: Promise; sdkSessionId: string; }) => { const tracked = trackedSessions.get(openclawSessionId); const stored = lookupStoredBinding(options?.sessionStore, openclawSessionId); const ownsTrackedSession = tracked?.sdkSessionId === sdkSessionId; const ownsStoredSession = stored?.sdkSessionId === sdkSessionId; - trackTimedOutCompactionCleanup({ + trackDeferredCompactionCleanup({ abort, cleanup, sessionId: openclawSessionId, @@ -636,9 +636,9 @@ export function createCopilotAgentHarness( if (!ownsTrackedSession && !ownsStoredSession) { return; } - // The timed-out attempt retains this SDK session until its - // background compaction resolves. Preserve its binding for a - // successful completion, but do not let a new turn resume it yet. + // The attempt retains this SDK session until its background + // compaction resolves. Preserve its binding for a successful + // completion, but do not let a new turn resume it yet. resetBlockedStoredSessions.add(openclawSessionId); void cleanup.then((outcome) => { const currentTracked = trackedSessions.get(openclawSessionId); @@ -681,7 +681,7 @@ export function createCopilotAgentHarness( if (!openclawSessionId) { return; } - await abortTimedOutCompactionCleanups(openclawSessionId); + await abortDeferredCompactionCleanups(openclawSessionId); const tracked = trackedSessions.get(openclawSessionId); if (deleteStoredBinding(options?.sessionStore, openclawSessionId)) { resetBlockedStoredSessions.delete(openclawSessionId); @@ -845,11 +845,11 @@ export function createCopilotAgentHarness( if (inFlight.size > 0) { await Promise.allSettled(inFlight); } - // Deferred compaction callbacks retain pooled clients after a timeout. + // Deferred compaction callbacks retain pooled clients after an attempt. // Cancel them before pool disposal so they cannot outlive this harness. - const cleanupSessionIds = [...timedOutCompactionCleanups.keys()]; + const cleanupSessionIds = [...deferredCompactionCleanups.keys()]; for (const sessionId of cleanupSessionIds) { - await abortTimedOutCompactionCleanups(sessionId); + await abortDeferredCompactionCleanups(sessionId); } trackedSessions.clear(); resetBlockedStoredSessions.clear(); diff --git a/extensions/copilot/src/attempt.test.ts b/extensions/copilot/src/attempt.test.ts index 3acfbf3493d3..7e4a1bc3f2d2 100644 --- a/extensions/copilot/src/attempt.test.ts +++ b/extensions/copilot/src/attempt.test.ts @@ -441,31 +441,38 @@ describe("runCopilotAttempt", () => { expect(beforeCompaction.mock.calls[0]?.[0]).not.toHaveProperty("messages"); }); - it("cancels cleanup when SDK compaction has no completion event", async () => { + it("defers and cancels compaction when the caller aborts after a turn result", async () => { const controller = new AbortController(); + const onDeferredCompaction = vi.fn(); let activeSession: FakeSession | undefined; const sdk = makeFakeSdk({ onCreateSession: (session) => { activeSession = session; session.sendAndWait.mockImplementationOnce(async () => { session.emit("session.compaction_start", {}); + setTimeout(() => controller.abort(), 0); return makeAssistantMessageEvent("done"); }); }, }); const attempt = runCopilotAttempt(makeParams({ abortSignal: controller.signal }), { + onDeferredCompaction, pool: makeFakePool(sdk), }); - await vi.waitFor(() => { - expect(activeSession?.sendAndWait).toHaveBeenCalled(); - }); - controller.abort(); const result = await attempt; expect(result.aborted).toBe(true); + expect(activeSession?.abort).not.toHaveBeenCalled(); expect(activeSession?.disconnect).toHaveBeenCalledTimes(1); + expect(activeSession?.rpc.history.cancelBackgroundCompaction).toHaveBeenCalledTimes(1); + expect(sdk.client.deleteSession).toHaveBeenCalledWith("sess-1"); + expect(onDeferredCompaction).toHaveBeenCalledWith( + expect.objectContaining({ + sdkSessionId: "sess-1", + }), + ); }); it("reports the native prompt hook's effective input through llm_input", async () => { diff --git a/extensions/copilot/src/attempt.ts b/extensions/copilot/src/attempt.ts index ffe28669c450..c09aba9c11f4 100644 --- a/extensions/copilot/src/attempt.ts +++ b/extensions/copilot/src/attempt.ts @@ -140,11 +140,11 @@ export interface CopilotAttemptDeps { sessionConfig: CopilotSessionConfig; }) => void; /** - * Called before a timed-out attempt retains its live SDK session to observe - * background compaction. The harness must prevent that session ID from being - * resumed until cleanup completes. + * Called before an attempt retains its live SDK session to observe background + * compaction. The harness must prevent that session ID from being resumed + * until cleanup completes. */ - onTimedOutCompaction?: (info: { + onDeferredCompaction?: (info: { abort: () => void; cleanup: Promise<"aborted" | "completed" | "deadline">; sdkSessionId: string; @@ -211,7 +211,7 @@ async function awaitCompactionCompletionOrAbort( } } -function deferTimedOutCompactionCleanup(params: { +function deferBackgroundCompactionCleanup(params: { abortSignal: AbortSignal | undefined; bridge: ReturnType; handle: PooledClient; @@ -220,8 +220,8 @@ function deferTimedOutCompactionCleanup(params: { session: SessionLike; timeoutMs: number; }): Promise<"aborted" | "completed" | "deadline"> { - // sendAndWait can time out while the SDK continues background compaction. - // Keep its bridge attached so after_compaction uses the originating run context. + // The SDK can compact after its turn result or a timeout. Keep the bridge + // attached so after_compaction uses the originating run context. return (async () => { let outcome: "aborted" | "completed" | "deadline" = "deadline"; try { @@ -781,8 +781,19 @@ export async function runCopilotAttempt( } } finally { settled = true; - if (timedOut && bridge?.isCompacting() && session && handle) { - timedOutDuringCompaction = true; + const compactionCompletionOutcome = + waitForCompactionCompletion && !aborted && !params.abortSignal?.aborted + ? await awaitCompactionCompletionOrAbort(bridge!, params.abortSignal) + : undefined; + const deferCompactionCleanup = + bridge?.isCompacting() && + session && + handle && + (timedOut || + compactionCompletionOutcome === "aborted" || + params.abortSignal?.aborted === true); + if (deferCompactionCleanup && bridge && session && handle) { + timedOutDuringCompaction ||= timedOut; const cleanupAbort = new AbortController(); const abortCleanup = () => cleanupAbort.abort(); if (params.abortSignal?.aborted) { @@ -790,7 +801,7 @@ export async function runCopilotAttempt( } else { params.abortSignal?.addEventListener("abort", abortCleanup, { once: true }); } - const cleanup = deferTimedOutCompactionCleanup({ + const cleanup = deferBackgroundCompactionCleanup({ abortSignal: cleanupAbort.signal, bridge, handle, @@ -806,7 +817,7 @@ export async function runCopilotAttempt( .catch(() => undefined); if (sdkSessionId) { try { - deps.onTimedOutCompaction?.({ + deps.onDeferredCompaction?.({ abort: () => cleanupAbort.abort(), cleanup, sdkSessionId, @@ -817,9 +828,7 @@ export async function runCopilotAttempt( } params.abortSignal?.removeEventListener("abort", onAbort); } else { - if (waitForCompactionCompletion && !aborted && !params.abortSignal?.aborted) { - await awaitCompactionCompletionOrAbort(bridge!, params.abortSignal); - } else { + if (compactionCompletionOutcome === undefined) { await bridge?.awaitCompactionChain(); } bridge?.detach();