🚰 fix: Record the Provider Drain When Shutdown Cuts a Generation Short (#15532)

* 🚰 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

@@ -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;

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();

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<void>();

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(

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<void> {
     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

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<void>; 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<void> {
+    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<void>((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<boolean> {
+    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<boolean> {
     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<void>((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<void> {
+  async destroy(options: { settlementBudgetMs?: number } = {}): Promise<void> {
+    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();

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();