From 370ec0071023bed1bb39da0baad38c62df6bb3aa Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 27 Oct 2026 07:33:56 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=A4=96=20refactor:=20move=20task=20waiter?= =?UTF-8?q?=20state=20into=20TaskWaiters?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move-only. TaskWaiters (taskWaiters.ts) owns the report waiters, start waiters, unpersisted launch-failure marks, foreground-await counts, backgroundable foreground waiters, user-backgrounded marks and the bounded completed-report cache, plus the methods that only touch them. ForegroundWaitBackgroundedError moves with it and taskService.ts re-exports it. TaskService keeps its public and interface methods as one-line facades, keeps resolveWaiters (ancestry, then cache, then drain) and rejectWaiters (spied on in tests) as entry points, and reads the cache and waiter lists through accessors. Deps are call-time closures (scheduleNotifyOnTerminalPersist, workspaceService.hasQueuedMessages). Test casts now point at `waiters.`. --- .../taskService.agentReports.waits.test.ts | 18 +- .../taskService.agentReports.wakes.test.ts | 4 +- .../services/taskService.settlement.test.ts | 4 +- .../taskService.settlementReceipts.test.ts | 4 +- src/node/services/taskService.ts | 323 +++-------------- src/node/services/taskWaiters.test.ts | 284 +++++++++++++++ src/node/services/taskWaiters.ts | 335 ++++++++++++++++++ .../workspaceService.removalCascade.test.ts | 5 +- 8 files changed, 698 insertions(+), 279 deletions(-) create mode 100644 src/node/services/taskWaiters.test.ts create mode 100644 src/node/services/taskWaiters.ts diff --git a/src/node/services/taskService.agentReports.waits.test.ts b/src/node/services/taskService.agentReports.waits.test.ts index c43a6ef9b67..62bca821fdb 100644 --- a/src/node/services/taskService.agentReports.waits.test.ts +++ b/src/node/services/taskService.agentReports.waits.test.ts @@ -112,11 +112,15 @@ describe("TaskService", () => { const hasQueuedMessages = mock(() => true); const { workspaceService } = createWorkspaceServiceMocks({ hasQueuedMessages }); const { taskService } = createTaskServiceHarness(config, { workspaceService }); - const internal = taskService as unknown as { - backgroundableForegroundWaitersByWorkspaceId: Map>; - pendingStartWaitersByTaskId: Map; - pendingWaitersByTaskId: Map; - }; + const internal = ( + taskService as unknown as { + waiters: { + backgroundableForegroundWaitersByWorkspaceId: Map>; + pendingStartWaitersByTaskId: Map; + pendingWaitersByTaskId: Map; + }; + } + ).waiters; const waitError = await taskService .waitForAgentReport(childId, { @@ -754,8 +758,8 @@ describe("TaskService", () => { // Simulate process restart / eviction. ( - taskService as unknown as { completedReportsByTaskId: Map } - ).completedReportsByTaskId.clear(); + taskService as unknown as { waiters: { completedReportsByTaskId: Map } } + ).waiters.completedReportsByTaskId.clear(); const report = await taskService.waitForAgentReport(childId, { timeoutMs: 10, diff --git a/src/node/services/taskService.agentReports.wakes.test.ts b/src/node/services/taskService.agentReports.wakes.test.ts index 097fcefb91c..1856169cfc3 100644 --- a/src/node/services/taskService.agentReports.wakes.test.ts +++ b/src/node/services/taskService.agentReports.wakes.test.ts @@ -609,7 +609,7 @@ describe("TaskService", () => { }); const internal = taskService as unknown as { - completedReportsByTaskId: Map; + waiters: { completedReportsByTaskId: Map }; }; await streamEnd(taskService, { @@ -644,7 +644,7 @@ describe("TaskService", () => { expect(ws?.taskStatus).toBe("reported"); // Validate report persistence path (not just in-memory cache). - internal.completedReportsByTaskId.clear(); + internal.waiters.completedReportsByTaskId.clear(); const persisted = await taskService.waitForAgentReport(childId, { timeoutMs: 10_000, requestingWorkspaceId: parentId, diff --git a/src/node/services/taskService.settlement.test.ts b/src/node/services/taskService.settlement.test.ts index ea032371485..78084e0a646 100644 --- a/src/node/services/taskService.settlement.test.ts +++ b/src/node/services/taskService.settlement.test.ts @@ -2220,7 +2220,7 @@ describe("TaskService", () => { interface AttemptLedgerInternals { ownedAttemptByTaskId: Map; attemptSettlementByTaskId: Map; - completedReportsByTaskId: Map; + waiters: { completedReportsByTaskId: Map }; shouldAllowLegacyInvalidWorkflowOutputSchema: (...args: unknown[]) => Promise; } const ledgerInternals = (taskService: TaskService) => @@ -2276,7 +2276,7 @@ describe("TaskService", () => { const reported = { kind: "reported", report: { reportMarkdown: "final answer" } }; expect(await taskService.readAttemptOutcome(taskId, requesting)).toMatchObject(reported); // Not just the in-memory cache: the persisted artifact answers, here and in a fresh process. - internals.completedReportsByTaskId.clear(); + internals.waiters.completedReportsByTaskId.clear(); expect(await taskService.readAttemptOutcome(taskId, requesting)).toMatchObject(reported); const { taskService: freshService } = createTaskServiceHarness(config); expect(await freshService.readAttemptOutcome(taskId, requesting)).toMatchObject(reported); diff --git a/src/node/services/taskService.settlementReceipts.test.ts b/src/node/services/taskService.settlementReceipts.test.ts index a79477a1e52..04fcfefac3a 100644 --- a/src/node/services/taskService.settlementReceipts.test.ts +++ b/src/node/services/taskService.settlementReceipts.test.ts @@ -865,8 +865,8 @@ describe("TaskService settlement receipt producers (G2)", () => { await ownEligibleAttempt(config, taskService, taskId); const pendingWaiters = () => ( - svc as unknown as { pendingWaitersByTaskId: Map } - ).pendingWaitersByTaskId.get(taskId)?.length ?? 0; + svc as unknown as { waiters: { pendingWaitersByTaskId: Map } } + ).waiters.pendingWaitersByTaskId.get(taskId)?.length ?? 0; const settledOf = (promise: Promise) => { const state: { settled?: "resolved" | "rejected"; error?: unknown } = {}; promise.then( diff --git a/src/node/services/taskService.ts b/src/node/services/taskService.ts index 7e4b84f0105..5b1713bce45 100644 --- a/src/node/services/taskService.ts +++ b/src/node/services/taskService.ts @@ -56,6 +56,7 @@ import { type WorkspaceTurnHost, } from "@/node/services/taskWorkspaceSeam"; export type { TaskKind } from "@/node/services/taskWorkspaceSeam"; +export { ForegroundWaitBackgroundedError } from "@/node/services/taskWaiters"; import { ACTIVE_AGENT_TASK_STATUSES, buildAgentTaskIndex, @@ -111,6 +112,11 @@ import { } from "@/constants/taskMessages"; import { log } from "@/node/services/log"; import { PeerSendParking } from "@/node/services/peerSendParking"; +import { + TaskWaiters, + type PendingTaskStartWaiter, + type PendingTaskWaiter, +} from "@/node/services/taskWaiters"; import { isProjectTrusted, isWorkspaceProjectTrustedIn } from "@/node/utils/projectTrust"; import { eventSpine } from "@/node/services/events/eventSpine"; import { sandboxHostService } from "@/node/services/sandbox/sandboxHostService"; @@ -1331,8 +1337,6 @@ function isWorkspaceBusyIdleOnlySend(error: unknown): boolean { const REMOVED_AGENT_TASKS_DIR = "removed-agent-tasks"; -const COMPLETED_REPORT_CACHE_MAX_ENTRIES = 128; - // Level-triggered backstop for workflow terminal wakes: the sweep re-derives owed wakes from // durable state (run records + settled markers), so lost terminal callbacks, deferred // (transiently unreadable) evaluations, and crashes all recover here without any per-failure @@ -1630,37 +1634,6 @@ interface InactiveWorkflowTaskOwner { type InterruptedTaskStatusMutation = "interrupted" | "preserved-completed-report"; -interface PendingTaskWaiter extends BackgroundableForegroundWaiter { - resolve: (report: { - reportMarkdown: string; - title?: string; - structuredOutput?: unknown; - planFilePath?: string; - model?: string; - thinkingLevel?: ThinkingLevel; - }) => void; -} - -interface PendingTaskStartWaiter { - start: () => void; - cleanup: () => void; -} - -interface CompletedAgentReportCacheEntry { - reportMarkdown: string; - planFilePath?: string; - structuredOutput?: unknown; - title?: string; - // Final settings the child reported with (post plan-to-exec handoff), not the launch snapshot. - model?: string; - thinkingLevel?: ThinkingLevel; - // Ancestor workspace IDs captured when the report was cached. - // Used to keep descendant-scope checks working even if the task workspace is cleaned up. - ancestorWorkspaceIds: string[]; - // Ancestors for which the task report must only be consumed through a workflow run. - workflowOwnedAncestorWorkspaceIds?: string[]; -} - interface ParentAutoResumeHint { agentId?: string; } @@ -1837,13 +1810,6 @@ async function readTaskBaseCommitShaByProjectPath(params: { return taskBaseCommitShaByProjectPath; } -export class ForegroundWaitBackgroundedError extends Error { - constructor() { - super("Foreground wait sent to background due to queued message"); - this.name = "ForegroundWaitBackgroundedError"; - } -} - function buildWorkflowTimeoutFinalizationPrompt( finalInstructions: string | undefined, completionKind: "final_response" | "propose_plan", @@ -1979,28 +1945,10 @@ export class TaskService implements AgentTaskIntegration { private workflowAttentionSweepTimer: ReturnType | null = null; /** The pending drain retry after a refused launch lease (QUEUED_LAUNCH_LEASE_RETRY_MS). */ private queuedLaunchLeaseRetryTimer: ReturnType | undefined; - private readonly pendingWaitersByTaskId = new Map(); - private readonly pendingStartWaitersByTaskId = new Map(); - // Tasks whose launch failed, or whose reservation was canceled or failed, but whose write - // ending the row rejected (#4747, #5028): the row still says queued/starting and nothing will - // move it to running, so waiters arm their timeout at once. A later move to running, a - // recorded launch failure or the workspace's removal clears the mark. - private readonly unpersistedLaunchFailureTaskIds = new Set(); - // Tracks workspaces currently blocked in a foreground wait (e.g. a task tool call awaiting - // agent_report). Used to avoid scheduler deadlocks when maxParallelAgentTasks is low and tasks - // spawn nested tasks in the foreground. - private readonly foregroundAwaitCountByWorkspaceId = new Map(); - private readonly backgroundableForegroundWaitersByWorkspaceId = new Map< - string, - Set - >(); + /** Report/start waiters, foreground-await state and the completed-report cache. */ + private readonly waiters: TaskWaiters; private readonly terminalAttentionStore: TerminalAttentionStore; - private readonly userBackgroundedTaskIds = new Set(); - - // Cache completed reports so callers can retrieve them without re-reading disk. - // Bounded by max entries; disk persistence is the source of truth for restart-safety. - private readonly completedReportsByTaskId = new Map(); // Task workspace removals that outlived their termination timeout. Retries must // await the ORIGINAL removal outcome: the host's remove() short-circuits Ok @@ -4263,7 +4211,7 @@ export class TaskService implements AgentTaskIntegration { kind: "indeterminate", reason, }); - const cached = this.completedReportsByTaskId.get(taskId); + const cached = this.waiters.getCompletedReport(taskId); if (cached) { return { kind: "reported", report: toAgentTaskReport(cached) }; } @@ -4282,7 +4230,7 @@ export class TaskService implements AgentTaskIntegration { taskId ); if (read.kind === "found") { - this.completedReportsByTaskId.set(taskId, { + this.waiters.cacheCompletedReport(taskId, { reportMarkdown: read.artifact.reportMarkdown, title: read.artifact.title, planFilePath: read.artifact.planFilePath, @@ -4292,7 +4240,6 @@ export class TaskService implements AgentTaskIntegration { workflowOwnedAncestorWorkspaceIds: read.artifact.workflowOwnedAncestorWorkspaceIds, ancestorWorkspaceIds: read.artifact.ancestorWorkspaceIds, }); - this.enforceCompletedReportCacheLimit(); return { kind: "reported", report: toAgentTaskReport(read.artifact) }; } if (read.kind === "unreadable") { @@ -4677,16 +4624,8 @@ export class TaskService implements AgentTaskIntegration { return true; } - private markTaskQueueBackgrounded(taskId: string): void { - this.userBackgroundedTaskIds.add(taskId); - } - markTaskForegroundRelevant(taskId: string): void { - this.userBackgroundedTaskIds.delete(taskId); - } - - private isTaskQueueBackgrounded(taskId: string): boolean { - return this.userBackgroundedTaskIds.has(taskId); + this.waiters.markTaskForegroundRelevant(taskId); } /** @@ -4730,6 +4669,12 @@ export class TaskService implements AgentTaskIntegration { workspaceService.countQueuedAgentPeerMessages(targetId) + this.peerSendParking.countParked(targetId), }); + this.waiters = new TaskWaiters({ + scheduleNotifyOnTerminalPersist: (taskId, ownerWorkspaceId) => + this.scheduleNotifyOnTerminalPersist(taskId, ownerWorkspaceId), + hasQueuedMessages: (workspaceId, dispatchMode) => + this.workspaceService.hasQueuedMessages(workspaceId, dispatchMode), + }); this.terminalAttentionStore = terminalAttentionStore ?? new TerminalAttentionStore(config); this.workflowRunAttention = new WorkflowRunAttention({ sessionsDir: () => this.config.sessionsDir, @@ -6819,7 +6764,7 @@ export class TaskService implements AgentTaskIntegration { } catch (error: unknown) { // Nothing launches this row now (#5028), unless the row already moved on (another // writer's attempt, whose waiters this cancel must not touch). - if (!rowMoved) this.armWaitersOfUnendedTask(plan.taskId); + if (!rowMoved) this.waiters.armWaitersOfUnendedTask(plan.taskId); fenceFailure ??= { error }; continue; } @@ -7062,7 +7007,7 @@ export class TaskService implements AgentTaskIntegration { error: fenceError, }); // Nothing launches this row now (#5028), unless it already moved on (see above). - if (!rowMoved) this.armWaitersOfUnendedTask(plan.taskId); + if (!rowMoved) this.waiters.armWaitersOfUnendedTask(plan.taskId); continue; } } @@ -7558,10 +7503,10 @@ export class TaskService implements AgentTaskIntegration { // unknown when the updater never ran. But a waiter that attached while the task was // queued/starting arms its timeout only on the move to running, which this failed launch // never makes; arm it now (and for later waiters) so they settle at their timeout. - this.armWaitersOfUnendedTask(taskId); + this.waiters.armWaitersOfUnendedTask(taskId); throw error; } - this.unpersistedLaunchFailureTaskIds.delete(taskId); + this.waiters.clearUnpersistedLaunchFailure(taskId); if (superseded) { log.info("Task launch failure not recorded: the record was re-admitted by another writer", { taskId, @@ -11257,7 +11202,7 @@ export class TaskService implements AgentTaskIntegration { for (const id of toTerminate) { this.beginWorkspaceStop(id); latched.push(id); - this.completedReportsByTaskId.delete(id); + this.waiters.deleteCompletedReport(id); this.rejectWaiters(id, terminationError); // Durable stop marker before removal: a workspace whose removal later fails or times out // must not read as a live running task once its latch settles. An entry already gone has @@ -11896,85 +11841,30 @@ export class TaskService implements AgentTaskIntegration { } isForegroundAwaiting(workspaceId: string): boolean { - const count = this.foregroundAwaitCountByWorkspaceId.get(workspaceId); - return typeof count === "number" && count > 0; + return this.waiters.isForegroundAwaiting(workspaceId); } startForegroundAwait(workspaceId: string): () => void { - assert(workspaceId.length > 0, "startForegroundAwait: workspaceId must be non-empty"); - - const current = this.foregroundAwaitCountByWorkspaceId.get(workspaceId) ?? 0; - assert( - Number.isInteger(current) && current >= 0, - "startForegroundAwait: expected non-negative integer counter" - ); - - this.foregroundAwaitCountByWorkspaceId.set(workspaceId, current + 1); - - return () => { - const current = this.foregroundAwaitCountByWorkspaceId.get(workspaceId) ?? 0; - assert( - Number.isInteger(current) && current > 0, - "startForegroundAwait cleanup: expected positive integer counter" - ); - if (current <= 1) { - this.foregroundAwaitCountByWorkspaceId.delete(workspaceId); - } else { - this.foregroundAwaitCountByWorkspaceId.set(workspaceId, current - 1); - } - }; + return this.waiters.startForegroundAwait(workspaceId); } registerBackgroundableForegroundWaiter( workspaceId: string, waiter: BackgroundableForegroundWaiter ): void { - let set = this.backgroundableForegroundWaitersByWorkspaceId.get(workspaceId); - if (!set) { - set = new Set(); - this.backgroundableForegroundWaitersByWorkspaceId.set(workspaceId, set); - } - set.add(waiter); + this.waiters.registerBackgroundableForegroundWaiter(workspaceId, waiter); } unregisterBackgroundableForegroundWaiter( workspaceId: string, waiter: BackgroundableForegroundWaiter ): void { - const set = this.backgroundableForegroundWaitersByWorkspaceId.get(workspaceId); - if (!set) return; - set.delete(waiter); - if (set.size === 0) { - this.backgroundableForegroundWaitersByWorkspaceId.delete(workspaceId); - } + this.waiters.unregisterBackgroundableForegroundWaiter(workspaceId, waiter); } - /** - * Reject all foreground task waiters for a workspace that opted into backgrounding - * when a new message is queued. Returns the number of waiters signaled. - * Safe to call repeatedly — already-cleaned-up waiters are skipped. - */ + /** See TaskWaiters.backgroundForegroundWaitsForWorkspace. */ backgroundForegroundWaitsForWorkspace(workspaceId: string): number { - const set = this.backgroundableForegroundWaitersByWorkspaceId.get(workspaceId); - if (!set || set.size === 0) return 0; - - const waiters = [...set]; - let count = 0; - for (const waiter of waiters) { - try { - this.markTaskQueueBackgrounded(waiter.taskId); - // A foreground wait detached by a queued message becomes durably non-blocking: - // persist notify_on_terminal so future stream-ends and restarts do not re-force the - // await. The in-memory mark above covers the immediate next stream-end while this - // persistence settles. Tracked so handleStreamEnd can await it before reading config. - this.scheduleNotifyOnTerminalPersist(waiter.taskId, waiter.requestingWorkspaceId); - waiter.reject(new ForegroundWaitBackgroundedError()); - count++; - } catch { - // waiter already resolved/rejected — ignore - } - } - return count; + return this.waiters.backgroundForegroundWaitsForWorkspace(workspaceId); } /** @@ -13688,24 +13578,15 @@ export class TaskService implements AgentTaskIntegration { await markPendingDelivered(); } - /** - * Background any registered foreground waits for the requesting workspace when a - * tool-end message is already queued. Shared by both wait-registration paths - * (workspace-turn and task await): the auto-backgrounding signal is edge-triggered - * on enqueue, so a message queued before the waiter registered must be re-checked - * here. No-op when backgrounding is disabled or no requesting workspace is set. - */ + /** See TaskWaiters.backgroundForegroundWaitIfQueued. */ backgroundForegroundWaitIfQueued( shouldBackgroundOnQueuedMessage: boolean, requestingWorkspaceId: string | undefined ): void { - if ( - shouldBackgroundOnQueuedMessage && - requestingWorkspaceId && - this.workspaceService.hasQueuedMessages(requestingWorkspaceId, "tool-end") - ) { - this.backgroundForegroundWaitsForWorkspace(requestingWorkspaceId); - } + this.waiters.backgroundForegroundWaitIfQueued( + shouldBackgroundOnQueuedMessage, + requestingWorkspaceId + ); } async reportAgentProgress( @@ -14096,7 +13977,7 @@ export class TaskService implements AgentTaskIntegration { if (!entry?.workspace.parentWorkspaceId) { return { status: "not_active" as const }; } - if (hasCompletedAgentReport(entry.workspace) || this.completedReportsByTaskId.has(taskId)) { + if (hasCompletedAgentReport(entry.workspace) || this.waiters.hasCompletedReport(taskId)) { return { status: "already_reported" as const }; } if (entry.workspace.taskStatus === "interrupted" && !this.aiService.isStreaming(taskId)) { @@ -14136,10 +14017,7 @@ export class TaskService implements AgentTaskIntegration { if (!freshEntry?.workspace.parentWorkspaceId) { return "not_active"; } - if ( - hasCompletedAgentReport(freshEntry.workspace) || - this.completedReportsByTaskId.has(taskId) - ) { + if (hasCompletedAgentReport(freshEntry.workspace) || this.waiters.hasCompletedReport(taskId)) { return "already_reported"; } let finalizationAccepted = false; @@ -14150,7 +14028,7 @@ export class TaskService implements AgentTaskIntegration { if (!entry?.workspace.parentWorkspaceId) { return; } - if (hasCompletedAgentReport(entry.workspace) || this.completedReportsByTaskId.has(taskId)) { + if (hasCompletedAgentReport(entry.workspace) || this.waiters.hasCompletedReport(taskId)) { return; } await this.editActiveWorkspaceEntry( @@ -14243,7 +14121,7 @@ export class TaskService implements AgentTaskIntegration { if (!entry?.workspace.parentWorkspaceId) { return; } - if (hasCompletedAgentReport(entry.workspace) || this.completedReportsByTaskId.has(taskId)) { + if (hasCompletedAgentReport(entry.workspace) || this.waiters.hasCompletedReport(taskId)) { return; } try { @@ -14290,7 +14168,7 @@ export class TaskService implements AgentTaskIntegration { // Report monotonicity invariant: check the in-memory cache before any status-based // interruption handling so a finalized report stays awaitable once observed. - const cached = this.completedReportsByTaskId.get(taskId); + const cached = this.waiters.getCompletedReport(taskId); if (cached) { return { reportMarkdown: cached.reportMarkdown, @@ -14462,15 +14340,7 @@ export class TaskService implements AgentTaskIntegration { this.unregisterBackgroundableForegroundWaiter(entry.requestingWorkspaceId, entry); } - const current = this.pendingWaitersByTaskId.get(taskId); - if (current) { - const next = current.filter((w) => w !== entry); - if (next.length === 0) { - this.pendingWaitersByTaskId.delete(taskId); - } else { - this.pendingWaitersByTaskId.set(taskId, next); - } - } + this.waiters.removePendingWaiter(taskId, entry); cleanupStartWaiter(); @@ -14494,9 +14364,7 @@ export class TaskService implements AgentTaskIntegration { }, }; - const list = this.pendingWaitersByTaskId.get(taskId) ?? []; - list.push(entry); - this.pendingWaitersByTaskId.set(taskId, list); + this.waiters.addPendingWaiter(taskId, entry); const shouldBackgroundOnQueuedMessage = Boolean( requestingWorkspaceId && (options?.backgroundOnMessageQueued ?? true) @@ -14521,22 +14389,12 @@ export class TaskService implements AgentTaskIntegration { const startWaiterEntry: PendingTaskStartWaiter = { start: startReportTimeout, cleanup: () => { - const currentStartWaiters = this.pendingStartWaitersByTaskId.get(taskId); - if (currentStartWaiters) { - const next = currentStartWaiters.filter((w) => w !== startWaiterEntry); - if (next.length === 0) { - this.pendingStartWaitersByTaskId.delete(taskId); - } else { - this.pendingStartWaitersByTaskId.set(taskId, next); - } - } + this.waiters.removePendingStartWaiter(taskId, startWaiterEntry); }, }; startWaiter = startWaiterEntry; - const currentStartWaiters = this.pendingStartWaitersByTaskId.get(taskId) ?? []; - currentStartWaiters.push(startWaiterEntry); - this.pendingStartWaitersByTaskId.set(taskId, currentStartWaiters); + this.waiters.addPendingStartWaiter(taskId, startWaiterEntry); // Close the race where the task starts between the initial config read and registering the waiter. // A launch whose failure could not be written never starts: arm the timeout now (#4747). @@ -14545,7 +14403,7 @@ export class TaskService implements AgentTaskIntegration { if ( (afterEntry?.workspace.taskStatus !== "queued" && afterEntry?.workspace.taskStatus !== "starting") || - this.unpersistedLaunchFailureTaskIds.has(taskId) + this.waiters.isLaunchFailureUnpersisted(taskId) ) { cleanupStartWaiter(); startReportTimeout(); @@ -15413,7 +15271,7 @@ export class TaskService implements AgentTaskIntegration { continue; } - const cached = this.completedReportsByTaskId.get(taskId); + const cached = this.waiters.getCompletedReport(taskId); if (hasAncestorWorkspaceId(cached, ancestorWorkspaceId)) { result.push(taskId); continue; @@ -15514,7 +15372,7 @@ export class TaskService implements AgentTaskIntegration { return indexResult; } - const cached = this.completedReportsByTaskId.get(taskId); + const cached = this.waiters.getCompletedReport(taskId); if (hasWorkflowOwnedAncestorWorkspaceId(cached, ancestorWorkspaceId)) { return true; } @@ -15579,7 +15437,7 @@ export class TaskService implements AgentTaskIntegration { // The task workspace may have been removed after it settled (cleanup/restart). Preserve scope // checks by consulting persisted report AND failure artifacts in the ancestor session dir — // a terminally-failed child must stay awaitable so its typed failure can be surfaced. - const cached = this.completedReportsByTaskId.get(taskId); + const cached = this.waiters.getCompletedReport(taskId); if (hasAncestorWorkspaceId(cached, ancestorWorkspaceId)) { return true; } @@ -16446,7 +16304,7 @@ export class TaskService implements AgentTaskIntegration { const workflowOwned = next.workflowOwned || entry?.workflowTask != null; const nonBlockingSubtree = this.resolveAgentTaskAttentionPolicy(next.taskId, index) === "notify_on_terminal" || - this.isTaskQueueBackgrounded(next.taskId); + this.waiters.isTaskQueueBackgrounded(next.taskId); if ( !nonBlockingSubtree && entry != null && @@ -16971,25 +16829,14 @@ export class TaskService implements AgentTaskIntegration { await this.emitWorkspaceMetadata(workspaceId); if (status === "running") { - this.unpersistedLaunchFailureTaskIds.delete(workspaceId); - this.startPendingStartWaiters(workspaceId); + this.waiters.clearUnpersistedLaunchFailure(workspaceId); + this.waiters.startPendingStartWaiters(workspaceId); } return true; } - /** - * A write that would have ended a queued/starting task rejected, and nothing will move the row - * to running (#4747, #5028). Fail closed: waiters are not rejected with the unpersisted outcome, - * whose owner is unknown when the updater never ran. Their timeout is armed instead, now and for - * waiters that attach later, so they settle at their timeout. - */ - private armWaitersOfUnendedTask(taskId: string): void { - this.unpersistedLaunchFailureTaskIds.add(taskId); - this.startPendingStartWaiters(taskId); - } - noteWorkspaceRemoved(workspaceId: string): void { - this.unpersistedLaunchFailureTaskIds.delete(workspaceId); + this.waiters.clearUnpersistedLaunchFailure(workspaceId); // #5334: the stop generations otherwise grow with every workspace ever stopped. Dropping // them resets the generation to 0, which keeps the latch sound only because this runs after // the workspace left the config: peer-send admission rereads both endpoints from the config @@ -17004,20 +16851,6 @@ export class TaskService implements AgentTaskIntegration { this.workspaceGoalService?.forgetStreamAccountingReceipts(workspaceId); } - /** Arms the report timeout of every waiter that attached while the task was queued/starting. */ - private startPendingStartWaiters(workspaceId: string): void { - const waiters = this.pendingStartWaitersByTaskId.get(workspaceId); - if (!waiters || waiters.length === 0) return; - this.pendingStartWaitersByTaskId.delete(workspaceId); - for (const waiter of waiters) { - try { - waiter.start(); - } catch (error: unknown) { - log.error("Task start waiter callback failed", { workspaceId, error }); - } - } - } - /** * Reset interrupt + auto-resume state for a workspace (called when user sends a real message). */ @@ -17249,7 +17082,7 @@ export class TaskService implements AgentTaskIntegration { return lost(TASK_REAWAKEN_LOST_SEND_BLOCKED_MESSAGE); } // The row names a new attempt now: the earlier attempt's cached report is not its (#7003). - this.completedReportsByTaskId.delete(workspaceId); + this.waiters.deleteCompletedReport(workspaceId); // Writers outside this mutex (a reactivation under its own locks, another backend) may have // rotated the row again during the commit's own awaits: never publish or own an identity // the durable row no longer names — the successor's ownership and closures stay untouched. @@ -17840,7 +17673,7 @@ export class TaskService implements AgentTaskIntegration { }); let activeTaskIds = [...listBlockingDescendantTaskIds(), ...activeWorkspaceTurnIds]; const queueBackgroundedTaskIds = new Set( - activeTaskIds.filter((id) => this.isTaskQueueBackgrounded(id)) + activeTaskIds.filter((id) => this.waiters.isTaskQueueBackgrounded(id)) ); // Durable `notify_on_terminal` work is non-blocking: it never forces the parent to // task_await and is not consumed at stream-end. Agent-task policy is applied by @@ -19423,7 +19256,7 @@ export class TaskService implements AgentTaskIntegration { // Captured before settlement: rejectWaiters consumes the pending waiters // that prove a parent turn is actively listening for this task. - const hadForegroundWaiters = (this.pendingWaitersByTaskId.get(workspaceId)?.length ?? 0) > 0; + const hadForegroundWaiters = this.waiters.hasPendingWaiters(workspaceId); await this.settleInterruptedTaskAtStreamEnd( workspaceId, @@ -20951,7 +20784,7 @@ export class TaskService implements AgentTaskIntegration { } // Cache for the current process (best-effort). Disk is the source of truth. - this.completedReportsByTaskId.set(taskId, { + this.waiters.cacheCompletedReport(taskId, { reportMarkdown: artifact.reportMarkdown, title: artifact.title, planFilePath: artifact.planFilePath, @@ -20961,7 +20794,6 @@ export class TaskService implements AgentTaskIntegration { workflowOwnedAncestorWorkspaceIds: artifact.workflowOwnedAncestorWorkspaceIds, ancestorWorkspaceIds: artifact.ancestorWorkspaceIds, }); - this.enforceCompletedReportCacheLimit(); // After the committed `reported` (ours above, or the other backend's): a report bound to the // row's attempt settles that attempt's stream-end decision (#7003); the notification then @@ -21044,14 +20876,6 @@ export class TaskService implements AgentTaskIntegration { return !rowMoved; } - private enforceCompletedReportCacheLimit(): void { - while (this.completedReportsByTaskId.size > COMPLETED_REPORT_CACHE_MAX_ENTRIES) { - const first = this.completedReportsByTaskId.keys().next(); - if (first.done) break; - this.completedReportsByTaskId.delete(first.value); - } - } - private resolveWaiters( taskId: string, report: { @@ -21078,7 +20902,7 @@ export class TaskService implements AgentTaskIntegration { true ); - this.completedReportsByTaskId.set(taskId, { + this.waiters.cacheCompletedReport(taskId, { reportMarkdown: report.reportMarkdown, title: report.title, planFilePath: report.planFilePath, @@ -21088,41 +20912,13 @@ export class TaskService implements AgentTaskIntegration { ancestorWorkspaceIds, workflowOwnedAncestorWorkspaceIds, }); - this.enforceCompletedReportCacheLimit(); - const waiters = this.pendingWaitersByTaskId.get(taskId); - if (!waiters || waiters.length === 0) { - return false; - } - - this.pendingWaitersByTaskId.delete(taskId); - for (const waiter of waiters) { - try { - waiter.cleanup(); - waiter.resolve(report); - } catch { - // ignore - } - } - - return true; + return this.waiters.resolvePendingWaiters(taskId, report); } + // Kept on TaskService: tests spy on it, and every internal rejection goes through here. private rejectWaiters(taskId: string, error: Error): void { - this.markTaskForegroundRelevant(taskId); - - const waiters = this.pendingWaitersByTaskId.get(taskId); - if (!waiters || waiters.length === 0) { - return; - } - - for (const waiter of [...waiters]) { - try { - waiter.reject(error); - } catch (rejectError: unknown) { - log.error("Task waiter reject callback failed", { taskId, error: rejectError }); - } - } + this.waiters.rejectWaiters(taskId, error); } private findProposePlanSuccessInParts(parts: readonly unknown[]): { planPath: string } | null { @@ -21727,8 +21523,7 @@ export class TaskService implements AgentTaskIntegration { // If someone is actively awaiting this report (foreground task tool call or task_await), // skip injecting a synthetic history message to avoid duplicating the report in context. if (childWorkspaceId) { - const waiters = this.pendingWaitersByTaskId.get(childWorkspaceId); - if (waiters && waiters.length > 0) { + if (this.waiters.hasPendingWaiters(childWorkspaceId)) { return []; } } diff --git a/src/node/services/taskWaiters.test.ts b/src/node/services/taskWaiters.test.ts new file mode 100644 index 00000000000..3efa6ad96e9 --- /dev/null +++ b/src/node/services/taskWaiters.test.ts @@ -0,0 +1,284 @@ +import { afterAll, afterEach, describe, expect, mock, spyOn, test } from "bun:test"; + +import { log } from "@/node/services/log"; + +import { + COMPLETED_REPORT_CACHE_MAX_ENTRIES, + ForegroundWaitBackgroundedError, + TaskWaiters, + type CompletedAgentReportCacheEntry, + type PendingTaskStartWaiter, + type PendingTaskWaiter, +} from "@/node/services/taskWaiters"; + +function createHarness(options: { queued?: boolean } = {}) { + const calls: string[] = []; + const scheduleNotifyOnTerminalPersist = mock((taskId: string, owner: string | undefined) => { + calls.push(`persist:${taskId}:${owner ?? "-"}`); + }); + const hasQueuedMessages = mock(() => options.queued ?? false); + const waiters = new TaskWaiters({ scheduleNotifyOnTerminalPersist, hasQueuedMessages }); + return { waiters, calls, scheduleNotifyOnTerminalPersist, hasQueuedMessages }; +} + +function createWaiter( + taskId: string, + calls: string[], + options: { requestingWorkspaceId?: string; throwOnReject?: boolean } = {} +): PendingTaskWaiter { + return { + taskId, + requestingWorkspaceId: options.requestingWorkspaceId, + backgroundOnMessageQueued: options.requestingWorkspaceId != null, + resolve: (report) => calls.push(`resolve:${taskId}:${report.reportMarkdown}`), + reject: (error) => { + calls.push(`reject:${taskId}:${error.name}`); + if (options.throwOnReject) throw new Error("already settled"); + }, + cleanup: () => calls.push(`cleanup:${taskId}`), + }; +} + +function cacheEntry(reportMarkdown: string): CompletedAgentReportCacheEntry { + return { reportMarkdown, ancestorWorkspaceIds: [] }; +} + +describe("TaskWaiters", () => { + const logError = spyOn(log, "error").mockImplementation(() => undefined); + afterEach(() => logError.mockClear()); + afterAll(() => logError.mockRestore()); + + test("foreground await counts nest and each release decrements once", () => { + const { waiters } = createHarness(); + expect(waiters.isForegroundAwaiting("parent")).toBe(false); + + const releaseA = waiters.startForegroundAwait("parent"); + const releaseB = waiters.startForegroundAwait("parent"); + releaseA(); + expect(waiters.isForegroundAwaiting("parent")).toBe(true); + releaseB(); + expect(waiters.isForegroundAwaiting("parent")).toBe(false); + // A release past zero is an impossible state and must crash, not go negative. + expect(() => releaseB()).toThrow(); + }); + + test("backgrounding marks, schedules persistence, then rejects each registered waiter", () => { + const { waiters, calls } = createHarness(); + const first = createWaiter("t1", calls, { requestingWorkspaceId: "parent" }); + // A waiter that already settled throws on reject: it is skipped and not counted. + const settled = createWaiter("t2", calls, { + requestingWorkspaceId: "parent", + throwOnReject: true, + }); + waiters.registerBackgroundableForegroundWaiter("parent", first); + waiters.registerBackgroundableForegroundWaiter("parent", settled); + + expect(waiters.backgroundForegroundWaitsForWorkspace("parent")).toBe(1); + expect(calls).toEqual([ + "persist:t1:parent", + `reject:t1:${new ForegroundWaitBackgroundedError().name}`, + "persist:t2:parent", + `reject:t2:${new ForegroundWaitBackgroundedError().name}`, + ]); + expect(waiters.isTaskQueueBackgrounded("t1")).toBe(true); + expect(waiters.isTaskQueueBackgrounded("t2")).toBe(true); + expect(waiters.backgroundForegroundWaitsForWorkspace("other")).toBe(0); + }); + + test("backgrounding iterates a copy so waiters may unregister while being rejected", () => { + const { waiters, calls } = createHarness(); + const register = (taskId: string) => { + const waiter = createWaiter(taskId, calls, { requestingWorkspaceId: "parent" }); + waiter.reject = () => { + calls.push(`reject:${taskId}`); + waiters.unregisterBackgroundableForegroundWaiter("parent", waiter); + }; + waiters.registerBackgroundableForegroundWaiter("parent", waiter); + }; + register("t1"); + register("t2"); + + expect(waiters.backgroundForegroundWaitsForWorkspace("parent")).toBe(2); + expect(calls.filter((call) => call.startsWith("reject:"))).toEqual(["reject:t1", "reject:t2"]); + expect(waiters.backgroundableForegroundWaitersByWorkspaceId.has("parent")).toBe(false); + expect(waiters.backgroundForegroundWaitsForWorkspace("parent")).toBe(0); + }); + + test("backgroundForegroundWaitIfQueued only backgrounds when enabled, targeted and queued", () => { + const queued = createHarness({ queued: true }); + queued.waiters.registerBackgroundableForegroundWaiter( + "parent", + createWaiter("t1", queued.calls, { requestingWorkspaceId: "parent" }) + ); + queued.waiters.backgroundForegroundWaitIfQueued(false, "parent"); + queued.waiters.backgroundForegroundWaitIfQueued(true, undefined); + expect(queued.hasQueuedMessages).not.toHaveBeenCalled(); + expect(queued.waiters.isTaskQueueBackgrounded("t1")).toBe(false); + + queued.waiters.backgroundForegroundWaitIfQueued(true, "parent"); + expect(queued.hasQueuedMessages).toHaveBeenCalledWith("parent", "tool-end"); + expect(queued.waiters.isTaskQueueBackgrounded("t1")).toBe(true); + + const idle = createHarness({ queued: false }); + idle.waiters.registerBackgroundableForegroundWaiter( + "parent", + createWaiter("t1", idle.calls, { requestingWorkspaceId: "parent" }) + ); + idle.waiters.backgroundForegroundWaitIfQueued(true, "parent"); + expect(idle.scheduleNotifyOnTerminalPersist).not.toHaveBeenCalled(); + }); + + test("resolvePendingWaiters drains once, cleaning up before resolving", () => { + const { waiters, calls } = createHarness(); + const failing = createWaiter("t1", calls); + failing.resolve = () => { + throw new Error("resolve failed"); + }; + waiters.addPendingWaiter("t1", failing); + waiters.addPendingWaiter("t1", createWaiter("t1", calls)); + expect(waiters.hasPendingWaiters("t1")).toBe(true); + + expect(waiters.resolvePendingWaiters("t1", { reportMarkdown: "done" })).toBe(true); + // A throwing waiter does not stop the rest from resolving. + expect(calls).toEqual(["cleanup:t1", "cleanup:t1", "resolve:t1:done"]); + expect(waiters.hasPendingWaiters("t1")).toBe(false); + expect(waiters.resolvePendingWaiters("t1", { reportMarkdown: "again" })).toBe(false); + }); + + test("a single pending waiter is resolved or rejected too", () => { + const { waiters, calls } = createHarness(); + expect(waiters.hasPendingWaiters("t1")).toBe(false); + waiters.addPendingWaiter("t1", createWaiter("t1", calls)); + expect(waiters.hasPendingWaiters("t1")).toBe(true); + waiters.rejectWaiters("t1", new Error("stopped")); + expect(calls).toEqual(["reject:t1:Error"]); + + expect(waiters.resolvePendingWaiters("t1", { reportMarkdown: "done" })).toBe(true); + expect(calls).toEqual(["reject:t1:Error", "cleanup:t1", "resolve:t1:done"]); + }); + + test("unregistering the last backgroundable waiter drops the workspace entry", () => { + const { waiters, calls } = createHarness(); + const a = createWaiter("t1", calls, { requestingWorkspaceId: "parent" }); + const b = createWaiter("t2", calls, { requestingWorkspaceId: "parent" }); + waiters.registerBackgroundableForegroundWaiter("parent", a); + waiters.registerBackgroundableForegroundWaiter("parent", b); + + waiters.unregisterBackgroundableForegroundWaiter("parent", a); + expect([...(waiters.backgroundableForegroundWaitersByWorkspaceId.get("parent") ?? [])]).toEqual( + [b] + ); + waiters.unregisterBackgroundableForegroundWaiter("parent", b); + expect(waiters.backgroundableForegroundWaitersByWorkspaceId.has("parent")).toBe(false); + }); + + test("removePendingWaiter drops only that waiter", () => { + const { waiters, calls } = createHarness(); + const a = createWaiter("t1", calls); + const b = createWaiter("t1", calls); + waiters.addPendingWaiter("t1", a); + waiters.addPendingWaiter("t1", b); + + waiters.removePendingWaiter("t1", a); + expect(waiters.pendingWaitersByTaskId.get("t1")).toEqual([b]); + waiters.removePendingWaiter("t1", b); + expect(waiters.pendingWaitersByTaskId.has("t1")).toBe(false); + waiters.removePendingWaiter("t1", b); + expect(waiters.pendingWaitersByTaskId.has("t1")).toBe(false); + }); + + test("rejectWaiters marks the task foreground-relevant and keeps waiters until they clean up", () => { + const { waiters, calls } = createHarness(); + waiters.registerBackgroundableForegroundWaiter( + "parent", + createWaiter("t1", calls, { requestingWorkspaceId: "parent" }) + ); + waiters.backgroundForegroundWaitsForWorkspace("parent"); + expect(waiters.isTaskQueueBackgrounded("t1")).toBe(true); + + // No pending waiters: still clears the backgrounded mark. + waiters.rejectWaiters("t1", new Error("stopped")); + expect(waiters.isTaskQueueBackgrounded("t1")).toBe(false); + + calls.length = 0; + const selfRemoving = createWaiter("t1", calls, { throwOnReject: true }); + selfRemoving.reject = (error) => { + calls.push(`reject:self:${error.message}`); + waiters.removePendingWaiter("t1", selfRemoving); + throw new Error("late"); + }; + waiters.addPendingWaiter("t1", selfRemoving); + waiters.addPendingWaiter("t1", createWaiter("t1", calls)); + + waiters.rejectWaiters("t1", new Error("stopped")); + expect(calls).toEqual(["reject:self:stopped", "reject:t1:Error"]); + expect(logError).toHaveBeenCalledTimes(1); + // The class never drops the entry itself: each waiter's cleanup does. + expect(waiters.pendingWaitersByTaskId.get("t1")?.length).toBe(1); + }); + + test("start waiters run once, deleted before each start, and failures do not stop the rest", () => { + const { waiters } = createHarness(); + const order: string[] = []; + const add = (start: () => void) => { + waiters.addPendingStartWaiter("t1", { start, cleanup: () => undefined }); + }; + add(() => { + order.push(`a:${String(waiters.pendingStartWaitersByTaskId.has("t1"))}`); + throw new Error("boom"); + }); + add(() => order.push("b")); + + waiters.startPendingStartWaiters("t1"); + waiters.startPendingStartWaiters("t1"); + expect(order).toEqual(["a:false", "b"]); + expect(logError).toHaveBeenCalledTimes(1); + }); + + test("removePendingStartWaiter drops only that waiter", () => { + const { waiters } = createHarness(); + const a: PendingTaskStartWaiter = { start: () => undefined, cleanup: () => undefined }; + const b: PendingTaskStartWaiter = { start: () => undefined, cleanup: () => undefined }; + waiters.addPendingStartWaiter("t1", a); + waiters.addPendingStartWaiter("t1", b); + + waiters.removePendingStartWaiter("t1", a); + expect(waiters.pendingStartWaitersByTaskId.get("t1")).toEqual([b]); + waiters.removePendingStartWaiter("t1", b); + expect(waiters.pendingStartWaitersByTaskId.has("t1")).toBe(false); + waiters.removePendingStartWaiter("t1", b); + expect(waiters.pendingStartWaitersByTaskId.has("t1")).toBe(false); + }); + + test("armWaitersOfUnendedTask marks the task and starts its waiters", () => { + const { waiters } = createHarness(); + const started: boolean[] = []; + waiters.addPendingStartWaiter("t1", { + start: () => started.push(waiters.isLaunchFailureUnpersisted("t1")), + cleanup: () => undefined, + }); + + waiters.armWaitersOfUnendedTask("t1"); + // The mark is set before the waiters start, so a start observes it. + expect(started).toEqual([true]); + waiters.clearUnpersistedLaunchFailure("t1"); + expect(waiters.isLaunchFailureUnpersisted("t1")).toBe(false); + }); + + test("completed-report cache evicts the oldest entries past its bound", () => { + const { waiters } = createHarness(); + for (let i = 0; i <= COMPLETED_REPORT_CACHE_MAX_ENTRIES; i++) { + waiters.cacheCompletedReport(`t${i}`, cacheEntry(`r${i}`)); + } + + expect(waiters.completedReportsByTaskId.size).toBe(COMPLETED_REPORT_CACHE_MAX_ENTRIES); + expect(waiters.hasCompletedReport("t0")).toBe(false); + expect(waiters.getCompletedReport("t1")?.reportMarkdown).toBe("r1"); + expect( + waiters.getCompletedReport(`t${COMPLETED_REPORT_CACHE_MAX_ENTRIES}`)?.reportMarkdown + ).toBe(`r${COMPLETED_REPORT_CACHE_MAX_ENTRIES}`); + + waiters.deleteCompletedReport("t1"); + expect(waiters.hasCompletedReport("t1")).toBe(false); + }); +}); diff --git a/src/node/services/taskWaiters.ts b/src/node/services/taskWaiters.ts new file mode 100644 index 00000000000..4bd135d94d4 --- /dev/null +++ b/src/node/services/taskWaiters.ts @@ -0,0 +1,335 @@ +import assert from "node:assert/strict"; +import type { ThinkingLevel } from "@/common/types/thinking"; +import { log } from "@/node/services/log"; +import type { BackgroundableForegroundWaiter } from "@/node/services/taskWorkspaceSeam"; + +export const COMPLETED_REPORT_CACHE_MAX_ENTRIES = 128; + +export class ForegroundWaitBackgroundedError extends Error { + constructor() { + super("Foreground wait sent to background due to queued message"); + this.name = "ForegroundWaitBackgroundedError"; + } +} + +/** The report a pending waiter resolves with. */ +export interface PendingTaskWaiterReport { + reportMarkdown: string; + title?: string; + structuredOutput?: unknown; + planFilePath?: string; + model?: string; + thinkingLevel?: ThinkingLevel; +} + +export interface PendingTaskWaiter extends BackgroundableForegroundWaiter { + resolve: (report: PendingTaskWaiterReport) => void; +} + +export interface PendingTaskStartWaiter { + start: () => void; + cleanup: () => void; +} + +export interface CompletedAgentReportCacheEntry { + reportMarkdown: string; + planFilePath?: string; + structuredOutput?: unknown; + title?: string; + // Final settings the child reported with (post plan-to-exec handoff), not the launch snapshot. + model?: string; + thinkingLevel?: ThinkingLevel; + // Ancestor workspace IDs captured when the report was cached. + // Used to keep descendant-scope checks working even if the task workspace is cleaned up. + ancestorWorkspaceIds: string[]; + // Ancestors for which the task report must only be consumed through a workflow run. + workflowOwnedAncestorWorkspaceIds?: string[]; +} + +// Report waiters and start waiters share one list-per-task bookkeeping: append in arrival order, +// and drop the key once its last waiter leaves so `has(taskId)` means "someone is waiting". +function appendToListMap(map: Map, key: string, item: T): void { + const list = map.get(key) ?? []; + list.push(item); + map.set(key, list); +} + +function removeFromListMap(map: Map, key: string, item: T): void { + const current = map.get(key); + if (current) { + const next = current.filter((w) => w !== item); + if (next.length === 0) { + map.delete(key); + } else { + map.set(key, next); + } + } +} + +export interface TaskWaitersDeps { + /** Durably mark a backgrounded wait's handle `notify_on_terminal` (TaskService owns persistence). */ + scheduleNotifyOnTerminalPersist: (taskId: string, ownerWorkspaceId: string | undefined) => void; + /** Whether the workspace has a queued message with the given dispatch mode. */ + hasQueuedMessages: (workspaceId: string, dispatchMode: "tool-end") => boolean; +} + +/** + * In-memory waiter state of TaskService: report waiters, start waiters (timeouts armed once a + * queued task runs), foreground-await counts, backgroundable foreground waiters, the + * user-backgrounded marks and the bounded completed-report cache. Nothing here is persisted. + */ +export class TaskWaiters { + readonly pendingWaitersByTaskId = new Map(); + readonly pendingStartWaitersByTaskId = new Map(); + // Tasks whose launch failed, or whose reservation was canceled or failed, but whose write + // ending the row rejected (#4747, #5028): the row still says queued/starting and nothing will + // move it to running, so waiters arm their timeout at once. A later move to running, a + // recorded launch failure or the workspace's removal clears the mark. + readonly unpersistedLaunchFailureTaskIds = new Set(); + // Tracks workspaces currently blocked in a foreground wait (e.g. a task tool call awaiting + // agent_report). Used to avoid scheduler deadlocks when maxParallelAgentTasks is low and tasks + // spawn nested tasks in the foreground. + private readonly foregroundAwaitCountByWorkspaceId = new Map(); + readonly backgroundableForegroundWaitersByWorkspaceId = new Map< + string, + Set + >(); + + private readonly userBackgroundedTaskIds = new Set(); + + // Cache completed reports so callers can retrieve them without re-reading disk. + // Bounded by max entries; disk persistence is the source of truth for restart-safety. + readonly completedReportsByTaskId = new Map(); + + constructor(private readonly deps: TaskWaitersDeps) {} + + private markTaskQueueBackgrounded(taskId: string): void { + this.userBackgroundedTaskIds.add(taskId); + } + + markTaskForegroundRelevant(taskId: string): void { + this.userBackgroundedTaskIds.delete(taskId); + } + + isTaskQueueBackgrounded(taskId: string): boolean { + return this.userBackgroundedTaskIds.has(taskId); + } + + isForegroundAwaiting(workspaceId: string): boolean { + const count = this.foregroundAwaitCountByWorkspaceId.get(workspaceId); + return typeof count === "number" && count > 0; + } + + startForegroundAwait(workspaceId: string): () => void { + assert(workspaceId.length > 0, "startForegroundAwait: workspaceId must be non-empty"); + + const current = this.foregroundAwaitCountByWorkspaceId.get(workspaceId) ?? 0; + assert( + Number.isInteger(current) && current >= 0, + "startForegroundAwait: expected non-negative integer counter" + ); + + this.foregroundAwaitCountByWorkspaceId.set(workspaceId, current + 1); + + return () => { + const current = this.foregroundAwaitCountByWorkspaceId.get(workspaceId) ?? 0; + assert( + Number.isInteger(current) && current > 0, + "startForegroundAwait cleanup: expected positive integer counter" + ); + if (current <= 1) { + this.foregroundAwaitCountByWorkspaceId.delete(workspaceId); + } else { + this.foregroundAwaitCountByWorkspaceId.set(workspaceId, current - 1); + } + }; + } + + registerBackgroundableForegroundWaiter( + workspaceId: string, + waiter: BackgroundableForegroundWaiter + ): void { + let set = this.backgroundableForegroundWaitersByWorkspaceId.get(workspaceId); + if (!set) { + set = new Set(); + this.backgroundableForegroundWaitersByWorkspaceId.set(workspaceId, set); + } + set.add(waiter); + } + + unregisterBackgroundableForegroundWaiter( + workspaceId: string, + waiter: BackgroundableForegroundWaiter + ): void { + const set = this.backgroundableForegroundWaitersByWorkspaceId.get(workspaceId); + if (!set) return; + set.delete(waiter); + if (set.size === 0) { + this.backgroundableForegroundWaitersByWorkspaceId.delete(workspaceId); + } + } + + /** + * Reject all foreground task waiters for a workspace that opted into backgrounding + * when a new message is queued. Returns the number of waiters signaled. + * Safe to call repeatedly — already-cleaned-up waiters are skipped. + */ + backgroundForegroundWaitsForWorkspace(workspaceId: string): number { + const set = this.backgroundableForegroundWaitersByWorkspaceId.get(workspaceId); + if (!set || set.size === 0) return 0; + + const waiters = [...set]; + let count = 0; + for (const waiter of waiters) { + try { + this.markTaskQueueBackgrounded(waiter.taskId); + // A foreground wait detached by a queued message becomes durably non-blocking: + // persist notify_on_terminal so future stream-ends and restarts do not re-force the + // await. The in-memory mark above covers the immediate next stream-end while this + // persistence settles. Tracked so handleStreamEnd can await it before reading config. + this.deps.scheduleNotifyOnTerminalPersist(waiter.taskId, waiter.requestingWorkspaceId); + waiter.reject(new ForegroundWaitBackgroundedError()); + count++; + } catch { + // waiter already resolved/rejected — ignore + } + } + return count; + } + + /** + * Background any registered foreground waits for the requesting workspace when a + * tool-end message is already queued. Shared by both wait-registration paths + * (workspace-turn and task await): the auto-backgrounding signal is edge-triggered + * on enqueue, so a message queued before the waiter registered must be re-checked + * here. No-op when backgrounding is disabled or no requesting workspace is set. + */ + backgroundForegroundWaitIfQueued( + shouldBackgroundOnQueuedMessage: boolean, + requestingWorkspaceId: string | undefined + ): void { + if ( + shouldBackgroundOnQueuedMessage && + requestingWorkspaceId && + this.deps.hasQueuedMessages(requestingWorkspaceId, "tool-end") + ) { + this.backgroundForegroundWaitsForWorkspace(requestingWorkspaceId); + } + } + + addPendingWaiter(taskId: string, waiter: PendingTaskWaiter): void { + appendToListMap(this.pendingWaitersByTaskId, taskId, waiter); + } + + removePendingWaiter(taskId: string, waiter: PendingTaskWaiter): void { + removeFromListMap(this.pendingWaitersByTaskId, taskId, waiter); + } + + hasPendingWaiters(taskId: string): boolean { + return (this.pendingWaitersByTaskId.get(taskId)?.length ?? 0) > 0; + } + + /** Drains and resolves the task's report waiters. Returns whether any waiter was pending. */ + resolvePendingWaiters(taskId: string, report: PendingTaskWaiterReport): boolean { + const waiters = this.pendingWaitersByTaskId.get(taskId); + if (!waiters || waiters.length === 0) { + return false; + } + + this.pendingWaitersByTaskId.delete(taskId); + for (const waiter of waiters) { + try { + waiter.cleanup(); + waiter.resolve(report); + } catch { + // ignore + } + } + + return true; + } + + rejectWaiters(taskId: string, error: Error): void { + this.markTaskForegroundRelevant(taskId); + + const waiters = this.pendingWaitersByTaskId.get(taskId); + if (!waiters || waiters.length === 0) { + return; + } + + for (const waiter of [...waiters]) { + try { + waiter.reject(error); + } catch (rejectError: unknown) { + log.error("Task waiter reject callback failed", { taskId, error: rejectError }); + } + } + } + + addPendingStartWaiter(taskId: string, waiter: PendingTaskStartWaiter): void { + appendToListMap(this.pendingStartWaitersByTaskId, taskId, waiter); + } + + removePendingStartWaiter(taskId: string, waiter: PendingTaskStartWaiter): void { + removeFromListMap(this.pendingStartWaitersByTaskId, taskId, waiter); + } + + /** + * A write that would have ended a queued/starting task rejected, and nothing will move the row + * to running (#4747, #5028). Fail closed: waiters are not rejected with the unpersisted outcome, + * whose owner is unknown when the updater never ran. Their timeout is armed instead, now and for + * waiters that attach later, so they settle at their timeout. + */ + armWaitersOfUnendedTask(taskId: string): void { + this.unpersistedLaunchFailureTaskIds.add(taskId); + this.startPendingStartWaiters(taskId); + } + + isLaunchFailureUnpersisted(taskId: string): boolean { + return this.unpersistedLaunchFailureTaskIds.has(taskId); + } + + clearUnpersistedLaunchFailure(taskId: string): void { + this.unpersistedLaunchFailureTaskIds.delete(taskId); + } + + /** Arms the report timeout of every waiter that attached while the task was queued/starting. */ + startPendingStartWaiters(workspaceId: string): void { + const waiters = this.pendingStartWaitersByTaskId.get(workspaceId); + if (!waiters || waiters.length === 0) return; + this.pendingStartWaitersByTaskId.delete(workspaceId); + for (const waiter of waiters) { + try { + waiter.start(); + } catch (error: unknown) { + log.error("Task start waiter callback failed", { workspaceId, error }); + } + } + } + + getCompletedReport(taskId: string): CompletedAgentReportCacheEntry | undefined { + return this.completedReportsByTaskId.get(taskId); + } + + hasCompletedReport(taskId: string): boolean { + return this.completedReportsByTaskId.has(taskId); + } + + /** Caches a completed report, evicting the oldest entries past the cache bound. */ + cacheCompletedReport(taskId: string, entry: CompletedAgentReportCacheEntry): void { + this.completedReportsByTaskId.set(taskId, entry); + this.enforceCompletedReportCacheLimit(); + } + + deleteCompletedReport(taskId: string): void { + this.completedReportsByTaskId.delete(taskId); + } + + private enforceCompletedReportCacheLimit(): void { + while (this.completedReportsByTaskId.size > COMPLETED_REPORT_CACHE_MAX_ENTRIES) { + const first = this.completedReportsByTaskId.keys().next(); + if (first.done) break; + this.completedReportsByTaskId.delete(first.value); + } + } +} diff --git a/src/node/services/workspaceService.removalCascade.test.ts b/src/node/services/workspaceService.removalCascade.test.ts index c2fd48882af..179d40573e9 100644 --- a/src/node/services/workspaceService.removalCascade.test.ts +++ b/src/node/services/workspaceService.removalCascade.test.ts @@ -172,8 +172,9 @@ describe("parent removal cascades over its sub-agent tree across two backends", // #5028: the mark a task gets when the write ending it rejects (#4747) is in-memory only and // has no reader once the row is gone. White-box: the leak is its only observable effect. test("removing tasks drops their unpersisted-end marks", async () => { - const marks = (a.taskService as unknown as { unpersistedLaunchFailureTaskIds: Set }) - .unpersistedLaunchFailureTaskIds; + const marks = ( + a.taskService as unknown as { waiters: { unpersistedLaunchFailureTaskIds: Set } } + ).waiters.unpersistedLaunchFailureTaskIds; marks.add(childId); marks.add(grandchildId);