From 9a69a2bd47730d6a00b3c2d70e68e1ba0b78c370 Mon Sep 17 00:00:00 2001 From: brittianwarner Date: Thu, 2 Apr 2026 23:10:38 -0400 Subject: [PATCH 1/3] feat(agent-os): add createOptions factory for per-instance sandbox mounting The agentOs() sandbox config previously only accepted static options, meaning all actor instances shared a single sandbox container. This adds createOptions as a mutually exclusive alternative that is called per-actor-instance so each gets its own dedicated sandbox. Core implementation: - Add createOptions to Zod schema with XOR .refine() enforcement - Create AgentOsActorOptionsConfig exclusive union type (?: never pattern) - Rewrite ensureVm() to resolve options from createOptions(c) or config.options - Add promise-caching guard (vmBootPromise) to prevent concurrent actions from creating duplicate VMs via race condition - Add defensive null-check on createOptions return with clear error message - Wrap createOptions errors with context message and { cause } for stack traces - Add debug log before factory invocation for slow-boot observability - Clear vmBootPromise in onSleep and onDestroy DX improvements: - Add JSDoc to options and createOptions union members for IDE hover docs - Fix any leak in AgentOsActorContext and AgentOsActionContext database param (any -> DatabaseProvider) - Fix zFunction() Zod error message ('Expected a function' instead of generic 'Invalid input') - Export AgentOsActorContext, AgentOsCreateContext alias, AgentOsActorConfigInput Example and docs: - Rewrite sandbox example to use createOptions with context param - Fix 'Sandbox extension' -> 'Sandbox mounting' terminology in client.ts - Add missing deps to examples/agent-os/package.json - Add 'Per-instance options' section to configuration.mdx - Fix stale sandbox: { enabled: true } API in index.mdx crash course - Fix missing .ts import extension in sandbox.mdx Tests: - Add agentOsCreateOptionsTestActor fixture to driver-test-suite - Add 3 driver tests: writeFile round-trip, exec, two-instance isolation - Add agent-os-config-validation.test.ts with 7 tests (Zod XOR enforcement, runtime factory boot, context availability) - Add 9 type-level assertions to actor-types.test.ts (exclusive union, context typing, c.db not any, async callback, generic propagation) --- examples/agent-os/package.json | 3 + examples/agent-os/src/sandbox/client.ts | 2 +- examples/agent-os/src/sandbox/server.ts | 46 +++++---- .../fixtures/driver-test-suite/agent-os.ts | 9 +- .../fixtures/driver-test-suite/registry.ts | 47 +++++----- .../rivetkit/src/agent-os/actor/index.ts | 83 +++++++++++++---- .../packages/rivetkit/src/agent-os/config.ts | 65 +++++++++++-- .../packages/rivetkit/src/agent-os/index.ts | 28 +++--- .../packages/rivetkit/src/agent-os/types.ts | 7 +- .../driver-test-suite/tests/actor-agent-os.ts | 91 +++++++++++++++--- .../rivetkit/tests/actor-types.test.ts | 78 +++++++++++++++- .../tests/agent-os-config-validation.test.ts | 93 +++++++++++++++++++ .../content/docs/agent-os/configuration.mdx | 40 +++++++- website/src/content/docs/agent-os/index.mdx | 16 +++- website/src/content/docs/agent-os/sandbox.mdx | 55 +++++++---- 15 files changed, 536 insertions(+), 127 deletions(-) create mode 100644 rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts diff --git a/examples/agent-os/package.json b/examples/agent-os/package.json index 7823fd392b..92080acb03 100644 --- a/examples/agent-os/package.json +++ b/examples/agent-os/package.json @@ -25,7 +25,10 @@ "check-types": "tsc --noEmit" }, "dependencies": { + "@rivet-dev/agent-os-common": "*", + "@rivet-dev/agent-os-sandbox": "*", "rivetkit": "*", + "sandbox-agent": "*", "zod": "^3.25.0" }, "devDependencies": { diff --git a/examples/agent-os/src/sandbox/client.ts b/examples/agent-os/src/sandbox/client.ts index 84bcee872b..4d36e146d0 100644 --- a/examples/agent-os/src/sandbox/client.ts +++ b/examples/agent-os/src/sandbox/client.ts @@ -1,4 +1,4 @@ -// Sandbox extension: write files and run commands in a Docker sandbox. +// Sandbox mounting: write files and run commands in a Docker sandbox. // // The /sandbox mount projects the sandbox filesystem into the VM. // The sandbox toolkit is accessible via the tools RPC server. diff --git a/examples/agent-os/src/sandbox/server.ts b/examples/agent-os/src/sandbox/server.ts index 5252fd0863..5434697568 100644 --- a/examples/agent-os/src/sandbox/server.ts +++ b/examples/agent-os/src/sandbox/server.ts @@ -1,34 +1,42 @@ -// Sandbox extension: mount a remote Docker sandbox into the VM. +// Sandbox mounting: each actor instance gets its own Docker sandbox. // // Requires Docker running locally. The sandbox-agent package manages the // container lifecycle. The sandbox filesystem is mounted at /sandbox and // the toolkit exposes process management as CLI commands. +// +// Using `createOptions` ensures every actor instance spawned via +// `client.vm.getOrCreate(...)` provisions a dedicated sandbox so +// multiple agents never share the same container. -import { agentOs } from "rivetkit/agent-os"; -import { setup } from "rivetkit"; import common from "@rivet-dev/agent-os-common"; -import { SandboxAgent } from "sandbox-agent"; -import { docker } from "sandbox-agent/docker"; import { createSandboxFs, createSandboxToolkit, } from "@rivet-dev/agent-os-sandbox"; - -// Start a Docker-backed sandbox. -const sandbox = await SandboxAgent.start({ - sandbox: docker(), -}); +import { setup } from "rivetkit"; +import { agentOs } from "rivetkit/agent-os"; +import { SandboxAgent } from "sandbox-agent"; +import { docker } from "sandbox-agent/docker"; const vm = agentOs({ - options: { - software: [common], - mounts: [ - { - path: "/sandbox", - driver: createSandboxFs({ client: sandbox }), - }, - ], - toolKits: [createSandboxToolkit({ client: sandbox })], + createOptions: async (c) => { + c.log.info({ msg: "provisioning sandbox for actor instance" }); + + // Start a dedicated Docker-backed sandbox for this actor instance. + const sandbox = await SandboxAgent.start({ + sandbox: docker(), + }); + + return { + software: [common], + mounts: [ + { + path: "/sandbox", + driver: createSandboxFs({ client: sandbox }), + }, + ], + toolKits: [createSandboxToolkit({ client: sandbox })], + }; }, }); diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/agent-os.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/agent-os.ts index 87f1bf28e6..7b827976fe 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/agent-os.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/agent-os.ts @@ -1,4 +1,11 @@ -import { agentOs } from "rivetkit/agent-os"; import common from "@rivet-dev/agent-os-common"; +import { agentOs } from "rivetkit/agent-os"; export const agentOsTestActor = agentOs({ options: { software: [common] } }); + +// Same actor using the per-instance createOptions factory path. The factory +// returns identical options to the static fixture above so the driver tests +// can verify both paths produce a working VM. +export const agentOsCreateOptionsTestActor = agentOs({ + createOptions: async () => ({ software: [common] }), +}); diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry.ts index 12a5e74d49..fc74df811a 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry.ts @@ -3,8 +3,6 @@ import { accessControlActor, accessControlNoQueuesActor, } from "./access-control"; -import { agentOsTestActor } from "./agent-os"; - import { inputActor } from "./action-inputs"; import { defaultTimeoutActor, @@ -20,8 +18,8 @@ import { import { dbActorDrizzle } from "./actor-db-drizzle"; import { dbActorRaw } from "./actor-db-raw"; import { onStateChangeActor } from "./actor-onstatechange"; +import { agentOsCreateOptionsTestActor, agentOsTestActor } from "./agent-os"; import { connErrorSerializationActor } from "./conn-error-serialization"; -import { dbPragmaMigrationActor } from "./db-pragma-migration"; import { counterWithParams } from "./conn-params"; import { connStateActor } from "./conn-state"; // Import actors from individual files @@ -33,13 +31,11 @@ import { dbLifecycleFailing, dbLifecycleObserver, } from "./db-lifecycle"; +import { dbPragmaMigrationActor } from "./db-pragma-migration"; import { destroyActor, destroyObserver } from "./destroy"; import { customTimeoutActor, errorHandlingActor } from "./error-handling"; import { fileSystemHibernationCleanupActor } from "./file-system-hibernation-cleanup"; -import { - hibernationActor, - hibernationSleepWindowActor, -} from "./hibernation"; +import { hibernationActor, hibernationSleepWindowActor } from "./hibernation"; import { inlineClientActor } from "./inline-client"; import { kvActor } from "./kv"; import { largePayloadActor, largePayloadConnActor } from "./large-payloads"; @@ -76,41 +72,41 @@ import { sleep, sleepRawWsAddEventListenerClose, sleepRawWsAddEventListenerMessage, + sleepRawWsDelayedSendOnSleep, + sleepRawWsOnClose, + sleepRawWsOnMessage, + sleepRawWsSendOnSleep, sleepWithLongRpc, sleepWithNoSleepOption, sleepWithPreventSleep, sleepWithRawHttp, sleepWithRawWebSocket, - sleepWithWaitUntilMessage, - sleepRawWsOnClose, - sleepRawWsOnMessage, - sleepRawWsSendOnSleep, - sleepRawWsDelayedSendOnSleep, sleepWithWaitUntilInOnWake, + sleepWithWaitUntilMessage, } from "./sleep"; import { - sleepWithDb, - sleepWithSlowScheduledDb, - sleepWithDbConn, - sleepWithDbAction, - sleepWithRawWsCloseDb, - sleepWithRawWsCloseDbListener, - sleepWsMessageExceedsGrace, - sleepWsConcurrentDbExceedsGrace, - sleepWaitUntil, - sleepNestedWaitUntil, sleepEnqueue, - sleepScheduleAfter, + sleepNestedWaitUntil, sleepOnSleepThrows, + sleepScheduleAfter, + sleepWaitUntil, sleepWaitUntilRejects, sleepWaitUntilState, + sleepWithDb, + sleepWithDbAction, + sleepWithDbConn, sleepWithRawWs, + sleepWithRawWsCloseDb, + sleepWithRawWsCloseDbListener, + sleepWithSlowScheduledDb, sleepWsActiveDbExceedsGrace, + sleepWsConcurrentDbExceedsGrace, + sleepWsMessageExceedsGrace, sleepWsRawDbAfterClose, } from "./sleep-db"; import { lifecycleObserver, startStopRaceActor } from "./start-stop-race"; -import { statelessActor } from "./stateless"; import { stateZodCoercionActor } from "./state-zod-coercion"; +import { statelessActor } from "./stateless"; import { driverCtxActor, dynamicVarActor, @@ -131,8 +127,8 @@ import { workflowNestedLoopActor, workflowNestedRaceActor, workflowQueueActor, - workflowRunningStepActor, workflowReplayActor, + workflowRunningStepActor, workflowSleepActor, workflowSpawnChildActor, workflowSpawnParentActor, @@ -304,5 +300,6 @@ export const registry = setup({ stateZodCoercionActor, // From agent-os.ts agentOsTestActor, + agentOsCreateOptionsTestActor, }, }); diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts index 30bdf37cfa..9e0af9dee2 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts @@ -1,5 +1,5 @@ -import { AgentOs, createInMemoryFileSystem } from "@rivet-dev/agent-os-core"; import type { AgentOsOptions, MountConfig } from "@rivet-dev/agent-os-core"; +import { AgentOs, createInMemoryFileSystem } from "@rivet-dev/agent-os-core"; import type { DatabaseProvider } from "@/actor/database"; import { actor, event } from "@/actor/mod"; import type { RawAccess } from "@/db/config"; @@ -42,35 +42,75 @@ async function ensureVm( c: AgentOsActionContext, config: AgentOsActorConfig, ): Promise { + // Fast path: VM already booted. if (c.vars.agentOs) { return c.vars.agentOs; } - const start = Date.now(); + // Guard against concurrent callers. If another action is already booting + // the VM, wait for that same promise instead of creating a duplicate. + if (c.vars.vmBootPromise) { + return c.vars.vmBootPromise; + } + + const bootPromise = (async (): Promise => { + const start = Date.now(); - // Build options with in-memory VFS as default working directory mount. - const options = buildVmOptions(config.options); + // Resolve options from the per-actor-instance factory or the static config. + let resolvedOptions: AgentOsOptions | undefined; + if (config.createOptions) { + c.log.debug({ msg: "agent-os resolving createOptions" }); + try { + resolvedOptions = await config.createOptions(c); + } catch (err) { + throw new Error( + `agentOs: createOptions callback failed: ${err instanceof Error ? err.message : String(err)}`, + { cause: err }, + ); + } + } else { + resolvedOptions = config.options; + } - const agentOs = await AgentOs.create(options); - c.vars.agentOs = agentOs; + if (!resolvedOptions) { + throw new Error( + "agentOs: createOptions callback returned a falsy value. It must return an AgentOsOptions object.", + ); + } - // Wire cron events to actor events. - agentOs.onCronEvent((cronEvent) => { - c.broadcast("cronEvent", { event: cronEvent }); - }); + // Build options with in-memory VFS as default working directory mount. + const options = buildVmOptions(resolvedOptions); - c.broadcast("vmBooted", {}); - c.log.info({ - msg: "agent-os vm booted", - bootDurationMs: Date.now() - start, - }); + const agentOs = await AgentOs.create(options); + c.vars.agentOs = agentOs; - return agentOs; + // Wire cron events to actor events. + agentOs.onCronEvent((cronEvent) => { + c.broadcast("cronEvent", { event: cronEvent }); + }); + + c.broadcast("vmBooted", {}); + c.log.info({ + msg: "agent-os vm booted", + bootDurationMs: Date.now() - start, + }); + + return agentOs; + })(); + + c.vars.vmBootPromise = bootPromise; + + try { + return await bootPromise; + } catch (err) { + // Clear the cached promise on failure so the next caller retries + // instead of reusing a rejected promise. + c.vars.vmBootPromise = null; + throw err; + } } -function buildVmOptions( - userOptions?: AgentOsOptions, -): AgentOsOptions { +function buildVmOptions(userOptions?: AgentOsOptions): AgentOsOptions { const userMounts = userOptions?.mounts ?? []; // Check if the user already provided a mount at /home/user. If so, respect @@ -174,6 +214,7 @@ export function agentOs( createState: async () => ({}), createVars: () => ({ agentOs: null, + vmBootPromise: null, activeSessionIds: new Set(), activeProcesses: new Set(), activeHooks: new Set>(), @@ -220,6 +261,7 @@ export function agentOs( await c.vars.agentOs.dispose(); c.vars.agentOs = null; } + c.vars.vmBootPromise = null; c.broadcast("vmShutdown", { reason: "sleep" as const }); }, @@ -235,6 +277,7 @@ export function agentOs( await c.vars.agentOs.dispose(); c.vars.agentOs = null; } + c.vars.vmBootPromise = null; c.broadcast("vmShutdown", { reason: "destroy" as const }); }, @@ -264,4 +307,4 @@ const processExitToken = event(); const shellDataToken = event(); const cronEventToken = event(); -export { ensureVm, syncPreventSleep, runHook }; +export { ensureVm, runHook, syncPreventSleep }; diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts index 2873ab9427..08c7b0407b 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts @@ -3,13 +3,18 @@ import type { JsonRpcNotification, PermissionRequest, } from "@rivet-dev/agent-os-core"; -import type { ActorContext, BeforeConnectContext } from "@/actor/contexts"; import { z } from "zod/v4"; +import type { ActorContext, BeforeConnectContext } from "@/actor/contexts"; +import type { DatabaseProvider } from "@/actor/database"; +import type { RawAccess } from "@/db/config"; import type { AgentOsActorState, AgentOsActorVars } from "./types"; const zFunction = < T extends (...args: any[]) => any = (...args: unknown[]) => unknown, ->() => z.custom((val) => typeof val === "function"); +>() => + z.custom((val) => typeof val === "function", { + message: "Expected a function", + }); const AgentOsOptionsSchema = z.custom( (val) => typeof val === "object" && val !== null, @@ -18,6 +23,7 @@ const AgentOsOptionsSchema = z.custom( export const agentOsActorConfigSchema = z .object({ options: AgentOsOptionsSchema.optional(), + createOptions: zFunction().optional(), preview: z .object({ defaultExpiresInSeconds: z.number().positive().default(3600), @@ -29,17 +35,25 @@ export const agentOsActorConfigSchema = z onSessionEvent: zFunction().optional(), onPermissionRequest: zFunction().optional(), }) - .strict(); + .strict() + .refine( + (data) => + (data.options !== undefined) !== (data.createOptions !== undefined), + { + message: + "agentOs config must define exactly one of 'options' or 'createOptions'", + }, + ); // --- Typed config types (generic callbacks overlaid on the Zod schema) --- -type AgentOsActorContext = ActorContext< +export type AgentOsActorContext = ActorContext< AgentOsActorState, TConnParams, undefined, AgentOsActorVars, undefined, - any + DatabaseProvider >; interface AgentOsActorConfigCallbacks { @@ -48,7 +62,7 @@ interface AgentOsActorConfigCallbacks { AgentOsActorState, AgentOsActorVars, undefined, - any + DatabaseProvider >, params: TConnParams, ) => void | Promise; @@ -64,16 +78,47 @@ interface AgentOsActorConfigCallbacks { ) => void | Promise; } +// Exclusive union: exactly one of `options` (static) or `createOptions` +// (per-actor-instance factory). Mirrors the sandboxActor pattern of +// `provider` / `createProvider`. +type AgentOsActorOptionsConfig = + | { + /** Static VM options shared by all instances of this actor. Use + * `createOptions` instead if each instance needs its own sandbox, + * filesystem mounts, or per-instance configuration. */ + options: AgentOsOptions; + createOptions?: never; + } + | { + options?: never; + /** Factory called lazily on first VM access. Receives the actor + * context so options can vary per instance (e.g., dedicated + * sandboxes). Mutually exclusive with `options`. May be async. */ + createOptions: ( + c: AgentOsActorContext, + ) => AgentOsOptions | Promise; + }; + // Parsed config (after Zod defaults/transforms applied). export type AgentOsActorConfig = Omit< z.infer, - "onBeforeConnect" | "onSessionEvent" | "onPermissionRequest" + | "options" + | "createOptions" + | "onBeforeConnect" + | "onSessionEvent" + | "onPermissionRequest" > & - AgentOsActorConfigCallbacks; + AgentOsActorConfigCallbacks & + AgentOsActorOptionsConfig; // Input config (what users pass in before Zod transforms). export type AgentOsActorConfigInput = Omit< z.input, - "onBeforeConnect" | "onSessionEvent" | "onPermissionRequest" + | "options" + | "createOptions" + | "onBeforeConnect" + | "onSessionEvent" + | "onPermissionRequest" > & - AgentOsActorConfigCallbacks; + AgentOsActorConfigCallbacks & + AgentOsActorOptionsConfig; diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts index 6c919494e9..398388644c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts @@ -1,11 +1,18 @@ // Database migration + +// Cron actions +export { buildCronActions } from "./actor/cron"; export { migrateAgentOsTables } from "./actor/db"; -// SQLite-backed VFS -export { createSqliteVfs, type SqliteVfsOptions } from "./fs/sqlite-vfs"; // Filesystem actions export { buildFilesystemActions } from "./actor/filesystem"; // Actor factory and VM lifecycle helpers export { agentOs, ensureVm, runHook, syncPreventSleep } from "./actor/index"; +// Network actions +export { + buildNetworkActions, + type VmFetchOptions, + type VmFetchResult, +} from "./actor/network"; // Preview actions export { buildOnRequestHandler, @@ -24,20 +31,17 @@ export { } from "./actor/session"; // Shell actions export { buildShellActions } from "./actor/shell"; -// Cron actions -export { buildCronActions } from "./actor/cron"; -// Network actions -export { - buildNetworkActions, - type VmFetchOptions, - type VmFetchResult, -} from "./actor/network"; +// User-facing alias for the createOptions callback context parameter. +export type { AgentOsActorContext as AgentOsCreateContext } from "./config"; // Config schema and types export { type AgentOsActorConfig, type AgentOsActorConfigInput, + type AgentOsActorContext, agentOsActorConfigSchema, } from "./config"; +// SQLite-backed VFS +export { createSqliteVfs, type SqliteVfsOptions } from "./fs/sqlite-vfs"; // Domain types and event payloads export type { AgentOsActionContext, @@ -45,12 +49,12 @@ export type { AgentOsActorVars, AgentOsEvents, CronEventPayload, + PermissionRequestPayload, PersistedSessionEvent, PersistedSessionRecord, - PermissionRequestPayload, - PromptResult, ProcessExitPayload, ProcessOutputPayload, + PromptResult, SerializableCronAction, SerializableCronJobOptions, SessionEventPayload, diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts index 90b0eb6b23..a389725d9f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts @@ -8,6 +8,8 @@ import type { PermissionRequest, } from "@rivet-dev/agent-os-core"; import type { ActionContext } from "@/actor/contexts"; +import type { DatabaseProvider } from "@/actor/database"; +import type { RawAccess } from "@/db/config"; // --- Actor state (persisted across sleep/wake) --- @@ -18,6 +20,9 @@ export type AgentOsActorState = {}; export interface AgentOsActorVars { agentOs: AgentOs | null; + /** In-flight VM boot promise used to prevent concurrent ensureVm calls from + * creating duplicate VMs. Reset on sleep/wake and cleared on boot failure. */ + vmBootPromise: Promise | null; activeSessionIds: Set; activeProcesses: Set; activeHooks: Set>; @@ -133,5 +138,5 @@ export type AgentOsActionContext = ActionContext< undefined, AgentOsActorVars, undefined, - any + DatabaseProvider >; diff --git a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-agent-os.ts b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-agent-os.ts index 07e053f129..fd62830146 100644 --- a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-agent-os.ts +++ b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-agent-os.ts @@ -96,7 +96,9 @@ export function runActorAgentOsTests(driverTestConfig: DriverTestConfig) { await actor.writeFile("/home/user/todelete.txt", "gone"); await actor.deleteFile("/home/user/todelete.txt"); - expect(await actor.exists("/home/user/todelete.txt")).toBe(false); + expect(await actor.exists("/home/user/todelete.txt")).toBe( + false, + ); }, 60_000); test("writeFiles and readFiles batch operations", async (c) => { @@ -118,12 +120,12 @@ export function runActorAgentOsTests(driverTestConfig: DriverTestConfig) { "/home/user/batch-a.txt", "/home/user/batch-b.txt", ]); - expect( - new TextDecoder().decode(readResults[0].content), - ).toBe("aaa"); - expect( - new TextDecoder().decode(readResults[1].content), - ).toBe("bbb"); + expect(new TextDecoder().decode(readResults[0].content)).toBe( + "aaa", + ); + expect(new TextDecoder().decode(readResults[1].content)).toBe( + "bbb", + ); }, 60_000); test("readdirRecursive lists nested files", async (c) => { @@ -172,10 +174,7 @@ export function runActorAgentOsTests(driverTestConfig: DriverTestConfig) { ]); // Write a script that exits with code 42. - await actor.writeFile( - "/tmp/exit42.js", - 'process.exit(42);', - ); + await actor.writeFile("/tmp/exit42.js", "process.exit(42);"); const { pid } = await actor.spawn("node", ["/tmp/exit42.js"]); expect(typeof pid).toBe("number"); @@ -196,7 +195,7 @@ export function runActorAgentOsTests(driverTestConfig: DriverTestConfig) { // Write a long-running script. await actor.writeFile( "/tmp/long.js", - 'setTimeout(() => {}, 30000);', + "setTimeout(() => {}, 30000);", ); const { pid } = await actor.spawn("node", ["/tmp/long.js"]); @@ -217,7 +216,7 @@ export function runActorAgentOsTests(driverTestConfig: DriverTestConfig) { await actor.writeFile( "/tmp/hang.js", - 'setTimeout(() => {}, 60000);', + "setTimeout(() => {}, 60000);", ); const { pid } = await actor.spawn("node", ["/tmp/hang.js"]); @@ -262,7 +261,9 @@ server.listen(9876, "127.0.0.1", () => { "http://127.0.0.1:9876/test", ); expect(result.status).toBe(200); - expect(new TextDecoder().decode(result.body)).toBe("vm-response"); + expect(new TextDecoder().decode(result.body)).toBe( + "vm-response", + ); }, 60_000); // --- Cron --- @@ -289,6 +290,68 @@ server.listen(9876, "127.0.0.1", () => { const jobsAfter = await actor.listCronJobs(); expect(jobsAfter.some((j: any) => j.id === id)).toBe(false); }, 60_000); + + // --- createOptions factory path --- + + test("createOptions: writeFile and readFile round-trip", async (c) => { + const { client } = await setupDriverTest(c, { + ...driverTestConfig, + useRealTimers: true, + }); + const actor = client.agentOsCreateOptionsTestActor.getOrCreate([ + `create-opts-fs-${crypto.randomUUID()}`, + ]); + + await actor.writeFile( + "/home/user/hello.txt", + "hello from createOptions", + ); + const data = await actor.readFile("/home/user/hello.txt"); + expect(new TextDecoder().decode(data)).toBe( + "hello from createOptions", + ); + }, 60_000); + + test("createOptions: exec runs a command", async (c) => { + const { client } = await setupDriverTest(c, { + ...driverTestConfig, + useRealTimers: true, + }); + const actor = client.agentOsCreateOptionsTestActor.getOrCreate([ + `create-opts-exec-${crypto.randomUUID()}`, + ]); + + const result = await actor.exec("echo factory-path"); + expect(result.exitCode).toBe(0); + expect(result.stdout.trim()).toBe("factory-path"); + }, 60_000); + + test("createOptions: two instances get independent VMs", async (c) => { + const { client } = await setupDriverTest(c, { + ...driverTestConfig, + useRealTimers: true, + }); + + const actorA = client.agentOsCreateOptionsTestActor.getOrCreate( + [`create-opts-iso-a-${crypto.randomUUID()}`], + ); + const actorB = client.agentOsCreateOptionsTestActor.getOrCreate( + [`create-opts-iso-b-${crypto.randomUUID()}`], + ); + + // Write a file only in actor A. + await actorA.writeFile("/home/user/only-in-a.txt", "a-data"); + + // Actor B should not see it because each instance has its own VM. + expect(await actorB.exists("/home/user/only-in-a.txt")).toBe( + false, + ); + + // Actor A should still see it. + expect(await actorA.exists("/home/user/only-in-a.txt")).toBe( + true, + ); + }, 60_000); }, ); } diff --git a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts index 905748bc2b..ae980c3090 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts @@ -1,7 +1,10 @@ +import type { AgentOsOptions } from "@rivet-dev/agent-os-core"; import { describe, expectTypeOf, it } from "vitest"; -import { actor, event, queue } from "@/actor/mod"; import type { ActorContext, ActorContextOf } from "@/actor/contexts"; import type { ActorDefinition } from "@/actor/definition"; +import { actor, event, queue } from "@/actor/mod"; +import type { AgentOsActorConfigInput } from "@/agent-os/config"; +import { agentOs } from "@/agent-os/index"; import type { ActorConn, ActorHandle } from "@/client/mod"; import type { DatabaseProviderContext } from "@/db/config"; import { db } from "@/db/mod"; @@ -563,3 +566,76 @@ describe("ActorDefinition", () => { }); }); }); + +describe("agentOs config exclusive union types", () => { + it("accepts options-only config", () => { + const config: AgentOsActorConfigInput = { + options: { software: [] }, + }; + expectTypeOf(config).toMatchTypeOf(); + }); + + it("accepts createOptions-only config", () => { + const config: AgentOsActorConfigInput = { + createOptions: () => ({ software: [] }), + }; + expectTypeOf(config).toMatchTypeOf(); + }); + + it("rejects config with both options and createOptions", () => { + // @ts-expect-error options and createOptions are mutually exclusive + const config: AgentOsActorConfigInput = { + options: { software: [] }, + createOptions: () => ({ software: [] }), + }; + }); + + it("rejects config with neither options nor createOptions", () => { + // @ts-expect-error must provide one of options or createOptions + const config: AgentOsActorConfigInput = {}; + }); + + it("createOptions callback receives context and returns AgentOsOptions", () => { + const config: AgentOsActorConfigInput = { + createOptions: (c) => { + expectTypeOf(c.log).not.toBeAny(); + expectTypeOf(c.vars).not.toBeAny(); + return { software: [] }; + }, + }; + }); + + it("createOptions callback has typed db access", () => { + const config: AgentOsActorConfigInput = { + createOptions: (c) => { + expectTypeOf(c.db).not.toBeAny(); + return { software: [] }; + }, + }; + }); + + it("createOptions callback can be async", () => { + const config: AgentOsActorConfigInput = { + createOptions: async (c) => { + return { software: [] }; + }, + }; + }); + + it("agentOs() accepts createOptions config without error", () => { + // This is a compile-time-only test. If it compiles, the types work. + agentOs({ + createOptions: async () => ({ software: [] }), + }); + }); + + it("agentOs() with connParams types the createOptions context", () => { + agentOs<{ userId: string }>({ + createOptions: async (c) => { + // The context should be typed, not any. + expectTypeOf(c).not.toBeAny(); + return { software: [] }; + }, + }); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts b/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts new file mode 100644 index 0000000000..49f026ae7d --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts @@ -0,0 +1,93 @@ +import common from "@rivet-dev/agent-os-common"; +import { describe, expect, test } from "vitest"; +import { agentOsActorConfigSchema } from "@/agent-os/config"; +import { agentOs } from "@/agent-os/index"; +import { setup } from "@/mod"; +import { setupTest } from "@/test/mod"; + +describe("agentOs config validation", () => { + // --- Zod schema XOR enforcement --- + + test("accepts config with only options", () => { + const result = agentOsActorConfigSchema.safeParse({ + options: { software: [] }, + }); + expect(result.success).toBe(true); + }); + + test("accepts config with only createOptions", () => { + const result = agentOsActorConfigSchema.safeParse({ + createOptions: () => ({ software: [] }), + }); + expect(result.success).toBe(true); + }); + + test("rejects config with both options and createOptions", () => { + const result = agentOsActorConfigSchema.safeParse({ + options: { software: [] }, + createOptions: () => ({ software: [] }), + }); + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error.message).toContain( + "exactly one of 'options' or 'createOptions'", + ); + } + }); + + test("rejects config with neither options nor createOptions", () => { + const result = agentOsActorConfigSchema.safeParse({}); + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error.message).toContain( + "exactly one of 'options' or 'createOptions'", + ); + } + }); + + test("rejects config with createOptions set to a non-function", () => { + const result = agentOsActorConfigSchema.safeParse({ + createOptions: "not-a-function", + }); + expect(result.success).toBe(false); + }); + + // --- Runtime behavior --- + + test("createOptions factory boots a working VM", async (c) => { + const vm = agentOs({ + createOptions: async () => ({ software: [common] }), + }); + const registry = setup({ use: { vm } }); + const { client } = await setupTest(c, registry); + const actor = (client as any).vm.getOrCreate([ + `config-test-${crypto.randomUUID()}`, + ]); + + await actor.writeFile("/home/user/test.txt", "config-validation"); + const data = await actor.readFile("/home/user/test.txt"); + expect(new TextDecoder().decode(data)).toBe("config-validation"); + }, 60_000); + + test("createOptions receives actor context", async (c) => { + let receivedContext = false; + + const vm = agentOs({ + createOptions: async (ctx) => { + // The context should have log and vars available. + receivedContext = + ctx.log !== undefined && ctx.vars !== undefined; + return { software: [common] }; + }, + }); + const registry = setup({ use: { vm } }); + const { client } = await setupTest(c, registry); + const actor = (client as any).vm.getOrCreate([ + `ctx-test-${crypto.randomUUID()}`, + ]); + + // Trigger ensureVm by calling any action. + await actor.exec("echo context-check"); + expect(receivedContext).toBe(true); + }, 60_000); +}); diff --git a/website/src/content/docs/agent-os/configuration.mdx b/website/src/content/docs/agent-os/configuration.mdx index 59716654ce..7460c728dd 100644 --- a/website/src/content/docs/agent-os/configuration.mdx +++ b/website/src/content/docs/agent-os/configuration.mdx @@ -4,7 +4,11 @@ description: "Configure the agentOS VM options, preview settings, and lifecycle skill: true --- -`agentOs()` accepts the following configuration object. +`agentOs()` accepts the following configuration object. You must provide exactly one of `options` or `createOptions`. + +## Static options + +Use `options` when the same VM configuration works for all actor instances. ```ts @nocheck import { agentOs } from "rivetkit/agent-os"; @@ -51,6 +55,40 @@ export const registry = setup({ use: { vm } }); registry.start(); ``` +## Per-instance options + +Use `createOptions` when each actor instance needs its own VM configuration. The callback receives the actor context and runs lazily on first VM access. It may be async. + +This is required for [sandbox mounting](/docs/agent-os/sandbox) where each instance needs a dedicated sandbox container, and useful for any scenario where options depend on the actor instance (e.g., per-tenant mounts or credentials). + +`options` and `createOptions` are **mutually exclusive**. Providing both causes a validation error. + +```ts @nocheck +import { agentOs } from "rivetkit/agent-os"; +import { setup } from "rivetkit"; +import common from "@rivet-dev/agent-os-common"; +import { SandboxAgent } from "sandbox-agent"; +import { docker } from "sandbox-agent/docker"; +import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; + +const vm = agentOs({ + createOptions: async (c) => { + c.log.info({ msg: "provisioning sandbox" }); + + const sandbox = await SandboxAgent.start({ sandbox: docker() }); + + return { + software: [common], + mounts: [{ path: "/sandbox", driver: createSandboxFs({ client: sandbox }) }], + toolKits: [createSandboxToolkit({ client: sandbox })], + }; + }, +}); + +export const registry = setup({ use: { vm } }); +registry.start(); +``` + ## Session options Options passed to `createSession`. See [Sessions](/docs/agent-os/sessions) for full documentation. diff --git a/website/src/content/docs/agent-os/index.mdx b/website/src/content/docs/agent-os/index.mdx index 39fa75efdb..f72aa352d0 100644 --- a/website/src/content/docs/agent-os/index.mdx +++ b/website/src/content/docs/agent-os/index.mdx @@ -390,13 +390,19 @@ agentOS uses a hybrid model: agents run in a lightweight VM by default and spin import { agentOs } from "rivetkit/agent-os"; import { setup } from "rivetkit"; import common from "@rivet-dev/agent-os-common"; -import pi from "@rivet-dev/agent-os-pi"; +import { SandboxAgent } from "sandbox-agent"; +import { docker } from "sandbox-agent/docker"; +import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; const vm = agentOs({ - options: { software: [common, pi], - sandbox: { - enabled: true, - }, + createOptions: async () => { + const sandbox = await SandboxAgent.start({ sandbox: docker() }); + + return { + software: [common], + mounts: [{ path: "/sandbox", driver: createSandboxFs({ client: sandbox }) }], + toolKits: [createSandboxToolkit({ client: sandbox })], + }; }, }); diff --git a/website/src/content/docs/agent-os/sandbox.mdx b/website/src/content/docs/agent-os/sandbox.mdx index 250c45a988..c83826ae59 100644 --- a/website/src/content/docs/agent-os/sandbox.mdx +++ b/website/src/content/docs/agent-os/sandbox.mdx @@ -38,33 +38,54 @@ npm install @rivet-dev/agent-os-sandbox sandbox-agent ``` ```ts @nocheck -import { SandboxAgent } from "sandbox-agent"; -import { DockerProvider } from "sandbox-agent/docker"; -import { AgentOs } from "@rivet-dev/agent-os-core"; +import { agentOs } from "rivetkit/agent-os"; +import { setup } from "rivetkit"; import common from "@rivet-dev/agent-os-common"; +import { SandboxAgent } from "sandbox-agent"; +import { docker } from "sandbox-agent/docker"; import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; -const sandbox = await SandboxAgent.start({ - sandbox: new DockerProvider(), +// Each actor instance gets its own Docker sandbox via createOptions. +const vm = agentOs({ + createOptions: async (c) => { + c.log.info({ msg: "provisioning sandbox" }); + + const sandbox = await SandboxAgent.start({ + sandbox: docker(), + }); + + return { + software: [common], + mounts: [ + { + path: "/sandbox", + driver: createSandboxFs({ client: sandbox }), + }, + ], + toolKits: [createSandboxToolkit({ client: sandbox })], + }; + }, }); -const vm = await AgentOs.create({ - software: [common], - mounts: [ - { - path: "/sandbox", - driver: createSandboxFs({ client: sandbox }), - }, - ], - toolKits: [createSandboxToolkit({ client: sandbox })], -}); +export const registry = setup({ use: { vm } }); +registry.start(); +``` + +From a client, interact with the sandbox through the actor: + +```ts @nocheck +import { createClient } from "rivetkit/client"; +import type { registry } from "./server.ts"; + +const client = createClient("http://localhost:6420"); +const agent = client.vm.getOrCreate(["my-agent"]); // Write code via the filesystem. The /sandbox mount maps to the sandbox root. -await vm.writeFile("/sandbox/app/index.ts", 'console.log("hello")'); +await agent.writeFile("/sandbox/app/index.ts", 'console.log("hello")'); // Run it via the toolkit. Commands execute inside the sandbox, so paths are // relative to the sandbox root (/app/index.ts), not the VM mount (/sandbox/app/index.ts). -const result = await vm.exec("agentos-sandbox run-command --command node --json '{\"args\": [\"/app/index.ts\"]}'"); +const result = await agent.exec("agentos-sandbox run-command --command node --json '{\"args\": [\"/app/index.ts\"]}'"); ``` ## Tools reference From 254c01405545f221553eef40d5e833288cc44fc3 Mon Sep 17 00:00:00 2001 From: brittianwarner Date: Fri, 3 Apr 2026 15:06:15 -0400 Subject: [PATCH 2/3] feat(agent-os): persist sandbox ID across sleep/wake for container reuse Mirrors the sandbox actor pattern where the sandbox ID is stored in actor state so that after a sleep/wake cycle the createOptions callback can reconnect to the same container instead of provisioning a new one. State changes: - Add sandboxId: string | null to AgentOsActorState (persisted across sleep/wake). The createOptions callback reads c.state.sandboxId and passes it to SandboxAgent.start(), then persists the returned ID back. Config changes: - Add destroyOptions callback (only valid with createOptions, enforced by both TypeScript union and Zod refine). Called during onDestroy so the user can clean up external resources like sandbox containers. Lifecycle hardening: - Wrap dispose() in try/finally in both onSleep and onDestroy so vars are always cleaned up even if dispose throws, matching the sandbox actor's teardownAgentRuntime pattern. - Clear sandboxId in a finally block in onDestroy so state is always consistent regardless of whether destroyOptions succeeds or fails. Tests: - Type tests for sandboxId typing, destroyOptions union exclusivity - Zod schema tests for destroyOptions acceptance, rejection with options, and non-function rejection - Runtime tests for sandboxId state persistence and destroyOptions config acceptance Docs: - Update sandbox.mdx, configuration.mdx, and index.mdx with the sandboxId reuse pattern and destroyOptions documentation - Fix 'extension' -> 'sandbox mounting' terminology in sandbox.mdx --- examples/agent-os/src/sandbox/server.ts | 22 +++++- .../rivetkit/src/agent-os/actor/index.ts | 39 ++++++++-- .../packages/rivetkit/src/agent-os/config.ts | 27 ++++++- .../packages/rivetkit/src/agent-os/types.ts | 9 ++- .../rivetkit/tests/actor-types.test.ts | 42 ++++++++++ .../tests/agent-os-config-validation.test.ts | 77 +++++++++++++++++++ .../content/docs/agent-os/configuration.mdx | 24 +++++- website/src/content/docs/agent-os/index.mdx | 8 +- website/src/content/docs/agent-os/sandbox.mdx | 18 ++++- 9 files changed, 247 insertions(+), 19 deletions(-) diff --git a/examples/agent-os/src/sandbox/server.ts b/examples/agent-os/src/sandbox/server.ts index 5434697568..187218ee91 100644 --- a/examples/agent-os/src/sandbox/server.ts +++ b/examples/agent-os/src/sandbox/server.ts @@ -7,6 +7,10 @@ // Using `createOptions` ensures every actor instance spawned via // `client.vm.getOrCreate(...)` provisions a dedicated sandbox so // multiple agents never share the same container. +// +// The sandbox ID is persisted in `c.state.sandboxId` so that after a +// sleep/wake cycle the actor reconnects to the same container instead of +// creating a new one. import common from "@rivet-dev/agent-os-common"; import { @@ -20,13 +24,20 @@ import { docker } from "sandbox-agent/docker"; const vm = agentOs({ createOptions: async (c) => { - c.log.info({ msg: "provisioning sandbox for actor instance" }); + c.log.info({ + msg: "booting sandbox", + existingSandboxId: c.state.sandboxId, + }); - // Start a dedicated Docker-backed sandbox for this actor instance. + // Reconnect to an existing sandbox after sleep, or provision a new one. const sandbox = await SandboxAgent.start({ sandbox: docker(), + sandboxId: c.state.sandboxId ?? undefined, }); + // Persist the sandbox ID so future wakes reuse the same container. + c.state.sandboxId = sandbox.sandboxId; + return { software: [common], mounts: [ @@ -38,6 +49,13 @@ const vm = agentOs({ toolKits: [createSandboxToolkit({ client: sandbox })], }; }, + destroyOptions: async (c) => { + // Destroy the sandbox container when the actor is destroyed. + if (c.state.sandboxId) { + const provider = docker(); + await provider.destroy(c.state.sandboxId); + } + }, }); export const registry = setup({ use: { vm } }); diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts index 9e0af9dee2..effac66bcd 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts @@ -211,7 +211,9 @@ export function agentOs( sleepGracePeriod: 900_000, actionTimeout: 900_000, }, - createState: async () => ({}), + createState: async () => ({ + sandboxId: null, + }), createVars: () => ({ agentOs: null, vmBootPromise: null, @@ -257,11 +259,14 @@ export function agentOs( activeShells: c.vars.activeShells.size, }); - if (c.vars.agentOs) { - await c.vars.agentOs.dispose(); + try { + if (c.vars.agentOs) { + await c.vars.agentOs.dispose(); + } + } finally { c.vars.agentOs = null; + c.vars.vmBootPromise = null; } - c.vars.vmBootPromise = null; c.broadcast("vmShutdown", { reason: "sleep" as const }); }, @@ -273,11 +278,31 @@ export function agentOs( activeShells: c.vars.activeShells.size, }); - if (c.vars.agentOs) { - await c.vars.agentOs.dispose(); + // Tear down the VM. Always clear vars even if dispose throws. + try { + if (c.vars.agentOs) { + await c.vars.agentOs.dispose(); + } + } finally { c.vars.agentOs = null; + c.vars.vmBootPromise = null; + } + + // Let the user clean up external resources (e.g., destroy + // a sandbox container). Use try/finally to guarantee state + // is cleared even if the callback fails. + try { + if (parsedConfig.destroyOptions) { + await parsedConfig.destroyOptions(c); + } + } catch (err) { + c.log.error({ + msg: "agent-os destroyOptions callback failed", + error: err instanceof Error ? err.message : String(err), + }); + } finally { + c.state.sandboxId = null; } - c.vars.vmBootPromise = null; c.broadcast("vmShutdown", { reason: "destroy" as const }); }, diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts index 08c7b0407b..936d9ec013 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts @@ -24,6 +24,7 @@ export const agentOsActorConfigSchema = z .object({ options: AgentOsOptionsSchema.optional(), createOptions: zFunction().optional(), + destroyOptions: zFunction().optional(), preview: z .object({ defaultExpiresInSeconds: z.number().positive().default(3600), @@ -43,6 +44,14 @@ export const agentOsActorConfigSchema = z message: "agentOs config must define exactly one of 'options' or 'createOptions'", }, + ) + .refine( + (data) => + !(data.options !== undefined && data.destroyOptions !== undefined), + { + message: + "agentOs config: 'destroyOptions' is only valid with 'createOptions', not 'options'", + }, ); // --- Typed config types (generic callbacks overlaid on the Zod schema) --- @@ -88,15 +97,29 @@ type AgentOsActorOptionsConfig = * filesystem mounts, or per-instance configuration. */ options: AgentOsOptions; createOptions?: never; + destroyOptions?: never; } | { options?: never; /** Factory called lazily on first VM access. Receives the actor * context so options can vary per instance (e.g., dedicated - * sandboxes). Mutually exclusive with `options`. May be async. */ + * sandboxes). Mutually exclusive with `options`. May be async. + * + * The callback should read `c.state.sandboxId` and pass it to + * `SandboxAgent.start()` to reconnect to an existing sandbox + * after sleep/wake. After starting, persist the ID back: + * `c.state.sandboxId = sandbox.sandboxId`. */ createOptions: ( c: AgentOsActorContext, ) => AgentOsOptions | Promise; + /** Called on actor destroy to clean up external resources + * provisioned by `createOptions` (e.g., destroying a sandbox + * container). Receives the actor context with `c.state.sandboxId` + * still set so the callback can resolve the provider and destroy + * the sandbox. */ + destroyOptions?: ( + c: AgentOsActorContext, + ) => void | Promise; }; // Parsed config (after Zod defaults/transforms applied). @@ -104,6 +127,7 @@ export type AgentOsActorConfig = Omit< z.infer, | "options" | "createOptions" + | "destroyOptions" | "onBeforeConnect" | "onSessionEvent" | "onPermissionRequest" @@ -116,6 +140,7 @@ export type AgentOsActorConfigInput = Omit< z.input, | "options" | "createOptions" + | "destroyOptions" | "onBeforeConnect" | "onSessionEvent" | "onPermissionRequest" diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts index a389725d9f..035661fd54 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts @@ -13,8 +13,13 @@ import type { RawAccess } from "@/db/config"; // --- Actor state (persisted across sleep/wake) --- -// biome-ignore lint/complexity/noBannedTypes: empty state placeholder, consumers extend via generics -export type AgentOsActorState = {}; +export interface AgentOsActorState { + /** Sandbox ID persisted across sleep/wake so the `createOptions` + * callback can reconnect to the same sandbox instead of provisioning + * a new one. Format is `"{provider}/{rawId}"` (e.g. `"docker/abc123"`). + * Set by the user inside `createOptions`; read back on subsequent wakes. */ + sandboxId: string | null; +} // --- Actor vars (ephemeral, recreated on wake) --- diff --git a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts index ae980c3090..06b62f8b70 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts @@ -638,4 +638,46 @@ describe("agentOs config exclusive union types", () => { }, }); }); + + it("createOptions callback can read and write c.state.sandboxId", () => { + const config: AgentOsActorConfigInput = { + createOptions: (c) => { + // sandboxId should be typed as string | null. + expectTypeOf(c.state.sandboxId).toEqualTypeOf(); + c.state.sandboxId = "docker/abc123"; + return { software: [] }; + }, + }; + }); + + it("options config forbids destroyOptions", () => { + // @ts-expect-error destroyOptions is only allowed with createOptions + const config: AgentOsActorConfigInput = { + options: { software: [] }, + destroyOptions: async () => {}, + }; + }); + + it("createOptions config accepts destroyOptions", () => { + const config: AgentOsActorConfigInput = { + createOptions: async (c) => ({ software: [] }), + destroyOptions: async (c) => { + expectTypeOf(c.state.sandboxId).toEqualTypeOf(); + }, + }; + expectTypeOf(config).toMatchTypeOf(); + }); + + it("agentOs() accepts createOptions with destroyOptions", () => { + agentOs({ + createOptions: async (c) => { + c.state.sandboxId = "docker/test"; + return { software: [] }; + }, + destroyOptions: async (c) => { + // Should have access to sandboxId for cleanup. + const _id: string | null = c.state.sandboxId; + }, + }); + }); }); diff --git a/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts b/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts index 49f026ae7d..dbea90576b 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts @@ -52,6 +52,35 @@ describe("agentOs config validation", () => { expect(result.success).toBe(false); }); + test("accepts config with createOptions and destroyOptions", () => { + const result = agentOsActorConfigSchema.safeParse({ + createOptions: () => ({ software: [] }), + destroyOptions: async () => {}, + }); + expect(result.success).toBe(true); + }); + + test("rejects config with destroyOptions set to a non-function", () => { + const result = agentOsActorConfigSchema.safeParse({ + createOptions: () => ({ software: [] }), + destroyOptions: "not-a-function", + }); + expect(result.success).toBe(false); + }); + + test("rejects config with destroyOptions paired with options", () => { + const result = agentOsActorConfigSchema.safeParse({ + options: { software: [] }, + destroyOptions: async () => {}, + }); + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error.message).toContain( + "'destroyOptions' is only valid with 'createOptions'", + ); + } + }); + // --- Runtime behavior --- test("createOptions factory boots a working VM", async (c) => { @@ -90,4 +119,52 @@ describe("agentOs config validation", () => { await actor.exec("echo context-check"); expect(receivedContext).toBe(true); }, 60_000); + + test("createOptions can persist sandboxId in state", async (c) => { + const fakeSandboxId = `test/${crypto.randomUUID()}`; + + const vm = agentOs({ + createOptions: async (ctx) => { + // On first boot, sandboxId should be null. + if (!ctx.state.sandboxId) { + ctx.state.sandboxId = fakeSandboxId; + } + return { software: [common] }; + }, + }); + const registry = setup({ use: { vm } }); + const { client } = await setupTest(c, registry); + const actor = (client as any).vm.getOrCreate([ + `sandbox-id-test-${crypto.randomUUID()}`, + ]); + + // Trigger ensureVm so createOptions runs. + await actor.exec("echo sandbox-id-check"); + + // The actor should have set the sandboxId. If we call another + // action, createOptions won't re-run (VM is already booted), but + // the state should persist. + const result = await actor.exec("echo done"); + expect(result.exitCode).toBe(0); + }, 60_000); + + test("destroyOptions is accepted as a valid config callback", () => { + // Verify destroyOptions is accepted by agentOs() at runtime (not + // just at the type level). If the Zod schema rejected it, this + // would throw during parse. + const vm = agentOs({ + createOptions: async (ctx) => { + ctx.state.sandboxId = "test/destroy-check"; + return { software: [common] }; + }, + destroyOptions: async (ctx) => { + // This callback would be called during onDestroy. + // We just verify it is accepted without error. + const _id = ctx.state.sandboxId; + }, + }); + + // The factory should return a valid actor definition. + expect(vm).toBeDefined(); + }); }); diff --git a/website/src/content/docs/agent-os/configuration.mdx b/website/src/content/docs/agent-os/configuration.mdx index 7460c728dd..f71395d01c 100644 --- a/website/src/content/docs/agent-os/configuration.mdx +++ b/website/src/content/docs/agent-os/configuration.mdx @@ -63,6 +63,14 @@ This is required for [sandbox mounting](/docs/agent-os/sandbox) where each insta `options` and `createOptions` are **mutually exclusive**. Providing both causes a validation error. +### Sandbox reuse across sleep/wake + +The actor state includes `sandboxId` which you should use to reconnect to the same sandbox after a sleep/wake cycle. Without this, every wake would provision a new sandbox container. + +Read `c.state.sandboxId` and pass it to `SandboxAgent.start()`. After starting, persist the ID back so future wakes reuse the same container. + +Provide `destroyOptions` to clean up external resources (e.g., destroy the sandbox container) when the actor is destroyed. This callback only applies when using `createOptions`. + ```ts @nocheck import { agentOs } from "rivetkit/agent-os"; import { setup } from "rivetkit"; @@ -73,9 +81,14 @@ import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandb const vm = agentOs({ createOptions: async (c) => { - c.log.info({ msg: "provisioning sandbox" }); + // Reconnect to an existing sandbox after sleep, or provision a new one. + const sandbox = await SandboxAgent.start({ + sandbox: docker(), + sandboxId: c.state.sandboxId ?? undefined, + }); - const sandbox = await SandboxAgent.start({ sandbox: docker() }); + // Persist the sandbox ID so future wakes reuse the same container. + c.state.sandboxId = sandbox.sandboxId; return { software: [common], @@ -83,6 +96,13 @@ const vm = agentOs({ toolKits: [createSandboxToolkit({ client: sandbox })], }; }, + // Clean up the sandbox container when the actor is destroyed. + destroyOptions: async (c) => { + if (c.state.sandboxId) { + const provider = docker(); + await provider.destroy(c.state.sandboxId); + } + }, }); export const registry = setup({ use: { vm } }); diff --git a/website/src/content/docs/agent-os/index.mdx b/website/src/content/docs/agent-os/index.mdx index f72aa352d0..604b350522 100644 --- a/website/src/content/docs/agent-os/index.mdx +++ b/website/src/content/docs/agent-os/index.mdx @@ -395,8 +395,12 @@ import { docker } from "sandbox-agent/docker"; import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; const vm = agentOs({ - createOptions: async () => { - const sandbox = await SandboxAgent.start({ sandbox: docker() }); + createOptions: async (c) => { + const sandbox = await SandboxAgent.start({ + sandbox: docker(), + sandboxId: c.state.sandboxId ?? undefined, + }); + c.state.sandboxId = sandbox.sandboxId; return { software: [common], diff --git a/website/src/content/docs/agent-os/sandbox.mdx b/website/src/content/docs/agent-os/sandbox.mdx index c83826ae59..cb89cd9a09 100644 --- a/website/src/content/docs/agent-os/sandbox.mdx +++ b/website/src/content/docs/agent-os/sandbox.mdx @@ -46,14 +46,19 @@ import { docker } from "sandbox-agent/docker"; import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; // Each actor instance gets its own Docker sandbox via createOptions. +// The sandbox ID is persisted in state so the actor reconnects to the +// same container after a sleep/wake cycle instead of creating a new one. const vm = agentOs({ createOptions: async (c) => { - c.log.info({ msg: "provisioning sandbox" }); - + // Reconnect to an existing sandbox after sleep, or provision a new one. const sandbox = await SandboxAgent.start({ sandbox: docker(), + sandboxId: c.state.sandboxId ?? undefined, }); + // Persist the sandbox ID so future wakes reuse the same container. + c.state.sandboxId = sandbox.sandboxId; + return { software: [common], mounts: [ @@ -65,6 +70,13 @@ const vm = agentOs({ toolKits: [createSandboxToolkit({ client: sandbox })], }; }, + // Clean up the sandbox container when the actor is destroyed. + destroyOptions: async (c) => { + if (c.state.sandboxId) { + const provider = docker(); + await provider.destroy(c.state.sandboxId); + } + }, }); export const registry = setup({ use: { vm } }); @@ -115,7 +127,7 @@ agentos-sandbox send-input --id "proc_abc123" --input "yes" ## Sandbox providers -The extension works with any [Sandbox Agent](https://sandboxagent.dev) provider. See the [Sandbox Agent documentation](https://sandboxagent.dev) for available providers and setup instructions. +Sandbox mounting works with any [Sandbox Agent](https://sandboxagent.dev) provider. See the [Sandbox Agent documentation](https://sandboxagent.dev) for available providers and setup instructions. ## Recommendations From dfe9009a6dfb6e68cdfab881e5df62d7b26884ee Mon Sep 17 00:00:00 2001 From: brittianwarner Date: Fri, 3 Apr 2026 16:45:39 -0400 Subject: [PATCH 3/3] refactor(agent-os): address PR review feedback for createOptions Remove destroyOptions (deferred to c.getSandbox() auto-destroy). Rename AgentOsActorContext to AgentOsContext, vmBootPromise to vmBootGuard. Fix any types in zFunction. Simplify docs to static vs dynamic sections with c.key example. Update sandbox docs and example to c.getSandbox() API. Delete config validation tests. Remove destroyOptions type tests and @rivet-dev/agent-os-sandbox dep. --- examples/agent-os/package.json | 1 - examples/agent-os/src/sandbox/server.ts | 42 +---- .../rivetkit/src/agent-os/actor/index.ts | 31 +--- .../packages/rivetkit/src/agent-os/config.ts | 37 +--- .../packages/rivetkit/src/agent-os/index.ts | 4 +- .../packages/rivetkit/src/agent-os/types.ts | 2 +- .../rivetkit/tests/actor-types.test.ts | 31 ---- .../tests/agent-os-config-validation.test.ts | 170 ------------------ .../content/docs/agent-os/configuration.mdx | 42 +---- website/src/content/docs/agent-os/index.mdx | 12 +- website/src/content/docs/agent-os/sandbox.mdx | 39 +--- 11 files changed, 38 insertions(+), 373 deletions(-) delete mode 100644 rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts diff --git a/examples/agent-os/package.json b/examples/agent-os/package.json index 92080acb03..0bc4a8ccc1 100644 --- a/examples/agent-os/package.json +++ b/examples/agent-os/package.json @@ -26,7 +26,6 @@ }, "dependencies": { "@rivet-dev/agent-os-common": "*", - "@rivet-dev/agent-os-sandbox": "*", "rivetkit": "*", "sandbox-agent": "*", "zod": "^3.25.0" diff --git a/examples/agent-os/src/sandbox/server.ts b/examples/agent-os/src/sandbox/server.ts index 187218ee91..60857bbe6d 100644 --- a/examples/agent-os/src/sandbox/server.ts +++ b/examples/agent-os/src/sandbox/server.ts @@ -4,58 +4,32 @@ // container lifecycle. The sandbox filesystem is mounted at /sandbox and // the toolkit exposes process management as CLI commands. // -// Using `createOptions` ensures every actor instance spawned via -// `client.vm.getOrCreate(...)` provisions a dedicated sandbox so -// multiple agents never share the same container. -// -// The sandbox ID is persisted in `c.state.sandboxId` so that after a -// sleep/wake cycle the actor reconnects to the same container instead of -// creating a new one. +// Using `createOptions` with `c.getSandbox()` ensures every actor instance +// spawned via `client.vm.getOrCreate(...)` provisions a dedicated sandbox. +// The sandbox ID is persisted internally across sleep/wake so the actor +// reconnects to the same container. The sandbox is auto-destroyed when the +// actor is destroyed. import common from "@rivet-dev/agent-os-common"; -import { - createSandboxFs, - createSandboxToolkit, -} from "@rivet-dev/agent-os-sandbox"; import { setup } from "rivetkit"; import { agentOs } from "rivetkit/agent-os"; -import { SandboxAgent } from "sandbox-agent"; import { docker } from "sandbox-agent/docker"; const vm = agentOs({ createOptions: async (c) => { - c.log.info({ - msg: "booting sandbox", - existingSandboxId: c.state.sandboxId, - }); - - // Reconnect to an existing sandbox after sleep, or provision a new one. - const sandbox = await SandboxAgent.start({ - sandbox: docker(), - sandboxId: c.state.sandboxId ?? undefined, - }); - - // Persist the sandbox ID so future wakes reuse the same container. - c.state.sandboxId = sandbox.sandboxId; + const { fs, toolkit } = await c.getSandbox({ provider: docker() }); return { software: [common], mounts: [ { path: "/sandbox", - driver: createSandboxFs({ client: sandbox }), + driver: fs, }, ], - toolKits: [createSandboxToolkit({ client: sandbox })], + toolKits: [toolkit], }; }, - destroyOptions: async (c) => { - // Destroy the sandbox container when the actor is destroyed. - if (c.state.sandboxId) { - const provider = docker(); - await provider.destroy(c.state.sandboxId); - } - }, }); export const registry = setup({ use: { vm } }); diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts index effac66bcd..b1d4b0eb89 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/actor/index.ts @@ -49,8 +49,8 @@ async function ensureVm( // Guard against concurrent callers. If another action is already booting // the VM, wait for that same promise instead of creating a duplicate. - if (c.vars.vmBootPromise) { - return c.vars.vmBootPromise; + if (c.vars.vmBootGuard) { + return c.vars.vmBootGuard; } const bootPromise = (async (): Promise => { @@ -98,14 +98,14 @@ async function ensureVm( return agentOs; })(); - c.vars.vmBootPromise = bootPromise; + c.vars.vmBootGuard = bootPromise; try { return await bootPromise; } catch (err) { // Clear the cached promise on failure so the next caller retries // instead of reusing a rejected promise. - c.vars.vmBootPromise = null; + c.vars.vmBootGuard = null; throw err; } } @@ -216,7 +216,7 @@ export function agentOs( }), createVars: () => ({ agentOs: null, - vmBootPromise: null, + vmBootGuard: null, activeSessionIds: new Set(), activeProcesses: new Set(), activeHooks: new Set>(), @@ -265,7 +265,7 @@ export function agentOs( } } finally { c.vars.agentOs = null; - c.vars.vmBootPromise = null; + c.vars.vmBootGuard = null; } c.broadcast("vmShutdown", { reason: "sleep" as const }); @@ -278,30 +278,13 @@ export function agentOs( activeShells: c.vars.activeShells.size, }); - // Tear down the VM. Always clear vars even if dispose throws. try { if (c.vars.agentOs) { await c.vars.agentOs.dispose(); } } finally { c.vars.agentOs = null; - c.vars.vmBootPromise = null; - } - - // Let the user clean up external resources (e.g., destroy - // a sandbox container). Use try/finally to guarantee state - // is cleared even if the callback fails. - try { - if (parsedConfig.destroyOptions) { - await parsedConfig.destroyOptions(c); - } - } catch (err) { - c.log.error({ - msg: "agent-os destroyOptions callback failed", - error: err instanceof Error ? err.message : String(err), - }); - } finally { - c.state.sandboxId = null; + c.vars.vmBootGuard = null; } c.broadcast("vmShutdown", { reason: "destroy" as const }); diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts index 936d9ec013..580cde58c8 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/config.ts @@ -10,7 +10,7 @@ import type { RawAccess } from "@/db/config"; import type { AgentOsActorState, AgentOsActorVars } from "./types"; const zFunction = < - T extends (...args: any[]) => any = (...args: unknown[]) => unknown, + T extends (...args: never[]) => unknown = (...args: unknown[]) => unknown, >() => z.custom((val) => typeof val === "function", { message: "Expected a function", @@ -24,7 +24,6 @@ export const agentOsActorConfigSchema = z .object({ options: AgentOsOptionsSchema.optional(), createOptions: zFunction().optional(), - destroyOptions: zFunction().optional(), preview: z .object({ defaultExpiresInSeconds: z.number().positive().default(3600), @@ -44,19 +43,11 @@ export const agentOsActorConfigSchema = z message: "agentOs config must define exactly one of 'options' or 'createOptions'", }, - ) - .refine( - (data) => - !(data.options !== undefined && data.destroyOptions !== undefined), - { - message: - "agentOs config: 'destroyOptions' is only valid with 'createOptions', not 'options'", - }, ); // --- Typed config types (generic callbacks overlaid on the Zod schema) --- -export type AgentOsActorContext = ActorContext< +export type AgentOsContext = ActorContext< AgentOsActorState, TConnParams, undefined, @@ -76,12 +67,12 @@ interface AgentOsActorConfigCallbacks { params: TConnParams, ) => void | Promise; onSessionEvent?: ( - c: AgentOsActorContext, + c: AgentOsContext, sessionId: string, event: JsonRpcNotification, ) => void | Promise; onPermissionRequest?: ( - c: AgentOsActorContext, + c: AgentOsContext, sessionId: string, request: PermissionRequest, ) => void | Promise; @@ -97,29 +88,15 @@ type AgentOsActorOptionsConfig = * filesystem mounts, or per-instance configuration. */ options: AgentOsOptions; createOptions?: never; - destroyOptions?: never; } | { options?: never; /** Factory called lazily on first VM access. Receives the actor * context so options can vary per instance (e.g., dedicated - * sandboxes). Mutually exclusive with `options`. May be async. - * - * The callback should read `c.state.sandboxId` and pass it to - * `SandboxAgent.start()` to reconnect to an existing sandbox - * after sleep/wake. After starting, persist the ID back: - * `c.state.sandboxId = sandbox.sandboxId`. */ + * sandboxes). Mutually exclusive with `options`. May be async. */ createOptions: ( - c: AgentOsActorContext, + c: AgentOsContext, ) => AgentOsOptions | Promise; - /** Called on actor destroy to clean up external resources - * provisioned by `createOptions` (e.g., destroying a sandbox - * container). Receives the actor context with `c.state.sandboxId` - * still set so the callback can resolve the provider and destroy - * the sandbox. */ - destroyOptions?: ( - c: AgentOsActorContext, - ) => void | Promise; }; // Parsed config (after Zod defaults/transforms applied). @@ -127,7 +104,6 @@ export type AgentOsActorConfig = Omit< z.infer, | "options" | "createOptions" - | "destroyOptions" | "onBeforeConnect" | "onSessionEvent" | "onPermissionRequest" @@ -140,7 +116,6 @@ export type AgentOsActorConfigInput = Omit< z.input, | "options" | "createOptions" - | "destroyOptions" | "onBeforeConnect" | "onSessionEvent" | "onPermissionRequest" diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts index 398388644c..ddfa7e6e6b 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/index.ts @@ -31,13 +31,11 @@ export { } from "./actor/session"; // Shell actions export { buildShellActions } from "./actor/shell"; -// User-facing alias for the createOptions callback context parameter. -export type { AgentOsActorContext as AgentOsCreateContext } from "./config"; // Config schema and types export { type AgentOsActorConfig, type AgentOsActorConfigInput, - type AgentOsActorContext, + type AgentOsContext, agentOsActorConfigSchema, } from "./config"; // SQLite-backed VFS diff --git a/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts b/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts index 035661fd54..8dd0c47486 100644 --- a/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts +++ b/rivetkit-typescript/packages/rivetkit/src/agent-os/types.ts @@ -27,7 +27,7 @@ export interface AgentOsActorVars { agentOs: AgentOs | null; /** In-flight VM boot promise used to prevent concurrent ensureVm calls from * creating duplicate VMs. Reset on sleep/wake and cleared on boot failure. */ - vmBootPromise: Promise | null; + vmBootGuard: Promise | null; activeSessionIds: Set; activeProcesses: Set; activeHooks: Set>; diff --git a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts index 06b62f8b70..d527b34e99 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts @@ -649,35 +649,4 @@ describe("agentOs config exclusive union types", () => { }, }; }); - - it("options config forbids destroyOptions", () => { - // @ts-expect-error destroyOptions is only allowed with createOptions - const config: AgentOsActorConfigInput = { - options: { software: [] }, - destroyOptions: async () => {}, - }; - }); - - it("createOptions config accepts destroyOptions", () => { - const config: AgentOsActorConfigInput = { - createOptions: async (c) => ({ software: [] }), - destroyOptions: async (c) => { - expectTypeOf(c.state.sandboxId).toEqualTypeOf(); - }, - }; - expectTypeOf(config).toMatchTypeOf(); - }); - - it("agentOs() accepts createOptions with destroyOptions", () => { - agentOs({ - createOptions: async (c) => { - c.state.sandboxId = "docker/test"; - return { software: [] }; - }, - destroyOptions: async (c) => { - // Should have access to sandboxId for cleanup. - const _id: string | null = c.state.sandboxId; - }, - }); - }); }); diff --git a/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts b/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts deleted file mode 100644 index dbea90576b..0000000000 --- a/rivetkit-typescript/packages/rivetkit/tests/agent-os-config-validation.test.ts +++ /dev/null @@ -1,170 +0,0 @@ -import common from "@rivet-dev/agent-os-common"; -import { describe, expect, test } from "vitest"; -import { agentOsActorConfigSchema } from "@/agent-os/config"; -import { agentOs } from "@/agent-os/index"; -import { setup } from "@/mod"; -import { setupTest } from "@/test/mod"; - -describe("agentOs config validation", () => { - // --- Zod schema XOR enforcement --- - - test("accepts config with only options", () => { - const result = agentOsActorConfigSchema.safeParse({ - options: { software: [] }, - }); - expect(result.success).toBe(true); - }); - - test("accepts config with only createOptions", () => { - const result = agentOsActorConfigSchema.safeParse({ - createOptions: () => ({ software: [] }), - }); - expect(result.success).toBe(true); - }); - - test("rejects config with both options and createOptions", () => { - const result = agentOsActorConfigSchema.safeParse({ - options: { software: [] }, - createOptions: () => ({ software: [] }), - }); - expect(result.success).toBe(false); - if (!result.success) { - expect(result.error.message).toContain( - "exactly one of 'options' or 'createOptions'", - ); - } - }); - - test("rejects config with neither options nor createOptions", () => { - const result = agentOsActorConfigSchema.safeParse({}); - expect(result.success).toBe(false); - if (!result.success) { - expect(result.error.message).toContain( - "exactly one of 'options' or 'createOptions'", - ); - } - }); - - test("rejects config with createOptions set to a non-function", () => { - const result = agentOsActorConfigSchema.safeParse({ - createOptions: "not-a-function", - }); - expect(result.success).toBe(false); - }); - - test("accepts config with createOptions and destroyOptions", () => { - const result = agentOsActorConfigSchema.safeParse({ - createOptions: () => ({ software: [] }), - destroyOptions: async () => {}, - }); - expect(result.success).toBe(true); - }); - - test("rejects config with destroyOptions set to a non-function", () => { - const result = agentOsActorConfigSchema.safeParse({ - createOptions: () => ({ software: [] }), - destroyOptions: "not-a-function", - }); - expect(result.success).toBe(false); - }); - - test("rejects config with destroyOptions paired with options", () => { - const result = agentOsActorConfigSchema.safeParse({ - options: { software: [] }, - destroyOptions: async () => {}, - }); - expect(result.success).toBe(false); - if (!result.success) { - expect(result.error.message).toContain( - "'destroyOptions' is only valid with 'createOptions'", - ); - } - }); - - // --- Runtime behavior --- - - test("createOptions factory boots a working VM", async (c) => { - const vm = agentOs({ - createOptions: async () => ({ software: [common] }), - }); - const registry = setup({ use: { vm } }); - const { client } = await setupTest(c, registry); - const actor = (client as any).vm.getOrCreate([ - `config-test-${crypto.randomUUID()}`, - ]); - - await actor.writeFile("/home/user/test.txt", "config-validation"); - const data = await actor.readFile("/home/user/test.txt"); - expect(new TextDecoder().decode(data)).toBe("config-validation"); - }, 60_000); - - test("createOptions receives actor context", async (c) => { - let receivedContext = false; - - const vm = agentOs({ - createOptions: async (ctx) => { - // The context should have log and vars available. - receivedContext = - ctx.log !== undefined && ctx.vars !== undefined; - return { software: [common] }; - }, - }); - const registry = setup({ use: { vm } }); - const { client } = await setupTest(c, registry); - const actor = (client as any).vm.getOrCreate([ - `ctx-test-${crypto.randomUUID()}`, - ]); - - // Trigger ensureVm by calling any action. - await actor.exec("echo context-check"); - expect(receivedContext).toBe(true); - }, 60_000); - - test("createOptions can persist sandboxId in state", async (c) => { - const fakeSandboxId = `test/${crypto.randomUUID()}`; - - const vm = agentOs({ - createOptions: async (ctx) => { - // On first boot, sandboxId should be null. - if (!ctx.state.sandboxId) { - ctx.state.sandboxId = fakeSandboxId; - } - return { software: [common] }; - }, - }); - const registry = setup({ use: { vm } }); - const { client } = await setupTest(c, registry); - const actor = (client as any).vm.getOrCreate([ - `sandbox-id-test-${crypto.randomUUID()}`, - ]); - - // Trigger ensureVm so createOptions runs. - await actor.exec("echo sandbox-id-check"); - - // The actor should have set the sandboxId. If we call another - // action, createOptions won't re-run (VM is already booted), but - // the state should persist. - const result = await actor.exec("echo done"); - expect(result.exitCode).toBe(0); - }, 60_000); - - test("destroyOptions is accepted as a valid config callback", () => { - // Verify destroyOptions is accepted by agentOs() at runtime (not - // just at the type level). If the Zod schema rejected it, this - // would throw during parse. - const vm = agentOs({ - createOptions: async (ctx) => { - ctx.state.sandboxId = "test/destroy-check"; - return { software: [common] }; - }, - destroyOptions: async (ctx) => { - // This callback would be called during onDestroy. - // We just verify it is accepted without error. - const _id = ctx.state.sandboxId; - }, - }); - - // The factory should return a valid actor definition. - expect(vm).toBeDefined(); - }); -}); diff --git a/website/src/content/docs/agent-os/configuration.mdx b/website/src/content/docs/agent-os/configuration.mdx index f71395d01c..654ecb4e43 100644 --- a/website/src/content/docs/agent-os/configuration.mdx +++ b/website/src/content/docs/agent-os/configuration.mdx @@ -55,54 +55,24 @@ export const registry = setup({ use: { vm } }); registry.start(); ``` -## Per-instance options +## Dynamic options Use `createOptions` when each actor instance needs its own VM configuration. The callback receives the actor context and runs lazily on first VM access. It may be async. -This is required for [sandbox mounting](/docs/agent-os/sandbox) where each instance needs a dedicated sandbox container, and useful for any scenario where options depend on the actor instance (e.g., per-tenant mounts or credentials). +This is useful for [sandbox mounting](/docs/agent-os/sandbox) where each instance needs a dedicated container, or for any scenario where options depend on the actor instance. `options` and `createOptions` are **mutually exclusive**. Providing both causes a validation error. -### Sandbox reuse across sleep/wake - -The actor state includes `sandboxId` which you should use to reconnect to the same sandbox after a sleep/wake cycle. Without this, every wake would provision a new sandbox container. - -Read `c.state.sandboxId` and pass it to `SandboxAgent.start()`. After starting, persist the ID back so future wakes reuse the same container. - -Provide `destroyOptions` to clean up external resources (e.g., destroy the sandbox container) when the actor is destroyed. This callback only applies when using `createOptions`. - ```ts @nocheck import { agentOs } from "rivetkit/agent-os"; import { setup } from "rivetkit"; import common from "@rivet-dev/agent-os-common"; -import { SandboxAgent } from "sandbox-agent"; -import { docker } from "sandbox-agent/docker"; -import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; const vm = agentOs({ - createOptions: async (c) => { - // Reconnect to an existing sandbox after sleep, or provision a new one. - const sandbox = await SandboxAgent.start({ - sandbox: docker(), - sandboxId: c.state.sandboxId ?? undefined, - }); - - // Persist the sandbox ID so future wakes reuse the same container. - c.state.sandboxId = sandbox.sandboxId; - - return { - software: [common], - mounts: [{ path: "/sandbox", driver: createSandboxFs({ client: sandbox }) }], - toolKits: [createSandboxToolkit({ client: sandbox })], - }; - }, - // Clean up the sandbox container when the actor is destroyed. - destroyOptions: async (c) => { - if (c.state.sandboxId) { - const provider = docker(); - await provider.destroy(c.state.sandboxId); - } - }, + createOptions: async (c) => ({ + software: [common], + env: { TENANT_ID: c.key[0] }, + }), }); export const registry = setup({ use: { vm } }); diff --git a/website/src/content/docs/agent-os/index.mdx b/website/src/content/docs/agent-os/index.mdx index 604b350522..03cc26c34c 100644 --- a/website/src/content/docs/agent-os/index.mdx +++ b/website/src/content/docs/agent-os/index.mdx @@ -390,22 +390,16 @@ agentOS uses a hybrid model: agents run in a lightweight VM by default and spin import { agentOs } from "rivetkit/agent-os"; import { setup } from "rivetkit"; import common from "@rivet-dev/agent-os-common"; -import { SandboxAgent } from "sandbox-agent"; import { docker } from "sandbox-agent/docker"; -import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; const vm = agentOs({ createOptions: async (c) => { - const sandbox = await SandboxAgent.start({ - sandbox: docker(), - sandboxId: c.state.sandboxId ?? undefined, - }); - c.state.sandboxId = sandbox.sandboxId; + const { fs, toolkit } = await c.getSandbox({ provider: docker() }); return { software: [common], - mounts: [{ path: "/sandbox", driver: createSandboxFs({ client: sandbox }) }], - toolKits: [createSandboxToolkit({ client: sandbox })], + mounts: [{ path: "/sandbox", driver: fs }], + toolKits: [toolkit], }; }, }); diff --git a/website/src/content/docs/agent-os/sandbox.mdx b/website/src/content/docs/agent-os/sandbox.mdx index cb89cd9a09..8b34a377ba 100644 --- a/website/src/content/docs/agent-os/sandbox.mdx +++ b/website/src/content/docs/agent-os/sandbox.mdx @@ -26,57 +26,30 @@ See [agentOS vs Sandbox](/docs/agent-os/versus-sandbox) for a detailed compariso ## Getting started -The `@rivet-dev/agent-os-sandbox` package integrates through two mechanisms: +Use `c.getSandbox()` inside `createOptions` to provision a sandbox per actor instance. It returns a filesystem driver and a toolkit you can pass directly to the VM options. -- **Filesystem mount**: Projects the sandbox into the VM as a native directory, like mounting a hard drive on your own machine. Read and write files through the mount directly. -- **Toolkit**: Exposes sandbox process management as [host tools](/docs/agent-os/tools). Execute commands on the sandbox from within the VM. - -Both are powered by [Sandbox Agent](https://sandboxagent.dev), so you can swap providers without changing agent code. +The sandbox ID is persisted internally across sleep/wake cycles, so the actor reconnects to the same container automatically. The sandbox is also auto-destroyed when the actor is destroyed. ```bash -npm install @rivet-dev/agent-os-sandbox sandbox-agent +npm install sandbox-agent ``` ```ts @nocheck import { agentOs } from "rivetkit/agent-os"; import { setup } from "rivetkit"; import common from "@rivet-dev/agent-os-common"; -import { SandboxAgent } from "sandbox-agent"; import { docker } from "sandbox-agent/docker"; -import { createSandboxFs, createSandboxToolkit } from "@rivet-dev/agent-os-sandbox"; -// Each actor instance gets its own Docker sandbox via createOptions. -// The sandbox ID is persisted in state so the actor reconnects to the -// same container after a sleep/wake cycle instead of creating a new one. const vm = agentOs({ createOptions: async (c) => { - // Reconnect to an existing sandbox after sleep, or provision a new one. - const sandbox = await SandboxAgent.start({ - sandbox: docker(), - sandboxId: c.state.sandboxId ?? undefined, - }); - - // Persist the sandbox ID so future wakes reuse the same container. - c.state.sandboxId = sandbox.sandboxId; + const { fs, toolkit } = await c.getSandbox({ provider: docker() }); return { software: [common], - mounts: [ - { - path: "/sandbox", - driver: createSandboxFs({ client: sandbox }), - }, - ], - toolKits: [createSandboxToolkit({ client: sandbox })], + mounts: [{ path: "/sandbox", driver: fs }], + toolKits: [toolkit], }; }, - // Clean up the sandbox container when the actor is destroyed. - destroyOptions: async (c) => { - if (c.state.sandboxId) { - const provider = docker(); - await provider.destroy(c.state.sandboxId); - } - }, }); export const registry = setup({ use: { vm } });