diff --git a/packages/agent-core-v2/src/agent/task/persist.ts b/packages/agent-core-v2/src/agent/task/persist.ts index 9fc37d94f..45e0bea00 100644 --- a/packages/agent-core-v2/src/agent/task/persist.ts +++ b/packages/agent-core-v2/src/agent/task/persist.ts @@ -20,8 +20,10 @@ import { join } from 'pathe'; +import { Error2, ErrorCodes } from '#/errors'; import type { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; import type { IFileSystemStorageService } from '#/persistence/interface/storage'; +import { IWriteAuthorityRegistry, sessionIdFromScope } from '#/persistence/interface/writeAuthority'; import type { AgentTaskInfo, AgentTaskStatus } from './types'; @@ -73,7 +75,8 @@ export class AgentTaskPersistence { private readonly agentScope: string, private readonly docs: IAtomicDocumentStore, private readonly bytes: IFileSystemStorageService, - private readonly fallbackRoot?: AgentTaskPersistenceRoot, + private readonly fallbackRoot: AgentTaskPersistenceRoot | undefined, + private readonly authorityRegistry: IWriteAuthorityRegistry, ) {} private primaryRoot(): AgentTaskPersistenceRoot { @@ -103,6 +106,7 @@ export class AgentTaskPersistence { async writeTask(task: PersistedTask): Promise { validateTaskId(task.taskId); + this.assertWritable(); await this.docs.set(this.tasksScope(), `${task.taskId}${JSON_SUFFIX}`, task); } @@ -122,9 +126,23 @@ export class AgentTaskPersistence { async appendTaskOutput(taskId: string, chunk: string): Promise { if (chunk.length === 0) return; + validateTaskId(taskId); + this.assertWritable(); await this.bytes.append(this.taskOutputScope(taskId), OUTPUT_LOG_KEY, textEncoder.encode(chunk)); } + private assertWritable(): void { + const sessionId = sessionIdFromScope(this.agentScope); + if (sessionId === undefined) return; + const authority = this.authorityRegistry.resolve(sessionId); + if (authority === undefined) { + throw new Error2(ErrorCodes.SESSION_LEASE_LOST, 'session has no registered write authority', { + details: { sessionId }, + }); + } + authority.assertWritable(); + } + async taskOutputSizeBytes(taskId: string): Promise { const output = await this.readTaskOutputData(taskId); return output?.data.byteLength ?? 0; diff --git a/packages/agent-core-v2/src/agent/task/task.ts b/packages/agent-core-v2/src/agent/task/task.ts index 9e091502f..8e11832d9 100644 --- a/packages/agent-core-v2/src/agent/task/task.ts +++ b/packages/agent-core-v2/src/agent/task/task.ts @@ -84,6 +84,7 @@ export interface IAgentTaskService { getTask(taskId: string): AgentTaskInfo | undefined; list(activeOnly?: boolean, limit?: number): readonly AgentTaskInfo[]; persistOutput(taskId: string): void; + flushPersistence(): Promise; getOutputSnapshot( taskId: string, maxPreviewBytes: number, diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index 774572c5c..b4e4b77fc 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/packages/agent-core-v2/src/agent/task/taskService.ts @@ -62,6 +62,7 @@ import { IConfigService } from '#/app/config/config'; import { ISessionContext } from '#/session/sessionContext/sessionContext'; import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; import { IFileSystemStorageService } from '#/persistence/interface/storage'; +import { IWriteAuthorityRegistry } from '#/persistence/interface/writeAuthority'; import { ITelemetryService } from '#/app/telemetry/telemetry'; import { defineModel } from '#/wire/model'; import { IWireService } from '#/wire/wire'; @@ -229,6 +230,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { @IConfigService private readonly config: IConfigService, @IAtomicDocumentStore atomicDocs: IAtomicDocumentStore, @IFileSystemStorageService byteStore: IFileSystemStorageService, + @IWriteAuthorityRegistry authorityRegistry: IWriteAuthorityRegistry, @ISessionContext session: ISessionContext, @IAgentScopeContext scopeContext: IAgentScopeContext, @ITaskService private readonly taskService: ITaskService, @@ -248,6 +250,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { atomicDocs, byteStore, fallbackRoot, + authorityRegistry, ); this._register( this.wire.hooks.onDidRestore.register('task', async (_ctx, next) => { @@ -480,6 +483,34 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { this.startOutputPersist(entry); } + async flushPersistence(): Promise { + let firstFailure: unknown; + for (;;) { + const entries = [...this.tasks.values()]; + for (const entry of entries) { + if (entry.pendingOutput.length > 0) this.startOutputPersist(entry); + } + const queues = entries.flatMap((entry) => [ + { entry, kind: 'state' as const, promise: entry.persistWriteQueue }, + { entry, kind: 'output' as const, promise: entry.outputWriteQueue }, + ]); + const results = await Promise.allSettled(queues.map(({ promise }) => promise)); + for (const result of results) { + if (result.status === 'rejected' && firstFailure === undefined) { + firstFailure = result.reason; + } + } + const changed = queues.some(({ entry, kind, promise }) => { + return kind === 'state' + ? entry.persistWriteQueue !== promise + : entry.outputWriteQueue !== promise; + }); + if (changed) continue; + if (firstFailure !== undefined) throw firstFailure; + return; + } + } + async loadFromDisk(options: AgentTaskLoadOptions = {}): Promise { const persistence = this.persistence; if (options.replace !== false) { @@ -876,8 +907,8 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { const persistence = this.persistence; const info = this.toInfo(entry); entry.persistWriteQueue = entry.persistWriteQueue - .then(() => persistence.writeTask(info)) - .catch(() => { }); + .then(() => persistence.writeTask(info)); + void entry.persistWriteQueue.catch(() => {}); return entry.persistWriteQueue; } @@ -911,8 +942,8 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { private appendTaskOutput(entry: ManagedTask, chunk: string): void { const persistence = this.persistence; entry.outputWriteQueue = entry.outputWriteQueue - .then(() => persistence.appendTaskOutput(entry.taskId, chunk)) - .catch(() => { }); + .then(() => persistence.appendTaskOutput(entry.taskId, chunk)); + void entry.outputWriteQueue.catch(() => {}); } private startOutputPersist(entry: ManagedTask): void { diff --git a/packages/agent-core-v2/src/app/sessionLifecycle/sessionLifecycleService.ts b/packages/agent-core-v2/src/app/sessionLifecycle/sessionLifecycleService.ts index d0d8f0335..f6811506f 100644 --- a/packages/agent-core-v2/src/app/sessionLifecycle/sessionLifecycleService.ts +++ b/packages/agent-core-v2/src/app/sessionLifecycle/sessionLifecycleService.ts @@ -93,6 +93,7 @@ import { OsLockErrors, } from '#/os/interface/crossProcessLock'; import { IAgentLifecycleService, MAIN_AGENT_ID } from '#/session/agentLifecycle/agentLifecycle'; +import { IAgentTaskService } from '#/agent/task/task'; import { ensureMainAgent } from '#/session/agentLifecycle/mainAgent'; import { ISessionMcpService } from '#/session/mcp/sessionMcp'; import { labelsFromAgentMeta } from '#/session/agentLifecycle/subagentMetadata'; @@ -153,6 +154,7 @@ interface SessionEntry { closeKind?: SessionCloseKind; closeStep: number; closePromise?: Promise; + dirtyAbortPromise?: Promise; } export class SessionLifecycleService extends Disposable implements ISessionLifecycleService { @@ -917,16 +919,31 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec } private dirtyAbortSession(entry: SessionEntry): void { + if (entry.dirtyAbortPromise !== undefined) return; if (this.entries.get(entry.handle.id) === entry) this.entries.delete(entry.handle.id); + + let taskServices: IAgentTaskService[] = []; + try { + taskServices = entry.handle + .accessor.get(IAgentLifecycleService) + .list() + .map((agent) => agent.accessor.get(IAgentTaskService)); + } catch { + } try { this.disposeSessionHandle(entry); } catch { } - try { - entry.registration.dispose(); - } catch { - } - entry.lease.release(); + entry.dirtyAbortPromise = Promise.allSettled( + taskServices.map((tasks) => tasks.flushPersistence()), + ).then(() => { + try { + entry.registration.dispose(); + } catch { + } + entry.lease.release(); + }); + void entry.dirtyAbortPromise.catch(() => {}); } private async readMetaFromDisk( diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index df61abfbd..b03935dac 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -344,7 +344,9 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle const handle = this.handles.get(agentId); if (handle === undefined) return; this.handles.delete(agentId); - await handle.accessor.get(IAgentTaskService).stopAllOnExit('Session closed'); + const tasks = handle.accessor.get(IAgentTaskService); + await tasks.stopAllOnExit('Session closed'); + await tasks.flushPersistence?.(); 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/persist.test.ts b/packages/agent-core-v2/test/agent/task/persist.test.ts index 8054fe687..10749d8a8 100644 --- a/packages/agent-core-v2/test/agent/task/persist.test.ts +++ b/packages/agent-core-v2/test/agent/task/persist.test.ts @@ -16,14 +16,17 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { SyncDescriptor } from '#/_base/di/descriptors'; import { DisposableStore } from '#/_base/di/lifecycle'; import { TestInstantiationService } from '#/_base/di/test'; +import { ErrorCodes } from '#/errors'; import { AgentTaskPersistence, type AgentTaskInfo, } from '#/agent/task/task'; import { JsonAtomicDocumentStore } from '#/persistence/backends/node-fs/atomicDocumentStore'; import { FileStorageService } from '#/persistence/backends/node-fs/fileStorageService'; +import { WriteAuthorityRegistryService } from '#/persistence/backends/node-fs/writeAuthorityRegistryService'; import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; import { IFileSystemStorageService } from '#/persistence/interface/storage'; +import type { IWriteAuthorityRegistry } from '#/persistence/interface/writeAuthority'; const SESSION_SCOPE = 'session'; const AGENT_SCOPE = `${SESSION_SCOPE}/agents/main`; @@ -32,6 +35,7 @@ let disposables: DisposableStore; let sessionDir: string; let docs: IAtomicDocumentStore; let bytes: IFileSystemStorageService; +let authorityRegistry: IWriteAuthorityRegistry; let persistence: AgentTaskPersistence; function sample(overrides: Partial> = {}): Extract { @@ -64,7 +68,15 @@ beforeEach(async () => { ix.set(IAtomicDocumentStore, new SyncDescriptor(JsonAtomicDocumentStore)); docs = ix.get(IAtomicDocumentStore); bytes = ix.get(IFileSystemStorageService); - persistence = new AgentTaskPersistence(sessionDir, SESSION_SCOPE, docs, bytes); + authorityRegistry = new WriteAuthorityRegistryService(); + persistence = new AgentTaskPersistence( + sessionDir, + SESSION_SCOPE, + docs, + bytes, + undefined, + authorityRegistry, + ); }); afterEach(async () => { @@ -77,7 +89,14 @@ describe('AgentTaskPersistence', () => { scope: string, fallbackRoot?: { readonly dir: string; readonly scope: string }, ): AgentTaskPersistence { - return new AgentTaskPersistence(join(sessionDir, scope), scope, docs, bytes, fallbackRoot); + return new AgentTaskPersistence( + join(sessionDir, scope), + scope, + docs, + bytes, + fallbackRoot, + authorityRegistry, + ); } function sessionRoot(): { readonly dir: string; readonly scope: string } { @@ -108,6 +127,26 @@ describe('AgentTaskPersistence', () => { }); }); + it('fails closed when the session write authority is unregistered', async () => { + const scope = 'sessions/workspace/test-session/agents/main'; + const fenced = rootedPersistence(scope); + const registration = authorityRegistry.register({ + sessionId: 'test-session', + assertWritable: () => {}, + }); + + await fenced.writeTask(sample()); + await fenced.appendTaskOutput(sample().taskId, 'before release'); + registration.dispose(); + + await expect(fenced.writeTask(sample({ status: 'completed', endedAt: 2 }))).rejects.toMatchObject({ + code: ErrorCodes.SESSION_LEASE_LOST, + }); + await expect(fenced.appendTaskOutput(sample().taskId, 'after release')).rejects.toMatchObject({ + code: ErrorCodes.SESSION_LEASE_LOST, + }); + }); + it('listTasks enumerates all persisted entries', async () => { await persistence.writeTask(sample({ taskId: 'bash-11111111' })); await persistence.writeTask(sample({ taskId: 'bash-22222222', command: 'pnpm test' })); diff --git a/packages/agent-core-v2/test/agent/task/stubs.ts b/packages/agent-core-v2/test/agent/task/stubs.ts index 66f59de32..b822521d7 100644 --- a/packages/agent-core-v2/test/agent/task/stubs.ts +++ b/packages/agent-core-v2/test/agent/task/stubs.ts @@ -14,6 +14,7 @@ import { } from '#/agent/task/task'; import { JsonAtomicDocumentStore } from '#/persistence/backends/node-fs/atomicDocumentStore'; import { FileStorageService } from '#/persistence/backends/node-fs/fileStorageService'; +import { WriteAuthorityRegistryService } from '#/persistence/backends/node-fs/writeAuthorityRegistryService'; export type TaskServiceTestManager = IAgentTaskService & { loadFromDisk(): Promise; @@ -26,10 +27,17 @@ export const TASK_TEST_AGENT_SCOPE = `${TASK_TEST_SESSION_SCOPE}/agents/main`; export function createAgentTaskPersistence(homedir: string): AgentTaskPersistence { const storage = new FileStorageService(homedir); + const authorityRegistry = new WriteAuthorityRegistryService(); + authorityRegistry.register({ + sessionId: 'test-session', + assertWritable: () => {}, + }); return new AgentTaskPersistence( join(homedir, TASK_TEST_AGENT_SCOPE), TASK_TEST_AGENT_SCOPE, new JsonAtomicDocumentStore(storage), storage, + undefined, + authorityRegistry, ); } 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 8d9f68dc6..45148cda5 100644 --- a/packages/agent-core-v2/test/agent/task/taskService.test.ts +++ b/packages/agent-core-v2/test/agent/task/taskService.test.ts @@ -36,6 +36,7 @@ import { IAgentScopeContext, makeAgentScopeContext } from '#/agent/scopeContext/ import { ISessionContext, makeSessionContext } from '#/session/sessionContext/sessionContext'; import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; import { IFileSystemStorageService } from '#/persistence/interface/storage'; +import { IWriteAuthorityRegistry } from '#/persistence/interface/writeAuthority'; import { ITelemetryService } from '#/app/telemetry/telemetry'; import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry'; import { createHooks } from '#/hooks'; @@ -148,6 +149,10 @@ describe('AgentTaskService', () => { flush: async () => {}, close: async () => {}, }); + ix.stub(IWriteAuthorityRegistry, { + resolve: () => ({ sessionId: 'test-session', assertWritable: () => {} }), + register: () => toDisposable(() => {}), + }); ix.set(IAgentTaskService, new SyncDescriptor(AgentTaskService)); }); afterEach(() => disposables.dispose()); @@ -316,6 +321,63 @@ describe('AgentTaskService', () => { await svc.stop(taskId); }); + it('flushPersistence waits for a delayed output append', async () => { + let releaseAppend!: () => void; + let markAppendStarted!: () => void; + const appendStarted = new Promise((resolve) => { + markAppendStarted = resolve; + }); + const appendReleased = new Promise((resolve) => { + releaseAppend = resolve; + }); + ix.stub(IFileSystemStorageService, { + read: async () => undefined, + readStream: async function* () {}, + write: async () => {}, + append: async () => { + markAppendStarted(); + await appendReleased; + }, + list: async () => [], + delete: async () => {}, + flush: async () => {}, + close: async () => {}, + }); + const svc = ix.get(IAgentTaskService); + let releaseTask!: () => void; + const taskRunning = new Promise((resolve) => { + releaseTask = resolve; + }); + const taskId = svc.registerTask({ + idPrefix: 'test', + kind: 'agent', + description: 'delayed output', + start: async (sink) => { + sink.appendOutput('delayed output'); + await taskRunning; + await sink.settle({ status: 'completed' }); + }, + toInfo: (base) => ({ ...base, kind: 'agent' }), + }); + await Promise.resolve(); + await Promise.resolve(); + svc.persistOutput(taskId); + await appendStarted; + + let flushed = false; + const flush = svc.flushPersistence().then(() => { + flushed = true; + }); + await Promise.resolve(); + expect(flushed).toBe(false); + + releaseAppend(); + await flush; + expect(flushed).toBe(true); + releaseTask(); + await svc.stop(taskId); + }); + it('dispose aborts live tasks as a last resort', async () => { const svc = ix.get(IAgentTaskService); let abortReason: unknown; @@ -496,6 +558,10 @@ describe('AgentTaskService', () => { ); ix.stub(IAtomicDocumentStore, docs); ix.stub(IFileSystemStorageService, bytes); + ix.stub(IWriteAuthorityRegistry, { + resolve: () => ({ sessionId: 'test-session', assertWritable: () => {} }), + register: () => toDisposable(() => {}), + }); ix.set(IAgentTaskService, new SyncDescriptor(AgentTaskService)); return ix; } 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 97051a38b..c6f2ec580 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 @@ -172,6 +172,8 @@ class FakeTaskService implements IAgentTaskService { persistOutput(_taskId: string): void {} + async flushPersistence(): Promise {} + async getOutputSnapshot( taskId: string, _maxPreviewBytes: number, diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index a8836c6b5..3793f1ae3 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -165,6 +165,12 @@ import { ISessionQuestionService, type QuestionResult } from '#/session/question import { ISessionSkillCatalog } from '#/session/sessionSkillCatalog/skillCatalog'; import { ISessionSwarmService } from '#/session/swarm/sessionSwarm'; import type { PathAccessOperation } from '#/session/workspaceContext/workspaceContext'; +import { WriteAuthorityRegistryService } from '#/persistence/backends/node-fs/writeAuthorityRegistryService'; +import { IWriteAuthorityRegistry } from '#/persistence/interface/writeAuthority'; +import { + IHostFsWatchService, + type HostFsChange, +} from '#/os/interface/hostFsWatch'; import { recordAgentEvents, type RecordedEventEntry } from '../snapshot/events'; import { createFakeHostFs, createFakeProcessRunner } from '../tools/fixtures/fake-exec'; @@ -523,6 +529,7 @@ export function homeDirServices(homeDir: string | undefined): TestAgentServiceOv if (homeDir !== undefined) { for (const [id, value] of bootstrapSeed({ homeDir, + osHomeDir: homeDir, cwd: process.cwd(), env: process.env, })) { @@ -645,6 +652,14 @@ const noopHookRunner: IExternalHooksRunnerService = { fireAndForgetTrigger: async () => [], }; +const noopHostFsWatchService: IHostFsWatchService = { + _serviceBrand: undefined, + watch: () => ({ + onDidChange: Event.None as Event, + dispose: () => {}, + }), +}; + export function permissionModeServices(mode: PermissionMode): TestAgentServiceOverride { return agentService(IAgentPermissionModeService, createPermissionModeService(mode)); } @@ -992,6 +1007,13 @@ export class AgentTestContext { })) { reg.defineInstance(id, value); } + const authorityRegistry = new WriteAuthorityRegistryService(); + authorityRegistry.register({ + sessionId, + assertWritable: () => {}, + }); + reg.defineInstance(IWriteAuthorityRegistry, authorityRegistry); + reg.defineInstance(IHostFsWatchService, noopHostFsWatchService); const memoryStorage = (): SyncDescriptor => new SyncDescriptor(InMemoryStorageService, [], true); reg.defineDescriptor(IFileSystemStorageService, memoryStorage()); 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 a4ca4eea3..4cab533e5 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 @@ -474,6 +474,7 @@ function createFakeTaskService(options: { maxRunningTasks?: number } = {}): { const service: IAgentTaskService = { _serviceBrand: undefined, + flushPersistence: async () => {}, track(): never { throw new Error('fake IAgentTaskService.track is not implemented'); },