openlatch-client 0.5.3

//! Bounded Tier-3 hold lifecycle owned by the daemon.

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use secrecy::ExposeSecret;
use tokio::sync::Notify;

use crate::cloud::CredentialProvider;
use crate::generated::types::Verdict;
use crate::zone_eval::{Event, HoldRequest};

pub const DEFAULT_MAX_PENDING: usize = 16;
pub const DEFAULT_HOST_TIMEOUT: Duration = Duration::from_secs(900);
const DAEMON_DEADLINE_MARGIN: Duration = Duration::from_secs(30);
const COMPLETED_RETENTION: Duration = Duration::from_secs(5 * 60);
const MAX_COMPLETED: usize = 64;
const MAX_PLATFORM_WAIT: Duration = Duration::from_secs(5);
const RETRY_DELAY: Duration = Duration::from_millis(250);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HoldAnswer {
    Approved,
    Rejected,
    Retired,
    Timeout,
}

/// Bounded evidence registered with the configured platform hold endpoint.
#[derive(Debug, Clone)]
pub struct HoldRegistration {
    pub tool_use_id: String,
    pub atom_id: Option<String>,
    pub policy_id: Option<String>,
    pub policy_public_id: Option<String>,
    pub effects: Vec<serde_json::Value>,
}

/// Platform answer seam. Production can implement the I-6 endpoint contract;
/// deterministic tests never need a network service.
pub trait HoldAnswerSource: Send + Sync {
    fn answer(
        &self,
        registration: HoldRegistration,
        timeout: Duration,
        deadline: tokio::time::Instant,
    ) -> Pin<Box<dyn Future<Output = HoldAnswer> + Send>>;
}

/// Honest fallback while the platform has no hold router: wait to the daemon
/// deadline and decide `on_timeout`, never allow early.
pub struct DeadlineAnswerSource;

impl HoldAnswerSource for DeadlineAnswerSource {
    fn answer(
        &self,
        _registration: HoldRegistration,
        _timeout: Duration,
        deadline: tokio::time::Instant,
    ) -> Pin<Box<dyn Future<Output = HoldAnswer> + Send>> {
        Box::pin(async move {
            tokio::time::sleep_until(deadline).await;
            HoldAnswer::Timeout
        })
    }
}

/// Authenticated production source for the I-6 hold lifecycle.
///
/// This source is called only from the queue's spawned task after evaluation;
/// the hook itself continues to wait exclusively on the daemon's loopback
/// route, so remote I/O never enters the verdict path.
pub struct CloudHoldAnswerSource {
    url: String,
    credentials: Arc<dyn CredentialProvider>,
    http: crate::egress::ClientHandle,
}

impl CloudHoldAnswerSource {
    pub fn new(
        api_url: &str,
        credentials: Arc<dyn CredentialProvider>,
        http: crate::egress::ClientHandle,
    ) -> Self {
        Self {
            url: format!("{}/api/v1/holds", api_url.trim_end_matches('/')),
            credentials,
            http,
        }
    }

    async fn register(
        &self,
        registration: &HoldRegistration,
        timeout: Duration,
        deadline: tokio::time::Instant,
    ) -> bool {
        let (Some(http), Some(token)) = (self.http.current(), self.credentials.retrieve()) else {
            return false;
        };
        let body = serde_json::json!({
            "tool_use_id": registration.tool_use_id,
            "atom_id": registration.atom_id,
            "policy_id": registration.policy_id,
            "policy_public_id": registration.policy_public_id,
            "effects": registration.effects,
            "timeout_s": timeout.as_secs().clamp(1, 3_600),
        });
        let request = http
            .post(&self.url)
            .bearer_auth(token.expose_secret())
            .json(&body)
            .send();
        match tokio::time::timeout_at(deadline, request).await {
            Ok(Ok(response)) => matches!(response.status().as_u16(), 201 | 409),
            Ok(Err(error)) => {
                tracing::debug!(target: "policy", %error, "hold registration failed within the local deadline");
                false
            }
            Err(_) => false,
        }
    }

    async fn poll(
        &self,
        tool_use_id: &str,
        wait: Duration,
        deadline: tokio::time::Instant,
    ) -> Result<Option<HoldAnswer>, ()> {
        let (Some(http), Some(token)) = (self.http.current(), self.credentials.retrieve()) else {
            return Ok(None);
        };
        let mut url = match reqwest::Url::parse(&self.url) {
            Ok(url) => url,
            Err(_) => return Ok(None),
        };
        let Ok(mut segments) = url.path_segments_mut() else {
            return Ok(None);
        };
        segments.push(tool_use_id);
        drop(segments);
        url.query_pairs_mut()
            .append_pair("wait", &wait.as_secs().to_string());
        let request = http.get(url).bearer_auth(token.expose_secret()).send();
        tokio::time::timeout_at(deadline, async move {
            let response = request.await.ok()?;
            if !response.status().is_success() {
                return None;
            }
            let body: serde_json::Value = response.json().await.ok()?;
            match body
                .pointer("/data/state")
                .and_then(serde_json::Value::as_str)
            {
                Some("approved") => Some(HoldAnswer::Approved),
                Some("rejected") => Some(HoldAnswer::Rejected),
                Some("retired") => Some(HoldAnswer::Retired),
                Some("timed_out") => Some(HoldAnswer::Timeout),
                _ => None,
            }
        })
        .await
        .map_err(|_| ())
    }

    async fn run(
        &self,
        registration: HoldRegistration,
        timeout: Duration,
        deadline: tokio::time::Instant,
    ) -> HoldAnswer {
        let mut registered = false;
        loop {
            let now = tokio::time::Instant::now();
            if now >= deadline {
                return HoldAnswer::Timeout;
            }
            if !registered {
                registered = self.register(&registration, timeout, deadline).await;
                if !registered {
                    tokio::time::sleep_until((now + RETRY_DELAY).min(deadline)).await;
                    continue;
                }
            }
            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
            let wait = remaining.min(MAX_PLATFORM_WAIT);
            match self.poll(&registration.tool_use_id, wait, deadline).await {
                Ok(Some(answer)) => return answer,
                Err(()) => return HoldAnswer::Timeout,
                Ok(None) => {}
            }
            tokio::time::sleep_until((tokio::time::Instant::now() + RETRY_DELAY).min(deadline))
                .await;
        }
    }
}

impl HoldAnswerSource for CloudHoldAnswerSource {
    fn answer(
        &self,
        registration: HoldRegistration,
        timeout: Duration,
        deadline: tokio::time::Instant,
    ) -> Pin<Box<dyn Future<Output = HoldAnswer> + Send>> {
        let source = Self {
            url: self.url.clone(),
            credentials: self.credentials.clone(),
            http: self.http.clone(),
        };
        Box::pin(async move { source.run(registration, timeout, deadline).await })
    }
}

#[derive(Debug, Clone)]
pub struct HoldResolution {
    pub answer: HoldAnswer,
    pub verdict: Verdict,
    pub result: &'static str,
    pub event: Event,
    pub correlation_id: String,
    pub original_time: String,
    pub subject: Option<String>,
}

struct PendingHold {
    request: HoldRequest,
    event: Event,
    correlation_id: String,
    original_time: String,
    subject: Option<String>,
    resolution: Mutex<Option<HoldResolution>>,
    completed_at: Mutex<Option<Instant>>,
    notify: Notify,
}

/// Registration outcome; a full queue is an immediate timeout decision.
pub enum RegisterResult {
    Pending { timeout: Duration },
    Full(Box<HoldResolution>),
}

pub struct HoldQueue {
    pending: Mutex<HashMap<String, Arc<PendingHold>>>,
    source: Mutex<Arc<dyn HoldAnswerSource>>,
}

impl Default for HoldQueue {
    fn default() -> Self {
        Self::new(Arc::new(DeadlineAnswerSource))
    }
}

impl HoldQueue {
    pub fn new(source: Arc<dyn HoldAnswerSource>) -> Self {
        Self {
            pending: Mutex::new(HashMap::new()),
            source: Mutex::new(source),
        }
    }

    /// Install the authenticated source once credentials and a routed client
    /// both exist. Entries already registered retain their captured source.
    pub fn set_source(&self, source: Arc<dyn HoldAnswerSource>) {
        *self.source.lock().unwrap_or_else(|e| e.into_inner()) = source;
    }

    #[allow(clippy::too_many_arguments)]
    pub fn register(
        self: &Arc<Self>,
        registration: HoldRegistration,
        request: HoldRequest,
        event: Event,
        correlation_id: String,
        original_time: String,
        subject: Option<String>,
        host_timeout: Duration,
        max_pending: usize,
    ) -> RegisterResult {
        let atom_timeout = Duration::from_secs(request.timeout_s.unwrap_or(0).max(0) as u64);
        let daemon_ceiling = host_timeout.saturating_sub(DAEMON_DEADLINE_MARGIN);
        let timeout = if atom_timeout.is_zero() {
            daemon_ceiling
        } else {
            atom_timeout.min(daemon_ceiling)
        };
        let mut pending = self.pending.lock().unwrap_or_else(|e| e.into_inner());
        purge_completed(&mut pending, Instant::now());
        if pending.contains_key(&registration.tool_use_id) {
            return RegisterResult::Pending { timeout };
        }
        let unresolved = pending
            .values()
            .filter(|entry| {
                entry
                    .resolution
                    .lock()
                    .map(|resolution| resolution.is_none())
                    .unwrap_or(true)
            })
            .count();
        if unresolved >= max_pending.max(1) {
            return RegisterResult::Full(Box::new(resolve(
                HoldAnswer::Timeout,
                &request,
                event,
                correlation_id,
                original_time,
                subject,
            )));
        }
        let entry = Arc::new(PendingHold {
            request,
            event,
            correlation_id,
            original_time,
            subject,
            resolution: Mutex::new(None),
            completed_at: Mutex::new(None),
            notify: Notify::new(),
        });
        pending.insert(registration.tool_use_id.clone(), entry.clone());
        drop(pending);

        let source = self
            .source
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .clone();
        let queue = Arc::downgrade(self);
        tokio::spawn(async move {
            let deadline = tokio::time::Instant::now() + timeout;
            let answer = source.answer(registration, timeout, deadline).await;
            let resolution = resolve(
                answer,
                &entry.request,
                entry.event.clone(),
                entry.correlation_id.clone(),
                entry.original_time.clone(),
                entry.subject.clone(),
            );
            *entry.resolution.lock().unwrap_or_else(|e| e.into_inner()) = Some(resolution);
            *entry.completed_at.lock().unwrap_or_else(|e| e.into_inner()) = Some(Instant::now());
            entry.notify.notify_waiters();
            if let Some(queue) = queue.upgrade() {
                queue.trim_completed();
            }
        });
        RegisterResult::Pending { timeout }
    }

    pub async fn wait(&self, tool_use_id: &str) -> Option<HoldResolution> {
        let entry = self.pending.lock().ok()?.get(tool_use_id).cloned()?;
        loop {
            // Register the waiter before inspecting the condition. A completion
            // between these two lines leaves a permit for this future instead
            // of becoming a lost notification and an unbounded wait.
            let notified = entry.notify.notified();
            if let Some(resolution) = entry.resolution.lock().ok()?.clone() {
                self.pending.lock().ok()?.remove(tool_use_id);
                return Some(resolution);
            }
            notified.await;
        }
    }

    pub fn pending_count(&self) -> usize {
        self.pending
            .lock()
            .map(|entries| {
                entries
                    .values()
                    .filter(|entry| {
                        entry
                            .resolution
                            .lock()
                            .map(|resolution| resolution.is_none())
                            .unwrap_or(true)
                    })
                    .count()
            })
            .unwrap_or(0)
    }

    fn trim_completed(&self) {
        let Ok(mut entries) = self.pending.lock() else {
            return;
        };
        purge_completed(&mut entries, Instant::now());
    }
}

fn purge_completed(entries: &mut HashMap<String, Arc<PendingHold>>, now: Instant) {
    entries.retain(|_, entry| {
        entry
            .completed_at
            .lock()
            .map(|completed| {
                completed.is_none_or(|at| now.saturating_duration_since(at) <= COMPLETED_RETENTION)
            })
            .unwrap_or(false)
    });
    let mut completed: Vec<(String, Instant)> = entries
        .iter()
        .filter_map(|(id, entry)| {
            entry
                .completed_at
                .lock()
                .ok()
                .and_then(|completed| completed.map(|at| (id.clone(), at)))
        })
        .collect();
    completed.sort_by_key(|(_, at)| *at);
    let excess = completed.len().saturating_sub(MAX_COMPLETED);
    for (id, _) in completed.into_iter().take(excess) {
        entries.remove(&id);
    }
}

fn resolve(
    answer: HoldAnswer,
    request: &HoldRequest,
    event: Event,
    correlation_id: String,
    original_time: String,
    subject: Option<String>,
) -> HoldResolution {
    let (verdict, result) = match answer {
        HoldAnswer::Approved => (
            request.verdict_on_approve.unwrap_or(Verdict::Allow),
            "held_approved",
        ),
        HoldAnswer::Rejected => (
            request.verdict_on_reject.unwrap_or(Verdict::Block),
            "held_rejected",
        ),
        HoldAnswer::Timeout => (request.on_timeout.unwrap_or(Verdict::Block), "held_timeout"),
        // The handler re-evaluates this against the then-current resident bundle.
        HoldAnswer::Retired => (Verdict::Ask, "deferred"),
    };
    HoldResolution {
        answer,
        verdict,
        result,
        event,
        correlation_id,
        original_time,
        subject,
    }
}

/// A CloudEvents time for a follow-up that can never predate its original.
pub fn resolution_time(original: &str) -> String {
    let now = chrono::Utc::now();
    let original = chrono::DateTime::parse_from_rfc3339(original)
        .map(|dt| dt.with_timezone(&chrono::Utc))
        .unwrap_or(now);
    now.max(original)
        .to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
}

#[cfg(test)]
mod tests {
    use super::*;
    use secrecy::SecretString;

    struct FixedCredential;

    impl CredentialProvider for FixedCredential {
        fn retrieve(&self) -> Option<SecretString> {
            Some(SecretString::from("ol_hold_test".to_string()))
        }
    }

    struct Immediate(HoldAnswer);
    impl HoldAnswerSource for Immediate {
        fn answer(
            &self,
            _: HoldRegistration,
            _: Duration,
            _: tokio::time::Instant,
        ) -> Pin<Box<dyn Future<Output = HoldAnswer> + Send>> {
            let answer = self.0;
            Box::pin(async move { answer })
        }
    }

    fn request() -> HoldRequest {
        HoldRequest {
            resolve: Some("ask_human".into()),
            directive_template_id: None,
            max_attempts: None,
            timeout_s: Some(120),
            on_timeout: Some(Verdict::Block),
            verdict_on_approve: Some(Verdict::Allow),
            verdict_on_reject: Some(Verdict::Block),
        }
    }

    fn registration(tool_use_id: String) -> HoldRegistration {
        HoldRegistration {
            tool_use_id,
            atom_id: None,
            policy_id: None,
            policy_public_id: None,
            effects: Vec::new(),
        }
    }

    #[tokio::test]
    async fn registration_is_idempotent_and_resolution_is_exact() {
        let queue = Arc::new(HoldQueue::new(Arc::new(Immediate(HoldAnswer::Approved))));
        for _ in 0..2 {
            assert!(matches!(
                queue.register(
                    registration("tool-1".into()),
                    request(),
                    Event::default(),
                    "corr".into(),
                    "2026-09-10T00:00:00Z".into(),
                    Some("s".into()),
                    DEFAULT_HOST_TIMEOUT,
                    16
                ),
                RegisterResult::Pending { .. }
            ));
        }
        let resolved = queue.wait("tool-1").await.unwrap();
        assert_eq!(resolved.verdict, Verdict::Allow);
        assert_eq!(resolved.result, "held_approved");
    }

    #[tokio::test]
    async fn cloud_source_registers_with_auth_and_uses_a_bounded_wait() {
        let mut server = mockito::Server::new_async().await;
        let post = server
            .mock("POST", "/api/v1/holds")
            .match_header("authorization", "Bearer ol_hold_test")
            .match_body(mockito::Matcher::PartialJson(serde_json::json!({
                "tool_use_id": "tool-cloud",
                "atom_id": "atom-cloud",
                "policy_id": "policy-cloud",
                "policy_public_id": "AIP-007",
                "effects": [{"verb": "write", "target_class": "repository"}],
                "timeout_s": 10,
            })))
            .with_status(201)
            .with_body(r#"{"data":{"state":"pending"}}"#)
            .create_async()
            .await;
        let get = server
            .mock("GET", "/api/v1/holds/tool-cloud")
            .match_header("authorization", "Bearer ol_hold_test")
            .match_query(mockito::Matcher::Regex(r"(^|&)wait=5(&|$)".into()))
            .with_status(200)
            .with_body(r#"{"data":{"state":"approved"}}"#)
            .create_async()
            .await;
        let source = CloudHoldAnswerSource::new(
            &server.url(),
            Arc::new(FixedCredential),
            crate::egress::ClientHandle::of(crate::egress::client()),
        );
        let answer = source
            .answer(
                HoldRegistration {
                    tool_use_id: "tool-cloud".into(),
                    atom_id: Some("atom-cloud".into()),
                    policy_id: Some("policy-cloud".into()),
                    policy_public_id: Some("AIP-007".into()),
                    effects: vec![serde_json::json!({
                        "verb": "write",
                        "target_class": "repository"
                    })],
                },
                Duration::from_secs(10),
                tokio::time::Instant::now() + Duration::from_secs(10),
            )
            .await;
        assert_eq!(answer, HoldAnswer::Approved);
        post.assert_async().await;
        get.assert_async().await;
    }

    #[tokio::test]
    async fn cloud_source_deadline_bounds_a_delayed_poll_body() {
        let mut server = mockito::Server::new_async().await;
        let post = server
            .mock("POST", "/api/v1/holds")
            .with_status(201)
            .create_async()
            .await;
        let get = server
            .mock("GET", "/api/v1/holds/tool-slow-body")
            .match_query(mockito::Matcher::UrlEncoded("wait".into(), "0".into()))
            .with_status(200)
            .with_chunked_body(|writer| {
                std::thread::sleep(Duration::from_millis(300));
                writer.write_all(br#"{"data":{"state":"approved"}}"#)
            })
            .create_async()
            .await;
        let source = CloudHoldAnswerSource::new(
            &server.url(),
            Arc::new(FixedCredential),
            crate::egress::ClientHandle::of(crate::egress::client()),
        );
        let started = tokio::time::Instant::now();
        let answer = source
            .answer(
                registration("tool-slow-body".into()),
                Duration::from_secs(10),
                started + Duration::from_millis(100),
            )
            .await;

        assert_eq!(answer, HoldAnswer::Timeout);
        assert!(started.elapsed() < Duration::from_millis(250));
        post.assert_async().await;
        get.assert_async().await;
    }

    #[tokio::test]
    async fn reject_and_retire_keep_their_distinct_resolution_semantics() {
        for (answer, verdict, result) in [
            (HoldAnswer::Rejected, Verdict::Block, "held_rejected"),
            (HoldAnswer::Retired, Verdict::Ask, "deferred"),
        ] {
            let queue = Arc::new(HoldQueue::new(Arc::new(Immediate(answer))));
            assert!(matches!(
                queue.register(
                    registration(format!("tool-{result}")),
                    request(),
                    Event::default(),
                    "corr".into(),
                    "2026-09-10T00:00:00Z".into(),
                    Some("s".into()),
                    DEFAULT_HOST_TIMEOUT,
                    16,
                ),
                RegisterResult::Pending { .. }
            ));
            let resolution = queue.wait(&format!("tool-{result}")).await.unwrap();
            assert_eq!(resolution.verdict, verdict);
            assert_eq!(resolution.result, result);
        }
    }

    #[tokio::test]
    async fn completed_unconsumed_holds_release_capacity_and_stay_bounded() {
        let queue = Arc::new(HoldQueue::new(Arc::new(Immediate(HoldAnswer::Approved))));
        for index in 0..(MAX_COMPLETED + 2) {
            let id = format!("tool-{index}");
            assert!(matches!(
                queue.register(
                    registration(id),
                    request(),
                    Event::default(),
                    "corr".into(),
                    "2026-09-10T00:00:00Z".into(),
                    Some("s".into()),
                    DEFAULT_HOST_TIMEOUT,
                    1,
                ),
                RegisterResult::Pending { .. }
            ));
            tokio::task::yield_now().await;
            assert_eq!(queue.pending_count(), 0, "completed holds are not pending");
        }
        tokio::task::yield_now().await;
        assert!(
            queue.pending.lock().unwrap().len() <= MAX_COMPLETED,
            "completed-but-unconsumed retention is bounded"
        );
    }

    #[test]
    fn resolution_time_never_predates_the_original() {
        let future = "2099-01-01T00:00:00Z";
        assert!(resolution_time(future).starts_with("2099-01-01"));
    }
}