feat(pi): support RivetKit 3.0.0-preview

examples/docs/package.json

@@ -18,14 +18,14 @@
 		"@rivet-dev/pi": "workspace:*",
 		"@rivet-dev/sandbox-adapter": "workspace:*",
 		"@rivet-dev/workflows": "1.0.0",
-		"@rivetkit/react": "2.3.30",
+		"@rivetkit/react": "3.0.0-preview.2",
 		"@slack/web-api": "^8.2.0",
 		"discord-interactions": "^4.4.0",
 		"e2b": "^2.51.0",
 		"freestyle": "^0.2.16",
 		"hono": "^4.13.12",
 		"react": "^19.3.0",
-		"rivetkit": "2.3.30"
+		"rivetkit": "3.0.0-preview.2"
 	},
 	"devDependencies": {
 		"@types/node": "^24.0.0",

packages/pi/package.json

@@ -55,7 +55,7 @@
 		"@earendil-works/pi-ai": "^1.4.0",
 		"@earendil-works/pi-durable": "^1.4.0",
 		"@earendil-works/pi-telemetry": "^1.4.0",
-		"rivetkit": ">=2.3.28 <2.4.0"
+		"rivetkit": ">=2.3.28 <2.4.0 || ^3.0.0-preview.2"
 	},
 	"devDependencies": {
 		"@earendil-works/chord": "1.4.0",
@@ -64,7 +64,7 @@
 		"@earendil-works/pi-telemetry": "1.4.0",
 		"@opentelemetry/sdk-trace-node": "2.11.0",
 		"@types/node": "^24.0.0",
-		"rivetkit": "2.3.30",
+		"rivetkit": "3.0.0-preview.2",
 		"typescript": "^5.9.2",
 		"vitest": "^3.2.4"
 	},

packages/pi/src/actor.ts

@@ -27,6 +27,7 @@ import {
 	type PiContext,
 	type PiOptions,
 	piRuntime,
+	piStopping,
 	withPiAccess,
 	withUnavailablePi,
 } from "./runtime.js";
@@ -240,6 +241,16 @@ export function pi<
 		},
 		onConnect: withPi(actorConfig.onConnect),
 		onDisconnect: async (c: PiContext, conn: { id: string }) => {
+			// Shutdown runs onDisconnect in the same window as onSleep and
+			// onDestroy, in no set order, and those close Pi Durable. So a
+			// disconnect during shutdown never uses Pi, and its watches end when
+			// the harness closes.
+			if (c.abortSignal.aborted || piRuntime(c).status !== "open") {
+				await withUnavailablePi(c, piStopping(), () =>
+					actorConfig.onDisconnect?.(c, conn),
+				);
+				return;
+			}
 			try {
 				await userOnDisconnect?.(c, conn);
 			} finally {

packages/pi/src/runtime.ts

@@ -719,6 +719,11 @@ export async function resumeDueSleeps(c: PiContext): Promise<void> {
 	(await piRuntime(c).harness)?.resumeDue();
 }
 
+/** What a shutdown hook gets from `c.pi` once the actor has started to stop. */
+export function piStopping() {
+	return new UserError("Pi Durable is closed because the actor is stopping.");
+}
+
 /** Runs `body` with a `c.pi` that throws `error` when read. */
 export function withUnavailablePi<T>(
 	c: PiContext,

packages/pi/tests/helpers/engine-root.ts

@@ -0,0 +1,11 @@
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+/** One directory for the engine data of every test file in this run, removed when the run ends. */
+export default function setup(): () => void {
+	const root = mkdtempSync(join(tmpdir(), "rivet-pi-test-"));
+	// Test files run in child processes, which inherit this.
+	process.env.RIVET_PI_TEST_ROOT = root;
+	return () => rmSync(root, { recursive: true, force: true });
+}

packages/pi/tests/helpers/isolated-engine.ts

@@ -0,0 +1,30 @@
+import { mkdtempSync } from "node:fs";
+import { createServer } from "node:net";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+// Each test file runs in its own process and starts its own engine. An engine
+// from the previous file can still hold its data directory while it shuts
+// down, and a new engine that finds it locked exits. So each file gets its own
+// data directory (under HOME) and port.
+// `engine-root.ts` removes the run's directory when the run ends.
+process.env.HOME = mkdtempSync(
+	join(process.env.RIVET_PI_TEST_ROOT ?? tmpdir(), "home-"),
+);
+process.env.RIVET_RUN_ENGINE_PORT = String(await freePort());
+
+/** A free port with room above it, because the engine also listens on the next ports. */
+async function freePort(): Promise<number> {
+	for (;;) {
+		const port = 20_000 + Math.floor(Math.random() * 400) * 100;
+		if (await isFree(port)) return port;
+	}
+}
+
+function isFree(port: number): Promise<boolean> {
+	return new Promise((resolve) => {
+		const server = createServer();
+		server.once("error", () => resolve(false));
+		server.listen(port, "127.0.0.1", () => server.close(() => resolve(true)));
+	});
+}

packages/pi/tests/upgrade.test.ts

@@ -509,18 +509,19 @@ describe("pi drain and upgrades", () => {
 		]);
 	});
 
-	test("the app's onDisconnect runs after a destroy closed Pi Durable, and only reading c.pi fails", async (c) => {
+	test("the app's onDisconnect runs when a destroy closes its connection, and c.pi is open or reports the stop", async (c) => {
 		const { client } = await setupTest(c, registry);
 		const key = ["closed", randomUUID()];
 		const conn = client.drainer.getOrCreate(key).connect();
 		await conn.harness.root();
 
 		await client.drainer.getOrCreate(key).destroySelf();
-		// Destroy closes Pi Durable, then disconnects the connection; the hook records what c.pi gave in this process.
+		// RivetKit 2.x runs onDestroy first; 3.0 disconnects before shutdown starts. Either way the hook runs, and c.pi never fails with another error. The hook records what c.pi gave in this process.
 		await vi.waitFor(() =>
-			expect(disconnectOutcomes.get(JSON.stringify(key))).toBe(
-				"Pi Durable is closed because the actor was destroyed.",
-			),
+			expect([
+				"open",
+				"Pi Durable is closed because the actor is stopping.",
+			]).toContain(disconnectOutcomes.get(JSON.stringify(key))),
 		);
 		await conn.dispose();
 	});

packages/pi/vitest.config.ts

@@ -5,6 +5,8 @@ export default defineConfig({
 		// Each file boots its own engine through setupTest. Running files in
 		// parallel makes those engines contend for ports and CPU.
 		fileParallelism: false,
+		globalSetup: ["./tests/helpers/engine-root.ts"],
+		setupFiles: ["./tests/helpers/isolated-engine.ts"],
 		sequence: { concurrent: false },
 		testTimeout: 60_000,
 		hookTimeout: 60_000,