fix(substrate): drain warm actors on shutdown to prevent RUNNING leak

internal/controller/registry.go

@@ -15,6 +15,7 @@
 package controller
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
@@ -73,7 +74,22 @@ func (r *Registry) SetDefaultHarness(id string) error {
 	return nil
 }
 
-// Close releases resources held by the registry.
+// Close releases resources held by the registry. It drains every registered
+// harness that implements the optional harness.Drainer capability so warm
+// actors awaiting deferred idle suspension are suspended rather than leaked on
+// process exit.
 func (r *Registry) Close() error {
+	r.mu.RLock()
+	drainers := make([]harness.Drainer, 0, len(r.harnesses))
+	for _, h := range r.harnesses {
+		if d, ok := h.(harness.Drainer); ok {
+			drainers = append(drainers, d)
+		}
+	}
+	r.mu.RUnlock()
+
+	for _, d := range drainers {
+		d.Shutdown(context.Background())
+	}
 	return nil
 }

internal/controller/registry_test.go

@@ -27,6 +27,38 @@ func (d *dummyHarness) Start(ctx context.Context, conversationID string, harness
 	return nil, nil
 }
 
+// drainableHarness implements the optional harness.Drainer capability so the
+// registry test can assert that Close drains it.
+type drainableHarness struct {
+	dummyHarness
+	shutdownCalls int
+}
+
+func (d *drainableHarness) Shutdown(ctx context.Context) {
+	d.shutdownCalls++
+}
+
+// Close must invoke Shutdown on every registered harness that implements the
+// Drainer capability so warm actors are suspended before the process exits.
+// Harnesses without the capability are skipped without error.
+func TestRegistry_CloseDrainsDrainerHarnesses(t *testing.T) {
+	r := NewRegistry()
+	drainable := &drainableHarness{}
+	if err := r.RegisterHarness("warm", drainable); err != nil {
+		t.Fatalf("RegisterHarness(warm): %v", err)
+	}
+	if err := r.RegisterHarness("plain", &dummyHarness{}); err != nil {
+		t.Fatalf("RegisterHarness(plain): %v", err)
+	}
+
+	if err := r.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	if drainable.shutdownCalls != 1 {
+		t.Fatalf("Shutdown calls = %d, want 1", drainable.shutdownCalls)
+	}
+}
+
 func TestRegistry_RegisterHarness(t *testing.T) {
 	r := NewRegistry()
 	h := &dummyHarness{}

internal/harness/harness.go

@@ -46,6 +46,16 @@ type Harness interface {
 	Start(ctx context.Context, conversationID string, harnessConfig []byte) (Execution, error)
 }
 
+// Drainer is an optional Harness capability. Shutdown releases per-conversation
+// resources that outlive individual turns -- for example warm actors kept
+// running between turns and awaiting a deferred idle suspension -- so they are
+// not leaked when the process exits. Harnesses that hold no such deferred state
+// need not implement it. Callers should invoke Shutdown after in-flight turns
+// have drained so no turn re-arms deferred state after the drain.
+type Drainer interface {
+	Shutdown(ctx context.Context)
+}
+
 // Execution represents an active interactive session with an agent or planner.
 type Execution interface {
 	// Run executes the session and streams events to the provided Handler.

internal/harness/substrate/substrate.go

@@ -513,6 +513,71 @@ func (h *SubstrateHarness) suspendWarmActor(conversationID, execID string, gener
 	h.idleMu.Unlock()
 }
 
+// Shutdown drains warm actors awaiting idle suspension so a process exit does
+// not leak them as RUNNING actors. A warm actor sits between turns with a
+// pending idle timer whose only home is this process's memory; if the process
+// dies before the timer fires, the actor is never suspended. Shutdown stops each
+// pending timer and suspends the actor synchronously.
+//
+// Actors with an active turn are left untouched: the turn owns the actor and
+// schedules its own suspension on Close. Callers should therefore invoke
+// Shutdown only after in-flight turns have drained (e.g. after the gRPC server's
+// GracefulStop returns) so no turn re-arms a timer after this drain.
+func (h *SubstrateHarness) Shutdown(ctx context.Context) {
+	if h.idleMode != idleModeWarmThenSuspend {
+		return
+	}
+
+	h.idleMu.Lock()
+	var (
+		drain      []string
+		inProgress []chan struct{}
+	)
+	for conversationID, state := range h.warmActors {
+		switch {
+		case state.inTurn:
+			// An active turn owns the actor; it will suspend on Close.
+			continue
+		case state.suspending != nil:
+			// A fired timer is already suspending this actor; wait for it below.
+			inProgress = append(inProgress, state.suspending)
+		case state.timer != nil:
+			state.timer.Stop()
+			state.timer = nil
+			// Neutralize a timer callback that already fired but is still blocked
+			// on idleMu: the generation bump makes suspendWarmActor a no-op so it
+			// cannot suspend again after we do.
+			state.generation++
+			if state.workerAddr == "" {
+				delete(h.warmActors, conversationID)
+				continue
+			}
+			drain = append(drain, conversationID)
+		default:
+			// No timer and not suspending: nothing deferred to clean up.
+			delete(h.warmActors, conversationID)
+		}
+	}
+	h.idleMu.Unlock()
+
+	for _, conversationID := range drain {
+		h.suspendActor(ctx, conversationID, "")
+		h.idleMu.Lock()
+		delete(h.warmActors, conversationID)
+		h.idleMu.Unlock()
+	}
+
+	// Wait for any timer-driven suspensions already in progress so they are not
+	// cut short by process exit.
+	for _, done := range inProgress {
+		select {
+		case <-ctx.Done():
+			return
+		case <-done:
+		}
+	}
+}
+
 func (h *SubstrateHarness) suspendActor(ctx context.Context, conversationID, execID string) {
 	// Suspend actor to return resource to standard standby pool
 	slog.InfoContext(ctx, "Suspending SubstrATE actor",

internal/harness/substrate/substrate_test.go

@@ -657,6 +657,98 @@ func TestSubstrateHarness_FailedWarmResetRetainsIdleCleanup(t *testing.T) {
 	}
 }
 
+// A warm actor sits with a pending idle-suspend timer between turns. If the
+// ax-server process exits before that timer fires, the timer dies with the
+// process and the actor is never suspended -- it leaks as a RUNNING actor
+// holding a worker slot. Shutdown must drain those pending timers by suspending
+// the warm actors synchronously.
+func TestSubstrateHarness_ShutdownDrainsWarmActorsAwaitingIdleSuspend(t *testing.T) {
+	ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"}
+	srv := &harnesstest.MockHarnessServer{}
+	h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv))
+	h.idleMode = idleModeWarmThenSuspend
+	h.idleTimeout = time.Hour // must NOT fire on its own during the test
+
+	ctx := context.Background()
+	exec, err := h.Start(ctx, "conv-drain", substrateHarnessConfig)
+	if err != nil {
+		t.Fatalf("Start: %v", err)
+	}
+	if err := exec.Run(ctx, &harnesstest.MockHandler{}); err != nil {
+		t.Fatalf("Run: %v", err)
+	}
+	if err := exec.Close(ctx); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+
+	// The actor is warm now with a one-hour pending timer; nothing suspended yet.
+	if _, _, suspend := ctrl.Calls(); len(suspend) != 0 {
+		t.Fatalf("suspend=%v, want none before shutdown", suspend)
+	}
+
+	h.Shutdown(ctx)
+
+	if _, _, suspend := ctrl.Calls(); !slices.Equal(suspend, []string{"conv-drain"}) {
+		t.Fatalf("suspend=%v, want the warm actor suspended on shutdown", suspend)
+	}
+
+	// The bookkeeping entry is gone, so a lingering timer cannot suspend twice.
+	h.idleMu.Lock()
+	_, present := h.warmActors["conv-drain"]
+	h.idleMu.Unlock()
+	if present {
+		t.Fatalf("warm actor entry still present after shutdown drain")
+	}
+}
+
+// Shutdown must not touch an actor that still has an active turn: the turn owns
+// the actor and will schedule its own suspension on Close.
+func TestSubstrateHarness_ShutdownLeavesActiveTurnAlone(t *testing.T) {
+	ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"}
+	srv := &harnesstest.MockHarnessServer{}
+	h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv))
+	h.idleMode = idleModeWarmThenSuspend
+	h.idleTimeout = time.Hour
+
+	ctx := context.Background()
+	exec, err := h.Start(ctx, "conv-active", substrateHarnessConfig)
+	if err != nil {
+		t.Fatalf("Start: %v", err)
+	}
+	t.Cleanup(func() { _ = exec.Close(ctx) })
+
+	// Turn is in-flight (Start ran, Close has not). Shutdown must skip it.
+	h.Shutdown(ctx)
+
+	if _, _, suspend := ctrl.Calls(); len(suspend) != 0 {
+		t.Fatalf("suspend=%v, want none while a turn is active", suspend)
+	}
+}
+
+// Shutdown is a no-op in immediate-suspend mode, which never tracks warm actors.
+func TestSubstrateHarness_ShutdownImmediateModeIsNoOp(t *testing.T) {
+	ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"}
+	srv := &harnesstest.MockHarnessServer{}
+	h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv))
+	// idleMode defaults to immediate-suspend.
+
+	ctx := context.Background()
+	exec, err := h.Start(ctx, "conv-immediate", substrateHarnessConfig)
+	if err != nil {
+		t.Fatalf("Start: %v", err)
+	}
+	if err := exec.Close(ctx); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+
+	// Close already suspended once (immediate mode); Shutdown adds nothing.
+	_, _, before := ctrl.Calls()
+	h.Shutdown(ctx)
+	if _, _, after := ctrl.Calls(); !slices.Equal(after, before) {
+		t.Fatalf("suspend calls changed across Shutdown: before=%v after=%v", before, after)
+	}
+}
+
 // The eager-close capability tells the controller whether an execution must be
 // closed before the next Start for the same conversation. Only warm mode needs
 // that (turn-slot bookkeeping); immediate mode must keep upstream's

internal/server/server.go

@@ -73,7 +73,6 @@ func (s *Server) Exec(req *proto.ExecRequest, stream grpc.ServerStreamingServer[
 	return s.controller.Exec(ctx, req, outputHandler)
 }
 
-
 func (s *Server) DeleteConversation(ctx context.Context, req *proto.DeleteConversationRequest) (*proto.DeleteConversationResponse, error) {
 	slog.InfoContext(ctx, "Deleting conversation...",
 		slog.String("conversation_id", req.ConversationId))
@@ -121,15 +120,20 @@ func (s *Server) Serve(address string, opts ...grpc.ServerOption) error {
 	return nil
 }
 
-// GracefulStop stops the gRPC server gracefully.
+// GracefulStop stops the gRPC server gracefully. The gRPC server is drained
+// first so in-flight turns run to completion before the controller releases
+// resources. This ordering also lets warm harnesses drain correctly: by the
+// time the controller closes, every turn has finished and parked its actor for
+// idle suspension, so the drain suspends those actors instead of racing a turn
+// that would re-arm the idle timer afterward.
 func (s *Server) GracefulStop() {
 	slog.Info("Stopping server gracefully...")
-	if s.controller != nil {
-		s.controller.Close()
-	}
 	if s.grpcServer != nil {
 		s.grpcServer.GracefulStop()
 	}
+	if s.controller != nil {
+		s.controller.Close()
+	}
 }
 
 func (s *Server) markInFlight(id string) (exists bool, cleanup func()) {