From 5d6b647d0326bb325965c94624c766663a92310c Mon Sep 17 00:00:00 2001 From: rmanjunath18 <242545101+rmanjunath18@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:42:28 +0530 Subject: [PATCH] feat(agents): snapshot, list snapshots and fork an agent Agents and AsyncAgents gain create_snapshot, list_snapshots and fork, and the agent handles gain snapshot, snapshots and fork. Fork returns the new agent, Provisioning until ready. HTTP 425 now raises TooEarlyError. Adds docs, a changeset and an end-to-end example. Regenerates the aiagent types from the updated spec: snapshots now carry a required trigger and can be filesystem-only; audit events add program and network records, device and cloud.org, and new actor type captions. Keep trigger optional for older servers, accept snapshot_type on create, update test data to the new shapes, and stop the audit examples reading api on events that have none. AIPLATFORM-1564 --- .changeset/agent-snapshot-fork.md | 11 ++ .changeset/audit-ocsf-events.md | 2 +- .changeset/sandbox-replicas-removed.md | 2 +- docs/api-inventory.md | 36 ++++- docs/api-reference.md | 15 +- docs/example-coverage.md | 6 +- examples/README.md | 17 +- examples/agent_snapshot_fork_rollback.py | 81 ++++++++++ specs/aiagent.yaml | 41 +++-- src/neevai/__init__.py | 2 + src/neevai/errors.py | 8 + src/neevai/generated/aiagent.py | 17 +- src/neevai/handles/agent.py | 121 ++++++++++++++ src/neevai/resources/agents.py | 118 ++++++++++++++ src/neevai/types.py | 29 +++- tests/conftest.py | 1 + tests/test_agent_snapshots.py | 198 +++++++++++++++++++++++ tests/test_audit.py | 7 +- tests/test_sandboxes.py | 3 +- 19 files changed, 667 insertions(+), 48 deletions(-) create mode 100644 .changeset/agent-snapshot-fork.md create mode 100644 examples/agent_snapshot_fork_rollback.py create mode 100644 tests/test_agent_snapshots.py diff --git a/.changeset/agent-snapshot-fork.md b/.changeset/agent-snapshot-fork.md new file mode 100644 index 0000000..e614fc3 --- /dev/null +++ b/.changeset/agent-snapshot-fork.md @@ -0,0 +1,11 @@ +--- +"neevai": minor +--- + +Agents can be snapshotted and forked. `client.agents.create_snapshot(id, params=None)` and `list_snapshots(id, page=None, limit=None)` snapshot the sandbox backing an agent and list those snapshots, with the same params and `Snapshot` shape as the sandbox calls. `client.agents.fork(id, name)` forks an agent into a new agent with the source's template and configuration; it returns the new `Agent`, which is `Provisioning` until `wait_until_ready()`, and the source keeps running. The `Agent` and `AsyncAgent` handles gain `snapshot(params=None, *, wait_until_ready=False, ...)`, `snapshots()` and `fork(name)`. + +`CreateSnapshotParams` takes `snapshot_type`: `"filesystem"` captures /workspace only, without process memory, so it can be restored onto any sandbox image; `"full"` captures both. + +A new `TooEarlyError` is raised for HTTP 425, such as calling an agent that is still provisioning. It was a plain `APIError` before. + +The audit trail now also records every program that runs in a sandbox, unless its `execution_trace` is turned off: a Process Activity event (`class_uid` 1007) for a program starting or ending, with `process` and `exit_code`, and a Network Activity event (4001) for a connection it opened. These events have no `api`, so check `e.api` before reading `e.api.operation`. Every event carries `device` (the sandbox) and `cloud.org`. `actor.user.type` is now `User`, `API Key`, `System` or `Unknown`; the `actor_type` filter still takes `"user"`, `"api_key"` and `"system"`. diff --git a/.changeset/audit-ocsf-events.md b/.changeset/audit-ocsf-events.md index c96bb2d..730a4dc 100644 --- a/.changeset/audit-ocsf-events.md +++ b/.changeset/audit-ocsf-events.md @@ -2,7 +2,7 @@ "neevai": minor --- -The audit trail returns events instead of records. `sandbox.audit()`, `agent.audit()` and the matching `sandboxes` / `agents` calls now return `AuditTrail.data: list[AuditEvent]`, each event saying which call ran (`api.operation`), who made it (`actor.user`, with the API key in `credential_uid`), how it ended (`status`), what it touched (`resources`) and when (`time`, in milliseconds since the Unix epoch). `AuditTrail.records`, `AuditTrail.sandbox_id` and the `AuditRecord` model are removed. +The audit trail returns events instead of records. `sandbox.audit()`, `agent.audit()` and the matching `sandboxes` / `agents` calls now return `AuditTrail.data: list[AuditEvent]`, each event saying which call ran (`api.operation`), who made it (`actor.user`, with the API key in `programmatic_credentials`), how it ended (`status`), what it touched (`resources`) and when (`time`, in milliseconds since the Unix epoch). `AuditTrail.records`, `AuditTrail.sandbox_id` and the `AuditRecord` model are removed. The audit calls take new filters: `operation` (one name or a list), `status` (`"success"` or `"failure"`), `actor_type` (`"user"`, `"api_key"`, `"system"`, or a list), `actor_uid` and `credential_uid`. diff --git a/.changeset/sandbox-replicas-removed.md b/.changeset/sandbox-replicas-removed.md index 772b18f..5970110 100644 --- a/.changeset/sandbox-replicas-removed.md +++ b/.changeset/sandbox-replicas-removed.md @@ -4,4 +4,4 @@ `Sandbox.replicas` and `AsyncSandbox.replicas` are removed; read `phase` instead, which is `"Paused"` while a sandbox is paused. The field is deprecated in the API and is now optional on the generated `Sandbox` model, so a sandbox response without it parses. -`snapshot_type` adds `"filesystem"` for a files-only snapshot that cannot be restored; a snapshot of that type no longer fails response validation. +`snapshot_type` adds `"filesystem"` for a files-only snapshot that can be restored onto any sandbox image; a snapshot of that type no longer fails response validation. diff --git a/docs/api-inventory.md b/docs/api-inventory.md index e963d05..d3d69b5 100644 --- a/docs/api-inventory.md +++ b/docs/api-inventory.md @@ -94,7 +94,7 @@ Everything in `neevai.__all__`: | `AuditEvent` / `AuditActor` / `AuditActorUser` / `AuditAPI` / `AuditResource` / `AuditProcess` / `AuditFile` / `AuditDstEndpoint` / `AuditDevice` | model | generated | | `SandboxEgressRule` | model | generated | | `SandboxEgressRules` | model | generated | -| `NeevAIError`, `APIConnectionError`, `APITimeoutError`, `APIError`, `BadRequestError`, `AuthenticationError`, `PermissionDeniedError`, `NotFoundError`, `ConflictError`, `PreconditionFailedError`, `RateLimitError`, `DeadlineExceededError`, `InternalServerError`, `ServiceUnavailableError` | exceptions | `errors.py` | +| `NeevAIError`, `APIConnectionError`, `APITimeoutError`, `APIError`, `BadRequestError`, `AuthenticationError`, `PermissionDeniedError`, `NotFoundError`, `ConflictError`, `PreconditionFailedError`, `TooEarlyError`, `RateLimitError`, `DeadlineExceededError`, `InternalServerError`, `ServiceUnavailableError` | exceptions | `errors.py` | Types exported from `neevai.types.__all__`: @@ -674,6 +674,27 @@ Pause and resume return updated `Agent` handles. Delete returns `None` (HTTP 204 progress. `rollback` rolls the agent's backing sandbox back in place to a snapshot from the same project. Both return the updated `Agent`. +### `client.agents.create_snapshot(id, params=None, ...)` / `client.agents.list_snapshots(id, page=None, limit=None, ...)` + +Snapshot the sandbox backing an agent, and list those snapshots. Same params and +`Snapshot` shape as `client.sandboxes.create_snapshot` / `.list_snapshots`; the +snapshot starts `Pending`. Wait on it with `client.sandboxes.wait_for_snapshot`, or +use `agent.snapshot(wait_until_ready=True)`. + +### `client.agents.fork(id, name, ...)` + +Forks an agent into a new agent with the source's template and configuration and its +own gateway token. Returns the new `Agent`, which is `Provisioning` until it is ready; +call `wait_until_ready()` on it before use. The source agent keeps running. + +```python +forked = client.agents.fork(agent.id, "web-fork") +forked.wait_until_ready() +``` + +A paused or still-provisioning agent raises `ConflictError` (409) or +`TooEarlyError` (425). + ### Agent preview ports `client.agents.expose_port(id, port, org_id=None, project_id=None, *, slug=None)`, @@ -783,6 +804,11 @@ Resolves the backing sandbox as a `Sandbox` handle via `Agents.get_sandbox()`. Handle wrappers for `client.agents.keepalive` / `.rollback` / `.audit`; the first two update the handle in place and return `self`. +### `agent.snapshot(params=None, *, wait_until_ready=False, timeout_ms=..., poll_interval_ms=...)` / `agent.snapshots()` / `agent.fork(name)` + +Same as `sandbox.snapshot()` / `.snapshots()`, for the sandbox backing this agent. +`fork` returns a new `Agent` handle for the fork; this handle is unchanged. + ### `agent.expose_port(port, *, slug=None)` / `agent.list_ports()` / `agent.revoke_port(port)` / `agent.get_url(port, ..., *, slug=None)` Same signatures as the sandbox handle's preview-port methods. @@ -1929,7 +1955,7 @@ for `egress_add` / `egress_remove`. | `api.operation` | `str` | API Activity only (e.g. `pty_command`, `ssh`, `exec`, `process.start`, `fs.read`) | | `status` | `str` | yes (`Success` or `Failure`) | | `status_code` | `str \| None` | no (`ok`, `permission_denied`, `internal`) | -| `actor.user.type` | `User` \| `API Key` \| `System` \| `Unknown` | yes (`Unknown` on program and connection events) | +| `actor.user.type` | `User` \| `API Key` \| `System` \| `Unknown` | yes (`Unknown` only for a program the sandbox ran on its own, such as a background process) | | `actor.user.uid` | `str` | yes (the user's or API key's id) | | `actor.user.name` | `str \| None` | no | | `actor.user.programmatic_credentials` | `list \| None` | no (the API key used, as `[0].uid`; absent for a user) | @@ -1938,7 +1964,7 @@ for `egress_add` / `egress_remove`. | `exit_code` | `int \| None` | no (on a Terminate, unless a signal stopped the program) | | `status_detail` | `str \| None` | no (`signal ` for a program a signal stopped; how a code-interpreter run ended) | | `dst_endpoint` | `AuditDstEndpoint \| None` | Network Activity only (`ip` and `port`, or `name` as `unix:`) | -| `device.uid` | `str \| None` | no (the sandbox id) | +| `device.uid` | `str` | yes (the sandbox id) | | `actor.session.uid` | `str \| None` | no (the console session) | | `resources` | `list[AuditResource] \| None` | API Activity only (the sandbox, then any file, process, program or terminal the call touched; `data.pid` on the process or terminal it started, which a program's pid or an ancestor's equals) | | `duration` | `int \| None` | no (milliseconds) | @@ -2110,6 +2136,7 @@ All SDK errors inherit from `NeevAIError`. Import from `neevai` or `neevai.error | `NotFoundError` | HTTP 404 | | `ConflictError` | HTTP 409 | | `PreconditionFailedError` | HTTP 412 | +| `TooEarlyError` | HTTP 425 (not ready yet, such as an agent still provisioning) | | `RateLimitError` | HTTP 429 | | `DeadlineExceededError` | HTTP 504 | | `InternalServerError` | HTTP 5xx | @@ -2270,6 +2297,8 @@ Compact reviewer index. Each symbol should also appear in | `Agents.delete` | method | `None` | | `Agents.get_sandbox` | method | `Sandbox` | | `Agents.keepalive` / `.rollback` | methods | `Agent` | +| `Agents.create_snapshot` / `.list_snapshots` | methods | `Snapshot` / `list[Snapshot]` | +| `Agents.fork` | method | `Agent` (the new agent) | | `Agents.expose_port` / `.list_ports` / `.revoke_port` / `.get_port_url` | methods | `SandboxPort` / `list[SandboxPort]` / `None` / `str` | | `Agents.audit` | method | `AuditTrail` | | `AsyncAgents.*` | methods | Same as sync with `await` | @@ -2293,6 +2322,7 @@ Compact reviewer index. Each symbol should also appear in | `Agent.sandbox` | method | Backing `Sandbox` handle | | `Agent.update` / `.pause` / `.resume` / `.delete` | methods | Scope-threaded | | `Agent.keepalive` / `.rollback` / `.audit` | methods | Scope-threaded | +| `Agent.snapshot` / `.snapshots` / `.fork` | methods | `Snapshot` / `list[Snapshot]` / new `Agent` | | `Agent.expose_port` / `.list_ports` / `.revoke_port` / `.get_url` | methods | Preview ports | | `Agent.idle_timeout_seconds` | property | `int \| None` | | `AsyncAgent.*` | mirror | Same surface with `await` | diff --git a/docs/api-reference.md b/docs/api-reference.md index 04e46f6..b704f64 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -55,7 +55,7 @@ A sandbox is an isolated machine for your code. `client.sandboxes` and `client.t | `audit(id, *, from_=None, to=None, cursor=None, limit=None, ...)` | `AuditTrail` | One page of the audit trail, newest first: calls, programs and connections. See [Audit trail](#audit-trail). | | `expose_port(id, port, ..., *, slug=None)` / `list_ports(id)` / `revoke_port(id, port)` | `SandboxPort` / `list[SandboxPort]` / `None` | Preview ports. A new `slug` on an exposed port rotates its URL. | | `get_port_url(id, port, wait_until_ready=True, ..., *, slug=None)` | `str` | Exposes a port and returns its preview URL once routable. | -| `create_snapshot(id, params=None, ...)` | `Snapshot` | Creates a filesystem snapshot (returns immediately with status Pending). | +| `create_snapshot(id, params=None, ...)` | `Snapshot` | Creates a snapshot of the sandbox: process memory and filesystem by default, or `/workspace` only with `snapshot_type="filesystem"` (returns immediately with status Pending). | | `list_snapshots(id, page=None, limit=None, ...)` | `list[Snapshot]` | Lists snapshots for a sandbox (unwraps pagination items). | | `get_snapshot(snapshot_id, ...)` | `Snapshot` | Fetches snapshot metadata by project-scoped ID. | | `delete_snapshot(snapshot_id, ...)` | `None` | Deletes a snapshot permanently. | @@ -300,6 +300,8 @@ An agent is a catalogue template (such as `"claude-code"`) running on its own sa | `get(id, org_id=None, project_id=None)` | `Agent` | Fetches the current record for an agent by ID or name. | | `update(id, params, org_id=None, project_id=None, *, allow_internet=None, allow_egress=None)` | `Agent` | In-place update of egress, cpu/memory (`resources`; defaults & bounds in [Agent resources](./api-inventory.md#agent-resources)) and `idle_timeout_seconds`. `allow_internet` / `allow_egress` are the same egress convenience as `create`; `egress_add` / `egress_remove` edit the allow-list in place. Rejects `{}` locally, naming every field. | | `keepalive(id, ...)` / `rollback(id, snapshot_id, ...)` | `Agent` | Resets the idle timer / rolls the backing sandbox back to a snapshot. | +| `create_snapshot(id, params=None, ...)` / `list_snapshots(id, page=None, limit=None, ...)` | `Snapshot` / `list[Snapshot]` | Snapshots the backing sandbox (starts `Pending`) / lists its snapshots. | +| `fork(id, name, ...)` | `Agent` | Forks into a new agent, `Provisioning` until ready; the source keeps running. | | `expose_port` / `list_ports` / `revoke_port` / `get_port_url` | as on `client.sandboxes` | Agent preview ports, same signatures. | | `audit(id, *, from_=None, to=None, cursor=None, limit=None, ...)` | `AuditTrail` | The agent's audit trail: the same events and filters as the [sandbox one](#audit-trail). An agent has no `execution_trace` setting. | | `pause(id, org_id=None, project_id=None)` | `Agent` | Pauses the agent and its backing sandbox. | @@ -324,6 +326,7 @@ Returned by `client.agents.create()`, `get()`, `list().items`, etc. | `wait_until_ready(timeout_ms=120000, poll_interval_ms=2000, on_poll=None)` | method — polls until `Ready`; fails fast on `Failed` / `Paused` | | `pause()` / `resume()` / `delete()` | methods | | `keepalive()` / `rollback(snapshot_id)` / `audit(...)` | methods | +| `snapshot(params=None, *, wait_until_ready=False, ...)` / `snapshots()` / `fork(name)` | methods — `fork` returns a new `Agent` handle | | `expose_port(port, *, slug=None)` / `list_ports()` / `revoke_port(port)` / `get_url(port, ..., *, slug=None)` | methods — preview ports | | `sandbox()` | method — backing `Sandbox` handle for runtime `exec`, `files`, and `processes` | | `to_json()` | method | @@ -420,6 +423,7 @@ Every failure is a `NeevAIError`. API failures are `APIError` subclasses with `s | `NotFoundError` | 404: no such resource | | `ConflictError` | 409: conflicting state | | `PreconditionFailedError` | 412: the sandbox is not in a state that allows the call | +| `TooEarlyError` | 425: not ready yet, such as an agent still provisioning; retry later | | `RateLimitError` | 429: too many requests | | `DeadlineExceededError` | 504: the call ran past its deadline | | `InternalServerError` | 5xx (`ServiceUnavailableError` for 503) | @@ -470,7 +474,10 @@ Minimal one-liners for each public API. Runnable examples link to repo paths. | `client.agents.pause(id)` / `.resume(id)` | `agent = client.agents.pause(agent_id)` / `client.agents.resume(agent_id)` | `agent = await client.agents.pause(agent_id)` / `await client.agents.resume(agent_id)` | [create_agent.py](../examples/create_agent.py) (pause via handle) | | `client.agents.delete(id)` | `client.agents.delete(agent_id)` | `await client.agents.delete(agent_id)` | [create_agent.py](../examples/create_agent.py) | | `client.agents.keepalive(id)` | `client.agents.keepalive(agent_id)` | `await client.agents.keepalive(agent_id)` | [agent_ports_audit.py](../examples/agent_ports_audit.py) | -| `client.agents.rollback(id, snapshot_id)` | `client.agents.rollback(agent_id, snap_id)` | `await client.agents.rollback(agent_id, snap_id)` | — | +| `client.agents.rollback(id, snapshot_id)` | `client.agents.rollback(agent_id, snap_id)` | `await client.agents.rollback(agent_id, snap_id)` | [agent_snapshot_fork_rollback.py](../examples/agent_snapshot_fork_rollback.py) (via `agent.rollback`) | +| `client.agents.create_snapshot(id, ...)` | `snap = client.agents.create_snapshot(agent_id, {"name": "demo"})` | `snap = await client.agents.create_snapshot(agent_id, {"name": "demo"})` | [agent_snapshot_fork_rollback.py](../examples/agent_snapshot_fork_rollback.py) (via `agent.snapshot`) | +| `client.agents.list_snapshots(id)` | `snaps = client.agents.list_snapshots(agent_id)` | `snaps = await client.agents.list_snapshots(agent_id)` | [agent_snapshot_fork_rollback.py](../examples/agent_snapshot_fork_rollback.py) (via `agent.snapshots`) | +| `client.agents.fork(id, name)` | `fork = client.agents.fork(agent_id, "fork-name")` | `fork = await client.agents.fork(agent_id, "fork-name")` | [agent_snapshot_fork_rollback.py](../examples/agent_snapshot_fork_rollback.py) (via `agent.fork`) | | `client.agents.expose_port(id, port, ..., *, slug)` | `p = client.agents.expose_port(agent_id, 3000)` | `p = await client.agents.expose_port(agent_id, 3000)` | [agent_ports_audit.py](../examples/agent_ports_audit.py) | | `client.agents.audit(id, ...)` | `trail = client.agents.audit(agent_id, limit=50)` | `trail = await client.agents.audit(agent_id, limit=50)` | [agent_ports_audit.py](../examples/agent_ports_audit.py) | | `client.agent_templates.list(...)` | `page = client.agent_templates.list()` | `page = await client.agent_templates.list()` | [create_agent.py](../examples/create_agent.py) | @@ -509,7 +516,9 @@ Minimal one-liners for each public API. Runnable examples link to repo paths. | `agent.sandbox()` | `sandbox = agent.sandbox()` — then `exec` / `files` / `processes` | `sandbox = await agent.sandbox()` | [create_agent.py](../examples/create_agent.py) | | `agent.update(params)` | `agent.update({"resources": {...}})` | `await agent.update({...})` | [create_agent.py](../examples/create_agent.py) | | `agent.pause()` / `.resume()` / `.delete()` | `agent.pause(); agent.resume(); agent.delete()` | `await agent.pause(); await agent.resume(); await agent.delete()` | [create_agent.py](../examples/create_agent.py) (pause/delete; resume via handle API) | -| `agent.keepalive()` / `.rollback(snapshot_id)` | `agent.keepalive()` | `await agent.keepalive()` | [agent_ports_audit.py](../examples/agent_ports_audit.py) | +| `agent.keepalive()` / `.rollback(snapshot_id)` | `agent.keepalive()` | `await agent.keepalive()` | [agent_ports_audit.py](../examples/agent_ports_audit.py), [agent_snapshot_fork_rollback.py](../examples/agent_snapshot_fork_rollback.py) (rollback) | +| `agent.snapshot(...)` / `.snapshots()` | `snap = agent.snapshot(wait_until_ready=True)` | `snap = await agent.snapshot(wait_until_ready=True)` | [agent_snapshot_fork_rollback.py](../examples/agent_snapshot_fork_rollback.py) | +| `agent.fork(name)` | `fork = agent.fork("fork-name"); fork.wait_until_ready()` | `fork = await agent.fork("fork-name")` | [agent_snapshot_fork_rollback.py](../examples/agent_snapshot_fork_rollback.py) | | `agent.expose_port(port, *, slug)` / `.list_ports()` / `.revoke_port(port)` / `.get_url(port)` | `p = agent.expose_port(3000)` | `p = await agent.expose_port(3000)` | [agent_ports_audit.py](../examples/agent_ports_audit.py) | | `agent.audit(...)` | `trail = agent.audit(limit=50)` | `trail = await agent.audit(limit=50)` | [agent_ports_audit.py](../examples/agent_ports_audit.py) | | `agent.idle_timeout_seconds` | `print(agent.idle_timeout_seconds)` | same | — | diff --git a/docs/example-coverage.md b/docs/example-coverage.md index 2ec4899..0f9d8fd 100644 --- a/docs/example-coverage.md +++ b/docs/example-coverage.md @@ -51,7 +51,10 @@ examples themselves, see [`examples/README.md`](../examples/README.md). | `client.agents.pause` / `.resume` | — | `agent = client.agents.pause(agent_id)` / `client.agents.resume(agent_id)` | | `client.agents.delete` | `create_agent.py` | `client.agents.delete(agent_id)` | | `client.agents.keepalive` / `agent.keepalive` | `agent_ports_audit.py` | `agent.keepalive()` | -| `client.agents.rollback` / `agent.rollback` | — | `agent.rollback(snapshot_id)` | +| `client.agents.rollback` / `agent.rollback` | `agent_snapshot_fork_rollback.py` | `agent.rollback(snapshot_id)` | +| `client.agents.create_snapshot` / `agent.snapshot` | `agent_snapshot_fork_rollback.py` | `snap = agent.snapshot({"name": "before-change"}, wait_until_ready=True)` | +| `client.agents.list_snapshots` / `agent.snapshots` | `agent_snapshot_fork_rollback.py` | `snaps = agent.snapshots()` | +| `client.agents.fork` / `agent.fork` | `agent_snapshot_fork_rollback.py` | `fork = agent.fork("web-fork"); fork.wait_until_ready()` | | `agent.expose_port` / `.list_ports` / `.revoke_port` / `.get_url` | `agent_ports_audit.py` | `port = agent.expose_port(3000, slug=random_slug())` | | `client.agents.audit` / `agent.audit` | `agent_ports_audit.py` | `trail = agent.audit(limit=20)` | | `agent.idle_timeout_seconds` | `agent_ports_audit.py` | `print(agent.idle_timeout_seconds)` | @@ -118,6 +121,7 @@ examples themselves, see [`examples/README.md`](../examples/README.md). | Resize & egress edits | `examples/update_resize_egress.py` | Resize + re-scope egress in one update → `egress_add` / `egress_remove` | | Audit trail | `examples/audit_trail.py` | Run a script → read calls, programs and connections, each program tied to its call | | Agent ports & audit | `examples/agent_ports_audit.py` | Agent preview port with slug → rotate → `keepalive` → `audit` → revoke | +| Agent snapshot, fork & rollback | `examples/agent_snapshot_fork_rollback.py` | Write state → snapshot → modify → fork → check fork has the modified file → rollback → check original state → cleanup | | Streaming exec | `examples/streaming_exec.py` | `exec_stream` with progress output | | Processes lifecycle | `examples/processes.py` | Create → connect_url wait → Ready → runtime probe → start → follow → logs → list → kill → wait | | Process pool | `examples/process_pool.py` | Create → connect_url wait → Ready → runtime probe → parallel start → list → kill_all → wait | diff --git a/examples/README.md b/examples/README.md index 73c1e66..e5f9344 100644 --- a/examples/README.md +++ b/examples/README.md @@ -91,6 +91,7 @@ These need only your credentials. | [`create_agent.py`](create_agent.py) | Create an agent, use its sandbox, update, pause, delete | | [`agent_ports_audit.py`](agent_ports_audit.py) | An agent's preview URL, keepalive and audit trail | | [`agent_auto_resume.py`](agent_auto_resume.py) | Use a paused agent without calling `resume()`; reports whether it woke by itself | +| [`agent_snapshot_fork_rollback.py`](agent_snapshot_fork_rollback.py) | Snapshot an agent, fork it into a new agent, roll it back | ## Agent tools @@ -311,23 +312,29 @@ uv run python examples/agent_auto_resume.py ``` → `created agent …, waiting for it to be Ready...` → `agent is Ready; wrote auto-resume-marker.txt into its sandbox …` → `agent paused (status: Paused)` → `reading auto-resume-marker.txt from the paused agent, without calling resume()...` → `file read back intact after N.Ns` → `agent status now: Ready` → `RESULT: woke automatically` (or `the read was refused with 429 and the agent is still Paused (…)` → `calling resume() and reading the file again...` → `file read back intact after N.Ns` → … → `RESULT: resumed explicitly` where automatic wake-up is off) → `deleted agent …`. +**31. Agent snapshot, fork and rollback** +```bash +uv run python examples/agent_snapshot_fork_rollback.py +``` +→ `ready: …` → `snapshot … is Ready` → `agent snapshots: ['before-change']` → `forked: … (…)` → `fork has: modified state` → `after rollback, snapshot-demo (…) has: original state`. + ### Agent tools These call a model. They use `NEEV_API_KEY` for inference unless you set `NEEV_INFERENCE_API_KEY`, and the model is `gpt-oss-120b` unless you set `NEEV_MODEL`. -**31. Minimal agent loop** +**32. Minimal agent loop** ```bash uv run python examples/agent_patterns/minimal_agent.py ``` → `AI code-interpreter ( -> Neev sandbox)` → `task: Print every prime number below 50, …` → `[sandbox] creating (template=…)…` → for each step: `━━━━━ step N ━━━━━` → `[model] calling … with N messages…` → `[model] tokens: prompt=… completion=… total=…` → `[model] requested N tool call(s)` → `[tool] run_shell` and ` $ ` → ` ┌─ live output ───…` → the command's output as it runs → ` └─ exit 0 ───…` → `[tool] returned N bytes to the model` … → `✅ final answer:` → the model's answer → `[sandbox] deleting …`. -**32. Ready-made agent tools** +**33. Ready-made agent tools** ```bash uv run python examples/agent_patterns/agent_tools_loop.py ``` → `[sandbox] …` → `[tool] ({…}) -> {…}` (one line per tool call, as the model writes and runs `fib.py`) → `[answer] …` (the model's reply with the last Fibonacci number) → `[sandbox] deleted …`. -**33. LangGraph agent** +**34. LangGraph agent** ```bash uv sync --extra agents uv run --extra agents python examples/agent_patterns/langchain_agent.py @@ -338,13 +345,13 @@ uv run --extra agents python examples/agent_patterns/langchain_agent.py These also call a model, and save their result under `examples/workflow_examples/output/`. -**34. Repository analyzer** +**35. Repository analyzer** ```bash uv run python examples/workflow_examples/repo_analyzer.py ``` → `[sandbox] creating (template=…)…` → `[sandbox] created …; waiting until ready…` → `[sandbox] ready: …` → `[bootstrap] git available` and `[bootstrap] ripgrep available` (or `installing … via …` lines if the template lacks them) → `[clone] fetching https://github.com/NeevCloudAI/neev-sdk-python…` → `[clone] done (…)` → for each step: `[model] calling … (step N)…` → ` tokens: …` → ` $ ` with its live output between `┌─ live output` and `└─ exit N` … → the four-section report → `[artifact] saved N bytes → …/output/repo-analysis.md` → `[sandbox] deleting …`. Pass `--repo ` to analyze another repository. -**35. Browser agent** +**36. Browser agent** ```bash uv run python examples/workflow_examples/browser_agent.py --query "AI" ``` diff --git a/examples/agent_snapshot_fork_rollback.py b/examples/agent_snapshot_fork_rollback.py new file mode 100644 index 0000000..a5b6346 --- /dev/null +++ b/examples/agent_snapshot_fork_rollback.py @@ -0,0 +1,81 @@ +""" +Snapshot a running agent, fork it, and roll it back. + +- ``agent.snapshot(wait_until_ready=True)`` snapshots the sandbox backing the agent + and waits until the snapshot is Ready. +- ``agent.fork(name)`` creates a **new agent** from the agent's current state. The + source keeps running; the fork is Provisioning until ``wait_until_ready()``. +- ``agent.rollback(snapshot_id)`` returns the agent's sandbox to the snapshot in + place. The agent keeps its id and name. + +A paused or still-provisioning agent raises ``ConflictError`` (409) or +``TooEarlyError`` (425) instead. + +Run:: + + NEEV_API_KEY=... NEEV_ORG_ID=... NEEV_PROJECT_ID=... \\ + uv run python examples/agent_snapshot_fork_rollback.py +""" + +from __future__ import annotations + +import os +import sys + +from neevai import NeevAI +from neevai.errors import NeevAIError + +AGENT_TEMPLATE = os.environ.get("NEEV_AGENT_TEMPLATE", "claude-code") +MESSAGE_PATH = "demo/message.txt" + + +def read_message(agent) -> str: + """Reads the marker file from the agent's backing sandbox.""" + sandbox = agent.sandbox() + sandbox.wait_until_ready() + return sandbox.files.read_text(MESSAGE_PATH).strip() + + +def main() -> None: + with NeevAI() as client: + agent = fork = snapshot = None + try: + agent = client.agents.create( + {"name": "snapshot-demo", "agent_template": AGENT_TEMPLATE} + ) + agent.wait_until_ready() + print(f"ready: {agent.id}", file=sys.stderr) + + sandbox = agent.sandbox() + sandbox.wait_until_ready() + sandbox.files.write(MESSAGE_PATH, "original state") + + snapshot = agent.snapshot({"name": "before-change"}, wait_until_ready=True) + print(f"snapshot {snapshot.id} is {snapshot.status.value}", file=sys.stderr) + print(f"agent snapshots: {[s.name for s in agent.snapshots()]}") + + sandbox.files.write(MESSAGE_PATH, "modified state") + + fork = agent.fork("snapshot-demo-fork") + print(f"forked: {fork.id} ({fork.status})", file=sys.stderr) + fork.wait_until_ready() + print(f"fork has: {read_message(fork)}") + + agent.rollback(snapshot.id) + agent.wait_until_ready() + print(f"after rollback, {agent.name} ({agent.id}) has: {read_message(agent)}") + except NeevAIError as exc: + print(f"error: {exc}", file=sys.stderr) + raise + finally: + # The snapshot goes before the agent it was taken from. + if snapshot is not None: + client.sandboxes.delete_snapshot(snapshot.id) + if fork is not None: + fork.delete() + if agent is not None: + agent.delete() + + +if __name__ == "__main__": + main() diff --git a/specs/aiagent.yaml b/specs/aiagent.yaml index 0cb5ed4..fa1f8be 100644 --- a/specs/aiagent.yaml +++ b/specs/aiagent.yaml @@ -1474,11 +1474,13 @@ paths: /api/v1beta1/orgs/{org_id}/projects/{project_id}/agents/{agent_id}/fork: post: tags: [Snapshots] - summary: Fork an agent's sandbox into a new sandbox + summary: Fork an agent into a new agent description: | - Atomically snapshots the sandbox backing this agent and creates a new - sandbox from that snapshot. The agent's sandbox keeps running. Returns - the new sandbox object. + Snapshots the running agent and creates a new agent from that + snapshot, with the source's template and configuration and its own + gateway token. The source agent keeps running. The new agent is + Provisioning until it is ready to use; poll it until its status is + Ready. Returns the new agent. operationId: forkAgent parameters: - {$ref: '#/components/parameters/OrgID'} @@ -1491,10 +1493,10 @@ paths: schema: {$ref: '#/components/schemas/ForkSandboxRequest'} responses: "201": - description: New forked sandbox + description: New forked agent content: application/json: - schema: {$ref: '#/components/schemas/Sandbox'} + schema: {$ref: '#/components/schemas/Agent'} "400": {$ref: '#/components/responses/BadRequest'} "401": {$ref: '#/components/responses/Unauthorized'} "403": {$ref: '#/components/responses/Forbidden'} @@ -1920,7 +1922,8 @@ components: name: status in: query required: false - description: Only calls that ended this way. + description: >- + Only records that ended this way. A recorded program fails when it exits non-zero or is killed; its start and its connections always succeed. schema: {$ref: '#/components/schemas/AuditStatusFilter'} AuditActorTypeFilter: name: actor_type @@ -2354,9 +2357,9 @@ components: SnapshotType: type: string description: | - Type of the snapshot. `full`: memory and filesystem, restorable. - `filesystem`: files only, taken when a full capture is not possible; - its `restorability` is `unsupported`. + Type of the snapshot. `full`: captures process memory and filesystem (/workspace). + `filesystem`: captures filesystem (/workspace) only without process memory, + allowing restore onto any sandbox image. enum: [full, filesystem] SnapshotRestorability: type: string @@ -2375,8 +2378,8 @@ components: enum: [manual, memory_pressure] Snapshot: type: object - required: [id, sandbox_id, org_id, project_id, status, snapshot_type, source_region, created_at, updated_at, sandbox_format_version, restorability] - description: Metadata for a sandbox memory+filesystem snapshot. + required: [id, sandbox_id, org_id, project_id, status, snapshot_type, trigger, source_region, created_at, updated_at, sandbox_format_version, restorability] + description: Metadata for a sandbox snapshot (full or filesystem-only). properties: id: {type: string, format: uuid, description: Snapshot UUID.} sandbox_id: {type: string, format: uuid, description: UUID of the source sandbox.} @@ -2439,6 +2442,8 @@ components: Optional name for this snapshot. Must be a valid DNS name (lowercase alphanumeric or '-', starting with a letter, max 63 characters) so it can name sandboxes created on fork and restore. Omit to leave unnamed. + snapshot_type: + $ref: '#/components/schemas/SnapshotType' RollbackSandboxRequest: type: object required: [snapshot_id] @@ -2612,7 +2617,7 @@ components: type: object description: >- One thing that happened in the sandbox, as an OCSF 1.9.0 event: an API Activity event (class_uid 6003) for a call to the sandbox. Unless its `execution_trace` is turned off, a sandbox also records every program that runs inside it, whether or not anything requested it: a Process Activity event (1007) for a program starting or ending, and a Network Activity event (4001) for a connection a program opened. Program names are recorded without their arguments, and input typed at a hidden password prompt is never recorded. - required: [class_uid, class_name, category_uid, category_name, activity_id, activity_name, type_uid, type_name, time, severity_id, severity, status_id, status, actor, cloud, osint, metadata] + required: [class_uid, class_name, category_uid, category_name, activity_id, activity_name, type_uid, type_name, time, severity_id, severity, status_id, status, actor, cloud, device, osint, metadata] properties: class_uid: type: integer @@ -2692,7 +2697,7 @@ components: type: object description: >- The sandbox the event happened in, as OCSF device. Present on every event. - required: [type_id] + required: [type_id, uid] properties: type_id: type: integer @@ -2716,6 +2721,7 @@ components: parent_process: {$ref: '#/components/schemas/AuditParentProcess'} AuditParentProcess: type: object + required: [pid] description: >- The process that started this one, so a run can be traced back to what caused it. Its own Launch event names it. properties: @@ -2785,7 +2791,7 @@ components: The process the call started. A Process Activity event whose pid, or any ancestor's pid through parent_process, equals it was started by this call, so its actor.user is who ran that program. AuditCloud: type: object - required: [provider] + required: [provider, account, org] properties: provider: {type: string, description: Always NeevCloud.} account: {$ref: '#/components/schemas/AuditAccount'} @@ -2849,12 +2855,13 @@ components: enum: [User, API Key, System, Unknown] x-enum-varnames: [AuditActorTypeUser, AuditActorTypeAPIKey, AuditActorTypeSystem, AuditActorTypeUnknown] description: >- - The OCSF caption of type_id: User for a user in the console, System for the platform acting on its own, such as pairing a device, Unknown when no principal can be named, and API Key (type_id 99, Other) for an API key. + The OCSF caption of type_id: User for a user in the console, System for the platform acting on its own, such as pairing a device, and API Key (type_id 99, Other) for an API key. Unknown is only for a program the sandbox ran on its own, such as a background process, since no caller started it. type_id: type: integer enum: [0, 1, 3, 99] x-enum-varnames: [AuditActorTypeIDUnknown, AuditActorTypeIDUser, AuditActorTypeIDSystem, AuditActorTypeIDOther] - description: "OCSF user.type_id: 0 (Unknown), 1 (User), 3 (System), 99 (Other) for an API key." + description: >- + OCSF user.type_id: 1 (User), 3 (System), 99 (Other) for an API key, the same values as the platform audit log, so use it to match actors across them. 0 (Unknown) is only for a program the sandbox ran on its own. uid: type: string description: >- diff --git a/src/neevai/__init__.py b/src/neevai/__init__.py index 083fe00..fbe38eb 100644 --- a/src/neevai/__init__.py +++ b/src/neevai/__init__.py @@ -17,6 +17,7 @@ PreconditionFailedError, RateLimitError, ServiceUnavailableError, + TooEarlyError, ) from neevai.handles import Agent, AsyncAgent, AsyncSandbox, Sandbox from neevai.resources.agent_templates import ( @@ -152,6 +153,7 @@ "NotFoundError", "ConflictError", "PreconditionFailedError", + "TooEarlyError", "RateLimitError", "DeadlineExceededError", "InternalServerError", diff --git a/src/neevai/errors.py b/src/neevai/errors.py index 47314ca..f29ce11 100644 --- a/src/neevai/errors.py +++ b/src/neevai/errors.py @@ -120,6 +120,12 @@ class PreconditionFailedError(APIError): pass +class TooEarlyError(APIError): + """425 - The resource is not ready yet (for example, an agent still provisioning); retry later.""" + + pass + + class RateLimitError(APIError): """429 - Rate limit exceeded.""" @@ -194,6 +200,8 @@ def error_from_status( return ConflictError(status_code, body, request_id, **kwargs) elif status_code == 412: return PreconditionFailedError(status_code, body, request_id, **kwargs) + elif status_code == 425: + return TooEarlyError(status_code, body, request_id, **kwargs) elif status_code == 429: return RateLimitError(status_code, body, request_id, **kwargs) elif status_code == 503: diff --git a/src/neevai/generated/aiagent.py b/src/neevai/generated/aiagent.py index 2cfe03f..118bb70 100644 --- a/src/neevai/generated/aiagent.py +++ b/src/neevai/generated/aiagent.py @@ -187,7 +187,7 @@ class Snapshot(BaseModel): name: str | None = Field(None, description="Optional customer-supplied name.") status: SnapshotStatus snapshot_type: SnapshotType - trigger: SnapshotTrigger | None = None + trigger: SnapshotTrigger source_region: str = Field(..., description="Region where the snapshot is stored.") size_bytes: int | None = Field( None, @@ -229,6 +229,7 @@ class CreateSnapshotRequest(BaseModel): None, description="Optional name for this snapshot. Must be a valid DNS name (lowercase\nalphanumeric or '-', starting with a letter, max 63 characters) so it\ncan name sandboxes created on fork and restore. Omit to leave unnamed.\n", ) + snapshot_type: SnapshotType | None = None class RollbackSandboxRequest(BaseModel): @@ -387,11 +388,11 @@ class TypeId(IntEnum): class AuditDevice(BaseModel): type_id: TypeId = Field(..., description="Always 6, Virtual.") type: str | None = Field(None, description="Always Virtual.") - uid: str | None = Field(None, description="The sandbox id.") + uid: str = Field(..., description="The sandbox id.") class AuditParentProcess(BaseModel): - pid: int | None = Field(None, description="Identifies the parent within the sandbox.") + pid: int = Field(..., description="Identifies the parent within the sandbox.") class TypeId1(IntEnum): @@ -594,8 +595,8 @@ class AuditResource(BaseModel): class AuditCloud(BaseModel): provider: str = Field(..., description="Always NeevCloud.") - account: AuditAccount | None = None - org: AuditOrg | None = None + account: AuditAccount + org: AuditOrg class AuditMetadata(BaseModel): @@ -616,11 +617,11 @@ class AuditMetadata(BaseModel): class AuditActorUser(BaseModel): type: Type1 = Field( ..., - description="The OCSF caption of type_id: User for a user in the console, System for the platform acting on its own, such as pairing a device, Unknown when no principal can be named, and API Key (type_id 99, Other) for an API key.", + description="The OCSF caption of type_id: User for a user in the console, System for the platform acting on its own, such as pairing a device, and API Key (type_id 99, Other) for an API key. Unknown is only for a program the sandbox ran on its own, such as a background process, since no caller started it.", ) type_id: TypeId2 = Field( ..., - description="OCSF user.type_id: 0 (Unknown), 1 (User), 3 (System), 99 (Other) for an API key.", + description="OCSF user.type_id: 1 (User), 3 (System), 99 (Other) for an API key, the same values as the platform audit log, so use it to match actors across them. 0 (Unknown) is only for a program the sandbox ran on its own.", ) uid: str = Field( ..., @@ -1197,7 +1198,7 @@ class AuditEvent(BaseModel): description="What a program returned when it ended. Zero means it succeeded. Present on a Process Activity Terminate event unless a signal stopped the program.", ) dst_endpoint: AuditDstEndpoint | None = None - device: AuditDevice | None = None + device: AuditDevice resources: list[AuditResource] | None = Field( None, description="What the call touched: the sandbox, then any file or process it named, the program it ran, and the terminal it ran in. Present on API Activity only; Process and Network events name the sandbox in device.", diff --git a/src/neevai/handles/agent.py b/src/neevai/handles/agent.py index 48e4648..1dab6d0 100644 --- a/src/neevai/handles/agent.py +++ b/src/neevai/handles/agent.py @@ -11,13 +11,18 @@ from neevai.resources.sandboxes import ( DEFAULT_PORT_POLL_INTERVAL_MS, DEFAULT_PORT_WAIT_TIMEOUT_MS, + DEFAULT_SNAPSHOT_POLL_INTERVAL_MS, + DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS, + check_snapshot_wait, ) from neevai.types import ( AgentData, AgentLastCrash, AuditTrail, + CreateSnapshotParams, SandboxPort, Scope, + Snapshot, UpdateAgentParams, ) @@ -228,6 +233,64 @@ def rollback(self, snapshot_id: str | UUID) -> Agent: self._state = next_state._state return self + def snapshot( + self, + params: CreateSnapshotParams | Mapping[str, Any] | None = None, + *, + wait_until_ready: bool = False, + timeout_ms: int = DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS, + poll_interval_ms: int = DEFAULT_SNAPSHOT_POLL_INTERVAL_MS, + ) -> Snapshot: + """Captures a snapshot of the sandbox backing this agent. + + Returns at once with status Pending; with ``wait_until_ready=True`` it waits + and returns the Ready snapshot. + """ + if self.agents is None: + raise NeevAIError("Cannot snapshot an agent handle with no client context.") + if wait_until_ready: + # Reject a bad timing before the snapshot exists, so none is left behind. + check_snapshot_wait(timeout_ms, poll_interval_ms) + org_id = self.scope.org_id if self.scope else None + project_id = self.scope.project_id if self.scope else None + snapshot = self.agents.create_snapshot( + self.id, params, org_id=org_id, project_id=project_id + ) + if not wait_until_ready: + return snapshot + return self.agents._sandboxes.wait_for_snapshot( + snapshot.id, + timeout_ms=timeout_ms, + poll_interval_ms=poll_interval_ms, + org_id=org_id, + project_id=project_id, + ) + + def snapshots(self) -> list[Snapshot]: + """Lists snapshots taken of the sandbox backing this agent.""" + if self.agents is None: + raise NeevAIError("Cannot list snapshots on an agent handle with no client context.") + return self.agents.list_snapshots( + self.id, + org_id=self.scope.org_id if self.scope else None, + project_id=self.scope.project_id if self.scope else None, + ) + + def fork(self, name: str) -> Agent: + """Forks this agent into a new agent seeded from its current state. + + This agent keeps running. The new agent is Provisioning until it is ready; + call ``wait_until_ready()`` on it before use. + """ + if self.agents is None: + raise NeevAIError("Cannot fork an agent handle with no client context.") + return self.agents.fork( + self.id, + name, + org_id=self.scope.org_id if self.scope else None, + project_id=self.scope.project_id if self.scope else None, + ) + def audit( self, *, @@ -519,6 +582,64 @@ async def rollback(self, snapshot_id: str | UUID) -> AsyncAgent: self._state = next_state._state return self + async def snapshot( + self, + params: CreateSnapshotParams | Mapping[str, Any] | None = None, + *, + wait_until_ready: bool = False, + timeout_ms: int = DEFAULT_SNAPSHOT_WAIT_TIMEOUT_MS, + poll_interval_ms: int = DEFAULT_SNAPSHOT_POLL_INTERVAL_MS, + ) -> Snapshot: + """Captures a snapshot of the sandbox backing this agent asynchronously. + + Returns at once with status Pending; with ``wait_until_ready=True`` it waits + and returns the Ready snapshot. + """ + if self.agents is None: + raise NeevAIError("Cannot snapshot an agent handle with no client context.") + if wait_until_ready: + # Reject a bad timing before the snapshot exists, so none is left behind. + check_snapshot_wait(timeout_ms, poll_interval_ms) + org_id = self.scope.org_id if self.scope else None + project_id = self.scope.project_id if self.scope else None + snapshot = await self.agents.create_snapshot( + self.id, params, org_id=org_id, project_id=project_id + ) + if not wait_until_ready: + return snapshot + return await self.agents._sandboxes.wait_for_snapshot( + snapshot.id, + timeout_ms=timeout_ms, + poll_interval_ms=poll_interval_ms, + org_id=org_id, + project_id=project_id, + ) + + async def snapshots(self) -> list[Snapshot]: + """Lists snapshots taken of the sandbox backing this agent asynchronously.""" + if self.agents is None: + raise NeevAIError("Cannot list snapshots on an agent handle with no client context.") + return await self.agents.list_snapshots( + self.id, + org_id=self.scope.org_id if self.scope else None, + project_id=self.scope.project_id if self.scope else None, + ) + + async def fork(self, name: str) -> AsyncAgent: + """Forks this agent into a new agent seeded from its current state asynchronously. + + This agent keeps running. The new agent is Provisioning until it is ready; + await ``wait_until_ready()`` on it before use. + """ + if self.agents is None: + raise NeevAIError("Cannot fork an agent handle with no client context.") + return await self.agents.fork( + self.id, + name, + org_id=self.scope.org_id if self.scope else None, + project_id=self.scope.project_id if self.scope else None, + ) + async def audit( self, *, diff --git a/src/neevai/resources/agents.py b/src/neevai/resources/agents.py index 9dfadb6..0187301 100644 --- a/src/neevai/resources/agents.py +++ b/src/neevai/resources/agents.py @@ -14,6 +14,7 @@ DEFAULT_PORT_POLL_INTERVAL_MS, DEFAULT_PORT_WAIT_TIMEOUT_MS, _await_for_preview_url, + _prepare_create_snapshot_body, _wait_for_preview_url, ) from neevai.types import ( @@ -21,7 +22,10 @@ AgentListResponse, AuditTrail, CreateAgentParams, + CreateSnapshotParams, SandboxPort, + Snapshot, + SnapshotListResponse, UpdateAgentParams, ) @@ -283,6 +287,63 @@ def rollback( data = coerce_model(AgentData, raw) return Agent(self, data, scope) + def create_snapshot( + self, + id: str, + params: CreateSnapshotParams | Mapping[str, Any] | None = None, + org_id: str | None = None, + project_id: str | None = None, + ) -> Snapshot: + """Snapshots the sandbox backing an agent (returns immediately with status Pending).""" + scope = self._client._resolve_scope(org_id=org_id, project_id=project_id) + raw = self._client._transport.request( + "POST", + f"{_agents_path(scope)}/{id}/snapshots", + body=_prepare_create_snapshot_body(params), + ) + return coerce_model(Snapshot, raw) + + def list_snapshots( + self, + id: str, + page: int | None = None, + limit: int | None = None, + org_id: str | None = None, + project_id: str | None = None, + ) -> builtins.list[Snapshot]: + """Lists snapshots taken of the sandbox backing an agent.""" + scope = self._client._resolve_scope(org_id=org_id, project_id=project_id) + query: dict[str, Any] = {} + if page is not None: + query["page"] = page + if limit is not None: + query["limit"] = limit + raw = self._client._transport.request( + "GET", f"{_agents_path(scope)}/{id}/snapshots", query=query + ) + return builtins.list(coerce_model(SnapshotListResponse, raw).items) + + def fork( + self, + id: str, + name: str, + org_id: str | None = None, + project_id: str | None = None, + ) -> Agent: + """Forks an agent into a new agent seeded from its current state. + + The source agent keeps running. The new agent is Provisioning until it is + ready; call ``wait_until_ready()`` on it before use. + """ + from neevai.handles.agent import Agent + + scope = self._client._resolve_scope(org_id=org_id, project_id=project_id) + raw = self._client._transport.request( + "POST", f"{_agents_path(scope)}/{id}/fork", body={"name": name} + ) + data = coerce_model(AgentData, raw) + return Agent(self, data, scope) + def audit( self, id: str, @@ -588,6 +649,63 @@ async def rollback( data = coerce_model(AgentData, raw) return AsyncAgent(self, data, scope) + async def create_snapshot( + self, + id: str, + params: CreateSnapshotParams | Mapping[str, Any] | None = None, + org_id: str | None = None, + project_id: str | None = None, + ) -> Snapshot: + """Snapshots the sandbox backing an agent asynchronously (returns with status Pending).""" + scope = self._client._resolve_scope(org_id=org_id, project_id=project_id) + raw = await self._client._transport.request( + "POST", + f"{_agents_path(scope)}/{id}/snapshots", + body=_prepare_create_snapshot_body(params), + ) + return coerce_model(Snapshot, raw) + + async def list_snapshots( + self, + id: str, + page: int | None = None, + limit: int | None = None, + org_id: str | None = None, + project_id: str | None = None, + ) -> builtins.list[Snapshot]: + """Lists snapshots taken of the sandbox backing an agent asynchronously.""" + scope = self._client._resolve_scope(org_id=org_id, project_id=project_id) + query: dict[str, Any] = {} + if page is not None: + query["page"] = page + if limit is not None: + query["limit"] = limit + raw = await self._client._transport.request( + "GET", f"{_agents_path(scope)}/{id}/snapshots", query=query + ) + return builtins.list(coerce_model(SnapshotListResponse, raw).items) + + async def fork( + self, + id: str, + name: str, + org_id: str | None = None, + project_id: str | None = None, + ) -> AsyncAgent: + """Forks an agent into a new agent seeded from its current state asynchronously. + + The source agent keeps running. The new agent is Provisioning until it is + ready; await ``wait_until_ready()`` on it before use. + """ + from neevai.handles.agent import AsyncAgent + + scope = self._client._resolve_scope(org_id=org_id, project_id=project_id) + raw = await self._client._transport.request( + "POST", f"{_agents_path(scope)}/{id}/fork", body={"name": name} + ) + data = coerce_model(AgentData, raw) + return AsyncAgent(self, data, scope) + async def audit( self, id: str, diff --git a/src/neevai/types.py b/src/neevai/types.py index b4602c1..035223f 100644 --- a/src/neevai/types.py +++ b/src/neevai/types.py @@ -180,10 +180,7 @@ SandboxTemplateListResponse as SandboxTemplateListResponse, ) from neevai.generated.aiagent import ( # noqa: F401 - Snapshot as Snapshot, -) -from neevai.generated.aiagent import ( # noqa: F401 - SnapshotListResponse as SnapshotListResponse, + Snapshot as _GeneratedSnapshot, ) from neevai.generated.aiagent import ( # noqa: F401 SnapshotStatus as SnapshotStatus, @@ -235,11 +232,33 @@ class SandboxData(_GeneratedSandbox): phase: str # type: ignore[assignment] # pyright: ignore[reportIncompatibleVariableOverride] +class Snapshot(_GeneratedSnapshot): + """Snapshot record with an optional ``trigger``. + + Older API versions send no ``trigger``; accept its absence at the SDK boundary. + """ + + trigger: SnapshotTrigger | None = None # type: ignore[assignment] # pyright: ignore[reportIncompatibleVariableOverride] + + +class SnapshotListResponse(BaseModel): + items: list[Snapshot] + total: int + page: int + limit: int + + class CreateSnapshotParams(BaseModel): - """Caller-facing params for creating a sandbox snapshot.""" + """Caller-facing params for creating a sandbox snapshot. + + ``snapshot_type`` ``"filesystem"`` captures /workspace only, without process + memory, so it can be restored onto any sandbox image; ``"full"`` (the server + default) captures both. + """ name: str | None = None retain_for: str | None = None + snapshot_type: Literal["full", "filesystem"] | None = None class SandboxListResponse(BaseModel): diff --git a/tests/conftest.py b/tests/conftest.py index fa5166d..4feb806 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -70,6 +70,7 @@ def _make_snapshot_record( "name": req.get("name"), "status": "Pending", "snapshot_type": "full", + "trigger": "manual", "sandbox_format_version": 1, "restorability": "restorable", "source_region": "as-south-1", diff --git a/tests/test_agent_snapshots.py b/tests/test_agent_snapshots.py new file mode 100644 index 0000000..626d5d7 --- /dev/null +++ b/tests/test_agent_snapshots.py @@ -0,0 +1,198 @@ +"""Tests for agent snapshots and fork, and the 409/425 errors they surface.""" + +import json + +import httpx +import pytest + +from neevai.client import AsyncNeevAI, NeevAI +from neevai.errors import ConflictError, NeevAIError, TooEarlyError + +AGENT_ID = "11111111-1111-1111-1111-111111111111" +FORK_ID = "44444444-4444-4444-4444-444444444444" +SNAPSHOT_ID = "33333333-3333-3333-3333-333333333333" +PROJECT = "/api/v1beta1/orgs/org1/projects/proj1" +BASE = f"{PROJECT}/agents/{AGENT_ID}" + +AGENT = { + "id": AGENT_ID, + "org_id": "org1", + "project_id": "proj1", + "name": "web", + "agent_template_id": "ag-claude-code", + "drive_mode": "http", + "sandbox_id": "22222222-2222-2222-2222-222222222222", + "status": "Ready", + "metrics_url": "https://metrics.example/agents/1", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", +} +FORKED = { + **AGENT, + "id": FORK_ID, + "name": "web-fork", + "sandbox_id": "55555555-5555-5555-5555-555555555555", + "status": "Provisioning", +} + + +def _snapshot(status: str) -> dict: + return { + "id": SNAPSHOT_ID, + "sandbox_id": AGENT["sandbox_id"], + "org_id": "org1", + "project_id": "proj1", + "name": "before-change", + "status": status, + "snapshot_type": "full", + "trigger": "manual", + "sandbox_format_version": 1, + "restorability": "restorable", + "source_region": "as-south-1", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + } + + +def _handler(captured: list[tuple[str, str, dict | None, dict]]): + """Records (method, path, json body, query) and answers each route with a canned body.""" + + def handle(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) if request.content else None + path = request.url.path + captured.append((request.method, path, body, dict(request.url.params))) + if path == f"{BASE}/snapshots" and request.method == "POST": + return httpx.Response(202, json=_snapshot("Pending")) + if path == f"{BASE}/snapshots": + snaps = [_snapshot("Ready")] + return httpx.Response(200, json={"items": snaps, "total": 1, "page": 2, "limit": 5}) + if path == f"{PROJECT}/snapshots/{SNAPSHOT_ID}": + return httpx.Response(200, json=_snapshot("Ready")) + if path == f"{BASE}/fork": + return httpx.Response(201, json=FORKED) + return httpx.Response(200, json=AGENT) + + return handle + + +def _client(captured) -> NeevAI: + return NeevAI( + api_key="k", + base_url="https://api.example.com", + org_id="org1", + project_id="proj1", + client=httpx.Client(transport=httpx.MockTransport(_handler(captured))), + ) + + +def _async_client(captured) -> AsyncNeevAI: + return AsyncNeevAI( + api_key="k", + base_url="https://api.example.com", + org_id="org1", + project_id="proj1", + client=httpx.AsyncClient(transport=httpx.MockTransport(_handler(captured))), + ) + + +def test_agents_create_and_list_snapshots(): + captured: list = [] + client = _client(captured) + snap = client.agents.create_snapshot(AGENT_ID, {"name": "before-change", "retain_for": "24h"}) + assert str(snap.id) == SNAPSHOT_ID and snap.status == "Pending" + snaps = client.agents.list_snapshots(AGENT_ID, page=2, limit=5) + assert [str(s.id) for s in snaps] == [SNAPSHOT_ID] + client.agents.create_snapshot(AGENT_ID) + assert captured == [ + ("POST", f"{BASE}/snapshots", {"name": "before-change", "retain_for": "24h"}, {}), + ("GET", f"{BASE}/snapshots", None, {"page": "2", "limit": "5"}), + ("POST", f"{BASE}/snapshots", {}, {}), + ] + client.close() + + +def test_agents_fork_returns_the_new_agent(): + captured: list = [] + client = _client(captured) + forked = client.agents.fork(AGENT_ID, "web-fork") + assert forked.id == FORK_ID and forked.status == "Provisioning" + assert forked.agents is client.agents + assert captured == [("POST", f"{BASE}/fork", {"name": "web-fork"}, {})] + client.close() + + +def test_agent_handle_snapshot_snapshots_and_fork(): + captured: list = [] + client = _client(captured) + agent = client.agents.get(AGENT_ID) + assert agent.snapshot({"name": "before-change"}).status == "Pending" + assert agent.snapshot(wait_until_ready=True, poll_interval_ms=1).status == "Ready" + assert [str(s.id) for s in agent.snapshots()] == [SNAPSHOT_ID] + forked = agent.fork("web-fork") + assert forked is not agent and forked.id == FORK_ID + assert agent.id == AGENT_ID and agent.status == "Ready" + assert [(m, p) for m, p, _, _ in captured[1:]] == [ + ("POST", f"{BASE}/snapshots"), + ("POST", f"{BASE}/snapshots"), + ("GET", f"{PROJECT}/snapshots/{SNAPSHOT_ID}"), + ("GET", f"{BASE}/snapshots"), + ("POST", f"{BASE}/fork"), + ] + client.close() + + +def test_agent_handle_snapshot_rejects_bad_timing_before_creating(): + captured: list = [] + client = _client(captured) + agent = client.agents.get(AGENT_ID) + with pytest.raises(NeevAIError, match="timeout_ms must be a positive number"): + agent.snapshot(wait_until_ready=True, timeout_ms=0) + assert [m for m, _, _, _ in captured] == ["GET"] + client.close() + + +@pytest.mark.parametrize("status, error", [(409, ConflictError), (425, TooEarlyError)]) +def test_agent_fork_and_snapshot_surface_409_and_425(status: int, error: type): + def handle(request: httpx.Request) -> httpx.Response: + return httpx.Response(status, json={"code": "agent_not_ready", "message": "not ready"}) + + client = NeevAI( + api_key="k", + base_url="https://api.example.com", + org_id="org1", + project_id="proj1", + client=httpx.Client(transport=httpx.MockTransport(handle)), + ) + with pytest.raises(error) as fork_err: + client.agents.fork(AGENT_ID, "web-fork") + assert fork_err.value.status_code == status and fork_err.value.code == "agent_not_ready" + with pytest.raises(error): + client.agents.create_snapshot(AGENT_ID) + client.close() + + +@pytest.mark.asyncio +async def test_async_agent_snapshots_and_fork(): + captured: list = [] + client = _async_client(captured) + snap = await client.agents.create_snapshot(AGENT_ID, {"name": "before-change"}) + assert snap.status == "Pending" + assert len(await client.agents.list_snapshots(AGENT_ID, limit=5)) == 1 + assert (await client.agents.fork(AGENT_ID, "web-fork")).id == FORK_ID + agent = await client.agents.get(AGENT_ID) + ready = await agent.snapshot(wait_until_ready=True, poll_interval_ms=1) + assert ready.status == "Ready" + assert len(await agent.snapshots()) == 1 + forked = await agent.fork("web-fork") + assert forked is not agent and forked.status == "Provisioning" + assert [(m, p) for m, p, _, _ in captured] == [ + ("POST", f"{BASE}/snapshots"), + ("GET", f"{BASE}/snapshots"), + ("POST", f"{BASE}/fork"), + ("GET", BASE), + ("POST", f"{BASE}/snapshots"), + ("GET", f"{PROJECT}/snapshots/{SNAPSHOT_ID}"), + ("GET", f"{BASE}/snapshots"), + ("POST", f"{BASE}/fork"), + ] + await client.aclose() diff --git a/tests/test_audit.py b/tests/test_audit.py index ddc0b71..3158c26 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -57,7 +57,8 @@ def _event(uid: str, operation: str, status_id: int, actor: dict, **extra) -> di "status": status, "actor": actor, "api": {"operation": operation, "request": {"uid": f"req-{uid}"}}, - "cloud": {"provider": "NeevCloud", "account": {"uid": "org1"}, "project_uid": "proj1"}, + "cloud": {"provider": "NeevCloud", "account": {"uid": "proj1"}, "org": {"uid": "org1"}}, + "device": {"type_id": 6, "type": "Virtual", "uid": SANDBOX_ID}, "src_endpoint": {"name": "unknown"}, "osint": [], "metadata": { @@ -325,7 +326,7 @@ def _program(activity_id: int, name: str, **extra) -> dict: "status": "Success", "actor": {"user": {"type": "Unknown", "type_id": 0, "uid": "unknown"}}, "device": {"type_id": 6, "type": "Virtual", "uid": SANDBOX_ID}, - "cloud": {"provider": "NeevCloud"}, + "cloud": {"provider": "NeevCloud", "account": {"uid": "proj1"}, "org": {"uid": "org1"}}, "osint": [], "metadata": { "version": "1.9.0", @@ -412,7 +413,7 @@ def handler(request: httpx.Request) -> httpx.Response: assert launch.process.file.path == "/usr/bin/curl" assert launch.process.parent_process is not None assert launch.process.parent_process.pid == 39 - assert launch.device is not None and launch.device.uid == SANDBOX_ID + assert launch.device.uid == SANDBOX_ID assert launch.actor.user.type == "Unknown" assert call.api is not None and call.api.operation == "exec" client.close() diff --git a/tests/test_sandboxes.py b/tests/test_sandboxes.py index 07daa5d..4f95ce4 100644 --- a/tests/test_sandboxes.py +++ b/tests/test_sandboxes.py @@ -309,7 +309,8 @@ def capturing_request(method, path, query=None, body=None): client._transport.request = capturing_request # type: ignore[method-assign] snap = client.sandboxes.create_snapshot(sb.id, {"name": "demo-snap"}) - assert captured_bodies == [{"name": "demo-snap"}] + client.sandboxes.create_snapshot(sb.id, {"snapshot_type": "filesystem"}) + assert captured_bodies == [{"name": "demo-snap"}, {"snapshot_type": "filesystem"}] assert str(snap.id) == _first_snapshot_id() assert snap.status == SnapshotStatus.Pending client.close()