tail-fin-daemon 0.7.4

//! `SessionDriver` impl for tail-fin daemon sessions.
//!
//! Each session is tied to one **site** (one browser tab) and one **Chrome host**.
//! `TailFinSession` owns a `Box<dyn SiteHandler>` created at launch time, so
//! stateful site adapters (e.g. Grok) can keep per-session client state across
//! multiple commands without rebuilding from scratch.

use std::collections::HashMap;
use std::sync::{Arc, OnceLock};

use anyhow::{anyhow, Context};
use night_fury_core::{BrowserSession, SessionBuilder};
use night_fury_daemon_core::{protocol::Response, SessionDriver};
use serde_json::Value;
use tail_fin_cli_core::browser_session;
use tail_fin_core::Site;

use crate::handlers::SiteHandler;

/// Bundles a `Site` (lifecycle) with a handler factory (command execution).
/// Register one per site via [`register_sites`].
pub struct SiteEntry {
    pub site: Arc<dyn Site>,
    make_handler: Arc<dyn Fn(BrowserSession) -> Box<dyn SiteHandler> + Send + Sync>,
}

impl SiteEntry {
    /// Register a site with a custom stateful handler factory.
    pub fn new<MH>(site: impl Site + 'static, make_handler: MH) -> Self
    where
        MH: Fn(BrowserSession) -> Box<dyn SiteHandler> + Send + Sync + 'static,
    {
        Self {
            site: Arc::new(site),
            make_handler: Arc::new(make_handler),
        }
    }

    /// Register a site that has no command handlers yet (only session
    /// lifecycle). Commands will return a clear "not implemented" error.
    pub fn stateless(site: impl Site + 'static) -> Self {
        let site_id = site.id();
        Self::new(site, move |_session| {
            Box::new(crate::handlers::NoopHandler::new(site_id))
        })
    }

    fn create_handler(&self, session: BrowserSession) -> Box<dyn SiteHandler> {
        (self.make_handler)(session)
    }
}

/// Global registry of known sites, keyed by `Site::id()`.
static SITE_REGISTRY: OnceLock<HashMap<&'static str, SiteEntry>> = OnceLock::new();

/// Register all sites this daemon instance can serve. Call once at startup
/// before `serve()`. Subsequent calls are no-ops.
pub fn register_sites(entries: Vec<SiteEntry>) {
    let map: HashMap<&'static str, SiteEntry> =
        entries.into_iter().map(|e| (e.site.id(), e)).collect();
    let _ = SITE_REGISTRY.set(map);
}

/// Look up a registered site by its identifier (e.g. `"twitter"`, `"sa"`).
pub fn lookup_site(id: &str) -> Option<Arc<dyn Site>> {
    SITE_REGISTRY.get()?.get(id).map(|e| e.site.clone())
}

/// Return all registered site ids, for CLI --help / error messages.
pub fn registered_site_ids() -> Vec<&'static str> {
    SITE_REGISTRY
        .get()
        .map(|m| m.keys().copied().collect())
        .unwrap_or_default()
}

/// One browser tab bound to a single site + host.
pub struct TailFinSession {
    pub site: Arc<dyn Site>,
    pub host: String,
    handler: Box<dyn SiteHandler>,
}

impl SessionDriver for TailFinSession {
    async fn launch(params: Value) -> anyhow::Result<Self> {
        let site_id = params
            .get("site")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow!("launch: missing 'site'"))?;

        // Warn in debug builds if register_sites() was never called — this is
        // almost always a startup wiring bug, not a missing site.
        debug_assert!(
            SITE_REGISTRY.get().is_some(),
            "SITE_REGISTRY is uninitialised: call register_sites() before serve()"
        );
        // Treat an uninitialized registry the same as an empty one: the site
        // is simply not registered, so return the canonical "unknown site" error.
        let empty_map = HashMap::new();
        let registry = SITE_REGISTRY.get().unwrap_or(&empty_map);

        let entry = registry.get(site_id).ok_or_else(|| {
            let known = registered_site_ids().join(", ");
            if known.is_empty() {
                anyhow!("unknown site '{site_id}' (no sites registered)")
            } else {
                anyhow!("unknown site '{site_id}' (registered: {known})")
            }
        })?;

        let host = params
            .get("host")
            .and_then(|v| v.as_str())
            .unwrap_or(crate::cli::DEFAULT_CDP_HOST)
            .to_string();

        if host == "auto" {
            return Err(anyhow!(
                "host='auto' (stealth launch) is not supported in MVP; use --connect"
            ));
        }

        let session = if site_id == "grok" {
            // Attach to an already-open Grok tab so Cloudflare/session state
            // from the user's cleared tab is reused. Opening a fresh tab can
            // lose challenge clearance and break statsig harvest.
            SessionBuilder::default()
                .connect_to(format!("ws://{host}"))
                .attach_to_existing_url_containing("grok.com")
                .headed(false)
                .build()
                .await
                .with_context(|| format!("connect+attach to grok tab at {host}"))?
        } else {
            browser_session(&host, false)
                .await
                .with_context(|| format!("connect to chrome at {host}"))?
        };

        let handler = entry.create_handler(session);

        Ok(Self {
            site: entry.site.clone(),
            host,
            handler,
        })
    }

    async fn handle(&self, id: &str, cmd: &str, params: &Value) -> (Response, bool) {
        let resp = self.handler.handle(id, cmd, params).await;
        (resp, false)
    }
}

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

    struct FakeSite;

    #[async_trait::async_trait]
    impl Site for FakeSite {
        fn id(&self) -> &'static str {
            "fake"
        }
        fn display_name(&self) -> &'static str {
            "Fake"
        }
        fn cookie_domain_patterns(&self) -> &'static [&'static str] {
            &["*.fake.test"]
        }
        fn refresh_url(&self) -> &'static str {
            "https://fake.test/"
        }
        async fn validate(
            &self,
            _session: &BrowserSession,
        ) -> Result<tail_fin_core::SessionStatus, tail_fin_core::SiteError> {
            Ok(tail_fin_core::SessionStatus::Valid)
        }
    }

    #[test]
    fn register_and_lookup() {
        register_sites(vec![SiteEntry::stateless(FakeSite)]);
        let s = lookup_site("fake").expect("fake site registered");
        assert_eq!(s.id(), "fake");
        assert!(registered_site_ids().contains(&"fake"));
        assert!(lookup_site("definitely-not-a-real-site").is_none());
    }

    #[test]
    fn site_entry_stateless_holds_correct_site_id() {
        let entry = SiteEntry::stateless(FakeSite);
        assert_eq!(entry.site.id(), "fake");
    }
}