Harden secure agent runtime transitions

cmd/pulse-agent-runner/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net"
 	"net/url"
 	"os"
 	"os/signal"
@@ -25,6 +26,7 @@ type runtimeConfig struct {
 	StateDir          string
 	HealthFile        string
 	AgentIDFile       string
+	Hostname          string
 	ServerFingerprint string
 	CAFile            string
 	Insecure          bool
@@ -37,6 +39,7 @@ func loadConfig() (runtimeConfig, error) {
 		StateDir:          strings.TrimSpace(os.Getenv("PULSE_AGENT_RUNNER_STATE_DIR")),
 		HealthFile:        strings.TrimSpace(os.Getenv("PULSE_AGENT_RUNNER_HEALTH_FILE")),
 		AgentIDFile:       strings.TrimSpace(os.Getenv("PULSE_AGENT_RUNNER_AGENT_ID_FILE")),
+		Hostname:          strings.TrimSpace(os.Getenv("PULSE_AGENT_RUNNER_HOSTNAME")),
 		ServerFingerprint: strings.TrimSpace(os.Getenv("PULSE_SERVER_FINGERPRINT")),
 		CAFile:            strings.TrimSpace(os.Getenv("SSL_CERT_FILE")),
 		Insecure:          strings.EqualFold(strings.TrimSpace(os.Getenv("PULSE_INSECURE")), "true"),
@@ -44,6 +47,13 @@ func loadConfig() (runtimeConfig, error) {
 	if config.PulseURL == "" || config.TokenFile == "" || config.StateDir == "" || config.HealthFile == "" || config.AgentIDFile == "" {
 		return runtimeConfig{}, errors.New("PULSE_URL and the action-runner token, state, health, and agent identity file settings are required")
 	}
+	if config.Hostname != "" {
+		hostname, err := normalizeRunnerHostname(config.Hostname)
+		if err != nil {
+			return runtimeConfig{}, err
+		}
+		config.Hostname = hostname
+	}
 	parsed, err := url.Parse(config.PulseURL)
 	if err != nil || parsed.Host == "" || (parsed.Scheme != "https" && !(config.Insecure && parsed.Scheme == "http")) {
 		return runtimeConfig{}, errors.New("PULSE_URL must be HTTPS unless PULSE_INSECURE=true explicitly permits HTTP")
@@ -83,9 +93,13 @@ func run() error {
 	if err != nil {
 		return err
 	}
-	hostname, err := os.Hostname()
-	if err != nil || strings.TrimSpace(hostname) == "" {
-		return errors.New("determine action-runner hostname")
+	hostname := config.Hostname
+	if hostname == "" {
+		hostname, err = os.Hostname()
+		hostname = strings.TrimSpace(hostname)
+		if err != nil || hostname == "" {
+			return errors.New("determine action-runner hostname")
+		}
 	}
 	logger := zerolog.New(os.Stderr).With().Timestamp().Str("component", "action-runner").Logger()
 	transportConfig := actionrunner.TransportConfig{
@@ -112,6 +126,32 @@ func run() error {
 	return nil
 }
 
+func normalizeRunnerHostname(value string) (string, error) {
+	hostname := strings.ToLower(strings.TrimSpace(value))
+	hostname = strings.TrimRight(hostname, ".")
+	if hostname == "" || len(hostname) > 253 {
+		return "", errors.New("PULSE_AGENT_RUNNER_HOSTNAME must be a valid hostname or IP address")
+	}
+	if net.ParseIP(hostname) != nil {
+		return hostname, nil
+	}
+	for _, label := range strings.Split(hostname, ".") {
+		if len(label) == 0 || len(label) > 63 || !runnerHostnameLabelByte(label[0]) || !runnerHostnameLabelByte(label[len(label)-1]) {
+			return "", errors.New("PULSE_AGENT_RUNNER_HOSTNAME must be a valid hostname or IP address")
+		}
+		for index := 1; index < len(label)-1; index++ {
+			if !runnerHostnameLabelByte(label[index]) && label[index] != '-' {
+				return "", errors.New("PULSE_AGENT_RUNNER_HOSTNAME must be a valid hostname or IP address")
+			}
+		}
+	}
+	return hostname, nil
+}
+
+func runnerHostnameLabelByte(value byte) bool {
+	return value >= 'a' && value <= 'z' || value >= '0' && value <= '9'
+}
+
 func readPrivateValue(path, label string) (string, error) {
 	info, err := os.Lstat(path)
 	if err != nil {

cmd/pulse-agent-runner/main_test.go

@@ -22,16 +22,44 @@ func TestLoadConfigUsesDedicatedEnvironmentAndPrivateTokenFile(t *testing.T) {
 	t.Setenv("PULSE_AGENT_RUNNER_STATE_DIR", filepath.Join(dir, "state"))
 	t.Setenv("PULSE_AGENT_RUNNER_HEALTH_FILE", filepath.Join(dir, "state", "health.json"))
 	t.Setenv("PULSE_AGENT_RUNNER_AGENT_ID_FILE", agentID)
+	t.Setenv("PULSE_AGENT_RUNNER_HOSTNAME", " Node.Example. ")
 	t.Setenv("PULSE_SERVER_FINGERPRINT", "sha256:test")
 	config, err := loadConfig()
 	if err != nil {
 		t.Fatal(err)
 	}
-	if config.TokenFile != token || config.ServerFingerprint != "sha256:test" || config.Insecure {
+	if config.TokenFile != token || config.ServerFingerprint != "sha256:test" || config.Hostname != "node.example" || config.Insecure {
 		t.Fatalf("config = %+v", config)
 	}
 }
 
+func TestLoadConfigRejectsInvalidCanonicalHostnameOverride(t *testing.T) {
+	dir := t.TempDir()
+	token := filepath.Join(dir, "runner.token")
+	if err := os.WriteFile(token, []byte("secret\n"), 0600); err != nil {
+		t.Fatal(err)
+	}
+	t.Setenv("PULSE_URL", "https://pulse.example")
+	t.Setenv("PULSE_AGENT_RUNNER_TOKEN_FILE", token)
+	t.Setenv("PULSE_AGENT_RUNNER_STATE_DIR", filepath.Join(dir, "state"))
+	t.Setenv("PULSE_AGENT_RUNNER_HEALTH_FILE", filepath.Join(dir, "state", "health.json"))
+	t.Setenv("PULSE_AGENT_RUNNER_AGENT_ID_FILE", filepath.Join(dir, "agent-id"))
+	for _, hostname := range []string{"bad host", "-node.example", "node/example", strings.Repeat("a", 64) + ".example"} {
+		t.Run(hostname, func(t *testing.T) {
+			t.Setenv("PULSE_AGENT_RUNNER_HOSTNAME", hostname)
+			if _, err := loadConfig(); err == nil || !strings.Contains(err.Error(), "valid hostname") {
+				t.Fatalf("loadConfig error = %v", err)
+			}
+		})
+	}
+}
+
+func TestNormalizeRunnerHostnameAcceptsIPLiteral(t *testing.T) {
+	if got, err := normalizeRunnerHostname(" 2001:DB8::1 "); err != nil || got != "2001:db8::1" {
+		t.Fatalf("normalizeRunnerHostname = %q, %v", got, err)
+	}
+}
+
 func TestLoadConfigRejectsTokenInArgvEquivalentAndInsecureHTTPByDefault(t *testing.T) {
 	dir := t.TempDir()
 	t.Setenv("PULSE_URL", "http://pulse.example")

cmd/pulse-agent/main.go

@@ -119,6 +119,87 @@ func wireUpdaterHooks(hostCfg *hostagent.Config, updater *agentupdate.Updater) {
 	hostCfg.OnServerVersion = updater.NudgeVersion
 }
 
+func pendingUpdatePreviousVersion(pending *agentupdate.PendingPrivilegedUpdate) string {
+	if pending == nil {
+		return ""
+	}
+	return pending.PreviousVersion
+}
+
+func supervisePendingPrivilegedUpdate(
+	ctx context.Context,
+	update agentupdate.PrivilegedUpdate,
+	pending *agentupdate.PendingPrivilegedUpdate,
+	stateDir string,
+	locallyReady func() bool,
+	reportAccepted <-chan struct{},
+	pollInterval time.Duration,
+	logger *zerolog.Logger,
+) error {
+	if pending == nil {
+		return nil
+	}
+	if update == nil || locallyReady == nil || pollInterval <= 0 {
+		return errors.New("pending helper update supervisor is not configured")
+	}
+	activation := pending.Activation
+	rollback := func(reason error) error {
+		rollbackCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancel()
+		result, err := update.Rollback(rollbackCtx, activation)
+		if err != nil {
+			return errors.Join(reason, fmt.Errorf("typed helper rollback failed: %w", err))
+		}
+		if result.Action != "rolled_back" || !strings.EqualFold(result.ActiveSHA256, activation.RollbackSHA256) {
+			return errors.Join(reason, errors.New("typed helper returned an invalid rollback result"))
+		}
+		if clearErr := agentupdate.ClearPendingPrivilegedUpdate(stateDir); clearErr != nil {
+			return errors.Join(fmt.Errorf("%w; pending update rolled back", reason), fmt.Errorf("clear pending update handoff: %w", clearErr))
+		}
+		return fmt.Errorf("%w; pending update rolled back", reason)
+	}
+
+	deadlineDelay := time.Until(activation.RollbackDeadline)
+	if deadlineDelay <= 0 {
+		return rollback(errors.New("pending update health deadline expired before startup"))
+	}
+	deadline := time.NewTimer(deadlineDelay)
+	defer deadline.Stop()
+	ticker := time.NewTicker(pollInterval)
+	defer ticker.Stop()
+	reportOK := false
+	for {
+		if reportOK && locallyReady() {
+			commitCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+			result, err := update.Commit(commitCtx, activation)
+			cancel()
+			if err == nil && result.Action == "committed" && strings.EqualFold(result.ActiveSHA256, activation.ActiveSHA256) {
+				if err := agentupdate.ClearPendingPrivilegedUpdate(stateDir); err != nil {
+					return fmt.Errorf("clear committed update handoff: %w", err)
+				}
+				if logger != nil {
+					logger.Info().Str("activation_id", activation.ActivationID).Msg("Committed helper-backed agent update after local readiness and server report")
+				}
+				return nil
+			}
+			if logger != nil && err != nil {
+				logger.Warn().Err(err).Str("activation_id", activation.ActivationID).Msg("Pending helper update commit failed; retrying until rollback deadline")
+			}
+		}
+
+		select {
+		case <-ctx.Done():
+			return rollback(fmt.Errorf("runtime stopped before pending update commit: %w", ctx.Err()))
+		case <-deadline.C:
+			return rollback(errors.New("pending update did not reach local readiness and accepted-report health floor before deadline"))
+		case <-reportAccepted:
+			reportOK = true
+			reportAccepted = nil
+		case <-ticker.C:
+		}
+	}
+}
+
 type multiValue []string
 
 func (m *multiValue) String() string {
@@ -383,7 +464,8 @@ func run(ctx context.Context, args []string, getenv func(string) string) error {
 
 	// 7. Start Auto-Updater
 	var privilegedUpdate agentupdate.PrivilegedUpdate
-	if helperSocket := strings.TrimSpace(os.Getenv("PULSE_AGENT_HELPER_SOCKET")); helperSocket != "" {
+	helperSocket := strings.TrimSpace(os.Getenv("PULSE_AGENT_HELPER_SOCKET"))
+	if helperSocket != "" {
 		privilegedUpdate, err = newPrivilegeHelperUpdate(helperSocket)
 		if err != nil {
 			return fmt.Errorf("configure typed privilege-helper updates: %w", err)
@@ -403,6 +485,18 @@ func run(ctx context.Context, args []string, getenv func(string) string) error {
 		Disabled:           cfg.DisableAutoUpdate,
 		PrivilegedUpdate:   privilegedUpdate,
 	})
+	var pendingUpdate *agentupdate.PendingPrivilegedUpdate
+	if helperSocket != "" {
+		pendingUpdate, err = agentupdate.LoadPendingPrivilegedUpdate(cfg.StateDir)
+		if err != nil {
+			return fmt.Errorf("load pending helper update handoff: %w", err)
+		}
+	}
+	if pendingUpdate != nil && privilegedUpdate == nil {
+		return errors.New("pending helper update cannot be verified without the typed privilege helper")
+	}
+	pendingReportAccepted := make(chan struct{})
+	var pendingReportOnce sync.Once
 
 	g.Go(func() error {
 		updater.RunLoop(ctx)
@@ -419,6 +513,14 @@ func run(ctx context.Context, args []string, getenv func(string) string) error {
 		if err != nil {
 			return fmt.Errorf("configure typed privilege helper: %w", err)
 		}
+		if privilegedTelemetry != nil {
+			healthCtx, cancelHealth := context.WithTimeout(ctx, 10*time.Second)
+			healthErr := privilegedTelemetry.Health(healthCtx)
+			cancelHealth()
+			if healthErr != nil {
+				return fmt.Errorf("verify typed privilege helper protocol: %w", healthErr)
+			}
+		}
 		hostCfg := hostagent.Config{
 			PulseURL:                cfg.PulseURL,
 			APIToken:                cfg.APIToken,
@@ -450,10 +552,16 @@ func run(ctx context.Context, args []string, getenv func(string) string) error {
 			ModuleStatus:            runtimeStatus.moduleStatuses,
 			Observers:               hostObserverTargets(cfg.Observers),
 			PrivilegedTelemetry:     privilegedTelemetry,
+			UpdatedFromVersion:      pendingUpdatePreviousVersion(pendingUpdate),
 
 			DockerContainerUpdater:           dockerUpdaterBridge,
 			DockerContainerLifecycleOperator: dockerUpdaterBridge,
 		}
+		if pendingUpdate != nil {
+			hostCfg.OnPrimaryReportAccepted = func() {
+				pendingReportOnce.Do(func() { close(pendingReportAccepted) })
+			}
+		}
 		wireUpdaterHooks(&hostCfg, updater)
 
 		agent, err := newHostAgent(hostCfg)
@@ -592,6 +700,20 @@ func run(ctx context.Context, args []string, getenv func(string) string) error {
 			return nil
 		})
 	}
+	if pendingUpdate != nil {
+		g.Go(func() error {
+			return supervisePendingPrivilegedUpdate(
+				ctx,
+				privilegedUpdate,
+				pendingUpdate,
+				cfg.StateDir,
+				ready.Load,
+				pendingReportAccepted,
+				250*time.Millisecond,
+				&logger,
+			)
+		})
+	}
 
 	// 11. Wait for all agents to exit
 	if err := g.Wait(); err != nil && err != context.Canceled {

cmd/pulse-agent/main_test.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"github.com/rcourtman/pulse-go-rewrite/internal/agentexec"
+	"github.com/rcourtman/pulse-go-rewrite/internal/agenthelper"
 	"github.com/rcourtman/pulse-go-rewrite/internal/agentupdate"
 	"github.com/rcourtman/pulse-go-rewrite/internal/dockeragent"
 	"github.com/rcourtman/pulse-go-rewrite/internal/hostagent"
@@ -1754,9 +1755,10 @@ func TestRunConfiguresTypedPrivilegeHelperFromInstallerEnvironment(t *testing.T)
 	t.Setenv("PULSE_AGENT_HELPER_SOCKET", socketPath)
 	configuredPath := ""
 	configuredUpdatePath := ""
+	helper := &helperHealthStub{}
 	newPrivilegeHelperTelemetry = func(path string) (hostagent.PrivilegedTelemetry, error) {
 		configuredPath = path
-		return nil, nil
+		return helper, nil
 	}
 	newPrivilegeHelperUpdate = func(path string) (agentupdate.PrivilegedUpdate, error) {
 		configuredUpdatePath = path
@@ -1786,6 +1788,224 @@ func TestRunConfiguresTypedPrivilegeHelperFromInstallerEnvironment(t *testing.T)
 	if configuredUpdatePath != socketPath {
 		t.Fatalf("helper update socket path = %q, want %q", configuredUpdatePath, socketPath)
 	}
+	if helper.healthCalls != 1 {
+		t.Fatalf("helper health calls = %d, want 1", helper.healthCalls)
+	}
+}
+
+type helperHealthStub struct {
+	hostagent.PrivilegedTelemetry
+	healthErr   error
+	healthCalls int
+}
+
+func (h *helperHealthStub) Health(context.Context) error {
+	h.healthCalls++
+	return h.healthErr
+}
+
+func TestRunRejectsUnhealthyTypedPrivilegeHelper(t *testing.T) {
+	originalHelper := newPrivilegeHelperTelemetry
+	originalUpdate := newPrivilegeHelperUpdate
+	originalUpdater := newUpdater
+	originalHost := newHostAgent
+	defer func() {
+		newPrivilegeHelperTelemetry = originalHelper
+		newPrivilegeHelperUpdate = originalUpdate
+		newUpdater = originalUpdater
+		newHostAgent = originalHost
+	}()
+
+	t.Setenv("PULSE_AGENT_HELPER_SOCKET", "/run/pulse-agent/helper.sock")
+	helper := &helperHealthStub{healthErr: errors.New("incompatible helper")}
+	newPrivilegeHelperTelemetry = func(string) (hostagent.PrivilegedTelemetry, error) {
+		return helper, nil
+	}
+	newPrivilegeHelperUpdate = func(string) (agentupdate.PrivilegedUpdate, error) {
+		return nil, nil
+	}
+	newUpdater = func(agentupdate.Config) *agentupdate.Updater {
+		return agentupdate.New(agentupdate.Config{Disabled: true})
+	}
+	hostCreated := false
+	newHostAgent = func(hostagent.Config) (Runnable, error) {
+		hostCreated = true
+		return &mockRunnable{}, nil
+	}
+
+	err := run(context.Background(), []string{
+		"-token", "deadbeef",
+		"-enable-docker=false",
+		"-enable-kubernetes=false",
+		"-health-addr", "",
+	}, func(string) string { return "" })
+	if err == nil || !strings.Contains(err.Error(), "verify typed privilege helper protocol") {
+		t.Fatalf("run error = %v", err)
+	}
+	if helper.healthCalls != 1 {
+		t.Fatalf("helper health calls = %d, want 1", helper.healthCalls)
+	}
+	if hostCreated {
+		t.Fatal("host agent was created after helper health failed")
+	}
+}
+
+type pendingUpdateSupervisorStub struct {
+	commitResult   agenthelper.UpdateResult
+	rollbackResult agenthelper.UpdateResult
+	commitErr      error
+	rollbackErr    error
+	commitCalls    chan agenthelper.UpdateResult
+	rollbackCalls  chan agenthelper.UpdateResult
+}
+
+func (s *pendingUpdateSupervisorStub) CreateQuarantinedArtifact() (string, *os.File, func() error, error) {
+	return "", nil, func() error { return nil }, errors.New("not implemented")
+}
+
+func (s *pendingUpdateSupervisorStub) WriteQuarantinedSignature(string, string) error {
+	return errors.New("not implemented")
+}
+
+func (s *pendingUpdateSupervisorStub) Stage(context.Context, string, string) (agenthelper.UpdateStageResult, error) {
+	return agenthelper.UpdateStageResult{}, errors.New("not implemented")
+}
+
+func (s *pendingUpdateSupervisorStub) Activate(context.Context, string, string) (agenthelper.UpdateResult, error) {
+	return agenthelper.UpdateResult{}, errors.New("not implemented")
+}
+
+func (s *pendingUpdateSupervisorStub) Commit(_ context.Context, activation agenthelper.UpdateResult) (agenthelper.UpdateResult, error) {
+	s.commitCalls <- activation
+	return s.commitResult, s.commitErr
+}
+
+func (s *pendingUpdateSupervisorStub) Rollback(_ context.Context, activation agenthelper.UpdateResult) (agenthelper.UpdateResult, error) {
+	s.rollbackCalls <- activation
+	return s.rollbackResult, s.rollbackErr
+}
+
+func testPendingUpdate(t *testing.T, stateDir string) *agentupdate.PendingPrivilegedUpdate {
+	t.Helper()
+	activation := agenthelper.UpdateResult{
+		Action:           "pending",
+		ActivationID:     "pulse-agent-0123456789abcdef0123456789abcdef:0123456789abcdef",
+		ActiveSHA256:     strings.Repeat("a", 64),
+		RollbackSHA256:   strings.Repeat("b", 64),
+		RollbackDeadline: time.Now().Add(2 * time.Second).UTC(),
+	}
+	if err := agentupdate.PersistPendingPrivilegedUpdate(stateDir, "1.0.0", activation); err != nil {
+		t.Fatal(err)
+	}
+	return &agentupdate.PendingPrivilegedUpdate{Activation: activation, PreviousVersion: "1.0.0"}
+}
+
+func TestPendingPrivilegedUpdateCommitsOnlyAfterReadinessAndAcceptedReport(t *testing.T) {
+	stateDir := t.TempDir()
+	if err := os.Chmod(stateDir, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	pending := testPendingUpdate(t, stateDir)
+	stub := &pendingUpdateSupervisorStub{
+		commitResult: agenthelper.UpdateResult{
+			Action:         "committed",
+			ActivationID:   pending.Activation.ActivationID,
+			ActiveSHA256:   pending.Activation.ActiveSHA256,
+			RollbackSHA256: pending.Activation.RollbackSHA256,
+		},
+		commitCalls:   make(chan agenthelper.UpdateResult, 1),
+		rollbackCalls: make(chan agenthelper.UpdateResult, 1),
+	}
+	reportAccepted := make(chan struct{})
+	close(reportAccepted)
+	var ready atomic.Bool
+	result := make(chan error, 1)
+	go func() {
+		result <- supervisePendingPrivilegedUpdate(context.Background(), stub, pending, stateDir, ready.Load, reportAccepted, time.Millisecond, nil)
+	}()
+	select {
+	case <-stub.commitCalls:
+		t.Fatal("pending update committed before local readiness")
+	case <-time.After(20 * time.Millisecond):
+	}
+	ready.Store(true)
+	select {
+	case activation := <-stub.commitCalls:
+		if activation != pending.Activation {
+			t.Fatalf("commit activation = %#v", activation)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("pending update was not committed after both health signals")
+	}
+	if err := <-result; err != nil {
+		t.Fatal(err)
+	}
+	if loaded, err := agentupdate.LoadPendingPrivilegedUpdate(stateDir); err != nil || loaded != nil {
+		t.Fatalf("committed handoff = %#v, %v", loaded, err)
+	}
+	select {
+	case <-stub.rollbackCalls:
+		t.Fatal("healthy pending update rolled back")
+	default:
+	}
+}
+
+func TestPendingPrivilegedUpdateCancellationRollsBackAndClearsHandoff(t *testing.T) {
+	stateDir := t.TempDir()
+	if err := os.Chmod(stateDir, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	pending := testPendingUpdate(t, stateDir)
+	stub := &pendingUpdateSupervisorStub{
+		rollbackResult: agenthelper.UpdateResult{
+			Action:         "rolled_back",
+			ActivationID:   pending.Activation.ActivationID,
+			ActiveSHA256:   pending.Activation.RollbackSHA256,
+			RollbackSHA256: pending.Activation.ActiveSHA256,
+		},
+		commitCalls:   make(chan agenthelper.UpdateResult, 1),
+		rollbackCalls: make(chan agenthelper.UpdateResult, 1),
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err := supervisePendingPrivilegedUpdate(ctx, stub, pending, stateDir, func() bool { return false }, make(chan struct{}), time.Millisecond, nil)
+	if err == nil || !strings.Contains(err.Error(), "pending update rolled back") {
+		t.Fatalf("supervisor error = %v", err)
+	}
+	select {
+	case activation := <-stub.rollbackCalls:
+		if activation != pending.Activation {
+			t.Fatalf("rollback activation = %#v", activation)
+		}
+	default:
+		t.Fatal("pending update was not rolled back")
+	}
+	if loaded, loadErr := agentupdate.LoadPendingPrivilegedUpdate(stateDir); loadErr != nil || loaded != nil {
+		t.Fatalf("rolled-back handoff = %#v, %v", loaded, loadErr)
+	}
+}
+
+func TestPendingPrivilegedUpdateRollbackFailurePreservesHandoff(t *testing.T) {
+	stateDir := t.TempDir()
+	if err := os.Chmod(stateDir, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	pending := testPendingUpdate(t, stateDir)
+	stub := &pendingUpdateSupervisorStub{
+		rollbackErr:   errors.New("helper unavailable"),
+		commitCalls:   make(chan agenthelper.UpdateResult, 1),
+		rollbackCalls: make(chan agenthelper.UpdateResult, 1),
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err := supervisePendingPrivilegedUpdate(ctx, stub, pending, stateDir, func() bool { return false }, make(chan struct{}), time.Millisecond, nil)
+	if err == nil || !strings.Contains(err.Error(), "typed helper rollback failed") {
+		t.Fatalf("supervisor error = %v", err)
+	}
+	loaded, loadErr := agentupdate.LoadPendingPrivilegedUpdate(stateDir)
+	if loadErr != nil || loaded == nil || loaded.Activation != pending.Activation {
+		t.Fatalf("failed-rollback handoff = %#v, %v", loaded, loadErr)
+	}
 }
 
 func TestRun_AgentFailure(t *testing.T) {