diff --git a/changelog.d/stop-budget-bounds-every-wait.changed.md b/changelog.d/stop-budget-bounds-every-wait.changed.md new file mode 100644 index 000000000..561dae38e --- /dev/null +++ b/changelog.d/stop-budget-bounds-every-wait.changed.md @@ -0,0 +1,32 @@ +- With a budget, `stop(teardownTimeoutMs)` now also bounds its waits on the + streams it cancels, on maintenance in flight and on the event loop, as it + did its waits on an operator's operations and on modules and MCPL servers + (#373). Each wait takes the whole budget, so a `stop()` with one can take + up to five budgets, besides the withheld-spill deadline of a guard's write + still under way. Before, a stream whose cancel couldn't reach it, a + maintenance tick that didn't settle, or a module's `onProcess` that never + returned kept `stop(t)` waiting without end, whatever `t` was. With no + budget, `stop()` waits for each as before. + - A wait whose budget runs out fails the `stop()` with a + `ShutdownTimeoutError` naming what it waited for (`stream:`, + `maintenance` or `event loop`), once teardown has gone on; with more + than one, or beside another failure, the `stop()` fails with an + `AggregateError` of them all. An owned store stays open, so that + `stop()` can be called again, as after a timed-out wait on an + operator's operations. + - A later `stop()` waits again, within its own budget, for a stream or an + event loop that an earlier one left running, and closes an owned store + only once they've settled. + - An agent whose stream outlasted its wait keeps its deferred writes + queued, logged once for that agent, where `stop()` would otherwise + store them ahead of what the stream may still write. The stream's own + end takes them, as at any turn's end, if the process lives on, and a + later `stop()` takes what that leaves. A memory-only queue loses them if + the process exits before then. + - Every ephemeral run still registered is ended with "refused: the + framework is stopping", even one whose stream is still running, and its + end stores what is queued for it, as at any run's end. Before, `stop()` + waited for that run's stream first, without end if it never settled. + - An event loop still in a call, such as a module's `onProcess`, when its + wait runs out finishes that call and the rest of its iteration, then + exits, even when a `start()` has been made meanwhile. diff --git a/src/framework.ts b/src/framework.ts index 7fe1bdda8..f1ef0ec3d 100644 --- a/src/framework.ts +++ b/src/framework.ts @@ -1154,13 +1154,13 @@ const DEFAULT_MAINTENANCE_INTERVAL_MS = 5000; /** Bound one pass so a large backlog yields to inference and other agents. */ const MAINTENANCE_TICKS_PER_PASS = 8; /** The teardownTimeoutMs a failed create() gives stop(): how long each of its - * cleanup waits, on operator operations and then on modules and MCPL servers - * together, may take. stop() leaves its deadline to its caller, who owns the + * waits may take. stop() leaves its deadline to its caller, who owns the * framework, but the caller of a failed create() never got one: unbounded, a * module whose stop() never settles would turn a failed boot that exits into - * one that hangs. stop()'s other waits are on streams, the event loop, - * maintenance and guard annotations, which a framework create() never - * started has none of. */ + * one that hangs. A framework create() never started has no streams, event + * loop, maintenance or guard annotations to wait on, so the waits that can + * take it are on operator operations and then on modules and MCPL servers + * together. */ const FAILED_CREATE_TEARDOWN_TIMEOUT_MS = 30_000; /** Local per-residence provider admission (AF #114 bounded first slice). */ @@ -1760,6 +1760,13 @@ export class AgentFramework { * start. stop() waits until this set is empty before it tears down, unless * its wait budget runs out. */ private readonly operatorOperations = new Set<{ label: string; promise: Promise }>(); + /** What a wait of stop()'s left running when its budget ran out, by its + * promise: a stream it cancelled or an event loop still in a call. A later + * stop() waits for each again, beside what it finds, whatever has happened + * since to the entry it was found under (a run's end removes its stream's, + * a start() replaces the loop), so a budget never lets the store close + * under work still running. Each leaves once it settles. */ + private readonly leftRunning = new Map, { label: string; kind: 'stream' | 'loop' }>(); private membrane: Membrane; private queue: ProcessQueueImpl; private agents: Map = new Map(); @@ -2329,11 +2336,11 @@ export class AgentFramework { * released: the framework it had begun is stopped, with * FAILED_CREATE_TEARDOWN_TIMEOUT_MS as stop()'s teardownTimeoutMs, and a * store create() opened is closed even when that stop failed or timed out. - * stop() keeps an owned store open when its waits on operator operations, - * modules or MCPL servers fail or time out, so that it can be called again, - * but no one can call it on a framework that create() never returned. When - * the release doesn't finish cleanly, the rejection carries create()'s error - * and then the release's, and its message says what didn't finish. + * stop() keeps an owned store open when any of its waits fails or times + * out, so that it can be called again, but no one can call it on a + * framework that create() never returned. When the release doesn't finish + * cleanly, the rejection carries create()'s error and then the release's, + * and its message says what didn't finish. */ private static async releaseFailedCreate( error: unknown, @@ -2881,16 +2888,34 @@ export class AgentFramework { } /** - * Stop the event loop. After a failed or timed-out wait on operator, module - * or MCPL cleanup, storage stays open and stop() can be called again. A - * timed-out wait lets teardown go on while that work may still act on what - * is being torn down. - * @param teardownTimeoutMs Optional budget for each wait on cleanup: an - * operator's operations already running, then modules and MCPL together. - * Omit to wait until cleanup settles, as required by existing module/child - * contracts. This does not bound earlier stream, event-loop or maintenance - * quiescence (a maintenance pass or over-budget drain kick in flight), or - * cancel cleanup. + * Stop the event loop. After a failed or timed-out wait, storage stays open + * and stop() can be called again. A timed-out wait lets teardown go on while + * what it waited for may still act on what is being torn down. + * @param teardownTimeoutMs Optional budget for each of stop()'s waits, in + * turn: on the streams it cancels, on maintenance in flight (a pass or an + * over-budget drain kick), on the event loop, on an operator's operations + * already running, then on modules and MCPL together. With a budget, a + * stop() can take up to five of them, besides the withheld-spill deadline + * of a guard's write still under way. Omit to wait until each settles, as + * required by existing module/child contracts. A wait whose budget runs out + * fails the stop() with a ShutdownTimeoutError naming what it waited for; + * with more than one, or beside another failure, stop() fails with an + * AggregateError of them all. After it: + * - an agent still registered whose stream is still running keeps its + * deferred writes queued, since the stream may still write. The stream's + * own turn-end flush takes them when it settles, if the process lives on, + * and a later stop() takes what that leaves; a memory-only queue loses + * them if the process exits first. + * - every ephemeral run still registered is ended, as after any stop(), and + * its end stores what is queued for it, as at any run's end, whatever its + * stream is doing: once a run has left, its stream's late events are + * dropped, though a write the stream already had under way can still + * land in the run's namespace after them. + * - an event loop still in a call, such as a module's onProcess, finishes + * that call and the rest of its iteration, and then exits. A start() + * made meanwhile runs its own loop beside that last iteration. + * - a later stop() waits again, within its own budget, for any stream or + * loop this one left running, so the store never closes under it. */ stop(teardownTimeoutMs = DEFAULT_SHUTDOWN_TIMEOUT_MS): Promise { try { validateShutdownTimeout(teardownTimeoutMs); } @@ -3027,10 +3052,52 @@ export class AgentFramework { if ('stream' in agent.state && agent.state.stream) agent.cancelStream(); } - // Wait for all stream iteration handles to settle - if (this.activeStreams.size > 0) { - await Promise.allSettled(this.activeStreams.values()); + // From here to the event loop, each wait is on work this stop() has + // asked to end, and each takes teardownTimeoutMs, as the waits on cleanup + // below do. One that runs out is reported with them once teardown has + // gone on, and keeps the store open for another stop(): what it waited + // for may still be running, and stays in leftRunning until it settles, + // for a later stop() to wait for again. A rejection settles a wait here + // as a resolution does: these are waits on work stop() has asked to end, + // not cleanup that can fail. + const waitsTimedOut: unknown[] = []; + type Running = { label: string; promise: Promise }; + const waitFor = async (attempts: Running[], kind?: 'stream' | 'loop'): Promise => { + if (attempts.length === 0) return; + try { + await waitForShutdown(attempts.map(({ label, promise }) => ({ label, promise: promise.then(() => {}, () => {}) })), teardownTimeoutMs); + } catch (error) { + waitsTimedOut.push(error); + // Each leaves leftRunning as it settles, those that already have at once. + if (kind) for (const { label, promise } of attempts) { + this.leftRunning.set(promise, { label, kind }); + const leave = () => { this.leftRunning.delete(promise); }; + void promise.then(leave, leave); + } + } + }; + // What this stop() found of a kind, beside what an earlier one's budget + // left running, each promise once. + const withLeftovers = (kind: 'stream' | 'loop', found: Running[]): Running[] => { + const all = new Map, string>(found.map(({ label, promise }) => [promise, label])); + for (const [promise, entry] of this.leftRunning) { + if (entry.kind === kind && !all.has(promise)) all.set(promise, entry.label); + } + return [...all].map(([promise, label]) => ({ label, promise })); + }; + + // Wait for all stream iteration handles to settle. An agent whose handle + // hasn't settled by the final flush keeps its entries queued there. Its + // activeStreams entry alone doesn't say that a stream is still running: + // a restart whose compile failed after taking the stream id leaves the + // old frame's entry, which that frame no longer owns. + const stillStreaming = new Set(); + for (const [name, stream] of this.activeStreams) { + stillStreaming.add(name); + const settled = () => { stillStreaming.delete(name); }; + void stream.then(settled, settled); } + await waitFor(withLeftovers('stream', [...this.activeStreams].map(([name, promise]) => ({ label: `stream:${name}`, promise }))), 'stream'); // Stop sync timer if (this.syncTimer) { @@ -3043,23 +3110,29 @@ export class AgentFramework { this.maintenanceTimer = null; } // The pass and any over-budget drain kick: a compression in flight lands - // before the store closes. - await this.settleMaintenancePass(); - - if (this.loopPromise) { - await this.loopPromise; - } - - // Nothing drives an ephemeral run once the loop has ended: no turn starts - // after stop() begins, every stream it found was cancelled and has - // settled, and a tool result that arrives now is dropped, since the queue - // is closed. So a run still registered could only wait for a watchdog, - // such as one whose wake was never taken up, whose turn was refused, - // whose retry was cancelled, or whose continuation is still queued. End - // each one now, while the store is open for its cleanup. (A start() + // before the store closes. Its bookkeeping outlives a timeout, since no + // pass starts while one runs and a drain leaves its set only as it + // settles, so it needs no leftRunning entry. + await waitFor(this.maintenanceRunning() ? [{ label: 'maintenance', promise: this.settleMaintenancePass() }] : []); + + await waitFor(withLeftovers('loop', this.loopPromise ? [{ label: 'event loop', promise: this.loopPromise }] : []), 'loop'); + + // Once the waits above have settled, nothing drives an ephemeral run: no + // turn starts after stop() begins, every stream it found was cancelled + // and has settled, and a tool result that arrives now is dropped, since + // the queue is closed. So a run still registered could only wait for a + // watchdog, such as one whose wake was never taken up, whose turn was + // refused, whose retry was cancelled, or whose continuation is still + // queued. End each one now, while the store is open for its cleanup. + // After a wait whose budget ran out, a run's stream or the loop may still + // be running, and every run is ended all the same: the caller asked for a + // bounded stop(), and a run whose stream is still running would otherwise + // wait for its idle watchdog. Its end stores its queue, as at any run's + // end; that stream's late events are dropped once the run has left, and a + // later stop() waits for the stream itself (leftRunning). (A start() // during a stop() is #315's. One made before the loop ended keeps a loop - // running, so this stop waits for it; one made just after, before this - // point, has its runs ended here too.) + // running, so this stop waits for it, up to its budget; one made just + // after, before this point, has its runs ended here too.) for (const [name, run] of this.ephemeralRuns) run.settle.reject(ephemeralRunStopped(name)); // Admission closed when stop() began. What admitted operations start now, @@ -3102,7 +3175,10 @@ export class AgentFramework { } // Timed-out or failed teardown retains its context and may still write. // Sync available state best-effort, but leave its store open for a retry. - const shutdownErrors: unknown[] = operatorOperationsTimedOut === undefined ? [] : [operatorOperationsTimedOut]; + const shutdownErrors: unknown[] = [ + ...waitsTimedOut, + ...(operatorOperationsTimedOut === undefined ? [] : [operatorOperationsTimedOut]), + ]; for (const result of await Promise.allSettled(shutdownPromises)) { if (result.status === 'rejected') shutdownErrors.push(result.reason); } @@ -3127,9 +3203,10 @@ export class AgentFramework { } } - // No stream is left: from the start of stop() none registers and no - // agent turn starts (providerAdmissionClosed, checked for every agent), - // and every stream registered before was cancelled and awaited above. + // No stream is left, unless its wait above ran out: from the start of + // stop() none registers and no agent turn starts (providerAdmissionClosed, + // checked for every agent), and every stream registered before was + // cancelled and awaited above. // (A start() made before this stop finishes reopens admission, so a turn // it begins can be alive here and take one of these entries mid-stream; // #315 asks how start() after stop() should behave.) @@ -3157,10 +3234,33 @@ export class AgentFramework { // disposed fork) is left: dropped with a memory-only queue at exit, and // logged and dropped from a persisted one by the boot or resume flush // that meets it. + // + // A registered agent's stream that outlasted its wait may still write, + // so that agent's entries stay queued, said once for it: no entry lands + // ahead of a stream that may still write. The stream counts as running + // while the handle this stop() waited on for it hasn't settled + // (stillStreaming). When it settles in a process that lives on, its own + // turn-end flush takes them, as at any turn's end, and a later stop() + // takes what that leaves. Until then a persisted queue keeps them for + // boot, and a memory-only one loses them at exit. if (this.deferredMessages.length > 0) { - const entries = [...this.agents.keys()].flatMap((name) => this.drainDeferredFor(name)); + const registered = [...this.agents.keys()]; + const entries = registered + .filter((name) => !stillStreaming.has(name)) + .flatMap((name) => this.drainDeferredFor(name)); this.readdDrained(entries, 'stop', { bypassDeferral: !this.quiesced }); this.ackDeferredWrites(); + for (const name of registered.filter((agent) => stillStreaming.has(agent))) { + const held = this.deferredMessages.filter((msg) => (msg.forAgent ?? this.primaryAgentName) === name).length; + if (held > 0) { + console.error( + `[deferral] stop: ${held} deferred write(s) for ${name} stay queued: its stream outlasted stop()'s wait and may still write; ` + + (this.deferredWritesPersisted || this.quiesced + ? 'the persisted queue keeps them for boot' + : 'this memory-only queue loses them if the process exits before that stream settles'), + ); + } + } } // Final sync before closing @@ -9459,7 +9559,13 @@ export class AgentFramework { } private async runLoop(): Promise { - while (this.running) { + // A loop ends with its lifecycle. A stop() whose budget ran out while + // this loop was in a call (a module's onProcess) returns before the call + // does, and a start() after it begins a lifecycle with a loop of its own; + // this one finishes its iteration and exits, rather than running beside + // that loop for good. + const lifecycle = this.lifecycleGeneration; + while (this.running && this.lifecycleGeneration === lifecycle) { try { await this.processNextEvent(); } catch (error) { diff --git a/test/ephemeral-settle.test.ts b/test/ephemeral-settle.test.ts index 884392c17..ebbeb2114 100644 --- a/test/ephemeral-settle.test.ts +++ b/test/ephemeral-settle.test.ts @@ -16,7 +16,7 @@ import type { ToolDefinition, ToolResult, } from '../src/index.js'; -import { AgentFramework } from '../src/index.js'; +import { AgentFramework, ShutdownTimeoutError } from '../src/index.js'; import { createMockResponse, HeldStream, MockMembrane, MockYieldingStream } from './helpers/mock-membrane.js'; class EphemeralToolModule implements Module { @@ -514,6 +514,76 @@ describe('runEphemeralToCompletion settle signal', () => { } }); + it("a run whose stream outlasts stop()'s budget is ended by that stop() with its queue stored, and a later stop() waits for its stream", { timeout: 60_000 }, async () => { + // The run's stream yields nothing until the test releases it, cancel + // included, so stop(100)'s wait on it runs out. The caller asked for a + // bounded stop(), so its sweep ends the run then, rather than leave it to + // its idle watchdog, fifteen minutes by default, and the run's end stores + // what was queued for it, as at any run's end. The run's end removes its + // stream's activeStreams entry, but the stream is still running, so a + // second stop() waits for it too and keeps the store open until it has + // settled. Once it does, what it yields late is dropped: the run has left. + const { tempDir, storePath } = tempStorePath('ephemeral-stop-budget-'); + let release!: () => void; + const released = new Promise((resolve) => { release = resolve; }); + const membrane = new MockMembrane(); + membrane.streamYielding = (request) => { + membrane.calls.push(request); + return new HeldStream(released, [createMockResponse([{ type: 'text', text: 'too late' }])]); + }; + const framework = await AgentFramework.create({ + storePath, + membrane: membrane.asMembrane(), + agents: [], + modules: [new EphemeralToolModule()], + syncIntervalMs: 0, + }); + const within = (promise: Promise): Promise => { + let timer: ReturnType | undefined; + return Promise.race([ + promise.then(() => 'resolved', (error: unknown) => error), + new Promise((resolve) => { timer = setTimeout(() => resolve('still pending'), 5_000); }), + ]).finally(() => clearTimeout(timer)); + }; + const internals = framework as unknown as { + activeStreams: Map; + leftRunning: Map, unknown>; + addMessage(participant: string, content: Array<{ type: 'text'; text: string }>, metadata?: undefined, opts?: { forAgent?: string }): string; + }; + try { + const { agent, contextManager } = await createEphemeral(framework); + const stored = (): string[] => contextManager.getAllMessages() + .flatMap((m) => m.content.map((b) => (b.type === 'text' ? (b as { text: string }).text : `<${b.type}>`))); + framework.start(); + const run = framework.runEphemeralToCompletion(agent, contextManager); + void run.catch(() => {}); + for (let i = 0; i < 400 && !internals.activeStreams.has(agent.name); i++) await new Promise((resolve) => setTimeout(resolve, 5)); + assert.ok(internals.activeStreams.has(agent.name), "the run's stream started"); + assert.equal(internals.addMessage('Antra', [{ type: 'text', text: 'look left' }], undefined, { forAgent: agent.name }), '', + 'deferred: the run is streaming'); + const stopped = await within(framework.stop(100)); + assert.ok(stopped instanceof ShutdownTimeoutError + && stopped.message === `Shutdown timed out after 100 ms; pending: stream:${agent.name}`, String(stopped)); + const ended = await within(run); + assert.ok(ended instanceof Error && /refused: the framework is stopping/.test(ended.message), String(ended)); + assert.equal(framework.getAgent(agent.name), null, 'the run has left'); + assert.ok(stored().includes('look left'), `its end stored its queue:\n${stored().join('\n')}`); + const again = await within(framework.stop(100)); + assert.ok(again instanceof ShutdownTimeoutError + && again.message === `Shutdown timed out after 100 ms; pending: stream:${agent.name}`, String(again)); + assert.equal(framework.getStore().isClosed(), false, 'the run\'s stream is still running'); + release(); + for (let i = 0; i < 400 && internals.leftRunning.size > 0; i++) await new Promise((resolve) => setTimeout(resolve, 5)); + assert.equal(internals.leftRunning.size, 0, 'the stream settled'); + assert.ok(!stored().includes('too late'), `what it yielded once the run had left was dropped:\n${stored().join('\n')}`); + await framework.stop(); + assert.equal(framework.getStore().isClosed(), true, 'and a stop() closes the store once it has'); + } finally { + release(); + await stopAndRemove(framework, tempDir); + } + }); + it('resolves plain text completion and deregisters the agent', async () => { const { tempDir, storePath } = tempStorePath('ephemeral-plain-'); const membrane = new MockMembrane(); diff --git a/test/maintenance-timer.test.ts b/test/maintenance-timer.test.ts index 9ab61b645..eab606a85 100644 --- a/test/maintenance-timer.test.ts +++ b/test/maintenance-timer.test.ts @@ -22,7 +22,7 @@ import type { ToolDefinition, ToolResult, } from '../src/index.js'; -import { AgentFramework } from '../src/index.js'; +import { AgentFramework, ShutdownTimeoutError } from '../src/index.js'; import { OperatorActionError } from '../src/operator-log.js'; import type { NormalizedRequest, NormalizedResponse } from '@animalabs/membrane'; import { MembraneError } from '@animalabs/membrane'; @@ -2234,6 +2234,100 @@ describe('queued context maintenance timer', () => { } }); + describe("stop(t) with a maintenance tick that outlasts its budget", () => { + /** stop(t)'s outcome: its rejection, 'resolved', or 'still pending' after + * 5 s, so a stop() that waits on the tick fails the test rather than + * hanging it. */ + async function stopWithin(framework: AgentFramework, budget?: number): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + framework.stop(budget).then(() => 'resolved', (error: unknown) => error), + new Promise((resolve) => { timer = setTimeout(() => resolve('still pending'), 5_000); }), + ]); + } finally { clearTimeout(timer); } + } + const maintenancePending = (outcome: unknown): boolean => outcome instanceof ShutdownTimeoutError + && outcome.message === 'Shutdown timed out after 100 ms; pending: maintenance'; + + it("a drain kick's tick: stop(t) fails naming maintenance, keeps the store open, and a later stop(t) waits for the same tick again", { timeout: 20_000 }, async () => { + const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-stop-budget-')); + const strategy = new GatedTickStrategy(); + const framework = await AgentFramework.create({ + storePath: join(dir, 'store'), + membrane, + agents: [{ name: 'agent', model: 'test-model', systemPrompt: 'test', strategy }], + modules: [], + syncIntervalMs: 0, + maintenanceIntervalMs: 0, + }); + strategy.framework = framework; + const internals = framework as unknown as { + noteInferenceExhausted(agentName: string, reason: string, retryable?: boolean, errorType?: string): void; + }; + const quiet = console.error; + console.error = () => {}; + try { + framework.start(); + internals.noteInferenceExhausted('agent', 'compile refused', false, 'over_budget'); + await waitFor(() => strategy.ticks === 1); + const outcome = await stopWithin(framework, 100); + assert.ok(maintenancePending(outcome), String(outcome)); + assert.equal(framework.getStore().isClosed(), false, 'the tick may still write'); + const again = await stopWithin(framework, 100); + assert.ok(maintenancePending(again), `a later stop() waits for the tick still held: ${String(again)}`); + assert.equal(framework.getStore().isClosed(), false, 'the tick is still held'); + strategy.release(); + await framework.stop(); + assert.equal(strategy.storeOpenAtFirstTickEnd, true, 'the tick ended before the store closed'); + assert.equal(framework.getStore().isClosed(), true, 'a later stop() closes the store'); + } finally { + console.error = quiet; + strategy.release(); + await framework.stop().catch(() => {}); + rmSync(dir, { recursive: true, force: true }); + } + }); + + it("a regular pass's tick held across a start(): no second pass starts beside it, and a later stop(t) waits for it again", { timeout: 20_000 }, async () => { + // start() runs the first pass at once; the interval never fires here. + const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-pass-stop-budget-')); + const strategy = new GatedTickStrategy(); + 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 quiet = console.error; + console.error = () => {}; + try { + framework.start(); + await waitFor(() => strategy.ticks === 1); + const outcome = await stopWithin(framework, 100); + assert.ok(maintenancePending(outcome), String(outcome)); + framework.start(); + await new Promise((resolve) => setTimeout(resolve, 200)); + assert.equal(strategy.ticks, 1, 'no pass started beside the held one'); + const again = await stopWithin(framework, 100); + assert.ok(maintenancePending(again), `a later stop() waits for the pass still held: ${String(again)}`); + assert.equal(framework.getStore().isClosed(), false, 'the tick is still held'); + strategy.release(); + await framework.stop(); + assert.equal(strategy.storeOpenAtFirstTickEnd, true, 'the tick ended before the store closed'); + assert.equal(framework.getStore().isClosed(), true, 'a later stop() closes the store'); + } finally { + console.error = quiet; + strategy.release(); + await framework.stop().catch(() => {}); + rmSync(dir, { recursive: true, force: true }); + } + }); + }); + it('no pass starts while a surgery holds the store', async () => { const dir = mkdtempSync(join(tmpdir(), 'af-maintenance-surgery-hold-')); const strategy = new QueuedStrategy(); diff --git a/test/stop-with-work-in-flight.test.ts b/test/stop-with-work-in-flight.test.ts index bfc28ddfa..ac513baf2 100644 --- a/test/stop-with-work-in-flight.test.ts +++ b/test/stop-with-work-in-flight.test.ts @@ -22,6 +22,13 @@ * retry follows its inference:failed, whether its stream failed to start or * failed once started. A turn that continues after that start() is marked * inferring again. + * Work that outlasts stop()'s budget fails the stop, with teardown gone on + * and the store left open (#373): a stream whose cancel can't reach it, + * whose agent's deferred writes wait for its own end, and an event loop in a + * module's onProcess, which finishes that iteration once the call returns + * and exits, rather than run beside the loop of a start() made meanwhile. + * Two such waits fail it together, and a later stop() waits for what an + * earlier one left running, a start() between them or not. * Each test waits on the event it needs (a log line, an answer), not on a * fixed delay, and then one more turn, so a late rejection would be seen. */ @@ -35,7 +42,7 @@ import { join } from 'node:path'; import { fileURLToPath } from 'node:url'; import { parse } from '@babel/parser'; import { JsStore } from '@animalabs/chronicle'; -import type { ContentBlock } from '@animalabs/membrane'; +import type { ContentBlock, StreamEvent } from '@animalabs/membrane'; import type { ErrorAction, Module, @@ -47,7 +54,7 @@ import type { ToolDefinition, ToolResult, } from '../src/index.js'; -import { AgentFramework } from '../src/index.js'; +import { AgentFramework, ShutdownTimeoutError } from '../src/index.js'; import { HeldStream, MockMembrane, MockYieldingStream, createMockResponse } from './helpers/mock-membrane.js'; /** Resolve once `cond` holds, polling each turn; fail, naming `what`, after `ms`. */ @@ -509,6 +516,312 @@ describe('stop() with work still in flight', () => { } }); + describe("stop()'s budget, with work that outlasts it", () => { + /** What `promise` has come to after `ms`: its rejection, 'resolved', or + * 'still pending', so a stop() that doesn't come back fails the test + * rather than hanging it. */ + async function settledWithin(promise: Promise, ms: number): Promise { + let timer: ReturnType | undefined; + const late = new Promise((resolve) => { timer = setTimeout(() => resolve('still pending'), ms); }); + try { + return await Promise.race([promise.then(() => 'resolved', (error: unknown) => error), late]); + } finally { clearTimeout(timer); } + } + + /** + * A resident whose stream, once started, yields nothing until the test + * releases it, cancel included (HeldStream), and a module whose onProcess + * waits on any text that starts with "hang" until the test lets it go. + * Any other text is added for the resident and asks for its inference. + */ + async function frameworkWithHeldWork() { + let release!: () => void; + const released = new Promise((resolve) => { release = resolve; }); + let streaming!: () => void; + const residentStreaming = new Promise((resolve) => { streaming = resolve; }); + const membrane = new MockMembrane(); + membrane.streamYielding = (request) => { + membrane.calls.push(request); + streaming(); + return new HeldStream(released, [createMockResponse([{ type: 'text', text: 'resident done' }] as ContentBlock[])]); + }; + const entered: string[] = []; + const hung = new Map void>(); + const order: string[] = []; + const events: Module = { + name: 'events', + async start() {}, + async stop() { order.push('modules stopped'); }, + getTools: () => [], + async handleToolCall() { return { success: true }; }, + async onProcess(event) { + if (event.type !== 'external-message') return {}; + const text = String((event as { content?: unknown }).content); + if (text.startsWith('hang')) { + entered.push(text); + await new Promise((resolve) => { hung.set(text, resolve); }); + return {}; + } + return { addMessages: [{ participant: 'User', content: [{ type: 'text', text }] }], requestInference: true }; + }, + }; + const framework = await AgentFramework.create({ + storePath: join(tempDir, 'store'), + membrane: membrane.asMembrane(), + agents: [{ name: 'assistant', model: 'test-model', systemPrompt: 'test' }], + modules: [events], + syncIntervalMs: 0, + }); + const push = (content: string) => + framework.pushEvent({ type: 'external-message', source: 'test', content, metadata: {} } as unknown as ProcessEvent); + const letGo = (text: string) => { hung.get(text)?.(); }; + const internals = framework as unknown as { deferredMessages: Array<{ forAgent?: string }>; loopPromise: Promise | null }; + const storedTexts = (): string[] => framework.getAgent('assistant')!.getContextManager().getAllMessages() + .flatMap((m) => m.content.map((b) => (b.type === 'text' ? (b as { text: string }).text : `<${b.type}>`))); + return { framework, release, residentStreaming, entered, letGo, order, push, internals, storedTexts }; + } + + it("a stream that outlasts the budget fails the stop, keeps the store open, and keeps its agent's deferred write queued until its own end stores it", { timeout: 30_000 }, async () => { + // The deferral's rule holds through the timeout: no entry lands ahead of + // a stream that may still write. + const { framework, release, residentStreaming, order, push, internals, storedTexts } = await frameworkWithHeldWork(); + try { + framework.start(); + push('go'); + await residentStreaming; + push('a note for later'); + await until(() => internals.deferredMessages.length === 1, 'the note to be deferred behind the stream'); + const outcome = await settledWithin(framework.stop(100), 5_000); + assert.ok(outcome instanceof ShutdownTimeoutError + && outcome.message === 'Shutdown timed out after 100 ms; pending: stream:assistant', String(outcome)); + assert.ok(order.includes('modules stopped'), 'teardown went on'); + assert.equal(framework.getStore().isClosed(), false, 'the stream may still write'); + assert.equal(internals.deferredMessages.length, 1, 'the note stays queued'); + assert.ok(!storedTexts().includes('a note for later'), 'nothing landed ahead of the stream'); + assert.ok(logged.includes( + "[deferral] stop: 1 deferred write(s) for assistant stay queued: its stream outlasted stop()'s wait and may still write; " + + 'this memory-only queue loses them if the process exits before that stream settles'), logged.join('\n')); + release(); + await until(() => internals.deferredMessages.length === 0, "the stream's own end to take the note"); + assert.ok(storedTexts().includes('a note for later'), `its turn-end flush stored the note:\n${storedTexts().join('\n')}`); + await framework.stop(); + assert.equal(framework.getStore().isClosed(), true, 'a later stop() closes the store'); + assert.deepEqual(rejections, []); + } finally { + release(); + await framework.stop().catch(() => {}); + } + }); + + it("a write deferred behind the entry a retry that failed to compile left is stored by stop(), with a budget or without", { timeout: 30_000 }, async () => { + // A resident stream fails and its error policy retries, but the retry's + // compile throws after taking the agent's stream id. The old frame no + // longer owns the activeStreams entry, so its finally leaves it, and no + // new handle replaces it: the entry outlives its stream. No stream that + // stop() waited on is still running, so the final flush stores a write + // deferred behind that entry, as on main. + class FailingStream extends MockYieldingStream { + constructor(private readonly released: Promise) { + super([createMockResponse([{ type: 'text', text: 'unused' }] as ContentBlock[])]); + } + override async *[Symbol.asyncIterator](): AsyncGenerator { + await this.released; + yield { type: 'error', error: new Error('the provider failed') } as unknown as StreamEvent; + } + } + const agents = [{ name: 'assistant', model: 'test-model', systemPrompt: 'test' }]; + const textsIn = (framework: AgentFramework): string[] => framework.getAgent('assistant')!.getContextManager().getAllMessages() + .flatMap((m) => m.content.map((b) => (b.type === 'text' ? (b as { text: string }).text : `<${b.type}>`))); + for (const budget of [undefined, 100]) { + let release!: () => void; + const released = new Promise((resolve) => { release = resolve; }); + const membrane = new MockMembrane(); + membrane.streamYielding = (request) => { + membrane.calls.push(request); + return new FailingStream(released); + }; + const events: Module = { + name: 'events', + async start() {}, + async stop() {}, + getTools: () => [], + async handleToolCall() { return { success: true }; }, + async onProcess(event) { + if (event.type !== 'external-message') return {}; + return { addMessages: [{ participant: 'User', content: [{ type: 'text', text: 'go' }] }], requestInference: true }; + }, + }; + const attempts: number[] = []; + const storePath = join(tempDir, `store-${budget ?? 'unbounded'}`); + const framework = await AgentFramework.create({ + storePath, + membrane: membrane.asMembrane(), + agents, + modules: [events], + syncIntervalMs: 0, + errorPolicy: { + maxRetries: 1, + onInferenceError: (_error: Error, _agent: string, attempt: number): ErrorAction => { + attempts.push(attempt); + return attempt < 1 ? { retry: true, delayMs: 0 } : { retry: false }; + }, + }, + }); + const internals = framework as unknown as { + activeStreams: Map>; + deferredMessages: unknown[]; + addMessage(participant: string, content: ContentBlock[]): string; + }; + try { + framework.start(); + framework.pushEvent({ type: 'external-message', source: 'test', content: 'go', metadata: {} } as unknown as ProcessEvent); + await until(() => internals.activeStreams.has('assistant'), 'the first stream'); + const handle = internals.activeStreams.get('assistant')!; + (framework.getAgent('assistant') as unknown as { buildActivationRequest: () => never }).buildActivationRequest = () => { + throw new Error('the compile failed'); + }; + release(); + await handle; + await until(() => attempts.length === 2, 'the retry to fail'); + assert.equal(internals.activeStreams.get('assistant'), handle, 'the old frame left its entry'); + assert.equal(internals.addMessage('Antra', [{ type: 'text', text: 'late note' }] as ContentBlock[]), '', 'deferred behind it'); + await framework.stop(budget); + assert.equal(internals.deferredMessages.length, 0, `budget ${budget}: the final flush took the write`); + assert.ok(!logged.some((line) => line.startsWith('[deferral] stop:')), logged.join('\n')); + } finally { + release(); + await framework.stop().catch(() => {}); + } + const reopened = await AgentFramework.create({ storePath, membrane: new MockMembrane().asMembrane(), agents, modules: [], syncIntervalMs: 0 }); + try { + assert.ok(textsIn(reopened).includes('late note'), `budget ${budget}: the store has the write:\n${textsIn(reopened).join('\n')}`); + } finally { + await reopened.stop(); + } + } + assert.deepEqual(rejections, []); + }); + + it("a stream handle that rejects settles stop()'s wait on it, as on main, with a budget or without", { timeout: 30_000 }, async () => { + // A driver that crashed rather than ended: stop() waited on the + // streams with Promise.allSettled, so a rejection never failed it. + for (const budget of [undefined, 100]) { + const { framework } = await frameworkWithHeldWork(); + const streams = (framework as unknown as { activeStreams: Map> }).activeStreams; + const crashed = Promise.reject(new Error('the driver crashed')); + crashed.catch(() => {}); + streams.set('assistant', crashed); + try { + framework.start(); + await framework.stop(budget); + assert.equal(framework.getStore().isClosed(), true, `budget ${budget}: the stop finished`); + } finally { + streams.delete('assistant'); + await framework.stop().catch(() => {}); + } + } + assert.deepEqual(rejections, []); + }); + + it("an event loop still in a module's onProcess when the budget runs out fails the stop, and teardown goes on", { timeout: 30_000 }, async () => { + const { framework, entered, letGo, order, push } = await frameworkWithHeldWork(); + try { + framework.start(); + push('hang'); + await until(() => entered.includes('hang'), 'the loop to enter the call'); + const outcome = await settledWithin(framework.stop(100), 5_000); + assert.ok(outcome instanceof ShutdownTimeoutError + && outcome.message === 'Shutdown timed out after 100 ms; pending: event loop', String(outcome)); + assert.ok(order.includes('modules stopped'), 'teardown went on, the module whose call is still running included'); + assert.equal(framework.getStore().isClosed(), false); + letGo('hang'); + await framework.stop(); + assert.equal(framework.getStore().isClosed(), true, 'once the call returns, the loop ends and a later stop() closes the store'); + assert.deepEqual(rejections, []); + } finally { + letGo('hang'); + await framework.stop().catch(() => {}); + } + }); + + it('a stream and the event loop that both outlast the budget each fail the stop, together as an AggregateError', { timeout: 30_000 }, async () => { + const { framework, release, residentStreaming, entered, letGo, push } = await frameworkWithHeldWork(); + try { + framework.start(); + push('go'); + await residentStreaming; + push('hang'); + await until(() => entered.includes('hang'), 'the loop to enter the call'); + const outcome = await settledWithin(framework.stop(100), 5_000); + assert.ok(outcome instanceof AggregateError && !(outcome instanceof ShutdownTimeoutError), String(outcome)); + assert.deepEqual(outcome.errors.map((error: unknown) => error instanceof ShutdownTimeoutError ? error.pending : String(error)), + [['stream:assistant'], ['event loop']]); + assert.equal(outcome.message, + 'Framework shutdown failed: Shutdown timed out after 100 ms; pending: stream:assistant; Shutdown timed out after 100 ms; pending: event loop'); + assert.equal(framework.getStore().isClosed(), false); + release(); + letGo('hang'); + await framework.stop(); + assert.equal(framework.getStore().isClosed(), true); + assert.deepEqual(rejections, []); + } finally { + release(); + letGo('hang'); + await framework.stop().catch(() => {}); + } + }); + + it("a stop() after a timed-out one and a start() waits for the old loop too, and keeps the store open until it has settled", { timeout: 30_000 }, async () => { + // The start() replaces the loop a stop() waits for. What the first + // stop()'s budget left running stays the framework's to wait for, so + // the second one doesn't close the store under the old loop's call. + const { framework, entered, letGo, push } = await frameworkWithHeldWork(); + try { + framework.start(); + push('hang'); + await until(() => entered.includes('hang'), 'the loop to enter the call'); + const first = await settledWithin(framework.stop(100), 5_000); + assert.ok(first instanceof ShutdownTimeoutError && first.pending.includes('event loop'), String(first)); + framework.start(); + const second = await settledWithin(framework.stop(100), 5_000); + assert.ok(second instanceof ShutdownTimeoutError + && second.message === 'Shutdown timed out after 100 ms; pending: event loop', String(second)); + assert.equal(framework.getStore().isClosed(), false, "the old loop's call is still running"); + letGo('hang'); + await framework.stop(); + assert.equal(framework.getStore().isClosed(), true, 'a stop() closes the store once the old loop has settled'); + assert.deepEqual(rejections, []); + } finally { + letGo('hang'); + await framework.stop().catch(() => {}); + } + }); + + it("a loop whose stop() came back while it was in a call exits once that call returns, though a start() was made meanwhile, and the new lifecycle's loop runs on", { timeout: 30_000 }, async () => { + const { framework, entered, letGo, push, internals } = await frameworkWithHeldWork(); + try { + framework.start(); + push('hang'); + await until(() => entered.includes('hang'), 'the loop to enter the call'); + const outcome = await settledWithin(framework.stop(100), 5_000); + assert.ok(outcome instanceof ShutdownTimeoutError && outcome.pending.includes('event loop'), String(outcome)); + const oldLoop = internals.loopPromise!; + framework.start(); + const newLoop = internals.loopPromise!; + assert.notEqual(newLoop, oldLoop, 'the start() runs a loop of its own'); + letGo('hang'); + assert.equal(await settledWithin(oldLoop, 5_000), 'resolved', 'the old loop exits once its call returns'); + assert.equal(await settledWithin(newLoop, 100), 'still pending', "the new lifecycle's loop runs on"); + await framework.stop(); + assert.equal(await settledWithin(newLoop, 5_000), 'resolved', 'and ends with its own stop()'); + assert.deepEqual(rejections, []); + } finally { + letGo('hang'); + await framework.stop().catch(() => {}); + } + }); + }); + describe("tune-out's subconscious tools", () => { /** A framework whose tune-out coordinator is a stand-in that settles when the test says. */ async function withTuneOut() {