Kubernetes-native Context & Memory Gateway Sidecar for Google AX (Agent Executor)
ax-context-gateway is a production reference controller and sidecar that bridges Swarm-Context-Commander's bounded context compilation and token admission with Google AX's declarative agent orchestrator on Kubernetes.
Google AX introduces declarative Kubernetes primitives (Task, Workspace, Gateway, Model) and the Agent Substrate runtime for high-density, stateful agent multiplexing. Agents run in gVisor (runsc) sandboxes under strict zero-trust network policies (egress: mode: DenyAll).
ax-context-gateway resolves the fundamental friction between zero-trust sandboxing and stateful memory retrieval by serving pre-flight compiled context and state synchronization over an in-pod Unix Domain Socket Model Context Protocol (MCP) stream.
flowchart TD
subgraph ControlPlane ["Kubernetes Control Plane"]
AX_CTRL["Google AX Controller\n(Task, Workspace, Model)"]
ACB_CRD["AxContextBinding CRD\n(Tenant, Agent, Token Budget)"]
GW_CTRL["AxContextController\n(Reconciliation & Admission)"]
AX_CTRL <--> ACB_CRD
ACB_CRD --> GW_CTRL
end
subgraph Node ["Kubernetes Worker Node (gVisor runsc)"]
subgraph Pod ["Agent Pod"]
subgraph Vol ["Shared Memory Volume (emptyDir: Memory)"]
SOCK["/var/run/ax-context/mcp.sock\n(Zero-Network AF_UNIX IPC)"]
end
subgraph Sidecar ["Sidecar Container: ax-context-gateway"]
COMP["Bounded Context Compiler\n(Token Bounding & Tombstones)"]
MCP_SRV["UnixSocketMcpSidecar\n(JSON-RPC 2.0 MCP Server)"]
SNAP["State Snapshotter\n(SHA-256 Checkpointing)"]
COMP --> MCP_SRV
SNAP --> MCP_SRV
end
subgraph Agent ["Workload Container: Agent Runner"]
RUNNER["Autonomous Agent Runner\n(Python SDK / LangChain / Swarm)"]
CLIENT["AxContextClient\n(Zero-Egress MCP Client)"]
RUNNER --> CLIENT
end
MCP_SRV <== "duplex stream" ==> SOCK
CLIENT <== "duplex stream" ==> SOCK
end
end
GW_CTRL -->|"Pre-flight Context & Volume Prep"| Sidecar
AX_CTRL -->|"Suspend / Resume Signal"| GW_CTRL
GW_CTRL -->|"Trigger State Checkpoint"| SNAP
The gateway aligns directly with the Google AX Agent Substrate lifecycle states:
stateDiagram-v2
[*] --> Admitted: AX Task Created
Admitted --> Bound: AxContextBinding Reconciled
state Bound {
[*] --> CompileContext
CompileContext --> EnforceTenantBoundary
EnforceTenantBoundary --> FilterTombstones
FilterTombstones --> BoundTokenCapacity
BoundTokenCapacity --> StartUnixSocketSidecar
}
Bound --> Running: Pod Scheduled & Agent Awakens
state Running {
AgentQueryGraph --> InPodMcpSocket
AgentMutateState --> RecordMemoryDelta
}
Running --> Suspending: AX Actor Preempted / Sleep Signal
state Suspending {
CaptureWorkingMemory --> ComputeSha256Digest
ComputeSha256Digest --> SealCheckpointFile
}
Suspending --> Suspended: Pod Memory Released
Suspended --> Resuming: Task Triggered with New Input
state Resuming {
LoadCheckpoint --> VerifyIntegrityHash
VerifyIntegrityHash --> FastRehydrationSub5ms
}
Resuming --> Running: Agent Continues Execution
Running --> Finalized: Task Completed / Terminated
Finalized --> [*]
Context is compiled deterministically prior to container unfreezing, ensuring zero cross-tenant contamination:
flowchart LR
Q["Task Initial Query & ContextPolicy"] --> T_GATE{"Tenant Isolation\nBoundary Gate"}
T_GATE -->|"Mismatch Tenant ID"| DROP["Drop Immediately"]
T_GATE -->|"Valid Tenant"| TOMB{"Tombstone\nFilter"}
TOMB -->|"Deleted / Tombstoned"| EXCLUDE["Exclude from Knowledge"]
TOMB -->|"Active Node"| SIM["Lexical & Seed Scoring\n(Jaccard Similarity)"]
SIM --> HOP["Multi-Hop Graph Expansion\n(max_graph_hops)"]
HOP --> CAP{"Token Capacity Gate\n(total <= max_tokens)"}
CAP -->|"Capacity Exceeded"| TRUNC["Set budget_exhausted = True\nStop Frontier"]
CAP -->|"Within Budget"| ADMIT["Admit Node & Edges"]
ADMIT --> HASH["Compute Canonical SHA-256 Digest\n(context_hash)"]
TRUNC --> HASH
HASH --> BUNDLE["CompiledContextPackage\nInjected to MCP Sidecar"]
In contrast to traditional agent architectures that make external network calls to retrieve memory, ax-context-gateway maintains an airtight perimeter:
graph TD
subgraph Internet ["Public Internet / Untrusted Network"]
EXT_DB[("External Vector DBs")]
EXT_API["Third-Party Web APIs"]
end
subgraph KubernetesCluster ["GKE / Kubernetes Cluster"]
subgraph NetworkPolicy ["Cilium / Calico Zero-Trust Egress"]
FW["Egress Rule: DenyAll (Port * Drop)"]
end
subgraph Sandbox ["gVisor (runsc) Sandbox Boundary"]
subgraph AgentApp ["Workload Container"]
A_EXEC["Agent Code / LLM Loop"]
end
subgraph IPC ["Kernel AF_UNIX IPC"]
USOCK["/var/run/ax-context/mcp.sock"]
end
subgraph SidecarApp ["Sidecar Container"]
GATEWAY["ax-context-gateway Sidecar"]
end
end
end
A_EXEC -.->|"TCP Connect (BLOCKED)"| FW
FW -.->|"DROPPED"| EXT_DB
FW -.->|"DROPPED"| EXT_API
A_EXEC ===|"Zero-Network Unix Domain Socket"| USOCK
USOCK ===|"Local In-Memory IPC"| GATEWAY
apiVersion: agent.google.com/v1alpha1
kind: AxContextBinding
metadata:
name: incident-analyst-task-42-binding
namespace: agent-workloads
spec:
targetTaskRef:
name: incident-analyst-task-42
tenantId: enterprise-acme-corp
agentId: secops-reconciliation-agent
contextPolicy:
maxContextTokens: 4096
maxGraphHops: 2
similarityThreshold: 0.75
enforceTombstones: true
snapshotConfig:
autoCheckpointOnSuspend: true
storageUri: "s3://acme-agent-snapshots/secops-reconciliation-agent/"apiVersion: agent.google.com/v1alpha1
kind: Task
metadata:
name: incident-analyst-task-42
namespace: agent-workloads
spec:
modelRef:
name: gemini-2-5-pro
workspaceRef:
name: analyst-scratchpad-ws
runtimeClass: gvisor # runsc sandbox isolation
networkPolicy:
egress:
mode: DenyAll # Strict Zero-Trust
template:
containers:
- name: agent-runner
image: ghcr.io/aah20/autonomous-analyst:v1.0.0
env:
- name: MCP_SOCKET_PATH
value: /var/run/ax-context/mcp.sock
volumeMounts:
- name: ax-context-socket
mountPath: /var/run/ax-context
volumes:
- name: ax-context-socket
emptyDir:
medium: MemoryAgents inside the container use AxContextClient with zero boilerplate:
from ax_context_gateway.client import AxContextClient
# Auto-detects $MCP_SOCKET_PATH
client = AxContextClient.from_env()
# 1. Fetch pre-flight compiled context & graph
context = client.get_bounded_context()
print(f"Admitted tokens: {context['total_tokens']} (Hash: {context['context_hash']})")
# 2. Query in-memory graph without network access
nodes = client.query_memory_graph(keyword="database credentials")
# 3. Record state mutations before AX suspension
client.record_memory_delta(key="incident_status", value="mitigation_in_progress")
# 4. Inspect session metrics & budget
metrics = client.get_session_metrics()
print("Session metrics:", metrics)The repository includes AxContextEvaluator to audit multi-tenant boundaries and benchmark rehydration latency:
PYTHONPATH=projects python3 -m ax_context_gateway.evaluator{
"suite": "ax_context_gateway_conformance_eval",
"all_passed": true,
"drills": [
{
"test_name": "tenant_isolation_leak_drill",
"passed": true,
"victim_tenant": "tenant-bank-corp",
"attacker_tenant": "tenant-adversary",
"leaked_count": 0,
"status": "SECURE"
},
{
"test_name": "token_starvation_bounding_drill",
"passed": true,
"max_budget": 500,
"total_admitted_tokens": 500,
"admitted_node_count": 10,
"budget_exhausted": true,
"status": "BOUNDED"
},
{
"test_name": "rehydration_latency_benchmark",
"passed": true,
"checkpoint_latency_microseconds": 360.83,
"avg_rehydration_microseconds": 96.56,
"min_rehydration_microseconds": 60.58,
"max_rehydration_microseconds": 401.33,
"iterations": 50,
"status": "OPTIMAL (<5ms)"
}
]
}PYTHONPATH=projects python3 -m unittest discover -s projects/ax_context_gateway/tests -vtest_client_duplex_session (test_evaluator.TestEvaluatorAndClient.test_client_duplex_session) ... ok
test_evaluator_drills (test_evaluator.TestEvaluatorAndClient.test_evaluator_drills) ... ok
test_suspension_checkpoint_and_rehydration (test_gateway.TestAxContextGateway.test_suspension_checkpoint_and_rehydration) ... ok
test_tampered_checkpoint_rejection (test_gateway.TestAxContextGateway.test_tampered_checkpoint_rejection) ... ok
test_tenant_isolation_and_tombstones (test_gateway.TestAxContextGateway.test_tenant_isolation_and_tombstones) ... ok
test_token_budget_enforcement (test_gateway.TestAxContextGateway.test_token_budget_enforcement) ... ok
test_unix_socket_mcp_zero_egress (test_gateway.TestAxContextGateway.test_unix_socket_mcp_zero_egress) ... ok
Ran 7 tests in 0.019s (OK)
Apache-2.0