Skip to content

Commit 30d4256

Browse files
anyideaanyidea
authored
Merge branch 'zed-industries:main' into main
2 parents 13d1efa + 5611cc2 commit 30d4256

61 files changed

Lines changed: 4451 additions & 3942 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎Cargo.lock‎

Lines changed: 17 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ members = [
5757
"crates/edit_prediction",
5858
"crates/edit_prediction_cli",
5959
"crates/edit_prediction_context",
60+
"crates/edit_prediction_metrics",
6061
"crates/edit_prediction_types",
6162
"crates/edit_prediction_ui",
6263
"crates/editor",
@@ -480,6 +481,7 @@ zed_actions = { path = "crates/zed_actions" }
480481
zed_credentials_provider = { path = "crates/zed_credentials_provider" }
481482
zed_env_vars = { path = "crates/zed_env_vars" }
482483
edit_prediction = { path = "crates/edit_prediction" }
484+
edit_prediction_metrics = { path = "crates/edit_prediction_metrics" }
483485
zeta_prompt = { path = "crates/zeta_prompt" }
484486
zlog = { path = "crates/zlog" }
485487
zlog_settings = { path = "crates/zlog_settings" }

‎crates/acp_thread/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ anyhow.workspace = true
2323
buffer_diff.workspace = true
2424
chrono.workspace = true
2525
collections.workspace = true
26+
feature_flags.workspace = true
2627
multi_buffer.workspace = true
2728
file_icons.workspace = true
2829
futures.workspace = true

‎crates/acp_thread/src/acp_thread.rs‎

Lines changed: 207 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use anyhow::{Context as _, Result, anyhow};
88
use collections::HashSet;
99
pub use connection::*;
1010
pub use diff::*;
11+
use feature_flags::{AcpBetaFeatureFlag, FeatureFlagAppExt as _};
1112
use futures::{FutureExt, channel::oneshot, future::BoxFuture};
1213
use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
1314
use itertools::Itertools;
@@ -972,7 +973,7 @@ impl PlanEntry {
972973
}
973974
}
974975

975-
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
976+
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
976977
pub struct TokenUsage {
977978
pub max_tokens: u64,
978979
pub used_tokens: u64,
@@ -981,6 +982,12 @@ pub struct TokenUsage {
981982
pub max_output_tokens: Option<u64>,
982983
}
983984

985+
#[derive(Debug, Clone)]
986+
pub struct SessionCost {
987+
pub amount: f64,
988+
pub currency: SharedString,
989+
}
990+
984991
pub const TOKEN_USAGE_WARNING_THRESHOLD: f32 = 0.8;
985992

986993
impl TokenUsage {
@@ -1043,6 +1050,7 @@ pub struct AcpThread {
10431050
running_turn: Option<RunningTurn>,
10441051
connection: Rc<dyn AgentConnection>,
10451052
token_usage: Option<TokenUsage>,
1053+
cost: Option<SessionCost>,
10461054
prompt_capabilities: acp::PromptCapabilities,
10471055
available_commands: Vec<acp::AvailableCommand>,
10481056
_observe_prompt_capabilities: Task<anyhow::Result<()>>,
@@ -1232,6 +1240,7 @@ impl AcpThread {
12321240
connection,
12331241
session_id,
12341242
token_usage: None,
1243+
cost: None,
12351244
prompt_capabilities,
12361245
available_commands: Vec::new(),
12371246
_observe_prompt_capabilities: task,
@@ -1348,6 +1357,10 @@ impl AcpThread {
13481357
self.token_usage.as_ref()
13491358
}
13501359

1360+
pub fn cost(&self) -> Option<&SessionCost> {
1361+
self.cost.as_ref()
1362+
}
1363+
13511364
pub fn has_pending_edit_tool_calls(&self) -> bool {
13521365
for entry in self.entries.iter().rev() {
13531366
match entry {
@@ -1463,6 +1476,18 @@ impl AcpThread {
14631476
config_options,
14641477
..
14651478
}) => cx.emit(AcpThreadEvent::ConfigOptionsUpdated(config_options)),
1479+
acp::SessionUpdate::UsageUpdate(update) if cx.has_flag::<AcpBetaFeatureFlag>() => {
1480+
let usage = self.token_usage.get_or_insert_with(Default::default);
1481+
usage.max_tokens = update.size;
1482+
usage.used_tokens = update.used;
1483+
if let Some(cost) = update.cost {
1484+
self.cost = Some(SessionCost {
1485+
amount: cost.amount,
1486+
currency: cost.currency.into(),
1487+
});
1488+
}
1489+
cx.emit(AcpThreadEvent::TokenUsageUpdated);
1490+
}
14661491
_ => {}
14671492
}
14681493
Ok(())
@@ -1759,6 +1784,9 @@ impl AcpThread {
17591784
}
17601785

17611786
pub fn update_token_usage(&mut self, usage: Option<TokenUsage>, cx: &mut Context<Self>) {
1787+
if usage.is_none() {
1788+
self.cost = None;
1789+
}
17621790
self.token_usage = usage;
17631791
cx.emit(AcpThreadEvent::TokenUsageUpdated);
17641792
}
@@ -2340,6 +2368,15 @@ impl AcpThread {
23402368
}
23412369
}
23422370

2371+
if cx.has_flag::<AcpBetaFeatureFlag>()
2372+
&& let Some(response_usage) = &r.usage
2373+
{
2374+
let usage = this.token_usage.get_or_insert_with(Default::default);
2375+
usage.input_tokens = response_usage.input_tokens;
2376+
usage.output_tokens = response_usage.output_tokens;
2377+
cx.emit(AcpThreadEvent::TokenUsageUpdated);
2378+
}
2379+
23432380
cx.emit(AcpThreadEvent::Stopped(r.stop_reason));
23442381
Ok(Some(r))
23452382
}
@@ -5297,4 +5334,173 @@ mod tests {
52975334
"session info title update should not propagate back to the connection"
52985335
);
52995336
}
5337+
5338+
#[gpui::test]
5339+
async fn test_usage_update_populates_token_usage_and_cost(cx: &mut TestAppContext) {
5340+
init_test(cx);
5341+
5342+
let fs = FakeFs::new(cx.executor());
5343+
let project = Project::test(fs, [], cx).await;
5344+
let connection = Rc::new(FakeAgentConnection::new());
5345+
let thread = cx
5346+
.update(|cx| {
5347+
connection.new_session(project, PathList::new(&[Path::new(path!("/test"))]), cx)
5348+
})
5349+
.await
5350+
.unwrap();
5351+
5352+
thread.update(cx, |thread, cx| {
5353+
thread
5354+
.handle_session_update(
5355+
acp::SessionUpdate::UsageUpdate(
5356+
acp::UsageUpdate::new(5000, 10000).cost(acp::Cost::new(0.42, "USD")),
5357+
),
5358+
cx,
5359+
)
5360+
.unwrap();
5361+
});
5362+
5363+
thread.read_with(cx, |thread, _| {
5364+
let usage = thread.token_usage().expect("token_usage should be set");
5365+
assert_eq!(usage.max_tokens, 10000);
5366+
assert_eq!(usage.used_tokens, 5000);
5367+
5368+
let cost = thread.cost().expect("cost should be set");
5369+
assert!((cost.amount - 0.42).abs() < f64::EPSILON);
5370+
assert_eq!(cost.currency.as_ref(), "USD");
5371+
});
5372+
}
5373+
5374+
#[gpui::test]
5375+
async fn test_usage_update_without_cost_preserves_existing_cost(cx: &mut TestAppContext) {
5376+
init_test(cx);
5377+
5378+
let fs = FakeFs::new(cx.executor());
5379+
let project = Project::test(fs, [], cx).await;
5380+
let connection = Rc::new(FakeAgentConnection::new());
5381+
let thread = cx
5382+
.update(|cx| {
5383+
connection.new_session(project, PathList::new(&[Path::new(path!("/test"))]), cx)
5384+
})
5385+
.await
5386+
.unwrap();
5387+
5388+
thread.update(cx, |thread, cx| {
5389+
thread
5390+
.handle_session_update(
5391+
acp::SessionUpdate::UsageUpdate(
5392+
acp::UsageUpdate::new(1000, 10000).cost(acp::Cost::new(0.10, "USD")),
5393+
),
5394+
cx,
5395+
)
5396+
.unwrap();
5397+
5398+
thread
5399+
.handle_session_update(
5400+
acp::SessionUpdate::UsageUpdate(acp::UsageUpdate::new(2000, 10000)),
5401+
cx,
5402+
)
5403+
.unwrap();
5404+
});
5405+
5406+
thread.read_with(cx, |thread, _| {
5407+
let usage = thread.token_usage().expect("token_usage should be set");
5408+
assert_eq!(usage.used_tokens, 2000);
5409+
5410+
let cost = thread.cost().expect("cost should be preserved");
5411+
assert!((cost.amount - 0.10).abs() < f64::EPSILON);
5412+
});
5413+
}
5414+
5415+
#[gpui::test]
5416+
async fn test_response_usage_does_not_clobber_session_usage(cx: &mut TestAppContext) {
5417+
init_test(cx);
5418+
5419+
let fs = FakeFs::new(cx.executor());
5420+
let project = Project::test(fs, [], cx).await;
5421+
let connection = Rc::new(FakeAgentConnection::new().on_user_message(
5422+
move |_, thread, mut cx| {
5423+
async move {
5424+
thread.update(&mut cx, |thread, cx| {
5425+
thread
5426+
.handle_session_update(
5427+
acp::SessionUpdate::UsageUpdate(
5428+
acp::UsageUpdate::new(3000, 10000)
5429+
.cost(acp::Cost::new(0.05, "EUR")),
5430+
),
5431+
cx,
5432+
)
5433+
.unwrap();
5434+
})?;
5435+
Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)
5436+
.usage(acp::Usage::new(500, 200, 300)))
5437+
}
5438+
.boxed_local()
5439+
},
5440+
));
5441+
5442+
let thread = cx
5443+
.update(|cx| {
5444+
connection.new_session(project, PathList::new(&[Path::new(path!("/test"))]), cx)
5445+
})
5446+
.await
5447+
.unwrap();
5448+
5449+
thread
5450+
.update(cx, |thread, cx| thread.send_raw("hello", cx))
5451+
.await
5452+
.unwrap();
5453+
5454+
thread.read_with(cx, |thread, _| {
5455+
let usage = thread.token_usage().expect("token_usage should be set");
5456+
assert_eq!(usage.max_tokens, 10000, "max_tokens from UsageUpdate");
5457+
assert_eq!(usage.used_tokens, 3000, "used_tokens from UsageUpdate");
5458+
assert_eq!(usage.input_tokens, 200, "input_tokens from response usage");
5459+
assert_eq!(
5460+
usage.output_tokens, 300,
5461+
"output_tokens from response usage"
5462+
);
5463+
5464+
let cost = thread.cost().expect("cost should be set");
5465+
assert!((cost.amount - 0.05).abs() < f64::EPSILON);
5466+
assert_eq!(cost.currency.as_ref(), "EUR");
5467+
});
5468+
}
5469+
5470+
#[gpui::test]
5471+
async fn test_clearing_token_usage_also_clears_cost(cx: &mut TestAppContext) {
5472+
init_test(cx);
5473+
5474+
let fs = FakeFs::new(cx.executor());
5475+
let project = Project::test(fs, [], cx).await;
5476+
let connection = Rc::new(FakeAgentConnection::new());
5477+
let thread = cx
5478+
.update(|cx| {
5479+
connection.new_session(project, PathList::new(&[Path::new(path!("/test"))]), cx)
5480+
})
5481+
.await
5482+
.unwrap();
5483+
5484+
thread.update(cx, |thread, cx| {
5485+
thread
5486+
.handle_session_update(
5487+
acp::SessionUpdate::UsageUpdate(
5488+
acp::UsageUpdate::new(1000, 10000).cost(acp::Cost::new(0.25, "USD")),
5489+
),
5490+
cx,
5491+
)
5492+
.unwrap();
5493+
5494+
assert!(thread.token_usage().is_some());
5495+
assert!(thread.cost().is_some());
5496+
5497+
thread.update_token_usage(None, cx);
5498+
5499+
assert!(thread.token_usage().is_none());
5500+
assert!(
5501+
thread.cost().is_none(),
5502+
"cost should be cleared when token usage is cleared"
5503+
);
5504+
});
5505+
}
53005506
}

0 commit comments

Comments
 (0)