Merge pull request #3053 from code-yeongyu/fix/p0-3-await-session-abort
fix(background-agent): await session abort completion paths
This commit is contained in:
@@ -23,6 +23,21 @@ import { selectFallbackProvider } from "../../shared/model-error-classifier"
|
||||
import { readProviderModelsCache } from "../../shared"
|
||||
import type { BackgroundTask } from "./types"
|
||||
import type { ConcurrencyManager } from "./concurrency"
|
||||
import type { OpencodeClient, QueueItem } from "./constants"
|
||||
|
||||
function createDeferredPromise(): {
|
||||
promise: Promise<void>
|
||||
resolve: () => void
|
||||
} {
|
||||
let resolvePromise = () => {}
|
||||
const promise = new Promise<void>((resolve) => {
|
||||
resolvePromise = resolve
|
||||
})
|
||||
return {
|
||||
promise,
|
||||
resolve: resolvePromise,
|
||||
}
|
||||
}
|
||||
|
||||
function createMockTask(overrides: Partial<BackgroundTask> = {}): BackgroundTask {
|
||||
return {
|
||||
@@ -53,20 +68,27 @@ function createMockConcurrencyManager(): ConcurrencyManager {
|
||||
} as unknown as ConcurrencyManager
|
||||
}
|
||||
|
||||
function createMockClient() {
|
||||
function createMockClient(): {
|
||||
client: OpencodeClient
|
||||
abortMock: ReturnType<typeof mock>
|
||||
} {
|
||||
const abortMock = mock(async () => ({}))
|
||||
return {
|
||||
session: {
|
||||
abort: mock(async () => ({})),
|
||||
},
|
||||
} as any
|
||||
client: {
|
||||
session: {
|
||||
abort: abortMock,
|
||||
},
|
||||
} as unknown as OpencodeClient,
|
||||
abortMock,
|
||||
}
|
||||
}
|
||||
|
||||
function createDefaultArgs(taskOverrides: Partial<BackgroundTask> = {}) {
|
||||
const processKeyFn = mock(() => {})
|
||||
const queuesByKey = new Map<string, Array<{ task: BackgroundTask; input: any }>>()
|
||||
const queuesByKey = new Map<string, QueueItem[]>()
|
||||
const idleDeferralTimers = new Map<string, ReturnType<typeof setTimeout>>()
|
||||
const concurrencyManager = createMockConcurrencyManager()
|
||||
const client = createMockClient()
|
||||
const { client, abortMock } = createMockClient()
|
||||
const task = createMockTask(taskOverrides)
|
||||
|
||||
return {
|
||||
@@ -75,6 +97,7 @@ function createDefaultArgs(taskOverrides: Partial<BackgroundTask> = {}) {
|
||||
source: "polling",
|
||||
concurrencyManager,
|
||||
client,
|
||||
abortMock,
|
||||
idleDeferralTimers,
|
||||
queuesByKey,
|
||||
processKey: processKeyFn,
|
||||
@@ -93,97 +116,118 @@ describe("tryFallbackRetry", () => {
|
||||
})
|
||||
|
||||
describe("#given retryable error with fallback chain", () => {
|
||||
test("returns true and enqueues retry", () => {
|
||||
test("returns true and enqueues retry", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
const result = tryFallbackRetry(args)
|
||||
const result = await tryFallbackRetry(args)
|
||||
|
||||
expect(result).toBe(true)
|
||||
})
|
||||
|
||||
test("resets task status to pending", () => {
|
||||
test("resets task status to pending", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.status).toBe("pending")
|
||||
})
|
||||
|
||||
test("increments attemptCount", () => {
|
||||
test("increments attemptCount", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.attemptCount).toBe(1)
|
||||
})
|
||||
|
||||
test("updates task model to fallback", () => {
|
||||
test("updates task model to fallback", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.model?.modelID).toBe("fallback-model-1")
|
||||
expect(args.task.model?.providerID).toBe("provider-a")
|
||||
})
|
||||
|
||||
test("clears sessionID and startedAt", () => {
|
||||
test("clears sessionID and startedAt", async () => {
|
||||
const args = createDefaultArgs({
|
||||
sessionID: "old-session",
|
||||
startedAt: new Date(),
|
||||
})
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.sessionID).toBeUndefined()
|
||||
expect(args.task.startedAt).toBeUndefined()
|
||||
})
|
||||
|
||||
test("clears error field", () => {
|
||||
test("clears error field", async () => {
|
||||
const args = createDefaultArgs({ error: "previous error" })
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.error).toBeUndefined()
|
||||
})
|
||||
|
||||
test("sets new queuedAt", () => {
|
||||
test("sets new queuedAt", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.queuedAt).toBeInstanceOf(Date)
|
||||
})
|
||||
|
||||
test("releases concurrency slot", () => {
|
||||
test("releases concurrency slot", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.concurrencyManager.release).toHaveBeenCalledWith("provider-a/original-model")
|
||||
})
|
||||
|
||||
test("clears concurrencyKey after release", () => {
|
||||
test("clears concurrencyKey after release", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.concurrencyKey).toBeUndefined()
|
||||
})
|
||||
|
||||
test("aborts existing session", () => {
|
||||
test("aborts existing session", async () => {
|
||||
const args = createDefaultArgs({ sessionID: "session-to-abort" })
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.client.session.abort).toHaveBeenCalledWith({
|
||||
expect(args.abortMock).toHaveBeenCalledWith({
|
||||
path: { id: "session-to-abort" },
|
||||
})
|
||||
})
|
||||
|
||||
test("adds retry input to queue and calls processKey", () => {
|
||||
test("waits for session abort before resolving", async () => {
|
||||
const args = createDefaultArgs({ sessionID: "session-to-abort" })
|
||||
const deferred = createDeferredPromise()
|
||||
args.abortMock.mockImplementationOnce(() => deferred.promise)
|
||||
|
||||
const retryPromise = tryFallbackRetry(args)
|
||||
let settled = false
|
||||
void retryPromise.then(() => {
|
||||
settled = true
|
||||
})
|
||||
|
||||
await Promise.resolve()
|
||||
|
||||
expect(settled).toBe(false)
|
||||
|
||||
deferred.resolve()
|
||||
await retryPromise
|
||||
|
||||
expect(settled).toBe(true)
|
||||
})
|
||||
|
||||
test("adds retry input to queue and calls processKey", async () => {
|
||||
const args = createDefaultArgs()
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
const key = `${args.task.model!.providerID}/${args.task.model!.modelID}`
|
||||
const queue = args.queuesByKey.get(key)
|
||||
@@ -195,81 +239,81 @@ describe("tryFallbackRetry", () => {
|
||||
})
|
||||
|
||||
describe("#given non-retryable error", () => {
|
||||
test("returns false when shouldRetryError returns false", () => {
|
||||
test("returns false when shouldRetryError returns false", async () => {
|
||||
;(shouldRetryError as any).mockImplementation(() => false)
|
||||
const args = createDefaultArgs()
|
||||
|
||||
const result = tryFallbackRetry(args)
|
||||
const result = await tryFallbackRetry(args)
|
||||
|
||||
expect(result).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("#given no fallback chain", () => {
|
||||
test("returns false when fallbackChain is undefined", () => {
|
||||
test("returns false when fallbackChain is undefined", async () => {
|
||||
const args = createDefaultArgs({ fallbackChain: undefined })
|
||||
|
||||
const result = tryFallbackRetry(args)
|
||||
const result = await tryFallbackRetry(args)
|
||||
|
||||
expect(result).toBe(false)
|
||||
})
|
||||
|
||||
test("returns false when fallbackChain is empty", () => {
|
||||
test("returns false when fallbackChain is empty", async () => {
|
||||
const args = createDefaultArgs({ fallbackChain: [] })
|
||||
|
||||
const result = tryFallbackRetry(args)
|
||||
const result = await tryFallbackRetry(args)
|
||||
|
||||
expect(result).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("#given exhausted fallbacks", () => {
|
||||
test("returns false when attemptCount exceeds chain length", () => {
|
||||
test("returns false when attemptCount exceeds chain length", async () => {
|
||||
const args = createDefaultArgs({ attemptCount: 5 })
|
||||
|
||||
const result = tryFallbackRetry(args)
|
||||
const result = await tryFallbackRetry(args)
|
||||
|
||||
expect(result).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("#given task without concurrency key", () => {
|
||||
test("skips concurrency release", () => {
|
||||
test("skips concurrency release", async () => {
|
||||
const args = createDefaultArgs({ concurrencyKey: undefined })
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.concurrencyManager.release).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("#given task without session", () => {
|
||||
test("skips session abort", () => {
|
||||
test("skips session abort", async () => {
|
||||
const args = createDefaultArgs({ sessionID: undefined })
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.client.session.abort).not.toHaveBeenCalled()
|
||||
expect(args.abortMock).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("#given active idle deferral timer", () => {
|
||||
test("clears the timer and removes from map", () => {
|
||||
test("clears the timer and removes from map", async () => {
|
||||
const args = createDefaultArgs()
|
||||
const timerId = setTimeout(() => {}, 10000)
|
||||
args.idleDeferralTimers.set("test-task-1", timerId)
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.idleDeferralTimers.has("test-task-1")).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("#given second attempt", () => {
|
||||
test("uses second fallback in chain", () => {
|
||||
test("uses second fallback in chain", async () => {
|
||||
const args = createDefaultArgs({ attemptCount: 1 })
|
||||
|
||||
tryFallbackRetry(args)
|
||||
await tryFallbackRetry(args)
|
||||
|
||||
expect(args.task.model?.modelID).toBe("fallback-model-2")
|
||||
expect(args.task.attemptCount).toBe(2)
|
||||
@@ -277,7 +321,7 @@ describe("tryFallbackRetry", () => {
|
||||
})
|
||||
|
||||
describe("#given disconnected fallback providers with connected preferred provider", () => {
|
||||
test("keeps fallback entry and selects connected preferred provider", () => {
|
||||
test("keeps fallback entry and selects connected preferred provider", async () => {
|
||||
;(readProviderModelsCache as any).mockReturnValueOnce({ connected: ["provider-a"] })
|
||||
;(selectFallbackProvider as any).mockImplementationOnce(
|
||||
(_providers: string[], preferredProviderID?: string) => preferredProviderID ?? "provider-b",
|
||||
@@ -288,7 +332,7 @@ describe("tryFallbackRetry", () => {
|
||||
model: { providerID: "provider-a", modelID: "original-model" },
|
||||
})
|
||||
|
||||
const result = tryFallbackRetry(args)
|
||||
const result = await tryFallbackRetry(args)
|
||||
|
||||
expect(result).toBe(true)
|
||||
expect(args.task.model?.providerID).toBe("provider-a")
|
||||
|
||||
@@ -11,7 +11,7 @@ import {
|
||||
} from "../../shared/model-error-classifier"
|
||||
import { transformModelForProvider } from "../../shared/provider-model-id-transform"
|
||||
|
||||
export function tryFallbackRetry(args: {
|
||||
export async function tryFallbackRetry(args: {
|
||||
task: BackgroundTask
|
||||
errorInfo: { name?: string; message?: string }
|
||||
source: string
|
||||
@@ -20,7 +20,7 @@ export function tryFallbackRetry(args: {
|
||||
idleDeferralTimers: Map<string, ReturnType<typeof setTimeout>>
|
||||
queuesByKey: Map<string, QueueItem[]>
|
||||
processKey: (key: string) => void
|
||||
}): boolean {
|
||||
}): Promise<boolean> {
|
||||
const { task, errorInfo, source, concurrencyManager, client, idleDeferralTimers, queuesByKey, processKey } = args
|
||||
const fallbackChain = task.fallbackChain
|
||||
const canRetry =
|
||||
@@ -84,16 +84,14 @@ export function tryFallbackRetry(args: {
|
||||
task.concurrencyKey = undefined
|
||||
}
|
||||
|
||||
if (task.sessionID) {
|
||||
client.session.abort({ path: { id: task.sessionID } }).catch(() => {})
|
||||
}
|
||||
|
||||
const idleTimer = idleDeferralTimers.get(task.id)
|
||||
if (idleTimer) {
|
||||
clearTimeout(idleTimer)
|
||||
idleDeferralTimers.delete(task.id)
|
||||
}
|
||||
|
||||
const previousSessionID = task.sessionID
|
||||
|
||||
task.attemptCount = selectedAttemptCount
|
||||
const transformedModelId = transformModelForProvider(providerID, nextFallback.model)
|
||||
task.model = {
|
||||
@@ -123,6 +121,11 @@ export function tryFallbackRetry(args: {
|
||||
category: task.category,
|
||||
isUnstableAgent: task.isUnstableAgent,
|
||||
}
|
||||
|
||||
if (previousSessionID) {
|
||||
await client.session.abort({ path: { id: previousSessionID } }).catch(() => {})
|
||||
}
|
||||
|
||||
queue.push({ task, input: retryInput })
|
||||
queuesByKey.set(key, queue)
|
||||
processKey(key)
|
||||
|
||||
@@ -6,6 +6,20 @@ import { SessionCategoryRegistry } from "../../shared/session-category-registry"
|
||||
import { BackgroundManager } from "./manager"
|
||||
import type { BackgroundTask } from "./types"
|
||||
|
||||
function createDeferredPromise(): {
|
||||
promise: Promise<void>
|
||||
resolve: () => void
|
||||
} {
|
||||
let resolvePromise = () => {}
|
||||
const promise = new Promise<void>((resolve) => {
|
||||
resolvePromise = resolve
|
||||
})
|
||||
return {
|
||||
promise,
|
||||
resolve: resolvePromise,
|
||||
}
|
||||
}
|
||||
|
||||
function createTask(overrides: Partial<BackgroundTask> & { id: string; sessionID: string }): BackgroundTask {
|
||||
return {
|
||||
parentSessionID: "parent-session",
|
||||
@@ -94,4 +108,48 @@ describe("BackgroundManager shutdown global cleanup", () => {
|
||||
expect(SessionCategoryRegistry.has(completedSessionID)).toBe(false)
|
||||
expect(SessionCategoryRegistry.has(unrelatedSessionID)).toBe(true)
|
||||
})
|
||||
|
||||
test("awaits running session aborts before shutdown resolves", async () => {
|
||||
// given
|
||||
const runningSessionID = "ses-running-await-shutdown"
|
||||
const deferred = createDeferredPromise()
|
||||
const manager = createBackgroundManager()
|
||||
const tasks = new Map<string, BackgroundTask>([
|
||||
[
|
||||
"task-running-await-shutdown",
|
||||
createTask({
|
||||
id: "task-running-await-shutdown",
|
||||
sessionID: runningSessionID,
|
||||
}),
|
||||
],
|
||||
])
|
||||
|
||||
Object.assign(manager, { tasks })
|
||||
Object.assign(manager, {
|
||||
client: {
|
||||
session: {
|
||||
abort: () => deferred.promise,
|
||||
prompt: async () => ({}),
|
||||
promptAsync: async () => ({}),
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
// when
|
||||
const shutdownPromise = manager.shutdown()
|
||||
let settled = false
|
||||
void shutdownPromise.then(() => {
|
||||
settled = true
|
||||
})
|
||||
|
||||
await Promise.resolve()
|
||||
|
||||
// then
|
||||
expect(settled).toBe(false)
|
||||
|
||||
deferred.resolve()
|
||||
await shutdownPromise
|
||||
|
||||
expect(settled).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -901,7 +901,12 @@ export class BackgroundManager {
|
||||
name: extractErrorName(assistantError),
|
||||
message: extractErrorMessage(assistantError),
|
||||
}
|
||||
this.tryFallbackRetry(task, errorInfo, "message.updated")
|
||||
void this.tryFallbackRetry(task, errorInfo, "message.updated").catch((error) => {
|
||||
log("[background-agent] Error handling message.updated fallback retry:", {
|
||||
error,
|
||||
taskId: task.id,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
if (event.type === "message.part.updated" || event.type === "message.part.delta") {
|
||||
@@ -1015,62 +1020,18 @@ export class BackgroundManager {
|
||||
const errorMessage = props ? getSessionErrorMessage(props) : undefined
|
||||
|
||||
const errorInfo = { name: errorName, message: errorMessage }
|
||||
if (this.tryFallbackRetry(task, errorInfo, "session.error")) return
|
||||
|
||||
// Original error handling (no retry)
|
||||
const errorMsg = errorMessage ?? "Session error"
|
||||
const canRetry =
|
||||
shouldRetryError(errorInfo) &&
|
||||
!!task.fallbackChain &&
|
||||
hasMoreFallbacks(task.fallbackChain, task.attemptCount ?? 0)
|
||||
log("[background-agent] Session error - no retry:", {
|
||||
taskId: task.id,
|
||||
void this.handleSessionErrorEvent({
|
||||
errorInfo,
|
||||
errorMessage,
|
||||
errorName,
|
||||
errorMessage: errorMsg?.slice(0, 100),
|
||||
hasFallbackChain: !!task.fallbackChain,
|
||||
canRetry,
|
||||
})
|
||||
|
||||
task.status = "error"
|
||||
task.error = errorMsg
|
||||
task.completedAt = new Date()
|
||||
if (task.rootSessionID) {
|
||||
this.unregisterRootDescendant(task.rootSessionID)
|
||||
}
|
||||
this.taskHistory.record(task.parentSessionID, { id: task.id, sessionID: task.sessionID, agent: task.agent, description: task.description, status: "error", category: task.category, startedAt: task.startedAt, completedAt: task.completedAt })
|
||||
|
||||
if (task.concurrencyKey) {
|
||||
this.concurrencyManager.release(task.concurrencyKey)
|
||||
task.concurrencyKey = undefined
|
||||
}
|
||||
|
||||
const completionTimer = this.completionTimers.get(task.id)
|
||||
if (completionTimer) {
|
||||
clearTimeout(completionTimer)
|
||||
this.completionTimers.delete(task.id)
|
||||
}
|
||||
|
||||
const idleTimer = this.idleDeferralTimers.get(task.id)
|
||||
if (idleTimer) {
|
||||
clearTimeout(idleTimer)
|
||||
this.idleDeferralTimers.delete(task.id)
|
||||
}
|
||||
|
||||
this.cleanupPendingByParent(task)
|
||||
this.clearNotificationsForTask(task.id)
|
||||
const toastManager = getTaskToastManager()
|
||||
if (toastManager) {
|
||||
toastManager.removeTask(task.id)
|
||||
}
|
||||
this.scheduleTaskRemoval(task.id)
|
||||
if (task.sessionID) {
|
||||
SessionCategoryRegistry.remove(task.sessionID)
|
||||
}
|
||||
|
||||
this.markForNotification(task)
|
||||
this.enqueueNotificationForParent(task.parentSessionID, () => this.notifyParentSession(task)).catch(err => {
|
||||
log("[background-agent] Error in notifyParentSession for errored task:", { taskId: task.id, error: err })
|
||||
task,
|
||||
}).catch((error) => {
|
||||
log("[background-agent] Error handling session.error event:", {
|
||||
error,
|
||||
taskId: task.id,
|
||||
})
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if (event.type === "session.deleted") {
|
||||
@@ -1141,15 +1102,87 @@ export class BackgroundManager {
|
||||
|
||||
const errorMessage = typeof status.message === "string" ? status.message : undefined
|
||||
const errorInfo = { name: "SessionRetry", message: errorMessage }
|
||||
this.tryFallbackRetry(task, errorInfo, "session.status")
|
||||
void this.tryFallbackRetry(task, errorInfo, "session.status").catch((error) => {
|
||||
log("[background-agent] Error handling session.status fallback retry:", {
|
||||
error,
|
||||
taskId: task.id,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
private async handleSessionErrorEvent(args: {
|
||||
task: BackgroundTask
|
||||
errorInfo: { name?: string; message?: string }
|
||||
errorName: string | undefined
|
||||
errorMessage: string | undefined
|
||||
}): Promise<void> {
|
||||
const { task, errorInfo, errorMessage, errorName } = args
|
||||
|
||||
if (await this.tryFallbackRetry(task, errorInfo, "session.error")) {
|
||||
return
|
||||
}
|
||||
|
||||
const errorMsg = errorMessage ?? "Session error"
|
||||
const canRetry =
|
||||
shouldRetryError(errorInfo) &&
|
||||
!!task.fallbackChain &&
|
||||
hasMoreFallbacks(task.fallbackChain, task.attemptCount ?? 0)
|
||||
log("[background-agent] Session error - no retry:", {
|
||||
taskId: task.id,
|
||||
errorName,
|
||||
errorMessage: errorMsg?.slice(0, 100),
|
||||
hasFallbackChain: !!task.fallbackChain,
|
||||
canRetry,
|
||||
})
|
||||
|
||||
task.status = "error"
|
||||
task.error = errorMsg
|
||||
task.completedAt = new Date()
|
||||
if (task.rootSessionID) {
|
||||
this.unregisterRootDescendant(task.rootSessionID)
|
||||
}
|
||||
this.taskHistory.record(task.parentSessionID, { id: task.id, sessionID: task.sessionID, agent: task.agent, description: task.description, status: "error", category: task.category, startedAt: task.startedAt, completedAt: task.completedAt })
|
||||
|
||||
if (task.concurrencyKey) {
|
||||
this.concurrencyManager.release(task.concurrencyKey)
|
||||
task.concurrencyKey = undefined
|
||||
}
|
||||
|
||||
const completionTimer = this.completionTimers.get(task.id)
|
||||
if (completionTimer) {
|
||||
clearTimeout(completionTimer)
|
||||
this.completionTimers.delete(task.id)
|
||||
}
|
||||
|
||||
const idleTimer = this.idleDeferralTimers.get(task.id)
|
||||
if (idleTimer) {
|
||||
clearTimeout(idleTimer)
|
||||
this.idleDeferralTimers.delete(task.id)
|
||||
}
|
||||
|
||||
this.cleanupPendingByParent(task)
|
||||
this.clearNotificationsForTask(task.id)
|
||||
const toastManager = getTaskToastManager()
|
||||
if (toastManager) {
|
||||
toastManager.removeTask(task.id)
|
||||
}
|
||||
this.scheduleTaskRemoval(task.id)
|
||||
if (task.sessionID) {
|
||||
SessionCategoryRegistry.remove(task.sessionID)
|
||||
}
|
||||
|
||||
this.markForNotification(task)
|
||||
this.enqueueNotificationForParent(task.parentSessionID, () => this.notifyParentSession(task)).catch(err => {
|
||||
log("[background-agent] Error in notifyParentSession for errored task:", { taskId: task.id, error: err })
|
||||
})
|
||||
}
|
||||
|
||||
private tryFallbackRetry(
|
||||
task: BackgroundTask,
|
||||
errorInfo: { name?: string; message?: string },
|
||||
source: string,
|
||||
): boolean {
|
||||
): Promise<boolean> {
|
||||
const previousSessionID = task.sessionID
|
||||
const result = tryFallbackRetry({
|
||||
task,
|
||||
@@ -1161,10 +1194,12 @@ export class BackgroundManager {
|
||||
queuesByKey: this.queuesByKey,
|
||||
processKey: (key: string) => this.processKey(key),
|
||||
})
|
||||
if (result && previousSessionID) {
|
||||
subagentSessions.delete(previousSessionID)
|
||||
}
|
||||
return result
|
||||
return result.then((retried) => {
|
||||
if (retried && previousSessionID) {
|
||||
subagentSessions.delete(previousSessionID)
|
||||
}
|
||||
return retried
|
||||
})
|
||||
}
|
||||
|
||||
markForNotification(task: BackgroundTask): void {
|
||||
@@ -1889,7 +1924,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
|
||||
? (sessionStatus as { message?: string }).message
|
||||
: undefined
|
||||
const errorInfo = { name: "SessionRetry", message: retryMessage }
|
||||
if (this.tryFallbackRetry(task, errorInfo, "polling:session.status")) {
|
||||
if (await this.tryFallbackRetry(task, errorInfo, "polling:session.status")) {
|
||||
continue
|
||||
}
|
||||
}
|
||||
@@ -1981,6 +2016,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
|
||||
log("[background-agent] Shutting down BackgroundManager")
|
||||
this.stopPolling()
|
||||
const trackedSessionIDs = new Set<string>()
|
||||
const abortRequests: Array<{ sessionID: string; promise: Promise<unknown> }> = []
|
||||
|
||||
// Abort all running sessions to prevent zombie processes (#1240)
|
||||
for (const task of this.tasks.values()) {
|
||||
@@ -1989,9 +2025,24 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
|
||||
}
|
||||
|
||||
if (task.status === "running" && task.sessionID) {
|
||||
this.client.session.abort({
|
||||
path: { id: task.sessionID },
|
||||
}).catch(() => {})
|
||||
abortRequests.push({
|
||||
sessionID: task.sessionID,
|
||||
promise: this.client.session.abort({
|
||||
path: { id: task.sessionID },
|
||||
}),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if (abortRequests.length > 0) {
|
||||
const abortResults = await Promise.allSettled(abortRequests.map((request) => request.promise))
|
||||
for (const [index, abortResult] of abortResults.entries()) {
|
||||
if (abortResult.status === "fulfilled") continue
|
||||
|
||||
log("[background-agent] Error aborting session during shutdown:", {
|
||||
error: abortResult.reason,
|
||||
sessionID: abortRequests[index]?.sessionID,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,6 +16,20 @@ describe("checkAndInterruptStaleTasks", () => {
|
||||
}
|
||||
const mockNotify = mock(() => Promise.resolve())
|
||||
|
||||
function createDeferredPromise(): {
|
||||
promise: Promise<void>
|
||||
resolve: () => void
|
||||
} {
|
||||
let resolvePromise = () => {}
|
||||
const promise = new Promise<void>((resolve) => {
|
||||
resolvePromise = resolve
|
||||
})
|
||||
return {
|
||||
promise,
|
||||
resolve: resolvePromise,
|
||||
}
|
||||
}
|
||||
|
||||
function createRunningTask(overrides: Partial<BackgroundTask> = {}): 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({
|
||||
|
||||
@@ -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<Promise<unknown>> = []
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user