voidcrawl_mcp/
state.rs

1//! Shared application state carried inside the MCP server.
2
3use std::{
4    env, fmt,
5    path::PathBuf,
6    str::FromStr,
7    sync::{Arc, Mutex as StdMutex},
8};
9
10use tokio::sync::OnceCell;
11use void_crawl_core::{
12    BrowserPool, BrowserSession, PoolConfig, ProfileHandle, Result, VoidCrawlError,
13};
14
15use crate::sessions::SessionRegistry;
16
17/// Pre-acquired profile that pins the whole server.
18///
19/// Built by `main.rs` when the user passes `--profile NAME` or sets
20/// `VOIDCRAWL_PROFILE`. The owned `BrowserSession` is moved into the
21/// pool on first `AppState::pool()` call so every `fetch` /
22/// `fetch_many` / `screenshot` tool call inherits the profile's
23/// cookies. The `ProfileHandle` is kept alive (inside the `AppState`)
24/// so its `fs2` advisory lock stays held for the server lifetime.
25pub struct PinnedProfile {
26    pub handle:         ProfileHandle,
27    /// Extracted once, at AppState::pool() first-call time.
28    pub session:        StdMutex<Option<BrowserSession>>,
29    /// Mirrors the profile's Chrome-facing name for user_data_dir
30    /// defaulting in `session_open`.
31    pub name:           String,
32    pub user_data_root: PathBuf,
33}
34
35impl fmt::Debug for PinnedProfile {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        f.debug_struct("PinnedProfile")
38            .field("name", &self.name)
39            .field("user_data_root", &self.user_data_root)
40            .finish_non_exhaustive()
41    }
42}
43
44/// Bundle of shared state passed into the `VoidCrawlServer`. Cheap to
45/// clone (two `Arc`s).
46///
47/// The `BrowserPool` is lazy: Chrome is not launched until the first tool
48/// call that needs it. This keeps the MCP `initialize` handshake fast so the
49/// harness can register tools without timing out on browser startup.
50#[derive(Debug, Clone)]
51pub struct AppState {
52    pool:         Arc<OnceCell<Arc<BrowserPool>>>,
53    pub sessions: Arc<SessionRegistry>,
54    pub pinned:   Option<Arc<PinnedProfile>>,
55}
56
57impl AppState {
58    pub fn new(sessions: Arc<SessionRegistry>) -> Self {
59        Self { pool: Arc::new(OnceCell::new()), sessions, pinned: None }
60    }
61
62    /// Construct with a pre-acquired profile that will back the pool.
63    pub fn with_pinned_profile(sessions: Arc<SessionRegistry>, pinned: PinnedProfile) -> Self {
64        Self { pool: Arc::new(OnceCell::new()), sessions, pinned: Some(Arc::new(pinned)) }
65    }
66
67    /// Construct with a pre-built pool already initialized. Useful for tests
68    /// that want to inject a specific `BrowserPool` instead of letting the
69    /// state launch Chrome itself.
70    pub fn with_pool(pool: Arc<BrowserPool>, sessions: Arc<SessionRegistry>) -> Self {
71        let cell = OnceCell::new_with(Some(pool));
72        Self { pool: Arc::new(cell), sessions, pinned: None }
73    }
74
75    /// Get the pool, launching Chrome on first call. Subsequent calls reuse
76    /// the same pool instance.
77    ///
78    /// When a profile is pinned, the pool is built from the pinned
79    /// session — `browsers=1` (Chrome's SingletonLock prevents a second
80    /// process on the same user_data_dir), `tabs_per_browser` still
81    /// gives real concurrency. Without a pinned profile, the old
82    /// `BrowserPool::from_env()` path runs unchanged.
83    pub async fn pool(&self) -> Result<Arc<BrowserPool>> {
84        let pinned = self.pinned.clone();
85        self.pool
86            .get_or_try_init(|| async move {
87                let pool = if let Some(p) = pinned {
88                    // Take the session out of the PinnedProfile. The
89                    // handle (and its lock guard) stays alive inside
90                    // `self.pinned` for the server lifetime.
91                    let session = p
92                        .session
93                        .lock()
94                        .map_err(|_| {
95                            VoidCrawlError::Other("PinnedProfile session mutex poisoned".into())
96                        })?
97                        .take()
98                        .ok_or_else(|| {
99                            VoidCrawlError::Other("PinnedProfile session already consumed".into())
100                        })?;
101                    Arc::new(BrowserPool::new(pool_config_from_env(), vec![session]))
102                } else {
103                    Arc::new(BrowserPool::from_env().await?)
104                };
105                Arc::clone(&pool).start_eviction_task();
106                Ok(pool)
107            })
108            .await
109            .map(Arc::clone)
110    }
111
112    /// Returns the initialized pool, if any. Used on shutdown to close without
113    /// forcing a launch.
114    pub fn pool_if_initialized(&self) -> Option<Arc<BrowserPool>> {
115        self.pool.get().map(Arc::clone)
116    }
117}
118
119/// Read the subset of `PoolConfig` values from env that apply to a
120/// pinned-profile pool. Matches `BrowserPool::from_env` for these
121/// knobs but ignores `BROWSER_COUNT` / `CHROME_WS_URLS` (single Chrome,
122/// no attach).
123fn pool_config_from_env() -> PoolConfig {
124    fn parse<T: FromStr>(key: &str, default: T) -> T {
125        env::var(key).ok().and_then(|s| s.parse().ok()).unwrap_or(default)
126    }
127    PoolConfig {
128        browsers:             1,
129        tabs_per_browser:     parse("TABS_PER_BROWSER", 4),
130        tab_max_uses:         parse("TAB_MAX_USES", 50),
131        tab_max_idle_secs:    parse("TAB_MAX_IDLE_SECS", 60),
132        acquire_timeout_secs: parse("ACQUIRE_TIMEOUT_SECS", 30),
133        auto_evict:           env::var("AUTO_EVICT").map_or(true, |v| v != "0"),
134    }
135}
voidcrawl_mcp/state.rs

voidcrawl_mcp/
state.rs
