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
904904 pub struct TokenUsage {
905905 pub max_tokens: u64,
906906 pub used_tokens: u64,
907+ pub input_tokens: u64,
907908 pub output_tokens: u64,
908909 }
909910
crates/agent/src/tests/mod.rs+4−0
@@ -2516,6 +2516,7 @@ async fn test_truncate_first_message(cx: &mut TestAppContext) { · 2515 unmodified lines
25162516 Some(acp_thread::TokenUsage {
25172517 used_tokens: 32_000 + 16_000,
25182518 max_tokens: 1_000_000,
2519+ input_tokens: 32_000,
25192520 output_tokens: 16_000,
25202521 })
25212522 );
@@ -2576,6 +2577,7 @@ async fn test_truncate_first_message(cx: &mut TestAppContext) { · 54 unmodified lines
25762577 Some(acp_thread::TokenUsage {
25772578 used_tokens: 40_000 + 20_000,
25782579 max_tokens: 1_000_000,
2580+ input_tokens: 40_000,
25792581 output_tokens: 20_000,
25802582 })
25812583 );
@@ -2625,6 +2627,7 @@ async fn test_truncate_second_message(cx: &mut TestAppContext) { · 43 unmodified lines
26252627 Some(acp_thread::TokenUsage {
26262628 used_tokens: 32_000 + 16_000,
26272629 max_tokens: 1_000_000,
2630+ input_tokens: 32_000,
26282631 output_tokens: 16_000,
26292632 })
26302633 );
@@ -2680,6 +2683,7 @@ async fn test_truncate_second_message(cx: &mut TestAppContext) { · 49 unmodified lines
26802683 Some(acp_thread::TokenUsage {
26812684 used_tokens: 40_000 + 20_000,
26822685 max_tokens: 1_000_000,
2686+ input_tokens: 40_000,
26832687 output_tokens: 20_000,
26842688 })
26852689 );
crates/agent/src/thread.rs+1−0
@@ -1291,6 +1291,7 @@ impl Thread { · 1290 unmodified lines
12911291 Some(acp_thread::TokenUsage {
12921292 max_tokens: model.max_token_count_for_mode(self.completion_mode.into()),
12931293 used_tokens: usage.total_tokens(),
1294+ input_tokens: usage.input_tokens,
12941295 output_tokens: usage.output_tokens,
12951296 })
12961297 }
crates/agent_ui/src/acp/thread_view.rs+108−37
@@ -79,10 +79,9 @@ use crate::{ · 78 unmodified lines
7979 ToggleBurnMode, ToggleProfileSelector,
8080 };
8181
82−/// Maximum number of lines to show for a collapsed terminal command preview.
8382 const MAX_COLLAPSED_LINES: usize = 3;
84−const STOPWATCH_THRESHOLD: Duration = Duration::from_secs(1);
85−const TOKEN_THRESHOLD: u64 = 1;
83+const STOPWATCH_THRESHOLD: Duration = Duration::from_secs(30);
84+const TOKEN_THRESHOLD: u64 = 250;
8685
8786 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
8887 enum ThreadFeedback {
@@ -6058,46 +6057,118 @@ impl AcpThreadView { · 5969 unmodified lines
60586057 .is_some_and(|model| model.provider_id() == language_model::ZED_CLOUD_PROVIDER_ID)
60596058 }
60606059
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+
60616066 fn render_token_usage(&self, cx: &mut Context<Self>) -> Option<Div> {
60626067 let thread = self.thread()?.read(cx);
60636068 let usage = thread.token_usage()?;
60646069 let is_generating = thread.status() != ThreadStatus::Idle;
6070+ let show_split = self.supports_split_token_display(cx);
60656071
6066− let used = crate::text_thread_editor::humanize_token_count(usage.used_tokens);
6067− let max = crate::text_thread_editor::humanize_token_count(usage.max_tokens);
6072+ let separator_color = Color::Custom(cx.theme().colors().text_muted.opacity(0.5));
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+ };
60686093
6069− Some(
6070− h_flex()
6071− .flex_shrink_0()
6072− .gap_0p5()
6073− .mr_1p5()
6074− .child(
6075− Label::new(used)
6076− .size(LabelSize::Small)
6077− .color(Color::Muted)
6078− .map(|label| {
6079− if is_generating {
6080− label
6081− .with_animation(
6082− "used-tokens-label",
6083− Animation::new(Duration::from_secs(2))
6084− .repeat()
6085− .with_easing(pulsating_between(0.3, 0.8)),
6086− |label, delta| label.alpha(delta),
6087− )
6088− .into_any()
6089− } else {
6090− label.into_any_element()
6091− }
6092− }),
6093− )
6094− .child(
6095− Label::new("/")
6096− .size(LabelSize::Small)
6097− .color(Color::Custom(cx.theme().colors().text_muted.opacity(0.5))),
6098− )
6099− .child(Label::new(max).size(LabelSize::Small).color(Color::Muted)),
6100− )
6094+ if show_split {
6095+ let max_output_tokens = self
6096+ .as_native_thread(cx)
6097+ .and_then(|thread| thread.read(cx).model())
6098+ .and_then(|model| model.max_output_tokens())
6099+ .unwrap_or(0);
6100+
6101+ let input = crate::text_thread_editor::humanize_token_count(usage.input_tokens);
6102+ let input_max = crate::text_thread_editor::humanize_token_count(
6103+ usage.max_tokens.saturating_sub(max_output_tokens),
6104+ );
6105+ let output = crate::text_thread_editor::humanize_token_count(usage.output_tokens);
6106+ let output_max = crate::text_thread_editor::humanize_token_count(max_output_tokens);
6107+
6108+ Some(
6109+ h_flex()
6110+ .flex_shrink_0()
6111+ .gap_1()
6112+ .mr_1p5()
6113+ .child(
6114+ h_flex()
6115+ .gap_0p5()
6116+ .child(
6117+ Icon::new(IconName::ArrowUp)
6118+ .size(IconSize::XSmall)
6119+ .color(Color::Muted),
6120+ )
6121+ .child(token_label(input, "input-tokens-label"))
6122+ .child(
6123+ Label::new("/")
6124+ .size(LabelSize::Small)
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+ }
61016172 }
61026173
61036174 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
617617 false
618618 }
619619
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+
620626 fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
621627 LanguageModelToolSchemaFormat::JsonSchema
622628 }
crates/language_models/src/provider/cloud.rs+9−0
@@ -624,6 +624,11 @@ impl LanguageModel for CloudLanguageModel { · 623 unmodified lines
624624 self.model.supports_max_mode
625625 }
626626
627+ fn supports_split_token_display(&self) -> bool {
628+ use cloud_llm_client::LanguageModelProvider::*;
629+ matches!(self.model.provider, OpenAi)
630+ }
631+
627632 fn telemetry_id(&self) -> String {
628633 format!("zed.dev/{}", self.model.id)
629634 }
@@ -652,6 +657,10 @@ impl LanguageModel for CloudLanguageModel { · 22 unmodified lines
652657 .map(|max_token_count| max_token_count as u64)
653658 }
654659
660+ fn max_output_tokens(&self) -> Option<u64> {
661+ Some(self.model.max_output_tokens as u64)
662+ }
663+
655664 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
656665 match &self.model.provider {
657666 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
329329 }
330330 }
331331
332+ fn supports_split_token_display(&self) -> bool {
333+ true
334+ }
335+
332336 fn telemetry_id(&self) -> String {
333337 format!("openai/{}", self.model.id())
334338 }
crates/language_models/src/provider/open_ai_compatible.rs+4−0
@@ -319,6 +319,10 @@ impl LanguageModel for OpenAiCompatibleLanguageModel { · 318 unmodified lines
319319 }
320320 }
321321
322+ fn supports_split_token_display(&self) -> bool {
323+ true
324+ }
325+
322326 fn telemetry_id(&self) -> String {
323327 format!("openai/{}", self.model.name)
324328 }