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