Skip to main content

tuitbot_server/
state.rs

1//! Shared application state for the tuitbot server.
2
3use std::collections::HashMap;
4use std::net::IpAddr;
5use std::path::PathBuf;
6use std::sync::Arc;
7use std::time::{Instant, SystemTime};
8
9use tokio::sync::{broadcast, Mutex, RwLock};
10use tokio_util::sync::CancellationToken;
11use tuitbot_core::automation::circuit_breaker::CircuitBreaker;
12use tuitbot_core::automation::Runtime;
13use tuitbot_core::config::{
14    effective_config, Config, ConnectorConfig, ContentSourcesConfig, DeploymentMode,
15};
16use tuitbot_core::content::ContentGenerator;
17use tuitbot_core::llm::factory::create_provider;
18use tuitbot_core::storage::accounts::{self, DEFAULT_ACCOUNT_ID};
19use tuitbot_core::storage::DbPool;
20use tuitbot_core::x_api::auth::TokenManager;
21
22use tuitbot_core::error::XApiError;
23use tuitbot_core::x_api::auth;
24
25use crate::ws::AccountWsEvent;
26
27/// Pending OAuth PKCE state for connector link flows.
28pub struct PendingOAuth {
29    /// The PKCE code verifier needed to complete the token exchange.
30    pub code_verifier: String,
31    /// When this entry was created (for 10-minute expiry).
32    pub created_at: Instant,
33    /// The account ID that initiated this OAuth flow (empty for connectors).
34    pub account_id: String,
35}
36
37/// Shared application state accessible by all route handlers.
38pub struct AppState {
39    /// SQLite connection pool.
40    pub db: DbPool,
41    /// Path to the configuration file.
42    pub config_path: PathBuf,
43    /// Data directory for media storage (parent of config file).
44    pub data_dir: PathBuf,
45    /// Broadcast channel sender for real-time WebSocket events.
46    pub event_tx: broadcast::Sender<AccountWsEvent>,
47    /// Local bearer token for API authentication.
48    pub api_token: String,
49    /// Bcrypt hash of the web login passphrase (None if not configured).
50    pub passphrase_hash: RwLock<Option<String>>,
51    /// Last-observed mtime of the `passphrase_hash` file (for detecting out-of-band resets).
52    pub passphrase_hash_mtime: RwLock<Option<SystemTime>>,
53    /// Host address the server is bound to.
54    pub bind_host: String,
55    /// Port the server is listening on.
56    pub bind_port: u16,
57    /// Per-IP login attempt tracking for rate limiting: (count, window_start).
58    pub login_attempts: Mutex<HashMap<IpAddr, (u32, Instant)>>,
59    /// Per-account automation runtimes (keyed by account_id).
60    pub runtimes: Mutex<HashMap<String, Runtime>>,
61    /// Per-account content generators for AI assist endpoints.
62    pub content_generators: Mutex<HashMap<String, Arc<ContentGenerator>>>,
63    /// Optional circuit breaker for X API rate-limit protection.
64    pub circuit_breaker: Option<Arc<CircuitBreaker>>,
65    /// Cancellation token for the Watchtower filesystem watcher (None if not running).
66    pub watchtower_cancel: Option<CancellationToken>,
67    /// Content sources configuration for the Watchtower.
68    pub content_sources: ContentSourcesConfig,
69    /// Connector configuration for remote source OAuth flows.
70    pub connector_config: ConnectorConfig,
71    /// Deployment mode (desktop, self_host, or cloud).
72    pub deployment_mode: DeploymentMode,
73    /// Pending OAuth PKCE challenges keyed by state parameter.
74    pub pending_oauth: Mutex<HashMap<String, PendingOAuth>>,
75    /// Per-account X API token managers for automatic token refresh.
76    pub token_managers: Mutex<HashMap<String, Arc<TokenManager>>>,
77    /// X API client ID from config (needed to create token managers).
78    pub x_client_id: String,
79}
80
81impl AppState {
82    /// Get a fresh X API access token for the given account.
83    ///
84    /// Lazily creates a `TokenManager` on first use (loading tokens from disk),
85    /// then returns a token that is automatically refreshed before expiry.
86    pub async fn get_x_access_token(
87        &self,
88        token_path: &std::path::Path,
89        account_id: &str,
90    ) -> Result<String, XApiError> {
91        // Fast path: token manager already exists.
92        {
93            let managers = self.token_managers.lock().await;
94            if let Some(tm) = managers.get(account_id) {
95                return tm.get_access_token().await;
96            }
97        }
98
99        // Load tokens from disk and create a new manager.
100        let tokens = auth::load_tokens(token_path)?.ok_or(XApiError::AuthExpired)?;
101
102        let tm = Arc::new(TokenManager::new(
103            tokens,
104            self.x_client_id.clone(),
105            token_path.to_path_buf(),
106        ));
107
108        let access_token = tm.get_access_token().await?;
109
110        self.token_managers
111            .lock()
112            .await
113            .insert(account_id.to_string(), tm);
114
115        Ok(access_token)
116    }
117
118    /// Load the effective config for a given account.
119    ///
120    /// Default account: reads config.toml directly (backward compat).
121    /// Non-default: merges config.toml base with account's `config_overrides` from DB.
122    pub async fn load_effective_config(&self, account_id: &str) -> Result<Config, String> {
123        let contents = std::fs::read_to_string(&self.config_path).unwrap_or_default();
124        let base: Config = toml::from_str(&contents).unwrap_or_default();
125
126        if account_id == DEFAULT_ACCOUNT_ID {
127            return Ok(base);
128        }
129
130        let account = accounts::get_account(&self.db, account_id)
131            .await
132            .map_err(|e| e.to_string())?
133            .ok_or_else(|| format!("account not found: {account_id}"))?;
134
135        effective_config(&base, &account.config_overrides)
136            .map(|r| r.config)
137            .map_err(|e| e.to_string())
138    }
139
140    /// Lazily create or return a cached `ContentGenerator` for the given account.
141    ///
142    /// Loads effective config, creates the LLM provider, and caches the generator.
143    pub async fn get_or_create_content_generator(
144        &self,
145        account_id: &str,
146    ) -> Result<Arc<ContentGenerator>, String> {
147        // Fast path: already cached.
148        {
149            let generators = self.content_generators.lock().await;
150            if let Some(gen) = generators.get(account_id) {
151                return Ok(gen.clone());
152            }
153        }
154
155        let config = self.load_effective_config(account_id).await?;
156
157        let provider =
158            create_provider(&config.llm).map_err(|e| format!("LLM not configured: {e}"))?;
159
160        let gen = Arc::new(ContentGenerator::new(provider, config.business));
161
162        self.content_generators
163            .lock()
164            .await
165            .insert(account_id.to_string(), gen.clone());
166
167        Ok(gen)
168    }
169}