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