From d346e5ab5eda04cfdd4d7543abe4f925802ae4e3 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 3 Sep 2026 07:11:38 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=B0=20fix:=20Record=20the=20Provider?= =?UTF-8?q?=20Drain=20When=20Shutdown=20Cuts=20a=20Generation=20Short=20(#?= =?UTF-8?q?15532)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🚰 fix: Record the Provider Drain When Shutdown Cuts a Generation Short * fix: Await generation settlement on shutdown instead of publishing the drain * fix: Spend the real shutdown budget and cover the initialization window * fix: Keep the shutdown registration literal the index spec asserts on * fix: Keep draining settlements registered while shutdown is waiting * fix: Track provider executions at the begin/drain chokepoints instead of per caller * fix: Open the provider tracker before the begin CAS and cap cluster workers to the primary deadline * fix: Propagate the cluster deadline and release trackers on reconfigure and abandonment * fix: Hold SIGTERM until the worker acknowledges the cluster deadline * fix: Refuse provider starts after shutdown begins and harden the cluster deadline handoff * test: Assert the multi-execution shutdown wait with both executions begun before shutdown --- api/server/experimental.js | 118 +++++++++- api/server/index.js | 19 +- .../api/src/agents/remote/lifecycle.spec.ts | 24 ++ packages/api/src/agents/remote/lifecycle.ts | 7 + packages/api/src/app/shutdown.ts | 26 +++ .../api/src/stream/GenerationJobManager.ts | 180 ++++++++++++++- .../api/src/stream/__tests__/startup.spec.ts | 218 ++++++++++++++++++ 7 files changed, 582 insertions(+), 10 deletions(-) diff --git a/api/server/experimental.js b/api/server/experimental.js index ce21fae0dd..fbaa62d904 100644 --- a/api/server/experimental.js +++ b/api/server/experimental.js @@ -1,4 +1,12 @@ require('../config/credentials'); +/** + * The primary kills every worker and then force-exits the whole cluster this long after its own + * shutdown signal, regardless of what the workers are still doing. + */ +const CLUSTER_FORCE_EXIT_MS = 10_000; +/** Absolute time the primary will force-exit the cluster, propagated to this worker over IPC. */ +let clusterShutdownDeadlineAt = null; + const fs = require('fs'); const path = require('path'); require('module-alias')({ base: path.resolve(__dirname, '..') }); @@ -32,6 +40,9 @@ const { requestContextMiddleware, configureServerTimeouts, setupGracefulShutdown, + registerShutdownTask, + getRemainingShutdownMs, + getShutdownElapsedMs, configureMessageFilterRegexValidator, configureFileConfigRegexEngine, configureAgentEventRuntime, @@ -261,6 +272,55 @@ if (cluster.isMaster) { cluster.fork(); }); + /** + * Deliver the absolute deadline, then SIGTERM only once the worker has acknowledged it. + * IPC is asynchronous and worker.kill() is immediate, so without the acknowledgement a + * busy worker can enter its shutdown handler before the deadline arrives and fall back to + * a worker-local estimate the primary will not honor. Bounded so a stalled worker cannot + * hold the others. + */ + const signalWorkerAfterDeadlineAck = (worker, deadlineAt) => { + let signaled = false; + const onMessage = (msg) => { + if (msg != null && msg.type === 'cluster-shutdown-ack') { + signal(); + } + }; + /** A closed IPC channel is reported asynchronously, not thrown from send(); without a + * listener it reaches the global uncaughtException handler and exits the primary, + * killing every other worker before it can record its drain. */ + const onError = (err) => { + logger.warn('Worker IPC error during the shutdown handoff; treating it as gone:', err); + signal(); + }; + const signal = () => { + if (signaled) { + return; + } + signaled = true; + worker.off('message', onMessage); + worker.off('error', onError); + try { + worker.kill(); + } catch (err) { + logger.debug('Worker already gone before SIGTERM:', err); + } + }; + worker.on('message', onMessage); + worker.on('error', onError); + /** Deliberately no separate timeout. SIGTERM is sent only after the worker has recorded + * the deadline; a worker that never acknowledges is ended by the primary's own + * force-exit, the one deadline it can honor. Signaling sooner would let a stalled + * worker's SIGTERM handler run before the queued deadline message and fall back to a + * local budget the primary will not honor, the exact case this handoff exists for. + * Handshakes are per worker, so a stalled one holds no other. */ + worker.send({ type: 'cluster-shutdown', deadlineAt }, (err) => { + if (err) { + onError(err); + } + }); + }; + /** Graceful shutdown on SIGTERM/SIGINT */ const shutdown = () => { if (shuttingDown) { @@ -274,13 +334,16 @@ if (cluster.isMaster) { process.exit(0); return; } + /** Workers derive their settlement budget from THIS deadline — not from their own + * coordinator, and not from whenever their signal handler happened to run. */ + const deadlineAt = Date.now() + CLUSTER_FORCE_EXIT_MS; for (const worker of liveWorkers) { - worker.kill(); + signalWorkerAfterDeadlineAck(worker, deadlineAt); } setTimeout(() => { logger.info('Forcing shutdown after timeout'); process.exit(0); - }, 10000); + }, CLUSTER_FORCE_EXIT_MS); }; process.on('SIGTERM', shutdown); @@ -301,6 +364,43 @@ if (cluster.isMaster) { }), ); GenerationJobManager.initialize(); + // Stop active generations and close their SSE streams while the HTTP server drains. + registerShutdownTask( + 'generation job manager prepare', + () => GenerationJobManager.prepareForShutdown(), + { + phase: 'pre-drain', + priority: 100, + }, + ); + /** Spend the shutdown budget that is actually left waiting for open provider executions to + * record their own drains — but the budget this worker actually has is the primary's, not + * its own 60s coordinator: the primary force-exits the whole cluster CLUSTER_FORCE_EXIT_MS + * after signalling. Measure against that, and hold back a reserve for the tasks after this + * one. Abandoning an unrecorded drain fences the next generation permanently. */ + const CLUSTER_TEARDOWN_RESERVE_MS = 3_000; + const destroyGenerationJobManager = () => { + const remaining = getRemainingShutdownMs(); + const elapsed = getShutdownElapsedMs(); + if (remaining == null || elapsed == null) { + return GenerationJobManager.destroy(); + } + /** Prefer the deadline the primary actually set. The elapsed-based estimate starts + * counting only when this worker's signal handler ran, which lags the primary's timer + * by however long the event loop was blocked. */ + const primaryRemaining = + clusterShutdownDeadlineAt != null + ? clusterShutdownDeadlineAt - Date.now() + : CLUSTER_FORCE_EXIT_MS - elapsed; + return GenerationJobManager.destroy({ + settlementBudgetMs: Math.max( + 0, + Math.min(remaining, primaryRemaining) - CLUSTER_TEARDOWN_RESERVE_MS, + ), + }); + }; + // Tear down stream resources before shared caches and telemetry exporters shut down. + registerShutdownTask('generation job manager', destroyGenerationJobManager, { priority: 100 }); /** * The master may assign the sweep worker before or after this worker has * loaded app config. These flags join the IPC assignment with config @@ -331,6 +431,20 @@ if (cluster.isMaster) { }; /** Handle inter-process messages from master */ + process.on('message', (msg) => { + if (msg != null && msg.type === 'cluster-shutdown' && Number.isFinite(msg.deadlineAt)) { + clusterShutdownDeadlineAt = msg.deadlineAt; + /** The primary holds SIGTERM until this arrives, so the deadline is in place before + * the shutdown handler can run. */ + if (typeof process.send === 'function') { + try { + process.send({ type: 'cluster-shutdown-ack' }); + } catch (err) { + logger.debug('Could not acknowledge the shutdown deadline to the primary:', err); + } + } + } + }); process.on('message', (msg) => { if (msg.type === 'file-retention-sweep-worker') { shouldStartExpiredFileSweep = true; diff --git a/api/server/index.js b/api/server/index.js index 5f50c0d04c..70e0024fcb 100644 --- a/api/server/index.js +++ b/api/server/index.js @@ -45,6 +45,7 @@ const { preAuthTenantMiddleware, requestContextMiddleware, registerShutdownTask, + getRemainingShutdownMs, configureServerTimeouts, setupGracefulShutdown, updateInterfacePermissions, @@ -143,12 +144,24 @@ const configureGenerationStreams = () => { priority: 100, }, ); + /** Spend the shutdown budget that is actually left waiting for detached generations to + * record their own provider drains, holding back a reserve for the tasks after this one. + * Abandoning an unrecorded drain fences the next generation permanently. */ + const destroyGenerationJobManager = () => { + const remaining = getRemainingShutdownMs(); + return GenerationJobManager.destroy( + remaining == null + ? undefined + : { settlementBudgetMs: Math.max(0, remaining - SHUTDOWN_TEARDOWN_RESERVE_MS) }, + ); + }; // Tear down stream resources before shared caches and telemetry exporters shut down. - registerShutdownTask('generation job manager', () => GenerationJobManager.destroy(), { - priority: 100, - }); + registerShutdownTask('generation job manager', destroyGenerationJobManager, { priority: 100 }); }; +/** Reserved for the shutdown tasks that run after the generation job manager. */ +const SHUTDOWN_TEARDOWN_RESERVE_MS = 10_000; + const startServer = async () => { await waitForKeyvRedisClient(); await configureSubagentTaskRouting(); diff --git a/packages/api/src/agents/remote/lifecycle.spec.ts b/packages/api/src/agents/remote/lifecycle.spec.ts index 26b32bed34..d3dd68407a 100644 --- a/packages/api/src/agents/remote/lifecycle.spec.ts +++ b/packages/api/src/agents/remote/lifecycle.spec.ts @@ -213,6 +213,30 @@ describe('Agent execution enrollment', () => { ]); }); + it('abandons the shutdown tracker when terminalization gives up, without recording a drain', async () => { + const enrollment = await enrollAgentExecution(enrollmentParams('chatcmpl-terminal-abandon'), { + manager, + }); + await enrollment.beginProviderExecution(); + const abandon = jest.spyOn(manager, 'abandonProviderExecution'); + const drained = jest.spyOn(manager, 'markProviderExecutionDrained'); + jest + .spyOn(manager, 'completeJob') + .mockRejectedValueOnce(new Error('terminal store unavailable')) + .mockRejectedValueOnce(new Error('terminal store still unavailable')); + + await expect(enrollment.settle()).rejects.toThrow('terminal store still unavailable'); + + /** No marker on purpose — a successor must still see the truth — but the shutdown + * tracker has no other release path, so it is abandoned explicitly. */ + expect(drained).not.toHaveBeenCalled(); + expect(abandon).toHaveBeenCalledWith( + 'chatcmpl-terminal-abandon', + expect.any(Number), + expect.any(String), + ); + }); + it('lets destructive cleanup abort the canonical signal and wait for trailing writes', async () => { const enrollment = await enrollAgentExecution(enrollmentParams('resp-delete'), { manager }); const tail = deferred(); diff --git a/packages/api/src/agents/remote/lifecycle.ts b/packages/api/src/agents/remote/lifecycle.ts index f1a37b4010..b80aabbdb4 100644 --- a/packages/api/src/agents/remote/lifecycle.ts +++ b/packages/api/src/agents/remote/lifecycle.ts @@ -159,6 +159,13 @@ export class AgentExecutionEnrollment { } let drainError: unknown; + if (this.providerStarted && terminalError != null) { + /** Provider work and trailing writes have settled, but terminalization gave up, so no + * drain marker is published on purpose — a successor must still see the truth. The + * shutdown tracker has no other path to release, though; abandon it explicitly or it + * outlives this run and every later graceful shutdown waits its full budget on it. */ + this.manager.abandonProviderExecution(this.runId, this.createdAt, this.providerExecutionId); + } if (this.providerStarted && terminalError == null) { try { const drained = await this.manager.markProviderExecutionDrained( diff --git a/packages/api/src/app/shutdown.ts b/packages/api/src/app/shutdown.ts index ef10189384..63216a1053 100644 --- a/packages/api/src/app/shutdown.ts +++ b/packages/api/src/app/shutdown.ts @@ -24,6 +24,7 @@ let nextRegistrationOrder = 0; let isShuttingDown = false; let httpServer: Server | null = null; let forceExitTimer: NodeJS.Timeout | null = null; +let shutdownStartedAt: number | null = null; /** * Register a cleanup task for graceful shutdown. Post-drain is the default phase. @@ -44,6 +45,29 @@ export function registerShutdownTask( } /** Whether graceful shutdown has started, for admission paths that must fail closed. */ +/** + * Milliseconds left before graceful shutdown force-exits, or `null` when not shutting down. + * Lets a shutdown task spend the budget it actually has instead of guessing at a fixed cutoff. + */ +/** + * Milliseconds since graceful shutdown began, or `null` when not shutting down. For a task whose + * real deadline is imposed from outside this process — a cluster primary that force-exits the + * whole group on its own timer — this is what lets it measure against that deadline instead. + */ +export function getShutdownElapsedMs(): number | null { + if (shutdownStartedAt == null) { + return null; + } + return Math.max(0, Date.now() - shutdownStartedAt); +} + +export function getRemainingShutdownMs(): number | null { + if (shutdownStartedAt == null) { + return null; + } + return Math.max(0, SHUTDOWN_TIMEOUT_MS - (Date.now() - shutdownStartedAt)); +} + export function isShutdownInProgress(): boolean { return isShuttingDown; } @@ -73,6 +97,7 @@ export function __resetShutdownStateForTests(): void { tasks.length = 0; nextRegistrationOrder = 0; isShuttingDown = false; + shutdownStartedAt = null; httpServer = null; /** A drain that never settles leaves this armed. It is `unref`'d, so it does * not hold the process open — but it does fire if anything else keeps the @@ -114,6 +139,7 @@ async function shutdown(signal: NodeJS.Signals): Promise { return; } isShuttingDown = true; + shutdownStartedAt = Date.now(); logger.info(`Received ${signal}, draining HTTP server...`); /** Owned locally so a late `finally` from a superseded drain cannot clear the diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index a2d7c91ca3..5fe8ff6d75 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -97,6 +97,14 @@ const REAPED_JOB_ERROR = 'Generation timed out'; * queued commands by pacing the producer instead of growing without limit. */ const MAX_OUTSTANDING_COALESCED_RECEIPTS = 256; const PROVIDER_DRAIN_POLL_MS = 50; + +/** + * The settlement budget is the opt-in. Graceful shutdown owns the real deadline and passes what + * is actually left, so open provider executions get the full remaining window to record their + * drains. A bare `destroy()` — test teardown, a direct reset — supplies nothing and must not + * block on executions that were begun and never drained. + */ +const SHUTDOWN_SETTLEMENT_FALLBACK_MS = 0; type ToolExecutionStatus = 'success' | 'error' | 'cancelled'; type ToolCallWithExecutionStatus = Agents.AgentToolCall & { executionStatus?: ToolExecutionStatus; @@ -776,6 +784,18 @@ class GenerationJobManagerClass { /** Rejects new jobs once graceful shutdown has started. */ private shuttingDown = false; + /** + * Provider executions this process has begun and not yet drained, awaited during shutdown so + * each can record its own drain once its trailing writes settle. Opened in + * `beginProviderExecution` and closed in `markProviderExecutionDrained` — the two chokepoints + * every entry point (interactive, resume, remote API) already passes through — so no caller + * has to remember to register, and a new entry point cannot silently opt out. + */ + private readonly openProviderExecutions = new Map< + string, + { promise: Promise; settle: () => void } + >(); + /** Whether we're using Redis stores */ private _isRedis = false; @@ -886,6 +906,9 @@ class GenerationJobManagerClass { } this.ownedJobs.clear(); + /** The trackers belong to the services being replaced; a stale one would make a later + * graceful destroy() wait its whole budget on an execution that can never drain. */ + this.releaseOpenProviderExecutions(); setGenerationJobsInFlight(previousStore, 0); this.jobStore = services.jobStore; @@ -7084,10 +7107,70 @@ class GenerationJobManagerClass { /** Records that one exact provider segment has completed every trailing write. * The opaque segment id prevents a paused controller from acknowledging a * later HITL resume that reuses the same generation epoch. */ + + /** + * Wait for tracked generations to settle, bounded so shutdown still finishes inside a + * container stop grace period. Anything still outstanding is left alone: asserting a drain + * it has not reached is the unsafe direction, so we log and let the deployment's + * hard-kill handling own that case. + */ + private async awaitGenerationSettlements(budget: number): Promise { + const deadline = Date.now() + budget; + /** Re-snapshot until the set drains. A settlement can be registered while an earlier + * one is being awaited — the controller registers its detached chain just before its + * initialization settlement resolves — and a one-time snapshot would return between + * the two, tearing down the store and transport under a still-running provider. */ + while (this.openProviderExecutions.size > 0) { + const remaining = deadline - Date.now(); + if (remaining <= 0) { + break; + } + const pending = [...this.openProviderExecutions.values()].map((entry) => entry.promise); + let timer: NodeJS.Timeout | undefined; + try { + await Promise.race([ + Promise.allSettled(pending), + new Promise((resolve) => { + timer = setTimeout(resolve, remaining); + }), + ]); + } finally { + if (timer) { + clearTimeout(timer); + } + } + } + const unsettled = this.openProviderExecutions.size; + if (unsettled > 0) { + logger.warn( + `[GenerationJobManager] ${unsettled} provider execution(s) did not drain within ${budget}ms of shutdown; their drains are unrecorded`, + ); + } + } + async markProviderExecutionDrained( streamId: string, expectedCreatedAt: number, providerExecutionId: string, + ): Promise { + try { + return await this.recordProviderExecutionDrain( + streamId, + expectedCreatedAt, + providerExecutionId, + ); + } finally { + /** Settles the shutdown wait whether or not the record succeeded: a drain that failed + * to write has nothing further to wait for, and holding shutdown open on it would only + * push the process into the force-exit with the same unrecorded marker. */ + this.closeProviderExecution(streamId, expectedCreatedAt, providerExecutionId); + } + } + + private async recordProviderExecutionDrain( + streamId: string, + expectedCreatedAt: number, + providerExecutionId: string, ): Promise { if (providerExecutionId.length === 0 || providerExecutionId.length > 128) { return false; @@ -7154,10 +7237,86 @@ class GenerationJobManagerClass { if (providerExecutionId.length === 0 || providerExecutionId.length > 128) { return false; } - return ( - this.jobStore.beginProviderExecution?.(streamId, expectedCreatedAt, providerExecutionId) ?? - Promise.resolve(false) - ); + /** A provider that has not started by the time shutdown begins must not start. The + * settlement wait only covers executions begun before it returned; one begun after it + * would commit `providerDrained: false` against a store that is about to close and + * fence its successor. Callers already treat `false` as "stopped before provider + * startup", and nothing has been committed. */ + if (this.shuttingDown) { + return false; + } + /** Opened before the CAS and closed only on a definitive `false`. A rejected reply is + * ambiguous — the store may already have committed `providerDrained: false` — and every + * caller treats it that way, running or cleaning up and recording the drain either way. + * Leaving the tracker open makes shutdown wait for that record instead of tearing the + * transport down underneath it; the deadline still bounds the wait if it never comes. */ + this.openProviderExecution(streamId, expectedCreatedAt, providerExecutionId); + const started = + (await this.jobStore.beginProviderExecution?.( + streamId, + expectedCreatedAt, + providerExecutionId, + )) ?? false; + if (!started) { + this.closeProviderExecution(streamId, expectedCreatedAt, providerExecutionId); + } + return started; + } + + private providerExecutionKey( + streamId: string, + createdAt: number, + providerExecutionId: string, + ): string { + return `${streamId}:${createdAt}:${providerExecutionId}`; + } + + private openProviderExecution( + streamId: string, + createdAt: number, + providerExecutionId: string, + ): void { + const key = this.providerExecutionKey(streamId, createdAt, providerExecutionId); + if (this.openProviderExecutions.has(key)) { + return; + } + let settle: () => void = () => undefined; + const promise = new Promise((resolve) => { + settle = resolve; + }); + this.openProviderExecutions.set(key, { promise, settle }); + } + + private closeProviderExecution( + streamId: string, + createdAt: number, + providerExecutionId: string, + ): void { + const key = this.providerExecutionKey(streamId, createdAt, providerExecutionId); + const entry = this.openProviderExecutions.get(key); + if (entry == null) { + return; + } + this.openProviderExecutions.delete(key); + entry.settle(); + } + + /** + * Release the tracker for an execution whose owner will never record its drain — a run + * that settled its provider work but failed terminalization and deliberately publishes no + * marker. Nothing is written: shutdown simply stops waiting on it, and the absent marker + * keeps telling a successor the truth. + */ + abandonProviderExecution(streamId: string, createdAt: number, providerExecutionId: string): void { + this.closeProviderExecution(streamId, createdAt, providerExecutionId); + } + + /** Settle and drop every open tracker; used wherever the services they belong to go away. */ + private releaseOpenProviderExecutions(): void { + for (const entry of this.openProviderExecutions.values()) { + entry.settle(); + } + this.openProviderExecutions.clear(); } private async waitForProviderExecutionDrain( @@ -8375,7 +8534,13 @@ class GenerationJobManagerClass { * Destroy the manager. * Cleans up all resources including runtime state, buffers, and stores. */ - async destroy(): Promise { + async destroy(options: { settlementBudgetMs?: number } = {}): Promise { + const settlementBudgetMs = Number.isFinite(options.settlementBudgetMs) + ? Math.max(0, options.settlementBudgetMs as number) + : SHUTDOWN_SETTLEMENT_FALLBACK_MS; + /** Anchored here, not after subscriber cleanup: that cleanup is unbounded, and the caller's + * budget was measured against the force-exit timer that keeps running through it. */ + const settlementDeadline = Date.now() + settlementBudgetMs; this.shuttingDown = true; this.cancelFencedRuntimeRetirements(); @@ -8392,9 +8557,14 @@ class GenerationJobManagerClass { } await this.drainSubscriberCleanups(); + await this.awaitGenerationSettlements(Math.max(0, settlementDeadline - Date.now())); await this.finalizeOwnedJobsForShutdown(); await this.jobStore.destroy(); this.eventTransport.destroy(); + /** Whatever the bounded wait left behind must not outlive this store: a later + * `configure()`/`initialize()` in the same process would otherwise wait a full budget + * on executions that can never drain. */ + this.releaseOpenProviderExecutions(); this.runtimeState.clear(); this.ownedJobs.clear(); this.syncRunningJobMetrics(); diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts index 63d7980df1..3afb487751 100644 --- a/packages/api/src/stream/__tests__/startup.spec.ts +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -4118,6 +4118,224 @@ describe('GenerationJobManager startup telemetry', () => { }); }); + const configureShutdownManager = () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false }); + manager.initialize(); + return { manager, jobStore, eventTransport }; + }; + + it('waits for a begun provider execution to record its drain before shutdown completes', async () => { + const { manager, eventTransport } = configureShutdownManager(); + const streamId = 'stream-shutdown-open-execution'; + const job = await manager.createJob(streamId, 'user-1'); + const providerExecutionId = job.metadata.providerExecutionId!; + await expect( + manager.beginProviderExecution(streamId, job.createdAt, providerExecutionId), + ).resolves.toBe(true); + const recordProviderDrain = jest.spyOn(eventTransport, 'recordProviderDrain'); + + manager.prepareForShutdown(); + let destroyed = false; + const destroying = manager.destroy({ settlementBudgetMs: 5_000 }).then(() => { + destroyed = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(destroyed).toBe(false); + + /** The execution records its own drain — from the controller's `.finally`, after its + * trailing writes — and that is what releases shutdown. No registration involved. */ + await manager.markProviderExecutionDrained(streamId, job.createdAt, providerExecutionId); + await destroying; + + expect(destroyed).toBe(true); + expect(recordProviderDrain).toHaveBeenCalledWith(streamId, job.createdAt, providerExecutionId); + }); + + it('keeps waiting for every execution begun before shutdown, not only the first to drain', async () => { + const { manager } = configureShutdownManager(); + const first = await manager.createJob('stream-shutdown-first', 'user-1'); + const second = await manager.createJob('stream-shutdown-second', 'user-1'); + const firstExecution = first.metadata.providerExecutionId!; + const secondExecution = second.metadata.providerExecutionId!; + await manager.beginProviderExecution('stream-shutdown-first', first.createdAt, firstExecution); + await manager.beginProviderExecution( + 'stream-shutdown-second', + second.createdAt, + secondExecution, + ); + + manager.prepareForShutdown(); + let destroyed = false; + const destroying = manager.destroy({ settlementBudgetMs: 5_000 }).then(() => { + destroyed = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + + /** The first draining must not release shutdown while the second is still open. */ + await manager.markProviderExecutionDrained( + 'stream-shutdown-first', + first.createdAt, + firstExecution, + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(destroyed).toBe(false); + + await manager.markProviderExecutionDrained( + 'stream-shutdown-second', + second.createdAt, + secondExecution, + ); + await destroying; + expect(destroyed).toBe(true); + }); + + it('does not delay shutdown for a provider execution that already drained', async () => { + const { manager } = configureShutdownManager(); + const streamId = 'stream-shutdown-drained'; + const job = await manager.createJob(streamId, 'user-1'); + const providerExecutionId = job.metadata.providerExecutionId!; + await manager.beginProviderExecution(streamId, job.createdAt, providerExecutionId); + await manager.markProviderExecutionDrained(streamId, job.createdAt, providerExecutionId); + + manager.prepareForShutdown(); + await expect(manager.destroy({ settlementBudgetMs: 5_000 })).resolves.toBeUndefined(); + }); + + it('keeps waiting when the provider-begin reply is lost after the store may have committed', async () => { + const { manager, jobStore } = configureShutdownManager(); + const streamId = 'stream-shutdown-ambiguous-begin'; + const job = await manager.createJob(streamId, 'user-1'); + const providerExecutionId = job.metadata.providerExecutionId!; + jest + .spyOn(jobStore, 'beginProviderExecution') + .mockRejectedValueOnce(new Error('reply lost after commit')); + await expect( + manager.beginProviderExecution(streamId, job.createdAt, providerExecutionId), + ).rejects.toThrow('reply lost after commit'); + + manager.prepareForShutdown(); + let destroyed = false; + const destroying = manager.destroy({ settlementBudgetMs: 5_000 }).then(() => { + destroyed = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + /** The caller treats the rejection as possibly-started and still records the drain from + * its cleanup; shutdown has to wait for that rather than assume nothing began. */ + expect(destroyed).toBe(false); + + await manager.markProviderExecutionDrained(streamId, job.createdAt, providerExecutionId); + await destroying; + expect(destroyed).toBe(true); + }); + + it('does not carry undrained executions into a re-initialized manager', async () => { + const { manager } = configureShutdownManager(); + const streamId = 'stream-shutdown-stale-execution'; + const job = await manager.createJob(streamId, 'user-1'); + await manager.beginProviderExecution( + streamId, + job.createdAt, + job.metadata.providerExecutionId!, + ); + + /** A bare reset: zero budget, so the execution is still open when teardown runs. */ + manager.prepareForShutdown(); + await manager.destroy(); + + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + manager.configure({ jobStore, eventTransport: new InMemoryEventTransport(), isRedis: false }); + manager.initialize(); + + manager.prepareForShutdown(); + let destroyed = false; + const destroying = manager.destroy({ settlementBudgetMs: 5_000 }).then(() => { + destroyed = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + /** Resolved without consuming the budget: nothing from the previous store is left to wait + * on. Asserted before awaiting, or a 5s stale wait would pass this test too. */ + expect(destroyed).toBe(true); + await destroying; + }); + + it('does not carry undrained executions across a reconfigure', async () => { + const { manager } = configureShutdownManager(); + const streamId = 'stream-shutdown-reconfigure'; + const job = await manager.createJob(streamId, 'user-1'); + await manager.beginProviderExecution( + streamId, + job.createdAt, + job.metadata.providerExecutionId!, + ); + + /** Reconfiguring after initialization replaces the store and transport the open + * execution belongs to; its tracker must go with them. */ + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + jest.spyOn(jobStore, 'destroy').mockResolvedValue(); + manager.configure({ jobStore, eventTransport: new InMemoryEventTransport(), isRedis: false }); + + manager.prepareForShutdown(); + let destroyed = false; + const destroying = manager.destroy({ settlementBudgetMs: 5_000 }).then(() => { + destroyed = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(destroyed).toBe(true); + await destroying; + }); + + it('releases an abandoned provider execution without recording a drain', async () => { + const { manager, eventTransport } = configureShutdownManager(); + const streamId = 'stream-shutdown-abandoned'; + const job = await manager.createJob(streamId, 'user-1'); + const providerExecutionId = job.metadata.providerExecutionId!; + await manager.beginProviderExecution(streamId, job.createdAt, providerExecutionId); + const recordProviderDrain = jest.spyOn(eventTransport, 'recordProviderDrain'); + + /** A remote run that settled its provider work but failed terminalization twice + * deliberately publishes no drain marker; it must still stop shutdown waiting on it. */ + manager.abandonProviderExecution(streamId, job.createdAt, providerExecutionId); + + manager.prepareForShutdown(); + let destroyed = false; + const destroying = manager.destroy({ settlementBudgetMs: 5_000 }).then(() => { + destroyed = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(destroyed).toBe(true); + await destroying; + expect(recordProviderDrain).not.toHaveBeenCalled(); + }); + + it('refuses to begin a provider execution once shutdown has started', async () => { + const { manager, jobStore } = configureShutdownManager(); + const streamId = 'stream-shutdown-late-begin'; + const job = await manager.createJob(streamId, 'user-1'); + const providerExecutionId = job.metadata.providerExecutionId!; + const storeBegin = jest.spyOn(jobStore, 'beginProviderExecution'); + + manager.prepareForShutdown(); + /** Began after the settlement wait could have observed it: must not start, must not + * commit `providerDrained: false`, and must leave nothing for shutdown to wait on. */ + await expect( + manager.beginProviderExecution(streamId, job.createdAt, providerExecutionId), + ).resolves.toBe(false); + expect(storeBegin).not.toHaveBeenCalled(); + + let destroyed = false; + const destroying = manager.destroy({ settlementBudgetMs: 5_000 }).then(() => { + destroyed = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(destroyed).toBe(true); + await destroying; + }); + it('does not let late success overwrite or delete a shutdown error', async () => { const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); jest.spyOn(jobStore, 'destroy').mockResolvedValue();