perf(pi): connect the pi() sandbox on first use

packages/pi/src/actor-types.ts

@@ -0,0 +1,31 @@
+/** `Omit` that keeps each member of a union, so `state` and `createState` stay alternatives. */
+export type DistributiveOmit<T, K extends PropertyKey> = T extends unknown
+	? Omit<T, K>
+	: never;
+
+/** The app's actions as clients see them: the same arguments, without `c.pi`. */
+export type ClientActions<T, TContext> = string extends keyof T
+	? Record<never, never>
+	: {
+			[K in keyof T]: T[K] extends (
+				c: any,
+				...args: infer TArgs
+			) => infer TResult
+				? (c: TContext, ...args: TArgs) => TResult
+				: ClientActions<T[K], TContext>;
+		};
+
+/** Wraps each of the app's actions, including nested ones, so it runs inside `run`. */
+export function wrapActions(
+	actions: Record<string, unknown>,
+	run: (c: any, action: () => unknown) => Promise<unknown>,
+): Record<string, unknown> {
+	return Object.fromEntries(
+		Object.entries(actions).map(([name, action]) => [
+			name,
+			typeof action === "function"
+				? (c: unknown, ...args: unknown[]) => run(c, () => action(c, ...args))
+				: wrapActions(action as Record<string, unknown>, run),
+		]),
+	);
+}

packages/pi/src/actor.ts

@@ -1,17 +1,31 @@
-import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent";
+import type {
+	AgentSession,
+	AgentSessionEvent,
+} from "@earendil-works/pi-coding-agent";
 import {
-	type Actions,
+	type ActionContext,
 	type ActorConfigInput,
 	type ActorDefinition,
 	actor,
+	type Conn,
+	type ConnectContext,
+	type DestroyContext,
+	type DisconnectContext,
 	type EventSchemaConfig,
 	event,
 	type QueueSchemaConfig,
+	type SleepContext,
 	type Type,
+	type WakeContext,
 } from "rivetkit";
 import { db } from "rivetkit/db";
 import { drainDeadline } from "./drain.js";
 import { createPiActions, type PiActions } from "./actions.js";
+import {
+	type ClientActions,
+	type DistributiveOmit,
+	wrapActions,
+} from "./actor-types.js";
 import {
 	closePiSession,
 	createPiRuntime,
@@ -21,6 +35,7 @@ import {
 	type PiSessionOptions,
 	resumeInterruptedRun,
 	stopRunWhenLost,
+	withPiSession,
 } from "./runtime.js";
 import { migratePiTables } from "./storage.js";
 
@@ -61,6 +76,151 @@ const piOptionKeys = [
 	"sandbox",
 ] as const satisfies readonly (keyof PiSessionOptions)[];
 
+/**
+ * `c.pi`: the actor's Pi session, in the app's own actions and hooks. Reading
+ * it throws when the session cannot open, such as while the sandbox provider
+ * is down. In `onSleep` and `onDestroy` it is set only when the session is
+ * already open in this actor generation.
+ */
+export interface PiAccess {
+	readonly pi: AgentSession;
+}
+
+/** The app's own actions, which get `c.pi`. Nested objects become dotted action names. */
+export interface PiUserActions<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> {
+	[action: string]:
+		| ((
+				c: ActionContext<
+					TState,
+					TConnParams,
+					TConnState,
+					TVars,
+					TInput,
+					PiDatabaseProvider,
+					TEvents,
+					TQueues
+				> &
+					PiAccess,
+				...args: any[]
+		  ) => any)
+		| PiUserActions<
+				TState,
+				TConnParams,
+				TConnState,
+				TVars,
+				TInput,
+				TEvents,
+				TQueues
+		  >;
+}
+
+/** Lifecycle hooks that get `c.pi`. */
+interface PiHooks<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> {
+	onWake?: (
+		c: WakeContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		> &
+			PiAccess,
+	) => void | Promise<void>;
+	onSleep?: (
+		c: SleepContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		> &
+			PiAccess,
+	) => void | Promise<void>;
+	onDestroy?: (
+		c: DestroyContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		> &
+			PiAccess,
+	) => void | Promise<void>;
+	onConnect?: (
+		c: ConnectContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		> &
+			PiAccess,
+		conn: Conn<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		>,
+	) => void | Promise<void>;
+	onDisconnect?: (
+		c: DisconnectContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		> &
+			PiAccess,
+		conn: Conn<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		>,
+	) => void | Promise<void>;
+}
+
+type WrappedKey = "actions" | keyof PiHooks<any, any, any, any, any, any, any>;
+
 export type PiActorConfigInput<
 	TState = undefined,
 	TConnParams = undefined,
@@ -69,41 +229,38 @@ export type PiActorConfigInput<
 	TInput = undefined,
 	TEvents extends EventSchemaConfig = Record<never, never>,
 	TQueues extends QueueSchemaConfig = Record<never, never>,
-	TUserActions extends Actions<
+	TUserActions extends PiUserActions<
 		TState,
 		TConnParams,
 		TConnState,
 		TVars,
 		TInput,
-		PiDatabaseProvider,
 		TEvents,
 		TQueues
 	> = Record<never, never>,
-> = ActorConfigInput<
-	TState,
-	TConnParams,
-	TConnState,
-	TVars,
-	TInput,
-	PiDatabaseProvider,
-	TEvents,
-	TQueues,
-	TUserActions
+> = DistributiveOmit<
+	ActorConfigInput<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		PiDatabaseProvider,
+		TEvents,
+		TQueues
+	>,
+	WrappedKey
 > &
-	PiSessionOptions;
-
-/**
- * The actions the user defined. With no `actions` option, TypeScript uses the
- * default, a record with an index signature, which would hide the Pi actions.
- */
-type UserActions<T> = string extends keyof T ? Record<never, never> : T;
+	PiHooks<TState, TConnParams, TConnState, TVars, TInput, TEvents, TQueues> &
+	PiSessionOptions & { actions?: TUserActions };
 
 /**
  * Defines a Rivet Actor that owns one Pi coding-agent session.
  *
  * The session transcript and settings live in the actor's SQLite database and
  * are restored when the actor wakes. Pi events are broadcast on `event`.
  * Ordinary actor config (state, vars, actions, events, hooks) is passed through.
+ * The app's own actions and lifecycle hooks get the session as `c.pi`.
  */
 export function pi<
 	TState,
@@ -113,22 +270,20 @@ export function pi<
 	TInput,
 	TEvents extends EventSchemaConfig = Record<never, never>,
 	TQueues extends QueueSchemaConfig = Record<never, never>,
-	TUserActions extends Actions<
+	TUserActions extends PiUserActions<
 		TState,
 		TConnParams,
 		TConnState,
 		TVars,
 		TInput,
-		PiDatabaseProvider,
 		TEvents,
 		TQueues
-	> = Actions<
+	> = PiUserActions<
 		TState,
 		TConnParams,
 		TConnState,
 		TVars,
 		TInput,
-		PiDatabaseProvider,
 		TEvents,
 		TQueues
 	>,
@@ -161,7 +316,20 @@ export function pi<
 	PiDatabaseProvider,
 	TEvents & PiEvents,
 	TQueues,
-	UserActions<TUserActions> & PiActions
+	ClientActions<
+		TUserActions,
+		ActionContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			PiDatabaseProvider,
+			TEvents,
+			TQueues
+		>
+	> &
+		PiActions
 > {
 	const { actorConfig, sessionOptions } = splitConfig(config);
 	if (actorConfig.db !== undefined) {
@@ -171,11 +339,20 @@ export function pi<
 	assertNoReservedKeys("action", actorConfig.actions, actions);
 	assertNoReservedKeys("event", actorConfig.events, piEvents);
 
+	const withPi = <TArgs extends unknown[]>(
+		hook: ((c: PiContext, ...args: TArgs) => unknown) | undefined,
+		opensSession = true,
+	) =>
+		hook &&
+		((c: PiContext, ...args: TArgs) =>
+			withPiSession(c, sessionOptions, () => hook(c, ...args), {
+				opensSession,
+			}));
 	const userVars = actorConfig.vars;
 	const userCreateVars = actorConfig.createVars;
-	const userOnWake = actorConfig.onWake;
-	const userOnSleep = actorConfig.onSleep;
-	const userOnDestroy = actorConfig.onDestroy;
+	const userOnWake = withPi(actorConfig.onWake);
+	const userOnSleep = withPi(actorConfig.onSleep, false);
+	const userOnDestroy = withPi(actorConfig.onDestroy, false);
 	delete actorConfig.vars;
 	const sleepGracePeriod: number =
 		actorConfig.options?.sleepGracePeriod ?? DEFAULT_SLEEP_GRACE_PERIOD_MS;
@@ -189,7 +366,12 @@ export function pi<
 		},
 		db: db({ onMigrate: migratePiTables }),
 		events: { ...(actorConfig.events ?? {}), ...piEvents },
-		actions: { ...(actorConfig.actions ?? {}), ...actions },
+		actions: {
+			...wrapActions(actorConfig.actions ?? {}, (c: PiContext, action) =>
+				withPiSession(c, sessionOptions, action),
+			),
+			...actions,
+		},
 		createVars: async (c: unknown, driverCtx: unknown) => {
 			const vars = userCreateVars
 				? await userCreateVars(c, driverCtx)
@@ -208,6 +390,8 @@ export function pi<
 				});
 			});
 		},
+		onConnect: withPi(actorConfig.onConnect),
+		onDisconnect: withPi(actorConfig.onDisconnect),
 		onSleep: async (c: PiContext) => {
 			const drainUntil = drainDeadline(sleepGracePeriod);
 			try {
@@ -226,17 +410,7 @@ export function pi<
 				await closePiSession(c, sessionOptions, { reason: "destroy" });
 			}
 		},
-	} as any) as ActorDefinition<
-		TState,
-		TConnParams,
-		TConnState,
-		TVars,
-		TInput,
-		PiDatabaseProvider,
-		TEvents & PiEvents,
-		TQueues,
-		UserActions<TUserActions> & PiActions
-	>;
+	} as any) as any;
 }
 
 function splitConfig(config: object): {