internal/store/hosts.go

@@ -388,11 +388,14 @@ func (s *Store) SetHostAlwaysOn(ctx context.Context, hostID string, alwaysOn boo
 	if alwaysOn {
 		v = 1
 	}
-	_, err := s.db.ExecContext(ctx,
+	res, err := s.db.ExecContext(ctx,
 		`UPDATE hosts SET always_on = ? WHERE id = ?`, v, hostID)
 	if err != nil {
 		return fmt.Errorf("store: set host always_on: %w", err)
 	}
+	if n, _ := res.RowsAffected(); n == 0 {
+		return ErrNotFound
+	}
 	return nil
 }
 

internal/store/hosts_always_on_test.go

@@ -43,4 +43,13 @@ func TestHostAlwaysOnDefaultAndToggle(t *testing.T) {
 	if len(hosts) != 1 || hosts[0].AlwaysOn {
 		t.Fatalf("ListHosts should report always_on=false, got %+v", hosts)
 	}
+
+	// Verify the agent hot-path (LookupHostByAgentToken) also reflects the toggle.
+	byToken, err := st.LookupHostByAgentToken(ctx, "tok-hash")
+	if err != nil {
+		t.Fatalf("lookup by agent token: %v", err)
+	}
+	if byToken.AlwaysOn {
+		t.Fatalf("LookupHostByAgentToken: expected always_on=false after toggle, got true")
+	}
 }