lib.rs

  1pub mod api;
  2pub mod auth;
  3pub mod db;
  4pub mod env;
  5pub mod executor;
  6pub mod rpc;
  7pub mod seed;
  8
  9#[cfg(test)]
 10mod tests;
 11
 12use anyhow::Context as _;
 13use aws_config::{BehaviorVersion, Region};
 14use axum::{
 15    http::{HeaderMap, StatusCode},
 16    response::IntoResponse,
 17};
 18use db::{ChannelId, Database};
 19use executor::Executor;
 20use serde::Deserialize;
 21use std::{path::PathBuf, sync::Arc};
 22use util::ResultExt;
 23
 24pub type Result<T, E = Error> = std::result::Result<T, E>;
 25
 26pub enum Error {
 27    Http(StatusCode, String, HeaderMap),
 28    Database(sea_orm::error::DbErr),
 29    Internal(anyhow::Error),
 30}
 31
 32impl From<anyhow::Error> for Error {
 33    fn from(error: anyhow::Error) -> Self {
 34        Self::Internal(error)
 35    }
 36}
 37
 38impl From<sea_orm::error::DbErr> for Error {
 39    fn from(error: sea_orm::error::DbErr) -> Self {
 40        Self::Database(error)
 41    }
 42}
 43
 44impl From<axum::Error> for Error {
 45    fn from(error: axum::Error) -> Self {
 46        Self::Internal(error.into())
 47    }
 48}
 49
 50impl From<axum::http::Error> for Error {
 51    fn from(error: axum::http::Error) -> Self {
 52        Self::Internal(error.into())
 53    }
 54}
 55
 56impl From<serde_json::Error> for Error {
 57    fn from(error: serde_json::Error) -> Self {
 58        Self::Internal(error.into())
 59    }
 60}
 61
 62impl Error {
 63    fn http(code: StatusCode, message: String) -> Self {
 64        Self::Http(code, message, HeaderMap::default())
 65    }
 66}
 67
 68impl IntoResponse for Error {
 69    fn into_response(self) -> axum::response::Response {
 70        match self {
 71            Error::Http(code, message, headers) => {
 72                log::error!("HTTP error {}: {}", code, &message);
 73                (code, headers, message).into_response()
 74            }
 75            Error::Database(error) => {
 76                log::error!(
 77                    "HTTP error {}: {:?}",
 78                    StatusCode::INTERNAL_SERVER_ERROR,
 79                    &error
 80                );
 81                (StatusCode::INTERNAL_SERVER_ERROR, format!("{}", &error)).into_response()
 82            }
 83            Error::Internal(error) => {
 84                log::error!(
 85                    "HTTP error {}: {:?}",
 86                    StatusCode::INTERNAL_SERVER_ERROR,
 87                    &error
 88                );
 89                (StatusCode::INTERNAL_SERVER_ERROR, format!("{}", &error)).into_response()
 90            }
 91        }
 92    }
 93}
 94
 95impl std::fmt::Debug for Error {
 96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 97        match self {
 98            Error::Http(code, message, _headers) => (code, message).fmt(f),
 99            Error::Database(error) => error.fmt(f),
100            Error::Internal(error) => error.fmt(f),
101        }
102    }
103}
104
105impl std::fmt::Display for Error {
106    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        match self {
108            Error::Http(code, message, _) => write!(f, "{code}: {message}"),
109            Error::Database(error) => error.fmt(f),
110            Error::Internal(error) => error.fmt(f),
111        }
112    }
113}
114
115impl std::error::Error for Error {}
116
117#[derive(Clone, Deserialize)]
118pub struct Config {
119    pub http_port: u16,
120    pub database_url: String,
121    pub migrations_path: Option<PathBuf>,
122    pub seed_path: Option<PathBuf>,
123    pub database_max_connections: u32,
124    pub api_token: String,
125    pub invite_link_prefix: String,
126    pub livekit_server: Option<String>,
127    pub livekit_key: Option<String>,
128    pub livekit_secret: Option<String>,
129    pub llm_database_url: Option<String>,
130    pub llm_database_max_connections: Option<u32>,
131    pub llm_database_migrations_path: Option<PathBuf>,
132    pub llm_api_secret: Option<String>,
133    pub rust_log: Option<String>,
134    pub log_json: Option<bool>,
135    pub blob_store_url: Option<String>,
136    pub blob_store_region: Option<String>,
137    pub blob_store_access_key: Option<String>,
138    pub blob_store_secret_key: Option<String>,
139    pub blob_store_bucket: Option<String>,
140    pub kinesis_region: Option<String>,
141    pub kinesis_stream: Option<String>,
142    pub kinesis_access_key: Option<String>,
143    pub kinesis_secret_key: Option<String>,
144    pub zed_environment: Arc<str>,
145    pub openai_api_key: Option<Arc<str>>,
146    pub google_ai_api_key: Option<Arc<str>>,
147    pub anthropic_api_key: Option<Arc<str>>,
148    pub anthropic_staff_api_key: Option<Arc<str>>,
149    pub llm_closed_beta_model_name: Option<Arc<str>>,
150    pub prediction_api_url: Option<Arc<str>>,
151    pub prediction_api_key: Option<Arc<str>>,
152    pub prediction_model: Option<Arc<str>>,
153    pub zed_client_checksum_seed: Option<String>,
154    pub auto_join_channel_id: Option<ChannelId>,
155    pub supermaven_admin_api_key: Option<Arc<str>>,
156}
157
158impl Config {
159    pub fn is_development(&self) -> bool {
160        self.zed_environment == "development".into()
161    }
162
163    /// Returns the base `zed.dev` URL.
164    pub fn zed_dot_dev_url(&self) -> &str {
165        match self.zed_environment.as_ref() {
166            "development" => "http://localhost:3000",
167            "staging" => "https://staging.zed.dev",
168            _ => "https://zed.dev",
169        }
170    }
171
172    #[cfg(test)]
173    pub fn test() -> Self {
174        Self {
175            http_port: 0,
176            database_url: "".into(),
177            database_max_connections: 0,
178            api_token: "".into(),
179            invite_link_prefix: "".into(),
180            livekit_server: None,
181            livekit_key: None,
182            livekit_secret: None,
183            llm_database_url: None,
184            llm_database_max_connections: None,
185            llm_database_migrations_path: None,
186            llm_api_secret: None,
187            rust_log: None,
188            log_json: None,
189            zed_environment: "test".into(),
190            blob_store_url: None,
191            blob_store_region: None,
192            blob_store_access_key: None,
193            blob_store_secret_key: None,
194            blob_store_bucket: None,
195            openai_api_key: None,
196            google_ai_api_key: None,
197            anthropic_api_key: None,
198            anthropic_staff_api_key: None,
199            llm_closed_beta_model_name: None,
200            prediction_api_url: None,
201            prediction_api_key: None,
202            prediction_model: None,
203            zed_client_checksum_seed: None,
204            auto_join_channel_id: None,
205            migrations_path: None,
206            seed_path: None,
207            supermaven_admin_api_key: None,
208            kinesis_region: None,
209            kinesis_access_key: None,
210            kinesis_secret_key: None,
211            kinesis_stream: None,
212        }
213    }
214}
215
216/// The service mode that collab should run in.
217#[derive(Debug, PartialEq, Eq, Clone, Copy, strum::Display)]
218#[strum(serialize_all = "snake_case")]
219pub enum ServiceMode {
220    Api,
221    Collab,
222    All,
223}
224
225impl ServiceMode {
226    pub fn is_collab(&self) -> bool {
227        matches!(self, Self::Collab | Self::All)
228    }
229
230    pub fn is_api(&self) -> bool {
231        matches!(self, Self::Api | Self::All)
232    }
233}
234
235pub struct AppState {
236    pub db: Arc<Database>,
237    pub livekit_client: Option<Arc<dyn livekit_api::Client>>,
238    pub blob_store_client: Option<aws_sdk_s3::Client>,
239    pub executor: Executor,
240    pub kinesis_client: Option<::aws_sdk_kinesis::Client>,
241    pub config: Config,
242}
243
244impl AppState {
245    pub async fn new(config: Config, executor: Executor) -> Result<Arc<Self>> {
246        let mut db_options = db::ConnectOptions::new(config.database_url.clone());
247        db_options.max_connections(config.database_max_connections);
248        let mut db = Database::new(db_options).await?;
249        db.initialize_notification_kinds().await?;
250
251        let livekit_client = if let Some(((server, key), secret)) = config
252            .livekit_server
253            .as_ref()
254            .zip(config.livekit_key.as_ref())
255            .zip(config.livekit_secret.as_ref())
256        {
257            Some(Arc::new(livekit_api::LiveKitClient::new(
258                server.clone(),
259                key.clone(),
260                secret.clone(),
261            )) as Arc<dyn livekit_api::Client>)
262        } else {
263            None
264        };
265
266        let db = Arc::new(db);
267        let this = Self {
268            db: db.clone(),
269            livekit_client,
270            blob_store_client: build_blob_store_client(&config).await.log_err(),
271            executor,
272            kinesis_client: if config.kinesis_access_key.is_some() {
273                build_kinesis_client(&config).await.log_err()
274            } else {
275                None
276            },
277            config,
278        };
279        Ok(Arc::new(this))
280    }
281}
282
283async fn build_blob_store_client(config: &Config) -> anyhow::Result<aws_sdk_s3::Client> {
284    let keys = aws_sdk_s3::config::Credentials::new(
285        config
286            .blob_store_access_key
287            .clone()
288            .context("missing blob_store_access_key")?,
289        config
290            .blob_store_secret_key
291            .clone()
292            .context("missing blob_store_secret_key")?,
293        None,
294        None,
295        "env",
296    );
297
298    let s3_config = aws_config::defaults(BehaviorVersion::latest())
299        .endpoint_url(
300            config
301                .blob_store_url
302                .as_ref()
303                .context("missing blob_store_url")?,
304        )
305        .region(Region::new(
306            config
307                .blob_store_region
308                .clone()
309                .context("missing blob_store_region")?,
310        ))
311        .credentials_provider(keys)
312        .load()
313        .await;
314
315    Ok(aws_sdk_s3::Client::new(&s3_config))
316}
317
318async fn build_kinesis_client(config: &Config) -> anyhow::Result<aws_sdk_kinesis::Client> {
319    let keys = aws_sdk_s3::config::Credentials::new(
320        config
321            .kinesis_access_key
322            .clone()
323            .context("missing kinesis_access_key")?,
324        config
325            .kinesis_secret_key
326            .clone()
327            .context("missing kinesis_secret_key")?,
328        None,
329        None,
330        "env",
331    );
332
333    let kinesis_config = aws_config::defaults(BehaviorVersion::latest())
334        .region(Region::new(
335            config
336                .kinesis_region
337                .clone()
338                .context("missing kinesis_region")?,
339        ))
340        .credentials_provider(keys)
341        .load()
342        .await;
343
344    Ok(aws_sdk_kinesis::Client::new(&kinesis_config))
345}