openlatch-client 0.6.19

//! The daemon side of `openlatch system auth switch` (PLT-7 C-8).
//!
//! [`SwitchCtl`] orchestrates a switch through the three local admin routes
//! (`/admin/switch/begin`, `/commit`, `/abort`) and applies it whenever a
//! credential load meets a matching pending-switch marker: at commit, at daemon
//! start, and when the worker's periodic reload or an auth refresh asks.
//!
//! **The principle.** The policy poller stays the only writer of the resident
//! bundle and a switch never discards it: apply re-binds the expected
//! organization, and the old organization's bundle keeps enforcing until the
//! poller accepts the new organization's bundle with the new key. No event
//! captured under the old organization ever leaves under the new key — the
//! switch generation fences every pre-switch event, and apply drops what is
//! still queued.
//!
//! **Why apply is driven from here, never from the credential read.**
//! `CredentialStoreAdapter::retrieve` stays pure: the cloud worker awaits it,
//! and apply waits on the worker, so an apply inside a credential read would
//! deadlock. A read that meets a matching marker only *asks* (via
//! `SwitchState::request_apply`) and refuses the key; this controller applies.

use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use secrecy::SecretString;
use tokio::sync::{mpsc, oneshot};

use crate::cloud::outbox::Outbox;
use crate::cloud::{CloudEvent, CloudState, CredentialProvider, WorkerCommand};
use crate::core::pending_switch::{self, MarkerAction, PendingSwitch};

/// `POST /admin/switch/begin` never holds uploads for longer than this.
pub const MAX_BEGIN_TIMEOUT_MS: u64 = 30_000;
/// The lease outlives the drain budget by this much: a CLI that dies after
/// `begin` cannot hold uploads for longer than `timeout_ms` + this.
pub const LEASE_GRACE: Duration = Duration::from_secs(60);
/// How long a worker command may take to be acknowledged.
const WORKER_ACK_TIMEOUT: Duration = Duration::from_secs(15);
/// Pause between drain rounds while `begin` waits for the queues to empty.
const DRAIN_ROUND: Duration = Duration::from_millis(200);
/// How soon a marker check that could not finish (an unreadable marker, a
/// marker delete that failed) is tried again.
const MARKER_RETRY: Duration = Duration::from_secs(5);

/// Why a switch operation failed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SwitchError {
    /// `commit` with no pending-switch marker on disk.
    NoMarker,
    /// The stored credential is not the marker's key.
    CredentialMismatch,
    /// The cloud worker did not acknowledge a command in time.
    Worker(&'static str),
    /// The marker or the policy pin could not be read or written.
    Storage(String),
}

impl SwitchError {
    /// A stable machine-readable reason for the route reply.
    pub fn reason(&self) -> &'static str {
        match self {
            SwitchError::NoMarker => "no_marker",
            SwitchError::CredentialMismatch => "credential_mismatch",
            SwitchError::Worker(_) => "worker_unavailable",
            SwitchError::Storage(_) => "storage",
        }
    }
}

impl std::fmt::Display for SwitchError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SwitchError::NoMarker => write!(f, "no organization switch is pending"),
            SwitchError::CredentialMismatch => write!(
                f,
                "the stored credential is not the key of the pending organization switch"
            ),
            SwitchError::Worker(what) => {
                write!(f, "the cloud worker did not acknowledge {what}")
            }
            SwitchError::Storage(e) => write!(f, "{e}"),
        }
    }
}

/// `POST /admin/switch/begin` reply.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
pub struct BeginReply {
    /// Events sent with the old key while draining.
    pub sent: u64,
    /// Events still queued when the drain stopped; apply drops them.
    pub remaining: u64,
}

/// `POST /admin/switch/commit` reply.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
pub struct CommitReply {
    /// Old-organization events dropped by apply.
    pub dropped: u64,
}

struct Lease {
    id: u64,
    timer: tokio::task::JoinHandle<()>,
}

/// Where a switch reads and writes. Production uses the OpenLatch directory for
/// every field; tests point them at a temp dir.
#[derive(Debug, Clone)]
pub struct SwitchPaths {
    /// Holds `state/pending-switch.json` and `identity.json`.
    pub openlatch_dir: PathBuf,
    /// The policy store base (`<base>/policy/bundle.meta.json`).
    pub policy_base: PathBuf,
    /// `<openlatch_dir>/logs`, home of `fallback.jsonl`.
    pub log_dir: PathBuf,
}

impl SwitchPaths {
    /// Every path under one OpenLatch directory.
    pub fn under(openlatch_dir: PathBuf) -> Self {
        Self {
            policy_base: openlatch_dir.clone(),
            log_dir: openlatch_dir.join("logs"),
            openlatch_dir,
        }
    }
}

/// The switch controller. One per daemon, present when cloud forwarding runs.
pub struct SwitchCtl {
    cloud: CloudState,
    provider: Arc<dyn CredentialProvider>,
    outbox: Option<Arc<Outbox>>,
    cloud_tx: mpsc::Sender<CloudEvent>,
    paths: SwitchPaths,
    /// `(agent_id, source)` of the install's identity event, re-enqueued under
    /// the new key by apply step 5. `None` before provisioning.
    identity_event: Option<(String, String)>,
    /// Serializes begin, commit, abort and every apply.
    op: tokio::sync::Mutex<()>,
    lease: std::sync::Mutex<Option<Lease>>,
    next_lease: AtomicU64,
    lease_grace: Duration,
    /// Uploads are paused waiting for a marker check (daemon start, a failed
    /// apply) rather than by a lease: the next check that does not apply
    /// releases them.
    held_until_checked: AtomicBool,
    /// A marker this controller already applied but could not delete. Seen
    /// again, it is only deleted: re-applying would drop the new
    /// organization's events and bump the generation a second time.
    applied: std::sync::Mutex<Option<PendingSwitch>>,
    now: fn() -> u64,
}

impl SwitchCtl {
    /// A controller over the daemon's cloud pipeline.
    pub fn new(
        cloud: CloudState,
        provider: Arc<dyn CredentialProvider>,
        outbox: Option<Arc<Outbox>>,
        cloud_tx: mpsc::Sender<CloudEvent>,
        paths: SwitchPaths,
        identity_event: Option<(String, String)>,
    ) -> Self {
        Self {
            cloud,
            provider,
            outbox,
            cloud_tx,
            paths,
            identity_event,
            op: tokio::sync::Mutex::new(()),
            lease: std::sync::Mutex::new(None),
            next_lease: AtomicU64::new(0),
            lease_grace: LEASE_GRACE,
            held_until_checked: AtomicBool::new(false),
            applied: std::sync::Mutex::new(None),
            now: pending_switch::now_unix_secs,
        }
    }

    /// Override the lease grace (tests).
    #[cfg(test)]
    fn with_lease_grace(mut self, grace: Duration) -> Self {
        self.lease_grace = grace;
        self
    }

    /// Override the marker clock (tests).
    #[cfg(test)]
    fn with_clock(mut self, now: fn() -> u64) -> Self {
        self.now = now;
        self
    }

    /// Daemon start with a marker on disk: pause uploads, replay and polls
    /// **before** any of them start, until [`check_marker`](Self::check_marker)
    /// has applied the switch or released the hold (C-8, AC-28).
    pub fn hold_for_startup(&self) {
        self.held_until_checked.store(true, Ordering::Release);
        self.cloud.switch.set_paused(true);
    }

    /// The controller's loop: one marker check at start, then one whenever a
    /// credential load asks. Runs for the daemon's lifetime.
    pub async fn run(self: Arc<Self>) {
        self.check_marker().await;
        loop {
            self.cloud.switch.apply_requested().await;
            self.check_marker().await;
        }
    }

    /// `POST /admin/switch/begin`: send what is queued with the old key until
    /// it is empty or `timeout_ms` (clamped to 30 000) runs out, then pause
    /// uploads and fallback replay under a lease that resumes them by itself
    /// after `timeout_ms` + 60 s unless `commit` or `abort` comes first.
    pub async fn begin(self: &Arc<Self>, timeout_ms: u64) -> Result<BeginReply, SwitchError> {
        let _op = self.op.lock().await;
        let timeout = Duration::from_millis(timeout_ms.min(MAX_BEGIN_TIMEOUT_MS));
        let deadline = tokio::time::Instant::now() + timeout;
        let before = self.cloud.forwarded_count();

        // Nothing can leave without a usable old key: the queue counts as
        // dropped (C-8 preconditions), so skip straight to the pause.
        let can_send = !self.cloud.is_no_credential() && !self.cloud.is_auth_error();
        if can_send {
            self.cloud.notify_drain();
            loop {
                let held = self
                    .ask("a drain", |ack| WorkerCommand::Drain { deadline, ack })
                    .await?;
                let remaining = held + self.durable_backlog();
                if remaining == 0 || tokio::time::Instant::now() >= deadline {
                    break;
                }
                self.cloud.notify_drain();
                let left = deadline.saturating_duration_since(tokio::time::Instant::now());
                tokio::time::sleep(left.min(DRAIN_ROUND)).await;
            }
        }

        let held = match self.pause().await {
            Ok(held) => held,
            Err(e) => {
                // `begin` failed: uploads must not stay held without a lease.
                self.resume().await;
                return Err(e);
            }
        };
        let reply = BeginReply {
            sent: self.cloud.forwarded_count().saturating_sub(before),
            remaining: held + self.durable_backlog(),
        };
        self.start_lease(timeout + self.lease_grace);
        tracing::info!(
            sent = reply.sent,
            remaining = reply.remaining,
            "organization switch: uploads drained with the old key and paused"
        );
        Ok(reply)
    }

    /// `POST /admin/switch/commit`: apply when the marker matches the stored
    /// credential. [`SwitchError::NoMarker`] when there is no marker — the
    /// switch was already applied, or never started.
    pub async fn commit(self: &Arc<Self>) -> Result<CommitReply, SwitchError> {
        let _op = self.op.lock().await;
        let marker = pending_switch::read(&self.paths.openlatch_dir)
            .map_err(|e| SwitchError::Storage(format!("cannot read the switch marker: {e}")))?
            .ok_or(SwitchError::NoMarker)?;
        if self.finish_applied(&marker).await {
            return Ok(CommitReply { dropped: 0 });
        }
        let key = self.stored_key().await;
        match (marker.rule(key.as_ref(), (self.now)()), key) {
            (MarkerAction::Apply, Some(key)) => {
                let dropped = self.apply_locked(&marker, key).await?;
                Ok(CommitReply { dropped })
            }
            _ => Err(SwitchError::CredentialMismatch),
        }
    }

    /// `POST /admin/switch/abort`: end the lease and resume uploads and
    /// fallback replay. Changes nothing else.
    pub async fn abort(self: &Arc<Self>) {
        let _op = self.op.lock().await;
        self.end_lease();
        self.resume().await;
        tracing::info!("organization switch: aborted, uploads resumed");
    }

    /// The C-8 marker rule, applied to the stored credential: apply, keep, or
    /// delete an abandoned marker. Releases a startup hold unless it applied.
    pub async fn check_marker(self: &Arc<Self>) {
        let _op = self.op.lock().await;
        let marker = match pending_switch::read(&self.paths.openlatch_dir) {
            Ok(Some(marker)) => marker,
            Ok(None) => {
                self.release_hold().await;
                return;
            }
            Err(e) => {
                // Its state is unknown: it may be the marker of a key already
                // stored, with the old organization's queue still held. Keep
                // it, keep any hold (credential loads refuse the key while it
                // is unreadable), and read it again shortly.
                tracing::warn!(
                    error = %e,
                    "organization switch: the pending-switch marker cannot be read; credential use stays refused until it can"
                );
                self.retry_later();
                return;
            }
        };
        if self.finish_applied(&marker).await {
            return;
        }
        let key = self.stored_key().await;
        match (marker.rule(key.as_ref(), (self.now)()), key) {
            (MarkerAction::Apply, Some(key)) => match self.apply_locked(&marker, key).await {
                Ok(dropped) => tracing::info!(
                    dropped,
                    organization_id = %marker.organization_id,
                    "organization switch: applied at credential load"
                ),
                Err(e) => {
                    tracing::warn!(error = %e, "organization switch: apply failed; uploads stay paused until it succeeds")
                }
            },
            (MarkerAction::Expire, _) => {
                if let Err(e) = pending_switch::delete(&self.paths.openlatch_dir) {
                    tracing::warn!(error = %e, "organization switch: could not delete an abandoned marker");
                } else {
                    tracing::info!("organization switch: abandoned marker deleted (1 h old)");
                }
                self.release_hold().await;
            }
            // Keep: a switch may be in progress. The old credential keeps working.
            _ => self.release_hold().await,
        }
    }

    /// C-8 Apply, steps 1–6 in order. The caller holds `op`.
    async fn apply_locked(
        self: &Arc<Self>,
        marker: &PendingSwitch,
        key: SecretString,
    ) -> Result<u64, SwitchError> {
        // Uploads paused, and proven quiet, before anything is dropped.
        self.pause()
            .await
            .inspect_err(|_| self.hold_after_failure())?;

        // 1. Drop every old-organization event still held: the worker's batch
        //    and channel and the durable outbox, then the unread fallback tail.
        let mut dropped = self
            .ask("dropping the queue", |ack| WorkerCommand::DropQueued {
                ack,
            })
            .await
            .and_then(|outcome| outcome.map_err(SwitchError::Storage))
            .inspect_err(|_| self.hold_after_failure())?;
        dropped += crate::daemon::fallback_replay::skip_unread_tail_in(&self.paths.log_dir)
            .map_err(SwitchError::Storage)
            .inspect_err(|_| self.hold_after_failure())?;

        // 2–3. Bump the generation and re-bind the expected organization under
        //      the commit guard, so no poller response to an old-generation
        //      request writes after the re-bind.
        let generation = {
            let _commit = self.cloud.switch.commit_guard().await;
            let generation = self.cloud.switch.bump();
            crate::core::policy::store::rebind_organization(
                &self.paths.policy_base,
                &marker.organization_id,
            )
            .map_err(|e| {
                self.hold_after_failure();
                SwitchError::Storage(format!("cannot re-bind the policy organization: {e}"))
            })?;
            self.cloud.switch.set_rebound_org(&marker.organization_id);
            generation
        };

        // 4. Forget the old organization's id, install the new key in the
        //    worker and wait for it to confirm.
        crate::telemetry::identity::clear_org_id(
            crate::telemetry::global(),
            &self.paths.openlatch_dir,
        );
        self.provider.invalidate();
        self.ask("the new credential", |ack| WorkerCommand::SwapCredential {
            key,
            ack,
        })
        .await
        .inspect_err(|_| self.hold_after_failure())?;

        // 5. Materialize the install in the new organization.
        if let Some((agent_id, source)) = self.identity_event.as_ref() {
            let event = crate::cloud::startup_identity_event_at(agent_id, source, generation);
            let queued = tokio::time::timeout(WORKER_ACK_TIMEOUT, self.cloud_tx.send(event)).await;
            if !matches!(queued, Ok(Ok(()))) {
                tracing::warn!(
                    code = "OL-1200",
                    "organization switch: the identity event could not be queued"
                );
            }
        }

        // 6. Delete the marker, resume uploads and replay, wake the poller.
        if let Err(e) = pending_switch::delete(&self.paths.openlatch_dir) {
            // The switch is applied either way. The marker left behind is
            // recorded as done, so the next check only retries its deletion.
            tracing::warn!(error = %e, "organization switch: could not delete the marker; retrying");
            if let Ok(mut applied) = self.applied.lock() {
                *applied = Some(marker.clone());
            }
            self.retry_later();
        }
        self.end_lease();
        self.held_until_checked.store(false, Ordering::Release);
        self.resume().await;
        tracing::info!(
            organization_id = %marker.organization_id,
            dropped,
            generation,
            "organization switch: applied; the resident bundle keeps enforcing until the new organization's is accepted"
        );
        Ok(dropped)
    }

    /// Set the pause flag, wait for every drain or replay pass to finish, and
    /// for the worker to confirm no send is in flight. Returns how many events
    /// the worker still holds.
    async fn pause(&self) -> Result<u64, SwitchError> {
        self.cloud.switch.set_paused(true);
        self.cloud.switch.quiesce().await;
        self.ask("the pause", |ack| WorkerCommand::Pause { ack })
            .await
    }

    /// Clear the pause, then wake everything that waits on it.
    async fn resume(&self) {
        self.cloud.switch.set_paused(false);
        let _ = self
            .cloud
            .switch
            .worker_commands()
            .send(WorkerCommand::Resume)
            .await;
        self.cloud.notify_drain();
        self.cloud.policy_refresh_notify.notify_one();
    }

    /// A marker this controller already applied: retry its deletion only.
    /// Returns `false` when `marker` is not that one.
    async fn finish_applied(self: &Arc<Self>, marker: &PendingSwitch) -> bool {
        let done = self
            .applied
            .lock()
            .is_ok_and(|applied| applied.as_ref() == Some(marker));
        if !done {
            return false;
        }
        match pending_switch::delete(&self.paths.openlatch_dir) {
            Ok(()) => {
                if let Ok(mut applied) = self.applied.lock() {
                    *applied = None;
                }
                tracing::info!("organization switch: the applied marker is deleted");
            }
            Err(e) => {
                tracing::warn!(error = %e, "organization switch: could not delete the applied marker; retrying");
                self.retry_later();
            }
        }
        true
    }

    /// Ask for another marker check after [`MARKER_RETRY`].
    fn retry_later(self: &Arc<Self>) {
        let switch = self.cloud.switch.clone();
        tokio::spawn(async move {
            tokio::time::sleep(MARKER_RETRY).await;
            switch.request_apply();
        });
    }

    /// Resume only when uploads are held waiting for a marker check.
    async fn release_hold(&self) {
        if self.held_until_checked.swap(false, Ordering::AcqRel) {
            self.resume().await;
        }
    }

    /// An apply that stopped half-way keeps uploads paused — resuming could
    /// send new-generation events with the old key — and lets the next marker
    /// check (a credential load, a restart) finish or release it.
    fn hold_after_failure(&self) {
        self.end_lease();
        self.held_until_checked.store(true, Ordering::Release);
    }

    fn durable_backlog(&self) -> u64 {
        self.outbox.as_ref().map_or(0, |o| o.pending_count())
            + crate::daemon::fallback_replay::unread_entries_in(&self.paths.log_dir)
    }

    /// The stored credential, read fresh: the CLI wrote it from another
    /// process, so the in-process memo is dropped first.
    async fn stored_key(&self) -> Option<SecretString> {
        self.provider.invalidate();
        let provider = self.provider.clone();
        tokio::task::spawn_blocking(move || provider.retrieve())
            .await
            .unwrap_or(None)
    }

    async fn ask<T>(
        &self,
        what: &'static str,
        make: impl FnOnce(oneshot::Sender<T>) -> WorkerCommand,
    ) -> Result<T, SwitchError> {
        let (ack, reply) = oneshot::channel();
        let sent = tokio::time::timeout(
            WORKER_ACK_TIMEOUT,
            self.cloud.switch.worker_commands().send(make(ack)),
        )
        .await;
        if !matches!(sent, Ok(Ok(()))) {
            return Err(SwitchError::Worker(what));
        }
        match tokio::time::timeout(WORKER_ACK_TIMEOUT, reply).await {
            Ok(Ok(value)) => Ok(value),
            _ => Err(SwitchError::Worker(what)),
        }
    }

    fn start_lease(self: &Arc<Self>, duration: Duration) {
        let id = self.next_lease.fetch_add(1, Ordering::Relaxed) + 1;
        let ctl = Arc::clone(self);
        let timer = tokio::spawn(async move {
            tokio::time::sleep(duration).await;
            ctl.expire_lease(id).await;
        });
        if let Ok(mut lease) = self.lease.lock() {
            if let Some(old) = lease.replace(Lease { id, timer }) {
                old.timer.abort();
            }
        }
    }

    fn end_lease(&self) {
        if let Ok(mut lease) = self.lease.lock() {
            if let Some(old) = lease.take() {
                old.timer.abort();
            }
        }
    }

    /// The lease ran out with neither `commit` nor `abort`: resume as `abort`
    /// would (C-8 step 4, AC-29).
    async fn expire_lease(&self, id: u64) {
        let _op = self.op.lock().await;
        let current = self
            .lease
            .lock()
            .ok()
            .and_then(|mut lease| match lease.as_ref() {
                Some(l) if l.id == id => lease.take(),
                _ => None,
            });
        if current.is_some() {
            tracing::warn!(
                "organization switch: no commit or abort within the lease; uploads resumed"
            );
            self.resume().await;
        }
    }

    /// Whether a lease currently holds uploads (tests).
    #[cfg(test)]
    fn lease_active(&self) -> bool {
        self.lease.lock().is_ok_and(|l| l.is_some())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cloud::CloudConfig;
    use std::sync::Mutex;

    struct Provider {
        key: Mutex<Option<String>>,
    }

    impl Provider {
        fn with(key: &str) -> Arc<Self> {
            Arc::new(Self {
                key: Mutex::new(Some(key.to_string())),
            })
        }
        fn set(&self, key: &str) {
            *self.key.lock().unwrap() = Some(key.to_string());
        }
    }

    impl CredentialProvider for Provider {
        fn retrieve(&self) -> Option<SecretString> {
            self.key.lock().unwrap().clone().map(SecretString::from)
        }
    }

    struct Harness {
        ctl: Arc<SwitchCtl>,
        cloud: CloudState,
        provider: Arc<Provider>,
        tx: mpsc::Sender<CloudEvent>,
        dir: tempfile::TempDir,
        outbox: Arc<Outbox>,
        _worker: tokio::task::JoinHandle<()>,
    }

    async fn harness(server: &mockito::Server, key: &str, grace: Duration) -> Harness {
        let dir = tempfile::tempdir().unwrap();
        let cloud = CloudState::with_isolated_switch();
        let provider = Provider::with(key);
        let outbox = Arc::new(Outbox::new(dir.path(), 0));
        let (tx, rx) = mpsc::channel(32);
        let worker = tokio::spawn(crate::cloud::worker::run_cloud_worker(
            rx,
            provider.clone(),
            CloudConfig {
                api_url: server.url(),
                batch_max_wait_ms: 30,
                retry_delay_ms: 10,
                ..Default::default()
            },
            crate::egress::EgressReporter::direct(),
            cloud.clone(),
            dir.path().to_path_buf(),
            Some(outbox.clone()),
            None,
            crate::cloud::worker::SourceFormats::new(),
            Arc::new(crate::model_relay::session::SessionRegistry::default()),
        ));
        // Wait for the worker to load its key before driving it.
        for _ in 0..200 {
            if !cloud.is_no_credential() {
                break;
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
        let ctl = Arc::new(
            SwitchCtl::new(
                cloud.clone(),
                provider.clone(),
                Some(outbox.clone()),
                tx.clone(),
                SwitchPaths::under(dir.path().to_path_buf()),
                Some(("agt_test".into(), "claude-code".into())),
            )
            .with_lease_grace(grace),
        );
        Harness {
            ctl,
            cloud,
            provider,
            tx,
            dir,
            outbox,
            _worker: worker,
        }
    }

    fn seed_pin(base: &std::path::Path, org: &str) {
        let meta = serde_json::json!({
            "digest": "sha256:00",
            "etag": "\"sha256:00\"",
            "revision": 1,
            "organization_id": org,
            "built_at": "2026-10-12T00:00:00Z",
        });
        let dir = crate::core::policy::store::policy_dir(base);
        std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(
            crate::core::policy::store::meta_path(base),
            serde_json::to_vec(&meta).unwrap(),
        )
        .unwrap();
        std::fs::write(crate::core::policy::store::bundle_path(base), b"{}").unwrap();
    }

    async fn ingest(
        server: &mut mockito::Server,
        bearer: &str,
        body: &str,
        n: usize,
    ) -> mockito::Mock {
        server
            .mock("POST", "/api/v1/events/ingest")
            .match_header("authorization", format!("Bearer {bearer}").as_str())
            .match_body(mockito::Matcher::Regex(body.into()))
            .with_status(200)
            .with_body("{}")
            .expect(n)
            .create_async()
            .await
    }

    async fn wait_until(mut cond: impl FnMut() -> bool) -> bool {
        for _ in 0..1000 {
            if cond() {
                return true;
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
        false
    }

    #[tokio::test]
    async fn daemon_switch_begin_drains_with_the_old_key_and_clamps_the_timeout() {
        let mut server = mockito::Server::new_async().await;
        let old = ingest(&mut server, "old-key", "evt_queued", 1).await;
        let h = harness(&server, "old-key", LEASE_GRACE).await;
        // Held back by an earlier pause, so only `begin` can send it.
        h.cloud.switch.set_paused(true);
        h.tx.send(CloudEvent::at_generation(
            serde_json::json!({"id": "evt_queued"}),
            "agt_test",
            h.cloud.switch.generation(),
        ))
        .await
        .unwrap();

        let started = tokio::time::Instant::now();
        // 10 minutes asked, 30 s at most honoured; the drain ends as soon as
        // the queue is empty.
        let reply = h.ctl.begin(600_000).await.unwrap();
        assert!(started.elapsed() < Duration::from_secs(30));
        assert_eq!(reply.remaining, 0);
        assert!(reply.sent >= 1, "{reply:?}");
        assert!(
            h.cloud.switch.is_paused(),
            "uploads stay paused under the lease"
        );
        assert!(h.ctl.lease_active());
        old.assert_async().await;
    }

    #[tokio::test]
    async fn daemon_switch_lease_resumes_uploads_by_itself() {
        let server = mockito::Server::new_async().await;
        let h = harness(&server, "old-key", Duration::from_millis(50)).await;
        h.ctl.begin(0).await.unwrap();
        assert!(h.cloud.switch.is_paused());
        assert!(
            wait_until(|| !h.cloud.switch.is_paused()).await,
            "the lease auto-aborts after timeout_ms + grace"
        );
        assert!(!h.ctl.lease_active());
    }

    #[tokio::test]
    async fn daemon_switch_commit_without_a_marker_is_no_marker() {
        let server = mockito::Server::new_async().await;
        let h = harness(&server, "old-key", LEASE_GRACE).await;
        assert_eq!(h.ctl.commit().await, Err(SwitchError::NoMarker));
    }

    #[tokio::test]
    async fn daemon_switch_commit_with_a_mismatched_key_changes_nothing() {
        let server = mockito::Server::new_async().await;
        let h = harness(&server, "old-key", LEASE_GRACE).await;
        seed_pin(h.dir.path(), "org-a");
        let marker = PendingSwitch::new(
            "org-b",
            &SecretString::from("new-key".to_string()),
            pending_switch::now_unix_secs(),
        );
        pending_switch::write(h.dir.path(), &marker).unwrap();
        assert_eq!(h.ctl.commit().await, Err(SwitchError::CredentialMismatch));
        let meta = crate::core::policy::store::read_meta(h.dir.path())
            .unwrap()
            .unwrap();
        assert_eq!(meta.organization_id, "org-a");
        assert!(pending_switch::read(h.dir.path()).unwrap().is_some());
    }

    #[tokio::test]
    async fn daemon_switch_apply_drops_rebinds_swaps_and_resumes_in_order() {
        let mut server = mockito::Server::new_async().await;
        // Created first so it wins any match: the in-handler event must never
        // leave, with either key.
        let in_handler = server
            .mock("POST", "/api/v1/events/ingest")
            .match_body(mockito::Matcher::Regex("evt_in_handler".into()))
            .with_status(200)
            .with_body("{}")
            .expect(0)
            .create_async()
            .await;
        let old_after = ingest(&mut server, "old-key", ".", 0).await;
        let identity = ingest(&mut server, "new-key", "ai.openlatch.client.started", 1).await;
        let h = harness(&server, "old-key", LEASE_GRACE).await;
        seed_pin(h.dir.path(), "org-a");
        h.outbox
            .append(&serde_json::json!({"id": "evt_spooled_old"}))
            .unwrap();
        // fallback.jsonl with one unread line
        let logs = h.dir.path().join("logs");
        std::fs::create_dir_all(&logs).unwrap();
        std::fs::write(
            logs.join("fallback.jsonl"),
            "{\"id\":\"evt_fallback_old\"}\n",
        )
        .unwrap();

        let begin = h.ctl.begin(0).await.unwrap();
        assert!(begin.remaining >= 2, "{begin:?}");
        // An event captured before apply, still in flight.
        let captured = h.cloud.switch.generation();

        // CLI: marker, then the new key.
        let new_key = SecretString::from("new-key".to_string());
        pending_switch::write(
            h.dir.path(),
            &PendingSwitch::new("org-b", &new_key, pending_switch::now_unix_secs()),
        )
        .unwrap();
        h.provider.set("new-key");

        let reply = h.ctl.commit().await.unwrap();
        assert!(
            reply.dropped >= 2,
            "outbox and fallback tail dropped: {reply:?}"
        );
        // Re-bound, bundle body untouched, marker gone, uploads resumed.
        let meta = crate::core::policy::store::read_meta(h.dir.path())
            .unwrap()
            .unwrap();
        assert_eq!(meta.organization_id, "org-b");
        assert!(crate::core::policy::store::bundle_path(h.dir.path()).exists());
        assert_eq!(h.cloud.switch.rebound_org().as_deref(), Some("org-b"));
        assert!(pending_switch::read(h.dir.path()).unwrap().is_none());
        assert!(!h.cloud.switch.is_paused());
        assert!(!h.ctl.lease_active());
        assert_eq!(h.outbox.pending_count(), 0);
        assert!(h.cloud.switch.generation() > captured);

        assert!(
            wait_until(|| h.cloud.forwarded_count() >= 1).await,
            "the identity event is forwarded under the new key"
        );
        // The event captured before apply reaches the worker only now, after
        // the identity batch left, so a leak would be a batch of its own that
        // only the `in_handler` mock matches.
        h.tx.send(CloudEvent::at_generation(
            serde_json::json!({"id": "evt_in_handler"}),
            "agt_test",
            captured,
        ))
        .await
        .unwrap();
        tokio::time::sleep(Duration::from_millis(100)).await;
        identity.assert_async().await;
        old_after.assert_async().await;
        in_handler.assert_async().await;
    }

    #[tokio::test]
    async fn daemon_switch_abort_resumes_and_changes_nothing_else() {
        let server = mockito::Server::new_async().await;
        let h = harness(&server, "old-key", LEASE_GRACE).await;
        seed_pin(h.dir.path(), "org-a");
        h.ctl.begin(0).await.unwrap();
        let generation = h.cloud.switch.generation();
        h.ctl.abort().await;
        assert!(!h.cloud.switch.is_paused());
        assert!(!h.ctl.lease_active());
        assert_eq!(h.cloud.switch.generation(), generation);
        let meta = crate::core::policy::store::read_meta(h.dir.path())
            .unwrap()
            .unwrap();
        assert_eq!(meta.organization_id, "org-a");
    }

    #[tokio::test]
    async fn daemon_switch_startup_hold_applies_before_release() {
        let mut server = mockito::Server::new_async().await;
        let _identity = ingest(&mut server, "new-key", "ai.openlatch.client.started", 1).await;
        let h = harness(&server, "new-key", LEASE_GRACE).await;
        seed_pin(h.dir.path(), "org-a");
        let new_key = SecretString::from("new-key".to_string());
        pending_switch::write(
            h.dir.path(),
            &PendingSwitch::new("org-b", &new_key, pending_switch::now_unix_secs()),
        )
        .unwrap();
        h.ctl.hold_for_startup();
        assert!(h.cloud.switch.is_paused());
        h.ctl.check_marker().await;
        assert!(!h.cloud.switch.is_paused());
        assert!(pending_switch::read(h.dir.path()).unwrap().is_none());
        let meta = crate::core::policy::store::read_meta(h.dir.path())
            .unwrap()
            .unwrap();
        assert_eq!(meta.organization_id, "org-b");
    }

    fn two_hours_later() -> u64 {
        pending_switch::now_unix_secs() + 7200
    }

    #[tokio::test]
    async fn daemon_switch_marker_rule_keeps_a_young_marker_and_expires_an_old_one() {
        let server = mockito::Server::new_async().await;
        let h = harness(&server, "old-key", LEASE_GRACE).await;
        seed_pin(h.dir.path(), "org-a");
        let marker = PendingSwitch::new(
            "org-b",
            &SecretString::from("new-key".to_string()),
            pending_switch::now_unix_secs(),
        );
        pending_switch::write(h.dir.path(), &marker).unwrap();

        // Young: the CLI died between the marker and the key — keep it, keep
        // the old key and pin, release the startup hold.
        h.ctl.hold_for_startup();
        h.ctl.check_marker().await;
        assert!(!h.cloud.switch.is_paused());
        assert!(pending_switch::read(h.dir.path()).unwrap().is_some());

        // An hour later: abandoned, deleted; the pin never moved.
        let later = Arc::new(
            SwitchCtl::new(
                h.cloud.clone(),
                h.provider.clone(),
                Some(h.outbox.clone()),
                h.tx.clone(),
                SwitchPaths::under(h.dir.path().to_path_buf()),
                None,
            )
            .with_clock(two_hours_later),
        );
        later.check_marker().await;
        assert!(pending_switch::read(h.dir.path()).unwrap().is_none());
        let meta = crate::core::policy::store::read_meta(h.dir.path())
            .unwrap()
            .unwrap();
        assert_eq!(meta.organization_id, "org-a");
    }

    /// An unreadable marker is kept, and a startup hold stays: its state is
    /// unknown, so nothing may treat it as abandoned.
    #[tokio::test]
    async fn daemon_switch_an_unreadable_marker_is_kept_and_holds_uploads() {
        let server = mockito::Server::new_async().await;
        let h = harness(&server, "old-key", LEASE_GRACE).await;
        let path = pending_switch::marker_path(h.dir.path());
        std::fs::create_dir_all(path.parent().unwrap()).unwrap();
        std::fs::write(&path, b"{not json").unwrap();
        h.ctl.hold_for_startup();

        h.ctl.check_marker().await;

        assert!(path.exists(), "the marker is kept");
        assert!(h.cloud.switch.is_paused(), "uploads stay held");
    }

    /// A marker already applied but not deleted is only deleted when seen
    /// again: no second drop, no second generation bump.
    #[tokio::test]
    async fn daemon_switch_an_applied_marker_is_never_applied_twice() {
        let server = mockito::Server::new_async().await;
        let h = harness(&server, "new-key", LEASE_GRACE).await;
        seed_pin(h.dir.path(), "org-b");
        let marker = PendingSwitch::new(
            "org-b",
            &SecretString::from("new-key".to_string()),
            pending_switch::now_unix_secs(),
        );
        pending_switch::write(h.dir.path(), &marker).unwrap();
        *h.ctl.applied.lock().unwrap() = Some(marker);
        h.outbox
            .append(&serde_json::json!({"id": "evt_new_org"}))
            .unwrap();
        let generation = h.cloud.switch.generation();

        h.ctl.check_marker().await;

        assert!(pending_switch::read(h.dir.path()).unwrap().is_none());
        assert_eq!(h.cloud.switch.generation(), generation, "no second bump");
        assert_eq!(
            h.outbox.pending_count(),
            1,
            "the new organization's event stays"
        );
        assert!(h.ctl.applied.lock().unwrap().is_none());
    }
}