fix(rivetkit): fire abort signal on shutdown finalize

NathanFlurry · NathanFlurry · commit 11daf9281f8d · 2026-05-04T08:35:25.000-07:00
diff --git a/docs-internal/engine/rivetkit-core-internals.md b/docs-internal/engine/rivetkit-core-internals.md
@@ -85,7 +85,7 @@ Two-phase:
 - `SleepGrace` fires `onSleep` immediately and keeps dispatch/save timers live.
 - `SleepFinalize` gates dispatch, suspends alarms, and runs teardown.
 
-Sleep grace must fire the actor abort signal on entry and wait for the run handler to exit before finalize. Destroy abort firing remains unchanged.
+Sleep grace waits for the run handler to exit before finalize. The actor abort signal fires once when final shutdown completes, not when sleep grace or destroy starts.
 
 Finalize:
 
diff --git a/docs-internal/engine/sleep-sequence.md b/docs-internal/engine/sleep-sequence.md
@@ -37,9 +37,9 @@ Removing `preventSleep` deleted both predicate branches. Any future sleep-affect
 
 ## Grace period and abort signals
 
-- `start_grace(reason)` fires at the start of `SleepGrace` / `DestroyGrace`. It cancels the sleep idle timer, cancels the actor abort signal (`actor_abort_signal`), installs a `SleepGraceState` with the effective grace deadline, and resets the sleep timer to arm the grace tick.
-- The actor abort signal is a soft signal: "shutdown has started, please wrap up." User code observes it via `c.abortSignal`. It does not force-stop work.
-- For destroy, the abort signal may fire earlier than grace entry because `ctx.destroy()` cancels the abort token immediately via `mark_destroy_requested(...)`.
+- `start_grace(reason)` fires at the start of `SleepGrace` / `DestroyGrace`. It cancels the sleep idle timer, installs a `SleepGraceState` with the effective grace deadline, and resets the sleep timer to arm the grace tick.
+- The actor abort signal is a finalization signal. User code observes it via `c.abortSignal`, and it fires once after shutdown work has finished or the grace deadline has forced finalization.
+- Destroy does not fire the actor abort signal early. `ctx.destroy()` only records the destroy request and starts the same final shutdown path.
 
 ## Grace deadline enforcement
 
diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs
@@ -478,19 +478,17 @@ impl ActorContext {
 		self.flush_on_shutdown();
 		self.0.destroy_requested.store(true, Ordering::SeqCst);
 		self.0.destroy_completed.store(false, Ordering::SeqCst);
-		self.0.abort_signal.lock().cancel();
 	}
 
 	#[cfg(feature = "wasm-runtime")]
 	fn mark_destroy_requested_without_spawn(&self) {
 		self.cancel_sleep_timer();
 		self.0.destroy_requested.store(true, Ordering::SeqCst);
 		self.0.destroy_completed.store(false, Ordering::SeqCst);
-		self.0.abort_signal.lock().cancel();
 	}
 
 	#[doc(hidden)]
-	pub fn cancel_abort_signal_for_sleep(&self) {
+	pub fn cancel_abort_signal_for_shutdown_finalize(&self) {
 		self.0.abort_signal.lock().cancel();
 	}
 
@@ -500,9 +498,8 @@ impl ActorContext {
 			return;
 		}
 
-		// Sleep cancels the generation abort signal to break queue waits and the
-		// run loop out of blocking calls. A restarted actor needs a fresh signal
-		// so the next generation can wait normally.
+		// Final shutdown cancels the generation abort signal. A restarted actor
+		// needs a fresh signal so the next generation can wait normally.
 		let next_signal = CancellationToken::new();
 		*abort_signal = next_signal.clone();
 		*self.0.queue_abort_signal.lock() = next_signal;
diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs
@@ -1528,9 +1528,6 @@ impl ActorTask {
 		let grace_period = self.factory.config().effective_sleep_grace_period();
 		self.sleep_deadline = None;
 		self.ctx.cancel_sleep_timer();
-		// Entering grace cancels the actor abort signal so user code blocked on
-		// queues or other actor scoped waits can unwind and let sleep finalize.
-		self.ctx.cancel_abort_signal_for_sleep();
 		self.sleep_grace = Some(SleepGraceState {
 			deadline: Instant::now() + grace_period,
 			reason,
@@ -1745,15 +1742,19 @@ impl ActorTask {
 			ShutdownKind::Sleep => LifecycleState::SleepFinalize,
 			ShutdownKind::Destroy => LifecycleState::Destroying,
 		});
-		self.save_final_state().await?;
-		self.close_actor_event_channel();
-		self.join_aborted_run_handle().await;
-		Self::finish_shutdown_cleanup_with_ctx(self.ctx.clone(), reason).await?;
-		if matches!(reason, ShutdownKind::Destroy) {
+		let result: Result<()> = async {
+			self.save_final_state().await?;
+			self.close_actor_event_channel();
+			self.join_aborted_run_handle().await;
+			Self::finish_shutdown_cleanup_with_ctx(self.ctx.clone(), reason).await
+		}
+		.await;
+		self.ctx.cancel_abort_signal_for_shutdown_finalize();
+		if result.is_ok() && matches!(reason, ShutdownKind::Destroy) {
 			self.ctx.mark_destroy_completed();
 		}
 		self.ctx.record_shutdown_wait(reason, started_at.elapsed());
-		Ok(())
+		result
 	}
 
 	async fn save_final_state(&mut self) -> Result<()> {
diff --git a/rivetkit-rust/packages/rivetkit-core/tests/queue.rs b/rivetkit-rust/packages/rivetkit-core/tests/queue.rs
@@ -162,7 +162,7 @@ mod moved_tests {
 		});
 
 		yield_now().await;
-		queue.mark_destroy_requested();
+		queue.cancel_abort_signal_for_shutdown_finalize();
 
 		let error = wait
 			.await
diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs
@@ -2968,7 +2968,7 @@ mod moved_tests {
 			.await
 			.expect("sleep stop should send");
 		wait_for_count(&begin_sleep_count, 1).await;
-		assert!(ctx.actor_aborted());
+		assert!(!ctx.actor_aborted());
 
 		let (sleep_again_tx, sleep_again_rx) = oneshot::channel();
 		lifecycle_tx
@@ -3009,6 +3009,7 @@ mod moved_tests {
 			.await
 			.expect("sleep reply should send")
 			.expect("sleep should succeed");
+		assert!(ctx.actor_aborted());
 		keep_awake
 			.await
 			.expect("keep-awake task should finish after release");
diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/config.ts b/rivetkit-typescript/packages/rivetkit/src/actor/config.ts
@@ -1227,7 +1227,7 @@ interface BaseActorConfig<
 	 * shutdown window (`sleepGracePeriod`) cover deferred work.
 	 *
 	 * The handler receives an abort signal via `c.abortSignal` and a
-	 * `c.aborted` alias for loop checks. Use these to gracefully exit.
+	 * `c.aborted` alias. The signal fires once when actor shutdown finalizes.
 	 *
 	 * If this handler exits, the actor will follow the normal idle sleep timeout
 	 * once it becomes idle.
@@ -1770,7 +1770,7 @@ export const DocActorOptionsSchema = z
 			.number()
 			.optional()
 			.describe(
-				`Max time in ms for the graceful shutdown window. Covers lifecycle hooks (onSleep, onDestroy), the run handler abort wait, async raw WebSocket handlers, disconnect callbacks, and final state serialization. Default: ${DEFAULT_SLEEP_GRACE_PERIOD}.`,
+				`Max time in ms for the graceful shutdown window. Covers lifecycle hooks (onSleep, onDestroy), the run handler wait, async raw WebSocket handlers, disconnect callbacks, and final state serialization. Default: ${DEFAULT_SLEEP_GRACE_PERIOD}.`,
 			),
 		onDestroyTimeout: z
 			.number()
diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts
@@ -958,14 +958,6 @@ export class ActorInstance<
 			// on wake via initializeAlarms().
 			this.driver.cancelAlarm?.(this.#actorId);
 
-			// Abort listeners in the canonical stop path.
-			// This must run for all stop modes, including sleep and remote stop.
-			// Destroy may have already triggered an early abort, but repeating abort
-			// is intentional and safe.
-			try {
-				this.#abortController.abort();
-			} catch {}
-
 			// The run-handler join, lifecycle hooks, and remaining shutdown
 			// tasks all share the single sleepGracePeriod budget.
 			const shutdownTaskDeadlineTs =
@@ -1006,6 +998,7 @@ export class ActorInstance<
 			await this.stateManager.waitForPendingWrites();
 			await this.#scheduleManager.waitForPendingAlarmWrites();
 		} finally {
+			this.#abortActorSignal();
 			this.#shutdownComplete = true;
 			await this.#cleanupDatabase();
 		}
@@ -1031,11 +1024,8 @@ export class ActorInstance<
 
 			this.driver.cancelAlarm?.(this.#actorId);
 			this.stateManager.clearPendingSaveTimeout();
-
-			try {
-				this.#abortController.abort();
-			} catch {}
 		} finally {
+			this.#abortActorSignal();
 			this.#shutdownComplete = true;
 			await this.#cleanupDatabase();
 		}
@@ -1087,14 +1077,6 @@ export class ActorInstance<
 		}
 		this.#destroyCalled = true;
 
-		// Abort immediately so in flight waits can exit before the driver stop
-		// handshake completes.
-		// The onStop path will call abort again as a safety net for all stop
-		// modes.
-		try {
-			this.#abortController.abort();
-		} catch {}
-
 		const destroy = this.driver.startDestroy.bind(
 			this.driver,
 			this.#actorId,
@@ -1108,6 +1090,15 @@ export class ActorInstance<
 		});
 	}
 
+	#abortActorSignal() {
+		if (this.#abortController.signal.aborted) {
+			return;
+		}
+		try {
+			this.#abortController.abort();
+		} catch {}
+	}
+
 	// MARK: - HTTP Request Tracking
 	beginHonoHttpRequest() {
 		this.#activeHonoHttpRequests++;
@@ -2070,7 +2061,7 @@ export class ActorInstance<
 
 		if (timeoutMs <= 0) {
 			this.#rLog.warn({
-				msg: "run handler did not complete in time, it may have leaked - ensure you use c.aborted (or the abort signal c.abortSignal) to exit gracefully",
+				msg: "run handler did not complete in time, it may have leaked - ensure long-running work settles before shutdown finalizes",
 				timeoutMs,
 			});
 			return;
@@ -2085,7 +2076,7 @@ export class ActorInstance<
 
 		if (timedOut) {
 			this.#rLog.warn({
-				msg: "run handler did not complete in time, it may have leaked - ensure you use c.aborted (or the abort signal c.abortSignal) to exit gracefully",
+				msg: "run handler did not complete in time, it may have leaked - ensure long-running work settles before shutdown finalizes",
 				timeoutMs,
 			});
 		} else {
diff --git a/website/src/content/docs/actors/actions.mdx b/website/src/content/docs/actors/actions.mdx
@@ -325,7 +325,7 @@ Actions have a single return value. To stream realtime data in response to an ac
 
 ## Canceling Long-Running Actions
 
-For operations that should be cancelable on-demand, create your own `AbortController` and chain it with `c.abortSignal` for automatic cleanup on actor shutdown.
+For operations that should be cancelable on-demand, create your own `AbortController`. Chain it with `c.abortSignal` only for final actor shutdown cleanup.
 
 ```typescript
 import { actor } from "rivetkit";
diff --git a/website/src/content/docs/actors/index.mdx b/website/src/content/docs/actors/index.mdx
@@ -466,7 +466,7 @@ const userAccount = actor({
 
 ### Lifecycle Hooks
 
-Actors support hooks for initialization, background processing, connections, networking, and state changes. Use `run` for long-lived background loops, and exit cleanly on shutdown with `c.aborted` or `c.abortSignal`.
+Actors support hooks for initialization, background processing, connections, networking, and state changes. Use `run` for long-lived background loops, and use `c.aborted` or `c.abortSignal` for final shutdown cleanup.
 
 ```ts
 import { actor, event, queue } from "rivetkit";
diff --git a/website/src/content/docs/actors/lifecycle.mdx b/website/src/content/docs/actors/lifecycle.mdx
@@ -282,13 +282,13 @@ The `run` hook is called after the actor starts and runs in the background witho
 - Custom workflow logic
 - Background processing
 
-The handler exposes `c.aborted` for loop checks and `c.abortSignal` for canceling operations when the actor is stopping. You should always check or listen for shutdown to exit gracefully.
+The handler exposes `c.aborted` and `c.abortSignal` for final shutdown cleanup. The signal fires once when shutdown finalizes, not when sleep or destroy is first requested.
 
 **Important behavior:**
 - The actor may go to sleep at any time during the `run` handler. Wrap work that must keep the actor awake with `c.keepAwake(promise)` to block idle sleep until the promise settles.
 - If the `run` handler exits (returns), the actor follows its normal idle sleep timeout once it becomes idle
 - If the `run` handler throws an error, the actor logs the error and then follows its normal idle sleep timeout once it becomes idle
-- On shutdown, `c.abortSignal` fires so the `run` handler can exit within the graceful shutdown window.
+- On shutdown, the `run` handler is awaited during the graceful shutdown window. `c.abortSignal` fires after that window finishes or times out.
 
 ```typescript
 import { actor } from "rivetkit";
@@ -304,7 +304,7 @@ const tickActor = actor({
       c.state.tickCount++;
       c.log.info({ msg: "tick", count: c.state.tickCount });
 
-      // Wait 1 second, but exit early if aborted
+      // Wait 1 second. Final shutdown also resolves this wait.
       await new Promise<void>((resolve) => {
         const timeout = setTimeout(resolve, 1000);
         c.abortSignal.addEventListener("abort", () => {
@@ -874,9 +874,10 @@ WebSocket connections are preserved across sleep cycles by default and transpare
 
 When an actor sleeps or is destroyed, it enters the graceful shutdown window:
 
-1. `c.abortSignal` fires and `c.aborted` becomes `true`. New connections and dispatch are rejected. Alarm timeouts are cancelled. On sleep, scheduled events are persisted and will be re-armed when the actor wakes.
+1. New connections and dispatch are rejected. Alarm timeouts are cancelled. On sleep, scheduled events are persisted and will be re-armed when the actor wakes.
 2. `onSleep` or `onDestroy` and `onDisconnect` for each closing connection run during the same window. User `waitUntil` promises and async raw WebSocket handlers are drained. Hibernatable WebSocket connections are preserved on sleep and closed on destroy.
-3. Once graceful work has completed, state is saved and the database is cleaned up.
+3. Once graceful work has completed, state is saved and final cleanup runs.
+4. `c.abortSignal` fires once and `c.aborted` becomes `true` as final shutdown completes.
 
 The entire window is bounded by `sleepGracePeriod` on both sleep and destroy. Defaults to 15 seconds. If the window is exceeded, the actor proceeds to state save anyway.
 
@@ -954,9 +955,9 @@ const myActor = actor({
 
 ### Actor Shutdown Abort Signal
 
-The `c.abortSignal` provides an `AbortSignal` that fires when the actor is stopping, and `c.aborted` is the shorthand boolean for loop checks. Use these to cancel ongoing operations when the actor sleeps or is destroyed.
+The `c.abortSignal` provides an `AbortSignal` that fires once when actor shutdown finalizes, and `c.aborted` is the shorthand boolean for checking that final state.
 
-The abort signal fires at the very start of the [shutdown sequence](#shutdown-sequence), before `onSleep` or `onDestroy` runs. This means `c.aborted` is already `true` inside those lifecycle hooks. The signal fires early so that the `run` handler can exit promptly, but the actor remains fully functional throughout the graceful shutdown window.
+The abort signal does not fire when sleep or destroy is requested. It fires at the end of the [shutdown sequence](#shutdown-sequence), after graceful shutdown work has finished or the grace window has timed out.
 
 ```typescript
 import { actor } from "rivetkit";
diff --git a/website/src/content/docs/actors/queues.mdx b/website/src/content/docs/actors/queues.mdx
@@ -428,7 +428,7 @@ This means you can run normal code in `run` without worrying about sleep interru
 - Implement connection auth in `onBeforeConnect`. See [Authentication](/docs/actors/authentication).
 - Route most state changes through one queue loop so ordering stays predictable.
 - If you need more complex multi-step run loops, consider using workflows.
-- Use `c.aborted` and `c.abortSignal` for graceful shutdown and cancellation.
+- Use `c.aborted` and `c.abortSignal` for final shutdown cleanup. Use your own `AbortController` for earlier loop cancellation.
 - Add `timeout` when callers need bounded wait behavior.
 - Use `wait: true` only when the caller actually needs a response.
 
