Skip to content

Commit 73b65a1

Browse files
eersningtonNathanFlurry
eersnington
authored and
NathanFlurry
committed
feat(pi): route pi built-in tools through a sandbox
1 parent a1ea812 commit 73b65a1

6 files changed

Lines changed: 392 additions & 28 deletions

File tree

‎integrations/pi/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
},
2727
"dependencies": {
2828
"@earendil-works/pi-coding-agent": "0.87.1",
29+
"@rivet-dev/sandbox-adapter": "workspace:*",
2930
"rivetkit": "workspace:*"
3031
},
3132
"devDependencies": {

‎integrations/pi/src/actions.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -156,12 +156,15 @@ export function createPiActions(options: PiSessionOptions): PiActions {
156156
setActiveToolsByName: (c, ...args) =>
157157
mutate(c, ({ session }) => session.setActiveToolsByName(...args)),
158158
executeBash: (c, command, bashOptions) =>
159-
mutate(c, ({ session, bashOperations }) =>
160-
session.executeBash(command, undefined, {
159+
mutate(c, ({ session, bashOperations }) => {
160+
if (!bashOperations) {
161+
throw new Error("executeBash needs a sandbox; this pi() actor has none configured");
162+
}
163+
return session.executeBash(command, undefined, {
161164
...bashOptions,
162165
operations: bashOperations,
163-
}),
164-
),
166+
});
167+
}),
165168
abortBash: (c) => read(c, ({ session }) => session.abortBash()),
166169

167170
setSteeringMode: (c, ...args) =>

‎integrations/pi/src/actor.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ const piOptionKeys = [
4747
"resourceLoader",
4848
"sessionStartEvent",
4949
"settings",
50+
"sandbox",
5051
] as const satisfies readonly (keyof PiSessionOptions)[];
5152

5253
export type PiActorConfigInput<
@@ -186,14 +187,14 @@ export function pi<
186187
try {
187188
await userOnSleep?.(c);
188189
} finally {
189-
await closePiSession(c);
190+
await closePiSession(c, sessionOptions, "sleep");
190191
}
191192
},
192193
onDestroy: async (c: PiContext) => {
193194
try {
194195
await userOnDestroy?.(c);
195196
} finally {
196-
await closePiSession(c);
197+
await closePiSession(c, sessionOptions, "destroy");
197198
}
198199
},
199200
} as any) as ActorDefinition<

‎integrations/pi/src/runtime.ts‎

Lines changed: 161 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,25 @@ import {
55
type BashOperations,
66
createAgentSession,
77
type CreateAgentSessionOptions,
8+
DefaultResourceLoader,
9+
getAgentDir,
810
ModelRuntime,
911
SessionManager,
1012
SettingsManager,
1113
} from "@earendil-works/pi-coding-agent";
14+
import type { Sandbox, SandboxProvider } from "@rivet-dev/sandbox-adapter";
1215
import type { ActorContext } from "rivetkit";
1316
import type { DatabaseProvider, RawAccess } from "rivetkit/db";
17+
import { createSandboxBashOperations, createSandboxTools } from "./sandbox.js";
1418
import {
1519
appendPiEntry,
1620
createPiSession,
21+
loadPiSandbox,
1722
loadPiSession,
1823
type PiSettings,
24+
savePiSandbox,
1925
savePiSettings,
26+
type StoredSandbox,
2027
toFileEntries,
2128
} from "./storage.js";
2229

@@ -39,21 +46,34 @@ export interface PiSessionOptions
3946
extends Omit<CreateAgentSessionOptions, "sessionManager" | "settingsManager"> {
4047
/** Initial Pi settings for a brand-new session. Later changes persist per actor. */
4148
settings?: Partial<PiSettings>;
49+
/**
50+
* Runs Pi's built-in file and shell tools in a sandbox. Without one, Pi has
51+
* no file or shell tools and only the `customTools` passed in.
52+
*/
53+
sandbox?: SandboxProvider;
4254
}
4355

4456
/** One open Pi session for a live actor generation. */
4557
export interface PiSession {
4658
session: AgentSession;
4759
settingsManager: SettingsManager;
4860
cwd: string;
49-
/** Command execution for `executeBash`. Undefined runs on the actor host. */
61+
sandbox: ConnectedSandbox | undefined;
62+
/** Command execution for `executeBash`. Undefined when there is no sandbox. */
5063
bashOperations: BashOperations | undefined;
5164
/** JSON of the settings last written to SQLite, to skip no-op writes. */
5265
persistedSettings: string;
5366
/** Entries (header excluded) already written to SQLite, in Pi's append order. */
5467
persistedEntryCount: number;
5568
}
5669

70+
/** The sandbox a session's tools run in for this actor generation. */
71+
export interface ConnectedSandbox {
72+
provider: SandboxProvider;
73+
id: string;
74+
sandbox: Sandbox;
75+
}
76+
5777
/** Per-actor-generation runtime state, stored on `c.vars` under `PI_RUNTIME`. */
5878
export interface PiRuntime {
5979
ready?: Promise<PiSession>;
@@ -94,6 +114,9 @@ export function ensurePiSession(
94114
return runtime.ready;
95115
}
96116

117+
/** Pi's built-in tools, which all run on the actor host. */
118+
const PI_BUILT_IN_TOOLS = ["read", "bash", "powershell", "edit", "write", "grep", "find", "ls"];
119+
97120
let defaultModelRuntime: Promise<ModelRuntime> | undefined;
98121

99122
/** One `ModelRuntime` per process. It loads provider catalogs and credentials, which are not per actor. */
@@ -107,12 +130,24 @@ async function openPiSession(
107130
runtime: PiRuntime,
108131
options: PiSessionOptions,
109132
): Promise<PiSession> {
110-
const { settings, ...sessionOptions } = options;
133+
const { settings, sandbox: sandboxProvider, ...sessionOptions } =
134+
options;
111135
const stored = await loadPiSession(c.db);
112-
const cwd = stored?.cwd ?? sessionOptions.cwd ?? process.cwd();
136+
const connected = sandboxProvider
137+
? await connectSandbox(c, sandboxProvider)
138+
: undefined;
139+
const sandbox = connected?.sandbox;
140+
const cwd = sandbox?.cwd ?? stored?.cwd ?? sessionOptions.cwd ?? process.cwd();
113141
if (!isAbsolute(cwd)) {
114142
throw new Error(`pi() cwd must be an absolute path, received ${cwd}`);
115143
}
144+
if (stored && stored.cwd !== cwd) {
145+
c.log.warn({
146+
msg: "pi session cwd changed since it was stored",
147+
storedCwd: stored.cwd,
148+
cwd,
149+
});
150+
}
116151

117152
const settingsManager = SettingsManager.inMemory(
118153
stored?.settings ?? settings ?? {},
@@ -121,13 +156,26 @@ async function openPiSession(
121156
? SessionManager.inMemory(cwd, undefined, toFileEntries(stored))
122157
: SessionManager.inMemory(cwd);
123158
const modelRuntime = sessionOptions.modelRuntime ?? (await sharedModelRuntime());
159+
const resourceLoader =
160+
sessionOptions.resourceLoader ??
161+
(await isolatedResourceLoader(cwd, sessionOptions.agentDir, settingsManager));
124162

125163
const { session, modelFallbackMessage } = await createAgentSession({
126164
...sessionOptions,
127165
cwd,
128166
modelRuntime,
129167
settingsManager,
130168
sessionManager,
169+
resourceLoader,
170+
customTools: sandbox
171+
? [...(sessionOptions.customTools ?? []), ...createSandboxTools(sandbox)]
172+
: sessionOptions.customTools,
173+
excludeTools: [
174+
...new Set([
175+
...(sessionOptions.excludeTools ?? []),
176+
...(sandbox ? ["powershell"] : PI_BUILT_IN_TOOLS),
177+
]),
178+
],
131179
});
132180
if (modelFallbackMessage) {
133181
c.log.warn({ msg: "pi model fallback", detail: modelFallbackMessage });
@@ -137,7 +185,8 @@ async function openPiSession(
137185
session,
138186
settingsManager,
139187
cwd,
140-
bashOperations: undefined,
188+
sandbox: connected,
189+
bashOperations: sandbox ? createSandboxBashOperations(sandbox) : undefined,
141190
persistedSettings: JSON.stringify(settingsManager.getGlobalSettings()),
142191
persistedEntryCount: stored?.entries.length ?? 0,
143192
};
@@ -176,6 +225,65 @@ async function openPiSession(
176225
return handle;
177226
}
178227

228+
/**
229+
* Connects to the actor's sandbox, creating one when none is stored or the
230+
* provider reports the stored one no longer exists. A new sandbox id is saved
231+
* as soon as `create` returns, so a failure later in the start reuses it. Any
232+
* other connect failure is thrown, so a temporary outage never replaces a
233+
* sandbox.
234+
*/
235+
async function connectSandbox(
236+
c: PiContext,
237+
provider: SandboxProvider,
238+
): Promise<ConnectedSandbox> {
239+
const existing = await loadPiSandbox(c.db);
240+
if (existing && existing.provider !== provider.name) {
241+
throw new Error(
242+
`pi sandbox was created by provider ${existing.provider}, but the actor now uses ${provider.name}`,
243+
);
244+
}
245+
if (existing) {
246+
const sandbox = await provider.connect(c, existing.id);
247+
if (sandbox) return { provider, id: existing.id, sandbox };
248+
c.log.warn({
249+
msg: "pi sandbox no longer exists, creating a new one; files from the previous sandbox are lost",
250+
provider: provider.name,
251+
sandboxId: existing.id,
252+
});
253+
}
254+
const id = await provider.create(c);
255+
await savePiSandbox(c.db, { provider: provider.name, id });
256+
const sandbox = await provider.connect(c, id);
257+
if (!sandbox) {
258+
throw new Error(`pi sandbox ${provider.name}/${id} was not found right after it was created`);
259+
}
260+
return { provider, id, sandbox };
261+
}
262+
263+
/**
264+
* Pi's resource discovery reads the actor host's filesystem and loads host
265+
* code as extensions. Extensions, skills, prompt templates, context files, and
266+
* themes stay off unless the developer passes a loader.
267+
*/
268+
async function isolatedResourceLoader(
269+
cwd: string,
270+
agentDir: string | undefined,
271+
settingsManager: SettingsManager,
272+
): Promise<DefaultResourceLoader> {
273+
const loader = new DefaultResourceLoader({
274+
cwd,
275+
agentDir: agentDir ?? getAgentDir(),
276+
settingsManager,
277+
noExtensions: true,
278+
noSkills: true,
279+
noPromptTemplates: true,
280+
noContextFiles: true,
281+
noThemes: true,
282+
});
283+
await loader.reload();
284+
return loader;
285+
}
286+
179287
/** Token and output deltas never append entries, so they skip the entry diff. */
180288
function isStreamingDelta(event: AgentSessionEvent): boolean {
181289
return (
@@ -239,40 +347,71 @@ export async function persistPiState(
239347

240348
/**
241349
* Stops the Pi session for this actor generation: aborts any run, lets
242-
* extensions shut down, flushes settings, and waits for queued entry writes.
350+
* extensions shut down, and flushes settings and entries. Then suspends the
351+
* sandbox on sleep, or destroys it on destroy.
243352
*/
244-
export async function closePiSession(c: PiContext): Promise<void> {
353+
export async function closePiSession(
354+
c: PiContext,
355+
options: PiSessionOptions,
356+
reason: "sleep" | "destroy",
357+
): Promise<void> {
245358
const runtime = piRuntime(c);
246359
const ready = runtime.ready;
247360
runtime.ready = undefined;
248-
if (!ready) return;
249-
let handle: PiSession;
361+
const errors: unknown[] = [];
362+
let handle: PiSession | undefined;
250363
try {
251364
handle = await ready;
252365
} catch {
253-
return;
254366
}
255367

256-
const errors: unknown[] = [];
257-
await attempt(errors, () => handle.session.abort());
258-
await attempt(errors, async () => {
259-
if (handle.session.hasExtensionHandlers("session_shutdown")) {
260-
await handle.session.extensionRunner.emit({
261-
type: "session_shutdown",
262-
reason: "quit",
263-
});
264-
}
265-
});
266-
await attempt(errors, () => persistPiState(c, handle));
267-
handle.session.dispose();
368+
if (handle) {
369+
const open = handle;
370+
await attempt(errors, () => open.session.abort());
371+
await attempt(errors, async () => {
372+
if (open.session.hasExtensionHandlers("session_shutdown")) {
373+
await open.session.extensionRunner.emit({
374+
type: "session_shutdown",
375+
reason: "quit",
376+
});
377+
}
378+
});
379+
await attempt(errors, () => persistPiState(c, open));
380+
open.session.dispose();
381+
c.log.info({ msg: "pi session closed", sessionId: open.session.sessionId });
382+
}
383+
384+
const provider = options.sandbox;
385+
if (provider) {
386+
await attempt(errors, async () => {
387+
if (reason === "sleep") {
388+
if (handle?.sandbox && provider.suspend) {
389+
await provider.suspend(c, handle.sandbox.id);
390+
}
391+
return;
392+
}
393+
const sandbox = handle?.sandbox ?? (await storedSandbox(c, provider));
394+
if (sandbox && provider.destroy) {
395+
await provider.destroy(c, sandbox.id);
396+
}
397+
});
398+
}
268399

269-
c.log.info({ msg: "pi session closed", sessionId: handle.session.sessionId });
270400
if (errors.length === 1) throw errors[0];
271401
if (errors.length > 1) {
272402
throw new AggregateError(errors, "pi session shutdown failed");
273403
}
274404
}
275405

406+
/** The stored sandbox, when it belongs to `provider`. */
407+
async function storedSandbox(
408+
c: PiContext,
409+
provider: SandboxProvider,
410+
): Promise<StoredSandbox | undefined> {
411+
const sandbox = await loadPiSandbox(c.db);
412+
return sandbox?.provider === provider.name ? sandbox : undefined;
413+
}
414+
276415
async function attempt(
277416
errors: unknown[],
278417
operation: () => void | Promise<void>,

0 commit comments

Comments
 (0)