diff --git a/src/tools/delegate-task/cancel-unstable-agent-task.ts b/src/tools/delegate-task/cancel-unstable-agent-task.ts new file mode 100644 index 000000000..22deed878 --- /dev/null +++ b/src/tools/delegate-task/cancel-unstable-agent-task.ts @@ -0,0 +1,19 @@ +import type { ExecutorContext } from "./executor-types" + +export async function cancelUnstableAgentTask( + manager: ExecutorContext["manager"], + taskID: string | undefined, + reason: string +): Promise { + if (!taskID || typeof manager.cancelTask !== "function") { + return + } + + await Promise.allSettled([ + manager.cancelTask(taskID, { + source: "unstable-agent-task", + reason, + skipNotification: true, + }), + ]) +} diff --git a/src/tools/delegate-task/unstable-agent-cleanup.test.ts b/src/tools/delegate-task/unstable-agent-cleanup.test.ts new file mode 100644 index 000000000..298bbdb32 --- /dev/null +++ b/src/tools/delegate-task/unstable-agent-cleanup.test.ts @@ -0,0 +1,176 @@ +declare const require: (name: string) => any +const { describe, test, expect, beforeEach, afterEach } = require("bun:test") + +import { __resetTimingConfig, __setTimingConfig } from "./timing" + +function createArgs() { + return { + description: "cleanup case", + prompt: "run", + category: "unspecified-low", + run_in_background: false, + load_skills: [], + command: undefined, + } +} + +function createToolContext(aborted = false) { + const controller = new AbortController() + if (aborted) { + controller.abort() + } + + return { + sessionID: "parent-session", + messageID: "parent-message", + agent: "test-agent", + abort: controller.signal, + metadata: () => Promise.resolve(), + } +} + +function createParentContext() { + return { + sessionID: "parent-session", + messageID: "parent-message", + model: "gpt-test", + agent: "test-agent", + } +} + +describe("executeUnstableAgentTask cleanup", () => { + beforeEach(() => { + __setTimingConfig({ + POLL_INTERVAL_MS: 10, + MIN_STABILITY_TIME_MS: 0, + STABILITY_POLLS_REQUIRED: 1, + WAIT_FOR_SESSION_TIMEOUT_MS: 100, + WAIT_FOR_SESSION_INTERVAL_MS: 10, + }) + }) + + afterEach(() => { + __resetTimingConfig() + }) + + test("cancels launched task when parent aborts during monitoring", async () => { + // given + const { executeUnstableAgentTask } = require("./unstable-agent-task") + const cancelCalls: Array<{ taskId: string; options?: Record }> = [] + + const mockManager = { + launch: async () => ({ id: "bg_abort_monitoring", sessionID: "ses_abort_monitoring", status: "running" }), + getTask: () => ({ id: "bg_abort_monitoring", sessionID: "ses_abort_monitoring", status: "running" }), + cancelTask: async (taskId: string, options?: Record) => { + cancelCalls.push({ taskId, options }) + return true + }, + } + + // when + const result = await executeUnstableAgentTask( + createArgs(), + createToolContext(true), + { + manager: mockManager, + client: { + session: { + status: async () => ({ data: {} }), + messages: async () => ({ data: [] }), + }, + }, + }, + createParentContext(), + "test-agent", + undefined, + undefined, + "gpt-test" + ) + + // then + expect(result).toContain("Task aborted (was running in background mode).") + expect(cancelCalls).toHaveLength(1) + expect(cancelCalls[0]?.taskId).toBe("bg_abort_monitoring") + }) + + test("cancels launched task when monitored timeout budget is exhausted", async () => { + // given + const { executeUnstableAgentTask } = require("./unstable-agent-task") + const cancelCalls: Array<{ taskId: string; options?: Record }> = [] + + const mockManager = { + launch: async () => ({ id: "bg_timeout_cleanup", sessionID: "ses_timeout_cleanup", status: "running" }), + getTask: () => ({ id: "bg_timeout_cleanup", sessionID: "ses_timeout_cleanup", status: "running" }), + cancelTask: async (taskId: string, options?: Record) => { + cancelCalls.push({ taskId, options }) + return true + }, + } + + // when + const result = await executeUnstableAgentTask( + createArgs(), + createToolContext(), + { + manager: mockManager, + client: { + session: { + status: async () => ({ data: { ses_timeout_cleanup: { type: "running" } } }), + messages: async () => ({ data: [] }), + }, + }, + syncPollTimeoutMs: 0, + }, + createParentContext(), + "test-agent", + undefined, + undefined, + "gpt-test" + ) + + // then + expect(result).toContain("SUPERVISED TASK TIMED OUT") + expect(cancelCalls).toHaveLength(1) + expect(cancelCalls[0]?.taskId).toBe("bg_timeout_cleanup") + }) + + test("cancels launched task when parent aborts while waiting for session start", async () => { + // given + const { executeUnstableAgentTask } = require("./unstable-agent-task") + const cancelCalls: Array<{ taskId: string; options?: Record }> = [] + + const mockManager = { + launch: async () => ({ id: "bg_wait_abort", status: "pending" }), + getTask: () => ({ id: "bg_wait_abort", status: "pending" }), + cancelTask: async (taskId: string, options?: Record) => { + cancelCalls.push({ taskId, options }) + return true + }, + } + + // when + const result = await executeUnstableAgentTask( + createArgs(), + createToolContext(true), + { + manager: mockManager, + client: { + session: { + status: async () => ({ data: {} }), + messages: async () => ({ data: [] }), + }, + }, + }, + createParentContext(), + "test-agent", + undefined, + undefined, + "gpt-test" + ) + + // then + expect(result).toContain("Task aborted while waiting for session to start.") + expect(cancelCalls).toHaveLength(1) + expect(cancelCalls[0]?.taskId).toBe("bg_wait_abort") + }) +}) diff --git a/src/tools/delegate-task/unstable-agent-task.ts b/src/tools/delegate-task/unstable-agent-task.ts index 5b92955bd..335699fda 100644 --- a/src/tools/delegate-task/unstable-agent-task.ts +++ b/src/tools/delegate-task/unstable-agent-task.ts @@ -2,6 +2,7 @@ import type { DelegateTaskArgs, ToolContextWithMetadata } from "./types" import type { ExecutorContext, ParentContext, SessionMessage } from "./executor-types" import { DEFAULT_SYNC_POLL_TIMEOUT_MS, getTimingConfig } from "./timing" import { buildTaskPrompt } from "./prompt-builder" +import { cancelUnstableAgentTask } from "./cancel-unstable-agent-task" import { storeToolMetadata } from "../../features/tool-metadata-store" import { formatDuration } from "./time-formatter" import { formatDetailedError } from "./error-formatting" @@ -20,6 +21,8 @@ export async function executeUnstableAgentTask( actualModel: string | undefined ): Promise { const { manager, client, syncPollTimeoutMs } = executorCtx + let cleanupReason: string | undefined + let launchedTaskID: string | undefined try { const effectivePrompt = buildTaskPrompt(args.prompt, agentToUse) @@ -38,12 +41,14 @@ export async function executeUnstableAgentTask( category: args.category, sessionPermission: QUESTION_DENIED_SESSION_PERMISSION, }) + launchedTaskID = task.id const timing = getTimingConfig() const waitStart = Date.now() let sessionID = task.sessionID while (!sessionID && Date.now() - waitStart < timing.WAIT_FOR_SESSION_TIMEOUT_MS) { if (ctx.abort?.aborted) { + cleanupReason = "Parent aborted while waiting for unstable task session start" return `Task aborted while waiting for session to start.\n\nTask ID: ${task.id}` } await new Promise(resolve => setTimeout(resolve, timing.WAIT_FOR_SESSION_INTERVAL_MS)) @@ -51,6 +56,7 @@ export async function executeUnstableAgentTask( sessionID = updated?.sessionID } if (!sessionID) { + cleanupReason = "Unstable task session start timed out before session became available" return formatDetailedError(new Error(`Task failed to start within timeout (30s). Task ID: ${task.id}, Status: ${task.status}`), { operation: "Launch monitored background task", args, @@ -88,6 +94,7 @@ export async function executeUnstableAgentTask( while (Date.now() - pollStart < (syncPollTimeoutMs ?? DEFAULT_SYNC_POLL_TIMEOUT_MS)) { if (ctx.abort?.aborted) { + cleanupReason = "Parent aborted while monitoring unstable background task" return `Task aborted (was running in background mode).\n\nSession ID: ${sessionID}` } @@ -148,6 +155,7 @@ session_id: ${sessionID} } if (!completedDuringMonitoring) { + cleanupReason = "Monitored unstable background task exceeded timeout budget" const duration = formatDuration(startTime) const timeoutBudgetMs = syncPollTimeoutMs ?? DEFAULT_SYNC_POLL_TIMEOUT_MS return `SUPERVISED TASK TIMED OUT @@ -221,5 +229,9 @@ session_id: ${sessionID} agent: agentToUse, category: args.category, }) + } finally { + if (cleanupReason) { + await cancelUnstableAgentTask(manager, launchedTaskID, cleanupReason) + } } }