1use anyhow::anyhow;
2use axum::headers::HeaderMapExt;
3use axum::{
4 extract::MatchedPath,
5 http::{Request, Response},
6 routing::get,
7 Extension, Router,
8};
9use collab::api::CloudflareIpCountryHeader;
10use collab::llm::{db::LlmDatabase, log_usage_periodically};
11use collab::migrations::run_database_migrations;
12use collab::user_backfiller::spawn_user_backfiller;
13use collab::{api::billing::poll_stripe_events_periodically, llm::LlmState, ServiceMode};
14use collab::{
15 api::fetch_extensions_from_blob_store_periodically, db, env, executor::Executor,
16 rpc::ResultExt, AppState, Config, RateLimiter, Result,
17};
18use db::Database;
19use std::{
20 env::args,
21 net::{SocketAddr, TcpListener},
22 path::Path,
23 sync::Arc,
24 time::Duration,
25};
26#[cfg(unix)]
27use tokio::signal::unix::SignalKind;
28use tower_http::trace::TraceLayer;
29use tracing_subscriber::{
30 filter::EnvFilter, fmt::format::JsonFields, util::SubscriberInitExt, Layer,
31};
32use util::ResultExt as _;
33
34const VERSION: &str = env!("CARGO_PKG_VERSION");
35const REVISION: Option<&'static str> = option_env!("GITHUB_SHA");
36
37#[tokio::main]
38async fn main() -> Result<()> {
39 if let Err(error) = env::load_dotenv() {
40 eprintln!(
41 "error loading .env.toml (this is expected in production): {}",
42 error
43 );
44 }
45
46 let mut args = args().skip(1);
47 match args.next().as_deref() {
48 Some("version") => {
49 println!("collab v{} ({})", VERSION, REVISION.unwrap_or("unknown"));
50 }
51 Some("migrate") => {
52 let config = envy::from_env::<Config>().expect("error loading config");
53 setup_app_database(&config).await?;
54 }
55 Some("seed") => {
56 let config = envy::from_env::<Config>().expect("error loading config");
57 let db_options = db::ConnectOptions::new(config.database_url.clone());
58
59 let mut db = Database::new(db_options, Executor::Production).await?;
60 db.initialize_notification_kinds().await?;
61
62 collab::seed::seed(&config, &db, false).await?;
63
64 if let Some(llm_database_url) = config.llm_database_url.clone() {
65 let db_options = db::ConnectOptions::new(llm_database_url);
66 let mut db = LlmDatabase::new(db_options.clone(), Executor::Production).await?;
67 db.initialize().await?;
68 collab::llm::db::seed_database(&config, &mut db, true).await?;
69 }
70 }
71 Some("serve") => {
72 let mode = match args.next().as_deref() {
73 Some("collab") => ServiceMode::Collab,
74 Some("api") => ServiceMode::Api,
75 Some("llm") => ServiceMode::Llm,
76 Some("all") => ServiceMode::All,
77 _ => {
78 return Err(anyhow!(
79 "usage: collab <version | migrate | seed | serve <api|collab|llm|all>>"
80 ))?;
81 }
82 };
83
84 let config = envy::from_env::<Config>().expect("error loading config");
85 init_tracing(&config);
86 let mut app = Router::new()
87 .route("/", get(handle_root))
88 .route("/healthz", get(handle_liveness_probe))
89 .layer(Extension(mode));
90
91 let listener = TcpListener::bind(&format!("0.0.0.0:{}", config.http_port))
92 .expect("failed to bind TCP listener");
93
94 let mut on_shutdown = None;
95
96 if mode.is_llm() {
97 setup_llm_database(&config).await?;
98
99 let state = LlmState::new(config.clone(), Executor::Production).await?;
100
101 log_usage_periodically(state.clone());
102
103 app = app
104 .merge(collab::llm::routes())
105 .layer(Extension(state.clone()));
106 }
107
108 if mode.is_collab() || mode.is_api() {
109 setup_app_database(&config).await?;
110
111 let state = AppState::new(config, Executor::Production).await?;
112
113 if mode.is_collab() {
114 state.db.purge_old_embeddings().await.trace_err();
115 RateLimiter::save_periodically(
116 state.rate_limiter.clone(),
117 state.executor.clone(),
118 );
119
120 let epoch = state
121 .db
122 .create_server(&state.config.zed_environment)
123 .await?;
124 let rpc_server = collab::rpc::Server::new(epoch, state.clone());
125 rpc_server.start().await?;
126
127 app = app
128 .merge(collab::api::routes(rpc_server.clone()))
129 .merge(collab::rpc::routes(rpc_server.clone()));
130
131 on_shutdown = Some(Box::new(move || rpc_server.teardown()));
132 }
133
134 if mode.is_api() {
135 poll_stripe_events_periodically(state.clone());
136 fetch_extensions_from_blob_store_periodically(state.clone());
137 spawn_user_backfiller(state.clone());
138
139 app = app
140 .merge(collab::api::events::router())
141 .merge(collab::api::extensions::router())
142 }
143
144 app = app.layer(Extension(state.clone()));
145 }
146
147 app = app.layer(
148 TraceLayer::new_for_http()
149 .make_span_with(|request: &Request<_>| {
150 let matched_path = request
151 .extensions()
152 .get::<MatchedPath>()
153 .map(MatchedPath::as_str);
154
155 let geoip_country_code = request
156 .headers()
157 .typed_get::<CloudflareIpCountryHeader>()
158 .map(|header| header.to_string());
159
160 tracing::info_span!(
161 "http_request",
162 method = ?request.method(),
163 matched_path,
164 geoip_country_code,
165 user_id = tracing::field::Empty,
166 login = tracing::field::Empty,
167 authn.jti = tracing::field::Empty,
168 is_staff = tracing::field::Empty
169 )
170 })
171 .on_response(
172 |response: &Response<_>, latency: Duration, _: &tracing::Span| {
173 let duration_ms = latency.as_micros() as f64 / 1000.;
174 tracing::info!(
175 duration_ms,
176 status = response.status().as_u16(),
177 "finished processing request"
178 );
179 },
180 ),
181 );
182
183 #[cfg(unix)]
184 let signal = async move {
185 let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
186 .expect("failed to listen for interrupt signal");
187 let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
188 .expect("failed to listen for interrupt signal");
189 let sigterm = sigterm.recv();
190 let sigint = sigint.recv();
191 futures::pin_mut!(sigterm, sigint);
192 futures::future::select(sigterm, sigint).await;
193 };
194
195 #[cfg(windows)]
196 let signal = async move {
197 // todo(windows):
198 // `ctrl_close` does not work well, because tokio's signal handler always returns soon,
199 // but system terminates the application soon after returning CTRL+CLOSE handler.
200 // So we should implement blocking handler to treat CTRL+CLOSE signal.
201 let mut ctrl_break = tokio::signal::windows::ctrl_break()
202 .expect("failed to listen for interrupt signal");
203 let mut ctrl_c = tokio::signal::windows::ctrl_c()
204 .expect("failed to listen for interrupt signal");
205 let ctrl_break = ctrl_break.recv();
206 let ctrl_c = ctrl_c.recv();
207 futures::pin_mut!(ctrl_break, ctrl_c);
208 futures::future::select(ctrl_break, ctrl_c).await;
209 };
210
211 axum::Server::from_tcp(listener)
212 .map_err(|e| anyhow!(e))?
213 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
214 .with_graceful_shutdown(async move {
215 signal.await;
216 tracing::info!("Received interrupt signal");
217
218 if let Some(on_shutdown) = on_shutdown {
219 on_shutdown();
220 }
221 })
222 .await
223 .map_err(|e| anyhow!(e))?;
224 }
225 _ => {
226 Err(anyhow!(
227 "usage: collab <version | migrate | seed | serve <api|collab|llm|all>>"
228 ))?;
229 }
230 }
231 Ok(())
232}
233
234async fn setup_app_database(config: &Config) -> Result<()> {
235 let db_options = db::ConnectOptions::new(config.database_url.clone());
236 let mut db = Database::new(db_options, Executor::Production).await?;
237
238 let migrations_path = config.migrations_path.as_deref().unwrap_or_else(|| {
239 #[cfg(feature = "sqlite")]
240 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations.sqlite");
241 #[cfg(not(feature = "sqlite"))]
242 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
243
244 Path::new(default_migrations)
245 });
246
247 let migrations = run_database_migrations(db.options(), migrations_path).await?;
248 for (migration, duration) in migrations {
249 log::info!(
250 "Migrated {} {} {:?}",
251 migration.version,
252 migration.description,
253 duration
254 );
255 }
256
257 db.initialize_notification_kinds().await?;
258
259 if config.seed_path.is_some() {
260 collab::seed::seed(&config, &db, false).await?;
261 }
262
263 Ok(())
264}
265
266async fn setup_llm_database(config: &Config) -> Result<()> {
267 let database_url = config
268 .llm_database_url
269 .as_ref()
270 .ok_or_else(|| anyhow!("missing LLM_DATABASE_URL"))?;
271
272 let db_options = db::ConnectOptions::new(database_url.clone());
273 let db = LlmDatabase::new(db_options, Executor::Production).await?;
274
275 let migrations_path = config
276 .llm_database_migrations_path
277 .as_deref()
278 .unwrap_or_else(|| {
279 #[cfg(feature = "sqlite")]
280 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations_llm.sqlite");
281 #[cfg(not(feature = "sqlite"))]
282 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations_llm");
283
284 Path::new(default_migrations)
285 });
286
287 let migrations = run_database_migrations(db.options(), migrations_path).await?;
288 for (migration, duration) in migrations {
289 log::info!(
290 "Migrated {} {} {:?}",
291 migration.version,
292 migration.description,
293 duration
294 );
295 }
296
297 Ok(())
298}
299
300async fn handle_root(Extension(mode): Extension<ServiceMode>) -> String {
301 format!("zed:{mode} v{VERSION} ({})", REVISION.unwrap_or("unknown"))
302}
303
304async fn handle_liveness_probe(
305 app_state: Option<Extension<Arc<AppState>>>,
306 llm_state: Option<Extension<Arc<LlmState>>>,
307) -> Result<String> {
308 if let Some(state) = app_state {
309 state.db.get_all_users(0, 1).await?;
310 }
311
312 if let Some(llm_state) = llm_state {
313 llm_state.db.list_providers().await?;
314 }
315
316 Ok("ok".to_string())
317}
318
319pub fn init_tracing(config: &Config) -> Option<()> {
320 use std::str::FromStr;
321 use tracing_subscriber::layer::SubscriberExt;
322
323 let filter = EnvFilter::from_str(config.rust_log.as_deref()?).log_err()?;
324
325 tracing_subscriber::registry()
326 .with(if config.log_json.unwrap_or(false) {
327 Box::new(
328 tracing_subscriber::fmt::layer()
329 .fmt_fields(JsonFields::default())
330 .event_format(
331 tracing_subscriber::fmt::format()
332 .json()
333 .flatten_event(true)
334 .with_span_list(false),
335 )
336 .with_filter(filter),
337 ) as Box<dyn Layer<_> + Send + Sync>
338 } else {
339 Box::new(
340 tracing_subscriber::fmt::layer()
341 .event_format(tracing_subscriber::fmt::format().pretty())
342 .with_filter(filter),
343 )
344 })
345 .init();
346
347 None
348}