diff --git a/internal/controller/registry.go b/internal/controller/registry.go index 7c1ab57b..5088760d 100644 --- a/internal/controller/registry.go +++ b/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 } diff --git a/internal/controller/registry_test.go b/internal/controller/registry_test.go index c3eb7516..8e952483 100644 --- a/internal/controller/registry_test.go +++ b/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{} diff --git a/internal/harness/harness.go b/internal/harness/harness.go index 716f0165..07031df6 100644 --- a/internal/harness/harness.go +++ b/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. diff --git a/internal/harness/substrate/substrate.go b/internal/harness/substrate/substrate.go index 31509fb4..43cf7b4b 100644 --- a/internal/harness/substrate/substrate.go +++ b/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", diff --git a/internal/harness/substrate/substrate_test.go b/internal/harness/substrate/substrate_test.go index 3f52668f..61dca9e8 100644 --- a/internal/harness/substrate/substrate_test.go +++ b/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 diff --git a/internal/server/server.go b/internal/server/server.go index 267460ff..aa50e0e8 100644 --- a/internal/server/server.go +++ b/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()) {