ax
CLI Toolgoogle/ax
Declarative orchestrator for running sandboxed AI agent workloads at scale.
Overview
AX is Google's declarative orchestration runtime for autonomous AI agent workloads. Tasks, Workspaces, Gateways and Models are defined as Kubernetes-style manifests, and agents run in isolated sandboxes with network egress allowlists. A kubectl-shaped CLI provides live watching, shell access, and suspend/resume.
README Preview
\n \n AX\n\n\n> [!WARNING]\n> We are still actively refining our core concepts, protocols,\n> and specifications. We will likely to introduce major breaking\n> changes prior to a stable release.\n\n**Declare an agentic task with workspaces and gateway specifications. AX sandboxes it, wires up its workspace, fences its network, and helps running it at scale.**\n\nAX is a high-throughput, declarative orchestrator to run billions of autonomous agent workloads in a cluster. It runs on top of [Agent Substrate](https://github.com/agent-substrate/substrate) for sandboxed execution and is built to run billions of tasks per cluster. If you have used Kubernetes, `ax` will feel similar.\n\n```yaml\n# task.yaml\napiVersion: ax.io/v1alpha1\nkind: Workspace\nmetadata:\n name: golang\nspec:\n git:\n - repo: https://github.com/golang/go.git\n branch: "my-fix"\n---\napiVersion: ax.io/v1alpha1\nkind: Task\nmetadata:\n name: test\nspec:\n workspaces:\n - name: golang\n goal: "Ensure that Go tool chain is available and is built from source"\n debug: true # lets you `ax ssh` into the sandbox\n```\n\nThen apply it, watch it come up, and look over the agent's shoulder:\n\n```bash\nax apply -f task.yaml\nax watch task test\nax ssh test -- ls -al /workspace\n```\n\n## Why?\n\nAgents are a new kind of workload. They are neither stateless microservices nor run-to-completion batch jobs. They accumulate state, need strict isolation, call out to model APIs and tool servers, and can burn money in a loop if nobody is watching. AX gives you four small primitives that handle all of that declaratively:\n\n| You want to... | AX gives you |\n|---|---|\n| Run untrusted agent code in an isolated sandbox with CPU/memory limits | **`Task`** |\n| Pre-wire Git repos, MCP servers, and skill packages so every agent starts warm | **`Workspace`** |\n| Lock outbound traffic down to an explicit host allowlist | **`Gateway`** |\n| Configure which LLM the platform itself uses, with credentia
FAQ (5)
TroubleshootingWhy does `ax ssh <task> -- <cmd>` drop or misinterpret flags like `-a`, `-n`, or `--server` in the remote command?
This is a bug in the global flag parser (cmd/ax/main.go, lines 58-93): it does not stop parsing at the -- separator, so it consumes -a/--atespace, -n/--namespace, --server, and --context from your remote command and applies them to ax itself. Symptoms include task not found in atespace ..., namespaces "50" not found, or silent failures (e.g. uname -a runs as plain uname). **Fix:** patch the parser loop to pass everything after -- through unchanged:
if arg == "--" {
cleanArgs = append(cleanArgs, args[i:]...)
break
} else if arg == "-a" || arg == "--atespace" {
A fixed branch is also linked from issue #373. **Workarounds until patched:** combine short flags so they don't match an ax flag (e.g. use ls -al instead of ls -a -l), and avoid using -a, -n, --server, or --context inside the remote command. Flags placed *before* -- (e.g. ax ssh mytask -a team -- ls) continue to work normally.
TroubleshootingWhy does my task stay Suspended after `ax resume` returns success?
This is a race condition in ax versions before commit ac23328. When ax resume runs right after ax apply (within milliseconds), the controller's processEvent (in internal/controller/worker.go) read the task, reconciled it, and wrote the whole status back via UpdateTaskStatus with no version check. The late create-event write clobbered the Running phase set by ResumeTask, leaving the task stuck Suspended even though the CLI reported success.
**Fix:** Upgrade to a build including commit ac23328 or later. The controller rewrite replaced the Redis Streams queue with direct execution under a per-task lock (a Redis lease), which serializes create/resume operations and eliminates the lost-update window. A concurrency regression test now pins this invariant.
**Workaround on older versions:** Insert a delay between ax apply and ax resume in automation scripts, and verify the task phase after resuming (e.g., poll task status); retry ax resume if the task is still Suspended.
**Note:** The per-task lock is a Redis lease with a 30s TTL and no renewal. CLI deadlines (30s for apply, 15s for suspend/resume) keep handlers within it, but ax delete task can run up to 2 minutes, so a resume may theoretically acquire the lock while a very slow delete still holds it — no state loss has been observed from this.
TroubleshootingWhy does ax-task-runner fail with "unknown field" error after upgrading ax-server or adding a new Task field?
The runner decodes the Task from the AX_TASK_YAML env var using strict protojson.UnmarshalOptions{} in pkg/apis/v1alpha1/types.go, so any Task field newer than the runner's image causes a parse error like unknown field "egress" and the runner exits before starting the agent. Fix (runner-side): add an UnmarshalYAMLLenient method to pkg/apis/v1alpha1 that runs the same per-kind normalizers as the existing UnmarshalYAML but decodes with protojson.UnmarshalOptions{DiscardUnknown: true}, then use it in cmd/ax-task-runner/main.go for AX_TASK_YAML/--task-file and AX_WORKSPACES_YAML/--workspace-file. Keep the existing UnmarshalYAML strict so manifests applied with ax apply still reject typos. Additionally, control-plane-only fields the runner never uses should be stripped from AX_TASK_YAML by the reconciler so even older runners stay compatible. Reference commit: mastersingh24/ax@4363c253.
TroubleshootingWhy does AX show a task as Suspended after a request wakes it, and how to fix ax ssh refusing to connect?
AX kept its own cached copy of the task phase and only updated it when AX itself suspended or resumed the task, so wakes triggered by the Agent Substrate router (or external suspends like kubectl ate suspend actor) were invisible to it. This caused ax get task to show Suspended with no worker address while the actor was actually ACTOR_STATE_RUNNING, and ax ssh to fail with task "..." is in phase "Suspended" (must be Running to ssh).
Fix: upgrade to a release including PR #515. AX no longer stores its own copy of the task state — ax-server reads the actor's live state from the Substrate/ATE control API on every ax get and ax list, so the phase and worker address always reflect reality regardless of who suspended or resumed the actor.
Workaround on older versions: force AX to reconcile by resuming the task through AX itself, e.g. ax resume task <task-name>, after which ax ssh works again.
TroubleshootingWhy does AX SuspendTask return success with phase Suspended while the Substrate actor is still running?
This is a bug in AX (observed on main at commit ac2332829f22360ff97b0ba34d94dd0dd782f17e). When Substrate's SuspendActor RPC fails — e.g. returns codes.Unavailable or codes.PermissionDenied — Client.SuspendActor correctly propagates the error, but the suspend branch of TaskReconciler.Reconcile (internal/controller/reconciler.go, ~L192-L204) only logs 'could not suspend actor on Substrate' and still returns task, nil after setting the suspended status. Server.SuspendTask (internal/server/server.go, ~L256-L306) then persists phase: Suspended, clears the stored worker IP, sets Ready=False with reason TaskSuspended, and the gRPC call returns OK — even though the actor remains ACTOR_STATE_RUNNING.
Fix (for maintainers/forks): patch the reconciler to propagate the SuspendActor error instead of swallowing it, and only persist confirmed lifecycle results. Keep returning success when the actor is already suspended, since the pinned Substrate implementation treats repeated suspension as idempotent. Add regression coverage for: successful suspension, rejected suspension, already-suspended actors, and Task creation (creation shares this reconciliation branch). For ambiguous timeout or cancellation failures, re-query the actual actor state from Substrate before persisting, since a failed RPC does not guarantee the actor is still running.
Operational workaround: do not trust the stored Task phase or the SuspendTask response alone — verify suspension by querying the actor state directly from Substrate before assuming execution has stopped.