chore(pi): require RivetKit 2.3.26

examples/docs/package.json

@@ -17,13 +17,13 @@
 		"@rivet-dev/pi": "workspace:*",
 		"@rivet-dev/sandbox-adapter": "workspace:*",
 		"@rivet-dev/workflows": "1.0.0",
-		"@rivetkit/react": "2.3.23",
+		"@rivetkit/react": "2.3.26",
 		"@slack/web-api": "^8.2.0",
 		"discord-interactions": "^4.4.0",
 		"e2b": "^2.51.0",
 		"hono": "^4.13.12",
 		"react": "^19.3.0",
-		"rivetkit": "2.3.23"
+		"rivetkit": "2.3.26"
 	},
 	"devDependencies": {
 		"@types/node": "^24.0.0",

packages/pi/package.json

@@ -53,15 +53,15 @@
 	"peerDependencies": {
 		"@earendil-works/chord": "^1.0.0",
 		"@earendil-works/pi-durable": "^1.0.0",
-		"rivetkit": ">=2.3.23 <2.4.0"
+		"rivetkit": ">=2.3.26 <2.4.0"
 	},
 	"devDependencies": {
 		"@earendil-works/chord": "^1.1.0",
 		"@earendil-works/pi-ai": "^1.1.0",
 		"@earendil-works/pi-durable": "^1.1.0",
 		"@opentelemetry/sdk-trace-node": "2.11.0",
 		"@types/node": "^24.0.0",
-		"rivetkit": "2.3.23",
+		"rivetkit": "2.3.26",
 		"typescript": "^5.9.2",
 		"vitest": "^3.2.4"
 	},

packages/pi/src/actions.ts

@@ -27,6 +27,7 @@ import {
 	documentsByKind,
 	type PiOptions,
 	piRuntime,
+	type Waiting,
 	withHarness,
 } from "./runtime.js";
 import type { WatchSpec } from "./storage.js";
@@ -94,7 +95,7 @@ export function createPiActions(options: PiOptions) {
 	const documents = documentsByKind(options);
 	const run = <T>(
 		c: PiContext,
-		body: (harness: Harness, context: Context) => Promise<T>,
+		body: (harness: Harness, context: Context, waiting: Waiting) => Promise<T>,
 	) => withHarness(c, options, body);
 	const agentChange = (change: PiAgentChange | undefined) =>
 		change === undefined ? undefined : resolveAgentChange(options, change);
@@ -193,12 +194,14 @@ export function createPiActions(options: PiOptions) {
 					return harness.abortTask(id, context);
 				}),
 			waitForTask: (c: PiContext, id: TaskId) =>
-				run(c, async (harness, context) => {
+				run(c, async (harness, context, waiting) => {
 					await requireTask(harness, id, context);
-					return harness.waitForTask(id, context);
+					return waiting(() => harness.waitForTask(id, context));
 				}),
 			waitForIdle: (c: PiContext) =>
-				run(c, (harness, context) => harness.waitForIdle(context)),
+				run(c, (harness, context, waiting) =>
+					waiting(() => harness.waitForIdle(context)),
+				),
 			inspect: (c: PiContext) =>
 				run(c, async (harness, context) => {
 					const inspection = await harness.inspect(context);
@@ -360,11 +363,14 @@ export function createPiActions(options: PiOptions) {
 					),
 				),
 			waitForIdle: (c: PiContext, conversationId: ConversationId) =>
-				run(c, async (harness, context) =>
-					(
-						await requireConversation(harness, conversationId, context)
-					).waitForIdle(context),
-				),
+				run(c, async (harness, context, waiting) => {
+					const conversation = await requireConversation(
+						harness,
+						conversationId,
+						context,
+					);
+					return waiting(() => conversation.waitForIdle(context));
+				}),
 			/** The structural view; then `pi.view` frames with Chord operations to this connection. */
 			watch: (c: PiContext, conversationId: ConversationId) =>
 				watchFor<{ value: ConversationView }>(c, {
@@ -389,9 +395,10 @@ export function createPiActions(options: PiOptions) {
 					(await requireSubmission(harness, id, context)).status(context),
 				),
 			wait: (c: PiContext, id: SubmissionId) =>
-				run(c, async (harness, context) =>
-					(await requireSubmission(harness, id, context)).wait(context),
-				),
+				run(c, async (harness, context, waiting) => {
+					const submission = await requireSubmission(harness, id, context);
+					return waiting(() => submission.wait(context));
+				}),
 			abort: (c: PiContext, id: SubmissionId) =>
 				run(c, async (harness, context) =>
 					(await requireSubmission(harness, id, context)).abort(context),
@@ -408,7 +415,7 @@ export function createPiActions(options: PiOptions) {
 			content: UserInput,
 			promptOptions?: PiPromptOptions,
 		) =>
-			run(c, async (harness, context): Promise<PiPromptResult> => {
+			run(c, async (harness, context, waiting): Promise<PiPromptResult> => {
 				const conversation =
 					promptOptions?.conversationId === undefined
 						? await harness.root(context)
@@ -431,7 +438,7 @@ export function createPiActions(options: PiOptions) {
 					c.actorId,
 					conversation.id,
 					model && `${model.provider}/${model.modelId}`,
-					() => submission.wait(context),
+					() => waiting(() => submission.wait(context)),
 				);
 				if (settled.status === "unanswered") {
 					return {

packages/pi/src/runtime.ts

@@ -48,6 +48,9 @@ export type PiContext = ActorContext<
 	any
 >;
 
+/** Whether Pi Durable is open in this actor generation, or why it closed. */
+export type PiStatus = "open" | "sleeping" | "destroyed";
+
 /** How an actor generation stops. A sleep drains until `drainUntil`; a destroy stops at once. */
 export type PiStop =
 	| { reason: "sleep"; drainUntil: number }
@@ -111,8 +114,8 @@ export interface PiRuntime {
 	/** The open harness, or the open in progress. */
 	harness?: Promise<OpenHarness>;
 	sandbox?: Promise<ConnectedSandbox>;
-	/** Set when the actor starts to sleep or is destroyed. */
-	closing: boolean;
+	/** Pi Durable is open until the actor starts to sleep or is destroyed. */
+	status: PiStatus;
 	/**
 	 * The `onWake` context. Background work (keep-awake, scheduled wakes, the
 	 * sandbox, reopening) uses it, because an action's context is cancelled
@@ -121,12 +124,19 @@ export interface PiRuntime {
 	actor?: PiContext;
 	/** Open Pi watches of connections, for this generation. */
 	watches: ConnectionWatches;
+	/**
+	 * Calls that are using the harness and are not waiting on a run. A sleep
+	 * drains them, so it never closes the harness under a call it admitted.
+	 */
+	calls: number;
+	/** Called when a counted call ends. A drain sets it. */
+	callEnded?: () => void;
 }
 
 export const PI_RUNTIME: unique symbol = Symbol.for("@rivet-dev/pi/runtime");
 
 export function createPiRuntime(): PiRuntime {
-	return { closing: false, watches: createConnectionWatches() };
+	return { status: "open", watches: createConnectionWatches(), calls: 0 };
 }
 
 export function piRuntime(c: PiContext): PiRuntime {
@@ -161,7 +171,10 @@ async function ensureHarness(
 	open: { resume: boolean } = { resume: true },
 ): Promise<Harness> {
 	const runtime = piRuntime(c);
-	if (runtime.closing) throw harnessClosed();
+	// Callers check the status first, so a closed generation never reopens Pi Durable.
+	if (runtime.status !== "open") {
+		throw new Error(`pi durable cannot open after it closed (${runtime.status})`);
+	}
 	runtime.harness ??= openHarness(
 		runtime.actor ?? c,
 		runtime,
@@ -185,18 +198,43 @@ async function ensureHarness(
 export async function withHarness<T>(
 	c: PiContext,
 	options: PiOptions,
-	run: (harness: Harness, context: Context) => Promise<T>,
+	run: (harness: Harness, context: Context, waiting: Waiting) => Promise<T>,
 ): Promise<T> {
+	const runtime = piRuntime(c);
+	if (runtime.status !== "open") throw piClosed(runtime.status);
 	const harness = await ensureHarness(c, options);
+	const release = countCall(runtime);
 	try {
-		return await run(harness, BACKGROUND_CONTEXT);
+		return await run(harness, BACKGROUND_CONTEXT, (wait) => {
+			release();
+			return wait();
+		});
 	} catch (error) {
-		const runtime = piRuntime(c);
 		await recoverHarness(runtime.actor ?? c, runtime, options);
 		throw toClientError(error);
+	} finally {
+		release();
 	}
 }
 
+/**
+ * Runs a wait on a run, such as `submission.wait`, without holding a sleep.
+ * While the run is busy, the run itself holds the sleep.
+ */
+export type Waiting = <W>(wait: () => Promise<W>) => Promise<W>;
+
+/** Counts a call until the returned function runs. Calling it again does nothing. */
+function countCall(runtime: PiRuntime): () => void {
+	runtime.calls += 1;
+	let counted = true;
+	return () => {
+		if (!counted) return;
+		counted = false;
+		runtime.calls -= 1;
+		runtime.callEnded?.();
+	};
+}
+
 /**
  * Runs one of the app's own actions or hooks with the harness on `c.pi`. A
  * failure that left Pi Durable unusable reopens it. The error reaches the
@@ -213,7 +251,9 @@ export async function withPiAccess<T>(
 ): Promise<T> {
 	const runtime = piRuntime(c);
 	// Hooks such as onDisconnect still run after Pi Durable closed; only using c.pi fails.
-	if (runtime.closing) return withUnavailablePi(c, harnessClosed(), body);
+	if (runtime.status !== "open") {
+		return withUnavailablePi(c, piClosed(runtime.status), body);
+	}
 	if (!opensHarness && !runtime.harness) {
 		return withUnavailablePi(c, harnessNotOpen(), body);
 	}
@@ -228,11 +268,14 @@ export async function withPiAccess<T>(
 		return withUnavailablePi(c, error, body);
 	}
 	Object.defineProperty(c, "pi", { value: harness, configurable: true });
+	const release = countCall(runtime);
 	try {
 		return await body();
 	} catch (error) {
 		await recoverHarness(runtime.actor ?? c, runtime, options);
 		throw toClientError(error);
+	} finally {
+		release();
 	}
 }
 
@@ -242,8 +285,13 @@ function harnessNotOpen() {
 	);
 }
 
-function harnessClosed() {
-	return new UserError("Pi Durable is closed because the actor is stopping.");
+/** What a call gets once Pi Durable is closed for this generation. */
+function piClosed(status: Exclude<PiStatus, "open">) {
+	return new UserError(
+		status === "sleeping"
+			? "Pi Durable is closed because the actor is going to sleep."
+			: "Pi Durable is closed because the actor was destroyed.",
+	);
 }
 
 async function openHarness(
@@ -316,14 +364,18 @@ async function recoverHarness(
 	options: PiOptions,
 ): Promise<void> {
 	const current = runtime.harness;
-	if (!current || runtime.closing) return;
+	if (!current || runtime.status !== "open") return;
 	let open: OpenHarness;
 	try {
 		open = await current;
 	} catch {
 		return;
 	}
-	if (runtime.harness !== current || runtime.closing || isUsable(open.harness))
+	if (
+		runtime.harness !== current ||
+		runtime.status !== "open" ||
+		isUsable(open.harness)
+	)
 		return;
 
 	c.log.warn({
@@ -578,19 +630,22 @@ function watchBusy(
 }
 
 /**
- * Waits until no run is active, or until `drainUntil` passes. Idle sleep
- * never has live work, so this only waits on a forced stop, such as an upgrade.
- * The deadline ends before RivetKit's, so closing still fits in the grace
- * period. Pi's records are already stored, and the next wake resumes the work.
+ * Waits until no run is active and no admitted call is using the harness, or
+ * until `drainUntil` passes. Idle sleep has no live work, so it waits only for
+ * a call that arrived as the sleep started. The deadline ends before
+ * RivetKit's, so closing still fits in the grace period. Pi's records are
+ * already stored, and the next wake resumes the work.
  */
 async function drain(
 	c: PiContext,
+	runtime: PiRuntime,
 	open: OpenHarness,
 	drainUntil: number,
 ): Promise<void> {
 	let changed = () => {};
 	let timedOut = false;
 	const unsubscribe = open.harness.subscribeCommits(() => changed());
+	runtime.callEnded = () => changed();
 	const timer = setTimeout(
 		() => {
 			timedOut = true;
@@ -604,7 +659,7 @@ async function drain(
 			const next = new Promise<void>((resolve) => {
 				changed = resolve;
 			});
-			if (!(await open.isBusy())) return;
+			if (runtime.calls === 0 && !(await open.isBusy())) return;
 			await next;
 		}
 		c.log.warn({
@@ -613,6 +668,7 @@ async function drain(
 	} finally {
 		clearTimeout(timer);
 		unsubscribe();
+		runtime.callEnded = undefined;
 	}
 }
 
@@ -635,11 +691,11 @@ export async function closeHarness(
 	} catch {}
 	// Destroy has no next generation to resume in, so it stops running work at once.
 	if (open && stop.reason === "sleep") {
-		await drain(c, open, stop.drainUntil).catch((error: unknown) =>
+		await drain(c, runtime, open, stop.drainUntil).catch((error: unknown) =>
 			errors.push(error),
 		);
 	}
-	runtime.closing = true;
+	runtime.status = stop.reason === "sleep" ? "sleeping" : "destroyed";
 	runtime.harness = undefined;
 	if (open) {
 		await closeOpenHarness(open, runtime.watches).catch((error: unknown) =>

packages/pi/tests/actor.test.ts

@@ -22,7 +22,7 @@ import type { SandboxProvider } from "@rivet-dev/sandbox-adapter";
 import { setup } from "rivetkit";
 import { setupTest } from "rivetkit/test";
 import { afterAll, beforeAll, describe, expect, test, vi } from "vitest";
-import { pi } from "../src/index.js";
+import { type PiEventsFrame, pi } from "../src/index.js";
 import { localSandboxProvider } from "./helpers/local-sandbox.js";
 import {
 	createMockModel,
@@ -500,6 +500,41 @@ describe("pi actor", () => {
 		expect(result).toMatchObject({ status: "done", text: "Hi there!" });
 	});
 
+	test("a watching connection gets a fresh snapshot when a failed commit reopens the harness, then the next run's batches", async (c) => {
+		const { client } = await setupTest(c, registry);
+		const conn = client.agent
+			.getOrCreate(["poison-watch", randomUUID()])
+			.connect();
+		const frames: PiEventsFrame[] = [];
+		// SAFETY: RivetKit types event payloads as unknown on the client, and pi() sends only PiEventsFrame on pi.events.
+		conn.on("pi.events", (frame) => void frames.push(frame as PiEventsFrame));
+		const root = await conn.harness.root();
+		await conn.conversation.watchEvents(root.id);
+		await conn.failCommitsOfPoison();
+
+		await expect(conn.writePoison()).rejects.toThrow();
+		expect(await conn.prompt("say hello")).toMatchObject({
+			status: "done",
+			text: "Hi there!",
+		});
+
+		// The reopened harness reattaches the watch: a seq 0 snapshot, then the run's batches in order.
+		await vi.waitFor(() =>
+			expect(
+				frames.some((frame) =>
+					frame.events.some((event) => event.type === "run_end"),
+				),
+			).toBe(true),
+		);
+		const reattached = frames.findIndex((frame) => frame.seq === 0);
+		expect(frames[reattached]?.events).toEqual([
+			expect.objectContaining({ type: "snapshot" }),
+		]);
+		expect(frames.slice(reattached).map((frame) => frame.seq)).toEqual(
+			frames.slice(reattached).map((_, index) => index),
+		);
+	});
+
 	test("two conversations prompted at the same time each get their own answer, and neither transcript holds the other's messages", async (c) => {
 		const { client } = await setupTest(c, registry);
 		const handle = client.agent.getOrCreate(["two", randomUUID()]);