openlatch-client 0.5.4

//! Relay endpoints — one extra loopback listener per wired provider slot.
//!
//! The main relay port answers "which upstream?" from the request's wire format,
//! and that is enough while every agent speaking a format goes to one place.
//! It stops being enough for an agent that carries many providers at once:
//! a developer with LM Studio and DeepSeek configured sends two
//! OpenAI-compatible streams to two different hosts, and a route cannot tell
//! them apart. A GUI-hosted agent also has no header or path we can mark — its
//! own client builds the path from the base URL we write.
//!
//! So a provider slot gets its **own port**. Everything else is the same relay:
//! the same handler, the same measurement, policy and cloud pipeline, run with
//! one extra piece of state — the [`RelayEndpoint`] the request arrived on.
//! That endpoint names the one origin its traffic goes to (an **origin swap**:
//! scheme, host and port replaced, path and query forwarded verbatim) and the
//! agent it belongs to.
//!
//! Nothing in this module decides WHICH slots exist or writes any agent
//! config. It owns the listeners and the arithmetic every other party must
//! agree on: the port block, and what counts as an origin.

use std::collections::BTreeMap;
use std::ops::RangeInclusive;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use arc_swap::ArcSwap;
use tokio::sync::watch;
use tokio::task::JoinHandle;

use crate::core::supervision::task::{supervise, RestartPolicy, TaskHealth, TaskSpec};
use crate::error::{OlError, ERR_MODEL_RELAY_ENDPOINT_PORTS};

use super::wire_format::WireFormat;
use super::ModelRelayState;

/// How many endpoint ports sit above the main relay port.
///
/// A block DERIVED from the main port rather than a tenth isolation seam: the
/// main port already moves per instance (`OPENLATCH_MODEL_RELAY_PORT`), so the
/// block moves with it and two instances that do not collide on their main
/// ports cannot collide on their blocks either, provided their main ports are
/// at least `ENDPOINT_BLOCK + 1` apart. Thirty-two slots is far more than any
/// real agent configures; a host that exhausts it is told so with a code.
pub const ENDPOINT_BLOCK: u16 = 32;

/// Supervisor name for an endpoint listener.
///
/// Endpoints come and go while the daemon runs, and the health registry is
/// append-only with `&'static` names, so these are supervised without being
/// registered. Their health is reported per endpoint on the status surface
/// instead.
pub const ENDPOINT_TASK: &str = "model-relay-endpoint";

/// How long a released listener gets to drain before it is abandoned.
const RELEASE_DRAIN: Duration = Duration::from_secs(5);

/// The endpoint ports for a relay whose main listener is on `main`:
/// `main + 1 ..= main + ENDPOINT_BLOCK`.
///
/// **Empty** when the block would run past `u16::MAX`, and for `main == 0` —
/// an ephemeral test listener has no fixed port for a block to hang off. An
/// empty range contains nothing, so every caller's `contains` stays correct
/// without a special case.
pub fn endpoint_port_block(main: u16) -> RangeInclusive<u16> {
    if main == 0 {
        return empty_block();
    }
    match (main.checked_add(1), main.checked_add(ENDPOINT_BLOCK)) {
        (Some(lo), Some(hi)) => lo..=hi,
        _ => empty_block(),
    }
}

#[allow(clippy::reversed_empty_ranges)]
fn empty_block() -> RangeInclusive<u16> {
    1..=0
}

/// The loopback ports this client listens on — the ones that may never be an
/// upstream, a proxy, or a customer's own endpoint.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RelayPorts {
    /// The hook daemon's HTTP port.
    pub daemon: u16,
    /// The main model relay port.
    pub main: u16,
}

impl RelayPorts {
    /// Whether `port` is the daemon, the main relay, or inside the endpoint
    /// block.
    pub fn contains(&self, port: u16) -> bool {
        port == self.daemon || port == self.main || endpoint_port_block(self.main).contains(&port)
    }

    /// Whether `port` is one this relay may hand to an endpoint.
    pub fn in_block(&self, port: u16) -> bool {
        endpoint_port_block(self.main).contains(&port)
    }
}

/// Why a URL cannot be an endpoint's origin.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum OriginRefused {
    /// Not `http` or `https` — there is nothing to forward to.
    #[error("scheme \"{0}\" is not http or https")]
    Scheme(String),
    /// No host to connect to.
    #[error("the URL names no host")]
    NoHost,
    /// One of our own listeners. Forwarding there is an infinite loop dressed
    /// as a configuration, and it is exactly what a relay URL read back off
    /// disk looks like — so this refusal is what stops a slot from adopting
    /// its own previous value as the provider it replaced.
    #[error("the URL names this client's own listener on port {0}")]
    OwnListener(u16),
}

/// An origin an endpoint forwards to: scheme, host and port, with path `/`.
///
/// A newtype so an unnormalised URL cannot reach [`RelayEndpoint::set_origin`].
/// The path matters more than it looks: `proxy::join_upstream` treats a base
/// carrying a path as an API root and rewrites the request's own version
/// segment, which is right for an operator's gateway and wrong for an origin
/// swap, where the agent built the whole path itself. With the path fixed at
/// `/` the join keeps the request's path and query exactly as sent.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Origin(reqwest::Url);

impl Origin {
    /// The origin as a URL (path `/`, no query, no userinfo).
    pub fn as_url(&self) -> &reqwest::Url {
        &self.0
    }
}

impl std::fmt::Display for Origin {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.0.as_str())
    }
}

/// Reduce `url` to its origin, refusing anything that is not a forwardable
/// origin or that is one of our own listeners.
///
/// Userinfo is dropped: an origin is scheme, host and port by definition, and
/// the agent's own HTTP client sends its credential in a header, not in the URL
/// it was given.
pub fn normalize_origin(url: &reqwest::Url, own: &RelayPorts) -> Result<Origin, OriginRefused> {
    let scheme = url.scheme();
    if scheme != "http" && scheme != "https" {
        return Err(OriginRefused::Scheme(scheme.to_string()));
    }
    let Some(host) = url.host() else {
        return Err(OriginRefused::NoHost);
    };
    let loopback = match host {
        url::Host::Ipv4(ip) => ip.is_loopback() || ip.is_unspecified(),
        url::Host::Ipv6(ip) => ip.is_loopback() || ip.is_unspecified(),
        url::Host::Domain(d) => d.eq_ignore_ascii_case("localhost"),
    };
    if let Some(port) = url.port_or_known_default() {
        if loopback && own.contains(port) {
            return Err(OriginRefused::OwnListener(port));
        }
    }
    let mut origin = url.clone();
    origin.set_path("/");
    origin.set_query(None);
    origin.set_fragment(None);
    // Both setters fail only on a URL that cannot carry userinfo (no host),
    // which was refused above.
    let _ = origin.set_username("");
    let _ = origin.set_password(None);
    Ok(Origin(origin))
}

/// What a caller asks [`EndpointListeners::ensure`] to serve.
#[derive(Clone, Debug)]
pub struct EndpointSpec {
    /// The slot key the endpoint is recorded under, e.g. `cline:gs:shared:geminiBaseUrl`.
    pub key: String,
    /// The agent whose slot this is. Endpoint traffic is attributed to it.
    pub agent: &'static str,
    /// The wire format the slot's provider speaks, when known. Used as a
    /// format HINT for a request whose route alone resolves to `Unknown`
    /// (a customer gateway with its own path prefix).
    pub family: Option<WireFormat>,
    /// The loopback port, from the relay's endpoint block.
    pub port: u16,
    /// Where the slot's traffic goes.
    pub origin: Origin,
}

/// One live endpoint: identity, where it forwards, and what it has seen.
///
/// Shared by the listener serving it and the wiring pass that owns the slot,
/// so an origin update lands on the next request without a rebind and the
/// traffic counters are readable while it serves.
#[derive(Debug)]
pub struct RelayEndpoint {
    key: String,
    agent: &'static str,
    family: Option<WireFormat>,
    port: u16,
    origin: ArcSwap<Origin>,
    requests: AtomicU64,
    last_request_unix: AtomicU64,
    upstream_failures: AtomicU64,
    contested: std::sync::atomic::AtomicBool,
}

impl RelayEndpoint {
    /// A fresh endpoint with zeroed counters.
    pub fn new(spec: EndpointSpec) -> Self {
        Self {
            key: spec.key,
            agent: spec.agent,
            family: spec.family,
            port: spec.port,
            origin: ArcSwap::from_pointee(spec.origin),
            requests: AtomicU64::new(0),
            last_request_unix: AtomicU64::new(0),
            upstream_failures: AtomicU64::new(0),
            contested: std::sync::atomic::AtomicBool::new(false),
        }
    }

    /// Whether the slot is being reverted often enough that the wiring pass
    /// has slowed its re-applies.
    pub fn contested(&self) -> bool {
        self.contested.load(Ordering::Relaxed)
    }

    /// Set by the wiring pass.
    pub fn set_contested(&self, contested: bool) {
        self.contested.store(contested, Ordering::Relaxed);
    }

    /// The slot key.
    pub fn key(&self) -> &str {
        &self.key
    }

    /// The agent this endpoint's traffic is attributed to.
    pub fn agent(&self) -> &'static str {
        self.agent
    }

    /// The provider family, used as a format hint.
    pub fn family(&self) -> Option<WireFormat> {
        self.family
    }

    /// The loopback port this endpoint serves.
    pub fn port(&self) -> u16 {
        self.port
    }

    /// Where this endpoint forwards, right now.
    pub fn origin(&self) -> Arc<Origin> {
        self.origin.load_full()
    }

    /// Point this endpoint at `origin`, returning whether it changed.
    ///
    /// Lock-free for the request path: a request in flight keeps the origin it
    /// loaded, and the next one sees the new value.
    pub fn set_origin(&self, origin: Origin) -> bool {
        if *self.origin.load_full() == origin {
            return false;
        }
        self.origin.store(Arc::new(origin));
        true
    }

    /// Requests that arrived on this port, OpenLatch's own preflight excluded.
    ///
    /// The proof an agent actually dials this endpoint: a written value only
    /// says where the agent WILL go the next time it reads its config.
    pub fn requests(&self) -> u64 {
        self.requests.load(Ordering::Relaxed)
    }

    /// Unix seconds of the most recent counted request, if there was one.
    pub fn last_request_unix(&self) -> Option<u64> {
        match self.last_request_unix.load(Ordering::Relaxed) {
            0 => None,
            t => Some(t),
        }
    }

    /// Forwards on this endpoint that produced no upstream response.
    ///
    /// Counted here and NOT in the process-wide counter the main port's
    /// wiring supervisor watches: a developer stopping their local Ollama must
    /// not make the relay re-probe Claude Code's plane.
    pub fn upstream_failures(&self) -> u64 {
        self.upstream_failures.load(Ordering::Relaxed)
    }

    pub(crate) fn note_request(&self) {
        self.requests.fetch_add(1, Ordering::Relaxed);
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0);
        self.last_request_unix.store(now, Ordering::Relaxed);
    }

    pub(crate) fn note_upstream_failure(&self) {
        self.upstream_failures.fetch_add(1, Ordering::Relaxed);
    }

    /// The endpoint as the status surface reports it. The origin is already
    /// reduced to scheme, host and port, so it carries no credential and no
    /// query to mask.
    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "key": self.key,
            "agent": self.agent,
            "port": self.port,
            "origin": self.origin().to_string(),
            "family": self.family.map(WireFormat::as_str),
            "requests": self.requests(),
            "last_request_unix": self.last_request_unix(),
            "upstream_failures": self.upstream_failures(),
            "contested": self.contested(),
        })
    }
}

/// Builds the relay state an endpoint listener serves with.
///
/// Supplied by the daemon, which alone holds the shared handles — registry,
/// cloud rail, policy, wiring, egress. Called on every serve attempt, like the
/// main listener's factory, so a restart never inherits a dead connection pool.
pub type StateFactory = Arc<dyn Fn(Arc<RelayEndpoint>) -> ModelRelayState + Send + Sync>;

struct Running {
    endpoint: Arc<RelayEndpoint>,
    stop: watch::Sender<bool>,
    join: JoinHandle<()>,
}

/// The daemon's set of endpoint listeners.
///
/// Held by the daemon outside every supervised factory, for the same reason the
/// wiring state is: a main-relay restart must not drop listeners an editor may
/// still be dialling.
pub struct EndpointListeners {
    make: StateFactory,
    running: std::sync::Mutex<BTreeMap<String, Running>>,
}

impl EndpointListeners {
    /// An empty set whose listeners serve state built by `make`.
    pub fn new(make: StateFactory) -> Self {
        Self {
            make,
            running: std::sync::Mutex::new(BTreeMap::new()),
        }
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, BTreeMap<String, Running>> {
        self.running.lock().unwrap_or_else(|e| e.into_inner())
    }

    /// Serve `spec`, returning the live endpoint.
    ///
    /// Idempotent. A slot already served on the same port keeps its listener
    /// and its counters; only its origin is updated. A slot moving to another
    /// port is released and bound afresh.
    ///
    /// A bind failure is returned, never fatal: the main listener is the one
    /// the daemon cannot run without, and one provider slot's port being taken
    /// must not cost every other agent its relay.
    pub async fn ensure(&self, spec: EndpointSpec) -> Result<Arc<RelayEndpoint>, OlError> {
        let moved = {
            let running = self.lock();
            match running.get(&spec.key) {
                Some(r) if r.endpoint.port() == spec.port && !r.join.is_finished() => {
                    r.endpoint.set_origin(spec.origin.clone());
                    return Ok(r.endpoint.clone());
                }
                Some(_) => true,
                None => false,
            }
        };
        if moved {
            self.release(&spec.key).await;
        }

        let listener = super::bind_pinned(spec.port).await.map_err(|e| {
            OlError::new(
                ERR_MODEL_RELAY_ENDPOINT_PORTS,
                format!(
                    "model relay endpoint {} could not bind loopback port {}: {}",
                    spec.key, spec.port, e.message
                ),
            )
        })?;

        let endpoint = Arc::new(RelayEndpoint::new(spec));
        let (stop, stop_rx) = watch::channel(false);
        let pre_bound = Arc::new(tokio::sync::Mutex::new(Some(listener)));
        let make = self.make.clone();
        let served = endpoint.clone();
        let serve_stop = stop_rx.clone();
        let join = supervise(
            TaskSpec::new(ENDPOINT_TASK, RestartPolicy::Always),
            Arc::new(TaskHealth::new(ENDPOINT_TASK, RestartPolicy::Always)),
            stop_rx,
            move || {
                let state = Arc::new(make(served.clone()).with_endpoint(served.clone()));
                super::serve_attempt(pre_bound.clone(), state, serve_stop.clone())
            },
        );
        tracing::info!(
            key = endpoint.key(),
            port = endpoint.port(),
            origin = %endpoint.origin(),
            "model relay endpoint serving (loopback only)"
        );
        self.lock().insert(
            endpoint.key().to_string(),
            Running {
                endpoint: endpoint.clone(),
                stop,
                join,
            },
        );
        Ok(endpoint)
    }

    /// Stop serving `key`, if it is served.
    pub async fn release(&self, key: &str) {
        let Some(running) = self.lock().remove(key) else {
            return;
        };
        stop_and_join(running).await;
    }

    /// Stop every endpoint listener, bounded per listener.
    pub async fn shutdown_all(&self) {
        let all: Vec<Running> = std::mem::take(&mut *self.lock()).into_values().collect();
        for running in all {
            stop_and_join(running).await;
        }
    }

    /// The live endpoint for `key`.
    pub fn get(&self, key: &str) -> Option<Arc<RelayEndpoint>> {
        self.lock().get(key).map(|r| r.endpoint.clone())
    }

    /// Every live endpoint, ordered by key.
    pub fn snapshot(&self) -> Vec<Arc<RelayEndpoint>> {
        self.lock().values().map(|r| r.endpoint.clone()).collect()
    }
}

impl Drop for EndpointListeners {
    fn drop(&mut self) {
        for running in self.lock().values() {
            let _ = running.stop.send(true);
        }
    }
}

async fn stop_and_join(mut running: Running) {
    let _ = running.stop.send(true);
    if tokio::time::timeout(RELEASE_DRAIN, &mut running.join)
        .await
        .is_err()
    {
        tracing::warn!(
            key = running.endpoint.key(),
            port = running.endpoint.port(),
            "model relay endpoint did not stop within the drain window — aborting it"
        );
        running.join.abort();
    }
}

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

    fn ports(main: u16) -> RelayPorts {
        RelayPorts { daemon: 7443, main }
    }

    fn url(s: &str) -> reqwest::Url {
        reqwest::Url::parse(s).expect("test url parses")
    }

    #[test]
    fn endpoint_port_block_follows_the_main_port() {
        assert_eq!(endpoint_port_block(7600), 7601..=7632);
        assert_eq!(endpoint_port_block(17601), 17602..=17633);
        let ports = ports(7600);
        assert!(ports.in_block(7601) && ports.in_block(7632));
        assert!(!ports.in_block(7600) && !ports.in_block(7633));
    }

    #[test]
    fn a_main_port_at_the_top_of_the_range_yields_no_block() {
        assert!(endpoint_port_block(u16::MAX - 10).is_empty());
        assert!(endpoint_port_block(u16::MAX).is_empty());
        assert!(!endpoint_port_block(u16::MAX - ENDPOINT_BLOCK).is_empty());
        // An ephemeral test listener has no block at all.
        assert!(endpoint_port_block(0).is_empty());
        assert!(!ports(0).in_block(1));
    }

    #[test]
    fn origin_is_normalised_to_scheme_host_port() {
        let own = ports(7600);
        let cases = [
            ("https://gw.corp/anthropic/v1?x=1#f", "https://gw.corp/"),
            ("http://127.0.0.1:11434", "http://127.0.0.1:11434/"),
            ("http://localhost:1234/v1", "http://localhost:1234/"),
            (
                "https://generativelanguage.googleapis.com/v1beta",
                "https://generativelanguage.googleapis.com/",
            ),
            ("https://user:pw@gw.corp:8443/x", "https://gw.corp:8443/"),
        ];
        for (input, want) in cases {
            let got = normalize_origin(&url(input), &own).expect(input);
            assert_eq!(got.as_url().as_str(), want, "{input}");
        }
    }

    #[test]
    fn origin_that_is_our_own_listener_is_refused() {
        let own = ports(7600);
        for (input, port) in [
            ("http://127.0.0.1:7600", 7600),
            ("http://127.0.0.1:7601/v1", 7601),
            ("http://localhost:7632", 7632),
            ("http://[::1]:7443", 7443),
            ("http://0.0.0.0:7610", 7610),
        ] {
            assert_eq!(
                normalize_origin(&url(input), &own),
                Err(OriginRefused::OwnListener(port)),
                "{input}"
            );
        }
        // A customer's own loopback server beside ours stays a legal origin.
        assert!(normalize_origin(&url("http://127.0.0.1:11434"), &own).is_ok());
        assert!(normalize_origin(&url("http://127.0.0.1:7633"), &own).is_ok());
        // A remote host on one of our port numbers is not us.
        assert!(normalize_origin(&url("http://gw.corp:7601"), &own).is_ok());
    }

    #[test]
    fn only_http_origins_are_forwardable() {
        let own = ports(7600);
        assert_eq!(
            normalize_origin(&url("ftp://gw.corp/"), &own),
            Err(OriginRefused::Scheme("ftp".into()))
        );
        assert_eq!(
            normalize_origin(&url("file:///tmp/x"), &own),
            Err(OriginRefused::Scheme("file".into()))
        );
    }

    #[test]
    fn an_origin_update_reports_whether_it_changed() {
        let own = ports(7600);
        let ep = RelayEndpoint::new(EndpointSpec {
            key: "test:slot".into(),
            agent: "cline",
            family: None,
            port: 7601,
            origin: normalize_origin(&url("http://127.0.0.1:11434"), &own).expect("origin"),
        });
        let same = normalize_origin(&url("http://127.0.0.1:11434/api"), &own).expect("origin");
        assert!(!ep.set_origin(same), "the same origin is not a change");
        let moved = normalize_origin(&url("https://gw.corp"), &own).expect("origin");
        assert!(ep.set_origin(moved.clone()));
        assert_eq!(*ep.origin(), moved);
    }

    #[test]
    fn counters_start_at_zero_and_record_the_request_time() {
        let own = ports(7600);
        let ep = RelayEndpoint::new(EndpointSpec {
            key: "test:slot".into(),
            agent: "cline",
            family: Some(WireFormat::OllamaNative),
            port: 7601,
            origin: normalize_origin(&url("http://127.0.0.1:11434"), &own).expect("origin"),
        });
        assert_eq!(ep.requests(), 0);
        assert_eq!(ep.last_request_unix(), None);
        ep.note_request();
        assert_eq!(ep.requests(), 1);
        assert!(ep.last_request_unix().is_some());
        let json = ep.to_json();
        assert_eq!(json["family"], "ollama-native");
        assert_eq!(json["origin"], "http://127.0.0.1:11434/");
    }
}