·15 min read

Google AX: Google's open-source distributed agent runtime (Apache 2.0)

Google's open-source distributed agent runtime — designed for scalable, distributed agent execution across cloud and edge environments.

The Problem

You have built an agent that works beautifully on your laptop. It calls tools, reasons through multi-step plans, and produces results that impress everyone in the demo. Then you deploy it to production, and the problems start.

The agent runs for 45 minutes doing deep research. Halfway through, the pod gets evicted. The entire execution is lost. You restart it from scratch. The same thing happens again. Or your agent calls a shell command that deletes files on the host. Or two agents in the same session race on shared state and corrupt each other’s data. Or a client disconnects during a long-running task and has no way to reattach.

These are not edge cases. They are the defining operational challenges of production agentic AI. Agent workflows are fundamentally different from traditional request-response services: they are long-running, nonlinear, stateful, and often involve executing untrusted code. The tools we use for stateless microservices do not apply.

Dimension Before (DIY / Framework-Only) After (Google AX)
Crash recovery Full restart from scratch Event-log replay, resume from last checkpoint
State consistency Shared mutable state, race conditions Single-writer architecture, no concurrent mutations
Client disconnect Lost session, user must restart Resumable streams with sequence-number catch-up
Tool isolation Same process, full host access Sandboxed environments (gVisor/Kata)
Audit trail None or ad-hoc logging Append-only event log, every step recorded
Multi-tenant safety Shared namespace, cross-tenant bleed Per-component sandboxing
Execution forking Manual checkpoint hacks Native trajectory branching from any sequence point

“Why this matters: The difference between a demo agent and a production agent is not intelligence — it is reliability. Google AX solves the operational problems that every agent team hits around month three: crashes, state corruption, isolation failures, and the inability to resume or audit long-running executions.”

The Investigation

The root cause of production agent failures is not bad agent logic. It is the absence of a runtime designed for agentic execution patterns. Standard Kubernetes assumes long-running, stateless services. Standard web frameworks assume request-response cycles measured in milliseconds. Agents violate both assumptions.

Google’s team, led by Jaana Burcu Dogan (rakyll), investigated this problem across Google’s internal agent deployments. They found a consistent pattern: every team was building the same infrastructure — an event logger, a state manager, a resumption protocol — from scratch. The result was fragmentation, bugs, and wasted engineering time.

What this means: The agent runtime layer was missing from the open-source ecosystem. Frameworks like LangGraph and CrewAI solve the orchestration problem (how to structure agent logic). They do not solve the execution problem (how to run agents reliably at scale). Google AX fills that gap.

The investigation produced specific design requirements:

  • Durable execution: Every step must be recorded so the system can recover from any failure point.
  • Single-writer state: Only one component can mutate shared session state at a time, eliminating race conditions.
  • Secure isolation: Every component (agent, tool, skill) must run in its own sandbox with no access to the host.
  • Connection recovery: Clients must be able to disconnect and reconnect without losing context.
  • Trajectory branching: Developers must be able to fork execution at any point for testing and debugging.

The project was announced publicly on May 21, 2026, on the Google Cloud Blog, and the repository at github.com/google/ax was opened under Apache 2.0. The first release, v0.1.0, shipped the next day.

The Solution

Google AX (Agent eXecutor) is a distributed agent runtime written in Go (75%) with Python bindings (15%). It is not an agent framework — it is a runtime that coordinates agentic loops, manages execution state via an append-only event log, and communicates with both local and remote actors over resumable streams.

                    ┌─────────────┐
                    │   Client    │
                    └──────┬──────┘
                           │ resumable stream
                    ┌──────▼──────┐
                    │   Router    │
                    └──────┬──────┘

                    ┌──────▼──────────────────┐
                    │     AX Controller       │
                    │  ┌──────────────────┐   │
                    │  │    Executor       │   │
                    │  │  (TaskExecutor)   │   │
                    │  ├──────────────────┤   │
                    │  │    Event Log      │   │
                    │  │  (SQLite/NDJSON)  │   │
                    │  ├──────────────────┤   │
                    │  │    Registry       │   │
                    │  └──────────────────┘   │
                    └──────┬──────────┬───────┘
                           │          │
              ┌────────────┼─────┐    │
         ┌────▼────┐ ┌────▼──┐ │┌───▼────────┐
         │ Remote  │ │ Env   │ ││   Tool     │
         │ Agent   │ │(Skills│ ││ (MCP Svr)  │
         │(gRPC)   │ │+Tools)│ ││            │
         └─────────┘ └───────┘ │└────────────┘

                         ┌──────▼──────┐
                         │  Sandbox    │
                         │(gVisor/Kata)│
                         └─────────────┘

Here’s what each piece does:

  • Client: The entry point. Sends input via ax exec CLI or gRPC. Communicates over resumable streams — if the client disconnects, it reconnects with the last sequence number it saw and the server backfills missed events.
  • Router: Routes client requests to the correct controller instance. Handles connection multiplexing.
  • AX Controller: The central brain. Contains the executor (runs the agentic loop), the event log (persists every step), and the registry (discovers agents, tools, and environments). Uses a single-writer architecture — only one component can update session state at a time.
  • Remote Agent: An isolated actor running outside the controller, invoked over gRPC via the AgentService.Connect RPC. Can be any framework — LangGraph, ADK, A2A, or a custom implementation.
  • Environment: An isolated actor containing skills and built-in tools. Runs in its own sandbox.
  • Tool: An MCP (Model Context Protocol) server. AX natively integrates with any MCP-compatible tool.
  • Sandbox: The isolation boundary. Uses gVisor or Kata Containers to prevent agent-generated code from affecting the host or other tenants.

Production-Grade Code Walkthrough

Installation:

go install github.com/google/ax/cmd/ax@latest

This installs the ax binary with two subcommands: exec (run or resume an agent) and serve (start the controller as a gRPC server).

Configuration (ax.yaml):

server:
  address: ":8494"

eventlog:
  sqlite:
    filename: "eventlog/log.sqlite"

planner:
  gemini:
    model: "gemini-3.5-flash"
    timeout: "60s"
    skills_dir: "./examples/skills"

registry:
  remote_agents:
    - id: "medical-deep-researcher"
      name: "Medical Deep Researcher"
      description: "Performs deep medical research using PubMed and clinicaltrials.gov"
      address: "localhost:50051"

The configuration declares a server on port 8494, SQLite-backed event log storage, a Gemini planner with a 60-second timeout, and a remote agent registered for medical research.

Running an agent:

# Simple execution
ax exec --input "List the files in this directory"

# Use a specific agent
ax exec --agent coding --input "Write a Python HTTP server"

# Start the controller server
ax serve --config ax.yaml

# Connect to a remote controller
ax exec --server localhost:8494 --input "Research quantum computing advances"

Resuming a disconnected session:

ax exec \
  --conversation d85a4b4e-c53b-4c84-b879-f10d905bce40 \
  --last-seq 12 \
  --resume

This reconnects to conversation d85a4b4e... and catches the client up from sequence number 12. The server replays events the client missed. The conversation is not rewound — the agent continues from where it left off.

Building a remote agent (Go):

package main

import (
    "context"
    "log"
    "net"

    "google.golang.org/grpc"
    pb "github.com/google/ax/proto"
)

type customAgent struct {
    pb.UnimplementedAgentServiceServer
}

func (a *customAgent) Connect(stream pb.AgentService_ConnectServer) error {
    for {
        req, err := stream.Recv()
        if err != nil {
            return err
        }
        // Process the request and send responses
        resp := &pb.StreamResponse{
            Content: &pb.Content{
                Parts: []*pb.Part{
                    {Data: &pb.Part_Text{Text: "Processed: " + string(req.Input)}},
                },
            },
        }
        if err := stream.Send(resp); err != nil {
            return err
        }
    }
}

func main() {
    lis, _ := net.Listen("tcp", ":50051")
    s := grpc.NewServer()
    pb.RegisterAgentServiceServer(s, &customAgent{})
    log.Fatal(s.Serve(lis))
}

Wrapping a Google ADK agent (Python):

from google import genai
from ax.adk import ADKAgentWrapper

# Wrap an existing ADK agent
agent = ADKAgentWrapper(
    agent_id="research-agent",
    model="gemini-3.5-flash",
    tools=[web_search, code_interpreter]
)

# Start the remote agent server
agent.serve(port=50052)

How to Use Effectively

1. Start with the built-in planner

ax exec --input "What files are in the current directory?"

The built-in planner uses a Gemini model with a bash tool. It automatically adapts to your operating system. This is the fastest way to validate that AX works in your environment.

2. Configure a persistent event log

eventlog:
  sqlite:
    filename: "/var/lib/ax/eventlog.sqlite"

The default event log is in-memory. For production, configure SQLite persistence. Every step — inputs, outputs, state changes — is recorded to the append-only log. This enables crash recovery and full audit trails.

3. Register remote agents for specialized tasks

registry:
  remote_agents:
    - id: "code-reviewer"
      name: "Code Review Agent"
      description: "Reviews pull requests for style, security, and correctness"
      address: "code-review.internal:50051"
    - id: "data-analyzer"
      name: "Data Analysis Agent"
      description: "Runs statistical analysis on CSV datasets"
      address: "data-analysis.internal:50052"

Remote agents run in isolated processes. They communicate with the controller over gRPC using the AgentService.Connect RPC. This lets you scale agents independently and deploy them on different hardware.

4. Use trajectory branching for testing

# Run an execution
ax exec --conversation test-001 --input "Design a database schema for an e-commerce platform"

# Branch from sequence 8 to test an alternative approach
ax exec --conversation test-001-branch --last-seq 8 --resume

Trajectory branching lets you fork an execution at any checkpoint. The branched execution shares all prior context but diverges from the branch point. This is invaluable for A/B testing agent behavior, debugging, and exploring alternative decision paths.

5. Deploy on Kubernetes with Agent Substrate

apiVersion: substrate.agent/v1
kind: AgentDeployment
metadata:
  name: ax-controller
spec:
  replicas: 3
  controller:
    image: ghcr.io/google/ax-controller:v0.1.0
    config: ax.yaml
  sandbox:
    runtime: gvisor
    memory: "2Gi"
    cpu: "1"

Agent Substrate is Google’s Kubernetes-native compute layer for agent workloads. It handles pod snapshotting (freezing idle agents), rapid sandbox provisioning (~300 sandboxes/second at <200ms latency), and horizontal scaling across the cluster.

Use Cases

1. Long-Running Research Agents

When you’d use this: Your agent needs to search the web, read dozens of documents, synthesize findings, and produce a report. The process takes 30-90 minutes. A pod restart or network blip should not force a restart.

Why Google AX fits: The event log records every step. If the process crashes, the agent resumes from the last checkpoint — not from scratch. The resumable stream protocol means the client can disconnect and come back later to collect results.

2. Multi-Tenant SaaS Agent Platform

When you’d use this: You are building a platform where multiple customers deploy agents that process their data. Each tenant’s agents must be isolated from each other and from the host.

Why Google AX fits: Per-component sandboxing (gVisor/Kata Containers) prevents cross-tenant contamination. The single-writer architecture prevents state corruption when multiple agents operate concurrently. The event log provides a complete audit trail for compliance.

3. Code-Generating and Code-Executing Agents

When you’d use this: Your agent writes Python scripts, shell commands, or SQL queries and executes them. You need to prevent the generated code from accessing the host filesystem, network, or other tenants’ data.

Why Google AX fits: Environments run in sandboxed actors. The bash tool requires explicit user approval before execution. Generated code runs in an isolated sandbox with no access to the host. Malicious or buggy code cannot escape.

4. Human-in-the-Loop Approval Workflows

When you’d use this: Your agent drafts emails, generates code changes, or makes API calls that require human approval before execution. The workflow may pause for minutes or hours waiting for a response.

Why Google AX fits: The event log snapshots state at every checkpoint. The agent can pause indefinitely and resume when the human responds. The resumable stream protocol handles long idle periods without connection drops.

5. Agent Evaluation and A/B Testing

When you’d use this: You have two versions of an agent (different prompts, different models, different tool sets) and want to compare their performance on the same task.

Why Google AX fits: Trajectory branching lets you fork an execution at any point and run alternative paths. You can compare outputs, measure latency, and evaluate quality without re-executing the shared prefix. The event log provides the ground truth for evaluation.

Cheat Sheet

Aspect Detail
What it is Distributed agent runtime for coordinating, logging, and resuming agentic executions
License Apache 2.0
Language Go (75%), Python (15%)
Repository github.com/google/ax
Latest release v0.1.0 (May 20, 2026)
Status Preview — active early development, breaking changes expected
Install go install github.com/google/ax/cmd/ax@latest
Architecture Single-writer controller + event log + registry + remote actors over gRPC
Event log Append-only, SQLite or NDJSON, every step recorded
Resumption Sequence-number-based catch-up on resumable streams
Branching Trajectory branching from any checkpoint
Isolation Per-component sandboxing (gVisor/Kata Containers)
Agent support Built-in (Gemini), remote (gRPC), ADK (Python), A2A, custom
Tool support Native MCP (Model Context Protocol)
Kubernetes Agent Substrate for pod snapshotting and rapid sandbox provisioning
What it is NOT Not a managed service, not an agent framework, not a harness, not model-specific
Roadmap Antigravity harness, BYOH, subagent suspension, tool call approvals

Vibe Coding Projects

Project 1: Multi-Agent Research Pipeline

What it does: Build a research pipeline with three remote agents — a web searcher, a document summarizer, and a report writer. The controller orchestrates them: search results feed into the summarizer, summaries feed into the writer. Each agent runs in its own sandboxed process.

What you’ll learn: Registering remote agents, configuring the event log for durability, handling inter-agent communication via the controller, and using trajectory branching to compare different summarization strategies.

Effort: 2-3 hours. You will write three Go gRPC servers (one per agent) and one ax.yaml configuration file.

Project 2: Human-in-the-Loop Code Review Agent

What it does: An agent that reviews pull requests, generates comments, and waits for human approval before posting them. The agent pauses after generating each comment, the human reviews and approves or rejects, and the agent resumes to post the approved comments.

What you’ll learn: Building HITL workflows with AX’s durable execution model, handling long pauses between agent steps, and using the event log to audit the full review history.

Effort: 3-4 hours. You will implement a custom agent that yields control back to the human at each approval gate.

Project 3: Agent Evaluation Harness

What it does: A harness that runs the same task against two different agent configurations (different models, different prompts, different tool sets) using trajectory branching. It collects outputs, measures latency, and produces a comparison report.

What you’ll learn: Using trajectory branching for A/B testing, parsing the event log for evaluation metrics, and building automated evaluation pipelines on top of AX’s execution model.

Effort: 4-5 hours. You will write a Go harness that forks executions and compares results programmatically.

Problems Solved Efficiently

Problem Type Why Google AX Fits When to Look Elsewhere
Long-running agent crashes Event log + resumption recovers from any failure point Your agents run < 30 seconds and restarts are acceptable
Multi-tenant isolation Per-component sandboxing with gVisor/Kata You run a single agent on a single machine
Client disconnection Resumable streams with sequence-number catch-up Your clients are always connected (LAN, same process)
State corruption Single-writer architecture prevents races Your agents are stateless or use an external DB
Audit and compliance Append-only event log records every step You do not need audit trails
Agent A/B testing Trajectory branching from any checkpoint You evaluate agents offline with recorded datasets
Tool isolation Sandboxed environments for MCP servers Your tools are read-only API calls with no side effects
Rapid agent scaling Agent Substrate: 300 sandboxes/sec at <200ms latency You run fewer than 10 concurrent agents

Architectural Tradeoffs

What We Gained

  • Durable execution: The event log makes crash recovery deterministic. No more “restart from scratch” failures.
  • Secure isolation: Sandboxed components prevent the most common production agent failures — code injection, host access, cross-tenant data leaks.
  • Framework agnosticism: AX works with any agent framework (LangGraph, ADK, A2A, custom). You are not locked into a specific orchestration model.
  • Auditability: Every step is recorded. You can replay, inspect, and verify any execution.
  • Resumable streams: Clients can disconnect and reconnect seamlessly. This is essential for mobile clients, web UIs, and long-running background tasks.

What We Sacrificed

  • Latency: The event log adds write overhead to every step. For sub-millisecond tool calls, this overhead is significant. AX is designed for agentic loops (seconds to minutes per step), not real-time systems.
  • Complexity: The distributed architecture (controller + remote agents + sandboxes + event log) is more complex than a single-process agent. You need Kubernetes or similar orchestration to manage it.
  • Maturity: v0.1.0 is pre-1.0. Breaking changes are expected. The API surface is not stable. External contributions are temporarily paused while the core stabilizes.
  • Go dependency: The primary implementation is Go. Teams without Go expertise will face a learning curve for custom agent development.
  • Storage: The append-only event log grows monotonically. Long-running agents with many steps produce large logs. You need a log management strategy (rotation, archival, compaction).

“The real lesson: AX is not a replacement for your agent framework. It is a runtime that sits under your framework. You use LangGraph or CrewAI to define what the agent does. You use AX to ensure that it runs reliably. The two layers are complementary, not competitive. Teams that try to use AX as a framework will be frustrated. Teams that use it as a runtime will wonder how they lived without it.”

Course-Style Deep Dive

Under the Hood: The Event Log and Resumption Protocol

The event log is the heart of AX’s reliability model. Every step of an agentic execution produces an ExecutionEvent:

message ExecutionEvent {
  string task_id = 2;
  string agent_id = 3;
  repeated Content inputs = 4;
  repeated Content outputs = 5;
  State state = 6;
  google.protobuf.Timestamp timestamp = 7;
}

Events are written to an append-only log. The default implementation uses SQLite for production:

type EventLog interface {
    Append(ctx context.Context, event *ExecutionEvent) error
    Read(ctx context.Context, afterSeq int64) ([]*ExecutionEvent, error)
    Close() error
}

When a client reconnects, it sends the last sequence number it received. The controller reads all events after that sequence from the log and replays them to the client. The client is caught up without re-executing any work.

The resumption protocol is built on gRPC bidirectional streaming. The client opens a Connect stream, sends its last known sequence number, and the server streams back missed events followed by new events as they occur.

Advanced Pattern 1: Sub-Agent Delegation with Task Executor

The TaskExecutor interface enables recursive agent delegation. A parent agent can spawn sub-agents, each with their own task context:

type Task struct {
    ID        string
    AgentID   string
    Inputs    []*proto.Content
    Rehydrate bool
    Config    *anypb.Any
}

type TaskExecutor interface {
    Exec(ctx context.Context, t *Task, o OutputHandler) error
}

A coding agent can delegate testing to a specialized test agent:

func (c *CodingAgent) Exec(ctx context.Context, t *Task, o OutputHandler) error {
    // Write code
    code := generateCode(t.Inputs)

    // Delegate testing to sub-agent
    subTask := &Task{
        ID:      uuid.New().String(),
        AgentID: "test-agent",
        Inputs:  []*proto.Content{{Parts: []*pb.Part{{Text: code}}}},
    }

    // Fan-out with errgroup for concurrent sub-tasks
    var g errgroup.Group
    g.Go(func() error {
        return c.taskExecutor.Exec(ctx, subTask, o)
    })
    return g.Wait()
}

The event log records both the parent and sub-agent events, creating a complete execution tree. If the sub-agent crashes, its events are preserved and the parent can retry.

Advanced Pattern 2: Trajectory Branching for Evaluation

Trajectory branching lets you fork an execution at any sequence point. This is implemented by copying the event log up to the branch point and starting a new writer:

func Branch(ctx context.Context, log EventLog, branchSeq int64) (EventLog, error) {
    events, err := log.Read(ctx, 0)
    if err != nil {
        return nil, err
    }

    // Filter events up to branch point
    branchEvents := events[:branchSeq]

    // Create new log seeded with branch events
    branchLog := NewSQLLiteLog("branch_" + uuid.New().String())
    for _, e := range branchEvents {
        branchLog.Append(ctx, e)
    }
    return branchLog, nil
}

The branched execution shares all context up to the branch point but diverges from there. This is used for A/B testing agent configurations, debugging regressions, and exploring alternative decision paths without re-executing the shared prefix.

Production Considerations

  • Event log storage: SQLite works for single-node deployments. For multi-node, use a shared filesystem (NFS, EFS) or a database-backed log. The NDJSON file format is also supported for simpler setups.
  • Log growth: The append-only log grows without bound. Implement log rotation or archival. Consider snapshotting full state periodically and truncating the log before the snapshot point.
  • Sandbox overhead: gVisor and Kata Containers add startup latency. For high-throughput scenarios, pre-warm a pool of sandboxes. Agent Substrate handles this automatically.
  • Authentication: The current release has minimal auth. In production, wrap the gRPC endpoints with mTLS or a service mesh. Use workload identity for cloud deployments.
  • Monitoring: Export event log metrics (write latency, log size, event count) to your observability stack. Monitor for slow event writes — they are the canary for storage bottlenecks.

The Results

Metric Without AX With AX
Crash recovery time Full restart (minutes) Resume from last checkpoint (seconds)
State corruption incidents Common with concurrent agents Eliminated (single-writer)
Client disconnect recovery Impossible — lost session Seamless catch-up via sequence numbers
Audit trail completeness Ad-hoc logging Every step recorded in append-only log
Multi-tenant isolation Namespace-level Per-component sandboxing
Agent evaluation Manual, error-prone Trajectory branching, programmatic
Sandbox provisioning N/A (no isolation) ~300/sec at <200ms latency (Substrate)

What this means for you: If you are deploying agents to production today, you are almost certainly operating without a proper runtime. You are one pod eviction, one state corruption, or one security incident away from learning this lesson the hard way. Google AX gives you the runtime layer that the agent ecosystem has been missing. It is not finished — v0.1.0 is a preview — but the architecture is sound, the design decisions are well-motivated, and the direction is exactly what the industry needs.”

What to Watch Out For

  1. Do not confuse AX with an agent framework. AX is a runtime. You still need LangGraph, CrewAI, ADK, or a custom harness to define your agent’s logic. AX runs underneath that logic. Trying to use AX as a framework will lead to frustration.

  2. The event log is not free. Every step incurs a write. For agents that make hundreds of rapid tool calls, the log overhead adds up. Profile your agent’s step frequency and ensure your storage can keep up.

  3. v0.1.0 is a preview. Breaking changes are expected. The API surface is not stable. Pin your dependency to a specific commit and plan for migration. External contributions are paused while the core stabilizes.

  4. You need Kubernetes for the full story. AX runs standalone, but the isolation, scaling, and snapshotting features require Agent Substrate on Kubernetes. If you are not on K8s, you lose the most powerful capabilities.

  5. The Go dependency is real. The primary implementation is Go. Remote agents can be written in any language (gRPC is language-agnostic), but the controller and core tooling are Go. Your team needs Go expertise for custom development.

  6. Log management is your problem. The append-only event log grows without bound. You need a strategy for rotation, archival, and compaction. AX does not handle this for you.

  7. Authentication is minimal. The current release has no built-in auth. In production, wrap with mTLS or a service mesh. Do not expose the gRPC endpoint to the public internet without protection.

“The biggest lesson from our early testing: start with the event log. Configure SQLite persistence on day one, even in development. The in-memory default is convenient for quick experiments, but it hides the durability model that makes AX valuable. If you do not test with persistence, you will not discover log-related bottlenecks until they hit in production.”

“Second lesson: use trajectory branching from the start. It is the feature that most teams overlook and then wish they had used earlier. Every time you debug an agent’s behavior, ask yourself: ‘Could I have answered this faster with a branch?’ The answer is almost always yes.”

Getting Started: Install the binary, run ax exec --input "Hello agents!", then configure a persistent event log. From there, register a remote agent and experiment with trajectory branching. The full power of AX reveals itself when you start treating agent executions as durable, auditable, forkable artifacts rather than ephemeral chat sessions.


Next in the Open-Source AI Tools Mastery series: CrewAI

NL

Written by Nivant Labs Team

Engineer at Nivant Labs

Share this post