vrr/kimi-code
2
0
Fork 0
mirror of https://github.com/MoonshotAI/kimi-code.git synced 2026-08-10 09:16:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

fix: harden the session close/archive and create/fork lifecycle

Browse source 
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.
This commit is contained in:
7Sageer 2026-07-31 14:20:34 +08:00
parent 9a95e1c02f
commit e7c397a7c8
2 changed files with 230 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

118
packages/agent-core-v2/src/workspace/workspaceHandler/workspaceHandlerService.ts
View file

@ -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<CreateSessionOptions, 'sessionId'> & {
readonly sessionId: string;
};
type ClosingEntry = {
readonly kind: 'close' | 'archive';
readonly run: Promise<void>;
};
export class WorkspaceHandlerService extends Disposable implements IWorkspaceHandlerService {
declare readonly _serviceBrand: undefined;
private readonly sessions = new Map<string, ISessionScopeHandle>();
@ -173,8 +182,9 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan
private readonly _onDidForkSession = this._register(new Emitter<SessionForkedEvent>());
readonly onDidForkSession: Event<SessionForkedEvent> = this._onDidForkSession.event;
private readonly resuming = new Map<string, Promise<ISessionScopeHandle | undefined>>();
private readonly closing = new Map<string, Promise<void>>();
private readonly closing = new Map<string, ClosingEntry>();
private readonly sessionLifetimes = new Map<string, AbortController>();
private readonly reservedTargets = new Set<string>();
constructor(
@IInstantiationService private readonly instantiation: IInstantiationService,
@ -219,6 +229,24 @@ export class WorkspaceHandlerService extends Disposable implements IWorkspaceHan
async create(opts: CreateSessionOptions): Promise<ISessionScopeHandle> {
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<ISessionScopeHandle> {
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<void> {
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<void> {
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<void> {
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<void> {
@ -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<void>): Promise<void> {
this.closing.set(sessionId, run);
private async markArchivedOnDisk(sessionId: string): Promise<void> {
const scope = sessionScopeOf(this.handlerScope, sessionId);
const meta = await this.docs.get<SessionMeta>(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<void>,
): Promise<void> {
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);
}
}
}

134
packages/agent-core-v2/test/workspace/workspaceHandler/workspaceHandler.test.ts
View file

@ -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<void>((resolve) => {
closeStarted = resolve;
});
const closeGate = new Promise<void>((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<string, unknown>();
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<void>((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' });