fix(maintenance): find a pass from inside it, and stop ticking a disposed agent

changelog.d/maintenance-stop-and-dispose.fixed.md

@@ -0,0 +1,2 @@
+- **A `stop()`, surgery or maintain begun from inside a context-maintenance pass now finds the pass.** The pass is registered before its body runs, so a `stop()`, `rollbackToMessage()`, `suppressMessages()`, `switchBranch()` or `maintenanceTick()` called from a trace listener on the `ops:alert` the pass raises while it prepares, or from a strategy's tick before that tick's first await, waits for the pass or joins it. A tick that awaits such a call waits on its own pass, as a tick already did for a call it made after its first await. Before, `stop()` closed the store while that tick was in flight, and the pass then met the closed store and reported it as the agent's maintenance failure ("Store has been closed"), with a `context-maintenance-failed` alert. A rollback moved the branch under the tick, and the pass went on ticking on the new branch. And `maintenanceTick()` started a second pass beside the first, so two ticks of one agent ran at once.
+- **A context-maintenance pass or an over-budget drain kick takes no further tick of an agent disposed while it runs,** such as a conversation fork closed after it went idle, or an ephemeral run that ended. A tick in flight at disposal finishes. Before, each went on ticking the disposed agent's context manager, up to eight ticks, for a context the framework runs no further turn on. The drain kick's log line says `stopped as the agent was disposed`.

src/framework.ts

@@ -3465,7 +3465,21 @@ export class AgentFramework {
 
   private startQueuedMaintenance(): void {
     if (this.maintenancePass || !this.running) return;
-    const pass = this.runQueuedMaintenance().catch((error) => {
+    // The pass is registered before its body runs, and the body still starts
+    // in this call. Its synchronous steps reach callers' code, an ops:alert
+    // listener while it prepares or a strategy's tick before that tick's
+    // first await, and a stop(), surgery or maintain begun there must find
+    // it: stop() and the surgery wait for it, and the maintain joins it
+    // rather than starting a second pass beside it. A tick that awaits such a
+    // call waits on its own pass, as a tick already did for a call it made
+    // after its first await.
+    let begin!: (body: Promise<void>) => void;
+    const pass = new Promise<void>((resolve) => { begin = resolve; });
+    this.maintenancePass = pass;
+    void pass.finally(() => {
+      if (this.maintenancePass === pass) this.maintenancePass = null;
+    });
+    begin(this.runQueuedMaintenance().catch((error) => {
       // The alert cuts its own copies from the whole message, so the
       // webhook's states the whole length.
       const wholeReason = error instanceof Error ? String(error.message) : failureDescription(error);
@@ -3480,11 +3494,7 @@ export class AgentFramework {
         wholeReason,
         { data: { scope: 'pass' } },
       );
-    });
-    this.maintenancePass = pass;
-    void pass.finally(() => {
-      if (this.maintenancePass === pass) this.maintenancePass = null;
-    });
+    }));
   }
 
   /** resume()'s refusal while a surgery holds the store. */
@@ -3601,12 +3611,15 @@ export class AgentFramework {
         // A tick whose provider call a surgery or stop() withdrew ends this
         // agent's work for the pass, as the hold or the stop would have before
         // its next tick: it is no failure, and what the failure reports
-        // remember stays as it is.
+        // remember stays as it is. An agent disposed while the pass runs, such
+        // as a conversation fork closed after it went idle, takes no further
+        // tick, since the framework runs no further turn for it.
         let withdrawn = false;
         try {
           for (
             let i = 0;
-            i < MAINTENANCE_TICKS_PER_PASS && this.running && !this.surgeryHold && !cm.isReady();
+            i < MAINTENANCE_TICKS_PER_PASS && this.running && !this.surgeryHold
+              && this.agents.get(agent.name) === agent && !cm.isReady();
             i++
           ) {
             await cm.tick();
@@ -16123,22 +16136,25 @@ export class AgentFramework {
     if (agent && drainState === 'kicking') {
       this.overBudgetDrainInFlight.add(agentName);
       // Like a maintenance pass: none starts once stop() has begun, it starts
-      // no tick while a surgery holds the store, and a surgery, a
-      // switchBranch or stop() waits for the tick in flight. It joins the
-      // drains they wait for before the alert is raised and before its body
-      // runs, so the alert's listeners find it: one that begins stop() or a
-      // surgery waits for it, and its body then ticks none and says why. So
-      // does a stop() made from inside its first tick, before that tick's
-      // first await.
+      // no tick while a surgery holds the store or once its agent has been
+      // disposed, and a surgery, a switchBranch or stop() waits for the tick
+      // in flight. It joins the drains they wait for before the alert is
+      // raised and before its body runs, so the alert's listeners find it:
+      // one that begins stop() or a surgery waits for it, and its body then
+      // ticks none and says why. So does a stop() made from inside its first
+      // tick, before that tick's first await.
       const drain = Promise.resolve().then(async () => {
         let ticks = 0;
         try {
           const cm = agent.getContextManager();
-          while (ticks < 8 && !this.providerAdmissionClosed && !this.surgeryHold) {
+          while (ticks < 8 && !this.providerAdmissionClosed && !this.surgeryHold && this.agents.get(agentName) === agent) {
             await cm.tick();
             ticks++;
           }
-          const stopped = ticks >= 8 ? '' : this.providerAdmissionClosed ? ', stopped as the framework stops' : ', stopped while a surgery holds the store';
+          const stopped = ticks >= 8 ? ''
+            : this.agents.get(agentName) !== agent ? ', stopped as the agent was disposed'
+            : this.providerAdmissionClosed ? ', stopped as the framework stops'
+            : ', stopped while a surgery holds the store';
           console.error(`[inference-failed] drain kicked for ${agentName} (OverBudget breaker, ${ticks} ticks${stopped})`);
         } catch (err) {
           if (err instanceof MaintenanceWithdrawnError) {

test/maintenance-timer.test.ts

@@ -399,6 +399,80 @@ class StopInTickStrategy implements ContextStrategy {
   }
 }
 
+/** A strategy whose first tick begins a rollback of `agent` before its first await, then notes whether the branch had moved when the tick ended. */
+class RollbackInTickStrategy implements ContextStrategy {
+  readonly name = 'rollback-in-tick-test';
+  ticks = 0;
+  framework?: AgentFramework;
+  messageId?: string;
+  moving?: Promise<unknown>;
+  movedBeforeTickEnd?: boolean;
+  branchAtTickEnd?: string;
+
+  checkReadiness(): ReadinessState {
+    return { ready: false, description: 'compression queued' };
+  }
+
+  async tick(_ctx: StrategyContext): Promise<void> {
+    this.ticks++;
+    if (this.ticks > 1) return;
+    this.moving = this.framework!.rollbackToMessage('agent', { messageId: this.messageId! });
+    // A rollback that doesn't wait for this tick settles first; one that
+    // does can't settle before the tick ends, so the timer wins.
+    const first = await Promise.race([
+      this.moving.then(() => 'moved' as const, () => 'moved' as const),
+      new Promise<'waited'>((resolve) => setTimeout(() => resolve('waited'), 1000)),
+    ]);
+    this.movedBeforeTickEnd = first === 'moved';
+    this.branchAtTickEnd = this.framework!.getAgent('agent')!.getContextManager().currentBranch().name;
+  }
+
+  select(
+    _store: MessageStoreView,
+    _log: ContextLogView,
+    _budget: TokenBudget,
+  ): ContextEntry[] {
+    return [];
+  }
+}
+
+/** A strategy that counts its ticks, and its ticks in flight at once. Its first tick waits on a gate, and it is ready after three. */
+class GatedCountingStrategy implements ContextStrategy {
+  readonly name = 'gated-counting-test';
+  ticks = 0;
+  inFlight = 0;
+  mostInFlight = 0;
+  private open?: () => void;
+  private readonly gate = new Promise<void>((resolve) => { this.open = resolve; });
+
+  release(): void {
+    this.open?.();
+  }
+
+  checkReadiness(): ReadinessState {
+    return { ready: this.ticks >= 3, description: 'counted ticks queued' };
+  }
+
+  async tick(_ctx: StrategyContext): Promise<void> {
+    this.ticks++;
+    this.inFlight++;
+    this.mostInFlight = Math.max(this.mostInFlight, this.inFlight);
+    try {
+      if (this.ticks === 1) await this.gate;
+    } finally {
+      this.inFlight--;
+    }
+  }
+
+  select(
+    _store: MessageStoreView,
+    _log: ContextLogView,
+    _budget: TokenBudget,
+  ): ContextEntry[] {
+    return [];
+  }
+}
+
 /** A strategy whose every tick meets the same failure, its second only once its gate opens. */
 class GatedRepeatFailureStrategy implements ContextStrategy {
   readonly name = 'gated-repeat-failure-test';
@@ -1801,6 +1875,209 @@ describe('queued context maintenance timer', () => {
     }
   });
 
+  it("a stop() made from a pass's first tick, before that tick's first await, waits for the tick", { timeout: 10_000 }, async () => {
+    const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-pass-stop-in-tick-'));
+    const strategy = new StopInTickStrategy();
+    const framework = await AgentFramework.create({
+      storePath: join(dir, 'store'),
+      membrane,
+      agents: [{ name: 'agent', model: 'test-model', systemPrompt: 'test', strategy }],
+      modules: [],
+      syncIntervalMs: 0,
+      maintenanceIntervalMs: 60_000,
+    });
+    strategy.framework = framework;
+    const alerts: string[] = [];
+    (framework as any).opsAlert = (kind: string) => { alerts.push(kind); };
+    const quiet = console.error;
+    console.error = () => {};
+    try {
+      // start() runs the first pass at once.
+      framework.start();
+      assert.ok(strategy.stopping, "the pass's first tick began stop()");
+      await strategy.stopping;
+      await waitFor(() => strategy.storeOpenAtTickEnd !== undefined, 3000);
+      assert.equal(strategy.stoppedBeforeTickEnd, false, 'stop() waited for the tick');
+      assert.equal(strategy.storeOpenAtTickEnd, true, 'the tick ended before the store closed');
+      assert.deepEqual(alerts, [], 'the pass met no failure');
+    } finally {
+      console.error = quiet;
+      await (strategy.stopping ?? framework.stop());
+      rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
+  it("a rollback begun from a pass's first tick, before that tick's first await, waits for the tick, and the pass takes no tick under its hold", { timeout: 10_000 }, async () => {
+    const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-pass-rollback-in-tick-'));
+    const strategy = new RollbackInTickStrategy();
+    const framework = await AgentFramework.create({
+      storePath: join(dir, 'store'),
+      membrane,
+      agents: [{ name: 'agent', model: 'test-model', systemPrompt: 'test', strategy }],
+      modules: [],
+      syncIntervalMs: 0,
+      maintenanceIntervalMs: 60_000,
+    });
+    strategy.framework = framework;
+    const cm = framework.getAgent('agent')!.getContextManager();
+    strategy.messageId = String(cm.addMessage('user', [{ type: 'text', text: 'm1' }]));
+    cm.addMessage('user', [{ type: 'text', text: 'm2' }]);
+    const quiet = console.error;
+    console.error = () => {};
+    try {
+      framework.start();
+      assert.ok(strategy.moving, "the pass's first tick began a rollback");
+      await strategy.moving;
+      await waitFor(() => strategy.branchAtTickEnd !== undefined, 3000);
+      assert.equal(strategy.movedBeforeTickEnd, false, 'the rollback waited for the tick');
+      assert.equal(strategy.branchAtTickEnd, 'main', 'the branch was where it was when the tick ended');
+      assert.notEqual(cm.currentBranch().name, 'main', 'the rollback went on once the tick ended');
+      assert.equal(strategy.ticks, 1, 'the pass took no tick under the hold');
+    } finally {
+      console.error = quiet;
+      await framework.stop();
+      rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
+  it('a maintain begun from an alert a pass raises while it prepares joins that pass, rather than starting a second beside it', { timeout: 10_000 }, async () => {
+    const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-maintain-in-alert-'));
+    const counting = new GatedCountingStrategy();
+    const framework = await AgentFramework.create({
+      storePath: join(dir, 'store'),
+      membrane,
+      agents: [
+        { name: 'unreadable', model: 'test-model', systemPrompt: 'test', strategy: new SayingStrategy() },
+        { name: 'agent', model: 'test-model', systemPrompt: 'test', strategy: counting },
+      ],
+      modules: [],
+      syncIntervalMs: 0,
+      maintenanceIntervalMs: 60_000,
+    });
+    let maintaining: Promise<unknown> | undefined;
+    // An ops:alert listener that runs maintain when the unreadable agent's
+    // failure is alerted, which the pass does while it prepares.
+    (framework as any).opsAlert = (kind: string, agentName: string) => {
+      if (kind === 'context-maintenance-failed' && agentName === 'unreadable') maintaining ??= framework.maintenanceTick();
+    };
+    const quiet = console.error;
+    console.error = () => {};
+    try {
+      framework.start();
+      assert.ok(maintaining, 'the pass alerted the failure while it prepared');
+      await waitFor(() => counting.ticks >= 1);
+      counting.release();
+      await maintaining;
+      assert.equal(counting.mostInFlight, 1, 'one tick of the agent at a time');
+      assert.equal(framework.getContextMaintenanceSnapshot().history.length, 1, 'one pass ran');
+    } finally {
+      console.error = quiet;
+      counting.release();
+      await framework.stop();
+      rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
+  it('an over-budget drain kick takes no tick of a conversation agent disposed before its first tick, and names the disposal though a surgery holds the store', { timeout: 10_000 }, async () => {
+    const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-drain-disposed-'));
+    const forks: GatedCountingStrategy[] = [];
+    const framework = await AgentFramework.create({
+      storePath: join(dir, 'store'),
+      membrane,
+      agents: [{ name: 'trunk', model: 'test-model', systemPrompt: 'test', strategy: new ReadyStrategy() }],
+      modules: [],
+      syncIntervalMs: 0,
+      maintenanceIntervalMs: 0,
+      conversations: {
+        templateAgent: 'trunk',
+        bind: { channel: 'mention' },
+        strategyFactory: () => { const fork = new GatedCountingStrategy(); forks.push(fork); return fork; },
+      },
+    });
+    const internals = framework as unknown as {
+      createConversationAgent(name: string, channelId: string): Promise<unknown>;
+      disposeConversationAgent(name: string): void;
+      noteInferenceExhausted(agentName: string, reason: string, retryable?: boolean, errorType?: string): void;
+      overBudgetDrains: Set<Promise<void>>;
+      surgeryHold: { verb: string; agentName: string; since: number } | null;
+    };
+    (framework as any).opsAlert = () => {};
+    const logged: string[] = [];
+    const quiet = console.error;
+    console.error = (...args: unknown[]) => { logged.push(args.map(String).join(' ')); };
+    try {
+      // The kick's body runs a microtask later, so each fork is disposed
+      // before the drain's first check.
+      await internals.createConversationAgent('trunk-chan-g1', 'world:chan');
+      internals.noteInferenceExhausted('trunk-chan-g1', 'compile refused', false, 'over_budget');
+      assert.equal(internals.overBudgetDrains.size, 1, 'a kick was started');
+      internals.disposeConversationAgent('trunk-chan-g1');
+      // Open the gate, so a drain that did tick shows its count.
+      forks[0]!.release();
+      await Promise.all([...internals.overBudgetDrains]);
+      assert.equal(forks[0]!.ticks, 0, 'no tick of the disposed agent');
+      assert.ok(
+        logged.some((line) => line.includes('drain kicked for trunk-chan-g1 (OverBudget breaker, 0 ticks, stopped as the agent was disposed)')),
+        logged.join('\n'),
+      );
+      // A surgery for another agent holding the store as well: the log names
+      // the disposal, the agent's own reason.
+      await internals.createConversationAgent('trunk-chan-g2', 'world:chan');
+      internals.noteInferenceExhausted('trunk-chan-g2', 'compile refused', false, 'over_budget');
+      internals.disposeConversationAgent('trunk-chan-g2');
+      internals.surgeryHold = { verb: 'rollback', agentName: 'trunk', since: Date.now() };
+      forks[1]!.release();
+      await Promise.all([...internals.overBudgetDrains]);
+      assert.equal(forks[1]!.ticks, 0, 'no tick of the second disposed agent');
+      assert.ok(
+        logged.some((line) => line.includes('drain kicked for trunk-chan-g2 (OverBudget breaker, 0 ticks, stopped as the agent was disposed)')),
+        logged.join('\n'),
+      );
+    } finally {
+      console.error = quiet;
+      internals.surgeryHold = null;
+      for (const fork of forks) fork.release();
+      await framework.stop();
+      rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
+  it('a pass takes no further tick of a conversation agent disposed while its tick was in flight', { timeout: 10_000 }, async () => {
+    const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-pass-disposed-'));
+    const fork = new GatedCountingStrategy();
+    const framework = await AgentFramework.create({
+      storePath: join(dir, 'store'),
+      membrane,
+      agents: [{ name: 'trunk', model: 'test-model', systemPrompt: 'test', strategy: new ReadyStrategy() }],
+      modules: [],
+      syncIntervalMs: 0,
+      maintenanceIntervalMs: 0,
+      conversations: { templateAgent: 'trunk', bind: { channel: 'mention' }, strategyFactory: () => fork },
+    });
+    const internals = framework as unknown as {
+      createConversationAgent(name: string, channelId: string): Promise<unknown>;
+      disposeConversationAgent(name: string): void;
+    };
+    const quiet = console.error;
+    console.error = () => {};
+    try {
+      await internals.createConversationAgent('trunk-chan-g1', 'world:chan');
+      framework.start();
+      const maintaining = framework.maintenanceTick();
+      await waitFor(() => fork.ticks === 1);
+      internals.disposeConversationAgent('trunk-chan-g1');
+      fork.release();
+      const snapshot = await maintaining;
+      assert.equal(fork.ticks, 1, 'no tick after the agent was disposed');
+      assert.equal(snapshot.history[0]?.agents.find((run) => run.agentName === 'trunk-chan-g1')?.ticks, 1, 'its run records the one tick');
+    } finally {
+      console.error = quiet;
+      fork.release();
+      await framework.stop();
+      rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
   it('no over-budget drain kick starts once stop() has begun', async () => {
     const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-drain-after-stop-'));
     const framework = await AgentFramework.create({