From e7c397a7c8ac6694c991fc7ea76afb4459642a36 Mon Sep 17 00:00:00 2001 From: 7Sageer <7sageer@djwcb.cn> Date: Fri, 31 Jul 2026 14:20:34 +0800 Subject: [PATCH] fix: harden the session close/archive and create/fork lifecycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The closing registry now records the operation kind: an archive arriving during a plain close waits it out and lands the archived flag on the persisted document instead of riding the close to success, and a failing close hook no longer strands a half-closed session — the teardown always completes while the hook error still reaches the caller. create and fork reserve their target id synchronously with the existence check, so a concurrent create/fork of the same id loses up front and can never tear down the winner's scope or directory. --- .../workspaceHandlerService.ts | 118 ++++++++++++--- .../workspaceHandler/workspaceHandler.test.ts | 134 ++++++++++++++++++ 2 files changed, 230 insertions(+), 22 deletions(-) diff --git a/packages/agent-core-v2/src/workspace/workspaceHandler/workspaceHandlerService.ts b/packages/agent-core-v2/src/workspace/workspaceHandler/workspaceHandlerService.ts index 52fc5a2ef..013d2a249 100644 --- a/packages/agent-core-v2/src/workspace/workspaceHandler/workspaceHandlerService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceHandler/workspaceHandlerService.ts @@ -11,8 +11,12 @@ * session work can cancel itself and drop its write-back. A close/archive * is tracked in a closing registry from that first synchronous step until * the scope is disposed: `get`/`list` hide the session while it is - * closing, and `resume` waits the close out and re-materializes instead - * of returning the doomed handle, + * closing, `resume` waits the close out and re-materializes instead of + * returning the doomed handle, and a failing close hook never strands the + * teardown. A duplicate close joins the in-flight one, while an archive + * arriving during a plain close waits it out and then lands the archived + * flag directly on the persisted document (it never rides the close to + * success without archiving), * tearing sessions down on close/archive — archiving flags the session's * `sessionMetadata`, removes its `agentLifecycle` agents, restoring clears * the archived flag, and broadcasts through `event`; session start and @@ -161,6 +165,11 @@ type MaterializeSessionOptions = Omit & { readonly sessionId: string; }; +type ClosingEntry = { + readonly kind: 'close' | 'archive'; + readonly run: Promise; +}; + export class WorkspaceHandlerService extends Disposable implements IWorkspaceHandlerService { declare readonly _serviceBrand: undefined; private readonly sessions = new Map(); @@ -173,8 +182,9 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan private readonly _onDidForkSession = this._register(new Emitter()); readonly onDidForkSession: Event = this._onDidForkSession.event; private readonly resuming = new Map>(); - private readonly closing = new Map>(); + private readonly closing = new Map(); private readonly sessionLifetimes = new Map(); + private readonly reservedTargets = new Set(); constructor( @IInstantiationService private readonly instantiation: IInstantiationService, @@ -219,6 +229,24 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan async create(opts: CreateSessionOptions): Promise { const sessionId = opts.sessionId ?? createSessionId(); + if (this.sessions.has(sessionId) || this.reservedTargets.has(sessionId)) { + throw new Error2( + ErrorCodes.SESSION_ALREADY_EXISTS, + `Session "${sessionId}" already exists`, + ); + } + this.reservedTargets.add(sessionId); + try { + return await this.createReserved(opts, sessionId); + } finally { + this.reservedTargets.delete(sessionId); + } + } + + private async createReserved( + opts: CreateSessionOptions, + sessionId: string, + ): Promise { const handle = await this.materializeSession({ ...opts, sessionId }); try { const main = @@ -376,7 +404,7 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan if (closing !== undefined) { // A close is already in flight — wait it out and re-materialize // instead of handing back the doomed handle. - return closing.then(() => this.resume(sessionId, opts)); + return closing.run.then(() => this.resume(sessionId, opts)); } const live = this.sessions.get(sessionId); if (live !== undefined) return Promise.resolve(live); @@ -429,27 +457,43 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan async close(sessionId: string): Promise { const inflight = this.closing.get(sessionId); - if (inflight !== undefined) return inflight; + if (inflight !== undefined) return inflight.run; const handle = this.sessions.get(sessionId); if (handle === undefined) return; - return this.trackClosing(sessionId, this.doClose(sessionId, handle)); + return this.trackClosing(sessionId, 'close', this.doClose(sessionId, handle)); } private async doClose(sessionId: string, handle: ISessionScopeHandle): Promise { this.invalidateSessionLifetime(sessionId); - await this.announceWillClose({ sessionId, handle, reason: 'exit' }); - this.sessions.delete(sessionId); - await this.drainAgents(handle); - handle.dispose(); - this._onDidCloseSession.fire({ sessionId }); + try { + await this.announceWillClose({ sessionId, handle, reason: 'exit' }); + } finally { + // A failing close hook must not strand the session half-open: the + // teardown always completes; the hook error still reaches the caller. + this.sessions.delete(sessionId); + await this.drainAgents(handle); + handle.dispose(); + this._onDidCloseSession.fire({ sessionId }); + } } async archive(sessionId: string): Promise { const inflight = this.closing.get(sessionId); - if (inflight !== undefined) return inflight; + if (inflight !== undefined) { + if (inflight.kind === 'archive') return inflight.run; + // A plain close is already tearing the session down — the archive + // must not ride it to success. Wait it out (tracked, so a resume + // queues behind the marker write), then land the archived flag + // directly on the persisted document of the now-closed session. + return this.trackClosing( + sessionId, + 'archive', + inflight.run.then(() => this.markArchivedOnDisk(sessionId)), + ); + } const handle = this.sessions.get(sessionId); if (handle === undefined) return; - return this.trackClosing(sessionId, this.doArchive(sessionId, handle)); + return this.trackClosing(sessionId, 'archive', this.doArchive(sessionId, handle)); } private async doArchive(sessionId: string, handle: ISessionScopeHandle): Promise { @@ -461,18 +505,32 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan type: 'event.session.archived', payload: { sessionId }, }); - await this.announceWillClose({ sessionId, handle, reason: 'exit' }); - this.sessions.delete(sessionId); - handle.dispose(); - this._onDidArchiveSession.fire({ sessionId }); + try { + await this.announceWillClose({ sessionId, handle, reason: 'exit' }); + } finally { + this.sessions.delete(sessionId); + handle.dispose(); + this._onDidArchiveSession.fire({ sessionId }); + } } - private async trackClosing(sessionId: string, run: Promise): Promise { - this.closing.set(sessionId, run); + private async markArchivedOnDisk(sessionId: string): Promise { + const scope = sessionScopeOf(this.handlerScope, sessionId); + const meta = await this.docs.get(scope, 'state.json'); + if (meta === undefined) return; + await this.docs.set(scope, 'state.json', { ...meta, archived: true, updatedAt: Date.now() }); + } + + private async trackClosing( + sessionId: string, + kind: ClosingEntry['kind'], + run: Promise, + ): Promise { + this.closing.set(sessionId, { kind, run }); try { await run; } finally { - if (this.closing.get(sessionId) === run) this.closing.delete(sessionId); + if (this.closing.get(sessionId)?.run === run) this.closing.delete(sessionId); } } @@ -509,7 +567,7 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan // Wait out an in-flight close of the source instead of capturing its // doomed handle (a close starting mid-fork remains an accepted, // pre-existing window — fork never quiesces the source). - await this.closing.get(sourceId); + await this.closing.get(sourceId)?.run; const sourceHandle = this.sessions.get(sourceId); const indexSummary = await this.index.get(sourceId); if ( @@ -530,6 +588,7 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan let targetId: string | undefined; let target: ISessionScopeHandle | undefined; let targetSessionDir: string | undefined; + let targetReserved = false; try { const sourceMeta = sourceHandle !== undefined @@ -537,7 +596,18 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan : await this.readMetaFromDisk(sourceId); targetId = opts.newSessionId ?? createSessionId(); - if (this.sessions.has(targetId) || (await this.index.get(targetId)) !== undefined) { + if (this.sessions.has(targetId) || this.reservedTargets.has(targetId)) { + throw new Error2( + ErrorCodes.SESSION_ALREADY_EXISTS, + `Session "${targetId}" already exists`, + ); + } + // The check and the reservation are one synchronous step: a concurrent + // create/fork of the same target id loses before either materializes, + // so the loser can never tear down the winner's scope or directory. + this.reservedTargets.add(targetId); + targetReserved = true; + if ((await this.index.get(targetId)) !== undefined) { throw new Error2( ErrorCodes.SESSION_ALREADY_EXISTS, `Session "${targetId}" already exists`, @@ -612,6 +682,10 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan await this.hostFs.remove(targetSessionDir).catch(() => {}); } throw error; + } finally { + if (targetReserved && targetId !== undefined) { + this.reservedTargets.delete(targetId); + } } } diff --git a/packages/agent-core-v2/test/workspace/workspaceHandler/workspaceHandler.test.ts b/packages/agent-core-v2/test/workspace/workspaceHandler/workspaceHandler.test.ts index 384168af5..d7982e558 100644 --- a/packages/agent-core-v2/test/workspace/workspaceHandler/workspaceHandler.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceHandler/workspaceHandler.test.ts @@ -602,6 +602,140 @@ describe('WorkspaceHandlerService', () => { expect(svc.get('s1')).toBeUndefined(); }); + it('resume waits out an in-flight close and returns a fresh scope', async () => { + const svc = await build([ + stubPair(IWorkspaceService, persistentWorkspaceStub()), + stubPair(ISessionIndex, sessionIndexWithSummary('s1', '/tmp/proj')), + stubPair(IAgentLifecycleService, agentLifecycleWithMainStub()), + ]); + const original = await svc.create({ sessionId: 's1', workDir: '/tmp/proj' }); + + let closeStarted!: () => void; + let releaseClose!: () => void; + const closeBegan = new Promise((resolve) => { + closeStarted = resolve; + }); + const closeGate = new Promise((resolve) => { + releaseClose = resolve; + }); + original.accessor + .get(ISessionLifecycleHooks) + .onWillCloseSession.register('test-block', async (_event, next) => { + closeStarted(); + await closeGate; + await next(); + }); + + const closing = svc.close('s1'); + await closeBegan; + // The closing session is hidden from get/list immediately. + expect(svc.get('s1')).toBeUndefined(); + expect(svc.list()).toEqual([]); + + let settled = false; + const resumed = svc.resume('s1').then((handle) => { + settled = true; + return handle; + }); + await tick(); + expect(settled).toBe(false); + + releaseClose(); + await closing; + const fresh = await resumed; + expect(fresh).toBeDefined(); + expect(fresh).not.toBe(original); + expect(svc.get('s1')).toBe(fresh); + }); + + it('archive during a plain close waits it out and lands the archived flag on disk', async () => { + const docs = new Map(); + const docStore = { + _serviceBrand: undefined, + get: (scope: string, key: string) => Promise.resolve(docs.get(`${scope}/${key}`)), + set: (scope: string, key: string, value: unknown) => { + docs.set(`${scope}/${key}`, value); + return Promise.resolve(); + }, + delete: () => Promise.resolve(), + list: () => Promise.resolve([]), + watch: () => (_listener: unknown) => ({ dispose: () => {} }), + acquire: () => ({ dispose: () => {} }), + } as unknown as IAtomicDocumentStore; + const svc = await build([stubPair(IAtomicDocumentStore, docStore)]); + const original = await svc.create({ sessionId: 's1', workDir: '/tmp/proj' }); + const metaKey = `${original.accessor.get(ISessionContext).metaScope}/state.json`; + docs.set(metaKey, { id: 's1', version: 2, createdAt: 1, updatedAt: 1, archived: false }); + + let releaseClose!: () => void; + const closeGate = new Promise((resolve) => { + releaseClose = resolve; + }); + original.accessor + .get(ISessionLifecycleHooks) + .onWillCloseSession.register('test-block', async (_event, next) => { + await closeGate; + await next(); + }); + + const closing = svc.close('s1'); + await tick(); + const archived = svc.archive('s1'); + releaseClose(); + await closing; + await archived; + + // The archive must not ride the close: the marker lands on the persisted + // document after the close completes. + expect(docs.get(metaKey)).toMatchObject({ archived: true }); + }); + + it('still completes the teardown when a close hook fails', async () => { + const svc = await build(); + const original = await svc.create({ sessionId: 's1', workDir: '/tmp/proj' }); + original.accessor + .get(ISessionLifecycleHooks) + .onWillCloseSession.register('test-boom', async () => { + throw new Error('hook boom'); + }); + + await expect(svc.close('s1')).rejects.toThrow('hook boom'); + // The hook error reaches the caller, but the session is gone, not + // stranded as a half-closed zombie. + expect(svc.get('s1')).toBeUndefined(); + }); + + it('rejects a concurrent create with the same reserved session id', async () => { + const svc = await build(); + + const [first, second] = await Promise.allSettled([ + svc.create({ sessionId: 's1', workDir: '/tmp/proj' }), + svc.create({ sessionId: 's1', workDir: '/tmp/proj' }), + ]); + + expect([first.status, second.status].sort()).toEqual(['fulfilled', 'rejected']); + const rejection = [first, second].find((result) => result.status === 'rejected'); + expect((rejection as PromiseRejectedResult).reason).toMatchObject({ + code: ErrorCodes.SESSION_ALREADY_EXISTS, + }); + }); + + it('rejects a concurrent fork racing for the same reserved target id', async () => { + const svc = await build(); + await svc.create({ sessionId: 'src', workDir: '/tmp/proj' }); + + const [first, second] = await Promise.allSettled([ + svc.fork({ sourceSessionId: 'src', newSessionId: 'dst' }), + svc.fork({ sourceSessionId: 'src', newSessionId: 'dst' }), + ]); + + expect([first.status, second.status].sort()).toEqual(['fulfilled', 'rejected']); + const rejection = [first, second].find((result) => result.status === 'rejected'); + expect((rejection as PromiseRejectedResult).reason).toMatchObject({ + code: ErrorCodes.SESSION_ALREADY_EXISTS, + }); + }); + it('create seeds identity and materializes metadata', async () => { const svc = await build(); const h = await svc.create({ sessionId: 's1', workDir: '/tmp/proj' });