MUL-4976: feat(daemon): add drain-aware restart for active tasks

apps/docs/content/docs/daemon-runtimes.mdx

@@ -31,7 +31,7 @@ Common commands:
 |---|---|
 | `multica daemon start` | Start (background by default; add `--foreground` to run in the foreground) |
 | `multica daemon stop` | Stop |
-| `multica daemon restart` | Restart |
+| `multica daemon restart` | Restart immediately (add `--drain` to let active tasks finish first) |
 | `multica daemon status` | Show status |
 | `multica daemon logs` | Show logs (add `-f` to follow) |
 

docs/plans/2026-07-19-001-daemon-drain-restart-plan.md

@@ -0,0 +1,65 @@
+---
+title: Drain-Aware Daemon Restart - Plan
+type: feat
+date: 2026-07-19
+topic: daemon-drain-restart
+artifact_contract: ce-unified-plan/v1
+artifact_readiness: implementation-ready
+execution: code
+---
+
+# Drain-Aware Daemon Restart - Plan
+
+## Goal
+
+Add `multica daemon restart --drain`: stop claiming, let accepted work finish, then restart through the invoking CLI. Plain restart remains immediate for compatibility.
+
+## Problem
+
+`daemon restart` currently calls the local `/shutdown` endpoint. That cancels the daemon root context, which is also passed to active agent tasks. The subsequent 30-second wait only waits for cancelled task goroutines to clean up; it does not let agent work finish.
+
+`runtime_recovery` can retry the interrupted row, but still records a failure, redispatches work, and can spend more model tokens. It is crash recovery, not safe restart.
+
+## Design
+
+1. Add an opt-in `--drain` flag to `daemon restart`.
+2. Use a dedicated local endpoint so an older daemon returns 404 rather than ignoring a query parameter and stopping immediately.
+3. Under `claimMu`, represent the shared claim barrier with one explicit owner: `none`, `drain`, or `update`. Any non-`none` owner pauses new claims. This single state replaces independent `pauseClaims`, `draining`, and `updating` flags so acquisition and release cannot disagree about which lifecycle operation is active.
+4. Manual drain acquires the `drain` owner even while claims or tasks are active, then waits for `claimsInFlight == 0 && activeTasks == 0` without cancelling active contexts. Preserve the existing handoff invariant: `activeTasks` increments before `claimsInFlight` decrements, so the drain cannot observe false idle.
+5. Periodic auto-update acquires the `update` owner only while fully idle. Heartbeat-triggered update acquires the same `update` owner without requiring idle, preserving its existing immediate-update behavior while preventing it from bypassing an active drain.
+6. A heartbeat update that cannot acquire ownership remains pending and is retried by a later heartbeat; claim-barrier contention is not reported as an update failure. A drain that finds `update` ownership already held returns `409 Conflict`.
+7. Release is owner-specific. A cancelled drain releases only `drain`; a failed update releases only `update`. A successful drain or update keeps ownership through root-context cancellation so no new claim can enter during shutdown.
+8. After drain reaches idle, cancel the daemon and use the existing CLI start path, preserving binary, profile, foreground mode, and overrides.
+
+## Ownership Transitions
+
+| Current owner | Request | Result |
+| --- | --- | --- |
+| `none` | new task claim | Claim enters and increments `claimsInFlight` |
+| `none` | manual drain | Acquire `drain`, pause new claims, wait for accepted work |
+| `none`, idle | periodic auto-update | Acquire `update` and run the upgrade |
+| `none` | heartbeat update | Acquire `update` and run the existing immediate upgrade path |
+| `drain` | task claim | Reject the claim attempt |
+| `drain` | periodic or heartbeat update | Defer without running or reporting failure |
+| `update` | manual drain | Return `409 Conflict` |
+| `update` | another update | Defer without starting a second upgrade |
+| `drain` | requester cancellation | Release `drain` and resume claims |
+| `update` | upgrade failure | Release `update` and resume claims |
+| `drain` or `update` | successful shutdown/restart | Retain ownership until process exit |
+
+Plain `multica daemon restart` remains an intentionally immediate shutdown and does not participate in this opt-in drain protocol.
+
+## Concurrency Verification
+
+Add deterministic tests for both reviewer-reported orderings using channel-gated operations rather than timing-only assertions:
+
+1. **Drain then heartbeat update:** keep one task active, start `/shutdown/drain`, wait until `drain` owns the barrier, then invoke `handleUpdate`. Assert that the update function and restart are not called, the drain retains ownership, and the update is not reported failed.
+2. **Heartbeat update then drain:** start `handleUpdate` with an update function blocked after `update` ownership is acquired, then call `/shutdown/drain`. Assert `409 Conflict`; release the update and verify its existing completion/restart path.
+
+Keep the existing auto-update, claim-handoff, drain cancellation, concurrent-drain, endpoint compatibility, and immediate-shutdown regression tests. Run the focused package tests under the race detector because the contract is specifically about cross-goroutine ownership and ordering.
+
+## Scope and Verification
+
+Change only daemon lifecycle/health code, focused tests, and the daemon command doc. Do not change default restart, server recovery, executable trust, or PR #5494 ownership semantics.
+
+Cover active work, claim handoff, client cancellation, all owner transition conflicts, dedicated endpoint selection, immediate-shutdown regression, focused package tests, race tests, vet, and formatting.

server/cmd/multica/cmd_daemon.go

@@ -119,6 +119,7 @@ func init() {
 	rf.Int("max-concurrent-tasks", 0, "Max tasks running in parallel (env: MULTICA_DAEMON_MAX_CONCURRENT_TASKS)")
 	rf.Bool("no-auto-update", false, "Disable periodic CLI self-update (env: MULTICA_DAEMON_AUTO_UPDATE=false)")
 	rf.Duration("auto-update-interval", 0, "How often to poll GitHub for a newer release (env: MULTICA_DAEMON_AUTO_UPDATE_INTERVAL)")
+	rf.Bool("drain", false, "Stop claiming new tasks and wait for active tasks to finish before restarting")
 
 	df := daemonDiskUsageCmd.Flags()
 	df.Bool("by-workspace", false, "Aggregate output by workspace instead of by task")
@@ -932,6 +933,7 @@ func requireDaemonRestartPreflight(cmd *cobra.Command, profile string) error {
 func runDaemonRestart(cmd *cobra.Command, args []string) error {
 	profile := resolveProfile(cmd)
 	healthPort := healthPortForProfile(profile)
+	drain, _ := cmd.Flags().GetBool("drain")
 
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()
@@ -949,10 +951,18 @@ func runDaemonRestart(cmd *cobra.Command, args []string) error {
 		}
 		pid, _ := health["pid"].(float64)
 		if pid > 0 {
-			fmt.Fprintf(os.Stderr, "Stopping daemon (pid %d)...\n", int(pid))
-			if err := requestDaemonShutdown(healthPort); err != nil {
-				if p, perr := os.FindProcess(int(pid)); perr == nil {
-					_ = p.Kill()
+			if drain {
+				active, _ := health["active_task_count"].(float64)
+				fmt.Fprintf(os.Stderr, "Draining daemon (pid %d); waiting for %d active task(s)...\n", int(pid), int(active))
+				if err := requestDaemonDrainShutdown(cmd.Context(), healthPort); err != nil {
+					return fmt.Errorf("drain daemon before restart: %w", err)
+				}
+			} else {
+				fmt.Fprintf(os.Stderr, "Stopping daemon (pid %d)...\n", int(pid))
+				if err := requestDaemonShutdown(healthPort); err != nil {
+					if p, perr := os.FindProcess(int(pid)); perr == nil {
+						_ = p.Kill()
+					}
 				}
 			}
 			// Wait until the port is fully released (not merely past "running"),
@@ -1055,6 +1065,29 @@ func requestDaemonShutdown(healthPort int) error {
 	return nil
 }
 
+// requestDaemonDrainShutdown asks the daemon to stop claiming new tasks, wait
+// for every in-flight claim and active task to finish naturally, and then shut
+// down. It deliberately has no fixed client timeout: a legitimate agent task
+// can run for hours. Cancelling ctx aborts the request and makes the daemon
+// release its claim barrier instead of leaving the runtime paused.
+func requestDaemonDrainShutdown(ctx context.Context, healthPort int) error {
+	url := fmt.Sprintf("http://127.0.0.1:%d/shutdown/drain", healthPort)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil)
+	if err != nil {
+		return err
+	}
+	resp, err := (&http.Client{}).Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<10))
+		return fmt.Errorf("unexpected status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
+	}
+	return nil
+}
+
 // --- daemon status ---
 
 func runDaemonStatus(cmd *cobra.Command, _ []string) error {

server/cmd/multica/cmd_daemon_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"net"
 	"net/http"
@@ -384,6 +385,57 @@ func TestDaemonRestartUnauthenticatedFailsBeforeStopping(t *testing.T) {
 	}
 }
 
+func TestDaemonRestartCommandExposesDrainFlag(t *testing.T) {
+	t.Parallel()
+
+	flag := daemonRestartCmd.Flags().Lookup("drain")
+	if flag == nil {
+		t.Fatal("daemon restart is missing the --drain flag")
+	}
+	if flag.DefValue != "false" {
+		t.Fatalf("--drain default = %q, want false for backwards compatibility", flag.DefValue)
+	}
+}
+
+func TestRequestDaemonDrainShutdownUsesDedicatedEndpointAndCallerContext(t *testing.T) {
+	t.Parallel()
+
+	requestSeen := make(chan struct{})
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			t.Errorf("method = %s, want POST", r.Method)
+		}
+		if got := r.URL.Path; got != "/shutdown/drain" {
+			t.Errorf("path = %q, want /shutdown/drain", got)
+		}
+		close(requestSeen)
+		<-r.Context().Done()
+	}))
+	defer srv.Close()
+
+	port := srv.Listener.Addr().(*net.TCPAddr).Port
+	ctx, cancel := context.WithCancel(context.Background())
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- requestDaemonDrainShutdown(ctx, port)
+	}()
+
+	select {
+	case <-requestSeen:
+	case <-time.After(time.Second):
+		t.Fatal("drain shutdown request was not delivered")
+	}
+	cancel()
+	select {
+	case err := <-errCh:
+		if err == nil || !strings.Contains(err.Error(), "context canceled") {
+			t.Fatalf("request error = %v, want caller context cancellation", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("drain shutdown request ignored caller cancellation")
+	}
+}
+
 // fakeRunningDaemon serves a fake healthy daemon on the given profile's health
 // port and reports any /shutdown request on the returned channel. The PID in
 // /health is our own so a kill-fallback would be visible as a test crash too.

server/internal/daemon/auto_update.go

@@ -83,7 +83,7 @@ func (d *Daemon) autoUpdateLoop(ctx context.Context) {
 }
 
 // tryAutoUpdate runs one check-and-maybe-upgrade cycle. Bails early on any of:
-// already updating (server-triggered upgrade in flight), active tasks (defer
+// update ownership already held, active tasks (defer
 // to next tick — we never interrupt running agents), version fetch failure,
 // or no newer release. The function never returns an error: a check that
 // fails today will be retried at the next tick, and we don't want a transient
@@ -96,15 +96,15 @@ func (d *Daemon) tryAutoUpdate(ctx context.Context) {
 	// the Runtimes page is already in flight, let it finish and re-check next
 	// tick (by which time we'll either be on the new binary or it failed and
 	// we can retry).
-	if d.updating.Load() {
+	if d.isUpdating() {
 		d.logger.Debug("auto-update: skip — update already in progress")
 		return
 	}
 	// Cheap pre-fetch idle check: the release-metadata fetch below makes an
 	// HTTPS call to GitHub, and there is no point paying that cost (or the
 	// rate-limit budget) when we already know we are going to defer. A task
 	// that starts between this load and the barrier check below is caught
-	// by the strict re-check under claimMu inside trySetClaimBarrier.
+	// by the strict re-check under claimMu inside tryBeginUpdate.
 	if running := d.activeTasks.Load(); running > 0 {
 		d.logger.Debug("auto-update: skip — tasks running", "active", running)
 		return
@@ -122,36 +122,19 @@ func (d *Daemon) tryAutoUpdate(ctx context.Context) {
 		return
 	}
 
-	// CAS the updating flag so a concurrent server-triggered handleUpdate
-	// dropped onto a heartbeat tick can't double-fire. Release on every exit
-	// path before triggerRestart — once that lands, the daemon ctx is
-	// cancelled and the flag dies with the process.
-	if !d.updating.CompareAndSwap(false, true) {
-		d.logger.Debug("auto-update: skip — update already in progress (raced)")
-		return
-	}
-	released := false
-	defer func() {
-		if !released {
-			d.updating.Store(false)
-		}
-	}()
-
 	// Strict barrier: between the cheap pre-fetch idle check and now the
 	// release fetch took anywhere from tens of milliseconds (typical) to
-	// seconds (slow link, GitHub hiccup), plenty of time for a poller to
-	// claim a fresh task. trySetClaimBarrier checks claimsInFlight +
-	// activeTasks under claimMu and only flips pauseClaims to true if both
-	// are zero, so once it returns true we can run the upgrade knowing that
-	// no in-flight task will be cancelled by triggerRestart.
-	if !d.trySetClaimBarrier() {
-		d.logger.Info("auto-update: deferring — task or claim in flight at barrier check")
+	// seconds (slow link, GitHub hiccup), plenty of time for a task, drain, or
+	// heartbeat-triggered update to acquire lifecycle ownership. The update
+	// owner atomically covers all three cases and requires full idle here.
+	if !d.tryBeginUpdate(true) {
+		d.logger.Info("auto-update: deferring — lifecycle barrier unavailable or daemon not idle")
 		return
 	}
-	barrierReleased := false
+	keepBarrier := false
 	defer func() {
-		if !barrierReleased {
-			d.releaseClaimBarrier()
+		if !keepBarrier {
+			d.releaseUpdate()
 		}
 	}()
 
@@ -166,11 +149,8 @@ func (d *Daemon) tryAutoUpdate(ctx context.Context) {
 
 	d.logger.Info("auto-update: upgrade completed, restarting", "target", release.TagName, "output", output)
 	// triggerRestart cancels the root context, which causes Run() to return
-	// and the parent (cmd_daemon.go) to re-exec the new binary. Leave both
-	// the updating flag and the claim barrier held — process exit is
-	// imminent and clearing either would open a window for new claims / a
-	// second auto-update tick to fire mid-shutdown.
-	released = true
-	barrierReleased = true
+	// and the parent (cmd_daemon.go) to re-exec the new binary. Keep ownership
+	// only if a restart was actually scheduled; otherwise resume claims.
 	d.triggerRestart()
+	keepBarrier = d.RestartBinary() != ""
 }