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