diff --git a/internal/store/hosts.go b/internal/store/hosts.go index 8dad905..9afa5c7 100644 --- a/internal/store/hosts.go +++ b/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 } diff --git a/internal/store/hosts_always_on_test.go b/internal/store/hosts_always_on_test.go index 8f9cf04..72063ae 100644 --- a/internal/store/hosts_always_on_test.go +++ b/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") + } }