mixa/zed
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into debugger

Browse Source
This commit is contained in:
Remco Smits
2024-10-09 11:42:59 +02:00
parent a728f9e751 9c54bd1bd4
commit 8a4f677119
92 changed files with 2169 additions and 1611 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/close_stale_issues.yml vendored
View File

@@ -17,7 +17,7 @@ jobs:
We're working to clean up our issue tracker by closing older issues that might not be relevant anymore. Are you able to reproduce this issue in the latest version of Zed? If so, please let us know by commenting on this issue and we will keep it open; otherwise, we'll close it in 7 days. Feel free to open a new issue if you're seeing this message after the issue has been closed.
Thanks for your help!
close-issue-message: "This issue was closed due to inactivity; feel free to open a new issue if you're still experiencing this problem!"
close-issue-message: "This issue was closed due to inactivity. If you're still experiencing this problem, feel free to ping a Zed team member to reopen this issue or open a new one."
# We will increase `days-before-stale` to 365 on or after Jan 24th,
# 2024. This date marks one year since migrating issues from
# 'community' to 'zed' repository. The migration added activity to all

613
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

15
Cargo.toml
View File

@@ -55,6 +55,7 @@ members = [
"crates/indexed_docs",
"crates/inline_completion_button",
"crates/install_cli",
"crates/isahc_http_client",
"crates/journal",
"crates/language",
"crates/language_model",
@@ -90,7 +91,6 @@ members = [
"crates/remote",
"crates/remote_server",
"crates/repl",
"crates/reqwest_client",
"crates/rich_text",
"crates/rope",
"crates/rpc",
@@ -125,7 +125,6 @@ members = [
"crates/ui",
"crates/ui_input",
"crates/ui_macros",
"crates/ureq_client",
"crates/util",
"crates/vcs_menu",
"crates/vim",
@@ -236,6 +235,7 @@ image_viewer = { path = "crates/image_viewer" }
indexed_docs = { path = "crates/indexed_docs" }
inline_completion_button = { path = "crates/inline_completion_button" }
install_cli = { path = "crates/install_cli" }
isahc_http_client = { path = "crates/isahc_http_client" }
journal = { path = "crates/journal" }
language = { path = "crates/language" }
language_model = { path = "crates/language_model" }
@@ -272,7 +272,6 @@ release_channel = { path = "crates/release_channel" }
remote = { path = "crates/remote" }
remote_server = { path = "crates/remote_server" }
repl = { path = "crates/repl" }
reqwest_client = { path = "crates/reqwest_client" }
rich_text = { path = "crates/rich_text" }
rope = { path = "crates/rope" }
rpc = { path = "crates/rpc" }
@@ -307,7 +306,6 @@ title_bar = { path = "crates/title_bar" }
ui = { path = "crates/ui" }
ui_input = { path = "crates/ui_input" }
ui_macros = { path = "crates/ui_macros" }
ureq_client = { path = "crates/ureq_client" }
util = { path = "crates/util" }
vcs_menu = { path = "crates/vcs_menu" }
vim = { path = "crates/vim" }
@@ -335,7 +333,7 @@ async-pipe = { git = "https://github.com/zed-industries/async-pipe-rs", rev = "8
async-recursion = "1.0.0"
async-tar = "0.5.0"
async-trait = "0.1"
async-tungstenite = "0.28"
async-tungstenite = "0.23"
async-watch = "0.3.1"
async_zip = { version = "0.0.17", features = ["deflate", "deflate64"] }
base64 = "0.22"
@@ -375,6 +373,10 @@ ignore = "0.4.22"
image = "0.25.1"
indexmap = { version = "1.6.2", features = ["serde"] }
indoc = "2"
# We explicitly disable http2 support in isahc.
isahc = { version = "1.7.2", default-features = false, features = [
"text-decoding",
] }
itertools = "0.13.0"
jsonwebtoken = "9.3"
libc = "0.2"
@@ -399,14 +401,13 @@ pulldown-cmark = { version = "0.12.0", default-features = false }
rand = "0.8.5"
regex = "1.5"
repair_json = "0.1.0"
reqwest = { git = "https://github.com/zed-industries/reqwest.git", rev = "fd110f6998da16bbca97b6dddda9be7827c50e29" }
rsa = "0.9.6"
runtimelib = { version = "0.15", default-features = false, features = [
"async-dispatcher-runtime",
] }
rustc-demangle = "0.1.23"
rust-embed = { version = "8.4", features = ["include-exclude"] }
rustls = "0.21.12"
rustls = "0.20.3"
rustls-native-certs = "0.8.0"
schemars = { version = "0.8", features = ["impl_json_schema"] }
semver = "1.0"

4
Dockerfile-collab
View File

@@ -31,10 +31,12 @@ ENV GITHUB_SHA=$GITHUB_SHA
# - Staging: `4f408ec65a3867278322a189b4eb20f1ab51f508`
# - Production: `fc4c533d0a8c489e5636a4249d2b52a80039fbd7`
#
# Also add `cmake`, since we need it to build `wasmtime`.
#
# Installing these as a temporary workaround, but I think ideally we'd want to figure
# out what caused them to be included in the first place.
RUN apt-get update; \
apt-get install -y --no-install-recommends libxkbcommon-dev libxkbcommon-x11-dev
apt-get install -y --no-install-recommends libxkbcommon-dev libxkbcommon-x11-dev cmake
RUN --mount=type=cache,target=./script/node_modules \
--mount=type=cache,target=/usr/local/cargo/registry \

3
assets/settings/default.json
View File

@@ -840,6 +840,9 @@
"allowed": true
}
},
"Dart": {
"tab_size": 2
},
"Elixir": {
"language_servers": ["elixir-ls", "!next-ls", "!lexical", "..."]
},

26
crates/activity_indicator/src/activity_indicator.rs
View File

@@ -10,7 +10,7 @@ use gpui::{
use language::{
LanguageRegistry, LanguageServerBinaryStatus, LanguageServerId, LanguageServerName,
};
use project::{LanguageServerProgress, Project};
use project::{EnvironmentErrorMessage, LanguageServerProgress, Project, WorktreeId};
use smallvec::SmallVec;
use std::{cmp::Reverse, fmt::Write, sync::Arc, time::Duration};
use ui::{prelude::*, ButtonLike, ContextMenu, PopoverMenu, PopoverMenuHandle};
@@ -175,7 +175,31 @@ impl ActivityIndicator {
.flatten()
}
fn pending_environment_errors<'a>(
&'a self,
cx: &'a AppContext,
) -> impl Iterator<Item = (&'a WorktreeId, &'a EnvironmentErrorMessage)> {
self.project.read(cx).shell_environment_errors(cx)
}
fn content_to_render(&mut self, cx: &mut ViewContext<Self>) -> Option<Content> {
// Show if any direnv calls failed
if let Some((&worktree_id, error)) = self.pending_environment_errors(cx).next() {
return Some(Content {
icon: Some(
Icon::new(IconName::Warning)
.size(IconSize::Small)
.into_any_element(),
),
message: error.0.clone(),
on_click: Some(Arc::new(move |this, cx| {
this.project.update(cx, |project, cx| {
project.remove_environment_error(cx, worktree_id);
});
cx.dispatch_action(Box::new(workspace::OpenLog));
})),
});
}
// Show any language server has pending activity.
let mut pending_work = self.pending_language_server_work(cx);
if let Some(PendingWork {

4
crates/anthropic/src/anthropic.rs
View File

@@ -521,6 +521,10 @@ pub struct Usage {
pub input_tokens: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_tokens: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_input_tokens: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_input_tokens: Option<u32>,
}
#[derive(Debug, Serialize, Deserialize)]

3
crates/client/Cargo.toml
View File

@@ -18,7 +18,6 @@ test-support = ["clock/test-support", "collections/test-support", "gpui/test-sup
[dependencies]
anyhow.workspace = true
async-recursion = "0.3"
async-tls = "0.13"
async-tungstenite = { workspace = true, features = ["async-std", "async-tls"] }
chrono = { workspace = true, features = ["serde"] }
clock.workspace = true
@@ -35,6 +34,8 @@ postage.workspace = true
rand.workspace = true
release_channel.workspace = true
rpc = { workspace = true, features = ["gpui"] }
rustls.workspace = true
rustls-native-certs.workspace = true
schemars.workspace = true
serde.workspace = true
serde_json.workspace = true

27
crates/client/src/client.rs
View File

@@ -1023,7 +1023,7 @@ impl Client {
&self,
http: Arc<HttpClientWithUrl>,
release_channel: Option<ReleaseChannel>,
) -> impl Future<Output = Result<url::Url>> {
) -> impl Future<Output = Result<Url>> {
#[cfg(any(test, feature = "test-support"))]
let url_override = self.rpc_url.read().clone();
@@ -1117,7 +1117,7 @@ impl Client {
// for us from the RPC URL.
//
// Among other things, it will generate and set a `Sec-WebSocket-Key` header for us.
let mut request = IntoClientRequest::into_client_request(rpc_url.as_str())?;
let mut request = rpc_url.into_client_request()?;
// We then modify the request to add our desired headers.
let request_headers = request.headers_mut();
@@ -1137,13 +1137,30 @@ impl Client {
match url_scheme {
Https => {
let client_config = {
let mut root_store = rustls::RootCertStore::empty();
let root_certs = rustls_native_certs::load_native_certs();
for error in root_certs.errors {
log::warn!("error loading native certs: {:?}", error);
}
root_store.add_parsable_certificates(
&root_certs
.certs
.into_iter()
.map(|cert| cert.as_ref().to_owned())
.collect::<Vec<_>>(),
);
rustls::ClientConfig::builder()
.with_safe_defaults()
.with_root_certificates(root_store)
.with_no_client_auth()
};
let (stream, _) =
async_tungstenite::async_tls::client_async_tls_with_connector(
request,
stream,
Some(async_tls::TlsConnector::from(
http_client::TLS_CONFIG.clone(),
)),
Some(client_config.into()),
)
.await?;
Ok(Connection::new(

13
crates/client/src/telemetry.rs
View File

@@ -364,6 +364,7 @@ impl Telemetry {
operation: &'static str,
copilot_enabled: bool,
copilot_enabled_for_language: bool,
is_via_ssh: bool,
) {
let event = Event::Editor(EditorEvent {
file_extension,
@@ -371,6 +372,7 @@ impl Telemetry {
operation: operation.into(),
copilot_enabled,
copilot_enabled_for_language,
is_via_ssh,
});
self.report_event(event)
@@ -456,7 +458,7 @@ impl Telemetry {
}))
}
pub fn log_edit_event(self: &Arc<Self>, environment: &'static str) {
pub fn log_edit_event(self: &Arc<Self>, environment: &'static str, is_via_ssh: bool) {
let mut state = self.state.lock();
let period_data = state.event_coalescer.log_event(environment);
drop(state);
@@ -465,6 +467,7 @@ impl Telemetry {
let event = Event::Edit(EditEvent {
duration: end.timestamp_millis() - start.timestamp_millis(),
environment: environment.to_string(),
is_via_ssh,
});
self.report_event(event);
@@ -485,7 +488,7 @@ impl Telemetry {
worktree_id: WorktreeId,
updated_entries_set: &UpdatedEntriesSet,
) {
let project_names: Vec<String> = {
let project_type_names: Vec<String> = {
let mut state = self.state.lock();
state
.worktree_id_map
@@ -521,8 +524,8 @@ impl Telemetry {
};
// Done on purpose to avoid calling `self.state.lock()` multiple times
for project_name in project_names {
self.report_app_event(format!("open {} project", project_name));
for project_type_name in project_type_names {
self.report_app_event(format!("open {} project", project_type_name));
}
}
@@ -650,7 +653,7 @@ impl Telemetry {
.build_zed_api_url("/telemetry/events", &[])?
.as_ref(),
)
.header("Content-Type", "text/plain")
.header("Content-Type", "application/json")
.header("x-zed-checksum", checksum)
.body(json_bytes.into());

9
crates/clock/src/clock.rs
View File

@@ -216,10 +216,11 @@ impl fmt::Debug for Global {
if timestamp.replica_id > 0 {
write!(f, ", ")?;
}
write!(f, "{}: {}", timestamp.replica_id, timestamp.value)?;
}
if self.local_branch_value > 0 {
write!(f, "<branch>: {}", self.local_branch_value)?;
if timestamp.replica_id == LOCAL_BRANCH_REPLICA_ID {
write!(f, "<branch>: {}", timestamp.value)?;
} else {
write!(f, "{}: {}", timestamp.replica_id, timestamp.value)?;
}
}
write!(f, "}}")
}

2
crates/collab/Cargo.toml
View File

@@ -37,6 +37,7 @@ futures.workspace = true
google_ai.workspace = true
hex.workspace = true
http_client.workspace = true
isahc_http_client.workspace = true
jsonwebtoken.workspace = true
live_kit_server.workspace = true
log.workspace = true
@@ -47,7 +48,6 @@ prometheus = "0.13"
prost.workspace = true
rand.workspace = true
reqwest = { version = "0.11", features = ["json"] }
reqwest_client.workspace = true
rpc.workspace = true
rustc-demangle.workspace = true
scrypt = "0.11"

11
crates/collab/migrations_llm/20241007173634_add_cache_token_counts.sql Normal file
View File

@@ -0,0 +1,11 @@
alter table models
add column price_per_million_cache_creation_input_tokens integer not null default 0,
add column price_per_million_cache_read_input_tokens integer not null default 0;
alter table usages
add column cache_creation_input_tokens_this_month bigint not null default 0,
add column cache_read_input_tokens_this_month bigint not null default 0;
alter table lifetime_usages
add column cache_creation_input_tokens bigint not null default 0,
add column cache_read_input_tokens bigint not null default 0;

3
crates/collab/migrations_llm/20241007220716_drop_incorrect_usages_columns.sql Normal file
View File

@@ -0,0 +1,3 @@
alter table usages
drop column cache_creation_input_tokens_this_month,
drop column cache_read_input_tokens_this_month;

13
crates/collab/migrations_llm/20241008155620_create_monthly_usages.sql Normal file
View File

@@ -0,0 +1,13 @@
create table monthly_usages (
id serial primary key,
user_id integer not null,
model_id integer not null references models (id) on delete cascade,
month integer not null,
year integer not null,
input_tokens bigint not null default 0,
cache_creation_input_tokens bigint not null default 0,
cache_read_input_tokens bigint not null default 0,
output_tokens bigint not null default 0
);
create unique index uix_monthly_usages_on_user_id_model_id_month_year on monthly_usages (user_id, model_id, month, year);

103
crates/collab/src/api/billing.rs
View File

@@ -22,12 +22,15 @@ use stripe::{
};
use util::ResultExt;
use crate::db::billing_subscription::StripeSubscriptionStatus;
use crate::db::billing_subscription::{self, StripeSubscriptionStatus};
use crate::db::{
billing_customer, BillingSubscriptionId, CreateBillingCustomerParams,
CreateBillingSubscriptionParams, CreateProcessedStripeEventParams, UpdateBillingCustomerParams,
UpdateBillingSubscriptionParams,
};
use crate::llm::db::LlmDatabase;
use crate::llm::MONTHLY_SPENDING_LIMIT_IN_CENTS;
use crate::rpc::ResultExt as _;
use crate::{AppState, Error, Result};
pub fn router() -> Router {
@@ -79,7 +82,7 @@ async fn list_billing_subscriptions(
.into_iter()
.map(|subscription| BillingSubscriptionJson {
id: subscription.id,
name: "Zed Pro".to_string(),
name: "Zed LLM Usage".to_string(),
status: subscription.stripe_subscription_status,
cancel_at: subscription.stripe_cancel_at.map(|cancel_at| {
cancel_at
@@ -117,7 +120,7 @@ async fn create_billing_subscription(
let Some((stripe_client, stripe_price_id)) = app
.stripe_client
.clone()
.zip(app.config.stripe_price_id.clone())
.zip(app.config.stripe_llm_usage_price_id.clone())
else {
log::error!("failed to retrieve Stripe client or price ID");
Err(Error::http(
@@ -150,7 +153,7 @@ async fn create_billing_subscription(
params.client_reference_id = Some(user.github_login.as_str());
params.line_items = Some(vec![CreateCheckoutSessionLineItems {
price: Some(stripe_price_id.to_string()),
quantity: Some(1),
quantity: Some(0),
..Default::default()
}]);
let success_url = format!("{}/account", app.config.zed_dot_dev_url());
@@ -631,3 +634,95 @@ async fn find_or_create_billing_customer(
Ok(Some(billing_customer))
}
const SYNC_LLM_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
pub fn sync_llm_usage_with_stripe_periodically(app: Arc<AppState>, llm_db: LlmDatabase) {
let Some(stripe_client) = app.stripe_client.clone() else {
log::warn!("failed to retrieve Stripe client");
return;
};
let Some(stripe_llm_usage_price_id) = app.config.stripe_llm_usage_price_id.clone() else {
log::warn!("failed to retrieve Stripe LLM usage price ID");
return;
};
let executor = app.executor.clone();
executor.spawn_detached({
let executor = executor.clone();
async move {
loop {
sync_with_stripe(
&app,
&llm_db,
&stripe_client,
stripe_llm_usage_price_id.clone(),
)
.await
.trace_err();
executor.sleep(SYNC_LLM_USAGE_WITH_STRIPE_INTERVAL).await;
}
}
});
}
async fn sync_with_stripe(
app: &Arc<AppState>,
llm_db: &LlmDatabase,
stripe_client: &stripe::Client,
stripe_llm_usage_price_id: Arc<str>,
) -> anyhow::Result<()> {
let subscriptions = app.db.get_active_billing_subscriptions().await?;
for (customer, subscription) in subscriptions {
update_stripe_subscription(
llm_db,
stripe_client,
&stripe_llm_usage_price_id,
customer,
subscription,
)
.await
.log_err();
}
Ok(())
}
async fn update_stripe_subscription(
llm_db: &LlmDatabase,
stripe_client: &stripe::Client,
stripe_llm_usage_price_id: &Arc<str>,
customer: billing_customer::Model,
subscription: billing_subscription::Model,
) -> Result<(), anyhow::Error> {
let monthly_spending = llm_db
.get_user_spending_for_month(customer.user_id, Utc::now())
.await?;
let subscription_id = SubscriptionId::from_str(&subscription.stripe_subscription_id)
.context("failed to parse subscription ID")?;
let monthly_spending_over_free_tier =
monthly_spending.saturating_sub(MONTHLY_SPENDING_LIMIT_IN_CENTS);
let new_quantity = (monthly_spending_over_free_tier as f32 / 100.).ceil();
Subscription::update(
stripe_client,
&subscription_id,
stripe::UpdateSubscription {
items: Some(vec![stripe::UpdateSubscriptionItems {
// TODO: Do we need to send up the `id` if a subscription item
// with this price already exists, or will Stripe take care of
// it?
id: None,
price: Some(stripe_llm_usage_price_id.to_string()),
quantity: Some(new_quantity as u64),
..Default::default()
}]),
..Default::default()
},
)
.await?;
Ok(())
}

6
crates/collab/src/api/events.rs
View File

@@ -429,8 +429,6 @@ pub async fn post_events(
country_code.clone(),
checksum_matched,
)),
// Needed for clients sending old copilot_event types
Event::Copilot(_) => {}
Event::InlineCompletion(event) => {
to_upload
.inline_completion_events
@@ -679,6 +677,7 @@ pub struct EditorEventRow {
minor: Option<i32>,
patch: Option<i32>,
checksum_matched: bool,
is_via_ssh: bool,
}
impl EditorEventRow {
@@ -720,6 +719,7 @@ impl EditorEventRow {
region_code: "".to_string(),
city: "".to_string(),
historical_event: false,
is_via_ssh: event.is_via_ssh,
}
}
}
@@ -1263,6 +1263,7 @@ pub struct EditEventRow {
period_start: i64,
period_end: i64,
environment: String,
is_via_ssh: bool,
}
impl EditEventRow {
@@ -1296,6 +1297,7 @@ impl EditEventRow {
period_start: period_start.timestamp_millis(),
period_end: period_end.timestamp_millis(),
environment: event.environment,
is_via_ssh: event.is_via_ssh,
}
}
}

23
crates/collab/src/db/queries/billing_subscriptions.rs
View File

@@ -112,6 +112,29 @@ impl Database {
.await
}
pub async fn get_active_billing_subscriptions(
&self,
) -> Result<Vec<(billing_customer::Model, billing_subscription::Model)>> {
self.transaction(|tx| async move {
let mut result = Vec::new();
let mut rows = billing_subscription::Entity::find()
.inner_join(billing_customer::Entity)
.select_also(billing_customer::Entity)
.order_by_asc(billing_subscription::Column::Id)
.stream(&*tx)
.await?;
while let Some(row) = rows.next().await {
if let (subscription, Some(customer)) = row? {
result.push((customer, subscription));
}
}
Ok(result)
})
.await
}
/// Returns whether the user has an active billing subscription.
pub async fn has_active_billing_subscription(&self, user_id: UserId) -> Result<bool> {
Ok(self.count_active_billing_subscriptions(user_id).await? > 0)

8
crates/collab/src/lib.rs
View File

@@ -174,7 +174,7 @@ pub struct Config {
pub slack_panics_webhook: Option<String>,
pub auto_join_channel_id: Option<ChannelId>,
pub stripe_api_key: Option<String>,
pub stripe_price_id: Option<Arc<str>>,
pub stripe_llm_usage_price_id: Option<Arc<str>>,
pub supermaven_admin_api_key: Option<Arc<str>>,
pub user_backfiller_github_access_token: Option<Arc<str>>,
}
@@ -193,6 +193,10 @@ impl Config {
}
}
pub fn is_llm_billing_enabled(&self) -> bool {
self.stripe_llm_usage_price_id.is_some()
}
#[cfg(test)]
pub fn test() -> Self {
Self {
@@ -231,7 +235,7 @@ impl Config {
migrations_path: None,
seed_path: None,
stripe_api_key: None,
stripe_price_id: None,
stripe_llm_usage_price_id: None,
supermaven_admin_api_key: None,
user_backfiller_github_access_token: None,
}

108
crates/collab/src/llm.rs
View File

@@ -22,8 +22,7 @@ use chrono::{DateTime, Duration, Utc};
use collections::HashMap;
use db::{usage_measure::UsageMeasure, ActiveUserCount, LlmDatabase};
use futures::{Stream, StreamExt as _};
use reqwest_client::ReqwestClient;
use isahc_http_client::IsahcHttpClient;
use rpc::ListModelsResponse;
use rpc::{
proto::Plan, LanguageModelProvider, PerformCompletionParams, EXPIRED_LLM_TOKEN_HEADER_NAME,
@@ -44,7 +43,7 @@ pub struct LlmState {
pub config: Config,
pub executor: Executor,
pub db: Arc<LlmDatabase>,
pub http_client: ReqwestClient,
pub http_client: IsahcHttpClient,
pub clickhouse_client: Option<clickhouse::Client>,
active_user_count_by_model:
RwLock<HashMap<(LanguageModelProvider, String), (DateTime<Utc>, ActiveUserCount)>>,
@@ -70,8 +69,11 @@ impl LlmState {
let db = Arc::new(db);
let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));
let http_client =
ReqwestClient::user_agent(&user_agent).context("failed to construct http client")?;
let http_client = IsahcHttpClient::builder()
.default_header("User-Agent", user_agent)
.build()
.map(IsahcHttpClient::from)
.context("failed to construct http client")?;
let this = Self {
executor,
@@ -318,22 +320,31 @@ async fn perform_completion(
chunks
.map(move |event| {
let chunk = event?;
let (input_tokens, output_tokens) = match &chunk {
let (
input_tokens,
output_tokens,
cache_creation_input_tokens,
cache_read_input_tokens,
) = match &chunk {
anthropic::Event::MessageStart {
message: anthropic::Response { usage, .. },
}
| anthropic::Event::MessageDelta { usage, .. } => (
usage.input_tokens.unwrap_or(0) as usize,
usage.output_tokens.unwrap_or(0) as usize,
usage.cache_creation_input_tokens.unwrap_or(0) as usize,
usage.cache_read_input_tokens.unwrap_or(0) as usize,
),
_ => (0, 0),
_ => (0, 0, 0, 0),
};
anyhow::Ok((
serde_json::to_vec(&chunk).unwrap(),
anyhow::Ok(CompletionChunk {
bytes: serde_json::to_vec(&chunk).unwrap(),
input_tokens,
output_tokens,
))
cache_creation_input_tokens,
cache_read_input_tokens,
})
})
.boxed()
}
@@ -359,11 +370,13 @@ async fn perform_completion(
chunk.usage.as_ref().map_or(0, |u| u.prompt_tokens) as usize;
let output_tokens =
chunk.usage.as_ref().map_or(0, |u| u.completion_tokens) as usize;
(
serde_json::to_vec(&chunk).unwrap(),
CompletionChunk {
bytes: serde_json::to_vec(&chunk).unwrap(),
input_tokens,
output_tokens,
)
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
}
})
})
.boxed()
@@ -387,13 +400,13 @@ async fn perform_completion(
.map(|event| {
event.map(|chunk| {
// TODO - implement token counting for Google AI
let input_tokens = 0;
let output_tokens = 0;
(
serde_json::to_vec(&chunk).unwrap(),
input_tokens,
output_tokens,
)
CompletionChunk {
bytes: serde_json::to_vec(&chunk).unwrap(),
input_tokens: 0,
output_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
}
})
})
.boxed()
@@ -407,6 +420,8 @@ async fn perform_completion(
model,
input_tokens: 0,
output_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
inner_stream: stream,
})))
}
@@ -423,6 +438,9 @@ fn normalize_model_name(known_models: Vec<String>, name: String) -> String {
}
}
/// The maximum monthly spending an individual user can reach before they have to pay.
pub const MONTHLY_SPENDING_LIMIT_IN_CENTS: usize = 5 * 100;
/// The maximum lifetime spending an individual user can reach before being cut off.
///
/// Represented in cents.
@@ -445,6 +463,18 @@ async fn check_usage_limit(
)
.await?;
if state.config.is_llm_billing_enabled() {
if usage.spending_this_month >= MONTHLY_SPENDING_LIMIT_IN_CENTS {
if !claims.has_llm_subscription.unwrap_or(false) {
return Err(Error::http(
StatusCode::PAYMENT_REQUIRED,
"Maximum spending limit reached for this month.".to_string(),
));
}
}
}
// TODO: Remove this once we've rolled out monthly spending limits.
if usage.lifetime_spending >= LIFETIME_SPENDING_LIMIT_IN_CENTS {
return Err(Error::http(
StatusCode::FORBIDDEN,
@@ -492,7 +522,6 @@ async fn check_usage_limit(
UsageMeasure::RequestsPerMinute => "requests_per_minute",
UsageMeasure::TokensPerMinute => "tokens_per_minute",
UsageMeasure::TokensPerDay => "tokens_per_day",
_ => "",
};
if let Some(client) = state.clickhouse_client.as_ref() {
@@ -551,6 +580,14 @@ async fn check_usage_limit(
Ok(())
}
struct CompletionChunk {
bytes: Vec<u8>,
input_tokens: usize,
output_tokens: usize,
cache_creation_input_tokens: usize,
cache_read_input_tokens: usize,
}
struct TokenCountingStream<S> {
state: Arc<LlmState>,
claims: LlmTokenClaims,
@@ -558,22 +595,26 @@ struct TokenCountingStream<S> {
model: String,
input_tokens: usize,
output_tokens: usize,
cache_creation_input_tokens: usize,
cache_read_input_tokens: usize,
inner_stream: S,
}
impl<S> Stream for TokenCountingStream<S>
where
S: Stream<Item = Result<(Vec<u8>, usize, usize), anyhow::Error>> + Unpin,
S: Stream<Item = Result<CompletionChunk, anyhow::Error>> + Unpin,
{
type Item = Result<Vec<u8>, anyhow::Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match Pin::new(&mut self.inner_stream).poll_next(cx) {
Poll::Ready(Some(Ok((mut bytes, input_tokens, output_tokens)))) => {
bytes.push(b'\n');
self.input_tokens += input_tokens;
self.output_tokens += output_tokens;
Poll::Ready(Some(Ok(bytes)))
Poll::Ready(Some(Ok(mut chunk))) => {
chunk.bytes.push(b'\n');
self.input_tokens += chunk.input_tokens;
self.output_tokens += chunk.output_tokens;
self.cache_creation_input_tokens += chunk.cache_creation_input_tokens;
self.cache_read_input_tokens += chunk.cache_read_input_tokens;
Poll::Ready(Some(Ok(chunk.bytes)))
}
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
Poll::Ready(None) => Poll::Ready(None),
@@ -590,6 +631,8 @@ impl<S> Drop for TokenCountingStream<S> {
let model = std::mem::take(&mut self.model);
let input_token_count = self.input_tokens;
let output_token_count = self.output_tokens;
let cache_creation_input_token_count = self.cache_creation_input_tokens;
let cache_read_input_token_count = self.cache_read_input_tokens;
self.state.executor.spawn_detached(async move {
let usage = state
.db
@@ -599,6 +642,8 @@ impl<S> Drop for TokenCountingStream<S> {
provider,
&model,
input_token_count,
cache_creation_input_token_count,
cache_read_input_token_count,
output_token_count,
Utc::now(),
)
@@ -630,11 +675,20 @@ impl<S> Drop for TokenCountingStream<S> {
model,
provider: provider.to_string(),
input_token_count: input_token_count as u64,
cache_creation_input_token_count: cache_creation_input_token_count
as u64,
cache_read_input_token_count: cache_read_input_token_count as u64,
output_token_count: output_token_count as u64,
requests_this_minute: usage.requests_this_minute as u64,
tokens_this_minute: usage.tokens_this_minute as u64,
tokens_this_day: usage.tokens_this_day as u64,
input_tokens_this_month: usage.input_tokens_this_month as u64,
cache_creation_input_tokens_this_month: usage
.cache_creation_input_tokens_this_month
as u64,
cache_read_input_tokens_this_month: usage
.cache_read_input_tokens_this_month
as u64,
output_tokens_this_month: usage.output_tokens_this_month as u64,
spending_this_month: usage.spending_this_month as u64,
lifetime_spending: usage.lifetime_spending as u64,

8
crates/collab/src/llm/db.rs
View File

@@ -97,6 +97,14 @@ impl LlmDatabase {
.ok_or_else(|| anyhow!("unknown model {provider:?}:{name}"))?)
}
pub fn model_by_id(&self, id: ModelId) -> Result<&model::Model> {
Ok(self
.models
.values()
.find(|model| model.id == id)
.ok_or_else(|| anyhow!("no model for ID {id:?}"))?)
}
pub fn options(&self) -> &ConnectOptions {
&self.options
}

239
crates/collab/src/llm/db/queries/usages.rs
View File

@@ -1,5 +1,5 @@
use crate::db::UserId;
use chrono::Duration;
use chrono::{Datelike, Duration};
use futures::StreamExt as _;
use rpc::LanguageModelProvider;
use sea_orm::QuerySelect;
@@ -14,6 +14,8 @@ pub struct Usage {
pub tokens_this_minute: usize,
pub tokens_this_day: usize,
pub input_tokens_this_month: usize,
pub cache_creation_input_tokens_this_month: usize,
pub cache_read_input_tokens_this_month: usize,
pub output_tokens_this_month: usize,
pub spending_this_month: usize,
pub lifetime_spending: usize,
@@ -138,6 +140,46 @@ impl LlmDatabase {
.await
}
pub async fn get_user_spending_for_month(
&self,
user_id: UserId,
now: DateTimeUtc,
) -> Result<usize> {
self.transaction(|tx| async move {
let month = now.date_naive().month() as i32;
let year = now.date_naive().year();
let mut monthly_usages = monthly_usage::Entity::find()
.filter(
monthly_usage::Column::UserId
.eq(user_id)
.and(monthly_usage::Column::Month.eq(month))
.and(monthly_usage::Column::Year.eq(year)),
)
.stream(&*tx)
.await?;
let mut monthly_spending_in_cents = 0;
while let Some(usage) = monthly_usages.next().await {
let usage = usage?;
let Ok(model) = self.model_by_id(usage.model_id) else {
continue;
};
monthly_spending_in_cents += calculate_spending(
model,
usage.input_tokens as usize,
usage.cache_creation_input_tokens as usize,
usage.cache_read_input_tokens as usize,
usage.output_tokens as usize,
);
}
Ok(monthly_spending_in_cents)
})
.await
}
pub async fn get_usage(
&self,
user_id: UserId,
@@ -160,17 +202,26 @@ impl LlmDatabase {
.all(&*tx)
.await?;
let (lifetime_input_tokens, lifetime_output_tokens) = lifetime_usage::Entity::find()
let month = now.date_naive().month() as i32;
let year = now.date_naive().year();
let monthly_usage = monthly_usage::Entity::find()
.filter(
monthly_usage::Column::UserId
.eq(user_id)
.and(monthly_usage::Column::ModelId.eq(model.id))
.and(monthly_usage::Column::Month.eq(month))
.and(monthly_usage::Column::Year.eq(year)),
)
.one(&*tx)
.await?;
let lifetime_usage = lifetime_usage::Entity::find()
.filter(
lifetime_usage::Column::UserId
.eq(user_id)
.and(lifetime_usage::Column::ModelId.eq(model.id)),
)
.one(&*tx)
.await?
.map_or((0, 0), |usage| {
(usage.input_tokens as usize, usage.output_tokens as usize)
});
.await?;
let requests_this_minute =
self.get_usage_for_measure(&usages, now, UsageMeasure::RequestsPerMinute)?;
@@ -178,21 +229,45 @@ impl LlmDatabase {
self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerMinute)?;
let tokens_this_day =
self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerDay)?;
let input_tokens_this_month =
self.get_usage_for_measure(&usages, now, UsageMeasure::InputTokensPerMonth)?;
let output_tokens_this_month =
self.get_usage_for_measure(&usages, now, UsageMeasure::OutputTokensPerMonth)?;
let spending_this_month =
calculate_spending(model, input_tokens_this_month, output_tokens_this_month);
let lifetime_spending =
calculate_spending(model, lifetime_input_tokens, lifetime_output_tokens);
let spending_this_month = if let Some(monthly_usage) = &monthly_usage {
calculate_spending(
model,
monthly_usage.input_tokens as usize,
monthly_usage.cache_creation_input_tokens as usize,
monthly_usage.cache_read_input_tokens as usize,
monthly_usage.output_tokens as usize,
)
} else {
0
};
let lifetime_spending = if let Some(lifetime_usage) = &lifetime_usage {
calculate_spending(
model,
lifetime_usage.input_tokens as usize,
lifetime_usage.cache_creation_input_tokens as usize,
lifetime_usage.cache_read_input_tokens as usize,
lifetime_usage.output_tokens as usize,
)
} else {
0
};
Ok(Usage {
requests_this_minute,
tokens_this_minute,
tokens_this_day,
input_tokens_this_month,
output_tokens_this_month,
input_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.input_tokens as usize),
cache_creation_input_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.cache_creation_input_tokens as usize),
cache_read_input_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.cache_read_input_tokens as usize),
output_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.output_tokens as usize),
spending_this_month,
lifetime_spending,
})
@@ -208,6 +283,8 @@ impl LlmDatabase {
provider: LanguageModelProvider,
model_name: &str,
input_token_count: usize,
cache_creation_input_tokens: usize,
cache_read_input_tokens: usize,
output_token_count: usize,
now: DateTimeUtc,
) -> Result<Usage> {
@@ -235,6 +312,10 @@ impl LlmDatabase {
&tx,
)
.await?;
let total_token_count = input_token_count
+ cache_read_input_tokens
+ cache_creation_input_tokens
+ output_token_count;
let tokens_this_minute = self
.update_usage_for_measure(
user_id,
@@ -243,7 +324,7 @@ impl LlmDatabase {
&usages,
UsageMeasure::TokensPerMinute,
now,
input_token_count + output_token_count,
total_token_count,
&tx,
)
.await?;
@@ -255,36 +336,73 @@ impl LlmDatabase {
&usages,
UsageMeasure::TokensPerDay,
now,
input_token_count + output_token_count,
total_token_count,
&tx,
)
.await?;
let input_tokens_this_month = self
.update_usage_for_measure(
user_id,
is_staff,
model.id,
&usages,
UsageMeasure::InputTokensPerMonth,
now,
input_token_count,
&tx,
let month = now.date_naive().month() as i32;
let year = now.date_naive().year();
// Update monthly usage
let monthly_usage = monthly_usage::Entity::find()
.filter(
monthly_usage::Column::UserId
.eq(user_id)
.and(monthly_usage::Column::ModelId.eq(model.id))
.and(monthly_usage::Column::Month.eq(month))
.and(monthly_usage::Column::Year.eq(year)),
)
.one(&*tx)
.await?;
let output_tokens_this_month = self
.update_usage_for_measure(
user_id,
is_staff,
model.id,
&usages,
UsageMeasure::OutputTokensPerMonth,
now,
output_token_count,
&tx,
)
.await?;
let spending_this_month =
calculate_spending(model, input_tokens_this_month, output_tokens_this_month);
let monthly_usage = match monthly_usage {
Some(usage) => {
monthly_usage::Entity::update(monthly_usage::ActiveModel {
id: ActiveValue::unchanged(usage.id),
input_tokens: ActiveValue::set(
usage.input_tokens + input_token_count as i64,
),
cache_creation_input_tokens: ActiveValue::set(
usage.cache_creation_input_tokens + cache_creation_input_tokens as i64,
),
cache_read_input_tokens: ActiveValue::set(
usage.cache_read_input_tokens + cache_read_input_tokens as i64,
),
output_tokens: ActiveValue::set(
usage.output_tokens + output_token_count as i64,
),
..Default::default()
})
.exec(&*tx)
.await?
}
None => {
monthly_usage::ActiveModel {
user_id: ActiveValue::set(user_id),
model_id: ActiveValue::set(model.id),
month: ActiveValue::set(month),
year: ActiveValue::set(year),
input_tokens: ActiveValue::set(input_token_count as i64),
cache_creation_input_tokens: ActiveValue::set(
cache_creation_input_tokens as i64,
),
cache_read_input_tokens: ActiveValue::set(cache_read_input_tokens as i64),
output_tokens: ActiveValue::set(output_token_count as i64),
..Default::default()
}
.insert(&*tx)
.await?
}
};
let spending_this_month = calculate_spending(
model,
monthly_usage.input_tokens as usize,
monthly_usage.cache_creation_input_tokens as usize,
monthly_usage.cache_read_input_tokens as usize,
monthly_usage.output_tokens as usize,
);
// Update lifetime usage
let lifetime_usage = lifetime_usage::Entity::find()
@@ -303,6 +421,12 @@ impl LlmDatabase {
input_tokens: ActiveValue::set(
usage.input_tokens + input_token_count as i64,
),
cache_creation_input_tokens: ActiveValue::set(
usage.cache_creation_input_tokens + cache_creation_input_tokens as i64,
),
cache_read_input_tokens: ActiveValue::set(
usage.cache_read_input_tokens + cache_read_input_tokens as i64,
),
output_tokens: ActiveValue::set(
usage.output_tokens + output_token_count as i64,
),
@@ -316,6 +440,10 @@ impl LlmDatabase {
user_id: ActiveValue::set(user_id),
model_id: ActiveValue::set(model.id),
input_tokens: ActiveValue::set(input_token_count as i64),
cache_creation_input_tokens: ActiveValue::set(
cache_creation_input_tokens as i64,
),
cache_read_input_tokens: ActiveValue::set(cache_read_input_tokens as i64),
output_tokens: ActiveValue::set(output_token_count as i64),
..Default::default()
}
@@ -327,6 +455,8 @@ impl LlmDatabase {
let lifetime_spending = calculate_spending(
model,
lifetime_usage.input_tokens as usize,
lifetime_usage.cache_creation_input_tokens as usize,
lifetime_usage.cache_read_input_tokens as usize,
lifetime_usage.output_tokens as usize,
);
@@ -334,8 +464,11 @@ impl LlmDatabase {
requests_this_minute,
tokens_this_minute,
tokens_this_day,
input_tokens_this_month,
output_tokens_this_month,
input_tokens_this_month: monthly_usage.input_tokens as usize,
cache_creation_input_tokens_this_month: monthly_usage.cache_creation_input_tokens
as usize,
cache_read_input_tokens_this_month: monthly_usage.cache_read_input_tokens as usize,
output_tokens_this_month: monthly_usage.output_tokens as usize,
spending_this_month,
lifetime_spending,
})
@@ -501,18 +634,28 @@ impl LlmDatabase {
fn calculate_spending(
model: &model::Model,
input_tokens_this_month: usize,
cache_creation_input_tokens_this_month: usize,
cache_read_input_tokens_this_month: usize,
output_tokens_this_month: usize,
) -> usize {
let input_token_cost =
input_tokens_this_month * model.price_per_million_input_tokens as usize / 1_000_000;
let cache_creation_input_token_cost = cache_creation_input_tokens_this_month
* model.price_per_million_cache_creation_input_tokens as usize
/ 1_000_000;
let cache_read_input_token_cost = cache_read_input_tokens_this_month
* model.price_per_million_cache_read_input_tokens as usize
/ 1_000_000;
let output_token_cost =
output_tokens_this_month * model.price_per_million_output_tokens as usize / 1_000_000;
input_token_cost + output_token_cost
input_token_cost
+ cache_creation_input_token_cost
+ cache_read_input_token_cost
+ output_token_cost
}
const MINUTE_BUCKET_COUNT: usize = 12;
const DAY_BUCKET_COUNT: usize = 48;
const MONTH_BUCKET_COUNT: usize = 30;
impl UsageMeasure {
fn bucket_count(&self) -> usize {
@@ -520,8 +663,6 @@ impl UsageMeasure {
UsageMeasure::RequestsPerMinute => MINUTE_BUCKET_COUNT,
UsageMeasure::TokensPerMinute => MINUTE_BUCKET_COUNT,
UsageMeasure::TokensPerDay => DAY_BUCKET_COUNT,
UsageMeasure::InputTokensPerMonth => MONTH_BUCKET_COUNT,
UsageMeasure::OutputTokensPerMonth => MONTH_BUCKET_COUNT,
}
}
@@ -530,8 +671,6 @@ impl UsageMeasure {
UsageMeasure::RequestsPerMinute => Duration::minutes(1),
UsageMeasure::TokensPerMinute => Duration::minutes(1),
UsageMeasure::TokensPerDay => Duration::hours(24),
UsageMeasure::InputTokensPerMonth => Duration::days(30),
UsageMeasure::OutputTokensPerMonth => Duration::days(30),
}
}

1
crates/collab/src/llm/db/tables.rs
View File

@@ -1,5 +1,6 @@
pub mod lifetime_usage;
pub mod model;
pub mod monthly_usage;
pub mod provider;
pub mod revoked_access_token;
pub mod usage;

2
crates/collab/src/llm/db/tables/lifetime_usage.rs
View File

@@ -9,6 +9,8 @@ pub struct Model {
pub user_id: UserId,
pub model_id: ModelId,
pub input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub cache_read_input_tokens: i64,
pub output_tokens: i64,
}

2
crates/collab/src/llm/db/tables/model.rs
View File

@@ -14,6 +14,8 @@ pub struct Model {
pub max_tokens_per_minute: i64,
pub max_tokens_per_day: i64,
pub price_per_million_input_tokens: i32,
pub price_per_million_cache_creation_input_tokens: i32,
pub price_per_million_cache_read_input_tokens: i32,
pub price_per_million_output_tokens: i32,
}

22
crates/collab/src/llm/db/tables/monthly_usage.rs Normal file
View File

@@ -0,0 +1,22 @@
use crate::{db::UserId, llm::db::ModelId};
use sea_orm::entity::prelude::*;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
#[sea_orm(table_name = "monthly_usages")]
pub struct Model {
#[sea_orm(primary_key)]
pub id: i32,
pub user_id: UserId,
pub model_id: ModelId,
pub month: i32,
pub year: i32,
pub input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub cache_read_input_tokens: i64,
pub output_tokens: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}

2
crates/collab/src/llm/db/tables/usage_measure.rs
View File

@@ -9,8 +9,6 @@ pub enum UsageMeasure {
RequestsPerMinute,
TokensPerMinute,
TokensPerDay,
InputTokensPerMonth,
OutputTokensPerMonth,
}
#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]

68
crates/collab/src/llm/db/tests/usage_tests.rs
View File

@@ -6,7 +6,7 @@ use crate::{
},
test_llm_db,
};
use chrono::{Duration, Utc};
use chrono::{DateTime, Duration, Utc};
use pretty_assertions::assert_eq;
use rpc::LanguageModelProvider;
@@ -29,16 +29,19 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
.await
.unwrap();
let t0 = Utc::now();
// We're using a fixed datetime to prevent flakiness based on the clock.
let t0 = DateTime::parse_from_rfc3339("2024-08-08T22:46:33Z")
.unwrap()
.with_timezone(&Utc);
let user_id = UserId::from_proto(123);
let now = t0;
db.record_usage(user_id, false, provider, model, 1000, 0, now)
db.record_usage(user_id, false, provider, model, 1000, 0, 0, 0, now)
.await
.unwrap();
let now = t0 + Duration::seconds(10);
db.record_usage(user_id, false, provider, model, 2000, 0, now)
db.record_usage(user_id, false, provider, model, 2000, 0, 0, 0, now)
.await
.unwrap();
@@ -50,6 +53,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 3000,
tokens_this_day: 3000,
input_tokens_this_month: 3000,
cache_creation_input_tokens_this_month: 0,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
@@ -65,6 +70,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 2000,
tokens_this_day: 3000,
input_tokens_this_month: 3000,
cache_creation_input_tokens_this_month: 0,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
@@ -72,7 +79,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
);
let now = t0 + Duration::seconds(60);
db.record_usage(user_id, false, provider, model, 3000, 0, now)
db.record_usage(user_id, false, provider, model, 3000, 0, 0, 0, now)
.await
.unwrap();
@@ -84,6 +91,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 5000,
tokens_this_day: 6000,
input_tokens_this_month: 6000,
cache_creation_input_tokens_this_month: 0,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
@@ -100,13 +109,15 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 0,
tokens_this_day: 5000,
input_tokens_this_month: 6000,
cache_creation_input_tokens_this_month: 0,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
}
);
db.record_usage(user_id, false, provider, model, 4000, 0, now)
db.record_usage(user_id, false, provider, model, 4000, 0, 0, 0, now)
.await
.unwrap();
@@ -118,22 +129,55 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 4000,
tokens_this_day: 9000,
input_tokens_this_month: 10000,
cache_creation_input_tokens_this_month: 0,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
}
);
let t2 = t0 + Duration::days(30);
let now = t2;
// We're using a fixed datetime to prevent flakiness based on the clock.
let now = DateTime::parse_from_rfc3339("2024-10-08T22:15:58Z")
.unwrap()
.with_timezone(&Utc);
// Test cache creation input tokens
db.record_usage(user_id, false, provider, model, 1000, 500, 0, 0, now)
.await
.unwrap();
let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
assert_eq!(
usage,
Usage {
requests_this_minute: 0,
tokens_this_minute: 0,
tokens_this_day: 0,
input_tokens_this_month: 9000,
requests_this_minute: 1,
tokens_this_minute: 1500,
tokens_this_day: 1500,
input_tokens_this_month: 1000,
cache_creation_input_tokens_this_month: 500,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
}
);
// Test cache read input tokens
db.record_usage(user_id, false, provider, model, 1000, 0, 300, 0, now)
.await
.unwrap();
let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
assert_eq!(
usage,
Usage {
requests_this_minute: 2,
tokens_this_minute: 2800,
tokens_this_day: 2800,
input_tokens_this_month: 2000,
cache_creation_input_tokens_this_month: 500,
cache_read_input_tokens_this_month: 300,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,

4
crates/collab/src/llm/telemetry.rs
View File

@@ -12,11 +12,15 @@ pub struct LlmUsageEventRow {
pub model: String,
pub provider: String,
pub input_token_count: u64,
pub cache_creation_input_token_count: u64,
pub cache_read_input_token_count: u64,
pub output_token_count: u64,
pub requests_this_minute: u64,
pub tokens_this_minute: u64,
pub tokens_this_day: u64,
pub input_tokens_this_month: u64,
pub cache_creation_input_tokens_this_month: u64,
pub cache_read_input_tokens_this_month: u64,
pub output_tokens_this_month: u64,
pub spending_this_month: u64,
pub lifetime_spending: u64,

12
crates/collab/src/llm/token.rs
View File

@@ -13,15 +13,15 @@ pub struct LlmTokenClaims {
pub exp: u64,
pub jti: String,
pub user_id: u64,
pub github_user_login: String,
pub is_staff: bool,
pub has_llm_closed_beta_feature_flag: bool,
// This field is temporarily optional so it can be added
// in a backwards-compatible way. We can make it required
// once all of the LLM tokens have cycled (~1 hour after
// this change has been deployed).
#[serde(default)]
pub github_user_login: Option<String>,
pub is_staff: bool,
#[serde(default)]
pub has_llm_closed_beta_feature_flag: bool,
pub has_llm_subscription: Option<bool>,
pub plan: rpc::proto::Plan,
}
@@ -33,6 +33,7 @@ impl LlmTokenClaims {
github_user_login: String,
is_staff: bool,
has_llm_closed_beta_feature_flag: bool,
has_llm_subscription: bool,
plan: rpc::proto::Plan,
config: &Config,
) -> Result<String> {
@@ -47,9 +48,10 @@ impl LlmTokenClaims {
exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
jti: uuid::Uuid::new_v4().to_string(),
user_id: user_id.to_proto(),
github_user_login: Some(github_user_login),
github_user_login,
is_staff,
has_llm_closed_beta_feature_flag,
has_llm_subscription: Some(has_llm_subscription),
plan,
};

25
crates/collab/src/main.rs
View File

@@ -6,6 +6,7 @@ use axum::{
routing::get,
Extension, Router,
};
use collab::api::billing::sync_llm_usage_with_stripe_periodically;
use collab::api::CloudflareIpCountryHeader;
use collab::llm::{db::LlmDatabase, log_usage_periodically};
use collab::migrations::run_database_migrations;
@@ -29,7 +30,7 @@ use tower_http::trace::TraceLayer;
use tracing_subscriber::{
filter::EnvFilter, fmt::format::JsonFields, util::SubscriberInitExt, Layer,
};
use util::ResultExt as _;
use util::{maybe, ResultExt as _};
const VERSION: &str = env!("CARGO_PKG_VERSION");
const REVISION: Option<&'static str> = option_env!("GITHUB_SHA");
@@ -136,6 +137,28 @@ async fn main() -> Result<()> {
fetch_extensions_from_blob_store_periodically(state.clone());
spawn_user_backfiller(state.clone());
let llm_db = maybe!(async {
let database_url = state
.config
.llm_database_url
.as_ref()
.ok_or_else(|| anyhow!("missing LLM_DATABASE_URL"))?;
let max_connections = state
.config
.llm_database_max_connections
.ok_or_else(|| anyhow!("missing LLM_DATABASE_MAX_CONNECTIONS"))?;
let mut db_options = db::ConnectOptions::new(database_url);
db_options.max_connections(max_connections);
LlmDatabase::new(db_options, state.executor.clone()).await
})
.await
.trace_err();
if let Some(llm_db) = llm_db {
sync_llm_usage_with_stripe_periodically(state.clone(), llm_db);
}
app = app
.merge(collab::api::events::router())
.merge(collab::api::extensions::router())

33
crates/collab/src/rpc.rs
View File

@@ -36,8 +36,8 @@ use collections::{HashMap, HashSet};
pub use connection_pool::{ConnectionPool, ZedVersion};
use core::fmt::{self, Debug, Formatter};
use http_client::HttpClient;
use isahc_http_client::IsahcHttpClient;
use open_ai::{OpenAiEmbeddingModel, OPEN_AI_API_URL};
use reqwest_client::ReqwestClient;
use sha2::Digest;
use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi};
@@ -191,16 +191,26 @@ impl Session {
}
}
pub async fn current_plan(&self, db: MutexGuard<'_, DbHandle>) -> anyhow::Result<proto::Plan> {
pub async fn has_llm_subscription(
&self,
db: &MutexGuard<'_, DbHandle>,
) -> anyhow::Result<bool> {
if self.is_staff() {
return Ok(proto::Plan::ZedPro);
return Ok(true);
}
let Some(user_id) = self.user_id() else {
return Ok(proto::Plan::Free);
return Ok(false);
};
if db.has_active_billing_subscription(user_id).await? {
Ok(db.has_active_billing_subscription(user_id).await?)
}
pub async fn current_plan(
&self,
_db: &MutexGuard<'_, DbHandle>,
) -> anyhow::Result<proto::Plan> {
if self.is_staff() {
Ok(proto::Plan::ZedPro)
} else {
Ok(proto::Plan::Free)
@@ -954,8 +964,8 @@ impl Server {
tracing::info!("connection opened");
let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));
let http_client = match ReqwestClient::user_agent(&user_agent) {
Ok(http_client) => Arc::new(http_client),
let http_client = match IsahcHttpClient::builder().default_header("User-Agent", user_agent).build() {
Ok(http_client) => Arc::new(IsahcHttpClient::from(http_client)),
Err(error) => {
tracing::error!(?error, "failed to create HTTP client");
return;
@@ -3471,7 +3481,7 @@ fn should_auto_subscribe_to_channels(version: ZedVersion) -> bool {
}
async fn update_user_plan(_user_id: UserId, session: &Session) -> Result<()> {
let plan = session.current_plan(session.db().await).await?;
let plan = session.current_plan(&session.db().await).await?;
session
.peer
@@ -4471,7 +4481,7 @@ async fn count_language_model_tokens(
};
authorize_access_to_legacy_llm_endpoints(&session).await?;
let rate_limit: Box<dyn RateLimit> = match session.current_plan(session.db().await).await? {
let rate_limit: Box<dyn RateLimit> = match session.current_plan(&session.db().await).await? {
proto::Plan::ZedPro => Box::new(ZedProCountLanguageModelTokensRateLimit),
proto::Plan::Free => Box::new(FreeCountLanguageModelTokensRateLimit),
};
@@ -4592,7 +4602,7 @@ async fn compute_embeddings(
let api_key = api_key.context("no OpenAI API key configured on the server")?;
authorize_access_to_legacy_llm_endpoints(&session).await?;
let rate_limit: Box<dyn RateLimit> = match session.current_plan(session.db().await).await? {
let rate_limit: Box<dyn RateLimit> = match session.current_plan(&session.db().await).await? {
proto::Plan::ZedPro => Box::new(ZedProComputeEmbeddingsRateLimit),
proto::Plan::Free => Box::new(FreeComputeEmbeddingsRateLimit),
};
@@ -4915,7 +4925,8 @@ async fn get_llm_api_token(
user.github_login.clone(),
session.is_staff(),
has_llm_closed_beta_feature_flag,
session.current_plan(db).await?,
session.has_llm_subscription(&db).await?,
session.current_plan(&db).await?,
&session.app_state.config,
)?;
response.send(proto::GetLlmTokenResponse { token })?;

2
crates/collab/src/tests/test_server.rs
View File

@@ -677,7 +677,7 @@ impl TestServer {
migrations_path: None,
seed_path: None,
stripe_api_key: None,
stripe_price_id: None,
stripe_llm_usage_price_id: None,
supermaven_admin_api_key: None,
user_backfiller_github_access_token: None,
},

27
crates/editor/src/editor.rs
View File

@@ -6466,14 +6466,22 @@ impl Editor {
fn apply_selected_diff_hunks(&mut self, _: &ApplyDiffHunk, cx: &mut ViewContext<Self>) {
let snapshot = self.buffer.read(cx).snapshot(cx);
let hunks = hunks_for_selections(&snapshot, &self.selections.disjoint_anchors());
let mut ranges_by_buffer = HashMap::default();
self.transact(cx, |editor, cx| {
for hunk in hunks {
if let Some(buffer) = editor.buffer.read(cx).buffer(hunk.buffer_id) {
buffer.update(cx, |buffer, cx| {
buffer.merge_into_base(Some(hunk.buffer_range.to_offset(buffer)), cx);
});
ranges_by_buffer
.entry(buffer.clone())
.or_insert_with(Vec::new)
.push(hunk.buffer_range.to_offset(buffer.read(cx)));
}
}
for (buffer, ranges) in ranges_by_buffer {
buffer.update(cx, |buffer, cx| {
buffer.merge_into_base(ranges, cx);
});
}
});
}
@@ -12538,9 +12546,14 @@ impl Editor {
}
let Some(project) = &self.project else { return };
let telemetry = project.read(cx).client().telemetry().clone();
let (telemetry, is_via_ssh) = {
let project = project.read(cx);
let telemetry = project.client().telemetry().clone();
let is_via_ssh = project.is_via_ssh();
(telemetry, is_via_ssh)
};
refresh_linked_ranges(self, cx);
telemetry.log_edit_event("editor");
telemetry.log_edit_event("editor", is_via_ssh);
}
multi_buffer::Event::ExcerptsAdded {
buffer,
@@ -12881,13 +12894,15 @@ impl Editor {
.settings_at(0, cx)
.show_inline_completions;
let telemetry = project.read(cx).client().telemetry().clone();
let project = project.read(cx);
let telemetry = project.client().telemetry().clone();
telemetry.report_editor_event(
file_extension,
vim_mode,
operation,
copilot_enabled,
copilot_enabled_for_language,
project.is_via_ssh(),
)
}

54
crates/editor/src/editor_tests.rs
View File

@@ -11715,6 +11715,60 @@ async fn test_toggle_diff_expand_in_multi_buffer(cx: &mut gpui::TestAppContext)
);
}
#[gpui::test]
async fn test_expand_diff_hunk_at_excerpt_boundary(cx: &mut gpui::TestAppContext) {
init_test(cx, |_| {});
let base = "aaa\nbbb\nccc\nddd\neee\nfff\nggg\n";
let text = "aaa\nBBB\nBB2\nccc\nDDD\nEEE\nfff\nggg\n";
let buffer = cx.new_model(|cx| {
let mut buffer = Buffer::local(text.to_string(), cx);
buffer.set_diff_base(Some(base.into()), cx);
buffer
});
let multi_buffer = cx.new_model(|cx| {
let mut multibuffer = MultiBuffer::new(ReadWrite);
multibuffer.push_excerpts(
buffer.clone(),
[
ExcerptRange {
context: Point::new(0, 0)..Point::new(2, 0),
primary: None,
},
ExcerptRange {
context: Point::new(5, 0)..Point::new(7, 0),
primary: None,
},
],
cx,
);
multibuffer
});
let editor = cx.add_window(|cx| Editor::new(EditorMode::Full, multi_buffer, None, true, cx));
let mut cx = EditorTestContext::for_editor(editor, cx).await;
cx.run_until_parked();
cx.update_editor(|editor, cx| editor.expand_all_hunk_diffs(&Default::default(), cx));
cx.executor().run_until_parked();
cx.assert_diff_hunks(
"
aaa
- bbb
+ BBB
- ddd
- eee
+ EEE
fff
"
.unindent(),
);
}
#[gpui::test]
async fn test_edits_around_expanded_insertion_hunks(
executor: BackgroundExecutor,

88
crates/editor/src/hunk_diff.rs
View File

@@ -209,19 +209,20 @@ impl Editor {
retain
});
for remaining_hunk in hunks_to_toggle {
let remaining_hunk_point_range =
Point::new(remaining_hunk.row_range.start.0, 0)
..Point::new(remaining_hunk.row_range.end.0, 0);
for hunk in hunks_to_toggle {
let remaining_hunk_point_range = Point::new(hunk.row_range.start.0, 0)
..Point::new(hunk.row_range.end.0, 0);
let hunk_start = snapshot
.buffer_snapshot
.anchor_before(remaining_hunk_point_range.start);
let hunk_end = snapshot
.buffer_snapshot
.anchor_in_excerpt(hunk_start.excerpt_id, hunk.buffer_range.end)
.unwrap();
hunks_to_expand.push(HoveredHunk {
status: hunk_status(&remaining_hunk),
multi_buffer_range: snapshot
.buffer_snapshot
.anchor_before(remaining_hunk_point_range.start)
..snapshot
.buffer_snapshot
.anchor_after(remaining_hunk_point_range.end),
diff_base_byte_range: remaining_hunk.diff_base_byte_range.clone(),
status: hunk_status(&hunk),
multi_buffer_range: hunk_start..hunk_end,
diff_base_byte_range: hunk.diff_base_byte_range.clone(),
});
}
@@ -246,33 +247,22 @@ impl Editor {
hunk: &HoveredHunk,
cx: &mut ViewContext<'_, Editor>,
) -> Option<()> {
let multi_buffer_snapshot = self.buffer().read(cx).snapshot(cx);
let buffer = self.buffer.clone();
let multi_buffer_snapshot = buffer.read(cx).snapshot(cx);
let hunk_range = hunk.multi_buffer_range.clone();
let hunk_point_range = hunk_range.to_point(&multi_buffer_snapshot);
let buffer = self.buffer().clone();
let snapshot = self.snapshot(cx);
let (diff_base_buffer, deleted_text_lines) = buffer.update(cx, |buffer, cx| {
let hunk = buffer_diff_hunk(&snapshot.buffer_snapshot, hunk_point_range.clone())?;
let mut buffer_ranges = buffer.range_to_buffer_ranges(hunk_point_range, cx);
if buffer_ranges.len() == 1 {
let (buffer, _, _) = buffer_ranges.pop()?;
let diff_base_buffer = diff_base_buffer
.or_else(|| self.current_diff_base_buffer(&buffer, cx))
.or_else(|| create_diff_base_buffer(&buffer, cx))?;
let buffer = buffer.read(cx);
let deleted_text_lines = buffer.diff_base().map(|diff_base| {
let diff_start_row = diff_base
.offset_to_point(hunk.diff_base_byte_range.start)
.row;
let diff_end_row = diff_base.offset_to_point(hunk.diff_base_byte_range.end).row;
diff_end_row - diff_start_row
})?;
Some((diff_base_buffer, deleted_text_lines))
} else {
None
}
let buffer = buffer.buffer(hunk_range.start.buffer_id?)?;
let diff_base_buffer = diff_base_buffer
.or_else(|| self.current_diff_base_buffer(&buffer, cx))
.or_else(|| create_diff_base_buffer(&buffer, cx))?;
let deleted_text_lines = buffer.read(cx).diff_base().map(|diff_base| {
let diff_start_row = diff_base
.offset_to_point(hunk.diff_base_byte_range.start)
.row;
let diff_end_row = diff_base.offset_to_point(hunk.diff_base_byte_range.end).row;
diff_end_row - diff_start_row
})?;
Some((diff_base_buffer, deleted_text_lines))
})?;
let block_insert_index = match self.expanded_hunks.hunks.binary_search_by(|probe| {
@@ -350,7 +340,7 @@ impl Editor {
.next()?;
buffer.update(cx, |branch_buffer, cx| {
branch_buffer.merge_into_base(Some(range), cx);
branch_buffer.merge_into_base(vec![range], cx);
});
None
@@ -360,7 +350,7 @@ impl Editor {
let buffers = self.buffer.read(cx).all_buffers();
for branch_buffer in buffers {
branch_buffer.update(cx, |branch_buffer, cx| {
branch_buffer.merge_into_base(None, cx);
branch_buffer.merge_into_base(Vec::new(), cx);
});
}
}
@@ -1128,21 +1118,6 @@ fn editor_with_deleted_text(
(editor_height, editor)
}
fn buffer_diff_hunk(
buffer_snapshot: &MultiBufferSnapshot,
row_range: Range<Point>,
) -> Option<MultiBufferDiffHunk> {
let mut hunks = buffer_snapshot.git_diff_hunks_in_range(
MultiBufferRow(row_range.start.row)..MultiBufferRow(row_range.end.row),
);
let hunk = hunks.next()?;
let second_hunk = hunks.next();
if second_hunk.is_none() {
return Some(hunk);
}
None
}
impl DisplayDiffHunk {
pub fn start_display_row(&self) -> DisplayRow {
match self {
@@ -1209,7 +1184,10 @@ pub fn diff_hunk_to_display(
let hunk_end_point = Point::new(hunk_end_row.0, 0);
let multi_buffer_start = snapshot.buffer_snapshot.anchor_before(hunk_start_point);
let multi_buffer_end = snapshot.buffer_snapshot.anchor_after(hunk_end_point);
let multi_buffer_end = snapshot
.buffer_snapshot
.anchor_in_excerpt(multi_buffer_start.excerpt_id, hunk.buffer_range.end)
.unwrap();
let end = hunk_end_point.to_display_point(snapshot).row();
DisplayDiffHunk::Unfolded {

2
crates/evals/Cargo.toml
View File

@@ -25,6 +25,7 @@ fs.workspace = true
git.workspace = true
gpui.workspace = true
http_client.workspace = true
isahc_http_client.workspace = true
language.workspace = true
languages.workspace = true
node_runtime.workspace = true
@@ -35,4 +36,3 @@ serde.workspace = true
serde_json.workspace = true
settings.workspace = true
smol.workspace = true
ureq_client.workspace = true

7
crates/evals/src/eval.rs
View File

@@ -32,7 +32,6 @@ use std::{
Arc,
},
};
use ureq_client::UreqClient;
const CODESEARCH_NET_DIR: &'static str = "target/datasets/code-search-net";
const EVAL_REPOS_DIR: &'static str = "target/datasets/eval-repos";
@@ -101,11 +100,7 @@ fn main() -> Result<()> {
gpui::App::headless().run(move |cx| {
let executor = cx.background_executor().clone();
let client = Arc::new(UreqClient::new(
None,
"Zed LLM evals".to_string(),
executor.clone(),
));
let client = isahc_http_client::IsahcHttpClient::new(None, None);
cx.set_http_client(client.clone());
match cli.command {
Commands::Fetch {} => {

3
crates/extension/Cargo.toml
View File

@@ -56,6 +56,7 @@ wit-component.workspace = true
workspace.workspace = true
[dev-dependencies]
isahc_http_client.workspace = true
ctor.workspace = true
env_logger.workspace = true
fs = { workspace = true, features = ["test-support"] }
@@ -63,7 +64,5 @@ gpui = { workspace = true, features = ["test-support"] }
language = { workspace = true, features = ["test-support"] }
parking_lot.workspace = true
project = { workspace = true, features = ["test-support"] }
reqwest_client.workspace = true
tokio.workspace = true
ureq_client.workspace = true
workspace = { workspace = true, features = ["test-support"] }

2
crates/extension/src/extension_builder.rs
View File

@@ -25,7 +25,7 @@ use wit_component::ComponentEncoder;
/// Once Rust 1.78 is released, there will be a `wasm32-wasip2` target available, so we will
/// not need the adapter anymore.
const RUST_TARGET: &str = "wasm32-wasip1";
pub const WASI_ADAPTER_URL: &str =
const WASI_ADAPTER_URL: &str =
"https://github.com/bytecodealliance/wasmtime/releases/download/v18.0.2/wasi_snapshot_preview1.reactor.wasm";
/// Compiling Tree-sitter parsers from C to WASM requires Clang 17, and a WASM build of libc

52
crates/extension/src/extension_store_test.rs
View File

@@ -1,4 +1,3 @@
use crate::extension_builder::WASI_ADAPTER_URL;
use crate::extension_manifest::SchemaVersion;
use crate::extension_settings::ExtensionSettings;
use crate::{
@@ -12,14 +11,14 @@ use collections::BTreeMap;
use fs::{FakeFs, Fs, RealFs};
use futures::{io::BufReader, AsyncReadExt, StreamExt};
use gpui::{Context, SemanticVersion, TestAppContext};
use http_client::{AsyncBody, FakeHttpClient, HttpClient, Response};
use http_client::{FakeHttpClient, Response};
use indexed_docs::IndexedDocsRegistry;
use isahc_http_client::IsahcHttpClient;
use language::{LanguageMatcher, LanguageRegistry, LanguageServerBinaryStatus, LanguageServerName};
use node_runtime::NodeRuntime;
use parking_lot::Mutex;
use project::{Project, DEFAULT_COMPLETION_CONTEXT};
use release_channel::AppVersion;
use reqwest_client::ReqwestClient;
use serde_json::json;
use settings::{Settings as _, SettingsStore};
use snippet_provider::SnippetRegistry;
@@ -29,7 +28,6 @@ use std::{
sync::Arc,
};
use theme::ThemeRegistry;
use ureq_client::UreqClient;
use util::test::temp_tree;
#[cfg(test)]
@@ -578,7 +576,7 @@ async fn test_extension_store_with_test_extension(cx: &mut TestAppContext) {
std::env::consts::ARCH
)
});
let builder_client = Arc::new(UreqClient::new(None, user_agent, cx.executor().clone()));
let builder_client = IsahcHttpClient::new(None, Some(user_agent));
let extension_store = cx.new_model(|cx| {
ExtensionStore::new(
@@ -771,50 +769,6 @@ async fn test_extension_store_with_test_extension(cx: &mut TestAppContext) {
assert!(fs.metadata(&expected_server_path).await.unwrap().is_none());
}
#[gpui::test]
async fn test_wasi_adapter_download(cx: &mut TestAppContext) {
let client = Arc::new(UreqClient::new(
None,
"zed-test-wasi-adapter-download".to_string(),
cx.executor().clone(),
));
let mut response = client
.get(WASI_ADAPTER_URL, AsyncBody::default(), true)
.await
.unwrap();
let mut content = Vec::new();
let mut body = BufReader::new(response.body_mut());
body.read_to_end(&mut content).await.unwrap();
assert!(wasmparser::Parser::is_core_wasm(&content));
assert_eq!(content.len(), 96801); // Determined by downloading this to my computer
wit_component::ComponentEncoder::default()
.adapter("wasi_snapshot_preview1", &content)
.unwrap();
}
#[tokio::test]
async fn test_wasi_adapter_download_tokio() {
let client = Arc::new(ReqwestClient::new());
let mut response = client
.get(WASI_ADAPTER_URL, AsyncBody::default(), true)
.await
.unwrap();
let mut content = Vec::new();
let mut body = BufReader::new(response.body_mut());
body.read_to_end(&mut content).await.unwrap();
assert!(wasmparser::Parser::is_core_wasm(&content));
assert_eq!(content.len(), 96801); // Determined by downloading this to my computer
wit_component::ComponentEncoder::default()
.adapter("wasi_snapshot_preview1", &content)
.unwrap();
}
fn init_test(cx: &mut TestAppContext) {
cx.update(|cx| {
let store = SettingsStore::test(cx);

2
crates/extension_cli/Cargo.toml
View File

@@ -18,9 +18,9 @@ clap = { workspace = true, features = ["derive"] }
env_logger.workspace = true
extension = { workspace = true, features = ["no-webrtc"] }
fs.workspace = true
isahc_http_client.workspace = true
language.workspace = true
log.workspace = true
reqwest_client.workspace = true
rpc.workspace = true
serde.workspace = true
serde_json.workspace = true

9
crates/extension_cli/src/main.rs
View File

@@ -13,8 +13,8 @@ use extension::{
extension_builder::{CompileExtensionOptions, ExtensionBuilder},
ExtensionManifest,
};
use isahc_http_client::IsahcHttpClient;
use language::LanguageConfig;
use reqwest_client::ReqwestClient;
use theme::ThemeRegistry;
use tree_sitter::{Language, Query, WasmStore};
@@ -66,7 +66,12 @@ async fn main() -> Result<()> {
std::env::consts::OS,
std::env::consts::ARCH
);
let http_client = Arc::new(ReqwestClient::user_agent(&user_agent)?);
let http_client = Arc::new(
IsahcHttpClient::builder()
.default_header("User-Agent", user_agent)
.build()
.map(IsahcHttpClient::from)?,
);
let builder = ExtensionBuilder::new(http_client, scratch_dir);
builder

1
crates/gpui/src/platform/mac/window.rs
View File

@@ -768,6 +768,7 @@ impl Drop for MacWindow {
unsafe {
this.native_window.setDelegate_(nil);
}
this.input_handler.take();
this.executor
.spawn(async move {
unsafe {

4
crates/http_client/Cargo.toml
View File

@@ -16,13 +16,11 @@ path = "src/http_client.rs"
doctest = true
[dependencies]
http = "0.2"
anyhow.workspace = true
derive_more.workspace = true
futures.workspace = true
http = "1.1"
log.workspace = true
rustls-native-certs.workspace = true
rustls.workspace = true
serde.workspace = true
serde_json.workspace = true
smol.workspace = true

29
crates/http_client/src/http_client.rs
View File

@@ -11,21 +11,13 @@ use http::request::Builder;
#[cfg(feature = "test-support")]
use std::fmt;
use std::{
sync::{Arc, LazyLock, Mutex},
sync::{Arc, Mutex},
time::Duration,
};
pub use url::Url;
#[derive(Clone)]
pub struct ReadTimeout(pub Duration);
impl Default for ReadTimeout {
fn default() -> Self {
Self(Duration::from_secs(5))
}
}
#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
#[derive(Default, Debug, Clone)]
pub enum RedirectPolicy {
#[default]
NoFollow,
@@ -34,23 +26,6 @@ pub enum RedirectPolicy {
}
pub struct FollowRedirects(pub bool);
pub static TLS_CONFIG: LazyLock<Arc<rustls::ClientConfig>> = LazyLock::new(|| {
let mut root_store = rustls::RootCertStore::empty();
let root_certs = rustls_native_certs::load_native_certs();
for error in root_certs.errors {
log::warn!("error loading native certs: {:?}", error);
}
root_store.add_parsable_certificates(&root_certs.certs);
Arc::new(
rustls::ClientConfig::builder()
.with_safe_defaults()
.with_root_certificates(root_store)
.with_no_client_auth(),
)
});
pub trait HttpRequestExt {
/// Set a read timeout on the request.
/// For isahc, this is the low_speed_timeout.

15
crates/ureq_client/Cargo.toml → crates/isahc_http_client/Cargo.toml
View File

@@ -1,5 +1,5 @@
[package]
name = "ureq_client"
name = "isahc_http_client"
version = "0.1.0"
edition = "2021"
publish = false
@@ -12,20 +12,11 @@ workspace = true
test-support = []
[lib]
path = "src/ureq_client.rs"
doctest = true
[[example]]
name = "client"
path = "examples/client.rs"
path = "src/isahc_http_client.rs"
[dependencies]
anyhow.workspace = true
futures.workspace = true
gpui.workspace = true
http_client.workspace = true
parking_lot.workspace = true
serde.workspace = true
smol.workspace = true
ureq = "=2.9.1"
isahc.workspace = true
util.workspace = true

1
crates/isahc_http_client/LICENSE-APACHE Symbolic link
View File

@@ -0,0 +1 @@
../../LICENSE-APACHE

105
crates/isahc_http_client/src/isahc_http_client.rs Normal file
View File

@@ -0,0 +1,105 @@
use std::{mem, sync::Arc, time::Duration};
use futures::future::BoxFuture;
use util::maybe;
pub use isahc::config::Configurable;
pub struct IsahcHttpClient(isahc::HttpClient);
pub use http_client::*;
impl IsahcHttpClient {
pub fn new(proxy: Option<Uri>, user_agent: Option<String>) -> Arc<IsahcHttpClient> {
let mut builder = isahc::HttpClient::builder()
.connect_timeout(Duration::from_secs(5))
.low_speed_timeout(100, Duration::from_secs(5))
.proxy(proxy.clone());
if let Some(agent) = user_agent {
builder = builder.default_header("User-Agent", agent);
}
Arc::new(IsahcHttpClient(builder.build().unwrap()))
}
pub fn builder() -> isahc::HttpClientBuilder {
isahc::HttpClientBuilder::new()
}
}
impl From<isahc::HttpClient> for IsahcHttpClient {
fn from(client: isahc::HttpClient) -> Self {
Self(client)
}
}
impl HttpClient for IsahcHttpClient {
fn proxy(&self) -> Option<&Uri> {
None
}
fn send(
&self,
req: http_client::http::Request<http_client::AsyncBody>,
) -> BoxFuture<'static, Result<http_client::Response<http_client::AsyncBody>, anyhow::Error>>
{
let redirect_policy = req
.extensions()
.get::<http_client::RedirectPolicy>()
.cloned()
.unwrap_or_default();
let read_timeout = req
.extensions()
.get::<http_client::ReadTimeout>()
.map(|t| t.0);
let req = maybe!({
let (mut parts, body) = req.into_parts();
let mut builder = isahc::Request::builder()
.method(parts.method)
.uri(parts.uri)
.version(parts.version);
if let Some(read_timeout) = read_timeout {
builder = builder.low_speed_timeout(100, read_timeout);
}
let headers = builder.headers_mut()?;
mem::swap(headers, &mut parts.headers);
let extensions = builder.extensions_mut()?;
mem::swap(extensions, &mut parts.extensions);
let isahc_body = match body.0 {
http_client::Inner::Empty => isahc::AsyncBody::empty(),
http_client::Inner::AsyncReader(reader) => isahc::AsyncBody::from_reader(reader),
http_client::Inner::SyncReader(reader) => {
isahc::AsyncBody::from_bytes_static(reader.into_inner())
}
};
builder
.redirect_policy(match redirect_policy {
http_client::RedirectPolicy::FollowAll => isahc::config::RedirectPolicy::Follow,
http_client::RedirectPolicy::FollowLimit(limit) => {
isahc::config::RedirectPolicy::Limit(limit)
}
http_client::RedirectPolicy::NoFollow => isahc::config::RedirectPolicy::None,
})
.body(isahc_body)
.ok()
});
let client = self.0.clone();
Box::pin(async move {
match req {
Some(req) => client
.send_async(req)
.await
.map_err(Into::into)
.map(|response| {
let (parts, body) = response.into_parts();
let body = http_client::AsyncBody::from_reader(body);
http_client::Response::from_parts(parts, body)
}),
None => Err(anyhow::anyhow!("Request was malformed")),
}
})
}
}

125
crates/language/src/buffer.rs
View File

@@ -18,6 +18,7 @@ use crate::{
};
use anyhow::{anyhow, Context, Result};
use async_watch as watch;
use clock::Lamport;
pub use clock::ReplicaId;
use futures::channel::oneshot;
use gpui::{
@@ -90,7 +91,7 @@ enum BufferDiffBase {
PastBufferVersion {
buffer: Model<Buffer>,
rope: Rope,
operations_to_ignore: Vec<clock::Lamport>,
merged_operations: Vec<Lamport>,
},
}
@@ -802,7 +803,7 @@ impl Buffer {
diff_base: Some(BufferDiffBase::PastBufferVersion {
buffer: this.clone(),
rope: self.as_rope().clone(),
operations_to_ignore: Vec::new(),
merged_operations: Default::default(),
}),
language: self.language.clone(),
has_conflict: self.has_conflict,
@@ -826,42 +827,60 @@ impl Buffer {
})
}
/// Applies all of the changes in this buffer that intersect the given `range`
/// to its base buffer. This buffer must be a branch buffer to call this method.
pub fn merge_into_base(&mut self, range: Option<Range<usize>>, cx: &mut ModelContext<Self>) {
/// Applies all of the changes in this buffer that intersect any of the
/// given `ranges` to its base buffer.
///
/// If `ranges` is empty, then all changes will be applied. This buffer must
/// be a branch buffer to call this method.
pub fn merge_into_base(&mut self, ranges: Vec<Range<usize>>, cx: &mut ModelContext<Self>) {
let Some(base_buffer) = self.diff_base_buffer() else {
debug_panic!("not a branch buffer");
return;
};
base_buffer.update(cx, |base_buffer, cx| {
let edits = self
.edits_since::<usize>(&base_buffer.version)
.filter_map(|edit| {
if range
.as_ref()
.map_or(true, |range| range.overlaps(&edit.new))
{
Some((edit.old, self.text_for_range(edit.new).collect::<String>()))
} else {
None
let mut ranges = if ranges.is_empty() {
&[0..usize::MAX]
} else {
ranges.as_slice()
}
.into_iter()
.peekable();
let mut edits = Vec::new();
for edit in self.edits_since::<usize>(&base_buffer.read(cx).version()) {
let mut is_included = false;
while let Some(range) = ranges.peek() {
if range.end < edit.new.start {
ranges.next().unwrap();
} else {
if range.start <= edit.new.end {
is_included = true;
}
})
.collect::<Vec<_>>();
let operation = base_buffer.edit(edits, None, cx);
// Prevent this operation from being reapplied to the branch.
if let Some(BufferDiffBase::PastBufferVersion {
operations_to_ignore,
..
}) = &mut self.diff_base
{
operations_to_ignore.extend(operation);
break;
}
}
if is_included {
edits.push((
edit.old.clone(),
self.text_for_range(edit.new.clone()).collect::<String>(),
));
}
}
let operation = base_buffer.update(cx, |base_buffer, cx| {
cx.emit(BufferEvent::DiffBaseChanged);
base_buffer.edit(edits, None, cx)
});
if let Some(operation) = operation {
if let Some(BufferDiffBase::PastBufferVersion {
merged_operations, ..
}) = &mut self.diff_base
{
merged_operations.push(operation);
}
}
}
fn on_base_buffer_event(
@@ -870,31 +889,34 @@ impl Buffer {
event: &BufferEvent,
cx: &mut ModelContext<Self>,
) {
if let BufferEvent::Operation { operation, .. } = event {
if let Some(BufferDiffBase::PastBufferVersion {
operations_to_ignore,
..
}) = &mut self.diff_base
{
let mut is_ignored = false;
if let Operation::Buffer(text::Operation::Edit(buffer_operation)) = &operation {
operations_to_ignore.retain(|operation_to_ignore| {
match buffer_operation.timestamp.cmp(&operation_to_ignore) {
Ordering::Less => true,
Ordering::Equal => {
is_ignored = true;
false
}
Ordering::Greater => false,
}
});
}
if !is_ignored {
self.apply_ops([operation.clone()], cx);
self.diff_base_version += 1;
}
let BufferEvent::Operation { operation, .. } = event else {
return;
};
let Some(BufferDiffBase::PastBufferVersion {
merged_operations, ..
}) = &mut self.diff_base
else {
return;
};
let mut operation_to_undo = None;
if let Operation::Buffer(text::Operation::Edit(operation)) = &operation {
if let Ok(ix) = merged_operations.binary_search(&operation.timestamp) {
merged_operations.remove(ix);
operation_to_undo = Some(operation.timestamp);
}
}
self.apply_ops([operation.clone()], cx);
if let Some(timestamp) = operation_to_undo {
let operation = self
.text
.undo_operations([(timestamp, u32::MAX)].into_iter().collect());
self.send_operation(Operation::Buffer(operation), true, cx);
}
self.diff_base_version += 1;
}
#[cfg(test)]
@@ -1134,7 +1156,6 @@ impl Buffer {
this.non_text_state_update_count += 1;
if let Some(BufferDiffBase::PastBufferVersion { rope, .. }) = &mut this.diff_base {
*rope = diff_base_rope;
cx.emit(BufferEvent::DiffBaseChanged);
}
cx.emit(BufferEvent::DiffUpdated);
})

68
crates/language/src/buffer_tests.rs
View File

@@ -2472,7 +2472,7 @@ fn test_branch_and_merge(cx: &mut TestAppContext) {
// Merging the branch applies all of its changes to the base.
branch_buffer.update(cx, |branch_buffer, cx| {
branch_buffer.merge_into_base(None, cx);
branch_buffer.merge_into_base(Vec::new(), cx);
});
branch_buffer.update(cx, |branch_buffer, cx| {
@@ -2485,15 +2485,73 @@ fn test_branch_and_merge(cx: &mut TestAppContext) {
}
#[gpui::test]
fn test_merge_into_base(cx: &mut AppContext) {
init_settings(cx, |_| {});
fn test_merge_into_base(cx: &mut TestAppContext) {
cx.update(|cx| init_settings(cx, |_| {}));
let base = cx.new_model(|cx| Buffer::local("abcdefghijk", cx));
let branch = base.update(cx, |buffer, cx| buffer.branch(cx));
// Make 3 edits, merge one into the base.
branch.update(cx, |branch, cx| {
branch.edit([(0..3, "ABC"), (7..9, "HI"), (11..11, "LMN")], None, cx);
branch.merge_into_base(vec![5..8], cx);
});
branch.read_with(cx, |branch, _| assert_eq!(branch.text(), "ABCdefgHIjkLMN"));
base.read_with(cx, |base, _| assert_eq!(base.text(), "abcdefgHIjk"));
// Undo the one already-merged edit. Merge that into the base.
branch.update(cx, |branch, cx| {
branch.edit([(7..9, "hi")], None, cx);
branch.merge_into_base(vec![5..8], cx);
});
base.read_with(cx, |base, _| assert_eq!(base.text(), "abcdefghijk"));
// Merge an insertion into the base.
branch.update(cx, |branch, cx| {
branch.merge_into_base(vec![11..11], cx);
});
branch.read_with(cx, |branch, _| assert_eq!(branch.text(), "ABCdefghijkLMN"));
base.read_with(cx, |base, _| assert_eq!(base.text(), "abcdefghijkLMN"));
// Deleted the inserted text and merge that into the base.
branch.update(cx, |branch, cx| {
branch.edit([(11..14, "")], None, cx);
branch.merge_into_base(vec![10..11], cx);
});
base.read_with(cx, |base, _| assert_eq!(base.text(), "abcdefghijk"));
}
#[gpui::test]
fn test_undo_after_merge_into_base(cx: &mut TestAppContext) {
cx.update(|cx| init_settings(cx, |_| {}));
let base = cx.new_model(|cx| Buffer::local("abcdefghijk", cx));
let branch = base.update(cx, |buffer, cx| buffer.branch(cx));
// Make 2 edits, merge one into the base.
branch.update(cx, |branch, cx| {
branch.edit([(0..3, "ABC"), (7..9, "HI")], None, cx);
branch.merge_into_base(Some(5..8), cx);
branch.merge_into_base(vec![7..7], cx);
});
assert_eq!(base.read(cx).text(), "abcdefgHIjk");
base.read_with(cx, |base, _| assert_eq!(base.text(), "abcdefgHIjk"));
branch.read_with(cx, |branch, _| assert_eq!(branch.text(), "ABCdefgHIjk"));
// Undo the merge in the base buffer.
base.update(cx, |base, cx| {
base.undo(cx);
});
base.read_with(cx, |base, _| assert_eq!(base.text(), "abcdefghijk"));
branch.read_with(cx, |branch, _| assert_eq!(branch.text(), "ABCdefgHIjk"));
// Merge that operation into the base again.
branch.update(cx, |branch, cx| {
branch.merge_into_base(vec![7..7], cx);
});
base.read_with(cx, |base, _| assert_eq!(base.text(), "abcdefgHIjk"));
branch.read_with(cx, |branch, _| assert_eq!(branch.text(), "ABCdefgHIjk"));
}
fn start_recalculating_diff(buffer: &Model<Buffer>, cx: &mut TestAppContext) {

2
crates/live_kit_server/Cargo.toml
View File

@@ -20,7 +20,7 @@ jsonwebtoken.workspace = true
log.workspace = true
prost.workspace = true
prost-types.workspace = true
reqwest.workspace = true
reqwest = "0.11"
serde.workspace = true
[build-dependencies]

72
crates/project/src/direnv.rs Normal file
View File

@@ -0,0 +1,72 @@
use crate::environment::EnvironmentErrorMessage;
use std::process::ExitStatus;
#[cfg(not(any(test, feature = "test-support")))]
use {collections::HashMap, std::path::Path, util::ResultExt};
#[derive(Clone)]
pub enum DirenvError {
NotFound,
FailedRun,
NonZeroExit(ExitStatus, Vec<u8>),
EmptyOutput,
InvalidJson,
}
impl From<DirenvError> for Option<EnvironmentErrorMessage> {
fn from(value: DirenvError) -> Self {
match value {
DirenvError::NotFound => None,
DirenvError::FailedRun | DirenvError::NonZeroExit(_, _) => {
Some(EnvironmentErrorMessage(String::from(
"Failed to run direnv. See logs for more info",
)))
}
DirenvError::EmptyOutput => None,
DirenvError::InvalidJson => Some(EnvironmentErrorMessage(String::from(
"Direnv returned invalid json. See logs for more info",
))),
}
}
}
#[cfg(not(any(test, feature = "test-support")))]
pub async fn load_direnv_environment(dir: &Path) -> Result<HashMap<String, String>, DirenvError> {
let Ok(direnv_path) = which::which("direnv") else {
return Err(DirenvError::NotFound);
};
let Some(direnv_output) = smol::process::Command::new(direnv_path)
.args(["export", "json"])
.env("TERM", "dumb")
.current_dir(dir)
.output()
.await
.log_err()
else {
return Err(DirenvError::FailedRun);
};
if !direnv_output.status.success() {
log::error!(
"Loading direnv environment failed ({}), stderr: {}",
direnv_output.status,
String::from_utf8_lossy(&direnv_output.stderr)
);
return Err(DirenvError::NonZeroExit(
direnv_output.status,
direnv_output.stderr,
));
}
let output = String::from_utf8_lossy(&direnv_output.stdout);
if output.is_empty() {
return Err(DirenvError::EmptyOutput);
}
let Some(env) = serde_json::from_str(&output).log_err() else {
return Err(DirenvError::InvalidJson);
};
Ok(env)
}

136
crates/project/src/environment.rs
View File

@@ -1,4 +1,3 @@
use anyhow::Result;
use futures::{future::Shared, FutureExt};
use std::{path::Path, sync::Arc};
use util::ResultExt;
@@ -17,6 +16,7 @@ pub struct ProjectEnvironment {
cli_environment: Option<HashMap<String, String>>,
get_environment_task: Option<Shared<Task<Option<HashMap<String, String>>>>>,
cached_shell_environments: HashMap<WorktreeId, HashMap<String, String>>,
environment_error_messages: HashMap<WorktreeId, EnvironmentErrorMessage>,
}
impl ProjectEnvironment {
@@ -37,6 +37,7 @@ impl ProjectEnvironment {
cli_environment,
get_environment_task: None,
cached_shell_environments: Default::default(),
environment_error_messages: Default::default(),
}
})
}
@@ -54,6 +55,7 @@ impl ProjectEnvironment {
pub(crate) fn remove_worktree_environment(&mut self, worktree_id: WorktreeId) {
self.cached_shell_environments.remove(&worktree_id);
self.environment_error_messages.remove(&worktree_id);
}
/// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
@@ -66,6 +68,18 @@ impl ProjectEnvironment {
}
}
/// Returns an iterator over all pairs `(worktree_id, error_message)` of
/// environment errors associated with this project environment.
pub(crate) fn environment_errors(
&self,
) -> impl Iterator<Item = (&WorktreeId, &EnvironmentErrorMessage)> {
self.environment_error_messages.iter()
}
pub(crate) fn remove_environment_error(&mut self, worktree_id: WorktreeId) {
self.environment_error_messages.remove(&worktree_id);
}
/// Returns the project environment, if possible.
/// If the project was opened from the CLI, then the inherited CLI environment is returned.
/// If it wasn't opened from the CLI, and a worktree is given, then a shell is spawned in
@@ -120,25 +134,31 @@ impl ProjectEnvironment {
let load_direnv = ProjectSettings::get_global(cx).load_direnv.clone();
cx.spawn(|this, mut cx| async move {
let mut shell_env = cx
let (mut shell_env, error_message) = cx
.background_executor()
.spawn({
let cwd = worktree_abs_path.clone();
async move { load_shell_environment(&cwd, &load_direnv).await }
})
.await
.ok();
.await;
if let Some(shell_env) = shell_env.as_mut() {
this.update(&mut cx, |this, _| {
this.cached_shell_environments
.insert(worktree_id, shell_env.clone())
.insert(worktree_id, shell_env.clone());
})
.log_err();
set_origin_marker(shell_env, EnvironmentOrigin::WorktreeShell);
}
if let Some(error) = error_message {
this.update(&mut cx, |this, _| {
this.environment_error_messages.insert(worktree_id, error);
})
.log_err();
}
shell_env
})
}
@@ -165,64 +185,62 @@ impl From<EnvironmentOrigin> for String {
}
}
pub struct EnvironmentErrorMessage(pub String);
impl EnvironmentErrorMessage {
#[allow(dead_code)]
fn from_str(s: &str) -> Self {
Self(String::from(s))
}
}
#[cfg(any(test, feature = "test-support"))]
async fn load_shell_environment(
_dir: &Path,
_load_direnv: &DirenvSettings,
) -> Result<HashMap<String, String>> {
Ok([("ZED_FAKE_TEST_ENV".into(), "true".into())]
) -> (
Option<HashMap<String, String>>,
Option<EnvironmentErrorMessage>,
) {
let fake_env = [("ZED_FAKE_TEST_ENV".into(), "true".into())]
.into_iter()
.collect())
.collect();
(Some(fake_env), None)
}
#[cfg(not(any(test, feature = "test-support")))]
async fn load_shell_environment(
dir: &Path,
load_direnv: &DirenvSettings,
) -> Result<HashMap<String, String>> {
use anyhow::{anyhow, Context};
) -> (
Option<HashMap<String, String>>,
Option<EnvironmentErrorMessage>,
) {
use crate::direnv::{load_direnv_environment, DirenvError};
use std::path::PathBuf;
use util::parse_env_output;
async fn load_direnv_environment(dir: &Path) -> Result<Option<HashMap<String, String>>> {
let Ok(direnv_path) = which::which("direnv") else {
return Ok(None);
};
let direnv_output = smol::process::Command::new(direnv_path)
.args(["export", "json"])
.current_dir(dir)
.output()
.await
.context("failed to spawn direnv to get local environment variables")?;
anyhow::ensure!(
direnv_output.status.success(),
"direnv exited with error {:?}. Stderr:\n{}",
direnv_output.status,
String::from_utf8_lossy(&direnv_output.stderr)
);
let output = String::from_utf8_lossy(&direnv_output.stdout);
if output.is_empty() {
return Ok(None);
}
Ok(Some(
serde_json::from_str(&output).context("failed to parse direnv output")?,
))
fn message<T>(with: &str) -> (Option<T>, Option<EnvironmentErrorMessage>) {
let message = EnvironmentErrorMessage::from_str(with);
(None, Some(message))
}
let direnv_environment = match load_direnv {
DirenvSettings::ShellHook => None,
DirenvSettings::Direct => load_direnv_environment(dir).await.log_err().flatten(),
}
.unwrap_or(HashMap::default());
let (direnv_environment, direnv_error) = match load_direnv {
DirenvSettings::ShellHook => (None, None),
DirenvSettings::Direct => match load_direnv_environment(dir).await {
Ok(env) => (Some(env), None),
Err(err) => (
None,
<Option<EnvironmentErrorMessage> as From<DirenvError>>::from(err),
),
},
};
let direnv_environment = direnv_environment.unwrap_or(HashMap::default());
let marker = "ZED_SHELL_START";
let shell = std::env::var("SHELL").context(
"SHELL environment variable is not assigned so we can't source login environment variables",
)?;
let Some(shell) = std::env::var("SHELL").log_err() else {
return message("Failed to get login environment. SHELL environment variable is not set");
};
// What we're doing here is to spawn a shell and then `cd` into
// the project directory to get the env in there as if the user
@@ -259,26 +277,26 @@ async fn load_shell_environment(
additional_command.unwrap_or("")
);
let output = smol::process::Command::new(&shell)
let Some(output) = smol::process::Command::new(&shell)
.args(["-l", "-i", "-c", &command])
.envs(direnv_environment)
.output()
.await
.context("failed to spawn login shell to source login environment variables")?;
.log_err()
else {
return message("Failed to spawn login shell to source login environment variables. See logs for details");
};
anyhow::ensure!(
output.status.success(),
"login shell exited with error {:?}",
output.status
);
if !output.status.success() {
log::error!("login shell exited with {}", output.status);
return message("Login shell exited with nonzero exit code. See logs for details");
}
let stdout = String::from_utf8_lossy(&output.stdout);
let env_output_start = stdout.find(marker).ok_or_else(|| {
anyhow!(
"failed to parse output of `env` command in login shell: {}",
stdout
)
})?;
let Some(env_output_start) = stdout.find(marker) else {
log::error!("failed to parse output of `env` command in login shell: {stdout}");
return message("Failed to parse stdout of env command. See logs for the output");
};
let mut parsed_env = HashMap::default();
let env_output = &stdout[env_output_start + marker.len()..];
@@ -287,5 +305,5 @@ async fn load_shell_environment(
parsed_env.insert(key, value);
});
Ok(parsed_env)
(Some(parsed_env), direnv_error)
}

25
crates/project/src/project.rs
View File

@@ -16,7 +16,9 @@ pub mod worktree_store;
#[cfg(test)]
mod project_tests;
mod direnv;
mod environment;
pub use environment::EnvironmentErrorMessage;
pub mod search_history;
mod yarn;
@@ -1431,6 +1433,23 @@ impl Project {
self.environment.read(cx).get_cli_environment()
}
pub fn shell_environment_errors<'a>(
&'a self,
cx: &'a AppContext,
) -> impl Iterator<Item = (&'a WorktreeId, &'a EnvironmentErrorMessage)> {
self.environment.read(cx).environment_errors()
}
pub fn remove_environment_error(
&mut self,
cx: &mut ModelContext<Self>,
worktree_id: WorktreeId,
) {
self.environment.update(cx, |environment, _| {
environment.remove_environment_error(worktree_id);
});
}
#[cfg(any(test, feature = "test-support"))]
pub fn has_open_buffer(&self, path: impl Into<ProjectPath>, cx: &AppContext) -> bool {
self.buffer_store
@@ -1490,8 +1509,10 @@ impl Project {
.clone()
}
pub fn ssh_is_connected(&self, cx: &AppContext) -> Option<bool> {
Some(!self.ssh_client.as_ref()?.read(cx).is_reconnect_underway())
pub fn ssh_connection_state(&self, cx: &AppContext) -> Option<remote::ConnectionState> {
self.ssh_client
.as_ref()
.map(|ssh| ssh.read(cx).connection_state())
}
pub fn replica_id(&self) -> ReplicaId {

56
crates/recent_projects/src/dev_servers.rs
View File

@@ -279,6 +279,13 @@ impl DevServerProjects {
match connection.await {
Some(_) => this
.update(&mut cx, |this, cx| {
let _ = this.workspace.update(cx, |workspace, _| {
workspace
.client()
.telemetry()
.report_app_event("create ssh server".to_string())
});
this.add_ssh_server(connection_options, cx);
this.mode = Mode::Default(None);
cx.notify()
@@ -422,7 +429,15 @@ impl DevServerProjects {
);
cx.new_view(|cx| {
Workspace::new(None, project.clone(), app_state.clone(), cx)
let workspace =
Workspace::new(None, project.clone(), app_state.clone(), cx);
workspace
.client()
.telemetry()
.report_app_event("create ssh project".to_string());
workspace
})
})
.log_err();
@@ -541,23 +556,28 @@ impl DevServerProjects {
.w_full()
.border_l_1()
.border_color(cx.theme().colors().border_variant)
.my_1()
.mb_1()
.mx_1p5()
.py_0p5()
.px_3()
.pl_2()
.child(
List::new()
.empty_message("No projects.")
.children(ssh_connection.projects.iter().enumerate().map(|(pix, p)| {
self.render_ssh_project(ix, &ssh_connection, pix, p, cx)
v_flex().gap_0p5().child(self.render_ssh_project(
ix,
&ssh_connection,
pix,
p,
cx,
))
}))
.child(
h_flex().child(
h_flex().mt_1().pl_1().child(
Button::new("new-remote_project", "Open Folder…")
.icon(IconName::Plus)
.size(ButtonSize::Default)
.style(ButtonStyle::Filled)
.layer(ElevationIndex::ModalSurface)
.icon(IconName::Plus)
.icon_color(Color::Muted)
.icon_position(IconPosition::Start)
.on_click(cx.listener(move |this, _, cx| {
this.create_ssh_project(ix, ssh_connection.clone(), cx);
@@ -578,9 +598,15 @@ impl DevServerProjects {
) -> impl IntoElement {
let project = project.clone();
let server = server.clone();
ListItem::new(("remote-project", ix))
.inset(true)
.spacing(ui::ListItemSpacing::Sparse)
.start_slot(Icon::new(IconName::Folder).color(Color::Muted))
.start_slot(
Icon::new(IconName::Folder)
.color(Color::Muted)
.size(IconSize::Small),
)
.child(Label::new(project.paths.join(", ")))
.on_click(cx.listener(move |this, _, cx| {
let Some(app_state) = this
@@ -620,7 +646,7 @@ impl DevServerProjects {
.on_click(
cx.listener(move |this, _, cx| this.delete_ssh_project(server_ix, ix, cx)),
)
.tooltip(|cx| Tooltip::text("Delete remote project", cx))
.tooltip(|cx| Tooltip::text("Delete Remote Project", cx))
.into_any_element(),
))
}
@@ -694,6 +720,7 @@ impl DevServerProjects {
})
});
let theme = cx.theme();
v_flex()
.id("create-dev-server")
.overflow_hidden()
@@ -748,6 +775,7 @@ impl DevServerProjects {
.child(
h_flex()
.bg(theme.colors().editor_background)
.rounded_b_md()
.w_full()
.map(|this| {
if let Some(ssh_prompt) = ssh_prompt {
@@ -758,9 +786,8 @@ impl DevServerProjects {
h_flex()
.p_2()
.w_full()
.content_center()
.gap_2()
.child(h_flex().w_full())
.justify_center()
.gap_1p5()
.child(
div().p_1().rounded_lg().bg(color).with_animation(
"pulse-ssh-waiting-for-connection",
@@ -773,8 +800,7 @@ impl DevServerProjects {
.child(
Label::new("Waiting for connection…")
.size(LabelSize::Small),
)
.child(h_flex().w_full()),
),
)
}
}),

6
crates/recent_projects/src/recent_projects.rs
View File

@@ -566,7 +566,7 @@ impl PickerDelegate for RecentProjectsDelegate {
.border_t_1()
.py_2()
.pr_2()
.border_color(cx.theme().colors().border)
.border_color(cx.theme().colors().border_variant)
.justify_end()
.gap_4()
.child(
@@ -574,7 +574,7 @@ impl PickerDelegate for RecentProjectsDelegate {
.when_some(KeyBinding::for_action(&OpenRemote, cx), |button, key| {
button.child(key)
})
.child(Label::new("Open remote folder…").color(Color::Muted))
.child(Label::new("Open Remote Folder…").color(Color::Muted))
.on_click(|_, cx| cx.dispatch_action(OpenRemote.boxed_clone())),
)
.child(
@@ -583,7 +583,7 @@ impl PickerDelegate for RecentProjectsDelegate {
KeyBinding::for_action(&workspace::Open, cx),
|button, key| button.child(key),
)
.child(Label::new("Open local folder…").color(Color::Muted))
.child(Label::new("Open Local Folder…").color(Color::Muted))
.on_click(|_, cx| cx.dispatch_action(workspace::Open.boxed_clone())),
)
.into_any(),

147
crates/recent_projects/src/ssh_connections.rs
View File

@@ -16,9 +16,9 @@ use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use settings::{Settings, SettingsSources};
use ui::{
div, h_flex, v_flex, ActiveTheme, ButtonCommon, Clickable, Color, FluentBuilder as _, Icon,
IconButton, IconName, IconSize, InteractiveElement, IntoElement, Label, LabelCommon, Styled,
StyledExt as _, Tooltip, ViewContext, VisualContext, WindowContext,
div, h_flex, prelude::*, v_flex, ActiveTheme, ButtonCommon, Clickable, Color, Icon, IconButton,
IconName, IconSize, InteractiveElement, IntoElement, Label, LabelCommon, Styled, Tooltip,
ViewContext, VisualContext, WindowContext,
};
use workspace::{AppState, ModalView, Workspace};
@@ -84,6 +84,7 @@ pub struct SshPrompt {
pub struct SshConnectionModal {
pub(crate) prompt: View<SshPrompt>,
}
impl SshPrompt {
pub fn new(connection_options: &SshConnectionOptions, cx: &mut ViewContext<Self>) -> Self {
let connection_string = connection_options.connection_string().into();
@@ -136,57 +137,72 @@ impl SshPrompt {
}
impl Render for SshPrompt {
fn render(&mut self, _: &mut ViewContext<Self>) -> impl IntoElement {
fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
let cx = cx.window_context();
let theme = cx.theme();
v_flex()
.w_full()
.key_context("PasswordPrompt")
.justify_start()
.size_full()
.justify_center()
.child(
v_flex()
.p_4()
.size_full()
.child(
h_flex()
.gap_2()
.justify_between()
.child(h_flex().w_full())
.child(if self.error_message.is_some() {
Icon::new(IconName::XCircle)
.size(IconSize::Medium)
.color(Color::Error)
.into_any_element()
} else {
Icon::new(IconName::ArrowCircle)
.size(IconSize::Medium)
.with_animation(
"arrow-circle",
Animation::new(Duration::from_secs(2)).repeat(),
|icon, delta| {
icon.transform(Transformation::rotate(percentage(
delta,
)))
},
)
.into_any_element()
})
.child(Label::new(format!(
"Connecting to {}…",
self.connection_string
)))
.child(h_flex().w_full()),
)
.when_some(self.error_message.as_ref(), |el, error| {
el.child(Label::new(error.clone()))
h_flex()
.p_2()
.justify_center()
.flex_wrap()
.child(if self.error_message.is_some() {
Icon::new(IconName::XCircle)
.size(IconSize::Medium)
.color(Color::Error)
.into_any_element()
} else {
Icon::new(IconName::ArrowCircle)
.size(IconSize::Medium)
.with_animation(
"arrow-circle",
Animation::new(Duration::from_secs(2)).repeat(),
|icon, delta| {
icon.transform(Transformation::rotate(percentage(delta)))
},
)
.into_any_element()
})
.when(
self.error_message.is_none() && self.status_message.is_some(),
|el| el.child(Label::new(self.status_message.clone().unwrap())),
.child(
div()
.ml_1()
.child(Label::new("SSH Connection").size(LabelSize::Small)),
)
.when_some(self.prompt.as_ref(), |el, prompt| {
el.child(Label::new(prompt.0.clone()))
.child(self.editor.clone())
}),
.child(
div()
.text_ellipsis()
.overflow_x_hidden()
.when_some(self.error_message.as_ref(), |el, error| {
el.child(Label::new(format!("－{}", error)).size(LabelSize::Small))
})
.when(
self.error_message.is_none() && self.status_message.is_some(),
|el| {
el.child(
Label::new(format!(
"－{}",
self.status_message.clone().unwrap()
))
.size(LabelSize::Small),
)
},
),
),
)
.child(div().when_some(self.prompt.as_ref(), |el, prompt| {
el.child(
h_flex()
.p_4()
.border_t_1()
.border_color(theme.colors().border_variant)
.font_buffer(cx)
.child(Label::new(prompt.0.clone()))
.child(self.editor.clone()),
)
}))
}
}
@@ -210,39 +226,54 @@ impl Render for SshConnectionModal {
fn render(&mut self, cx: &mut ui::ViewContext<Self>) -> impl ui::IntoElement {
let connection_string = self.prompt.read(cx).connection_string.clone();
let theme = cx.theme();
let header_color = theme.colors().element_background;
let body_color = theme.colors().background;
let mut header_color = cx.theme().colors().text;
header_color.fade_out(0.96);
let body_color = theme.colors().editor_background;
v_flex()
.elevation_3(cx)
.on_action(cx.listener(Self::dismiss))
.on_action(cx.listener(Self::confirm))
.w(px(400.))
.w(px(500.))
.border_1()
.border_color(theme.colors().border)
.child(
h_flex()
.relative()
.p_1()
.rounded_t_md()
.border_b_1()
.border_color(theme.colors().border)
.bg(header_color)
.justify_between()
.child(
IconButton::new("ssh-connection-cancel", IconName::ArrowLeft)
.icon_size(IconSize::XSmall)
.on_click(|_, cx| cx.dispatch_action(menu::Cancel.boxed_clone()))
.tooltip(|cx| Tooltip::for_action("Back", &menu::Cancel, cx)),
div().absolute().left_0p5().top_0p5().child(
IconButton::new("ssh-connection-cancel", IconName::ArrowLeft)
.icon_size(IconSize::XSmall)
.on_click(|_, cx| cx.dispatch_action(menu::Cancel.boxed_clone()))
.tooltip(|cx| Tooltip::for_action("Back", &menu::Cancel, cx)),
),
)
.child(
h_flex()
.w_full()
.gap_2()
.justify_center()
.child(Icon::new(IconName::Server).size(IconSize::XSmall))
.child(
Label::new(connection_string)
.size(ui::LabelSize::Small)
.single_line(),
),
)
.child(div()),
),
)
.child(
h_flex()
.rounded_b_md()
.bg(body_color)
.w_full()
.child(self.prompt.clone()),
)
.child(h_flex().bg(body_color).w_full().child(self.prompt.clone()))
}
}

4
crates/remote/src/remote.rs
View File

@@ -2,4 +2,6 @@ pub mod json_log;
pub mod protocol;
pub mod ssh_session;
pub use ssh_session::{SshClientDelegate, SshConnectionOptions, SshPlatform, SshRemoteClient};
pub use ssh_session::{
ConnectionState, SshClientDelegate, SshConnectionOptions, SshPlatform, SshRemoteClient,
};

520
crates/remote/src/ssh_session.rs
View File

@@ -31,7 +31,8 @@ use smol::{
use std::{
any::TypeId,
ffi::OsStr,
mem,
fmt,
ops::ControlFlow,
path::{Path, PathBuf},
sync::{
atomic::{AtomicU32, Ordering::SeqCst},
@@ -40,7 +41,7 @@ use std::{
time::{Duration, Instant},
};
use tempfile::TempDir;
use util::maybe;
use util::ResultExt;
#[derive(
Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,
@@ -234,19 +235,157 @@ impl ChannelForwarder {
}
}
struct SshRemoteClientState {
ssh_connection: SshRemoteConnection,
delegate: Arc<dyn SshClientDelegate>,
forwarder: ChannelForwarder,
multiplex_task: Task<Result<()>>,
heartbeat_task: Task<Result<()>>,
const MAX_MISSED_HEARTBEATS: usize = 5;
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
const MAX_RECONNECT_ATTEMPTS: usize = 3;
enum State {
Connecting,
Connected {
ssh_connection: SshRemoteConnection,
delegate: Arc<dyn SshClientDelegate>,
forwarder: ChannelForwarder,
multiplex_task: Task<Result<()>>,
heartbeat_task: Task<Result<()>>,
},
HeartbeatMissed {
missed_heartbeats: usize,
ssh_connection: SshRemoteConnection,
delegate: Arc<dyn SshClientDelegate>,
forwarder: ChannelForwarder,
multiplex_task: Task<Result<()>>,
heartbeat_task: Task<Result<()>>,
},
Reconnecting,
ReconnectFailed {
ssh_connection: SshRemoteConnection,
delegate: Arc<dyn SshClientDelegate>,
forwarder: ChannelForwarder,
error: anyhow::Error,
attempts: usize,
},
ReconnectExhausted,
}
impl fmt::Display for State {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Connecting => write!(f, "connecting"),
Self::Connected { .. } => write!(f, "connected"),
Self::Reconnecting => write!(f, "reconnecting"),
Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),
Self::ReconnectExhausted => write!(f, "reconnect exhausted"),
Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),
}
}
}
impl State {
fn ssh_connection(&self) -> Option<&SshRemoteConnection> {
match self {
Self::Connected { ssh_connection, .. } => Some(ssh_connection),
Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection),
Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection),
_ => None,
}
}
fn can_reconnect(&self) -> bool {
matches!(
self,
Self::Connected { .. } | Self::HeartbeatMissed { .. } | Self::ReconnectFailed { .. }
)
}
fn heartbeat_recovered(self) -> Self {
match self {
Self::HeartbeatMissed {
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
..
} => Self::Connected {
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
},
_ => self,
}
}
fn heartbeat_missed(self) -> Self {
match self {
Self::Connected {
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
} => Self::HeartbeatMissed {
missed_heartbeats: 1,
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
},
Self::HeartbeatMissed {
missed_heartbeats,
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
} => Self::HeartbeatMissed {
missed_heartbeats: missed_heartbeats + 1,
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
},
_ => self,
}
}
}
/// The state of the ssh connection.
#[derive(Clone, Copy, Debug)]
pub enum ConnectionState {
Connecting,
Connected,
HeartbeatMissed,
Reconnecting,
Disconnected,
}
impl From<&State> for ConnectionState {
fn from(value: &State) -> Self {
match value {
State::Connecting => Self::Connecting,
State::Connected { .. } => Self::Connected,
State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,
State::HeartbeatMissed { .. } => Self::HeartbeatMissed,
State::ReconnectExhausted => Self::Disconnected,
}
}
}
pub struct SshRemoteClient {
client: Arc<ChannelClient>,
unique_identifier: String,
connection_options: SshConnectionOptions,
inner_state: Arc<Mutex<Option<SshRemoteClientState>>>,
state: Arc<Mutex<Option<State>>>,
}
impl Drop for SshRemoteClient {
@@ -266,6 +405,7 @@ impl SshRemoteClient {
let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
let client = cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx))?;
let this = cx.new_model(|cx| {
cx.on_app_quit(|this: &mut Self, _| {
this.shutdown_processes();
@@ -273,47 +413,49 @@ impl SshRemoteClient {
})
.detach();
let client = ChannelClient::new(incoming_rx, outgoing_tx, cx);
Self {
client,
client: client.clone(),
unique_identifier: unique_identifier.clone(),
connection_options: SshConnectionOptions::default(),
inner_state: Arc::new(Mutex::new(None)),
connection_options: connection_options.clone(),
state: Arc::new(Mutex::new(Some(State::Connecting))),
}
})?;
let inner_state = {
let (proxy, proxy_incoming_tx, proxy_outgoing_rx) =
ChannelForwarder::new(incoming_tx, outgoing_rx, &mut cx);
let (proxy, proxy_incoming_tx, proxy_outgoing_rx) =
ChannelForwarder::new(incoming_tx, outgoing_rx, &mut cx);
let (ssh_connection, ssh_proxy_process) = Self::establish_connection(
unique_identifier,
connection_options,
delegate.clone(),
&mut cx,
)
.await?;
let (ssh_connection, ssh_proxy_process) = Self::establish_connection(
unique_identifier,
connection_options,
delegate.clone(),
&mut cx,
)
.await?;
let multiplex_task = Self::multiplex(
this.downgrade(),
ssh_proxy_process,
proxy_incoming_tx,
proxy_outgoing_rx,
&mut cx,
);
let multiplex_task = Self::multiplex(
this.downgrade(),
ssh_proxy_process,
proxy_incoming_tx,
proxy_outgoing_rx,
&mut cx,
);
SshRemoteClientState {
if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
log::error!("failed to establish connection: {}", error);
delegate.set_error(error.to_string(), &mut cx);
return Err(error);
}
let heartbeat_task = Self::heartbeat(this.downgrade(), &mut cx);
this.update(&mut cx, |this, _| {
*this.state.lock() = Some(State::Connected {
ssh_connection,
delegate,
forwarder: proxy,
multiplex_task,
heartbeat_task: Self::heartbeat(this.downgrade(), &mut cx),
}
};
this.update(&mut cx, |this, cx| {
this.inner_state.lock().replace(inner_state);
cx.notify();
heartbeat_task,
});
})?;
Ok(this)
@@ -321,78 +463,192 @@ impl SshRemoteClient {
}
fn shutdown_processes(&self) {
let Some(mut state) = self.inner_state.lock().take() else {
let Some(state) = self.state.lock().take() else {
return;
};
log::info!("shutting down ssh processes");
// Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
// child of master_process.
let task = mem::replace(&mut state.multiplex_task, Task::ready(Ok(())));
drop(task);
// Now drop the rest of state, which kills master process.
drop(state);
}
fn reconnect(&self, cx: &ModelContext<Self>) -> Result<()> {
log::info!("Trying to reconnect to ssh server...");
let Some(state) = self.inner_state.lock().take() else {
return Err(anyhow!("reconnect is already in progress"));
};
let workspace_identifier = self.unique_identifier.clone();
let SshRemoteClientState {
mut ssh_connection,
delegate,
forwarder: proxy,
let State::Connected {
multiplex_task,
heartbeat_task,
} = state;
..
} = state
else {
return;
};
// Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
// child of master_process.
drop(multiplex_task);
// Now drop the rest of state, which kills master process.
drop(heartbeat_task);
}
cx.spawn(|this, mut cx| async move {
let (incoming_tx, outgoing_rx) = proxy.into_channels().await;
fn reconnect(&mut self, cx: &mut ModelContext<Self>) -> Result<()> {
let mut lock = self.state.lock();
ssh_connection.master_process.kill()?;
ssh_connection
let can_reconnect = lock
.as_ref()
.map(|state| state.can_reconnect())
.unwrap_or(false);
if !can_reconnect {
let error = if let Some(state) = lock.as_ref() {
format!("invalid state, cannot reconnect while in state {state}")
} else {
"no state set".to_string()
};
return Err(anyhow!(error));
}
let state = lock.take().unwrap();
let (attempts, mut ssh_connection, delegate, forwarder) = match state {
State::Connected {
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
}
| State::HeartbeatMissed {
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task,
..
} => {
drop(multiplex_task);
drop(heartbeat_task);
(0, ssh_connection, delegate, forwarder)
}
State::ReconnectFailed {
attempts,
ssh_connection,
delegate,
forwarder,
..
} => (attempts, ssh_connection, delegate, forwarder),
State::Connecting | State::Reconnecting | State::ReconnectExhausted => unreachable!(),
};
let attempts = attempts + 1;
if attempts > MAX_RECONNECT_ATTEMPTS {
log::error!(
"Failed to reconnect to after {} attempts, giving up",
MAX_RECONNECT_ATTEMPTS
);
*lock = Some(State::ReconnectExhausted);
return Ok(());
}
*lock = Some(State::Reconnecting);
drop(lock);
log::info!("Trying to reconnect to ssh server... Attempt {}", attempts);
let identifier = self.unique_identifier.clone();
let client = self.client.clone();
let reconnect_task = cx.spawn(|this, mut cx| async move {
macro_rules! failed {
($error:expr, $attempts:expr, $ssh_connection:expr, $delegate:expr, $forwarder:expr) => {
return State::ReconnectFailed {
error: anyhow!($error),
attempts: $attempts,
ssh_connection: $ssh_connection,
delegate: $delegate,
forwarder: $forwarder,
};
};
}
if let Err(error) = ssh_connection.master_process.kill() {
failed!(error, attempts, ssh_connection, delegate, forwarder);
};
if let Err(error) = ssh_connection
.master_process
.status()
.await
.context("Failed to kill ssh process")?;
.context("Failed to kill ssh process")
{
failed!(error, attempts, ssh_connection, delegate, forwarder);
}
let connection_options = ssh_connection.socket.connection_options.clone();
let (ssh_connection, ssh_process) = Self::establish_connection(
workspace_identifier,
let (incoming_tx, outgoing_rx) = forwarder.into_channels().await;
let (forwarder, proxy_incoming_tx, proxy_outgoing_rx) =
ChannelForwarder::new(incoming_tx, outgoing_rx, &mut cx);
let (ssh_connection, ssh_process) = match Self::establish_connection(
identifier,
connection_options,
delegate.clone(),
&mut cx,
)
.await?;
let (proxy, proxy_incoming_tx, proxy_outgoing_rx) =
ChannelForwarder::new(incoming_tx, outgoing_rx, &mut cx);
let inner_state = SshRemoteClientState {
ssh_connection,
delegate,
forwarder: proxy,
multiplex_task: Self::multiplex(
this.clone(),
ssh_process,
proxy_incoming_tx,
proxy_outgoing_rx,
&mut cx,
),
heartbeat_task: Self::heartbeat(this.clone(), &mut cx),
.await
{
Ok((ssh_connection, ssh_process)) => (ssh_connection, ssh_process),
Err(error) => {
failed!(error, attempts, ssh_connection, delegate, forwarder);
}
};
this.update(&mut cx, |this, _| {
this.inner_state.lock().replace(inner_state);
let multiplex_task = Self::multiplex(
this.clone(),
ssh_process,
proxy_incoming_tx,
proxy_outgoing_rx,
&mut cx,
);
if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
failed!(error, attempts, ssh_connection, delegate, forwarder);
};
State::Connected {
ssh_connection,
delegate,
forwarder,
multiplex_task,
heartbeat_task: Self::heartbeat(this.clone(), &mut cx),
}
});
cx.spawn(|this, mut cx| async move {
let new_state = reconnect_task.await;
this.update(&mut cx, |this, cx| {
match &new_state {
State::Connecting
| State::Reconnecting { .. }
| State::HeartbeatMissed { .. } => {}
State::Connected { .. } => {
log::info!("Successfully reconnected");
}
State::ReconnectFailed {
error, attempts, ..
} => {
log::error!(
"Reconnect attempt {} failed: {:?}. Starting new attempt...",
attempts,
error
);
}
State::ReconnectExhausted => {
log::error!("Reconnect attempt failed and all attempts exhausted");
}
}
let reconnect_failed = matches!(new_state, State::ReconnectFailed { .. });
*this.state.lock() = Some(new_state);
cx.notify();
if reconnect_failed {
this.reconnect(cx)
} else {
Ok(())
}
})
})
.detach();
.detach_and_log_err(cx);
Ok(())
}
@@ -403,10 +659,6 @@ impl SshRemoteClient {
cx.spawn(|mut cx| {
let this = this.clone();
async move {
const MAX_MISSED_HEARTBEATS: usize = 5;
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
let mut missed_heartbeats = 0;
let mut timer = Timer::interval(HEARTBEAT_INTERVAL);
@@ -415,19 +667,7 @@ impl SshRemoteClient {
log::info!("Sending heartbeat to server...");
let result = smol::future::or(
async {
client.request(proto::Ping {}).await?;
Ok(())
},
async {
smol::Timer::after(HEARTBEAT_TIMEOUT).await;
Err(anyhow!("Timeout detected"))
},
)
.await;
let result = client.ping(HEARTBEAT_TIMEOUT).await;
if result.is_err() {
missed_heartbeats += 1;
log::warn!(
@@ -440,17 +680,10 @@ impl SshRemoteClient {
missed_heartbeats = 0;
}
if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
log::error!(
"Missed last {} hearbeats. Reconnecting...",
missed_heartbeats
);
this.update(&mut cx, |this, cx| {
this.reconnect(cx)
.context("failed to reconnect after missing heartbeats")
})
.context("failed to update weak reference, SshRemoteClient lost?")??;
let result = this.update(&mut cx, |this, mut cx| {
this.handle_heartbeat_result(missed_heartbeats, &mut cx)
})?;
if result.is_break() {
return Ok(());
}
}
@@ -458,6 +691,34 @@ impl SshRemoteClient {
})
}
fn handle_heartbeat_result(
&mut self,
missed_heartbeats: usize,
cx: &mut ModelContext<Self>,
) -> ControlFlow<()> {
let state = self.state.lock().take().unwrap();
self.state.lock().replace(if missed_heartbeats > 0 {
state.heartbeat_missed()
} else {
state.heartbeat_recovered()
});
cx.notify();
if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
log::error!(
"Missed last {} heartbeats. Reconnecting...",
missed_heartbeats
);
self.reconnect(cx)
.context("failed to start reconnect process after missing heartbeats")
.log_err();
ControlFlow::Break(())
} else {
ControlFlow::Continue(())
}
}
fn multiplex(
this: WeakModel<Self>,
mut ssh_proxy_process: Child,
@@ -611,10 +872,11 @@ impl SshRemoteClient {
}
pub fn ssh_args(&self) -> Option<Vec<String>> {
let state = self.inner_state.lock();
state
self.state
.lock()
.as_ref()
.map(|state| state.ssh_connection.socket.ssh_args())
.and_then(|state| state.ssh_connection())
.map(|ssh_connection| ssh_connection.socket.ssh_args())
}
pub fn to_proto_client(&self) -> AnyProtoClient {
@@ -625,8 +887,12 @@ impl SshRemoteClient {
self.connection_options.connection_string()
}
pub fn is_reconnect_underway(&self) -> bool {
maybe!({ Some(self.inner_state.try_lock()?.is_none()) }).unwrap_or_default()
pub fn connection_state(&self) -> ConnectionState {
self.state
.lock()
.as_ref()
.map(ConnectionState::from)
.unwrap_or(ConnectionState::Disconnected)
}
#[cfg(any(test, feature = "test-support"))]
@@ -646,7 +912,7 @@ impl SshRemoteClient {
client,
unique_identifier: "fake".to_string(),
connection_options: SshConnectionOptions::default(),
inner_state: Arc::new(Mutex::new(None)),
state: Arc::new(Mutex::new(None)),
})
}),
server_cx.update(|cx| ChannelClient::new(client_to_server_rx, server_to_client_tx, cx)),
@@ -1046,6 +1312,20 @@ impl ChannelClient {
}
}
pub async fn ping(&self, timeout: Duration) -> Result<()> {
smol::future::or(
async {
self.request(proto::Ping {}).await?;
Ok(())
},
async {
smol::Timer::after(timeout).await;
Err(anyhow!("Timeout detected"))
},
)
.await
}
pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
log::debug!("ssh send name:{}", T::NAME);
self.send_dynamic(payload.into_envelope(0, None, None))

1
crates/remote_server/Cargo.toml
View File

@@ -22,6 +22,7 @@ test-support = ["fs/test-support"]
[dependencies]
anyhow.workspace = true
backtrace = "0.3"
clap.workspace = true
client.workspace = true
env_logger.workspace = true

6
crates/remote_server/src/main.rs
View File

@@ -37,7 +37,7 @@ fn main() {
#[cfg(not(windows))]
fn main() -> Result<()> {
use remote_server::unix::{execute_proxy, execute_run, init_logging};
use remote_server::unix::{execute_proxy, execute_run, init};
let cli = Cli::parse();
@@ -48,11 +48,11 @@ fn main() -> Result<()> {
stdin_socket,
stdout_socket,
}) => {
init_logging(Some(log_file))?;
init(Some(log_file))?;
execute_run(pid_file, stdin_socket, stdout_socket)
}
Some(Commands::Proxy { identifier }) => {
init_logging(None)?;
init(None)?;
execute_proxy(identifier)
}
Some(Commands::Version) => {

62
crates/remote_server/src/unix.rs
View File

@@ -20,7 +20,13 @@ use std::{
sync::Arc,
};
pub fn init_logging(log_file: Option<PathBuf>) -> Result<()> {
pub fn init(log_file: Option<PathBuf>) -> Result<()> {
init_logging(log_file)?;
init_panic_hook();
Ok(())
}
fn init_logging(log_file: Option<PathBuf>) -> Result<()> {
if let Some(log_file) = log_file {
let target = Box::new(if log_file.exists() {
std::fs::OpenOptions::new()
@@ -46,6 +52,45 @@ pub fn init_logging(log_file: Option<PathBuf>) -> Result<()> {
Ok(())
}
fn init_panic_hook() {
std::panic::set_hook(Box::new(|info| {
let payload = info
.payload()
.downcast_ref::<&str>()
.map(|s| s.to_string())
.or_else(|| info.payload().downcast_ref::<String>().cloned())
.unwrap_or_else(|| "Box<Any>".to_string());
let backtrace = backtrace::Backtrace::new();
let mut backtrace = backtrace
.frames()
.iter()
.flat_map(|frame| {
frame
.symbols()
.iter()
.filter_map(|frame| Some(format!("{:#}", frame.name()?)))
})
.collect::<Vec<_>>();
// Strip out leading stack frames for rust panic-handling.
if let Some(ix) = backtrace
.iter()
.position(|name| name == "rust_begin_unwind")
{
backtrace.drain(0..=ix);
}
log::error!(
"server: panic occurred: {}\nBacktrace:\n{}",
payload,
backtrace.join("\n")
);
std::process::abort();
}));
}
fn start_server(
stdin_listener: UnixListener,
stdout_listener: UnixListener,
@@ -61,6 +106,7 @@ fn start_server(
cx.on_app_quit(move |_| {
let mut app_quit_tx = app_quit_tx.clone();
async move {
log::info!("app quitting. sending signal to server main loop");
app_quit_tx.send(()).await.ok();
}
})
@@ -150,6 +196,13 @@ fn start_server(
}
pub fn execute_run(pid_file: PathBuf, stdin_socket: PathBuf, stdout_socket: PathBuf) -> Result<()> {
log::info!(
"server: starting up. pid_file: {:?}, stdin_socket: {:?}, stdout_socket: {:?}",
pid_file,
stdin_socket,
stdout_socket
);
write_pid_file(&pid_file)
.with_context(|| format!("failed to write pid file: {:?}", &pid_file))?;
@@ -157,10 +210,12 @@ pub fn execute_run(pid_file: PathBuf, stdin_socket: PathBuf, stdout_socket: Path
let stdout_listener =
UnixListener::bind(stdout_socket).context("failed to bind stdout socket")?;
log::debug!("server: starting gpui app");
gpui::App::headless().run(move |cx| {
settings::init(cx);
HeadlessProject::init(cx);
log::info!("server: gpui app started, initializing server");
let session = start_server(stdin_listener, stdout_listener, cx);
let project = cx.new_model(|cx| {
HeadlessProject::new(session, Arc::new(RealFs::new(Default::default(), None)), cx)
@@ -298,8 +353,9 @@ fn write_pid_file(path: &Path) -> Result<()> {
if path.exists() {
std::fs::remove_file(path)?;
}
std::fs::write(path, std::process::id().to_string()).context("Failed to write PID file")
let pid = std::process::id().to_string();
log::debug!("server: writing PID {} to file {:?}", pid, path);
std::fs::write(path, pid).context("Failed to write PID file")
}
async fn handle_io<R, W>(mut reader: R, mut writer: W, socket_name: &str) -> Result<()>

31
crates/reqwest_client/Cargo.toml
View File

@@ -1,31 +0,0 @@
[package]
name = "reqwest_client"
version = "0.1.0"
edition = "2021"
publish = false
license = "Apache-2.0"
[lints]
workspace = true
[features]
test-support = []
[lib]
path = "src/reqwest_client.rs"
doctest = true
[[example]]
name = "client"
path = "examples/client.rs"
[dependencies]
anyhow.workspace = true
bytes = "1.0"
futures.workspace = true
http_client.workspace = true
serde.workspace = true
smol.workspace = true
tokio.workspace = true
reqwest = { workspace = true, features = ["rustls-tls-manual-roots", "stream"] }

1
crates/reqwest_client/LICENSE-GPL
View File

@@ -1 +0,0 @@
../../LICENSE-GPL

16
crates/reqwest_client/examples/client.rs
View File

@@ -1,16 +0,0 @@
use futures::AsyncReadExt as _;
use http_client::AsyncBody;
use http_client::HttpClient;
use reqwest_client::ReqwestClient;
#[tokio::main]
async fn main() {
let resp = ReqwestClient::new()
.get("http://zed.dev", AsyncBody::empty(), true)
.await
.unwrap();
let mut body = String::new();
resp.into_body().read_to_string(&mut body).await.unwrap();
println!("{}", &body);
}

232
crates/reqwest_client/src/reqwest_client.rs
View File

@@ -1,232 +0,0 @@
use std::{borrow::Cow, io::Read, pin::Pin, task::Poll};
use anyhow::anyhow;
use bytes::{BufMut, Bytes, BytesMut};
use futures::{AsyncRead, TryStreamExt};
use http_client::{http, AsyncBody, ReadTimeout};
use reqwest::header::{HeaderMap, HeaderValue};
use smol::future::FutureExt;
const DEFAULT_CAPACITY: usize = 4096;
pub struct ReqwestClient {
client: reqwest::Client,
}
impl ReqwestClient {
pub fn new() -> Self {
Self {
client: reqwest::Client::new(),
}
}
pub fn user_agent(agent: &str) -> anyhow::Result<Self> {
let mut map = HeaderMap::new();
map.insert(http::header::USER_AGENT, HeaderValue::from_str(agent)?);
Ok(Self {
client: reqwest::Client::builder().default_headers(map).build()?,
})
}
}
impl From<reqwest::Client> for ReqwestClient {
fn from(client: reqwest::Client) -> Self {
Self { client }
}
}
// This struct is essentially a re-implementation of
// https://docs.rs/tokio-util/0.7.12/tokio_util/io/struct.ReaderStream.html
// except outside of Tokio's aegis
struct ReaderStream {
reader: Option<Pin<Box<dyn futures::AsyncRead + Send + Sync>>>,
buf: BytesMut,
capacity: usize,
}
impl ReaderStream {
fn new(reader: Pin<Box<dyn futures::AsyncRead + Send + Sync>>) -> Self {
Self {
reader: Some(reader),
buf: BytesMut::new(),
capacity: DEFAULT_CAPACITY,
}
}
}
impl futures::Stream for ReaderStream {
type Item = std::io::Result<Bytes>;
fn poll_next(
mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut();
let mut reader = match this.reader.take() {
Some(r) => r,
None => return Poll::Ready(None),
};
if this.buf.capacity() == 0 {
let capacity = this.capacity;
this.buf.reserve(capacity);
}
match poll_read_buf(&mut reader, cx, &mut this.buf) {
Poll::Pending => Poll::Pending,
Poll::Ready(Err(err)) => {
self.reader = None;
Poll::Ready(Some(Err(err)))
}
Poll::Ready(Ok(0)) => {
self.reader = None;
Poll::Ready(None)
}
Poll::Ready(Ok(_)) => {
let chunk = this.buf.split();
self.reader = Some(reader);
Poll::Ready(Some(Ok(chunk.freeze())))
}
}
}
}
/// Implementation from https://docs.rs/tokio-util/0.7.12/src/tokio_util/util/poll_buf.rs.html
/// Specialized for this use case
pub fn poll_read_buf(
io: &mut Pin<Box<dyn futures::AsyncRead + Send + Sync>>,
cx: &mut std::task::Context<'_>,
buf: &mut BytesMut,
) -> Poll<std::io::Result<usize>> {
if !buf.has_remaining_mut() {
return Poll::Ready(Ok(0));
}
let n = {
let dst = buf.chunk_mut();
// Safety: `chunk_mut()` returns a `&mut UninitSlice`, and `UninitSlice` is a
// transparent wrapper around `[MaybeUninit<u8>]`.
let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>]) };
let mut buf = tokio::io::ReadBuf::uninit(dst);
let ptr = buf.filled().as_ptr();
let unfilled_portion = buf.initialize_unfilled();
// SAFETY: Pin projection
let io_pin = unsafe { Pin::new_unchecked(io) };
std::task::ready!(io_pin.poll_read(cx, unfilled_portion)?);
// Ensure the pointer does not change from under us
assert_eq!(ptr, buf.filled().as_ptr());
buf.filled().len()
};
// Safety: This is guaranteed to be the number of initialized (and read)
// bytes due to the invariants provided by `ReadBuf::filled`.
unsafe {
buf.advance_mut(n);
}
Poll::Ready(Ok(n))
}
enum WrappedBodyInner {
None,
SyncReader(std::io::Cursor<Cow<'static, [u8]>>),
Stream(ReaderStream),
}
struct WrappedBody(WrappedBodyInner);
impl WrappedBody {
fn new(body: AsyncBody) -> Self {
match body.0 {
http_client::Inner::Empty => Self(WrappedBodyInner::None),
http_client::Inner::SyncReader(cursor) => Self(WrappedBodyInner::SyncReader(cursor)),
http_client::Inner::AsyncReader(pin) => {
Self(WrappedBodyInner::Stream(ReaderStream::new(pin)))
}
}
}
}
impl futures::stream::Stream for WrappedBody {
type Item = Result<Bytes, std::io::Error>;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
match &mut self.0 {
WrappedBodyInner::None => Poll::Ready(None),
WrappedBodyInner::SyncReader(cursor) => {
let mut buf = Vec::new();
match cursor.read_to_end(&mut buf) {
Ok(_) => {
return Poll::Ready(Some(Ok(Bytes::from(buf))));
}
Err(e) => return Poll::Ready(Some(Err(e))),
}
}
WrappedBodyInner::Stream(stream) => {
// SAFETY: Pin projection
let stream = unsafe { Pin::new_unchecked(stream) };
futures::Stream::poll_next(stream, cx)
}
}
}
}
impl http_client::HttpClient for ReqwestClient {
fn proxy(&self) -> Option<&http::Uri> {
None
}
fn send(
&self,
req: http::Request<http_client::AsyncBody>,
) -> futures::future::BoxFuture<
'static,
Result<http_client::Response<http_client::AsyncBody>, anyhow::Error>,
> {
let (parts, body) = req.into_parts();
let mut request = self.client.request(parts.method, parts.uri.to_string());
request = request.headers(parts.headers);
if let Some(redirect_policy) = parts.extensions.get::<http_client::RedirectPolicy>() {
request = request.redirect_policy(match redirect_policy {
http_client::RedirectPolicy::NoFollow => reqwest::redirect::Policy::none(),
http_client::RedirectPolicy::FollowLimit(limit) => {
reqwest::redirect::Policy::limited(*limit as usize)
}
http_client::RedirectPolicy::FollowAll => reqwest::redirect::Policy::limited(100),
});
}
if let Some(ReadTimeout(timeout)) = parts.extensions.get::<ReadTimeout>() {
request = request.timeout(*timeout);
}
let body = WrappedBody::new(body);
let request = request.body(reqwest::Body::wrap_stream(body));
async move {
let response = request.send().await.map_err(|e| anyhow!(e))?;
let status = response.status();
let mut builder = http::Response::builder().status(status.as_u16());
for (name, value) in response.headers() {
builder = builder.header(name, value);
}
let bytes = response.bytes_stream();
let bytes = bytes
.map_err(|e| futures::io::Error::new(futures::io::ErrorKind::Other, e))
.into_async_read();
let body = http_client::AsyncBody::from_reader(bytes);
builder.body(body).map_err(|e| anyhow!(e))
}
.boxed()
}
}

2
crates/semantic_index/Cargo.toml
View File

@@ -51,6 +51,7 @@ workspace.workspace = true
worktree.workspace = true
[dev-dependencies]
isahc_http_client.workspace = true
client = { workspace = true, features = ["test-support"] }
env_logger.workspace = true
fs = { workspace = true, features = ["test-support"] }
@@ -61,7 +62,6 @@ language = { workspace = true, features = ["test-support"] }
languages.workspace = true
project = { workspace = true, features = ["test-support"] }
tempfile.workspace = true
ureq_client.workspace = true
util = { workspace = true, features = ["test-support"] }
workspace = { workspace = true, features = ["test-support"] }
worktree = { workspace = true, features = ["test-support"] }

7
crates/semantic_index/examples/index.rs
View File

@@ -2,6 +2,7 @@ use client::Client;
use futures::channel::oneshot;
use gpui::App;
use http_client::HttpClientWithUrl;
use isahc_http_client::IsahcHttpClient;
use language::language_settings::AllLanguageSettings;
use project::Project;
use semantic_index::{OpenAiEmbeddingModel, OpenAiEmbeddingProvider, SemanticDb};
@@ -28,11 +29,7 @@ fn main() {
let clock = Arc::new(FakeSystemClock::default());
let http = Arc::new(HttpClientWithUrl::new(
Arc::new(ureq_client::UreqClient::new(
None,
"Zed semantic index example".to_string(),
cx.background_executor().clone(),
)),
IsahcHttpClient::new(None, None),
"http://localhost:11434",
None,
));

2
crates/storybook/Cargo.toml
View File

@@ -22,6 +22,7 @@ editor.workspace = true
fuzzy.workspace = true
gpui.workspace = true
indoc.workspace = true
isahc_http_client.workspace = true
language.workspace = true
log.workspace = true
menu.workspace = true
@@ -35,7 +36,6 @@ strum = { workspace = true, features = ["derive"] }
theme.workspace = true
title_bar = { workspace = true, features = ["stories"] }
ui = { workspace = true, features = ["stories"] }
ureq_client.workspace = true
[dev-dependencies]
gpui = { workspace = true, features = ["test-support"] }

12
crates/storybook/src/storybook.rs
View File

@@ -4,14 +4,13 @@ mod assets;
mod stories;
mod story_selector;
use std::sync::Arc;
use clap::Parser;
use dialoguer::FuzzySelect;
use gpui::{
div, px, size, AnyView, AppContext, Bounds, Render, ViewContext, VisualContext, WindowBounds,
WindowOptions,
};
use isahc_http_client::IsahcHttpClient;
use log::LevelFilter;
use project::Project;
use settings::{KeymapFile, Settings};
@@ -19,7 +18,6 @@ use simplelog::SimpleLogger;
use strum::IntoEnumIterator;
use theme::{ThemeRegistry, ThemeSettings};
use ui::prelude::*;
use ureq_client::UreqClient;
use crate::app_menus::app_menus;
use crate::assets::Assets;
@@ -68,12 +66,8 @@ fn main() {
gpui::App::new().with_assets(Assets).run(move |cx| {
load_embedded_fonts(cx).unwrap();
let http_client = UreqClient::new(
None,
"zed_storybook".to_string(),
cx.background_executor().clone(),
);
cx.set_http_client(Arc::new(http_client));
let http_client = IsahcHttpClient::new(None, Some("zed_storybook".to_string()));
cx.set_http_client(http_client);
settings::init(cx);
theme::init(theme::LoadThemes::All(Box::new(Assets)), cx);

16
crates/telemetry_events/src/telemetry_events.rs
View File

@@ -91,7 +91,6 @@ impl Display for AssistantPhase {
#[serde(tag = "type")]
pub enum Event {
Editor(EditorEvent),
Copilot(CopilotEvent), // Needed for clients sending old copilot_event types
InlineCompletion(InlineCompletionEvent),
Call(CallEvent),
Assistant(AssistantEvent),
@@ -117,15 +116,9 @@ pub struct EditorEvent {
pub copilot_enabled: bool,
/// Whether the user has copilot enabled for the language of the file opened or saved
pub copilot_enabled_for_language: bool,
}
/// Deprecated since Zed v0.137.0 (2024-05-29). Replaced by InlineCompletionEvent.
// Needed for clients sending old copilot_event types
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CopilotEvent {
pub suggestion_id: Option<String>,
pub suggestion_accepted: bool,
pub file_extension: Option<String>,
/// Whether the client is opening/saving a local file or a remote file via SSH
#[serde(default)]
pub is_via_ssh: bool,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
@@ -182,6 +175,9 @@ pub struct ActionEvent {
pub struct EditEvent {
pub duration: i64,
pub environment: String,
/// Whether the edits occurred locally or remotely via SSH
#[serde(default)]
pub is_via_ssh: bool,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

6
crates/terminal_view/src/terminal_element.rs
View File

@@ -1019,9 +1019,9 @@ impl InputHandler for TerminalInputHandler {
self.workspace
.update(cx, |this, cx| {
cx.invalidate_character_coordinates();
let telemetry = this.project().read(cx).client().telemetry().clone();
telemetry.log_edit_event("terminal");
let project = this.project().read(cx);
let telemetry = project.client().telemetry().clone();
telemetry.log_edit_event("terminal", project.is_via_ssh());
})
.ok();
}

18
crates/text/src/text.rs
View File

@@ -1430,16 +1430,22 @@ impl Buffer {
counts.insert(edit_id, self.undo_map.undo_count(edit_id) + 1);
}
let operation = self.undo_operations(counts);
self.history.push(operation.clone());
operation
}
pub fn undo_operations(&mut self, counts: HashMap<clock::Lamport, u32>) -> Operation {
let timestamp = self.lamport_clock.tick();
let version = self.version();
self.snapshot.version.observe(timestamp);
let undo = UndoOperation {
timestamp: self.lamport_clock.tick(),
version: self.version(),
timestamp,
version,
counts,
};
self.apply_undo(&undo);
self.snapshot.version.observe(undo.timestamp);
let operation = Operation::Undo(undo);
self.history.push(operation.clone());
operation
Operation::Undo(undo)
}
pub fn push_transaction(&mut self, transaction: Transaction, now: Instant) {

1
crates/title_bar/Cargo.toml
View File

@@ -41,6 +41,7 @@ gpui.workspace = true
notifications.workspace = true
project.workspace = true
recent_projects.workspace = true
remote.workspace = true
rpc.workspace = true
serde.workspace = true
smallvec.workspace = true

24
crates/title_bar/src/title_bar.rs
View File

@@ -24,8 +24,8 @@ use smallvec::SmallVec;
use std::sync::Arc;
use theme::ActiveTheme;
use ui::{
h_flex, prelude::*, Avatar, Button, ButtonLike, ButtonStyle, ContextMenu, Icon, IconName,
Indicator, PopoverMenu, Tooltip,
h_flex, prelude::*, Avatar, Button, ButtonLike, ButtonStyle, ContextMenu, Icon,
IconButtonShape, IconName, IconSize, Indicator, PopoverMenu, Tooltip,
};
use util::ResultExt;
use vcs_menu::{BranchList, OpenRecent as ToggleVcsMenu};
@@ -265,25 +265,28 @@ impl TitleBar {
fn render_ssh_project_host(&self, cx: &mut ViewContext<Self>) -> Option<AnyElement> {
let host = self.project.read(cx).ssh_connection_string(cx)?;
let meta = SharedString::from(format!("Connected to: {host}"));
let indicator_color = if self.project.read(cx).ssh_is_connected(cx)? {
Color::Success
} else {
Color::Warning
let indicator_color = match self.project.read(cx).ssh_connection_state(cx)? {
remote::ConnectionState::Connecting => Color::Info,
remote::ConnectionState::Connected => Color::Success,
remote::ConnectionState::HeartbeatMissed => Color::Warning,
remote::ConnectionState::Reconnecting => Color::Warning,
remote::ConnectionState::Disconnected => Color::Error,
};
let indicator = div()
.absolute()
.w_1_4()
.h_1_4()
.size_1p5()
.right_0p5()
.bottom_0p5()
.p_1()
.rounded_2xl()
.rounded_full()
.bg(indicator_color.color(cx));
Some(
div()
.relative()
.child(
IconButton::new("ssh-server-icon", IconName::Server)
.icon_size(IconSize::Small)
.shape(IconButtonShape::Square)
.tooltip(move |cx| {
Tooltip::with_meta(
"Remote Project",
@@ -292,7 +295,6 @@ impl TitleBar {
cx,
)
})
.shape(ui::IconButtonShape::Square)
.on_click(|_, cx| {
cx.dispatch_action(OpenRemote.boxed_clone());
}),

1
crates/ureq_client/LICENSE-GPL
View File

@@ -1 +0,0 @@
../../LICENSE-GPL

24
crates/ureq_client/examples/client.rs
View File

@@ -1,24 +0,0 @@
use futures::AsyncReadExt;
use http_client::{AsyncBody, HttpClient};
use ureq_client::UreqClient;
fn main() {
gpui::App::headless().run(|cx| {
println!("{:?}", std::thread::current().id());
cx.spawn(|cx| async move {
let resp = UreqClient::new(
None,
"Conrad's bot".to_string(),
cx.background_executor().clone(),
)
.get("http://zed.dev", AsyncBody::empty(), true)
.await
.unwrap();
let mut body = String::new();
resp.into_body().read_to_string(&mut body).await.unwrap();
println!("{}", body);
})
.detach();
})
}

187
crates/ureq_client/src/ureq_client.rs
View File

@@ -1,187 +0,0 @@
use std::collections::HashMap;
use std::io::Read;
use std::sync::Arc;
use std::time::Duration;
use std::{pin::Pin, task::Poll};
use anyhow::Error;
use futures::channel::mpsc;
use futures::future::BoxFuture;
use futures::{AsyncRead, SinkExt, StreamExt};
use http_client::{http, AsyncBody, HttpClient, RedirectPolicy, Uri};
use smol::future::FutureExt;
use util::ResultExt;
pub struct UreqClient {
// Note in ureq 2.x the options are stored on the Agent.
// In ureq 3.x we'll be able to set these on the request.
// In practice it's probably "fine" to have many clients, the number of distinct options
// is low; and most requests to the same connection will have the same options so the
// connection pool will work.
clients: Arc<parking_lot::Mutex<HashMap<(Duration, RedirectPolicy), ureq::Agent>>>,
proxy_url: Option<Uri>,
proxy: Option<ureq::Proxy>,
user_agent: String,
background_executor: gpui::BackgroundExecutor,
}
impl UreqClient {
pub fn new(
proxy_url: Option<Uri>,
user_agent: String,
background_executor: gpui::BackgroundExecutor,
) -> Self {
Self {
clients: Arc::default(),
proxy_url: proxy_url.clone(),
proxy: proxy_url.and_then(|url| ureq::Proxy::new(url.to_string()).log_err()),
user_agent,
background_executor,
}
}
fn agent_for(&self, redirect_policy: RedirectPolicy, timeout: Duration) -> ureq::Agent {
let mut clients = self.clients.lock();
// in case our assumption of distinct options is wrong, we'll sporadically clean it out.
if clients.len() > 50 {
clients.clear()
}
clients
.entry((timeout, redirect_policy.clone()))
.or_insert_with(|| {
let mut builder = ureq::AgentBuilder::new()
.timeout_connect(Duration::from_secs(5))
.timeout_read(timeout)
.timeout_write(timeout)
.user_agent(&self.user_agent)
.tls_config(http_client::TLS_CONFIG.clone())
.redirects(match redirect_policy {
RedirectPolicy::NoFollow => 0,
RedirectPolicy::FollowLimit(limit) => limit,
RedirectPolicy::FollowAll => 100,
});
if let Some(proxy) = &self.proxy {
builder = builder.proxy(proxy.clone());
}
builder.build()
})
.clone()
}
}
impl HttpClient for UreqClient {
fn proxy(&self) -> Option<&Uri> {
self.proxy_url.as_ref()
}
fn send(
&self,
request: http::Request<AsyncBody>,
) -> BoxFuture<'static, Result<http::Response<AsyncBody>, Error>> {
let agent = self.agent_for(
request
.extensions()
.get::<RedirectPolicy>()
.cloned()
.unwrap_or_default(),
request
.extensions()
.get::<http_client::ReadTimeout>()
.cloned()
.unwrap_or_default()
.0,
);
let mut req = agent.request(&request.method().as_ref(), &request.uri().to_string());
for (name, value) in request.headers().into_iter() {
req = req.set(name.as_str(), value.to_str().unwrap());
}
let body = request.into_body();
let executor = self.background_executor.clone();
self.background_executor
.spawn(async move {
let response = req.send(body)?;
let mut builder = http::Response::builder()
.status(response.status())
.version(http::Version::HTTP_11);
for name in response.headers_names() {
if let Some(value) = response.header(&name) {
builder = builder.header(name, value);
}
}
let body = AsyncBody::from_reader(UreqResponseReader::new(executor, response));
let http_response = builder.body(body)?;
Ok(http_response)
})
.boxed()
}
}
struct UreqResponseReader {
receiver: mpsc::Receiver<std::io::Result<Vec<u8>>>,
buffer: Vec<u8>,
idx: usize,
_task: gpui::Task<()>,
}
impl UreqResponseReader {
fn new(background_executor: gpui::BackgroundExecutor, response: ureq::Response) -> Self {
let (mut sender, receiver) = mpsc::channel(1);
let mut reader = response.into_reader();
let task = background_executor.spawn(async move {
let mut buffer = vec![0; 8192];
loop {
let n = match reader.read(&mut buffer) {
Ok(0) => break,
Ok(n) => n,
Err(e) => {
let _ = sender.send(Err(e)).await;
break;
}
};
let _ = sender.send(Ok(buffer[..n].to_vec())).await;
}
});
UreqResponseReader {
_task: task,
receiver,
buffer: Vec::new(),
idx: 0,
}
}
}
impl AsyncRead for UreqResponseReader {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &mut [u8],
) -> Poll<std::io::Result<usize>> {
if self.buffer.is_empty() {
match self.receiver.poll_next_unpin(cx) {
Poll::Ready(Some(Ok(data))) => self.buffer = data,
Poll::Ready(Some(Err(e))) => {
return Poll::Ready(Err(e));
}
Poll::Ready(None) => {
return Poll::Ready(Ok(0));
}
Poll::Pending => {
return Poll::Pending;
}
}
}
let n = std::cmp::min(buf.len(), self.buffer.len() - self.idx);
buf[..n].copy_from_slice(&self.buffer[self.idx..self.idx + n]);
self.idx += n;
if self.idx == self.buffer.len() {
self.buffer.clear();
self.idx = 0;
}
Poll::Ready(Ok(n))
}
}

2
crates/vim/Cargo.toml
View File

@@ -17,7 +17,7 @@ neovim = ["nvim-rs", "async-compat", "async-trait", "tokio"]
[dependencies]
anyhow.workspace = true
async-compat = { workspace = true, "optional" = true }
async-compat = { version = "0.2.1", "optional" = true }
async-trait = { workspace = true, "optional" = true }
collections.workspace = true
command_palette.workspace = true

6
crates/workspace/src/workspace.rs
View File

@@ -5605,6 +5605,12 @@ pub fn open_ssh_project(
cx.replace_root_view(|cx| {
let mut workspace =
Workspace::new(Some(workspace_id), project, app_state.clone(), cx);
workspace
.client()
.telemetry()
.report_app_event("open ssh project".to_string());
workspace.set_serialized_ssh_project(serialized_ssh_project);
workspace
});

1
crates/worktree/Cargo.toml
View File

@@ -56,4 +56,5 @@ gpui = { workspace = true, features = ["test-support"] }
http_client.workspace = true
pretty_assertions.workspace = true
rand.workspace = true
rpc = { workspace = true, features = ["test-support"] }
settings = { workspace = true, features = ["test-support"] }

2
crates/zed/Cargo.toml
View File

@@ -58,6 +58,7 @@ http_client.workspace = true
image_viewer.workspace = true
inline_completion_button.workspace = true
install_cli.workspace = true
isahc_http_client.workspace = true
journal.workspace = true
language.workspace = true
language_model.workspace = true
@@ -108,7 +109,6 @@ theme.workspace = true
theme_selector.workspace = true
time.workspace = true
ui.workspace = true
ureq_client.workspace = true
url.workspace = true
urlencoding = "2.1.2"
util.workspace = true

12
crates/zed/src/main.rs
View File

@@ -24,9 +24,9 @@ use gpui::{
UpdateGlobal as _, VisualContext,
};
use http_client::{read_proxy_from_env, Uri};
use isahc_http_client::IsahcHttpClient;
use language::LanguageRegistry;
use log::LevelFilter;
use ureq_client::UreqClient;
use assets::Assets;
use node_runtime::{NodeBinaryOptions, NodeRuntime};
@@ -334,7 +334,9 @@ fn main() {
log::info!("========== starting zed ==========");
let app = App::new().with_assets(Assets);
let app = App::new()
.with_assets(Assets)
.with_http_client(IsahcHttpClient::new(None, None));
let system_id = app.background_executor().block(system_id()).ok();
let installation_id = app.background_executor().block(installation_id()).ok();
@@ -468,8 +470,8 @@ fn main() {
.ok()
})
.or_else(read_proxy_from_env);
let http = UreqClient::new(proxy_url, user_agent, cx.background_executor().clone());
cx.set_http_client(Arc::new(http));
let http = IsahcHttpClient::new(proxy_url, Some(user_agent));
cx.set_http_client(http);
<dyn Fs>::set_global(fs.clone(), cx);
@@ -529,6 +531,8 @@ fn main() {
session_id,
cx,
);
// We should rename these in the future to `first app open`, `first app open for release channel`, and `app open`
if let (Some(system_id), Some(installation_id)) = (&system_id, &installation_id) {
match (&system_id, &installation_id) {
(IdType::New(_), IdType::New(_)) => {

4
docs/src/assistant/configuration.md
View File

@@ -142,8 +142,8 @@ Depending on your hardware or use-case you may wish to limit or increase the con
"low_speed_timeout_in_seconds": 120,
"available_models": [
{
"provider": "ollama",
"name": "mistral:latest",
"name": "qwen2.5-coder",
"display_name": "qwen 2.5 coder 32K",
"max_tokens": 32768
}
]

63
extensions/README.md Normal file
View File

@@ -0,0 +1,63 @@
# Zed Extensions
This directory contains extensions for Zed that are largely maintained by the Zed team. They currently live in the Zed repository for ease of maintenance.
If you are looking for the Zed extension registry, see the [`zed-industries/extensions`](https://github.com/zed-industries/extensions) repo.
## Structure
Currently, Zed includes support for a number of languages without requiring installing an extension. Those languages can be found under [`crates/languages/src`](https://github.com/zed-industries/zed/tree/main/crates/languages/src).
Support for all other languages is done via extensions. This directory ([extensions/](https://github.com/zed-industries/zed/tree/main/extensions/)) contains a number of officially maintained extensions. These extensions use the same [zed_extension_api](https://docs.rs/zed_extension_api/latest/zed_extension_api/) available to all [Zed Extensions](https://zed.dev/extensions) for providing [language servers](https://zed.dev/docs/extensions/languages#language-servers), [tree-sitter grammars](https://zed.dev/docs/extensions/languages#grammar) and [tree-sitter queries](https://zed.dev/docs/extensions/languages#tree-sitter-queries).
## Dev Extensions
See the docs for [Developing an Extension Locally](https://zed.dev/docs/extensions/developing-extensions#developing-an-extension-locally) for how to work with one of these extensions.
## Updating
> [!NOTE]
> This update process is usually handled by Zed staff.
> Community contributors should just submit a PR (step 1) and we'll take it from there.
The process for updating an extension in this directory has three parts.
1. Create a PR with your changes. (Merge it)
2. Bump the extension version in:
- extensions/{language_name}/extension.toml
- extensions/{language_name}/Cargo.toml
- Cargo.lock
You can do this manually, or with a script:
```sh
# Output the current version for a given language
./script/language-extension-version <langname>
# Update the version in `extension.toml` and `Cargo.toml` and trigger a `cargo check`
./script/language-extension-version <langname> <new_version>
```
Commit your changes to a branch, push a PR and merge it.
3. Open a PR to [`zed-industries/extensions`](https://github.com/zed-industries/extensions) repo that updates the extension in question
Edit [`extensions.toml`](https://github.com/zed-industries/extensions/blob/main/extensions.toml) in the extensions repo to reflect the new version you set above and update the submodule latest Zed commit.
```sh
# Go into your clone of the extensions repo
cd ../extensions
# Update
git checkout main
git pull
just init-submodule extensions/zed
# Update the Zed submodule
cd extensions/zed
git checkout main
git pull
cd -
git add extensions.toml extensions/zed
```

2
extensions/csharp/languages/csharp/config.toml
View File

@@ -2,7 +2,7 @@ name = "CSharp"
code_fence_block_name = "csharp"
grammar = "c_sharp"
path_suffixes = ["cs"]
line_comments = ["// "]
line_comments = ["// ", "/// "]
autoclose_before = ";:.,=}])>"
brackets = [
{ start = "{", end = "}", close = true, newline = true },

2
extensions/dart/Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "zed_dart"
version = "0.1.0"
version = "0.1.1"
edition = "2021"
publish = false
license = "Apache-2.0"

2
extensions/dart/extension.toml
View File

@@ -1,7 +1,7 @@
id = "dart"
name = "Dart"
description = "Dart support."
version = "0.1.0"
version = "0.1.1"
schema_version = 1
authors = ["Abdullah Alsigar <abdullah.alsigar@gmail.com>", "Flo <flo80@users.noreply.github.com>", "ybbond <hi@ybbond.id>"]
repository = "https://github.com/zed-industries/zed"

21
extensions/dart/languages/dart/indents.scm
View File

@@ -1,18 +1,3 @@
(class_definition
"class" @context
name: (_) @name) @item
(function_signature
name: (_) @name) @item
(getter_signature
"get" @context
name: (_) @name) @item
(setter_signature
"set" @context
name: (_) @name) @item
(enum_declaration
"enum" @context
name: (_) @name) @item
(_ "[" "]" @end) @indent
(_ "{" "}" @end) @indent
(_ "(" ")" @end) @indent

7
script/bundle-mac
View File

@@ -83,7 +83,10 @@ if [ "$local_arch" = true ]; then
cargo build ${build_flag} --package zed --package cli --package remote_server
else
echo "Compiling zed binaries"
cargo build ${build_flag} --package zed --package cli --package remote_server --target aarch64-apple-darwin --target x86_64-apple-darwin
cargo build ${build_flag} --package zed --package cli --target aarch64-apple-darwin --target x86_64-apple-darwin
# Build remote_server in separate invocation to prevent feature unification from other crates
# from influencing dynamic libraries required by it.
cargo build ${build_flag} --package remote_server --target aarch64-apple-darwin --target x86_64-apple-darwin
fi
echo "Creating application bundle"
@@ -358,7 +361,7 @@ function sign_binary() {
if [[ $can_code_sign = true ]]; then
echo "Code signing executable $binary_path"
/usr/bin/codesign --deep --force --timestamp --options runtime --sign "$IDENTITY" "${binary_path}" -v
/usr/bin/codesign --deep --force --timestamp --options runtime --entitlements crates/zed/resources/zed.entitlements --sign "$IDENTITY" "${binary_path}" -v
fi
}

29
script/language-extension-version Executable file
View File

@@ -0,0 +1,29 @@
#!/usr/bin/env bash
set -euox pipefail
if [ "$#" -lt 1 ]; then
echo "Usage: $0 <language> [version]"
exit 1
fi
LANGUAGE=$1
VERSION=${2:-}
EXTENSION_DIR="extensions/$LANGUAGE"
EXTENSION_TOML="$EXTENSION_DIR/extension.toml"
CARGO_TOML="$EXTENSION_DIR/Cargo.toml"
if [ ! -d "$EXTENSION_DIR" ]; then
echo "Directory $EXTENSION_DIR does not exist."
exit 1
fi
if [ -z "$VERSION" ]; then
grep -m 1 'version =' "$EXTENSION_TOML" | awk -F\" '{print $2}'
exit 0
fi
sed -i '' -e "s/^version = \".*\"/version = \"$VERSION\"/" "$EXTENSION_TOML"
sed -i '' -e "s/^version = \".*\"/version = \"$VERSION\"/" "$CARGO_TOML"
cargo check