rivetkit 2.1.4 → 2.1.6-rc.1

package/src/sandbox/config.ts

@@ -0,0 +1,151 @@
+import { z } from "zod/v4";
+import type {
+	ActorContext,
+	BeforeConnectContext,
+} from "@/actor/contexts";
+import type { AnyDatabaseProvider } from "@/actor/database";
+import type {
+	PermissionRequestListener,
+	SessionEventListener,
+	SandboxProvider,
+} from "sandbox-agent";
+import type {
+	SandboxActorVars,
+	SandboxActorState,
+} from "./types";
+
+const zFunction = <
+	T extends (...args: any[]) => any = (...args: unknown[]) => unknown,
+>() => z.custom<T>((val) => typeof val === "function");
+
+const SandboxProviderSchema = z.object({
+	name: z.string(),
+	create: zFunction<SandboxProvider["create"]>(),
+	destroy: zFunction<SandboxProvider["destroy"]>(),
+	getUrl: zFunction<NonNullable<SandboxProvider["getUrl"]>>().optional(),
+	getFetch: zFunction<NonNullable<SandboxProvider["getFetch"]>>().optional(),
+	ensureServer: zFunction<NonNullable<SandboxProvider["ensureServer"]>>().optional(),
+});
+
+export const SandboxActorOptionsSchema = z
+	.object({
+		// Log if the actor still thinks a turn is active but no new session event
+		// has arrived for this long.
+		warningAfterMs: z.number().nonnegative().default(30_000),
+		// Clear active-turn state after this timeout so a missing terminal event
+		// cannot keep the actor awake forever.
+		staleAfterMs: z.number().positive().default(5 * 60_000),
+	})
+	.strict()
+	.prefault(() => ({}))
+	.transform((value) => ({
+		...value,
+		warningAfterMs: Math.min(value.warningAfterMs, value.staleAfterMs),
+	}));
+
+export type SandboxActorOptions = z.input<typeof SandboxActorOptionsSchema>;
+export type SandboxActorOptionsRuntime = z.infer<
+	typeof SandboxActorOptionsSchema
+>;
+
+// This schema validates the config at runtime. Generic callback types are
+// defined separately below following the same pattern as ActorConfigSchema:
+// infer from the schema, omit function keys, then intersect typed callbacks.
+export const SandboxActorConfigSchema = z
+	.object({
+		provider: SandboxProviderSchema.optional(),
+		createProvider: zFunction().optional(),
+		persistRawEvents: z.boolean().optional(),
+		destroyActor: z.boolean().default(false),
+		options: SandboxActorOptionsSchema,
+		onBeforeConnect: zFunction().optional(),
+		onSessionEvent: zFunction().optional(),
+		onPermissionRequest: zFunction().optional(),
+	})
+	.strict()
+	.refine(
+		(data) =>
+			(data.provider !== undefined) !==
+			(data.createProvider !== undefined),
+		{
+			message:
+				"Sandbox actor config must define exactly one of 'provider' or 'createProvider'",
+		},
+	);
+
+// --- Typed config types (generic callbacks overlaid on the Zod schema) ---
+
+type SandboxActorContext<TConnParams> = ActorContext<
+	SandboxActorState,
+	TConnParams,
+	undefined,
+	SandboxActorVars,
+	undefined,
+	AnyDatabaseProvider
+>;
+
+interface SandboxActorConfigCallbacks<TConnParams> {
+	onBeforeConnect?: (
+		c: BeforeConnectContext<
+			SandboxActorState,
+			SandboxActorVars,
+			undefined,
+			AnyDatabaseProvider
+		>,
+		params: TConnParams,
+	) => void | Promise<void>;
+	onSessionEvent?: (
+		c: SandboxActorContext<TConnParams>,
+		sessionId: string,
+		event: Parameters<SessionEventListener>[0],
+	) => void | Promise<void>;
+	onPermissionRequest?: (
+		c: SandboxActorContext<TConnParams>,
+		sessionId: string,
+		request: Parameters<PermissionRequestListener>[0],
+	) => void | Promise<void>;
+}
+
+type SandboxActorProviderConfig<TConnParams> =
+	| {
+			provider: SandboxProvider;
+			createProvider?: never;
+	  }
+	| {
+			provider?: never;
+			createProvider: (
+				c: SandboxActorContext<TConnParams>,
+			) => SandboxProvider | Promise<SandboxProvider>;
+	  };
+
+// Parsed config (after Zod defaults/transforms applied).
+export type SandboxActorConfig<TConnParams = undefined> = Omit<
+	z.infer<typeof SandboxActorConfigSchema>,
+	| "provider"
+	| "createProvider"
+	| "onBeforeConnect"
+	| "onSessionEvent"
+	| "onPermissionRequest"
+> &
+	SandboxActorConfigCallbacks<TConnParams> &
+	SandboxActorProviderConfig<TConnParams>;
+
+// Input config (what users pass in before Zod transforms).
+export type SandboxActorConfigInput<TConnParams = undefined> = Omit<
+	z.input<typeof SandboxActorConfigSchema>,
+	| "provider"
+	| "createProvider"
+	| "onBeforeConnect"
+	| "onSessionEvent"
+	| "onPermissionRequest"
+> &
+	SandboxActorConfigCallbacks<TConnParams> &
+	SandboxActorProviderConfig<TConnParams>;
+
+export type SandboxActorBeforeConnectContext<TConnParams = undefined> =
+	BeforeConnectContext<
+		SandboxActorState,
+		SandboxActorVars,
+		undefined,
+		AnyDatabaseProvider
+	>;

package/src/sandbox/index.ts

@@ -0,0 +1,41 @@
+export { sandboxActor } from "./actor/index";
+export * from "./client";
+export {
+	type SandboxActorBeforeConnectContext,
+	type SandboxActorConfig,
+	type SandboxActorConfigInput,
+	type SandboxActorOptions,
+	type SandboxActorOptionsRuntime,
+	SandboxActorConfigSchema,
+	SandboxActorOptionsSchema,
+} from "./config";
+export {
+	type SandboxActionContext,
+	type SandboxActorActions,
+	type SandboxActorProvider,
+	type SandboxActorVars,
+	type SandboxActorRuntime,
+	type SandboxActorState,
+	type SandboxSessionEvent,
+	SANDBOX_AGENT_ACTION_METHODS,
+	SANDBOX_AGENT_HOOK_METHODS,
+} from "./types";
+export type {
+	PermissionReply,
+	ProcessLogFollowQuery,
+	ProcessLogListener,
+	ProcessLogSubscription,
+	ProcessTerminalConnectOptions,
+	ProcessTerminalSession,
+	ProcessTerminalSessionOptions,
+	ProcessTerminalWebSocketUrlOptions,
+	SandboxAgent,
+	SandboxProvider,
+	Session,
+	SessionCreateRequest,
+	SessionEvent,
+	SessionPermissionRequest,
+	SessionRecord,
+	SessionResumeOrCreateRequest,
+	SessionSendOptions,
+} from "sandbox-agent";

package/src/sandbox/providers/computesdk.ts

@@ -0,0 +1 @@
+export * from "sandbox-agent/computesdk";

package/src/sandbox/providers/daytona.ts

@@ -0,0 +1 @@
+export * from "sandbox-agent/daytona";

package/src/sandbox/providers/docker.ts

@@ -0,0 +1 @@
+export * from "sandbox-agent/docker";

package/src/sandbox/providers/e2b.ts

@@ -0,0 +1 @@
+export * from "sandbox-agent/e2b";

package/src/sandbox/providers/local.ts

@@ -0,0 +1 @@
+export * from "sandbox-agent/local";

package/src/sandbox/providers/modal.ts

@@ -0,0 +1 @@
+export * from "sandbox-agent/modal";

package/src/sandbox/providers/vercel.ts

@@ -0,0 +1 @@
+export * from "sandbox-agent/vercel";

package/src/sandbox/session-persist-driver.ts

@@ -0,0 +1,180 @@
+import type { RawAccess } from "@/db/config";
+import type {
+	ListEventsRequest,
+	ListPage,
+	ListPageRequest,
+	SessionEvent,
+	SessionPersistDriver,
+	SessionRecord,
+} from "sandbox-agent";
+
+type PersistSessionRow = {
+	record_json: string;
+};
+
+type PersistEventRow = {
+	id: string;
+	event_index: number;
+	session_id: string;
+	created_at: number;
+	connection_id: string;
+	sender: SessionEvent["sender"];
+	payload_json: string;
+};
+
+function parseCursor(cursor?: string): number {
+	if (!cursor) {
+		return 0;
+	}
+
+	const value = Number(cursor);
+	if (!Number.isFinite(value) || value < 0) {
+		return 0;
+	}
+
+	return Math.floor(value);
+}
+
+function nextCursor(
+	offset: number,
+	limit: number,
+	itemCount: number,
+): string | undefined {
+	if (itemCount < limit) {
+		return undefined;
+	}
+
+	return String(offset + itemCount);
+}
+
+export class SqliteSessionPersistDriver implements SessionPersistDriver {
+	#db: RawAccess;
+	#persistRawEvents: boolean;
+
+	constructor(db: RawAccess, persistRawEvents: boolean) {
+		this.#db = db;
+		this.#persistRawEvents = persistRawEvents;
+	}
+
+	async getSession(id: string): Promise<SessionRecord | undefined> {
+		const rows = await this.#db.execute<PersistSessionRow>(
+			"SELECT record_json FROM sandbox_agent_sessions WHERE id = ? LIMIT 1",
+			id,
+		);
+		const row = rows[0];
+		if (!row) {
+			return undefined;
+		}
+		return JSON.parse(row.record_json) as SessionRecord;
+	}
+
+	async listSessions(
+		request: ListPageRequest = {},
+	): Promise<ListPage<SessionRecord>> {
+		const limit = request.limit ?? 50;
+		const offset = parseCursor(request.cursor);
+		const rows = await this.#db.execute<PersistSessionRow>(
+			`
+				SELECT record_json
+				FROM sandbox_agent_sessions
+				ORDER BY created_at DESC, id DESC
+				LIMIT ? OFFSET ?
+			`,
+			limit,
+			offset,
+		);
+
+		return {
+			items: rows.map((row) => JSON.parse(row.record_json) as SessionRecord),
+			nextCursor: nextCursor(offset, limit, rows.length),
+		};
+	}
+
+	async updateSession(session: SessionRecord): Promise<void> {
+		await this.#db.execute(
+			`
+				INSERT INTO sandbox_agent_sessions (id, created_at, record_json)
+				VALUES (?, ?, ?)
+				ON CONFLICT(id) DO UPDATE SET
+					created_at = excluded.created_at,
+					record_json = excluded.record_json
+			`,
+			session.id,
+			session.createdAt,
+			JSON.stringify(session),
+		);
+	}
+
+	async listEvents(
+		request: ListEventsRequest,
+	): Promise<ListPage<SessionEvent>> {
+		const limit = request.limit ?? 200;
+		const offset = parseCursor(request.cursor);
+		const rows = await this.#db.execute<PersistEventRow>(
+			`
+				SELECT
+					id,
+					event_index,
+					session_id,
+					created_at,
+					connection_id,
+					sender,
+					payload_json
+				FROM sandbox_agent_events
+				WHERE session_id = ?
+				ORDER BY event_index ASC
+				LIMIT ? OFFSET ?
+			`,
+			request.sessionId,
+			limit,
+			offset,
+		);
+
+		return {
+			items: rows.map((row) => ({
+				id: row.id,
+				eventIndex: row.event_index,
+				sessionId: row.session_id,
+				createdAt: row.created_at,
+				connectionId: row.connection_id,
+				sender: row.sender,
+				payload: JSON.parse(row.payload_json),
+			} satisfies SessionEvent)),
+			nextCursor: nextCursor(offset, limit, rows.length),
+		};
+	}
+
+	async insertEvent(_sessionId: string, event: SessionEvent): Promise<void> {
+		const payload = JSON.stringify(event.payload);
+		await this.#db.execute(
+			`
+				INSERT INTO sandbox_agent_events (
+					id,
+					session_id,
+					event_index,
+					created_at,
+					connection_id,
+					sender,
+					payload_json,
+					raw_payload_json
+				)
+				VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+				ON CONFLICT(session_id, event_index) DO UPDATE SET
+					id = excluded.id,
+					created_at = excluded.created_at,
+					connection_id = excluded.connection_id,
+					sender = excluded.sender,
+					payload_json = excluded.payload_json,
+					raw_payload_json = excluded.raw_payload_json
+			`,
+			event.id,
+			event.sessionId,
+			event.eventIndex,
+			event.createdAt,
+			event.connectionId,
+			event.sender,
+			payload,
+			this.#persistRawEvents ? payload : null,
+		);
+	}
+}

package/src/sandbox/types.ts

@@ -0,0 +1,138 @@
+import type { ActionContext } from "@/actor/contexts";
+import type { DatabaseProvider } from "@/actor/database";
+import type { RawAccess } from "@/db/config";
+import type {
+	SandboxAgent,
+	SandboxProvider,
+} from "sandbox-agent";
+
+export type { SandboxProvider };
+
+/** @deprecated Use `SandboxProvider` from `sandbox-agent` instead. */
+export type SandboxActorProvider = SandboxProvider;
+
+// Keep this split in lockstep with the sandbox-agent SDK. Hooks should match the
+// SDK callback methods, and actions should match every other SDK instance
+// method. Update these lists and the parity test together when sandbox-agent
+// changes.
+export const SANDBOX_AGENT_HOOK_METHODS = [
+	"onSessionEvent",
+	"onPermissionRequest",
+] as const;
+
+export type SandboxActorHookMethodName =
+	(typeof SANDBOX_AGENT_HOOK_METHODS)[number];
+
+export const SANDBOX_AGENT_ACTION_METHODS = [
+	"dispose",
+	"listSessions",
+	"getSession",
+	"getEvents",
+	"createSession",
+	"resumeSession",
+	"resumeOrCreateSession",
+	"destroySandbox",
+	"destroySession",
+	"setSessionMode",
+	"setSessionConfigOption",
+	"setSessionModel",
+	"setSessionThoughtLevel",
+	"getSessionConfigOptions",
+	"getSessionModes",
+	"rawSendSessionMethod",
+	"respondPermission",
+	"rawRespondPermission",
+	"getHealth",
+	"listAgents",
+	"getAgent",
+	"installAgent",
+	"listAcpServers",
+	"listFsEntries",
+	"readFsFile",
+	"writeFsFile",
+	"deleteFsEntry",
+	"mkdirFs",
+	"moveFs",
+	"statFs",
+	"uploadFsBatch",
+	"getMcpConfig",
+	"setMcpConfig",
+	"deleteMcpConfig",
+	"getSkillsConfig",
+	"setSkillsConfig",
+	"deleteSkillsConfig",
+	"getProcessConfig",
+	"setProcessConfig",
+	"createProcess",
+	"runProcess",
+	"listProcesses",
+	"getProcess",
+	"stopProcess",
+	"killProcess",
+	"deleteProcess",
+	"getProcessLogs",
+	"followProcessLogs",
+	"sendProcessInput",
+	"resizeProcessTerminal",
+	"buildProcessTerminalWebSocketUrl",
+	"connectProcessTerminalWebSocket",
+	"connectProcessTerminal",
+] as const;
+
+export type SandboxAgentActionMethodName =
+	(typeof SANDBOX_AGENT_ACTION_METHODS)[number];
+
+export type SandboxActorActions = Pick<
+	SandboxAgent,
+	SandboxAgentActionMethodName
+>;
+
+export type SandboxSessionEvent = Parameters<
+	Parameters<SandboxAgent["onSessionEvent"]>[1]
+>[0];
+
+export interface SandboxActorState {
+	sandboxId: string | null;
+	providerName: string | null;
+	/** Persisted so that on wake, the actor knows which sessions to
+	 * re-subscribe to when reconnecting to the sandbox agent. Without
+	 * this, event listeners would be lost after a sleep/wake cycle. */
+	subscribedSessionIds: string[];
+	sandboxDestroyed: boolean;
+}
+
+export interface SandboxActorVars {
+	sandboxAgentClient: SandboxAgent | null;
+	provider: SandboxProvider | null;
+	activeSessionIds: Set<string>;
+	activePromptRequestIdsBySessionId: Map<string, string[]>;
+	lastEventAtBySessionId: Map<string, number>;
+	unsubscribeBySessionId: Map<
+		string,
+		{
+			event?: () => void;
+			permission?: () => void;
+		}
+	>;
+	/** Tracks in-flight hook promises. Size is used instead of a counter
+	 * to avoid increment/decrement mismatch bugs. */
+	activeHooks: Set<Promise<void>>;
+	warningTimeoutBySessionId: Map<string, ReturnType<typeof setTimeout>>;
+	staleTimeoutBySessionId: Map<string, ReturnType<typeof setTimeout>>;
+}
+
+/** @deprecated Use `SandboxActorVars` instead. */
+export type SandboxActorRuntime = SandboxActorVars;
+
+/**
+ * Action context type used by the sandbox actor implementation for session
+ * management, proxy actions, and lifecycle hooks.
+ */
+export type SandboxActionContext<TConnParams = undefined> = ActionContext<
+	SandboxActorState,
+	TConnParams,
+	undefined,
+	SandboxActorVars,
+	undefined,
+	DatabaseProvider<RawAccess>
+>;

package/src/serverless/configure.ts

@@ -49,15 +49,17 @@ export async function configureServerlessRunner(
 			serverless: {
 				url: customConfig.url,
 				headers: customConfig.headers ?? {},
-				max_runners: customConfig.maxRunners ?? 1000,
+				max_runners: customConfig.maxRunners ?? 100_000,
 				min_runners: customConfig.minRunners ?? 0,
 				request_lifespan: customConfig.requestLifespan ?? 15 * 60,
 				runners_margin: customConfig.runnersMargin ?? 0,
 				slots_per_runner: customConfig.slotsPerRunner ?? 1,
-				metadata_poll_interval: customConfig.metadataPollInterval ?? 1000,
+				metadata_poll_interval:
+					customConfig.metadataPollInterval ?? 1000,
 			},
 			metadata: customConfig.metadata ?? {},
-			drain_on_version_upgrade: customConfig.drainOnVersionUpgrade ?? true,
+			drain_on_version_upgrade:
+				customConfig.drainOnVersionUpgrade ?? true,
 			metadataPollInterval: customConfig.metadataPollInterval ?? 1000,
 		};
 		await updateRunnerConfig(clientConfig, runnerName, {

package/src/serverless/router.test.ts

@@ -78,9 +78,9 @@ describe("normalizeEndpointUrl", () => {
 
 	describe("regional endpoint normalization", () => {
 		test("normalizes api-us-west-1.rivet.dev to api.rivet.dev", () => {
-			expect(normalizeEndpointUrl("https://api-us-west-1.rivet.dev")).toBe(
-				"https://api.rivet.dev/",
-			);
+			expect(
+				normalizeEndpointUrl("https://api-us-west-1.rivet.dev"),
+			).toBe("https://api.rivet.dev/");
 		});
 
 		test("normalizes api-lax.staging.rivet.dev to api.staging.rivet.dev", () => {
@@ -102,14 +102,16 @@ describe("normalizeEndpointUrl", () => {
 		});
 
 		test("does not normalize non-rivet.dev domains", () => {
-			expect(normalizeEndpointUrl("https://api-us-west-1.example.com")).toBe(
-				"https://api-us-west-1.example.com/",
-			);
+			expect(
+				normalizeEndpointUrl("https://api-us-west-1.example.com"),
+			).toBe("https://api-us-west-1.example.com/");
 		});
 
 		test("preserves path when normalizing regional endpoint", () => {
 			expect(
-				normalizeEndpointUrl("https://api-us-west-1.rivet.dev/v1/actors"),
+				normalizeEndpointUrl(
+					"https://api-us-west-1.rivet.dev/v1/actors",
+				),
 			).toBe("https://api.rivet.dev/v1/actors");
 		});
 
@@ -136,7 +138,10 @@ describe("endpointsMatch", () => {
 
 	test("matches URLs with paths ignoring trailing slash", () => {
 		expect(
-			endpointsMatch("http://example.com/api/v1", "http://example.com/api/v1/"),
+			endpointsMatch(
+				"http://example.com/api/v1",
+				"http://example.com/api/v1/",
+			),
 		).toBe(true);
 	});
 
@@ -178,7 +183,10 @@ describe("endpointsMatch", () => {
 
 	test("does not match different paths", () => {
 		expect(
-			endpointsMatch("http://example.com/api/v1", "http://example.com/api/v2"),
+			endpointsMatch(
+				"http://example.com/api/v1",
+				"http://example.com/api/v2",
+			),
 		).toBe(false);
 	});
 

package/src/serverless/router.ts

@@ -38,7 +38,7 @@ export function buildServerlessRouter(
 			if (!parseResult.success) {
 				throw new InvalidRequest(
 					parseResult.error.issues[0]?.message ??
-						"invalid serverless start headers",
+					"invalid serverless start headers",
 				);
 			}
 			const { endpoint, token, totalSlots, runnerName, namespace } =
@@ -68,34 +68,41 @@ export function buildServerlessRouter(
 				}
 			}
 
-			// Convert config to runner config
-			const newConfig: RegistryConfig = {
+			const sharedConfig: RegistryConfig = {
 				...config,
-				endpoint: endpoint,
-				namespace: namespace,
-				token: token,
+				endpoint,
+				namespace,
 				runner: {
 					...config.runner,
-					totalSlots: totalSlots,
-					runnerName: runnerName,
-					// Not supported on serverless
-					runnerKey: undefined,
+					totalSlots,
+					runnerName,
 				},
 			};
+			const runnerConfig: RegistryConfig = {
+				...sharedConfig,
+				token,
+			};
+			const clientConfig: RegistryConfig = {
+				...sharedConfig,
+				// Preserve the configured application token for actor-to-actor
+				// calls. The start token is only needed for the runner
+				// connection and may not have gateway permissions.
+				token: config.token ?? token,
+			};
 
 			// Create manager driver on demand based on the properties provided
 			// by headers
 			//
-			// NOTE: This relies on the `newConfig.runner.runnerName` to
+			// NOTE: This relies on the `runnerConfig.runner.runnerName` to
 			// configure which runner to create actors on.
 			const managerDriver = new RemoteManagerDriver(
-				convertRegistryConfigToClientConfig(newConfig),
+				convertRegistryConfigToClientConfig(clientConfig),
 			);
 			const client = createClientWithDriver(managerDriver);
 
 			// Create new actor driver with updated config
 			const actorDriver = driverConfig.actor(
-				newConfig,
+				runnerConfig,
 				managerDriver,
 				client,
 			);

package/src/test/mod.ts

@@ -23,10 +23,9 @@ export async function setupTest<A extends Registry<any>>(
 	registry.config.test = { ...registry.config.test, enabled: true };
 
 	// Create driver
-	const driver = createFileSystemOrMemoryDriver(
-		true,
-		{ path: `/tmp/rivetkit-test-${crypto.randomUUID()}` },
-	);
+	const driver = createFileSystemOrMemoryDriver(true, {
+		path: `/tmp/rivetkit-test-${crypto.randomUUID()}`,
+	});
 
 	// Build driver config
 	// biome-ignore lint/style/useConst: Assigned later