sibling of 6a8be1714e
This commit is contained in:
Kirill Bulatov
GitHub
parent
a2aea00253
commit
dac29f2d73
302 changed files with 6480 additions and 19664 deletions
|
|
@ -23,7 +23,6 @@ use anyhow::{Context as _, anyhow, bail};
|
|||
use async_tungstenite::tungstenite::{
|
||||
Message as TungsteniteMessage, protocol::CloseFrame as TungsteniteCloseFrame,
|
||||
};
|
||||
use axum::headers::UserAgent;
|
||||
use axum::{
|
||||
Extension, Router, TypedHeader,
|
||||
body::Body,
|
||||
|
|
@ -434,8 +433,6 @@ impl Server {
|
|||
.add_request_handler(forward_mutating_project_request::<proto::SynchronizeContexts>)
|
||||
.add_request_handler(forward_mutating_project_request::<proto::Stage>)
|
||||
.add_request_handler(forward_mutating_project_request::<proto::Unstage>)
|
||||
.add_request_handler(forward_mutating_project_request::<proto::Stash>)
|
||||
.add_request_handler(forward_mutating_project_request::<proto::StashPop>)
|
||||
.add_request_handler(forward_mutating_project_request::<proto::Commit>)
|
||||
.add_request_handler(forward_mutating_project_request::<proto::GitInit>)
|
||||
.add_request_handler(forward_read_only_project_request::<proto::GetRemotes>)
|
||||
|
|
@ -751,7 +748,6 @@ impl Server {
|
|||
address: String,
|
||||
principal: Principal,
|
||||
zed_version: ZedVersion,
|
||||
user_agent: Option<String>,
|
||||
geoip_country_code: Option<String>,
|
||||
system_id: Option<String>,
|
||||
send_connection_id: Option<oneshot::Sender<ConnectionId>>,
|
||||
|
|
@ -764,14 +760,9 @@ impl Server {
|
|||
user_id=field::Empty,
|
||||
login=field::Empty,
|
||||
impersonator=field::Empty,
|
||||
user_agent=field::Empty,
|
||||
geoip_country_code=field::Empty
|
||||
);
|
||||
principal.update_span(&span);
|
||||
if let Some(user_agent) = user_agent {
|
||||
span.record("user_agent", user_agent);
|
||||
}
|
||||
|
||||
if let Some(country_code) = geoip_country_code.as_ref() {
|
||||
span.record("geoip_country_code", country_code);
|
||||
}
|
||||
|
|
@ -838,7 +829,7 @@ impl Server {
|
|||
// This arrangement ensures we will attempt to process earlier messages first, but fall
|
||||
// back to processing messages arrived later in the spirit of making progress.
|
||||
let mut foreground_message_handlers = FuturesUnordered::new();
|
||||
let concurrent_handlers = Arc::new(Semaphore::new(512));
|
||||
let concurrent_handlers = Arc::new(Semaphore::new(256));
|
||||
loop {
|
||||
let next_message = async {
|
||||
let permit = concurrent_handlers.clone().acquire_owned().await.unwrap();
|
||||
|
|
@ -1179,7 +1170,6 @@ pub async fn handle_websocket_request(
|
|||
ConnectInfo(socket_address): ConnectInfo<SocketAddr>,
|
||||
Extension(server): Extension<Arc<Server>>,
|
||||
Extension(principal): Extension<Principal>,
|
||||
user_agent: Option<TypedHeader<UserAgent>>,
|
||||
country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
|
||||
system_id_header: Option<TypedHeader<SystemIdHeader>>,
|
||||
ws: WebSocketUpgrade,
|
||||
|
|
@ -1235,7 +1225,6 @@ pub async fn handle_websocket_request(
|
|||
socket_address,
|
||||
principal,
|
||||
version,
|
||||
user_agent.map(|header| header.to_string()),
|
||||
country_code_header.map(|header| header.to_string()),
|
||||
system_id_header.map(|header| header.to_string()),
|
||||
None,
|
||||
|
|
@ -2868,12 +2857,12 @@ async fn make_update_user_plan_message(
|
|||
}
|
||||
|
||||
fn model_requests_limit(
|
||||
plan: cloud_llm_client::Plan,
|
||||
plan: zed_llm_client::Plan,
|
||||
feature_flags: &Vec<String>,
|
||||
) -> cloud_llm_client::UsageLimit {
|
||||
) -> zed_llm_client::UsageLimit {
|
||||
match plan.model_requests_limit() {
|
||||
cloud_llm_client::UsageLimit::Limited(limit) => {
|
||||
let limit = if plan == cloud_llm_client::Plan::ZedProTrial
|
||||
zed_llm_client::UsageLimit::Limited(limit) => {
|
||||
let limit = if plan == zed_llm_client::Plan::ZedProTrial
|
||||
&& feature_flags
|
||||
.iter()
|
||||
.any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG)
|
||||
|
|
@ -2883,9 +2872,9 @@ fn model_requests_limit(
|
|||
limit
|
||||
};
|
||||
|
||||
cloud_llm_client::UsageLimit::Limited(limit)
|
||||
zed_llm_client::UsageLimit::Limited(limit)
|
||||
}
|
||||
cloud_llm_client::UsageLimit::Unlimited => cloud_llm_client::UsageLimit::Unlimited,
|
||||
zed_llm_client::UsageLimit::Unlimited => zed_llm_client::UsageLimit::Unlimited,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -2895,21 +2884,21 @@ fn subscription_usage_to_proto(
|
|||
feature_flags: &Vec<String>,
|
||||
) -> proto::SubscriptionUsage {
|
||||
let plan = match plan {
|
||||
proto::Plan::Free => cloud_llm_client::Plan::ZedFree,
|
||||
proto::Plan::ZedPro => cloud_llm_client::Plan::ZedPro,
|
||||
proto::Plan::ZedProTrial => cloud_llm_client::Plan::ZedProTrial,
|
||||
proto::Plan::Free => zed_llm_client::Plan::ZedFree,
|
||||
proto::Plan::ZedPro => zed_llm_client::Plan::ZedPro,
|
||||
proto::Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
|
||||
};
|
||||
|
||||
proto::SubscriptionUsage {
|
||||
model_requests_usage_amount: usage.model_requests as u32,
|
||||
model_requests_usage_limit: Some(proto::UsageLimit {
|
||||
variant: Some(match model_requests_limit(plan, feature_flags) {
|
||||
cloud_llm_client::UsageLimit::Limited(limit) => {
|
||||
zed_llm_client::UsageLimit::Limited(limit) => {
|
||||
proto::usage_limit::Variant::Limited(proto::usage_limit::Limited {
|
||||
limit: limit as u32,
|
||||
})
|
||||
}
|
||||
cloud_llm_client::UsageLimit::Unlimited => {
|
||||
zed_llm_client::UsageLimit::Unlimited => {
|
||||
proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {})
|
||||
}
|
||||
}),
|
||||
|
|
@ -2917,12 +2906,12 @@ fn subscription_usage_to_proto(
|
|||
edit_predictions_usage_amount: usage.edit_predictions as u32,
|
||||
edit_predictions_usage_limit: Some(proto::UsageLimit {
|
||||
variant: Some(match plan.edit_predictions_limit() {
|
||||
cloud_llm_client::UsageLimit::Limited(limit) => {
|
||||
zed_llm_client::UsageLimit::Limited(limit) => {
|
||||
proto::usage_limit::Variant::Limited(proto::usage_limit::Limited {
|
||||
limit: limit as u32,
|
||||
})
|
||||
}
|
||||
cloud_llm_client::UsageLimit::Unlimited => {
|
||||
zed_llm_client::UsageLimit::Unlimited => {
|
||||
proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {})
|
||||
}
|
||||
}),
|
||||
|
|
@ -2935,21 +2924,21 @@ fn make_default_subscription_usage(
|
|||
feature_flags: &Vec<String>,
|
||||
) -> proto::SubscriptionUsage {
|
||||
let plan = match plan {
|
||||
proto::Plan::Free => cloud_llm_client::Plan::ZedFree,
|
||||
proto::Plan::ZedPro => cloud_llm_client::Plan::ZedPro,
|
||||
proto::Plan::ZedProTrial => cloud_llm_client::Plan::ZedProTrial,
|
||||
proto::Plan::Free => zed_llm_client::Plan::ZedFree,
|
||||
proto::Plan::ZedPro => zed_llm_client::Plan::ZedPro,
|
||||
proto::Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
|
||||
};
|
||||
|
||||
proto::SubscriptionUsage {
|
||||
model_requests_usage_amount: 0,
|
||||
model_requests_usage_limit: Some(proto::UsageLimit {
|
||||
variant: Some(match model_requests_limit(plan, feature_flags) {
|
||||
cloud_llm_client::UsageLimit::Limited(limit) => {
|
||||
zed_llm_client::UsageLimit::Limited(limit) => {
|
||||
proto::usage_limit::Variant::Limited(proto::usage_limit::Limited {
|
||||
limit: limit as u32,
|
||||
})
|
||||
}
|
||||
cloud_llm_client::UsageLimit::Unlimited => {
|
||||
zed_llm_client::UsageLimit::Unlimited => {
|
||||
proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {})
|
||||
}
|
||||
}),
|
||||
|
|
@ -2957,12 +2946,12 @@ fn make_default_subscription_usage(
|
|||
edit_predictions_usage_amount: 0,
|
||||
edit_predictions_usage_limit: Some(proto::UsageLimit {
|
||||
variant: Some(match plan.edit_predictions_limit() {
|
||||
cloud_llm_client::UsageLimit::Limited(limit) => {
|
||||
zed_llm_client::UsageLimit::Limited(limit) => {
|
||||
proto::usage_limit::Variant::Limited(proto::usage_limit::Limited {
|
||||
limit: limit as u32,
|
||||
})
|
||||
}
|
||||
cloud_llm_client::UsageLimit::Unlimited => {
|
||||
zed_llm_client::UsageLimit::Unlimited => {
|
||||
proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {})
|
||||
}
|
||||
}),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue