diff --git a/packages/agent-core-v2/src/agent/task/task.ts b/packages/agent-core-v2/src/agent/task/task.ts index 631f417afc..2a362988b0 100644 --- a/packages/agent-core-v2/src/agent/task/task.ts +++ b/packages/agent-core-v2/src/agent/task/task.ts @@ -85,6 +85,7 @@ export interface IAgentTaskService { ): Promise; readOutput(taskId: string, tail?: number): Promise; suppressTerminalNotification(taskId: string): Promise; + suppressAllTerminalNotifications(): Promise; markTasksDeliveredViaWait(tasks: readonly AgentTaskWaitDelivery[]): void; detach(taskId: string): AgentTaskInfo | undefined; stop(taskId: string, reason?: string): Promise; diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index d347c18990..69e4254e18 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/packages/agent-core-v2/src/agent/task/taskService.ts @@ -205,6 +205,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { declare readonly _serviceBrand: undefined; private readonly tasks = new Map(); + private exitSuppressionArmed = false; private readonly buildingNotificationKeys = new Set(); private readonly pendingNotificationRequests = new Map(); private readonly persistence: AgentTaskPersistence; @@ -779,21 +780,16 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { return results.filter((info): info is AgentTaskInfo => info !== undefined); } + async suppressAllTerminalNotifications(): Promise { + this.exitSuppressionArmed = true; + for (const [, request] of Array.from(this.pendingNotificationRequests)) { + request.drop(); + } + } + async stopAllOnExit(reason: string): Promise { + await this.suppressAllTerminalNotifications(); if (this.keepAliveOnExit()) return []; - const active = this.list(true); - await Promise.allSettled( - active - .filter((task) => task.detached === true) - .map((task) => - this.suppressTerminalNotification(task.taskId).catch((error: unknown) => { - this.log.error('terminal notification suppression failed', { - taskId: task.taskId, - error, - }); - }), - ), - ); return this.stopAll(reason); } @@ -834,6 +830,10 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { return this.sessionEventBus.isAgentActive(this.scopeContext.agentContext); } + private marksTerminalNotificationSuppressed(entry: ManagedTask): boolean { + return this.exitSuppressionArmed && !this.keepAliveOnExit() && this.isDetached(entry); + } + async wait( taskId: string, timeoutMs = 30_000, @@ -1035,12 +1035,22 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { entry.timeoutHandle = undefined; } const foregroundRelease = entry.foregroundRelease; + if (this.marksTerminalNotificationSuppressed(entry)) { + entry.terminalNotificationSuppressed = true; + } if (entry.outputPersistStarted) { await this.persistLive(entry); } else { entry.pendingOutput = []; entry.pendingOutputBytes = 0; } + if ( + this.marksTerminalNotificationSuppressed(entry) && + entry.terminalNotificationSuppressed !== true + ) { + entry.terminalNotificationSuppressed = true; + await this.persistLive(entry); + } this.fireTerminalEffects(entry); foregroundRelease?.resolve('terminal'); this.resolveWaiters(entry); @@ -1095,7 +1105,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { if (!this.lifecycleActive()) return; const context = await this.buildAgentTaskNotificationContext(info); if (context === undefined) return; - if (!this.lifecycleActive()) return; + if (!this.lifecycleActive() || this.isTerminalNotificationSuppressed(info.taskId)) return; const key = notificationKey(context.origin); if (this.deliveredNotificationKeys.has(key)) return; const handle = this.loop.notify({ @@ -1302,6 +1312,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { private isTerminalNotificationSuppressed(taskId: string): boolean { return ( + this.exitSuppressionArmed || this.tasks.get(taskId)?.terminalNotificationSuppressed === true || this.ghosts.get(taskId)?.terminalNotificationSuppressed === true ); diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index 25ac64ab0e..1696ffecce 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -388,6 +388,7 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle managed.closing = true; this.onWillCloseEmitter.fire(agent); const handle = managed.handle; + await handle.accessor.get(IAgentTaskService).suppressAllTerminalNotifications(); const loop = handle.accessor.get(IAgentLoopService); const compaction = handle.accessor.get(IAgentFullCompactionService).compacting; const compactionSettled = compaction?.promise.catch(() => undefined) ?? Promise.resolve(); diff --git a/packages/agent-core-v2/test/agent/task/taskService.test.ts b/packages/agent-core-v2/test/agent/task/taskService.test.ts index 847e40bfd1..3354771485 100644 --- a/packages/agent-core-v2/test/agent/task/taskService.test.ts +++ b/packages/agent-core-v2/test/agent/task/taskService.test.ts @@ -287,15 +287,37 @@ describe('AgentTaskService', () => { } } - it('enqueues a terminal notification for a finished detached task', async () => { - const svc = ix.get(IAgentTaskService); + it('enqueues a terminal notification for a finished detached task, but not when suppression arms mid-build', async () => { + let armOnRead = false; + let svc!: IAgentTaskService; + ix.stub(IFileSystemStorageService, { + read: async () => { + if (armOnRead) await svc.suppressAllTerminalNotifications(); + return undefined; + }, + readStream: async function* () {}, + write: async () => {}, + writeStream: async () => {}, + append: async () => {}, + list: async () => [], + delete: async () => {}, + flush: async () => {}, + }); + svc = ix.get(IAgentTaskService); const taskId = svc.registerTask(outputtingTask('done\n')); await svc.wait(taskId, 1000); const loop = stubLoop(); await waitForCondition(() => loop.hasPendingRequests()); - expect(loop.hasPendingRequests()).toBe(true); + + loop.drainNextBatch({ append: () => {} }); + armOnRead = true; + const second = svc.registerTask(outputtingTask('done\n')); + await svc.wait(second, 1000); + await new Promise((resolve) => setTimeout(resolve, 10)); + + expect(loop.hasPendingRequests()).toBe(false); }); it('markTasksDeliveredViaWait suppresses the automatic terminal notification', async () => { @@ -317,7 +339,7 @@ describe('AgentTaskService', () => { expect(states.get(taskNotificationDeliveryKey)).toContain(deliveryKey); }); - it('aborts an already-enqueued terminal notification when the task is marked delivered via wait', async () => { + it('aborts an already-enqueued terminal notification when the task is marked delivered via wait or suppression arms', async () => { const svc = ix.get(IAgentTaskService); const taskId = svc.registerTask(outputtingTask('done\n')); @@ -329,6 +351,15 @@ describe('AgentTaskService', () => { svc.markTasksDeliveredViaWait([{ taskId, status: 'completed' }]); expect(loop.hasPendingRequests()).toBe(false); + + const second = svc.registerTask(outputtingTask('done\n')); + await svc.wait(second, 1000); + await waitForCondition(() => loop.hasPendingRequests()); + expect(loop.hasPendingRequests()).toBe(true); + + await svc.suppressAllTerminalNotifications(); + + expect(loop.hasPendingRequests()).toBe(false); }); it('suppresses only the notification whose status was reported via wait', async () => { @@ -507,26 +538,25 @@ describe('AgentTaskService', () => { const first = svc.registerTask(fakeProcessTask()); const second = svc.registerTask(fakeProcessTask()); + await svc.suppressAllTerminalNotifications(); + const third = svc.registerTask(fakeProcessTask()); + const stopped = await svc.stopAllOnExit('Session closed'); - expect(stopped.map((info) => info.taskId).toSorted()).toEqual([first, second].toSorted()); - for (const taskId of [first, second]) { + expect(stopped.map((info) => info.taskId).toSorted()).toEqual( + [first, second, third].toSorted(), + ); + for (const taskId of [first, second, third]) { const info = svc.getTask(taskId); expect(info?.status).toBe('killed'); expect(info?.stopReason).toBe('Session closed'); expect(info?.terminalNotificationSuppressed).toBe(true); - const persisted = writes.filter((write) => write.taskId === taskId); - expect( - persisted.some( - (write) => - write.status === 'running' && write.terminalNotificationSuppressed === true, - ), - ).toBe(true); - expect(persisted.at(-1)).toMatchObject({ + expect(writes.filter((write) => write.taskId === taskId).at(-1)).toMatchObject({ status: 'killed', terminalNotificationSuppressed: true, }); } + expect(stubLoop().hasPendingRequests()).toBe(false); }); it('stopAllOnExit does not persist a foreground-only task', async () => { @@ -544,7 +574,7 @@ describe('AgentTaskService', () => { }); }); - it('stopAllOnExit still stops tasks when suppression persistence fails', async () => { + it('stopAllOnExit still stops tasks when persistence fails', async () => { let writes = 0; ix.stub(IAtomicDocumentStore, { get: async () => undefined, @@ -566,7 +596,7 @@ describe('AgentTaskService', () => { expect(svc.getTask(second)?.status).toBe('killed'); }); - it('stopAllOnExit leaves tasks running when keepAliveOnExit is set', async () => { + it('stopAllOnExit leaves tasks running and suppresses in flight without persisting the marker when keepAliveOnExit is set', async () => { stubTaskConfig({ keepAliveOnExit: true }); const svc = ix.get(IAgentTaskService); const taskId = svc.registerTask(fakeProcessTask()); @@ -577,6 +607,10 @@ describe('AgentTaskService', () => { expect(svc.getTask(taskId)?.status).toBe('running'); await svc.stop(taskId); + + expect(svc.getTask(taskId)?.status).toBe('killed'); + expect(svc.getTask(taskId)?.terminalNotificationSuppressed).toBeUndefined(); + expect(stubLoop().hasPendingRequests()).toBe(false); }); it('dispose aborts live tasks as a last resort', async () => { diff --git a/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts b/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts index 7584a8348d..7b095e4567 100644 --- a/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts +++ b/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts @@ -200,6 +200,11 @@ class FakeTaskService implements IAgentTaskService { } as AgentTaskInfo; } + async suppressAllTerminalNotifications(): Promise { + const active = this.list(true).filter((info) => info.detached === true); + await Promise.all(active.map((info) => this.suppressTerminalNotification(info.taskId))); + } + markTasksDeliveredViaWait(tasks: readonly AgentTaskWaitDelivery[]): void { this.waitDeliveries.push(tasks); } diff --git a/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts b/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts index a439a88813..161be6e186 100644 --- a/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts +++ b/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts @@ -593,6 +593,9 @@ function createFakeTaskService( async suppressTerminalNotification(): Promise { }, + async suppressAllTerminalNotifications(): Promise { + }, + markTasksDeliveredViaWait(): void { }, diff --git a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts index e9af2df4a1..d37297be07 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -204,6 +204,7 @@ describe('AgentLifecycleService', () => { let atomicDocs: Map; let permissionModeSetMode: ReturnType; let stopAllOnExit: ReturnType; + let suppressAllTerminalNotifications: ReturnType; let loopActiveTurnId: number | undefined; let loopPendingPromptIds: string[]; let loopCancel: ReturnType>; @@ -463,9 +464,11 @@ describe('AgentLifecycleService', () => { isBaselineServer: () => true, } satisfies ISessionMcpHandle); stopAllOnExit = vi.fn(async () => []); + suppressAllTerminalNotifications = vi.fn(async () => {}); ix.stub(IAgentTaskService, { _serviceBrand: undefined, stopAllOnExit, + suppressAllTerminalNotifications, } as unknown as IAgentTaskService); ix.stub(IAgentFullCompactionService, { _serviceBrand: undefined, @@ -709,6 +712,10 @@ describe('AgentLifecycleService', () => { expect(stopAllOnExit).toHaveBeenCalledWith('Session closed'); expect(promptDrain).toHaveBeenCalledOnce(); + expect(suppressAllTerminalNotifications).toHaveBeenCalledOnce(); + expect(suppressAllTerminalNotifications.mock.invocationCallOrder[0]).toBeLessThan( + promptDrain.mock.invocationCallOrder[0]!, + ); expect(stopAllOnExit.mock.invocationCallOrder[0]).toBeGreaterThan( promptDrain.mock.invocationCallOrder[0]!, );