feat: return text from session

Author: NathanFlurry

CLAUDE.md

@@ -31,7 +31,7 @@ agentOS wraps the kernel and adds: a high-level filesystem/process API, ACP agen
 - **Registry types**: `@rivet-dev/agent-os-registry-types` in `packages/registry-types/` -- shared type definitions for WASM command package descriptors. The registry software packages link to this package. When changing descriptor types, update here and rebuild the registry.
 - **npm scope**: `@rivet-dev/agent-os-*`
 - **Actor integration** lives in the Rivet repo at `rivetkit-typescript/packages/rivetkit/src/agent-os/`, not as a separate package
-- **The actor layer must maintain 1:1 feature parity with AgentOs.** Every public method on the `AgentOs` class (`packages/core/src/agent-os.ts`) must have a corresponding actor action in the Rivet repo's `rivetkit-typescript/packages/rivetkit/src/agent-os/`. Subscription methods (onProcessStdout, onShellData, onCronEvent, etc.) are wired through actor events. Lifecycle methods (dispose) are handled by the actor's onSleep/onDestroy hooks. When adding a new public method to AgentOs, add the corresponding actor action in the same change.
+- **The actor layer must maintain 1:1 feature parity with AgentOs.** Every public method on the `AgentOs` class (`packages/core/src/agent-os.ts`) must have a corresponding actor action in the Rivet repo's `rivetkit-typescript/packages/rivetkit/src/agent-os/`. Subscription methods (onProcessStdout, onShellData, onCronEvent, etc.) are wired through actor events. Lifecycle methods (dispose) are handled by the actor's onSleep/onDestroy hooks. When adding a new public method to AgentOs, add the corresponding actor action in the same change. This includes changes to method signatures, option types, return types, and configuration interfaces -- any API surface change in AgentOs must be mirrored in the actor layer. **Always ask the user which Rivet repo/path to update** (e.g., `~/r-aos`, `~/r16`, etc.) before making changes there.
 - **The RivetKit driver test suite must have full feature coverage of all agent-os actor actions.** Tests live in the Rivet repo's `rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/`. When adding a new actor action, add a corresponding driver test in the same change.
 - **The core quickstart (`examples/quickstart/`) and the RivetKit example (in the Rivet repo at `examples/agent-os/`) must stay in sync.** Both cover the same set of features (hello-world, filesystem, processes, network, cron, tools, agent-session, sandbox) with identical behavior, just different APIs. Core uses `AgentOs.create()` directly; RivetKit uses `agentOs()` actor with client-server split. When adding or changing a quickstart example, update both.
 

examples/quickstart/package.json

@@ -18,7 +18,8 @@
 		"nodejs": "node --import tsx src/nodejs.ts",
 		"python": "node --import tsx src/python.ts",
 		"bash": "node --import tsx src/bash.ts",
-		"s3-filesystem": "node --import tsx src/s3-filesystem.ts"
+		"s3-filesystem": "node --import tsx src/s3-filesystem.ts",
+		"pi-extensions": "node --import tsx src/pi-extensions.ts"
 	},
 	"dependencies": {
 		"@rivet-dev/agent-os-core": "workspace:*",

examples/quickstart/src/pi-extensions.ts

@@ -0,0 +1,75 @@
+// Pi extensions: write a custom extension into the VM before creating a
+// session and verify Pi discovers and loads it.
+//
+// The adapter scans ~/.pi/agent/extensions/ and <cwd>/.pi/extensions/ for
+// .js files at session start. Each file exports a function that receives
+// Pi's ExtensionAPI, which can register tools, modify the system prompt,
+// subscribe to lifecycle events, and more.
+//
+// Extensions must use CommonJS syntax (module.exports = function(pi) { ... }).
+//
+// NOTE: Requires ANTHROPIC_API_KEY to be set.
+
+import { AgentOs } from "@rivet-dev/agent-os-core";
+import common from "@rivet-dev/agent-os-common";
+import pi from "@rivet-dev/agent-os-pi";
+
+const ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY;
+if (!ANTHROPIC_API_KEY) {
+	console.error("Set ANTHROPIC_API_KEY to run this example.");
+	process.exit(1);
+}
+
+// ── Extension source code ──────────────────────────────────────────
+//
+// This extension hooks Pi's before_agent_start event to append a custom
+// instruction to the system prompt. No imports needed — the ExtensionAPI
+// is passed as a parameter.
+
+const extensionSource = `
+module.exports = function(pi) {
+  pi.on("before_agent_start", async (event) => {
+    return {
+      systemPrompt: event.systemPrompt +
+        "\\n\\nCRITICAL INSTRUCTION: You MUST begin every response with " +
+        "exactly the phrase 'EXTENSION_OK: ' followed by your answer. " +
+        "This is mandatory and non-negotiable."
+    };
+  });
+};
+`;
+
+// ── Create VM and write extension ──────────────────────────────────
+
+const vm = await AgentOs.create({ software: [common, pi] });
+
+// Write the extension into Pi's global extensions directory.
+// In the VM, HOME is /home/user, so ~/.pi/agent/extensions/ resolves there.
+const extensionsDir = "/home/user/.pi/agent/extensions";
+await vm.mkdir(extensionsDir, { recursive: true });
+await vm.writeFile(`${extensionsDir}/custom-greeting.js`, extensionSource);
+
+console.log("Extension written. Creating Pi session...\n");
+
+// ── Create session and prompt ──────────────────────────────────────
+
+const { sessionId } = await vm.createSession("pi", {
+	env: { ANTHROPIC_API_KEY },
+});
+console.log("Session created:", sessionId);
+
+// Ask a simple question — if the extension loaded, the agent will
+// prefix its response with "EXTENSION_OK: "
+const { text } = await vm.prompt(sessionId, "What is 2 + 2? Reply with just the number.");
+console.log("Agent:", text);
+
+// ── Verify ─────────────────────────────────────────────────────────
+
+if (text.includes("EXTENSION_OK:")) {
+	console.log("SUCCESS — Pi extension loaded and modified the system prompt.");
+} else {
+	console.log("FAIL — Response did not include the expected prefix.");
+}
+
+vm.closeSession(sessionId);
+await vm.dispose();

packages/core/src/agent-os.ts

@@ -273,6 +273,14 @@ export interface SessionInfo {
 	agentType: string;
 }
 
+/** Result from AgentOs.prompt(). */
+export interface PromptResult {
+	/** Raw JSON-RPC response from the ACP adapter. */
+	response: JsonRpcResponse;
+	/** Accumulated agent text output from streamed message chunks. */
+	text: string;
+}
+
 /** Information about a process spawned via AgentOs.spawn(). */
 export interface SpawnedProcessInfo {
 	pid: number;
@@ -978,7 +986,13 @@ export class AgentOs {
 		}
 	}
 
-	async mkdir(path: string): Promise<void> {
+	async mkdir(
+		path: string,
+		options?: { recursive?: boolean },
+	): Promise<void> {
+		if (options?.recursive) {
+			return this._mkdirp(path);
+		}
 		this._assertSafeAbsolutePath(path);
 		return this.kernel.mkdir(path);
 	}
@@ -1854,12 +1868,32 @@ export class AgentOs {
 
 	// ── Flat session API (ID-based) ───────────────────────────────
 
-	/** Send a prompt to the agent and wait for the final response. */
+	/** Send a prompt to the agent and wait for the final response.
+	 *  Returns the raw JSON-RPC response and the accumulated agent text. */
 	async prompt(
 		sessionId: string,
 		text: string,
-	): Promise<JsonRpcResponse> {
-		return this._requireSession(sessionId).prompt(text);
+	): Promise<PromptResult> {
+		const session = this._requireSession(sessionId);
+
+		// Collect streamed text while the prompt is running
+		let agentText = "";
+		const handler: SessionEventHandler = (event) => {
+			const params = event.params as Record<string, unknown> | undefined;
+			const update = params?.update as Record<string, unknown> | undefined;
+			if (update?.sessionUpdate === "agent_message_chunk") {
+				const content = update.content as { text?: string } | undefined;
+				if (content?.text) agentText += content.text;
+			}
+		};
+		session.onSessionEvent(handler);
+
+		try {
+			const response = await session.prompt(text);
+			return { response, text: agentText };
+		} finally {
+			session.removeSessionEventHandler(handler);
+		}
 	}
 
 	/** Cancel ongoing agent work for a session. */

packages/playground/vendor/monaco

@@ -1 +1 @@
-/home/nathan/a2/packages/playground/node_modules/monaco-editor/min
+/home/nathan/a1/packages/playground/node_modules/monaco-editor/min

packages/playground/vendor/pyodide

@@ -1 +1 @@
-/home/nathan/a2/packages/playground/node_modules/pyodide
+/home/nathan/a1/packages/playground/node_modules/pyodide

packages/playground/vendor/typescript.js

@@ -1 +1 @@
-/home/nathan/a2/packages/playground/node_modules/typescript/lib/typescript.js
+/home/nathan/a1/packages/playground/node_modules/typescript/lib/typescript.js

registry/agent/pi/src/adapter.ts

@@ -37,8 +37,10 @@ import {
 	createAgentSession,
 } from "@mariozechner/pi-coding-agent";
 import type { AgentSession } from "@mariozechner/pi-coding-agent";
-import { isAbsolute, resolve as resolvePath } from "node:path";
-import { readFileSync } from "node:fs";
+import { isAbsolute, join, resolve as resolvePath } from "node:path";
+import { existsSync, readFileSync, readdirSync } from "node:fs";
+import { homedir } from "node:os";
+import type { ExtensionFactory } from "@mariozechner/pi-coding-agent";
 
 // ── CLI argument parsing ────────────────────────────────────────────
 
@@ -51,6 +53,46 @@ for (let i = 0; i < argv.length; i++) {
 	}
 }
 
+// ── Extension discovery ────────────────────────────────────────────
+// Manually discover and load Pi extensions from standard directories.
+// Pi's built-in jiti loader requires performance.now() which the VM's
+// V8 runtime doesn't provide, so we load extensions ourselves via
+// require() and pass them as extensionFactories.
+
+function discoverExtensionFactories(cwd: string): ExtensionFactory[] {
+	const factories: ExtensionFactory[] = [];
+	const dirs = [
+		join(cwd, ".pi", "extensions"),
+		join(homedir(), ".pi", "agent", "extensions"),
+	];
+
+	for (const dir of dirs) {
+		if (!existsSync(dir)) continue;
+		let entries: string[];
+		try {
+			entries = readdirSync(dir);
+		} catch {
+			continue;
+		}
+		for (const name of entries) {
+			if (!name.endsWith(".js") && !name.endsWith(".ts")) continue;
+			const filePath = join(dir, name);
+			try {
+				// biome-ignore lint/security/noGlobalEval: needed to load extensions without jiti
+				const mod = eval(`require(${JSON.stringify(filePath)})`);
+				const factory = mod?.default ?? mod;
+				if (typeof factory === "function") {
+					factories.push(factory);
+				}
+			} catch {
+				// Skip extensions that fail to load
+			}
+		}
+	}
+
+	return factories;
+}
+
 // ── Agent implementation ────────────────────────────────────────────
 
 class PiSdkAgent implements Agent {
@@ -95,19 +137,25 @@ class PiSdkAgent implements Agent {
 	): Promise<NewSessionResponse> {
 		this.cwd = params.cwd;
 
-		const { session } = await createAgentSession({
+		// Discover extensions from standard Pi directories and load them
+		// manually (bypasses jiti which requires performance.now).
+		const extensionFactories = discoverExtensionFactories(params.cwd);
+
+		const { DefaultResourceLoader } = await import(
+			"@mariozechner/pi-coding-agent"
+		);
+		const resourceLoader = new DefaultResourceLoader({
+			cwd: params.cwd,
+			...(appendSystemPrompt ? { appendSystemPrompt } : {}),
+			noExtensions: true, // skip jiti-based discovery
+			extensionFactories,
+		});
+		await resourceLoader.reload();
+
+		const { session, extensionsResult } = await createAgentSession({
 			cwd: params.cwd,
 			sessionManager: SessionManager.inMemory(),
-			...(appendSystemPrompt
-				? {
-						resourceLoader: new (
-							await import("@mariozechner/pi-coding-agent")
-						).DefaultResourceLoader({
-							cwd: params.cwd,
-							appendSystemPrompt,
-						}),
-					}
-				: {}),
+			resourceLoader,
 		});
 
 		this.session = session;