diff --git a/changelog.d/maintenance-stop-and-dispose.fixed.md b/changelog.d/maintenance-stop-and-dispose.fixed.md new file mode 100644 index 000000000..f55c657b0 --- /dev/null +++ b/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`. diff --git a/src/framework.ts b/src/framework.ts index 5b4f3785f..9e461bdeb 100644 --- a/src/framework.ts +++ b/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; + const pass = new Promise((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) { diff --git a/test/maintenance-timer.test.ts b/test/maintenance-timer.test.ts index 138a9b0ba..9ab61b645 100644 --- a/test/maintenance-timer.test.ts +++ b/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; + movedBeforeTickEnd?: boolean; + branchAtTickEnd?: string; + + checkReadiness(): ReadinessState { + return { ready: false, description: 'compression queued' }; + } + + async tick(_ctx: StrategyContext): Promise { + 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((resolve) => { this.open = resolve; }); + + release(): void { + this.open?.(); + } + + checkReadiness(): ReadinessState { + return { ready: this.ticks >= 3, description: 'counted ticks queued' }; + } + + async tick(_ctx: StrategyContext): Promise { + 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 | 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; + disposeConversationAgent(name: string): void; + noteInferenceExhausted(agentName: string, reason: string, retryable?: boolean, errorType?: string): void; + overBudgetDrains: Set>; + 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; + 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({