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