1mod authorization;
2pub mod db;
3mod telemetry;
4mod token;
5
6use crate::{
7 api::CloudflareIpCountryHeader, build_clickhouse_client, db::UserId, executor::Executor,
8 Config, Error, Result,
9};
10use anyhow::{anyhow, Context as _};
11use authorization::authorize_access_to_language_model;
12use axum::{
13 body::Body,
14 http::{self, HeaderName, HeaderValue, Request, StatusCode},
15 middleware::{self, Next},
16 response::{IntoResponse, Response},
17 routing::post,
18 Extension, Json, Router, TypedHeader,
19};
20use chrono::{DateTime, Duration, Utc};
21use collections::HashMap;
22use db::{usage_measure::UsageMeasure, ActiveUserCount, LlmDatabase};
23use futures::{Stream, StreamExt as _};
24use http_client::IsahcHttpClient;
25use rpc::{
26 proto::Plan, LanguageModelProvider, PerformCompletionParams, EXPIRED_LLM_TOKEN_HEADER_NAME,
27};
28use std::{
29 pin::Pin,
30 sync::Arc,
31 task::{Context, Poll},
32};
33use strum::IntoEnumIterator;
34use telemetry::{report_llm_rate_limit, report_llm_usage, LlmRateLimitEventRow, LlmUsageEventRow};
35use tokio::sync::RwLock;
36use util::ResultExt;
37
38pub use token::*;
39
40pub struct LlmState {
41 pub config: Config,
42 pub executor: Executor,
43 pub db: Arc<LlmDatabase>,
44 pub http_client: IsahcHttpClient,
45 pub clickhouse_client: Option<clickhouse::Client>,
46 active_user_count_by_model:
47 RwLock<HashMap<(LanguageModelProvider, String), (DateTime<Utc>, ActiveUserCount)>>,
48}
49
50const ACTIVE_USER_COUNT_CACHE_DURATION: Duration = Duration::seconds(30);
51
52impl LlmState {
53 pub async fn new(config: Config, executor: Executor) -> Result<Arc<Self>> {
54 let database_url = config
55 .llm_database_url
56 .as_ref()
57 .ok_or_else(|| anyhow!("missing LLM_DATABASE_URL"))?;
58 let max_connections = config
59 .llm_database_max_connections
60 .ok_or_else(|| anyhow!("missing LLM_DATABASE_MAX_CONNECTIONS"))?;
61
62 let mut db_options = db::ConnectOptions::new(database_url);
63 db_options.max_connections(max_connections);
64 let mut db = LlmDatabase::new(db_options, executor.clone()).await?;
65 db.initialize().await?;
66
67 let db = Arc::new(db);
68
69 let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));
70 let http_client = IsahcHttpClient::builder()
71 .default_header("User-Agent", user_agent)
72 .build()
73 .context("failed to construct http client")?;
74
75 let this = Self {
76 executor,
77 db,
78 http_client,
79 clickhouse_client: config
80 .clickhouse_url
81 .as_ref()
82 .and_then(|_| build_clickhouse_client(&config).log_err()),
83 active_user_count_by_model: RwLock::new(HashMap::default()),
84 config,
85 };
86
87 Ok(Arc::new(this))
88 }
89
90 pub async fn get_active_user_count(
91 &self,
92 provider: LanguageModelProvider,
93 model: &str,
94 ) -> Result<ActiveUserCount> {
95 let now = Utc::now();
96
97 {
98 let active_user_count_by_model = self.active_user_count_by_model.read().await;
99 if let Some((last_updated, count)) =
100 active_user_count_by_model.get(&(provider, model.to_string()))
101 {
102 if now - *last_updated < ACTIVE_USER_COUNT_CACHE_DURATION {
103 return Ok(*count);
104 }
105 }
106 }
107
108 let mut cache = self.active_user_count_by_model.write().await;
109 let new_count = self.db.get_active_user_count(provider, model, now).await?;
110 cache.insert((provider, model.to_string()), (now, new_count));
111 Ok(new_count)
112 }
113}
114
115pub fn routes() -> Router<(), Body> {
116 Router::new()
117 .route("/completion", post(perform_completion))
118 .layer(middleware::from_fn(validate_api_token))
119}
120
121async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoResponse {
122 let token = req
123 .headers()
124 .get(http::header::AUTHORIZATION)
125 .and_then(|header| header.to_str().ok())
126 .ok_or_else(|| {
127 Error::http(
128 StatusCode::BAD_REQUEST,
129 "missing authorization header".to_string(),
130 )
131 })?
132 .strip_prefix("Bearer ")
133 .ok_or_else(|| {
134 Error::http(
135 StatusCode::BAD_REQUEST,
136 "invalid authorization header".to_string(),
137 )
138 })?;
139
140 let state = req.extensions().get::<Arc<LlmState>>().unwrap();
141 match LlmTokenClaims::validate(&token, &state.config) {
142 Ok(claims) => {
143 if state.db.is_access_token_revoked(&claims.jti).await? {
144 return Err(Error::http(
145 StatusCode::UNAUTHORIZED,
146 "unauthorized".to_string(),
147 ));
148 }
149
150 tracing::Span::current()
151 .record("user_id", claims.user_id)
152 .record("login", claims.github_user_login.clone())
153 .record("authn.jti", &claims.jti)
154 .record("is_staff", &claims.is_staff);
155
156 req.extensions_mut().insert(claims);
157 Ok::<_, Error>(next.run(req).await.into_response())
158 }
159 Err(ValidateLlmTokenError::Expired) => Err(Error::Http(
160 StatusCode::UNAUTHORIZED,
161 "unauthorized".to_string(),
162 [(
163 HeaderName::from_static(EXPIRED_LLM_TOKEN_HEADER_NAME),
164 HeaderValue::from_static("true"),
165 )]
166 .into_iter()
167 .collect(),
168 )),
169 Err(_err) => Err(Error::http(
170 StatusCode::UNAUTHORIZED,
171 "unauthorized".to_string(),
172 )),
173 }
174}
175
176async fn perform_completion(
177 Extension(state): Extension<Arc<LlmState>>,
178 Extension(claims): Extension<LlmTokenClaims>,
179 country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
180 Json(params): Json<PerformCompletionParams>,
181) -> Result<impl IntoResponse> {
182 let model = normalize_model_name(
183 state.db.model_names_for_provider(params.provider),
184 params.model,
185 );
186
187 authorize_access_to_language_model(
188 &state.config,
189 &claims,
190 country_code_header.map(|header| header.to_string()),
191 params.provider,
192 &model,
193 )?;
194
195 check_usage_limit(&state, params.provider, &model, &claims).await?;
196
197 let stream = match params.provider {
198 LanguageModelProvider::Anthropic => {
199 let api_key = if claims.is_staff {
200 state
201 .config
202 .anthropic_staff_api_key
203 .as_ref()
204 .context("no Anthropic AI staff API key configured on the server")?
205 } else {
206 state
207 .config
208 .anthropic_api_key
209 .as_ref()
210 .context("no Anthropic AI API key configured on the server")?
211 };
212
213 let mut request: anthropic::Request =
214 serde_json::from_str(¶ms.provider_request.get())?;
215
216 // Override the model on the request with the latest version of the model that is
217 // known to the server.
218 //
219 // Right now, we use the version that's defined in `model.id()`, but we will likely
220 // want to change this code once a new version of an Anthropic model is released,
221 // so that users can use the new version, without having to update Zed.
222 request.model = match model.as_str() {
223 "claude-3-5-sonnet" => anthropic::Model::Claude3_5Sonnet.id().to_string(),
224 "claude-3-opus" => anthropic::Model::Claude3Opus.id().to_string(),
225 "claude-3-haiku" => anthropic::Model::Claude3Haiku.id().to_string(),
226 "claude-3-sonnet" => anthropic::Model::Claude3Sonnet.id().to_string(),
227 _ => request.model,
228 };
229
230 let (chunks, rate_limit_info) = anthropic::stream_completion_with_rate_limit_info(
231 &state.http_client,
232 anthropic::ANTHROPIC_API_URL,
233 api_key,
234 request,
235 None,
236 )
237 .await
238 .map_err(|err| match err {
239 anthropic::AnthropicError::ApiError(ref api_error) => match api_error.code() {
240 Some(anthropic::ApiErrorCode::RateLimitError) => Error::http(
241 StatusCode::TOO_MANY_REQUESTS,
242 "Upstream Anthropic rate limit exceeded.".to_string(),
243 ),
244 Some(anthropic::ApiErrorCode::InvalidRequestError) => {
245 Error::http(StatusCode::BAD_REQUEST, api_error.message.clone())
246 }
247 Some(anthropic::ApiErrorCode::OverloadedError) => {
248 Error::http(StatusCode::SERVICE_UNAVAILABLE, api_error.message.clone())
249 }
250 Some(_) => {
251 Error::http(StatusCode::INTERNAL_SERVER_ERROR, api_error.message.clone())
252 }
253 None => Error::Internal(anyhow!(err)),
254 },
255 anthropic::AnthropicError::Other(err) => Error::Internal(err),
256 })?;
257
258 if let Some(rate_limit_info) = rate_limit_info {
259 tracing::info!(
260 target: "upstream rate limit",
261 is_staff = claims.is_staff,
262 provider = params.provider.to_string(),
263 model = model,
264 tokens_remaining = rate_limit_info.tokens_remaining,
265 requests_remaining = rate_limit_info.requests_remaining,
266 requests_reset = ?rate_limit_info.requests_reset,
267 tokens_reset = ?rate_limit_info.tokens_reset,
268 );
269 }
270
271 chunks
272 .map(move |event| {
273 let chunk = event?;
274 let (input_tokens, output_tokens) = match &chunk {
275 anthropic::Event::MessageStart {
276 message: anthropic::Response { usage, .. },
277 }
278 | anthropic::Event::MessageDelta { usage, .. } => (
279 usage.input_tokens.unwrap_or(0) as usize,
280 usage.output_tokens.unwrap_or(0) as usize,
281 ),
282 _ => (0, 0),
283 };
284
285 anyhow::Ok((
286 serde_json::to_vec(&chunk).unwrap(),
287 input_tokens,
288 output_tokens,
289 ))
290 })
291 .boxed()
292 }
293 LanguageModelProvider::OpenAi => {
294 let api_key = state
295 .config
296 .openai_api_key
297 .as_ref()
298 .context("no OpenAI API key configured on the server")?;
299 let chunks = open_ai::stream_completion(
300 &state.http_client,
301 open_ai::OPEN_AI_API_URL,
302 api_key,
303 serde_json::from_str(¶ms.provider_request.get())?,
304 None,
305 )
306 .await?;
307
308 chunks
309 .map(|event| {
310 event.map(|chunk| {
311 let input_tokens =
312 chunk.usage.as_ref().map_or(0, |u| u.prompt_tokens) as usize;
313 let output_tokens =
314 chunk.usage.as_ref().map_or(0, |u| u.completion_tokens) as usize;
315 (
316 serde_json::to_vec(&chunk).unwrap(),
317 input_tokens,
318 output_tokens,
319 )
320 })
321 })
322 .boxed()
323 }
324 LanguageModelProvider::Google => {
325 let api_key = state
326 .config
327 .google_ai_api_key
328 .as_ref()
329 .context("no Google AI API key configured on the server")?;
330 let chunks = google_ai::stream_generate_content(
331 &state.http_client,
332 google_ai::API_URL,
333 api_key,
334 serde_json::from_str(¶ms.provider_request.get())?,
335 )
336 .await?;
337
338 chunks
339 .map(|event| {
340 event.map(|chunk| {
341 // TODO - implement token counting for Google AI
342 let input_tokens = 0;
343 let output_tokens = 0;
344 (
345 serde_json::to_vec(&chunk).unwrap(),
346 input_tokens,
347 output_tokens,
348 )
349 })
350 })
351 .boxed()
352 }
353 LanguageModelProvider::Zed => {
354 let api_key = state
355 .config
356 .qwen2_7b_api_key
357 .as_ref()
358 .context("no Qwen2-7B API key configured on the server")?;
359 let api_url = state
360 .config
361 .qwen2_7b_api_url
362 .as_ref()
363 .context("no Qwen2-7B URL configured on the server")?;
364 let chunks = open_ai::stream_completion(
365 &state.http_client,
366 &api_url,
367 api_key,
368 serde_json::from_str(¶ms.provider_request.get())?,
369 None,
370 )
371 .await?;
372
373 chunks
374 .map(|event| {
375 event.map(|chunk| {
376 let input_tokens =
377 chunk.usage.as_ref().map_or(0, |u| u.prompt_tokens) as usize;
378 let output_tokens =
379 chunk.usage.as_ref().map_or(0, |u| u.completion_tokens) as usize;
380 (
381 serde_json::to_vec(&chunk).unwrap(),
382 input_tokens,
383 output_tokens,
384 )
385 })
386 })
387 .boxed()
388 }
389 };
390
391 Ok(Response::new(Body::wrap_stream(TokenCountingStream {
392 state,
393 claims,
394 provider: params.provider,
395 model,
396 input_tokens: 0,
397 output_tokens: 0,
398 inner_stream: stream,
399 })))
400}
401
402fn normalize_model_name(known_models: Vec<String>, name: String) -> String {
403 if let Some(known_model_name) = known_models
404 .iter()
405 .filter(|known_model_name| name.starts_with(known_model_name.as_str()))
406 .max_by_key(|known_model_name| known_model_name.len())
407 {
408 known_model_name.to_string()
409 } else {
410 name
411 }
412}
413
414/// The maximum lifetime spending an individual user can reach before being cut off.
415///
416/// Represented in cents.
417const LIFETIME_SPENDING_LIMIT_IN_CENTS: usize = 1_000 * 100;
418
419async fn check_usage_limit(
420 state: &Arc<LlmState>,
421 provider: LanguageModelProvider,
422 model_name: &str,
423 claims: &LlmTokenClaims,
424) -> Result<()> {
425 let model = state.db.model(provider, model_name)?;
426 let usage = state
427 .db
428 .get_usage(
429 UserId::from_proto(claims.user_id),
430 provider,
431 model_name,
432 Utc::now(),
433 )
434 .await?;
435
436 if usage.lifetime_spending >= LIFETIME_SPENDING_LIMIT_IN_CENTS {
437 return Err(Error::http(
438 StatusCode::FORBIDDEN,
439 "Maximum spending limit reached.".to_string(),
440 ));
441 }
442
443 let active_users = state.get_active_user_count(provider, model_name).await?;
444
445 let users_in_recent_minutes = active_users.users_in_recent_minutes.max(1);
446 let users_in_recent_days = active_users.users_in_recent_days.max(1);
447
448 let per_user_max_requests_per_minute =
449 model.max_requests_per_minute as usize / users_in_recent_minutes;
450 let per_user_max_tokens_per_minute =
451 model.max_tokens_per_minute as usize / users_in_recent_minutes;
452 let per_user_max_tokens_per_day = model.max_tokens_per_day as usize / users_in_recent_days;
453
454 let checks = [
455 (
456 usage.requests_this_minute,
457 per_user_max_requests_per_minute,
458 UsageMeasure::RequestsPerMinute,
459 ),
460 (
461 usage.tokens_this_minute,
462 per_user_max_tokens_per_minute,
463 UsageMeasure::TokensPerMinute,
464 ),
465 (
466 usage.tokens_this_day,
467 per_user_max_tokens_per_day,
468 UsageMeasure::TokensPerDay,
469 ),
470 ];
471
472 for (used, limit, usage_measure) in checks {
473 // Temporarily bypass rate-limiting for staff members.
474 if claims.is_staff {
475 continue;
476 }
477
478 if used > limit {
479 let resource = match usage_measure {
480 UsageMeasure::RequestsPerMinute => "requests_per_minute",
481 UsageMeasure::TokensPerMinute => "tokens_per_minute",
482 UsageMeasure::TokensPerDay => "tokens_per_day",
483 _ => "",
484 };
485
486 if let Some(client) = state.clickhouse_client.as_ref() {
487 tracing::info!(
488 target: "user rate limit",
489 user_id = claims.user_id,
490 login = claims.github_user_login,
491 authn.jti = claims.jti,
492 is_staff = claims.is_staff,
493 provider = provider.to_string(),
494 model = model.name,
495 requests_this_minute = usage.requests_this_minute,
496 tokens_this_minute = usage.tokens_this_minute,
497 tokens_this_day = usage.tokens_this_day,
498 users_in_recent_minutes = users_in_recent_minutes,
499 users_in_recent_days = users_in_recent_days,
500 max_requests_per_minute = per_user_max_requests_per_minute,
501 max_tokens_per_minute = per_user_max_tokens_per_minute,
502 max_tokens_per_day = per_user_max_tokens_per_day,
503 );
504
505 report_llm_rate_limit(
506 client,
507 LlmRateLimitEventRow {
508 time: Utc::now().timestamp_millis(),
509 user_id: claims.user_id as i32,
510 is_staff: claims.is_staff,
511 plan: match claims.plan {
512 Plan::Free => "free".to_string(),
513 Plan::ZedPro => "zed_pro".to_string(),
514 },
515 model: model.name.clone(),
516 provider: provider.to_string(),
517 usage_measure: resource.to_string(),
518 requests_this_minute: usage.requests_this_minute as u64,
519 tokens_this_minute: usage.tokens_this_minute as u64,
520 tokens_this_day: usage.tokens_this_day as u64,
521 users_in_recent_minutes: users_in_recent_minutes as u64,
522 users_in_recent_days: users_in_recent_days as u64,
523 max_requests_per_minute: per_user_max_requests_per_minute as u64,
524 max_tokens_per_minute: per_user_max_tokens_per_minute as u64,
525 max_tokens_per_day: per_user_max_tokens_per_day as u64,
526 },
527 )
528 .await
529 .log_err();
530 }
531
532 return Err(Error::http(
533 StatusCode::TOO_MANY_REQUESTS,
534 format!("Rate limit exceeded. Maximum {} reached.", resource),
535 ));
536 }
537 }
538
539 Ok(())
540}
541
542struct TokenCountingStream<S> {
543 state: Arc<LlmState>,
544 claims: LlmTokenClaims,
545 provider: LanguageModelProvider,
546 model: String,
547 input_tokens: usize,
548 output_tokens: usize,
549 inner_stream: S,
550}
551
552impl<S> Stream for TokenCountingStream<S>
553where
554 S: Stream<Item = Result<(Vec<u8>, usize, usize), anyhow::Error>> + Unpin,
555{
556 type Item = Result<Vec<u8>, anyhow::Error>;
557
558 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
559 match Pin::new(&mut self.inner_stream).poll_next(cx) {
560 Poll::Ready(Some(Ok((mut bytes, input_tokens, output_tokens)))) => {
561 bytes.push(b'\n');
562 self.input_tokens += input_tokens;
563 self.output_tokens += output_tokens;
564 Poll::Ready(Some(Ok(bytes)))
565 }
566 Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
567 Poll::Ready(None) => Poll::Ready(None),
568 Poll::Pending => Poll::Pending,
569 }
570 }
571}
572
573impl<S> Drop for TokenCountingStream<S> {
574 fn drop(&mut self) {
575 let state = self.state.clone();
576 let claims = self.claims.clone();
577 let provider = self.provider;
578 let model = std::mem::take(&mut self.model);
579 let input_token_count = self.input_tokens;
580 let output_token_count = self.output_tokens;
581 self.state.executor.spawn_detached(async move {
582 let usage = state
583 .db
584 .record_usage(
585 UserId::from_proto(claims.user_id),
586 claims.is_staff,
587 provider,
588 &model,
589 input_token_count,
590 output_token_count,
591 Utc::now(),
592 )
593 .await
594 .log_err();
595
596 if let Some(usage) = usage {
597 tracing::info!(
598 target: "user usage",
599 user_id = claims.user_id,
600 login = claims.github_user_login,
601 authn.jti = claims.jti,
602 is_staff = claims.is_staff,
603 requests_this_minute = usage.requests_this_minute,
604 tokens_this_minute = usage.tokens_this_minute,
605 );
606
607 if let Some(clickhouse_client) = state.clickhouse_client.as_ref() {
608 report_llm_usage(
609 clickhouse_client,
610 LlmUsageEventRow {
611 time: Utc::now().timestamp_millis(),
612 user_id: claims.user_id as i32,
613 is_staff: claims.is_staff,
614 plan: match claims.plan {
615 Plan::Free => "free".to_string(),
616 Plan::ZedPro => "zed_pro".to_string(),
617 },
618 model,
619 provider: provider.to_string(),
620 input_token_count: input_token_count as u64,
621 output_token_count: output_token_count as u64,
622 requests_this_minute: usage.requests_this_minute as u64,
623 tokens_this_minute: usage.tokens_this_minute as u64,
624 tokens_this_day: usage.tokens_this_day as u64,
625 input_tokens_this_month: usage.input_tokens_this_month as u64,
626 output_tokens_this_month: usage.output_tokens_this_month as u64,
627 spending_this_month: usage.spending_this_month as u64,
628 lifetime_spending: usage.lifetime_spending as u64,
629 },
630 )
631 .await
632 .log_err();
633 }
634 }
635 })
636 }
637}
638
639pub fn log_usage_periodically(state: Arc<LlmState>) {
640 state.executor.clone().spawn_detached(async move {
641 loop {
642 state
643 .executor
644 .sleep(std::time::Duration::from_secs(30))
645 .await;
646
647 for provider in LanguageModelProvider::iter() {
648 for model in state.db.model_names_for_provider(provider) {
649 if let Some(active_user_count) = state
650 .get_active_user_count(provider, &model)
651 .await
652 .log_err()
653 {
654 tracing::info!(
655 target: "active user counts",
656 provider = provider.to_string(),
657 model = model,
658 users_in_recent_minutes = active_user_count.users_in_recent_minutes,
659 users_in_recent_days = active_user_count.users_in_recent_days,
660 );
661 }
662 }
663 }
664
665 if let Some(usages) = state
666 .db
667 .get_application_wide_usages_by_model(Utc::now())
668 .await
669 .log_err()
670 {
671 for usage in usages {
672 tracing::info!(
673 target: "computed usage",
674 provider = usage.provider.to_string(),
675 model = usage.model,
676 requests_this_minute = usage.requests_this_minute,
677 tokens_this_minute = usage.tokens_this_minute,
678 );
679 }
680 }
681 }
682 })
683}