fix(rivetkit-core): close attached connections at the sleep grace deadline

docs-internal/engine/sleep-sequence.md

@@ -41,6 +41,7 @@ Removing `preventSleep` deleted both predicate branches. Any future sleep-affect
 - 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.
 - Destroy requests also use the normal grace path. The actor abort signal fires when destroy grace starts.
 - Sleep grace runs `onSleep` and connection disconnects only after tracked work that drains shutdown grace (actions, `keepAwake`, `waitUntil`, registered tasks, WebSocket callbacks) finishes, or at the hooks deadline, which keeps the last quarter of the grace period for the hooks and the final save. Raw `onRequest` handlers and `run` are not waited for. Destroy grace runs its hooks at once.
+- Non-hibernatable connections stay attached until the sleep hooks run, so in-flight results still reach clients. Destroy grace closes them at grace start. The grace deadline closes any that are still attached.
 - Sleep grace admits new work until the hooks run. After that, dispatch refuses with `actor/restarting` and the client retries against the next generation. `SleepFinalize` refuses the same way.
 - Core owns SQLite teardown: it syncs the alarm and then closes SQLite in final cleanup. The TypeScript sleep cleanup closes SQLite itself only when tracked work outlives the grace deadline, so that work cannot write after shutdown.
 

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs

@@ -841,9 +841,7 @@ impl ActorTask {
 	fn emit_grace_events(&mut self, reason: ShutdownKind) {
 		let conns: Vec<_> = self.ctx.conns().collect();
 		for conn in conns {
-			let hibernatable_sleep =
-				matches!(reason, ShutdownKind::Sleep) && conn.is_hibernatable();
-			if hibernatable_sleep {
+			if !closes_on_grace(reason, &conn) {
 				self.ctx.request_hibernation_transport_save(conn.id());
 				continue;
 			}
@@ -1718,6 +1716,18 @@ impl ActorTask {
 		let Some(grace) = self.sleep_grace.take() else {
 			return None;
 		};
+		// Hooks dispatched now would never be awaited, so close transports directly.
+		if let Err(error) = self
+			.ctx
+			.disconnect_conns(|conn| closes_on_grace(grace.reason, conn))
+			.await
+		{
+			tracing::warn!(
+				actor_id = %self.ctx.actor_id(),
+				?error,
+				"failed to close connections at the grace deadline"
+			);
+		}
 		if let Some(run_handle) = self.run_handle.as_mut() {
 			run_handle.abort();
 		}
@@ -2415,6 +2425,11 @@ fn shutdown_reason_label(reason: ShutdownKind) -> &'static str {
 	}
 }
 
+/// Hibernatable connections survive sleep, so sleep grace leaves them open.
+fn closes_on_grace(reason: ShutdownKind, conn: &ConnHandle) -> bool {
+	!(matches!(reason, ShutdownKind::Sleep) && conn.is_hibernatable())
+}
+
 #[cfg(all(unix, target_env = "gnu"))]
 fn trim_native_allocator_after_shutdown(actor_id: &str, reason: &str) {
 	unsafe extern "C" {

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-drain.ts

@@ -11,5 +11,10 @@ export const connDrainActor = actor({
 	},
 	actions: {
 		ping: () => "pong",
+		work: async (c) => {
+			c.broadcast("workStarted");
+			await new Promise((resolve) => setTimeout(resolve, 300));
+			return "done";
+		},
 	},
 });

rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-drain.test.ts

@@ -1,4 +1,4 @@
-import { expect, test } from "vitest";
+import { expect, type TestContext, test } from "vitest";
 import { CONN_DRAIN_DISCONNECT_MARKER } from "../../fixtures/driver-test-suite/conn-drain";
 import type { registry } from "../../fixtures/driver-test-suite/registry-static";
 import { createClient } from "@/client/mod";
@@ -8,7 +8,7 @@ import { describeDriverMatrix } from "./shared-matrix";
 describeDriverMatrix(
 	"Connection Drain",
 	(driverTestConfig) => {
-		test("closes connections with ws.going_away and runs onDisconnect once", async (c) => {
+		async function startDrainingRuntime(c: TestContext) {
 			const runtime = await driverTestConfig.start({
 				env: { RIVETKIT_TEST_DRAIN_ON_SIGTERM: "1" },
 			});
@@ -23,7 +23,11 @@ describeDriverMatrix(
 				await client.dispose();
 				await runtime.cleanup();
 			});
+			return { runtime, client };
+		}
 
+		test("closes connections with ws.going_away and runs onDisconnect once", async (c) => {
+			const { runtime, client } = await startDrainingRuntime(c);
 			const key = crypto.randomUUID();
 			const conn = client.connDrainActor.getOrCreate([key]).connect();
 			const errors: ActorError[] = [];
@@ -41,9 +45,26 @@ describeDriverMatrix(
 			).toEqual(["ws.going_away"]);
 			const output = runtime.getRuntimeOutput?.() ?? "";
 			const hookRuns =
-				output.split(`${CONN_DRAIN_DISCONNECT_MARKER} ${key}`).length - 1;
+				output.split(`${CONN_DRAIN_DISCONNECT_MARKER} ${key}`).length -
+				1;
 			expect(hookRuns).toBe(1);
 		});
+
+		test("delivers an in-flight action result before disconnecting", async (c) => {
+			const { runtime, client } = await startDrainingRuntime(c);
+			const conn = client.connDrainActor
+				.getOrCreate([crypto.randomUUID()])
+				.connect();
+			const started = new Promise<void>((resolve) =>
+				conn.once("workStarted", () => resolve()),
+			);
+
+			const result = conn.work();
+			await started;
+			await runtime.cleanup();
+
+			expect(await result).toBe("done");
+		});
 	},
 	{
 		encodings: ["bare"],