acp_thread: Log token usage when receiving `StopReason::MaxTokens` (#49343)

Bennet Bo Fenner created 1 month ago

Before you mark this PR as ready for review, make sure that you have:
- [ ] Added a solid test coverage and/or screenshots from doing manual
testing
- [x] Done a self-review taking into account security and performance
aspects
- [x] Aligned any UI changes with the [UI
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)

Release Notes:

- N/A

Change summary

Cargo.lock                          | 1 +
crates/acp_thread/Cargo.toml        | 1 +
crates/acp_thread/src/acp_thread.rs | 3 +++
crates/agent/src/tests/mod.rs       | 4 ++++
crates/agent/src/thread.rs          | 1 +
5 files changed, 10 insertions(+)

Detailed changes

Cargo.lock

@@ -22,6 +22,7 @@ dependencies = [
  "itertools 0.14.0",
  "language",
  "language_model",
+ "log",
  "markdown",
  "multi_buffer",
  "parking_lot",

crates/acp_thread/Cargo.toml

@@ -30,6 +30,7 @@ gpui.workspace = true
 itertools.workspace = true
 language.workspace = true
 language_model.workspace = true
+log.workspace = true
 markdown.workspace = true
 parking_lot = { workspace = true, optional = true }
 image = { workspace = true, optional = true }

crates/acp_thread/src/acp_thread.rs

@@ -885,6 +885,7 @@ pub struct TokenUsage {
     pub used_tokens: u64,
     pub input_tokens: u64,
     pub output_tokens: u64,
+    pub max_output_tokens: Option<u64>,
 }
 
 impl TokenUsage {
@@ -1953,11 +1954,13 @@ impl AcpThread {
                     Ok(Err(e)) => {
                         this.send_task.take();
                         cx.emit(AcpThreadEvent::Error);
+                        log::error!("Error in run turn: {:?}", e);
                         Err(e)
                     }
                     Ok(Ok(r)) if r.stop_reason == acp::StopReason::MaxTokens => {
                         this.send_task.take();
                         cx.emit(AcpThreadEvent::Error);
+                        log::error!("Max tokens reached. Usage: {:?}", this.token_usage);
                         Err(anyhow!("Max tokens reached"))
                     }
                     result => {

crates/agent/src/tests/mod.rs

@@ -2736,6 +2736,7 @@ async fn test_truncate_first_message(cx: &mut TestAppContext) {
             Some(acp_thread::TokenUsage {
                 used_tokens: 32_000 + 16_000,
                 max_tokens: 1_000_000,
+                max_output_tokens: None,
                 input_tokens: 32_000,
                 output_tokens: 16_000,
             })
@@ -2797,6 +2798,7 @@ async fn test_truncate_first_message(cx: &mut TestAppContext) {
             Some(acp_thread::TokenUsage {
                 used_tokens: 40_000 + 20_000,
                 max_tokens: 1_000_000,
+                max_output_tokens: None,
                 input_tokens: 40_000,
                 output_tokens: 20_000,
             })
@@ -2847,6 +2849,7 @@ async fn test_truncate_second_message(cx: &mut TestAppContext) {
                 Some(acp_thread::TokenUsage {
                     used_tokens: 32_000 + 16_000,
                     max_tokens: 1_000_000,
+                    max_output_tokens: None,
                     input_tokens: 32_000,
                     output_tokens: 16_000,
                 })
@@ -2903,6 +2906,7 @@ async fn test_truncate_second_message(cx: &mut TestAppContext) {
             Some(acp_thread::TokenUsage {
                 used_tokens: 40_000 + 20_000,
                 max_tokens: 1_000_000,
+                max_output_tokens: None,
                 input_tokens: 40_000,
                 output_tokens: 20_000,
             })

crates/agent/src/thread.rs

@@ -1507,6 +1507,7 @@ impl Thread {
         let model = self.model.clone()?;
         Some(acp_thread::TokenUsage {
             max_tokens: model.max_token_count(),
+            max_output_tokens: model.max_output_tokens(),
             used_tokens: usage.total_tokens(),
             input_tokens: usage.input_tokens,
             output_tokens: usage.output_tokens,