Google AX and the Runtime for Stateful AI Agents

Picture a coding agent working through a repository. It reads a file, runs a test, waits eight seconds for a model response, asks for approval, then continues from the same files and memory. That rhythm does not look like a traditional web request or a batch job. It is an agentic workload: software that can plan steps, call tools, retain state, and keep going over time.

Google AX is an open-source, declarative orchestrator for that shape of work. A declarative system lets you describe the result you want instead of scripting every machine-level action. AX turns an agent session into named resources for execution, environment setup, network access, and model configuration, then runs those resources on top of Agent Substrate. The project is Kubernetes-shaped, but its important idea is not another YAML wrapper: it treats a living agent session as infrastructure. (github.com)

Why a normal container starts to feel wrong

A microservice is usually expected to stay available and answer requests. A batch job starts, does its work, and exits. An agent does something awkward in between: it may use a CPU heavily while running code, then spend minutes waiting on a model API, a browser, a tool server, or a human approval. During that wait, a conventional container may still reserve memory and machine capacity.

The practical question is: how do you keep a long-running agent responsive without paying to keep an idle process alive? AX's answer depends on a runtime that can pause the session without throwing away its working memory and files. This is why the design talks about actors and workers rather than only containers.

The four pieces AX wants you to name

Task: one isolated unit of execution

A Task is the smallest AX unit that gets a sandbox, meaning a fenced environment for running code. It declares the container image, command, CPU and memory limits, environment variables, attached workspaces, and network gateway. A task can be the whole agent, or one child task in a larger tree when an agent delegates work. Either way, it gets the same lifecycle: create, run, suspend, resume, or delete.

Workspace: make the sandbox useful

A fresh sandbox is not useful if it opens with an empty directory and no tools. A Workspace describes the material an agent needs before its command starts: Git repositories, MCP servers, and skill packages. MCP, short for Model Context Protocol, is a common way for models and agent programs to discover and call external tools. A workspace can also carry a plain-language goal, such as preparing a Python environment and installing dependencies. On first boot, AX can use that goal to guide environment setup.

Gateway and Model: control the edges

A Gateway defines the task's network boundary. Instead of allowing every outbound connection, you can permit only the hosts and ports the agent needs, such as a model provider and a source-code mirror. That matters when generated code can execute shell commands or download packages.

A Model is not the model itself. It is a named configuration containing the provider, model identifier, generation parameters, and a reference to a Kubernetes Secret holding credentials. Centralizing this configuration means one change can update many tasks without copying keys through agent images.

A manifest becomes the packing slip

This small example gives an agent a repository, restricted egress, and a setup goal:

apiVersion: ax.io/v1alpha1
kind: Workspace
metadata:
 name: analytics
spec:
 git:
 - name: source
 repo: analytics-repository
 branch: main
---
apiVersion: ax.io/v1alpha1
kind: Gateway
metadata:
 name: restricted-egress
spec:
 egress:
 allowlist:
 hosts:
 - host: api.example.com
 port: 443
 - host: model-provider.example
 port: 443
---
apiVersion: ax.io/v1alpha1
kind: Task
metadata:
 name: report-agent
spec:
 image: ghcr.io/example/report-agent
 command: ['python', 'agent.py']
 resources:
 requests:
 cpu: '500m'
 memory: '1Gi'
 limits:
 cpu: '2'
 memory: '4Gi'
 workspaces:
 - name: analytics
 path: /workspace
 goal: 'Install dependencies and prepare a reproducible Python analysis environment'
 gateway:
 name: restricted-egress
 debug: true

Applying the file with ax apply -f agent.yaml sends the desired state to AX. ax watch follows readiness and lifecycle changes, while ax suspend and ax resume make the pause explicit. The debug switch is useful during development because it enables shell access into the sandbox; it is a deliberate capability, not something every production task needs. (github.com)

What happens behind the command

The ax command talks to an AX server over gRPC, a typed protocol for calling remote services. The server validates the manifest and stores state in Redis, a fast data store. Redis Streams then act as the work queue for horizontally scaled controllers, which compare the requested state with reality and make corrections.

That choice is more than an implementation detail. AX's design avoids making Kubernetes track every short-lived task as a custom resource in etcd, its control-plane database. For a system creating and suspending huge numbers of small sessions, that can reduce pressure on the normal Kubernetes API path. The controller asks Agent Substrate to create the actor, assign it to a worker, prepare the sandbox, and apply the egress policy.

The pause is the feature

Agent Substrate supplies the execution layer underneath. An actor is the logical agent session; a worker is a ready sandbox where an active actor runs. Because the actor is not permanently tied to one worker, the runtime can snapshot the session's RAM and local files when it goes idle, release the worker, and restore the actor onto another warm worker when a new event arrives.

That is different from restarting a container and rebuilding state from a database. The process can return with its working files and in-memory context intact. Google describes Agent Substrate as supporting sub-500-millisecond resume operations and hundreds of suspend/resume activations per second in its GKE announcement, while its documentation describes restoring suspended agents in under a second. Those figures are platform targets and published product claims, not a promise for every workload. (cloud.google.com)

The result is a form of dense multiplexing: many mostly idle sessions share a smaller pool of physical capacity, while active sessions receive compute when they need it. This is the piece that makes agent orchestration economically different from keeping one container running for every user. It does not make model-token costs disappear, and it does not remove the need for quotas, logging, and cancellation.

Where AX fits, and where it is still moving

AX is not an agent framework that decides how an agent reasons. You still bring the harness, the program that runs the model-and-tool loop, along with its image, tools, prompts, and application policy. AX handles the surrounding operational questions: where the code runs, what it can reach, what environment it starts with, and how its state survives a pause.

As of September 21, 2026, the public AX repository warns that its core concepts and specifications are still changing and may break before a stable release. Google describes Agent Executor as available in preview, while GKE documentation lists Agent Substrate for evaluation and non-production use, with production support offered through an allowlist. That makes AX interesting for platform experiments and research, but it also means teams should pin revisions, isolate test clusters, and expect the APIs to move.

The shift in mental model

The useful mental shift is small but important: an agent is not merely a request handler inside a container. It is a stateful actor that wakes, thinks, uses tools, waits, and wakes again. Google AX gives that actor a declarative home, and Agent Substrate supplies the pause-and-resume machinery underneath. Once those layers are separated, sandboxing, workspace preparation, network policy, and efficient scheduling stop looking like unrelated chores and start looking like parts of one runtime. (docs.cloud.google.com)