refactor(pi): move databaseEnv into pi and version pi_file

packages/sandbox-adapter/src/pi/database-env.ts → packages/pi/src/database-env.ts

@@ -12,50 +12,22 @@ import {
 	type Result,
 	toError,
 } from "@earendil-works/pi-durable/env";
-
-/** The SQLite access `databaseEnv` needs, such as a Rivet Actor's `c.db`. */
-export interface FileDatabase {
-	execute<TRow extends Record<string, unknown> = Record<string, unknown>>(
-		sql: string,
-		...args: unknown[]
-	): Promise<TRow[]>;
-}
+import type { PiDatabase } from "./storage.js";
 
 const DEFAULT_CWD = "/workspace";
 
 /** Bytes per query when a scan walks a whole file. */
 const SCAN_CHUNK_BYTES = 64 * 1024;
 
-const created = new WeakMap<FileDatabase, Promise<void>>();
-
-/** Creates the `pi_file` table once per database object. */
-function ensureTable(db: FileDatabase): Promise<void> {
-	let ready = created.get(db);
-	if (!ready) {
-		ready = db
-			.execute(
-				`CREATE TABLE IF NOT EXISTS pi_file (
-					path TEXT PRIMARY KEY,
-					content BLOB NOT NULL,
-					mtime_ms INTEGER NOT NULL
-				) STRICT`,
-			)
-			.then(() => {});
-		ready.catch(() => created.delete(db));
-		created.set(db, ready);
-	}
-	return ready;
-}
-
 /**
- * Pi Durable's `ExecutionEnv` over files stored in a SQLite database, in the
- * `pi_file` table, for Pi's `read`, `write`, and `edit` tools. Directories
+ * Pi Durable's `ExecutionEnv` over files stored in the actor's database, in
+ * the `pi_file` table, for Pi's `read`, `write`, and `edit` tools. Directories
  * exist only as the paths of their files. There is no shell, so `exec` fails
  * with `shell_unavailable`, and operations the tools do not use return
  * `not_supported`.
  */
 export function databaseEnv(
-	db: FileDatabase,
+	db: PiDatabase,
 	options: { id: string; cwd?: string },
 ): ExecutionEnv {
 	const env: ExecutionEnv = {
@@ -142,7 +114,6 @@ export function databaseEnv(
 		const resolved = posix.resolve(env.cwd, path);
 		if (context.abortSignal?.aborted) return err(aborted(resolved));
 		try {
-			await ensureTable(db);
 			return ok(await operation(resolved));
 		} catch (error) {
 			if (error instanceof FileError) return err(error);
@@ -156,7 +127,7 @@ export function databaseEnv(
 
 /** The file or directory at `path`. A directory exists when a file below it does. */
 async function entry(
-	db: FileDatabase,
+	db: PiDatabase,
 	path: string,
 ): Promise<FileInfo | undefined> {
 	const [file] = await db.execute<{ size: number; mtime_ms: number }>(
@@ -189,7 +160,7 @@ async function entry(
 	};
 }
 
-async function requireFile(db: FileDatabase, path: string): Promise<FileInfo> {
+async function requireFile(db: PiDatabase, path: string): Promise<FileInfo> {
 	const info = await entry(db, path);
 	if (!info) throw notFound(path);
 	if (info.kind === "directory") {
@@ -198,7 +169,7 @@ async function requireFile(db: FileDatabase, path: string): Promise<FileInfo> {
 	return info;
 }
 
-async function readFile(db: FileDatabase, path: string): Promise<Uint8Array> {
+async function readFile(db: PiDatabase, path: string): Promise<Uint8Array> {
 	const [row] = await db.execute<{ content: Uint8Array }>(
 		`SELECT content FROM pi_file WHERE path = ?`,
 		path,
@@ -209,7 +180,7 @@ async function readFile(db: FileDatabase, path: string): Promise<Uint8Array> {
 }
 
 /** Reads byte ranges of the stored file with SQL, so a large file is never loaded whole. */
-function fileReader(db: FileDatabase, path: string): BinaryReader {
+function fileReader(db: PiDatabase, path: string): BinaryReader {
 	const range = async (offset: number, length: number) => {
 		const [row] = await db.execute<{ chunk: Uint8Array | null }>(
 			`SELECT substr(content, ?, ?) AS chunk FROM pi_file WHERE path = ?`,

packages/pi/src/runtime.ts

@@ -14,11 +14,12 @@ import {
 import type { ExecutionEnv } from "@earendil-works/pi-durable/env";
 import { SqliteStorage } from "@earendil-works/pi-durable/storage/sqlite";
 import type { Sandbox, SandboxProvider } from "@rivet-dev/sandbox-adapter";
-import { databaseEnv, sandboxEnv } from "@rivet-dev/sandbox-adapter/pi";
+import { sandboxEnv } from "@rivet-dev/sandbox-adapter/pi";
 import { type ActorContext, UserError } from "rivetkit";
 import type { DatabaseProvider, RawAccess } from "rivetkit/db";
 import { actorSqlite } from "./actor-sqlite.js";
 import { SourceCredentialStore } from "./credentials.js";
+import { databaseEnv } from "./database-env.js";
 import { toClientError } from "./errors.js";
 import {
 	createActorModelRuntime,

packages/pi/src/storage.ts

@@ -33,6 +33,17 @@ export const migratePiTables = migrations({
 				) STRICT, WITHOUT ROWID;
 			`,
 		},
+		{
+			version: 3,
+			// Files of agents without a sandbox. Agents from 0.5.1 created this table on first use.
+			sql: `
+				CREATE TABLE IF NOT EXISTS pi_file (
+					path TEXT PRIMARY KEY,
+					content BLOB NOT NULL,
+					mtime_ms INTEGER NOT NULL
+				) STRICT;
+			`,
+		},
 	],
 });
 

packages/pi/tests/actor.test.ts

@@ -181,6 +181,12 @@ function buildRegistry(mock: MockModel, root: string) {
 			nap: (c) => {
 				c.sleep();
 			},
+			// What 0.5.1 stored: `pi_file` with its rows, at schema version 2.
+			storeAs051: async (c) => {
+				await c.db.execute(
+					"UPDATE pi_durable_schema_version SET schema_version = 2",
+				);
+			},
 		},
 	});
 	const backoff = pi({
@@ -605,6 +611,26 @@ describe("pi actor", () => {
 		});
 	});
 
+	test("files an agent without a sandbox stored on 0.5.1 stay readable after the upgrade", async (c) => {
+		const { client } = await setupTest(c, registry);
+		const key = ["files-051", randomUUID()];
+		const handle = client.files.getOrCreate(key);
+		await handle.prompt("create hello.txt");
+		await handle.storeAs051();
+		await handle.nap();
+		await sleeps.waitFor(key, 1);
+
+		await handle.prompt("read line two of hello.txt");
+
+		const root = await handle.harness.root();
+		const { messages } = await handle.conversation.context(root.id);
+		const read = messages.findLast(
+			(message) => message.role === "toolResult" && message.toolName === "read",
+		);
+		expect(read).toMatchObject({ isError: false });
+		expect(JSON.stringify(read?.content)).toContain("line two");
+	});
+
 	test("sandbox tools cannot write outside the sandbox or read the actor host's environment", async (c) => {
 		const { client } = await setupTest(c, registry);
 		const handle = client.coder.getOrCreate(["escape", randomUUID()]);

packages/sandbox-adapter/src/pi/index.ts

@@ -1,2 +1 @@
-export { databaseEnv, type FileDatabase } from "./database-env.js";
 export { sandboxEnv } from "./sandbox-env.js";