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