fix(agent-core-v2): keep agent lifecycle context active through scope teardown

.changeset/fix-restore-crash-loop.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix repeated server crashes when resuming a session that was interrupted in the middle of a turn.

packages/agent-core-v2/src/_base/di/instantiation.ts

@@ -196,6 +196,7 @@ export interface IInstantiationService {
   provideAll(entries: ReadonlyArray<ProvideAllEntry>): void;
   unprovide<T>(id: ServiceIdentifier<T>): void;
   dispose(): void;
+  disposeAsync(): Promise<void>;
 }
 
 export const IInstantiationService: ServiceIdentifier<IInstantiationService> =

packages/agent-core-v2/src/_base/di/instantiationService.ts

@@ -487,6 +487,10 @@ export class InstantiationService implements IInstantiationService {
     return this._ledger.register(disposer, label);
   }
 
+  anchorKernelFinalizer(disposer: Disposer, label: string): LedgerEntry {
+    return this._ledger.registerFinalizer(disposer, label);
+  }
+
   private _getFiberHost(): FiberHost {
     this._fiberHost ??= {
       mintUid: () => ++this._root()._nextUnitUid,
@@ -648,18 +652,31 @@ export class InstantiationService implements IInstantiationService {
     return new InstantiationService(services, this._strict, this, this._enableTracing);
   }
 
+  private _disposePromise: Promise<void> | undefined;
+
   dispose(): void {
+    void this.disposeAsync();
+  }
+
+  disposeAsync(): Promise<void> {
+    this._disposePromise ??= this.disposeCore();
+    return this._disposePromise;
+  }
+
+  private disposeCore(): Promise<void> {
     if (this._disposed) {
-      return;
+      return Promise.resolve();
     }
     this._disposed = true;
 
+    const childTeardowns: Promise<void>[] = [];
+    let teardown: void | Promise<void> = undefined;
     try {
       for (const child of Array.from(this._children)) {
-        child.dispose();
+        childTeardowns.push(child.disposeAsync());
       }
       this._children.clear();
-      void this._ledger.teardown('scope-close');
+      teardown = this._ledger.teardown('scope-close');
       this._services.dispose();
       this.cascade.dispose();
       for (const view of this._collectionViews.values()) {
@@ -674,6 +691,7 @@ export class InstantiationService implements IInstantiationService {
         this._parent._children.delete(this);
       }
     }
+    return Promise.all([...childTeardowns, Promise.resolve(teardown)]).then(() => undefined);
   }
 
   private _createInstance<T>(ctor: any, args: unknown[], _trace: Trace, unit?: {

packages/agent-core-v2/src/_base/di/scope.ts

@@ -112,7 +112,7 @@ export interface IScopeHandle<K extends ScopeKind = ScopeKind> {
   readonly id: string;
   readonly kind: K;
   readonly accessor: ServicesAccessor;
-  dispose(): void;
+  dispose(): void | Promise<void>;
 }
 
 export type IAppScopeHandle = IScopeHandle<'app'>;
@@ -171,7 +171,7 @@ export function createScopedChildHandle(
     get: <T>(serviceId: ServiceIdentifier<T>): T =>
       child.invokeFunction((a) => a.get(serviceId)),
   };
-  return { id, kind, accessor, dispose: () => child.dispose() };
+  return { id, kind, accessor, dispose: () => child.disposeAsync() };
 }
 
 export class Scope implements IDisposable {

packages/agent-core-v2/src/_base/di/scopeUnits.ts

@@ -22,7 +22,7 @@ export function watchScopeUnits(container: InstantiationService, kind: ScopeKind
   const foldLedger = new Ledger(`scope-units:${kind}`);
   container.anchorKernelEntry((reason) => foldLedger.teardown(reason), `scope-units:${kind}`);
 
-  const materialized = new Map<number, () => void>();
+  const materialized = new Map<number, () => void | Promise<void>>();
 
   const materialize = (record: StoredRecord): void => {
     const recipe = record.value as ServiceRecipe;
@@ -32,7 +32,7 @@ export function watchScopeUnits(container: InstantiationService, kind: ScopeKind
       if (isClassRecipe(recipe)) {
         const instance = host.constructService(recipe, undefined) as Partial<IDisposable>;
         unitLedger.register(() => {
-          instance.dispose?.();
+          return instance.dispose?.();
         }, `unit:${name}`);
       } else {
         const facade = new FiberRuntime(
@@ -57,23 +57,23 @@ export function watchScopeUnits(container: InstantiationService, kind: ScopeKind
     }
 
     let retracted = false;
-    const retract = (): void => {
+    const retract = (): void | Promise<void> => {
       if (retracted) {
-        return;
+        return undefined;
       }
       retracted = true;
       materialized.delete(record.id);
-      void unitLedger.teardown('unload');
+      return unitLedger.teardown('unload');
     };
     if (!record.providerBook.isActive) {
-      retract();
+      void retract();
       return;
     }
     record.providerBook.register(() => {
-      retract();
+      void retract();
     }, `scope-units:${kind}`);
     foldLedger.register(() => {
-      retract();
+      return retract();
     }, `record:${name}`);
     materialized.set(record.id, retract);
   };
@@ -92,7 +92,7 @@ export function watchScopeUnits(container: InstantiationService, kind: ScopeKind
     }
     for (const [id, retract] of Array.from(materialized)) {
       if (!seen.has(id)) {
-        retract();
+        void retract();
       }
     }
   };

packages/agent-core-v2/src/_base/di/test.ts

@@ -29,7 +29,9 @@ export function createScopedTestHost(appStubs: ScopeSeed = []): ScopedTestHost {
           id: handle.id,
           kind: handle.kind,
           accessor: handle.accessor,
-          dispose: () => handle.dispose(),
+          dispose: () => {
+            void handle.dispose();
+          },
         } as Scope;
       }
       return app.createChild(kind, id, { seeds: stubs });

packages/agent-core-v2/src/_base/di/testInstantiationService.ts

@@ -262,6 +262,14 @@ export class TestInstantiationService extends InstantiationService implements ID
       super.dispose();
     }
   }
+
+  public override disposeAsync(): Promise<void> {
+    sinon.restore();
+    if (this._properDispose) {
+      return super.disposeAsync();
+    }
+    return Promise.resolve();
+  }
 }
 
 interface SinonOptions {

packages/agent-core-v2/src/_base/lifecycle/ledger.ts

@@ -65,6 +65,11 @@ export class Ledger {
     return this._push({ label, kind: 'disposer', active: true, run: disposer });
   }
 
+  registerFinalizer(disposer: Disposer, label: string = 'finalizer'): LedgerEntry {
+    this._assertActive('registerFinalizer');
+    return this._push({ label, kind: 'disposer', active: true, run: disposer }, true);
+  }
+
   effect(body: EffectBody, label: string = 'effect'): LedgerEntry {
     this._assertActive('effect');
     const out = body();
@@ -151,11 +156,15 @@ export class Ledger {
     return infos;
   }
 
-  private _push(record: EntryRecord): LedgerEntry {
+  private _push(record: EntryRecord, front = false): LedgerEntry {
     if (Ledger.captureStacks) {
       record.stack = new Error('Ledger registration').stack;
     }
-    this._records.push(record);
+    if (front) {
+      this._records.unshift(record);
+    } else {
+      this._records.push(record);
+    }
     return {
       label: record.label,
       get disposed() {

packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts

@@ -222,6 +222,7 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
     eventBus?.activateAgent(agent);
     let managed: ManagedAgent | undefined;
     let didCreate = false;
+    let finalizerArmed = false;
     try {
       const handle = createScopedChildHandle(
         this.instantiation,
@@ -237,13 +238,17 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
             }],
           ],
           configureContainer: (container) => {
+            container.anchorKernelFinalizer(() => {
+              eventBus?.deactivateAgent(agent);
+            }, 'agent-event-bus-deactivate');
+            finalizerArmed = true;
             this.adopt({
               id: agentId,
               kind: LifecycleScope.Agent,
               accessor: {
                 get: (id) => container.invokeFunction((accessor) => accessor.get(id)),
               },
-              dispose: () => { container.dispose(); },
+              dispose: () => container.disposeAsync(),
             });
             managed = this.roster.get(agentId);
           },
@@ -274,10 +279,10 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
         await managed.runtimeSet.close().catch(() => undefined);
         managed.killSpace();
         try {
-          managed.handle.dispose();
+          await managed.handle.dispose();
         } catch { }
       }
-      eventBus?.deactivateAgent(agent);
+      if (!finalizerArmed) eventBus?.deactivateAgent(agent);
       if (didCreate) this.onDidCloseEmitter.fire(agent);
       throw error;
     }
@@ -446,10 +451,7 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
     await Promise.all([loop.settled(), compactionSettled, prompt.drain(reason)]);
     await managed.runtimeSet.close();
     managed.killSpace();
-    handle.dispose();
-    this.instantiation.invokeFunction((accessor) =>
-      (accessor.get(ISessionEventBus) as ISessionEventBus | undefined)?.deactivateAgent(agent),
-    );
+    await handle.dispose();
     if (this.roster.get(agent.agentId) === managed) this.roster.delete(agent.agentId);
     this.onDidCloseEmitter.fire(agent);
   }

packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts

@@ -208,7 +208,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec
       const sessionDir = handle.accessor.get(ISessionContext).sessionDir;
       this.sessions.delete(sessionId);
       await this.drainAgents(handle).catch(() => {});
-      handle.dispose();
+      void handle.dispose();
       await this.hostFs.remove(sessionDir).catch(() => {});
       throw error;
     }
@@ -282,7 +282,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec
         this.pluginAgentProfileLoader.ready,
       ]);
     } catch (error) {
-      handle.dispose();
+      void handle.dispose();
       void this.explicitAgentProfileLoader.reload().catch(() => undefined);
       throw error;
     }
@@ -364,7 +364,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec
       await this.announceCreated({ sessionId, handle, source: 'resume' });
     } catch (error) {
       this.sessions.delete(sessionId);
-      handle.dispose();
+      void handle.dispose();
       throw error;
     }
     return handle;
@@ -387,7 +387,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec
     await this.appendLogStore.drainRetirements();
     await drainSessionMetadataWrites();
     await this.indexMirror.drain();
-    handle.dispose();
+    void handle.dispose();
     await drainLogCloses();
     this._onDidCloseSession.fire({ sessionId });
   }
@@ -408,7 +408,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec
     this.sessions.delete(sessionId);
     await drainSessionMetadataWrites();
     await this.indexMirror.drain();
-    handle.dispose();
+    void handle.dispose();
     await drainLogCloses();
     this._onDidArchiveSession.fire({ sessionId });
   }
@@ -589,7 +589,7 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec
       }
       if (target !== undefined) {
         try {
-          target.dispose();
+          void target.dispose();
         } catch {
         }
       }

packages/agent-core-v2/test/_base/di/child.test.ts

@@ -214,6 +214,64 @@ describe('InstantiationService.createChild', () => {
     expect(events).toEqual(['disposed']);
   });
 
+  it('repeated disposeAsync returns the in-flight teardown promise', async () => {
+    const events: string[] = [];
+    let releaseGate!: () => void;
+    const ix = new InstantiationService(new ServiceCollection());
+    ix.anchorKernelEntry(() => {
+      events.push('finalizer');
+    }, 'finalizer');
+    ix.anchorKernelEntry(() => {
+      events.push('gate-entered');
+      return new Promise<void>((resolve) => {
+        releaseGate = resolve;
+      });
+    }, 'gate');
+
+    const first = ix.disposeAsync();
+    const second = ix.disposeAsync();
+    let secondSettled = false;
+    void second.then(() => {
+      secondSettled = true;
+    });
+    await new Promise((resolve) => setTimeout(resolve, 10));
+    expect(events).toEqual(['gate-entered']);
+    expect(secondSettled).toBe(false);
+    releaseGate();
+    await Promise.all([first, second]);
+    expect(events).toEqual(['gate-entered', 'finalizer']);
+  });
+
+  it('disposeAsync awaits asynchronous child container teardown', async () => {
+    const events: string[] = [];
+    let releaseChildGate!: () => void;
+    const parent = new InstantiationService(new ServiceCollection());
+    const child = parent.createChild(new ServiceCollection()) as InstantiationService;
+    child.anchorKernelEntry(() => {
+      events.push('child-finalizer');
+    }, 'child-finalizer');
+    child.anchorKernelEntry(() => {
+      events.push('child-gate-entered');
+      return new Promise<void>((resolve) => {
+        releaseChildGate = resolve;
+      });
+    }, 'child-gate');
+    parent.anchorKernelEntry(() => {
+      events.push('parent-finalizer');
+    }, 'parent-finalizer');
+
+    let settled = false;
+    const disposal = parent.disposeAsync().then(() => {
+      settled = true;
+    });
+    await new Promise((resolve) => setTimeout(resolve, 10));
+    expect(events).toEqual(['child-gate-entered', 'parent-finalizer']);
+    expect(settled).toBe(false);
+    releaseChildGate();
+    await disposal;
+    expect(events).toEqual(['child-gate-entered', 'parent-finalizer', 'child-finalizer']);
+  });
+
   it('parent dispose propagates to children', () => {
     const events: string[] = [];
     interface IParentSvc {

packages/agent-core-v2/test/features/skill/workspace/skillCatalog.test.ts

@@ -185,7 +185,7 @@ function makeHost(
   });
   const disposeHost = host.dispose.bind(host);
   host.dispose = () => {
-    workspaceHandle.dispose();
+    void workspaceHandle.dispose();
     disposeHost();
   };
   return { host, workspace: workspaceHandle, config };