server, store: move to Postgres and read actor state from Substrate

.github/CODEOWNERS

@@ -0,0 +1,6 @@
+# Changes under these paths need a review from their owners.
+# https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
+
+/internal/server/     @rakyll
+/internal/store/      @rakyll
+/internal/controller/ @rakyll

.github/workflows/go.yml

@@ -26,6 +26,20 @@ permissions:
 jobs:
   build-and-test:
     runs-on: ubuntu-latest
+    services:
+      postgres:
+        image: postgres:17-alpine
+        env:
+          POSTGRES_USER: ax
+          POSTGRES_PASSWORD: ax
+          POSTGRES_DB: ax
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U ax -d ax"
+          --health-interval 5s
+          --health-timeout 5s
+          --health-retries 10
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
 
@@ -46,4 +60,6 @@ jobs:
         go build -v ./cmd/ax-task-runner
 
     - name: Test
+      env:
+        AX_TEST_POSTGRES_URL: postgres://ax:ax@localhost:5432/ax?sslmode=disable
       run: go test -v ./...

CONTRIBUTING.md

@@ -90,7 +90,7 @@ make push-task-runner TASK_RUNNER_REPO=gcr.io/<your-project>/ax-task-runner
 
 ### Deploying to Kubernetes
 
-Deploy Redis and server components to your cluster in the `ax-system` namespace:
+Deploy PostgreSQL and server components to your cluster in the `ax-system` namespace:
 
 ```bash
 make deploy AX_IMAGE_REPO=gcr.io/<your-project>/ax-images

DESIGN.md

@@ -2,7 +2,9 @@
 
 ## Architecture
 
-Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its comfort zone (single-digit GB storage limits, write-rate bottlenecks, control plane degradation). AX stores its state in Redis and reconciles directly with Agent Substrate under fine-grained distributed locks.
+Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its comfort zone (single-digit GB storage limits, write-rate bottlenecks, control plane degradation). AX stores its state in PostgreSQL and reconciles directly with Agent Substrate under fine-grained distributed locks.
+
+PostgreSQL holds resources as protojson documents alongside the leases that back AX's distributed locks. Watches and lock waiters poll; nothing depends on `LISTEN/NOTIFY`, so any connection pooler works. `ax-server` creates any missing tables when it starts; there are no versioned migrations yet. The bundled manifest runs a single-instance StatefulSet with a persistent volume; any reachable PostgreSQL works. Storage sits behind the `store.Store` and `lock.Locker` interfaces. The store holds only what AX owns: specs, metadata, and the parts of a task's status that describe AX's own work (`Pending`, `Failed`, `Terminating`, the task ID, workspace setup). Everything that belongs to the task's actor (its phase while provisioned, `actor`, `workerIP`, and the `Ready` condition) is read from Agent Substrate whenever a task is fetched, listed, or watched, so AX never keeps a copy that can drift from the actor. The split lives in a few helpers in `internal/server/server.go` (`getTask`, `listTasks`, `saveTask`, `updateTask`, `deleteTask`), so storage implementations never need to know about Substrate.
 
 ```
                       ax apply -f task.yaml
@@ -13,9 +15,9 @@ Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its co
                                 │
          ┌──────────────────────┴──────────────────────┐
          ▼                                             ▼
-       Redis                                    Agent Substrate
+     PostgreSQL                                 Agent Substrate
  (Resource Store,                       ┌───────────────────────────────┐
-  Locks, PubSub)                        │ • Atespace Provisioning       │
+  Locks)                                │ • Atespace Provisioning       │
                                         │ • Actor Creation & Activation │
                                         │ • Worker Assignment           │
                                         └───────────────────────────────┘
@@ -26,7 +28,7 @@ Storing millions of short-lived tasks as Kubernetes CRDs pushes etcd past its co
 | Binary | Role |
 |---|---|
 | `ax` | Developer CLI. Applies manifests, inspects and watches resources, tunnels to the cluster. |
-| `ax-server` | Direct-execution gRPC API on port 8080. Validates manifests, manages distributed locks, reconciles directly with Agent Substrate, and persists state to Redis. |
+| `ax-server` | Direct-execution gRPC API on port 8080. Validates manifests, manages distributed locks, reconciles directly with Agent Substrate, and persists state to PostgreSQL. |
 | `ax-task-runner` | Entrypoint inside every task container. Bootstraps the workspace, serves metadata, and runs the agent command. A thin wrapper over the `runner` package, which custom images can embed directly. |
 
 ## API reference

Makefile

@@ -19,7 +19,7 @@ AX_IMAGE_REPO ?= gcr.io/ax-substrate/ate-images
 TASK_RUNNER_REPO ?= $(AX_IMAGE_REPO)/ax-task-runner
 CONTAINER_CLI ?= $(shell which podman 2>/dev/null || which docker 2>/dev/null)
 
-.PHONY: all build build-binaries build-task-runner install push push-task-runner deploy deploy-server deploy-redis apply-example test test-e2e test-e2e-head e2e-setup e2e-run e2e-cleanup clean
+.PHONY: all build build-binaries build-task-runner install push push-task-runner deploy deploy-server deploy-postgres apply-example test test-e2e test-e2e-head e2e-setup e2e-run e2e-cleanup clean
 
 all: build
 
@@ -63,12 +63,19 @@ push: push-task-runner
 ## Deployment Targets
 ## --------------------------------------
 
-# Deploy all AX components to Kubernetes (Redis, ax-server)
-deploy: deploy-redis deploy-server
-
-deploy-redis:
-	@echo "==> Deploying Redis to ax-system namespace..."
-	kubectl apply -f deploy/redis.yaml
+# Deploy all AX components to Kubernetes (Postgres, ax-server)
+deploy: deploy-postgres deploy-server
+
+deploy-postgres:
+	@echo "==> Deploying Postgres to ax-system namespace..."
+	kubectl create namespace ax-system --dry-run=client -o yaml | kubectl apply -f -
+	@# Create the password once. Postgres only reads it when initializing an
+	@# empty volume, so it must not change on later deploys.
+	@kubectl -n ax-system get secret ax-postgres > /dev/null 2>&1 || \
+		kubectl -n ax-system create secret generic ax-postgres \
+			--from-literal=password="$$(openssl rand -hex 24)"
+	kubectl apply -f deploy/postgres.yaml
+	kubectl -n ax-system rollout status statefulset/ax-postgres --timeout=180s
 
 deploy-server:
 	@echo "==> Building and deploying ax-server using ko..."

README.md

@@ -88,7 +88,7 @@ With the [prerequisites](#prerequisites) in place — most importantly a reachab
 make deploy AX_IMAGE_REPO=<your-registry>
 ```
 
-This deploys Redis, then builds and deploys the control plane images with `ko`. Everything lands in the `ax-system` namespace.
+This deploys PostgreSQL, then builds and deploys the control plane images with `ko`. Everything lands in the `ax-system` namespace.
 
 ### 3. Run your first task
 

cmd/ax-server/main.go

@@ -19,25 +19,23 @@ import (
 	"flag"
 	"log/slog"
 	"net/http"
+	"net/url"
 	"os"
 	"os/signal"
 	"syscall"
 	"time"
 
 	"github.com/google/ax/internal/controller"
-	"github.com/google/ax/internal/lock"
 	"github.com/google/ax/internal/server"
-	"github.com/google/ax/internal/store/redis"
+	"github.com/google/ax/internal/store/postgres"
 	"github.com/google/ax/internal/substrate"
 	"github.com/google/ax/pkg/apis/v1alpha1"
-	goredis "github.com/redis/go-redis/v9"
 )
 
 func main() {
 	var (
 		listenAddr              string
-		redisAddr               string
-		redisPassword           string
+		storeURL                string
 		substrateEndpoint       string
 		substrateAuthority      string
 		substrateTokenFile      string
@@ -50,8 +48,7 @@ func main() {
 	)
 
 	flag.StringVar(&listenAddr, "addr", ":8080", "HTTP listen address")
-	flag.StringVar(&redisAddr, "redis-addr", "localhost:6379", "Redis server address")
-	flag.StringVar(&redisPassword, "redis-password", "", "Redis password")
+	flag.StringVar(&storeURL, "store", "postgres://ax@localhost:5432/ax?sslmode=disable", "PostgreSQL connection URL (env: AX_STORE). libpq variables such as PGPASSWORD fill in what it leaves out")
 	flag.StringVar(&substrateEndpoint, "substrate-endpoint", "api.ate-system.svc.cluster.local:443", "Agent Substrate Control API endpoint")
 	flag.StringVar(&substrateAuthority, "substrate-authority", "api.ate-system.svc", "Authority / TLS ServerName for Substrate endpoint")
 	flag.StringVar(&substrateTokenFile, "substrate-token-file", "", "Path to bearer token file for Substrate auth")
@@ -66,11 +63,11 @@ func main() {
 	if envAddr := os.Getenv("ADDR"); envAddr != "" {
 		listenAddr = envAddr
 	}
-	if envRedis := os.Getenv("REDIS_ADDR"); envRedis != "" {
-		redisAddr = envRedis
-	}
-	if envPass := os.Getenv("REDIS_PASSWORD"); envPass != "" {
-		redisPassword = envPass
+	// TODO: Reject or warn on passwords embedded in the store URL. They end up
+	// in process args and pod specs; credentials should come from PGPASSWORD,
+	// a password file, or a Secret instead.
+	if envStore := os.Getenv("AX_STORE"); envStore != "" {
+		storeURL = envStore
 	}
 	if envImage := os.Getenv("AX_DEFAULT_TASK_IMAGE"); envImage != "" {
 		defaultTaskImage = envImage
@@ -81,20 +78,20 @@ func main() {
 
 	slog.Info("starting ax-server",
 		"listenAddr", listenAddr,
-		"redisAddr", redisAddr,
+		"store", redactURL(storeURL),
 		"substrateEndpoint", substrateEndpoint,
 		"template", defaultTemplate,
 		"defaultTaskImage", defaultTaskImage,
 	)
 
-	rClient := goredis.NewClient(&goredis.Options{
-		Addr:     redisAddr,
-		Password: redisPassword,
-	})
-	defer rClient.Close()
-
-	rStore := redis.NewStore(rClient, redis.Options{})
-	rLocker := lock.NewRedisLocker(rClient, lock.RedisLockerOptions{})
+	openCtx, openCancel := context.WithTimeout(context.Background(), time.Minute)
+	backend, err := postgres.Open(openCtx, storeURL, postgres.Options{})
+	openCancel()
+	if err != nil {
+		slog.Error("opening store", "store", redactURL(storeURL), "error", err)
+		os.Exit(1)
+	}
+	defer backend.Close()
 
 	var reconciler server.Reconciler
 	subClient, err := substrate.NewClientWithOptions(substrate.ClientOptions{
@@ -112,9 +109,10 @@ func main() {
 		reconciler = controller.NewTaskReconciler(subClient, defaultTemplate, defaultTemplateAtespace, defaultTaskImage)
 	}
 
-	srv := server.NewServer(rStore, server.Options{
-		Locker:     rLocker,
+	srv := server.NewServer(backend, server.Options{
+		Locker:     backend,
 		Reconciler: reconciler,
+		Substrate:  subClient,
 	})
 
 	httpServer := &http.Server{
@@ -142,3 +140,12 @@ func main() {
 	defer shutdownCancel()
 	_ = httpServer.Shutdown(shutdownCtx)
 }
+
+// redactURL hides any password in a store URL so that it can be logged.
+func redactURL(raw string) string {
+	u, err := url.Parse(raw)
+	if err != nil {
+		return "<unparseable>"
+	}
+	return u.Redacted()
+}

cmd/ax/main.go

@@ -64,7 +64,7 @@ func main() {
 	// Commands that don't require an AX server connection
 	switch cmd {
 	case "version":
-		fmt.Println("ax version v1alpha1 (standalone redis engine)")
+		fmt.Println("ax version v1alpha1")
 		return
 	case "help", "-h", "--help":
 		printUsage()
@@ -417,7 +417,7 @@ func runGet(serverURL, atespace string, args []string) error {
 		for _, t := range tasks {
 			workerIP := ""
 			actor := ""
-			phase := "Pending"
+			phase := v1alpha1.PhasePending
 			if t.Status != nil {
 				workerIP = t.Status.WorkerIp
 				actor = t.Status.Actor
@@ -769,7 +769,7 @@ func runWatch(serverURL, atespace string, args []string) error {
 				actor,
 				workerIP,
 			)
-			if phase == "Completed" || phase == "Failed" {
+			if phase == v1alpha1.PhaseFailed {
 				fmt.Printf("Task reached terminal phase %q.\n", phase)
 				break
 			}
@@ -871,9 +871,9 @@ func manifestFromArgs(args []string) (data []byte, ok bool, err error) {
 // lifecycleRPCTimeout bounds suspend/resume RPCs, which wait for Substrate to
 // snapshot or restore the task's sandbox. MicroVM-backed sandboxes can take
 // well over a minute to snapshot/restore, so this is deliberately generous.
-// The server renews its per-task Redis lock for as long as the RPC is in
+// The server renews its per-task lock lease for as long as the RPC is in
 // flight, so this deadline is not constrained by the lock TTL (see
-// lock.RedisLocker).
+// postgres.Store.Lock).
 const lifecycleRPCTimeout = 5 * time.Minute
 
 func runSuspend(serverURL, atespace string, args []string) error {
@@ -1091,7 +1091,7 @@ func runSSH(serverURL, atespace, kubeContext string, args []string) error {
 		return fmt.Errorf("task %q has no status available", taskName)
 	}
 
-	if task.Status.Phase != "Running" {
+	if task.Status.Phase != v1alpha1.PhaseRunning {
 		return fmt.Errorf("task %q is in phase %q (must be Running to ssh)", taskName, task.Status.Phase)
 	}
 

deploy/ax-server.yaml

@@ -72,14 +72,19 @@ spec:
           imagePullPolicy: IfNotPresent
           args:
             - "--addr=:8080"
-            - "--redis-addr=ax-redis.ax-system.svc.cluster.local:6379"
+            - "--store=postgres://ax@ax-postgres.ax-system.svc.cluster.local:5432/ax?sslmode=disable"
             - "--substrate-endpoint=api.ate-system.svc.cluster.local:443"
             - "--substrate-authority=api.ate-system.svc"
             - "--substrate-token-file=/var/run/secrets/ateapi/token"
             - "--substrate-ca-file=/run/servicedns-ca/trust-bundle.pem"
             - "--template=default-template"
             - "--template-atespace=ax-system"
           env:
+            - name: PGPASSWORD
+              valueFrom:
+                secretKeyRef:
+                  name: ax-postgres
+                  key: password
             - name: ATENET_ROUTER_ADDR
               value: "atenet-router.ate-system.svc.cluster.local:80"
             - name: AX_SNAPSHOTS_BUCKET