feat(agent): let the daemon re-authenticate an enrolled device (#175)

cmd/waired-agent/login.go

@@ -29,8 +29,13 @@ type enrollFunc func(ctx context.Context, opts setup.EnrollOptions) (*setup.Enro
 type loginController struct {
 	sb       *switchboard
 	activate func(parent context.Context) error
-	enroll   enrollFunc
-	rootCtx  context.Context
+	// reactivate tears the live session down and rebuilds it from the
+	// state a re-auth just rewrote. activate alone cannot do this: it
+	// refuses to publish over a session that is already current, which is
+	// precisely the state every re-auth starts from.
+	reactivate func(parent context.Context) error
+	enroll     enrollFunc
+	rootCtx    context.Context
 	// enrollHTTPFor builds the HTTP client enrollment talks to the control
 	// plane with, given the control URL that run() resolved. It is a
 	// factory, not a client: under --bypass-cp-iam the transport mints a
@@ -63,12 +68,17 @@ type loginControllerConfig struct {
 	Endpoint          string
 	RootCtx           context.Context
 	Activate          func(parent context.Context) error
-	Logger            *slog.Logger
+	// Reactivate replaces a live session with one built from the state on
+	// disk. Required for re-auth; optional otherwise (a nil one makes a
+	// re-auth fail loudly rather than half-succeed with stale tokens in
+	// the running session).
+	Reactivate func(parent context.Context) error
 	// EnrollHTTPFor is optional; nil enrolls with the default client. Set
 	// it when the control plane is behind something that needs a
 	// per-request credential — the IAM-gated Cloud Run service the testnet
 	// runs against answers an unauthenticated POST with a 403 HTML page.
 	EnrollHTTPFor func(ctx context.Context, controlURL string) *http.Client
+	Logger        *slog.Logger
 	// Enroll is optional; nil uses setup.Enroll.
 	Enroll enrollFunc
 }
@@ -81,6 +91,7 @@ func newLoginController(sb *switchboard, cfg loginControllerConfig) *loginContro
 	return &loginController{
 		sb:                sb,
 		activate:          cfg.Activate,
+		reactivate:        cfg.Reactivate,
 		enroll:            enroll,
 		enrollHTTPFor:     cfg.EnrollHTTPFor,
 		rootCtx:           cfg.RootCtx,
@@ -95,10 +106,26 @@ func (lc *loginController) Start(ctx context.Context, req management.LoginStartR
 	lc.mu.Lock()
 	defer lc.mu.Unlock()
 
-	// Already enrolled + active: idempotent no-op.
-	if lc.sb.current() != nil {
+	// Already enrolled + active: idempotent no-op. `waired init` run twice
+	// must not re-enrol, and the tray's start-on-click must not either.
+	//
+	// Reauth is the one caller that means it. It is how an enrolled device
+	// renews credentials the refresh loop can no longer renew for itself
+	// (#175) — the control plane matches the machine key and renews the
+	// same device row, so this replaces tokens, it does not add a device.
+	// Read the session once: lc.mu does not cover the switchboard, and the
+	// node-key rotator can swap it underneath. Two reads could disagree and
+	// leave this taking the no-op branch with reauth already true.
+	live := lc.sb.current() != nil
+	reauth := req.Reauth && live
+	if live && !req.Reauth {
 		return management.LoginStatus{Phase: management.LoginPhaseActive}, nil
 	}
+	if reauth && lc.reactivate == nil {
+		// Enrolling would rewrite the tokens on disk while the live session
+		// kept using the old ones — worse than refusing, and invisible.
+		return management.LoginStatus{}, errors.New("login: this daemon cannot re-authenticate a live session")
+	}
 	// A login is already in flight: single-flight — return its status
 	// rather than spawning a second browser OAuth.
 	if lc.session != nil {
@@ -131,7 +158,7 @@ func (lc *loginController) Start(ctx context.Context, req management.LoginStartR
 		phase:  management.LoginPhaseLoggingIn,
 		cancel: cancel,
 	}
-	go lc.run(loginCtx, sessID, controlURL, deviceName, req.AuthKey)
+	go lc.run(loginCtx, sessID, controlURL, deviceName, req.AuthKey, reauth)
 
 	return lc.snapshotLocked(), nil
 }
@@ -152,8 +179,10 @@ func (lc *loginController) Status(ctx context.Context, sessionID string) (manage
 }
 
 // run executes enrollment then live activation on a background
-// goroutine, advancing the session's phase as it goes.
-func (lc *loginController) run(ctx context.Context, sessID, controlURL, deviceName, authKey string) {
+// goroutine, advancing the session's phase as it goes. reauth selects
+// the activation that replaces a live session instead of publishing a
+// first one.
+func (lc *loginController) run(ctx context.Context, sessID, controlURL, deviceName, authKey string, reauth bool) {
 	// Resolve a port-0 login endpoint (default "udp4:127.0.0.1:0") to a
 	// concrete free UDP port before enrolling. The endpoint is persisted into
 	// identity.json and later parsed by udpListenPortFromEndpoint (which
@@ -203,7 +232,15 @@ func (lc *loginController) run(ctx context.Context, sessID, controlURL, deviceNa
 
 	// Live activation. Runs on rootCtx (process lifetime): the resulting
 	// session must outlive both this goroutine and the login context.
-	if err := lc.activate(lc.rootCtx); err != nil {
+	// A re-auth has a session already running on the credentials we just
+	// replaced, so it tears that down first — activate refuses to publish
+	// over a live one, and leaving it up would mean the daemon kept using
+	// tokens the control plane has already rotated away from.
+	activate := lc.activate
+	if reauth {
+		activate = lc.reactivate
+	}
+	if err := activate(lc.rootCtx); err != nil {
 		lc.fail(sessID, err)
 		return
 	}

cmd/waired-agent/login_test.go

@@ -65,6 +65,22 @@ func newTestLoginController(sb *switchboard, enroll enrollFunc, activate func(co
 	})
 }
 
+// newTestReauthController is the same, plus the reactivate hook a re-auth
+// needs. Separate so the tests above keep proving that a controller
+// WITHOUT one still serves every non-reauth login (#175).
+func newTestReauthController(sb *switchboard, enroll enrollFunc, activate, reactivate func(context.Context) error) *loginController {
+	return newLoginController(sb, loginControllerConfig{
+		StateDir:          "/tmp/does-not-matter",
+		DefaultControlURL: "https://cp.example",
+		Endpoint:          "udp4:127.0.0.1:0",
+		RootCtx:           context.Background(),
+		Activate:          activate,
+		Reactivate:        reactivate,
+		Logger:            testLogger(),
+		Enroll:            enroll,
+	})
+}
+
 func testLogger() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) }
 
 func waitPhase(t *testing.T, lc *loginController, sessID string, want management.LoginPhase) management.LoginStatus {
@@ -304,6 +320,92 @@ func TestLoginIdempotentWhenAlreadyActive(t *testing.T) {
 	}
 }
 
+// PRODUCT CONTRACT (#175): Reauth is the one request that means "yes, I
+// know this device is enrolled — enrol it again anyway". It is what
+// replaced the standalone re-auth path, so if this stops working there is
+// no other way for an enrolled device to renew credentials its refresh
+// loop can no longer renew.
+//
+// This does not invert TestLoginIdempotentWhenAlreadyActive above: that
+// pins the no-op for a request that did NOT ask, and still does.
+func TestLoginReauthReenrollsALiveDevice(t *testing.T) {
+	sb := &switchboard{}
+	sb.publish(&session{provider: &agentProvider{id: &identity.Identity{DeviceID: "d1"}}})
+	fe := &fakeEnroll{result: &setup.EnrollResult{DeviceID: "d1", AccountEmail: "ops@example.com"}}
+
+	var activates, reactivates int32
+	lc := newTestReauthController(sb, fe.fn,
+		func(context.Context) error { atomic.AddInt32(&activates, 1); return nil },
+		func(context.Context) error { atomic.AddInt32(&reactivates, 1); return nil })
+
+	st, err := lc.Start(context.Background(), management.LoginStartRequest{Reauth: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if st.SessionID == "" {
+		t.Fatal("a re-auth must open a real session, not report the no-op status")
+	}
+	got := waitPhase(t, lc, st.SessionID, management.LoginPhaseActive)
+	if got.AccountEmail != "ops@example.com" {
+		t.Errorf("account email = %q, want the one the re-enrolment returned", got.AccountEmail)
+	}
+	if n := atomic.LoadInt32(&fe.calls); n != 1 {
+		t.Errorf("enroll ran %d times, want exactly 1", n)
+	}
+	// The live session is running on the credentials that were just
+	// replaced, so it has to be rebuilt — activate alone refuses to
+	// publish over a current session, and leaving it up would keep the
+	// daemon on tokens the control plane has rotated away from.
+	if atomic.LoadInt32(&reactivates) != 1 {
+		t.Errorf("reactivate ran %d times, want 1", reactivates)
+	}
+	if atomic.LoadInt32(&activates) != 0 {
+		t.Errorf("plain activate ran %d times on a re-auth, want 0", activates)
+	}
+}
+
+// A fresh daemon has no session to tear down, so a Reauth request there is
+// just a login. Worth pinning because the CLI sets the flag from "does
+// identity.json exist", which the daemon has no reason to agree with — a
+// state dir restored from backup, or an identity the daemon failed to
+// activate, both land here.
+func TestLoginReauthOnUnenrolledDaemonIsAPlainLogin(t *testing.T) {
+	sb := &switchboard{}
+	fe := &fakeEnroll{result: &setup.EnrollResult{DeviceID: "d1"}}
+
+	var activates, reactivates int32
+	lc := newTestReauthController(sb, fe.fn,
+		func(context.Context) error { atomic.AddInt32(&activates, 1); return nil },
+		func(context.Context) error { atomic.AddInt32(&reactivates, 1); return nil })
+
+	st, err := lc.Start(context.Background(), management.LoginStartRequest{Reauth: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	waitPhase(t, lc, st.SessionID, management.LoginPhaseActive)
+	if atomic.LoadInt32(&activates) != 1 || atomic.LoadInt32(&reactivates) != 0 {
+		t.Errorf("activate=%d reactivate=%d; an unenrolled daemon has nothing to rebuild",
+			activates, reactivates)
+	}
+}
+
+// Refusing beats half-succeeding: with no way to rebuild the session, a
+// re-auth would rewrite the tokens on disk and leave the running session
+// using the old ones — an inconsistency nothing would report.
+func TestLoginReauthRefusedWithoutARebuildHook(t *testing.T) {
+	sb := &switchboard{}
+	sb.publish(&session{provider: &agentProvider{id: &identity.Identity{DeviceID: "d1"}}})
+	fe := &fakeEnroll{}
+	lc := newTestLoginController(sb, fe.fn, func(context.Context) error { return nil })
+
+	if _, err := lc.Start(context.Background(), management.LoginStartRequest{Reauth: true}); err == nil {
+		t.Fatal("want an error when the controller cannot rebuild the session")
+	}
+	if atomic.LoadInt32(&fe.calls) != 0 {
+		t.Error("enroll must not run when the rebuild it depends on is impossible")
+	}
+}
+
 func TestLoginStatusUnknownSessionResting(t *testing.T) {
 	sb := &switchboard{}
 	lc := newTestLoginController(sb, (&fakeEnroll{}).fn, func(context.Context) error { return nil })

cmd/waired-agent/main.go

@@ -1398,23 +1398,28 @@ func run(ctx context.Context, args []string) error {
 		return nil
 	}
 
-	// reactivate rebuilds the live session from the (now rotated) node key
-	// on disk: it tears down the current session and re-runs activate,
-	// which re-loads node.key and reconstructs the engine / multiplex-bind
-	// / relay factory / disco around it (#228). Serialised by its own mutex
-	// so a rotation cannot race a second rotation; the once-per-~150d
-	// cadence makes a race with a concurrent login implausible. Runs on a
-	// detached goroutine (the rotator triggers it via `go reactivate()`)
-	// because teardown cancels the rotator's own context.
+	// rebuildSession replaces the live session with one built from the
+	// state now on disk: it tears the current one down and re-runs
+	// activate, which re-loads node.key and reconstructs the engine /
+	// multiplex-bind / relay factory / disco around it (#228). Two callers
+	// need it — node-key rotation and a re-auth through the login
+	// controller (#175) — and both must be serialised against each other,
+	// so they share one mutex rather than each holding their own.
 	var reactivateMu sync.Mutex
-	reactivate = func() {
+	rebuildSession := func(parent context.Context) error {
 		reactivateMu.Lock()
 		defer reactivateMu.Unlock()
 		if s := sb.current(); s != nil {
 			s.teardown()
 		}
 		sb.reset()
-		if err := activate(ctx); err != nil {
+		return activate(parent)
+	}
+	// The rotator's entry point: same rebuild, error logged rather than
+	// returned. Runs on a detached goroutine (the rotator triggers it via
+	// `go reactivate()`) because teardown cancels the rotator's own context.
+	reactivate = func() {
+		if err := rebuildSession(ctx); err != nil {
 			logger.Error("re-activate after node-key rotation failed; device unenrolled until restart", "err", err)
 		}
 	}
@@ -1442,6 +1447,7 @@ func run(ctx context.Context, args []string) error {
 		Endpoint:          "udp4:" + *loginListen,
 		RootCtx:           ctx,
 		Activate:          activate,
+		Reactivate:        rebuildSession,
 		EnrollHTTPFor:     enrollHTTPFor,
 		Logger:            logger,
 	})

cmd/waired/init_benchmark.go

@@ -327,6 +327,22 @@ func waitForBenchmark(mgmtURL string, out io.Writer) (resp *management.Benchmark
 			case "pull_failed":
 				writePrompt(out, "Model download failed; skipping the interactive-performance check.")
 				return nil, false
+			case "disabled", "stopped":
+				// Terminal, the same way waitForBundledModel already treats
+				// them (init_pull.go): a subsystem that is off or parked will
+				// never report a ready model, so waiting is waiting for
+				// something nobody has asked to happen.
+				//
+				// This used to fall into the default arm below, which reads
+				// any unrecognised state as "engine is up, download in
+				// flight" — so `waired init --inference-enabled=false` sat on
+				// "Waiting for the model to finish downloading…" for the full
+				// ten-minute deadline and then reported it had given up, on a
+				// host with no model and no intention of getting one. Found
+				// by #175's installtest migration, which made this the path
+				// CI takes: ten minutes per leg, three legs, every PR.
+				writePrompt(out, "Local inference is off on this device; skipping the performance check.")
+				return nil, false
 			case "no_engine":
 				// On a fresh bundled install the engine is still being
 				// brought up at the first polls, so `no_engine` is transient

cmd/waired/init_benchmark_test.go

@@ -341,6 +341,56 @@ func TestPromptBenchmark_TerminalStateSkips(t *testing.T) {
 	}
 }
 
+// PRODUCT CONTRACT: a subsystem that is switched off or parked never
+// produces a ready model, so the wait must end at the first poll — the
+// same call waitForBundledModel already makes (init_pull.go).
+//
+// These two states used to fall through to "engine is up, a download must
+// be in flight", so `waired init --inference-enabled=false` on the daemon
+// path printed "Waiting for the model to finish downloading…" and held the
+// terminal for the full ten-minute deadline before reporting that it had
+// given up — on a host with no model and no intention of getting one.
+// Measured: three installtest legs, ten minutes each, every PR.
+func TestPromptBenchmark_OffSubsystemSkipsImmediately(t *testing.T) {
+	for _, state := range []string{"disabled", "stopped"} {
+		t.Run(state, func(t *testing.T) {
+			stub := &benchStub{ready: false, state: state}
+			srv := stub.server()
+			defer srv.Close()
+
+			var out strings.Builder
+			done := make(chan error, 1)
+			go func() {
+				done <- promptBenchmarkRecommendation(srv.URL, false, &out,
+					bufio.NewScanner(strings.NewReader("")), false)
+			}()
+			select {
+			case err := <-done:
+				if err != nil {
+					t.Fatalf("prompt: %v", err)
+				}
+			case <-time.After(10 * time.Second):
+				t.Fatal("still waiting: an off subsystem must not be waited out")
+			}
+
+			if got := out.String(); !strings.Contains(got, "Local inference is off") {
+				t.Errorf("expected the off-subsystem skip notice, got: %q", got)
+			}
+			if got := out.String(); strings.Contains(got, "finish downloading") {
+				t.Errorf("announced a model download for state %q: %q", state, got)
+			}
+			// One probe of each endpoint is enough to decide; anything more
+			// means the loop kept going after a terminal answer.
+			stub.mu.Lock()
+			calls := stub.benchCalls
+			stub.mu.Unlock()
+			if calls != 1 {
+				t.Errorf("/benchmark polled %d times, want 1", calls)
+			}
+		})
+	}
+}
+
 // When the only lighter step-down is the tiny 0.5B, declining (default No)
 // disables local inference rather than switching / dismissing.
 func TestPromptBenchmark_TinyDeclineDisables(t *testing.T) {