api,cli: report task residency, actor UID, and outcome

DESIGN.md

@@ -4,7 +4,7 @@
 
 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. 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 or listed, 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.
+PostgreSQL holds resources as protojson documents alongside the leases that back AX's distributed locks. 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 or listed, so AX never keeps a copy that can drift from the actor. The exceptions are `residency` and `actorUid`: the store keeps the values Substrate last returned to ax-server, so that each change AX makes is a stored change, while reads still report the actor's current state. 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

README.md

@@ -120,8 +120,8 @@ From your ax checkout:
 ax apply -f examples/simple.yaml         # the smallest useful Task
 ax resume task simple-task               # new tasks start Suspended; resume schedules the actor and starts it
 ax get tasks
-# NAME          ATESPACE   PHASE     ACTOR         WORKER-IP    AGE
-# simple-task   default    Running   simple-task   10.20.3.67   1m
+# NAME          ATESPACE   PHASE     RESIDENCY   OUTCOME   ACTOR         WORKER-IP    AGE
+# simple-task   default    Running   running     <none>    simple-task   10.20.3.67   1m
 
 ax ssh simple-task -- ls -la /workspace  # poke around inside the sandbox
 ax suspend task simple-task              # checkpoint and pause

cmd/ax/main.go

@@ -414,7 +414,7 @@ func runGet(serverURL, atespace string, args []string) error {
 		tasks := resp.Tasks
 
 		w := tabwriter.NewWriter(os.Stdout, 0, 8, 3, ' ', 0)
-		fmt.Fprintln(w, "NAME\tATESPACE\tPHASE\tACTOR\tWORKER-IP\tAGE")
+		fmt.Fprintln(w, "NAME\tATESPACE\tPHASE\tRESIDENCY\tOUTCOME\tACTOR\tWORKER-IP\tAGE")
 		for _, t := range tasks {
 			workerIP := ""
 			actor := ""
@@ -426,12 +426,6 @@ func runGet(serverURL, atespace string, args []string) error {
 					phase = t.Status.Phase
 				}
 			}
-			if workerIP == "" {
-				workerIP = "<none>"
-			}
-			if actor == "" {
-				actor = "<none>"
-			}
 			age := "<unknown>"
 			name := ""
 			tAtespace := ""
@@ -442,12 +436,14 @@ func runGet(serverURL, atespace string, args []string) error {
 					age = formatAge(time.Since(t.Metadata.CreateTime.AsTime()))
 				}
 			}
-			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
+			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
 				name,
 				tAtespace,
 				phase,
-				actor,
-				workerIP,
+				orNone(t.GetStatus().GetResidency()),
+				orNone(t.GetStatus().GetOutcome()),
+				orNone(actor),
+				orNone(workerIP),
 				age,
 			)
 		}
@@ -695,7 +691,10 @@ func runDescribe(serverURL, atespace string, args []string) error {
 	fmt.Printf("Name:         %s\n", taskName)
 	fmt.Printf("Atespace:     %s\n", taskAtespace)
 	fmt.Printf("Phase:        %s\n", phase)
+	fmt.Printf("Residency:    %s\n", orNone(task.GetStatus().GetResidency()))
+	fmt.Printf("Outcome:      %s\n", orNone(task.GetStatus().GetOutcome()))
 	fmt.Printf("Actor:        %s\n", actor)
+	fmt.Printf("Actor UID:    %s\n", orNone(task.GetStatus().GetActorUid()))
 	fmt.Printf("Worker IP:    %s\n", workerIP)
 	if task.Spec != nil {
 		if ref := task.Spec.GetWorkspace(); ref != nil {
@@ -730,6 +729,14 @@ func runDescribe(serverURL, atespace string, args []string) error {
 	return nil
 }
 
+// orNone returns s, or "<none>" when s is empty, for table and describe output.
+func orNone(s string) string {
+	if s == "" {
+		return "<none>"
+	}
+	return s
+}
+
 // runDelete removes one resource by kind and name.
 func runDelete(serverURL, atespace string, args []string) error {
 	if len(args) < 2 {

cmd/ax/main_test.go

@@ -312,6 +312,75 @@ func TestRunGetResourceAliases(t *testing.T) {
 	}
 }
 
+// captureStdout runs fn with os.Stdout redirected to a file and returns what
+// fn wrote.
+func captureStdout(t *testing.T, fn func() error) string {
+	t.Helper()
+	output, err := os.CreateTemp(t.TempDir(), "stdout")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer output.Close()
+	stdout := os.Stdout
+	os.Stdout = output
+	runErr := fn()
+	os.Stdout = stdout
+	if runErr != nil {
+		t.Fatal(runErr)
+	}
+	data, err := os.ReadFile(output.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+	return string(data)
+}
+
+func TestRunGetAndDescribeShowResidencyAndOutcome(t *testing.T) {
+	sub := substratetest.NewServer(t)
+	srv := server.NewServer(memory.NewStore(), newTestReconciler(sub))
+	if _, err := srv.CreateTask(context.Background(), &v1alpha1.CreateTaskRequest{Task: &v1alpha1.Task{
+		Metadata: &v1alpha1.ObjectMeta{Name: "resident", Atespace: "default"},
+		Spec:     &v1alpha1.TaskSpec{Image: "alpine"},
+	}}); err != nil {
+		t.Fatal(err)
+	}
+	actor, _ := sub.Actor("default", "resident")
+	uid := actor.GetMetadata().GetUid()
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	grpcServer := srv.GRPCServer()
+	t.Cleanup(grpcServer.Stop)
+	go func() { _ = grpcServer.Serve(listener) }()
+	addr := listener.Addr().String()
+
+	out := captureStdout(t, func() error { return runGet(addr, "default", []string{"tasks"}) })
+	rows := strings.Split(strings.TrimSpace(out), "\n")
+	if len(rows) != 2 {
+		t.Fatalf("expected a header and one task, got %q", out)
+	}
+	header, row := strings.Fields(rows[0]), strings.Fields(rows[1])
+	want := map[string]string{"NAME": "resident", "PHASE": "Suspended", "RESIDENCY": "suspended", "OUTCOME": "<none>"}
+	for i, col := range header {
+		if w, ok := want[col]; ok && (i >= len(row) || row[i] != w) {
+			t.Errorf("column %s = %q, want %q; output:\n%s", col, row[i], w, out)
+		}
+		delete(want, col)
+	}
+	if len(want) != 0 {
+		t.Errorf("ax get tasks has no columns %v; output:\n%s", want, out)
+	}
+
+	out = captureStdout(t, func() error { return runDescribe(addr, "default", []string{"task", "resident"}) })
+	for _, line := range []string{"Residency:    suspended", "Outcome:      <none>", "Actor UID:    " + uid} {
+		if !strings.Contains(out, line+"\n") {
+			t.Errorf("ax describe task output has no line %q:\n%s", line, out)
+		}
+	}
+}
+
 // newTestReconciler returns a TaskReconciler that runs tasks on sub.
 func newTestReconciler(sub *substratetest.Server) *controller.TaskReconciler {
 	return controller.NewTaskReconciler(sub.Client(), "test-template", v1alpha1.DefaultSystemNamespace, "")

docs/concepts.md

@@ -17,6 +17,8 @@ The unit is deliberately small. An agent is not one process that runs to complet
 | `WorkspaceReady` | The workspace has finished setting up. Stays True afterwards. |
 | `Ready` | The task is running and `WorkspaceReady` is True. This is the one to wait on. |
 
+`status.residency` says where the task's actor is on Agent Substrate: `queued` (no actor yet), `resuming`, `running`, `suspending`, `suspended`, `reverting`, `crashed`, `deleting`, or `deleted`. Unlike the phase, it tells a crashed or reverting actor apart from a task AX failed to provision. `status.actorUid` is the actor's Substrate UID, the `ate.actor.uid` that Substrate's logs and telemetry carry. `status.outcome` will say how the command ended (`succeeded`, `failed`, or `cancelled`); AX does not detect command exit yet ([#346](https://github.com/google/ax/issues/346)), so it is empty for now. `ax get tasks` shows the residency and outcome; `ax describe task` also shows the actor UID.
+
 Two transitions are worth knowing. Suspending a task sets `Ready` to False with reason `Suspended`; resuming sets it back. Deleting a task tears down the sandbox on Agent Substrate and removes the record. `ax delete` blocks until that has completed.
 
 ## Workspace

internal/controller/reconciler.go

@@ -32,6 +32,7 @@ import (
 
 	"github.com/google/ax/internal/substrate"
 	"github.com/google/ax/pkg/apis/v1alpha1"
+	ateapipb "go.ate.dev/substrate/api/ateapi/v1"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
@@ -123,11 +124,11 @@ func newProbeHTTPClient() *http.Client {
 }
 
 // ensureSubstrateResources fills in the task's defaults and makes sure its
-// atespace, ActorTemplate and actor exist on Substrate. It records only what AX
-// stores in the task's status: the task ID, and on failure the reason in a
-// False Ready condition. The caller sets the phase, and the server reads the
-// rest of the status from the actor.
-func (r *TaskReconciler) ensureSubstrateResources(ctx context.Context, task *v1alpha1.Task, ws *v1alpha1.Workspace) error {
+// atespace, ActorTemplate and actor exist on Substrate, and returns the actor.
+// It records only what AX stores in the task's status: the task ID, and on
+// failure the reason in a False Ready condition. The caller sets the phase,
+// and the server reads the rest of the status from the actor.
+func (r *TaskReconciler) ensureSubstrateResources(ctx context.Context, task *v1alpha1.Task, ws *v1alpha1.Workspace) (*ateapipb.Actor, error) {
 	if task.Metadata == nil {
 		task.Metadata = &v1alpha1.ObjectMeta{}
 	}
@@ -155,10 +156,10 @@ func (r *TaskReconciler) ensureSubstrateResources(ctx context.Context, task *v1a
 	if err := r.client.EnsureAtespace(ctx, atespace); err != nil {
 		if substrate.IsTransientError(err) {
 			r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now)
-			return fmt.Errorf("ensuring atespace: %w", err)
+			return nil, fmt.Errorf("ensuring atespace: %w", err)
 		}
 		r.setNotReady(task, v1alpha1.ReasonAtespaceCreationFailed, err.Error(), now)
-		return fmt.Errorf("ensuring atespace: %w", err)
+		return nil, fmt.Errorf("ensuring atespace: %w", err)
 	}
 
 	// 2. The actor is always named after the task, so the two can be used
@@ -205,7 +206,7 @@ func (r *TaskReconciler) ensureSubstrateResources(ctx context.Context, task *v1a
 	// here as well as at apply time: a task must never run without limits it asked for.
 	if err := v1alpha1.ValidateResources(task.Spec.Resources); err != nil {
 		r.setNotReady(task, v1alpha1.ReasonInvalidResources, err.Error(), now)
-		return fmt.Errorf("validating resources: %w", err)
+		return nil, fmt.Errorf("validating resources: %w", err)
 	}
 	resources := substrate.ResourceLimits(task.Spec.Resources)
 
@@ -220,14 +221,14 @@ func (r *TaskReconciler) ensureSubstrateResources(ctx context.Context, task *v1a
 		if err != nil {
 			if substrate.IsTransientError(err) {
 				r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now)
-				return fmt.Errorf("ensuring actor template: %w", err)
+				return nil, fmt.Errorf("ensuring actor template: %w", err)
 			}
 			// Falling back to the default template would run a different image than
 			// the user asked for, without the limits it asked for, or fail later with
 			// a misleading "actor template not found". Surface the real error on the
 			// task and stop here.
 			r.setNotReady(task, v1alpha1.ReasonTemplateCreationFailed, substrateErrorMessage(err), now)
-			return fmt.Errorf("ensuring actor template: %w", err)
+			return nil, fmt.Errorf("ensuring actor template: %w", err)
 		}
 		if tmpl != nil && tmpl.Metadata != nil {
 			templateAtespace = tmpl.Metadata.Atespace
@@ -236,47 +237,67 @@ func (r *TaskReconciler) ensureSubstrateResources(ctx context.Context, task *v1a
 		}
 	}
 
-	_, err := r.client.EnsureActor(ctx, atespace, actorName, templateAtespace, templateName)
+	actor, err := r.client.EnsureActor(ctx, atespace, actorName, templateAtespace, templateName)
 	if err != nil {
 		if substrate.IsTransientError(err) {
 			r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now)
-			return fmt.Errorf("ensuring actor: %w", err)
+			return nil, fmt.Errorf("ensuring actor: %w", err)
 		}
 		r.setNotReady(task, v1alpha1.ReasonActorCreationFailed, err.Error(), now)
-		return fmt.Errorf("ensuring actor: %w", err)
+		return nil, fmt.Errorf("ensuring actor: %w", err)
+	}
+	return actor, nil
+}
+
+// recordActor records the UID and residency of the task's actor, as Substrate
+// returned them, in the task's status. It leaves both unchanged when Substrate
+// returned no actor.
+func recordActor(task *v1alpha1.Task, actor *ateapipb.Actor) {
+	if uid := actor.GetMetadata().GetUid(); uid != "" {
+		task.Status.ActorUid = uid
+	}
+	if residency := substrate.Residency(actor); residency != "" {
+		task.Status.Residency = residency
 	}
-	return nil
 }
 
-// Suspend runs ensureSubstrateResources, then suspends the task's actor. A
-// failure to suspend the actor is returned without changing the task's status.
+// Suspend runs ensureSubstrateResources, then suspends the task's actor and
+// records the actor's UID and residency in the task's status. A failure to
+// suspend the actor is returned without changing the task's status.
 func (r *TaskReconciler) Suspend(ctx context.Context, task *v1alpha1.Task, ws *v1alpha1.Workspace) error {
-	if err := r.ensureSubstrateResources(ctx, task, ws); err != nil {
+	ensured, err := r.ensureSubstrateResources(ctx, task, ws)
+	if err != nil {
 		return err
 	}
 	atespace, actorName := task.Metadata.Atespace, task.Metadata.Name
 	slog.Info("suspending actor on Substrate", "actor", actorName)
-	if err := r.client.SuspendActor(ctx, atespace, actorName); err != nil {
+	suspended, err := r.client.SuspendActor(ctx, atespace, actorName)
+	if err != nil {
 		return fmt.Errorf("suspending actor: %w", err)
 	}
+	recordActor(task, ensured)
+	recordActor(task, suspended)
 	slog.Info("task successfully suspended", "name", task.Metadata.Name, "actor", actorName)
 	return nil
 }
 
-// Resume runs ensureSubstrateResources, resumes the task's actor, and waits up
-// to WorkspaceReadyTimeout for the task's workspace to be set up, recording the
+// Resume runs ensureSubstrateResources, resumes the task's actor, records the
+// actor's UID and residency in the task's status, and waits up to
+// WorkspaceReadyTimeout for the task's workspace to be set up, recording the
 // result in the WorkspaceReady condition. A failure to resume the actor is
 // recorded the way ensureSubstrateResources records its own.
 func (r *TaskReconciler) Resume(ctx context.Context, task *v1alpha1.Task, ws *v1alpha1.Workspace) error {
-	if err := r.ensureSubstrateResources(ctx, task, ws); err != nil {
+	ensured, err := r.ensureSubstrateResources(ctx, task, ws)
+	if err != nil {
 		return err
 	}
+	recordActor(task, ensured)
 	atespace, actorName := task.Metadata.Atespace, task.Metadata.Name
 	now := time.Now()
 
 	// Resume the Actor to activate container execution
 	slog.Info("resuming actor on Substrate worker", "actor", actorName)
-	_, workerIP, err := r.client.ResumeActor(ctx, atespace, actorName)
+	actor, workerIP, err := r.client.ResumeActor(ctx, atespace, actorName)
 	if err != nil {
 		if substrate.IsTransientError(err) {
 			r.setNotReady(task, v1alpha1.ReasonSubstrateUnavailable, err.Error(), now)
@@ -285,6 +306,7 @@ func (r *TaskReconciler) Resume(ctx context.Context, task *v1alpha1.Task, ws *v1
 		r.setNotReady(task, v1alpha1.ReasonActorResumeFailed, err.Error(), now)
 		return fmt.Errorf("resuming actor: %w", err)
 	}
+	recordActor(task, actor)
 
 	// Check if workspace setup inside the actor has completed
 	host := workerIP

internal/server/actor_status_test.go

@@ -126,3 +126,42 @@ func TestSetActorStatus_SubstrateErrorReportsUnknown(t *testing.T) {
 		t.Errorf("phase = %q, want %q to be kept", got, v1alpha1.PhaseTerminating)
 	}
 }
+
+func TestSetActorStatus_ResidencyAndActorUID(t *testing.T) {
+	withUID := func(a *ateapipb.Actor, uid string) *ateapipb.Actor {
+		a.Metadata.Uid = uid
+		return a
+	}
+	recorded := func(name, phase, residency, uid string) *v1alpha1.Task {
+		task := storedTask(name, phase)
+		task.Status.Residency, task.Status.ActorUid = residency, uid
+		return task
+	}
+
+	tests := []struct {
+		name          string
+		task          *v1alpha1.Task
+		actor         *ateapipb.Actor
+		err           error
+		wantResidency string
+		wantUID       string
+	}{
+		{"live state wins over the recorded one", recorded("a", "", "suspended", "uid-1"), withUID(testActor("a", ateapipb.ActorState_ACTOR_STATE_RUNNING, ""), "uid-1"), nil, "running", "uid-1"},
+		{"crashed", recorded("a", "", "running", "uid-1"), withUID(testActor("a", ateapipb.ActorState_ACTOR_STATE_CRASHED, ""), "uid-1"), nil, "crashed", "uid-1"},
+		{"reverting", recorded("a", "", "crashed", "uid-1"), withUID(testActor("a", ateapipb.ActorState_ACTOR_STATE_REVERTING, ""), "uid-1"), nil, "reverting", "uid-1"},
+		{"recreated actor has a new UID", recorded("a", "", "crashed", "uid-1"), withUID(testActor("a", ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ""), "uid-2"), nil, "suspended", "uid-2"},
+		{"actor without a UID keeps the recorded one", recorded("a", "", "suspended", "uid-1"), testActor("a", ateapipb.ActorState_ACTOR_STATE_RUNNING, ""), nil, "running", "uid-1"},
+		{"actor that existed is gone", recorded("a", "", "running", "uid-1"), nil, nil, "deleted", "uid-1"},
+		{"actor not created yet", recorded("a", v1alpha1.PhasePending, "queued", ""), nil, nil, "queued", ""},
+		{"substrate unreachable keeps the recorded values", recorded("a", "", "suspended", "uid-1"), nil, errors.New("unavailable"), "suspended", "uid-1"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			setActorStatus(tt.task, tt.actor, tt.err)
+			st := tt.task.GetStatus()
+			if st.GetResidency() != tt.wantResidency || st.GetActorUid() != tt.wantUID {
+				t.Errorf("residency=%q actorUID=%q, want %q %q", st.GetResidency(), st.GetActorUid(), tt.wantResidency, tt.wantUID)
+			}
+		})
+	}
+}