diff --git a/src/features/background-agent/task-poller.test.ts b/src/features/background-agent/task-poller.test.ts index 0343f99c0..ad08265f6 100644 --- a/src/features/background-agent/task-poller.test.ts +++ b/src/features/background-agent/task-poller.test.ts @@ -16,6 +16,20 @@ describe("checkAndInterruptStaleTasks", () => { } const mockNotify = mock(() => Promise.resolve()) + function createDeferredPromise(): { + promise: Promise + resolve: () => void + } { + let resolvePromise = () => {} + const promise = new Promise((resolve) => { + resolvePromise = resolve + }) + return { + promise, + resolve: resolvePromise, + } + } + function createRunningTask(overrides: Partial = {}): BackgroundTask { return { id: "task-1", @@ -114,6 +128,39 @@ describe("checkAndInterruptStaleTasks", () => { expect(task.error).toContain("no activity") }) + it("should await abort before resolving for no-progress stale interruption", async () => { + //#given + const task = createRunningTask({ + startedAt: new Date(Date.now() - 15 * 60 * 1000), + progress: undefined, + }) + const deferred = createDeferredPromise() + mockClient.session.abort.mockImplementationOnce(() => deferred.promise) + + //#when + const interruptPromise = checkAndInterruptStaleTasks({ + tasks: [task], + client: mockClient as never, + config: { messageStalenessTimeoutMs: 600_000 }, + concurrencyManager: mockConcurrencyManager as never, + notifyParentSession: mockNotify, + }) + let settled = false + void interruptPromise.then(() => { + settled = true + }) + + await Promise.resolve() + + //#then + expect(settled).toBe(false) + + deferred.resolve() + await interruptPromise + + expect(settled).toBe(true) + }) + it("should NOT interrupt tasks with NO progress.lastUpdate that are within messageStalenessTimeoutMs", async () => { //#given — task started 5 minutes ago, default timeout is 10 minutes const task = createRunningTask({ @@ -407,6 +454,45 @@ describe("checkAndInterruptStaleTasks", () => { expect(task.error).toContain("session gone from status registry") }) + it("should await abort before resolving for session-gone interruption", async () => { + //#given + const task = createRunningTask({ + startedAt: new Date(Date.now() - 300_000), + progress: { + toolCalls: 1, + lastUpdate: new Date(Date.now() - 120_000), + }, + consecutiveMissedPolls: 2, + }) + const deferred = createDeferredPromise() + mockClient.session.get.mockRejectedValue(new Error("missing")) + mockClient.session.abort.mockImplementationOnce(() => deferred.promise) + + //#when + const interruptPromise = checkAndInterruptStaleTasks({ + tasks: [task], + client: mockClient as never, + config: { staleTimeoutMs: 180_000, sessionGoneTimeoutMs: 60_000 }, + concurrencyManager: mockConcurrencyManager as never, + notifyParentSession: mockNotify, + sessionStatuses: {}, + }) + let settled = false + void interruptPromise.then(() => { + settled = true + }) + + await Promise.resolve() + + //#then + expect(settled).toBe(false) + + deferred.resolve() + await interruptPromise + + expect(settled).toBe(true) + }) + it("should use session-gone timeout when session is missing from status map (no progress)", async () => { //#given — task started 2min ago, no progress, session completely gone const task = createRunningTask({ diff --git a/src/features/background-agent/task-poller.ts b/src/features/background-agent/task-poller.ts index 1b32a55f4..0f2c6e2ce 100644 --- a/src/features/background-agent/task-poller.ts +++ b/src/features/background-agent/task-poller.ts @@ -119,6 +119,7 @@ export async function checkAndInterruptStaleTasks(args: { const staleTimeoutMs = config?.staleTimeoutMs ?? DEFAULT_STALE_TIMEOUT_MS const sessionGoneTimeoutMs = config?.sessionGoneTimeoutMs ?? DEFAULT_SESSION_GONE_TIMEOUT_MS const now = Date.now() + const abortPromises: Array> = [] const messageStalenessMs = config?.messageStalenessTimeoutMs ?? DEFAULT_MESSAGE_STALENESS_TIMEOUT_MS @@ -166,7 +167,7 @@ export async function checkAndInterruptStaleTasks(args: { onTaskInterrupted(task) - client.session.abort({ path: { id: sessionID } }).catch(() => {}) + abortPromises.push(client.session.abort({ path: { id: sessionID } })) log(`[background-agent] Task ${task.id} interrupted: no progress since start`) try { @@ -204,7 +205,7 @@ export async function checkAndInterruptStaleTasks(args: { onTaskInterrupted(task) - client.session.abort({ path: { id: sessionID } }).catch(() => {}) + abortPromises.push(client.session.abort({ path: { id: sessionID } })) log(`[background-agent] Task ${task.id} interrupted: stale timeout`) try { @@ -213,4 +214,8 @@ export async function checkAndInterruptStaleTasks(args: { log("[background-agent] Error in notifyParentSession for stale task:", { taskId: task.id, error: err }) } } + + if (abortPromises.length > 0) { + await Promise.allSettled(abortPromises) + } }