History

Split token display for OpenAI (#46829)

This feature cost $15. Up -> Tokens we're sending to the model Down -> Tokens we've received from the model. <img width="377" height="69" alt="Screenshot 2026-01-14 at 12 31 01 PM" src="https://github.com/user-attachments/assets/fc15824f-de5d-466b-8cc1-329f3c1940bb" /> Release Notes: - Changed the display of tokens for OpenAI models to reflect the input/output limits. ---------

Mikayla Maki committed 10 months ago

Commit 9c5fc6e1 parent b5b13d18 files · +137 −37Browse files

Unified · Split

8 files
  1. crates/acp_thread/src/acp_thread.rs +1 −0
  2. crates/agent/src/tests/mod.rs +4 −0
  3. crates/agent/src/thread.rs +1 −0
  4. crates/agent_ui/src/acp/thread_view.rs +108 −37
  5. crates/language_model/src/language_model.rs +6 −0
  6. crates/language_models/src/provider/cloud.rs +9 −0
  7. crates/language_models/src/provider/open_ai.rs +4 −0
  8. crates/language_models/src/provider/open_ai_compatible.rs +4 −0
crates/acp_thread/src/acp_thread.rs+1−0
@@ -904,6 +904,7 @@ impl PlanEntry { · 903 unmodified lines
904pub struct TokenUsage {904pub struct TokenUsage {
905 pub max_tokens: u64,905 pub max_tokens: u64,
906 pub used_tokens: u64,906 pub used_tokens: u64,
907 pub input_tokens: u64,
907 pub output_tokens: u64,908 pub output_tokens: u64,
908}909}
909 910
crates/agent/src/tests/mod.rs+4−0
@@ -2516,6 +2516,7 @@ async fn test_truncate_first_message(cx: &mut TestAppContext) { · 2515 unmodified lines
2516 Some(acp_thread::TokenUsage {2516 Some(acp_thread::TokenUsage {
2517 used_tokens: 32_000 + 16_000,2517 used_tokens: 32_000 + 16_000,
2518 max_tokens: 1_000_000,2518 max_tokens: 1_000_000,
2519 input_tokens: 32_000,
2519 output_tokens: 16_000,2520 output_tokens: 16_000,
2520 })2521 })
2521 );2522 );
@@ -2576,6 +2577,7 @@ async fn test_truncate_first_message(cx: &mut TestAppContext) { · 54 unmodified lines
2576 Some(acp_thread::TokenUsage {2577 Some(acp_thread::TokenUsage {
2577 used_tokens: 40_000 + 20_000,2578 used_tokens: 40_000 + 20_000,
2578 max_tokens: 1_000_000,2579 max_tokens: 1_000_000,
2580 input_tokens: 40_000,
2579 output_tokens: 20_000,2581 output_tokens: 20_000,
2580 })2582 })
2581 );2583 );
@@ -2625,6 +2627,7 @@ async fn test_truncate_second_message(cx: &mut TestAppContext) { · 43 unmodified lines
2625 Some(acp_thread::TokenUsage {2627 Some(acp_thread::TokenUsage {
2626 used_tokens: 32_000 + 16_000,2628 used_tokens: 32_000 + 16_000,
2627 max_tokens: 1_000_000,2629 max_tokens: 1_000_000,
2630 input_tokens: 32_000,
2628 output_tokens: 16_000,2631 output_tokens: 16_000,
2629 })2632 })
2630 );2633 );
@@ -2680,6 +2683,7 @@ async fn test_truncate_second_message(cx: &mut TestAppContext) { · 49 unmodified lines
2680 Some(acp_thread::TokenUsage {2683 Some(acp_thread::TokenUsage {
2681 used_tokens: 40_000 + 20_000,2684 used_tokens: 40_000 + 20_000,
2682 max_tokens: 1_000_000,2685 max_tokens: 1_000_000,
2686 input_tokens: 40_000,
2683 output_tokens: 20_000,2687 output_tokens: 20_000,
2684 })2688 })
2685 );2689 );
crates/agent/src/thread.rs+1−0
@@ -1291,6 +1291,7 @@ impl Thread { · 1290 unmodified lines
1291 Some(acp_thread::TokenUsage {1291 Some(acp_thread::TokenUsage {
1292 max_tokens: model.max_token_count_for_mode(self.completion_mode.into()),1292 max_tokens: model.max_token_count_for_mode(self.completion_mode.into()),
1293 used_tokens: usage.total_tokens(),1293 used_tokens: usage.total_tokens(),
1294 input_tokens: usage.input_tokens,
1294 output_tokens: usage.output_tokens,1295 output_tokens: usage.output_tokens,
1295 })1296 })
1296 }1297 }
crates/agent_ui/src/acp/thread_view.rs+108−37
@@ -79,10 +79,9 @@ use crate::{ · 78 unmodified lines
79 ToggleBurnMode, ToggleProfileSelector,79 ToggleBurnMode, ToggleProfileSelector,
80};80};
81 81
82/// Maximum number of lines to show for a collapsed terminal command preview.
83const MAX_COLLAPSED_LINES: usize = 3;82const MAX_COLLAPSED_LINES: usize = 3;
84const STOPWATCH_THRESHOLD: Duration = Duration::from_secs(1);83const STOPWATCH_THRESHOLD: Duration = Duration::from_secs(30);
85const TOKEN_THRESHOLD: u64 = 1;84const TOKEN_THRESHOLD: u64 = 250;
86 85
87#[derive(Copy, Clone, Debug, PartialEq, Eq)]86#[derive(Copy, Clone, Debug, PartialEq, Eq)]
88enum ThreadFeedback {87enum ThreadFeedback {
@@ -6058,46 +6057,118 @@ impl AcpThreadView { · 5969 unmodified lines
6058 .is_some_and(|model| model.provider_id() == language_model::ZED_CLOUD_PROVIDER_ID)6057 .is_some_and(|model| model.provider_id() == language_model::ZED_CLOUD_PROVIDER_ID)
6059 }6058 }
6060 6059
6060 fn supports_split_token_display(&self, cx: &App) -> bool {
6061 self.as_native_thread(cx)
6062 .and_then(|thread| thread.read(cx).model())
6063 .is_some_and(|model| model.supports_split_token_display())
6064 }
6065
6061 fn render_token_usage(&self, cx: &mut Context<Self>) -> Option<Div> {6066 fn render_token_usage(&self, cx: &mut Context<Self>) -> Option<Div> {
6062 let thread = self.thread()?.read(cx);6067 let thread = self.thread()?.read(cx);
6063 let usage = thread.token_usage()?;6068 let usage = thread.token_usage()?;
6064 let is_generating = thread.status() != ThreadStatus::Idle;6069 let is_generating = thread.status() != ThreadStatus::Idle;
6070 let show_split = self.supports_split_token_display(cx);
6065 6071
6066 let used = crate::text_thread_editor::humanize_token_count(usage.used_tokens);6072 let separator_color = Color::Custom(cx.theme().colors().text_muted.opacity(0.5));
6067 let max = crate::text_thread_editor::humanize_token_count(usage.max_tokens);6073 let token_label = |text: String, animation_id: &'static str| {
6074 Label::new(text)
6075 .size(LabelSize::Small)
6076 .color(Color::Muted)
6077 .map(|label| {
6078 if is_generating {
6079 label
6080 .with_animation(
6081 animation_id,
6082 Animation::new(Duration::from_secs(2))
6083 .repeat()
6084 .with_easing(pulsating_between(0.3, 0.8)),
6085 |label, delta| label.alpha(delta),
6086 )
6087 .into_any()
6088 } else {
6089 label.into_any_element()
6090 }
6091 })
6092 };
6068 6093
6069 Some(6094 if show_split {
6070 h_flex()6095 let max_output_tokens = self
6071 .flex_shrink_0()6096 .as_native_thread(cx)
6072 .gap_0p5()6097 .and_then(|thread| thread.read(cx).model())
6073 .mr_1p5()6098 .and_then(|model| model.max_output_tokens())
6074 .child(6099 .unwrap_or(0);
6075 Label::new(used)6100
6076 .size(LabelSize::Small)6101 let input = crate::text_thread_editor::humanize_token_count(usage.input_tokens);
6077 .color(Color::Muted)6102 let input_max = crate::text_thread_editor::humanize_token_count(
6078 .map(|label| {6103 usage.max_tokens.saturating_sub(max_output_tokens),
6079 if is_generating {6104 );
6080 label6105 let output = crate::text_thread_editor::humanize_token_count(usage.output_tokens);
6081 .with_animation(6106 let output_max = crate::text_thread_editor::humanize_token_count(max_output_tokens);
6082 "used-tokens-label",6107
6083 Animation::new(Duration::from_secs(2))6108 Some(
6084 .repeat()6109 h_flex()
6085 .with_easing(pulsating_between(0.3, 0.8)),6110 .flex_shrink_0()
6086 |label, delta| label.alpha(delta),6111 .gap_1()
6087 )6112 .mr_1p5()
6088 .into_any()6113 .child(
6089 } else {6114 h_flex()
6090 label.into_any_element()6115 .gap_0p5()
6091 }6116 .child(
6092 }),6117 Icon::new(IconName::ArrowUp)
6093 )6118 .size(IconSize::XSmall)
6094 .child(6119 .color(Color::Muted),
6095 Label::new("/")6120 )
6096 .size(LabelSize::Small)6121 .child(token_label(input, "input-tokens-label"))
6097 .color(Color::Custom(cx.theme().colors().text_muted.opacity(0.5))),6122 .child(
6098 )6123 Label::new("/")
6099 .child(Label::new(max).size(LabelSize::Small).color(Color::Muted)),6124 .size(LabelSize::Small)
6100 )6125 .color(separator_color),
6126 )
6127 .child(
6128 Label::new(input_max)
6129 .size(LabelSize::Small)
6130 .color(Color::Muted),
6131 ),
6132 )
6133 .child(
6134 h_flex()
6135 .gap_0p5()
6136 .child(
6137 Icon::new(IconName::ArrowDown)
6138 .size(IconSize::XSmall)
6139 .color(Color::Muted),
6140 )
6141 .child(token_label(output, "output-tokens-label"))
6142 .child(
6143 Label::new("/")
6144 .size(LabelSize::Small)
6145 .color(separator_color),
6146 )
6147 .child(
6148 Label::new(output_max)
6149 .size(LabelSize::Small)
6150 .color(Color::Muted),
6151 ),
6152 ),
6153 )
6154 } else {
6155 let used = crate::text_thread_editor::humanize_token_count(usage.used_tokens);
6156 let max = crate::text_thread_editor::humanize_token_count(usage.max_tokens);
6157
6158 Some(
6159 h_flex()
6160 .flex_shrink_0()
6161 .gap_0p5()
6162 .mr_1p5()
6163 .child(token_label(used, "used-tokens-label"))
6164 .child(
6165 Label::new("/")
6166 .size(LabelSize::Small)
6167 .color(separator_color),
6168 )
6169 .child(Label::new(max).size(LabelSize::Small).color(Color::Muted)),
6170 )
6171 }
6101 }6172 }
6102 6173
6103 fn toggle_burn_mode(6174 fn toggle_burn_mode(
crates/language_model/src/language_model.rs+6−0
@@ -617,6 +617,12 @@ pub trait LanguageModel: Send + Sync { · 616 unmodified lines
617 false617 false
618 }618 }
619 619
620 /// Returns whether this model/provider reports accurate split input/output token counts.
621 /// When true, the UI may show separate input/output token indicators.
622 fn supports_split_token_display(&self) -> bool {
623 false
624 }
625
620 fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {626 fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
621 LanguageModelToolSchemaFormat::JsonSchema627 LanguageModelToolSchemaFormat::JsonSchema
622 }628 }
crates/language_models/src/provider/cloud.rs+9−0
@@ -624,6 +624,11 @@ impl LanguageModel for CloudLanguageModel { · 623 unmodified lines
624 self.model.supports_max_mode624 self.model.supports_max_mode
625 }625 }
626 626
627 fn supports_split_token_display(&self) -> bool {
628 use cloud_llm_client::LanguageModelProvider::*;
629 matches!(self.model.provider, OpenAi)
630 }
631
627 fn telemetry_id(&self) -> String {632 fn telemetry_id(&self) -> String {
628 format!("zed.dev/{}", self.model.id)633 format!("zed.dev/{}", self.model.id)
629 }634 }
@@ -652,6 +657,10 @@ impl LanguageModel for CloudLanguageModel { · 22 unmodified lines
652 .map(|max_token_count| max_token_count as u64)657 .map(|max_token_count| max_token_count as u64)
653 }658 }
654 659
660 fn max_output_tokens(&self) -> Option<u64> {
661 Some(self.model.max_output_tokens as u64)
662 }
663
655 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {664 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
656 match &self.model.provider {665 match &self.model.provider {
657 cloud_llm_client::LanguageModelProvider::Anthropic => {666 cloud_llm_client::LanguageModelProvider::Anthropic => {
crates/language_models/src/provider/open_ai.rs+4−0
@@ -329,6 +329,10 @@ impl LanguageModel for OpenAiLanguageModel { · 328 unmodified lines
329 }329 }
330 }330 }
331 331
332 fn supports_split_token_display(&self) -> bool {
333 true
334 }
335
332 fn telemetry_id(&self) -> String {336 fn telemetry_id(&self) -> String {
333 format!("openai/{}", self.model.id())337 format!("openai/{}", self.model.id())
334 }338 }
crates/language_models/src/provider/open_ai_compatible.rs+4−0
@@ -319,6 +319,10 @@ impl LanguageModel for OpenAiCompatibleLanguageModel { · 318 unmodified lines
319 }319 }
320 }320 }
321 321
322 fn supports_split_token_display(&self) -> bool {
323 true
324 }
325
322 fn telemetry_id(&self) -> String {326 fn telemetry_id(&self) -> String {
323 format!("openai/{}", self.model.name)327 format!("openai/{}", self.model.name)
324 }328 }