From 052e81a7f08b3f288d4dd74c85165386d999eb01 Mon Sep 17 00:00:00 2001
From: hahuy <quanghuyha098@gmail.com>
Date: Sat, 11 Apr 2026 22:32:31 +0700
Subject: [PATCH] acp_thread: Forward token usage from PromptResponse to thread
 UI
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Extract r.usage from each PromptResponse turn and call update_token_usage()
so the token count ring in the agent panel shows live data.

Guard against max_tokens == 0 (ACP protocol has no context-window field)
by rendering a plain "N tokens" label instead of a broken "0% • N / 0" ring.

Add two gpui tests: usage populated and usage absent.
---
 crates/acp_thread/src/acp_thread.rs           | 89 +++++++++++++++++++
 .../src/conversation_view/thread_view.rs      | 17 ++++
 2 files changed, 106 insertions(+)

diff --git a/crates/acp_thread/src/acp_thread.rs b/crates/acp_thread/src/acp_thread.rs
index 7fb48c132f971f..5fe8835b0bbba1 100644
--- a/crates/acp_thread/src/acp_thread.rs
+++ b/crates/acp_thread/src/acp_thread.rs
@@ -2261,6 +2261,19 @@ impl AcpThread {
                     Ok(r) => {
                         Self::flush_streaming_text(&mut this.streaming_text_buffer, cx);
 
+                        if let Some(usage) = &r.usage {
+                            this.update_token_usage(
+                                Some(TokenUsage {
+                                    max_tokens: 0,
+                                    used_tokens: usage.total_tokens,
+                                    input_tokens: usage.input_tokens,
+                                    output_tokens: usage.output_tokens,
+                                    max_output_tokens: None,
+                                }),
+                                cx,
+                            );
+                        }
+
                         if r.stop_reason == acp::StopReason::MaxTokens {
                             this.had_error = true;
                             cx.emit(AcpThreadEvent::Error);
@@ -5247,4 +5260,80 @@ mod tests {
             "session info title update should not propagate back to the connection"
         );
     }
+
+    #[gpui::test]
+    async fn test_token_usage_from_prompt_response(cx: &mut gpui::TestAppContext) {
+        init_test(cx);
+
+        let fs = FakeFs::new(cx.executor());
+        let project = Project::test(fs, [], cx).await;
+        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
+            |_, _, _| {
+                async move {
+                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)
+                        .usage(acp::Usage::new(1500, 1000, 500)))
+                }
+                .boxed_local()
+            },
+        ));
+
+        let thread = cx
+            .update(|cx| {
+                connection.new_session(project, PathList::new(&[Path::new(path!("/test"))]), cx)
+            })
+            .await
+            .unwrap();
+
+        assert!(
+            thread.read_with(cx, |thread, _| thread.token_usage().cloned()).is_none(),
+            "token_usage should be None before any turn"
+        );
+
+        thread
+            .update(cx, |thread, cx| thread.send_raw("Hello!", cx))
+            .await
+            .unwrap();
+
+        thread.read_with(cx, |thread, _| {
+            let usage = thread.token_usage().expect("token_usage should be set after a turn");
+            assert_eq!(usage.used_tokens, 1500);
+            assert_eq!(usage.input_tokens, 1000);
+            assert_eq!(usage.output_tokens, 500);
+            assert_eq!(usage.max_tokens, 0);
+            assert_eq!(usage.max_output_tokens, None);
+        });
+    }
+
+    #[gpui::test]
+    async fn test_token_usage_absent_when_response_has_no_usage(cx: &mut gpui::TestAppContext) {
+        init_test(cx);
+
+        let fs = FakeFs::new(cx.executor());
+        let project = Project::test(fs, [], cx).await;
+        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
+            |_, _, _| {
+                async move { Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)) }
+                    .boxed_local()
+            },
+        ));
+
+        let thread = cx
+            .update(|cx| {
+                connection.new_session(project, PathList::new(&[Path::new(path!("/test"))]), cx)
+            })
+            .await
+            .unwrap();
+
+        thread
+            .update(cx, |thread, cx| thread.send_raw("Hello!", cx))
+            .await
+            .unwrap();
+
+        thread.read_with(cx, |thread, _| {
+            assert!(
+                thread.token_usage().is_none(),
+                "token_usage should remain None when PromptResponse has no usage"
+            );
+        });
+    }
 }
diff --git a/crates/agent_ui/src/conversation_view/thread_view.rs b/crates/agent_ui/src/conversation_view/thread_view.rs
index 412778e054cab1..98b50a34926768 100644
--- a/crates/agent_ui/src/conversation_view/thread_view.rs
+++ b/crates/agent_ui/src/conversation_view/thread_view.rs
@@ -3448,6 +3448,23 @@ impl ThreadView {
     fn render_token_usage(&self, cx: &mut Context<Self>) -> Option<impl IntoElement> {
         let thread = self.thread.read(cx);
         let usage = thread.token_usage()?;
+
+        if usage.max_tokens == 0 {
+            let label = crate::humanize_token_count(usage.used_tokens);
+            return Some(
+                h_flex()
+                    .id("token_count_label")
+                    .mt_px()
+                    .mr_1()
+                    .child(
+                        Label::new(format!("{label} tokens"))
+                            .size(LabelSize::Small)
+                            .color(Color::Muted),
+                    )
+                    .into_any_element(),
+            );
+        }
+
         let show_split = self.supports_split_token_display(cx);
 
         let progress_color = |ratio: f32| -> Hsla {