agent_ui: Fix max tokens error not being shown (#49098)

Bennet Bo Fenner created 2 months ago

Closes #ISSUE

Before you mark this PR as ready for review, make sure that you have:
- [x] 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:

- Fixed an issue where no error was shown when an LLM request exceeded
the maximum tokens supported by the model

Change summary

crates/acp_thread/src/acp_thread.rs    |  5 ++
crates/agent_ui/src/acp/thread_view.rs | 52 ++++++++++++++++++++++++++++
2 files changed, 57 insertions(+)

Detailed changes

crates/acp_thread/src/acp_thread.rs

@@ -1955,6 +1955,11 @@ impl AcpThread {
                         cx.emit(AcpThreadEvent::Error);
                         Err(e)
                     }
+                    Ok(Ok(r)) if r.stop_reason == acp::StopReason::MaxTokens => {
+                        this.send_task.take();
+                        cx.emit(AcpThreadEvent::Error);
+                        Err(anyhow!("Max tokens reached"))
+                    }
                     result => {
                         let canceled = matches!(
                             result,

crates/agent_ui/src/acp/thread_view.rs

@@ -5825,4 +5825,56 @@ pub(crate) mod tests {
             );
         });
     }
+
+    #[gpui::test]
+    async fn test_max_tokens_error_is_rendered(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let connection = StubAgentConnection::new();
+
+        let (thread_view, cx) =
+            setup_thread_view(StubAgentServer::new(connection.clone()), cx).await;
+
+        let message_editor = message_editor(&thread_view, cx);
+        message_editor.update_in(cx, |editor, window, cx| {
+            editor.set_text("Some prompt", window, cx);
+        });
+        active_thread(&thread_view, cx).update_in(cx, |view, window, cx| view.send(window, cx));
+
+        let session_id = thread_view.read_with(cx, |view, cx| {
+            view.active_thread()
+                .unwrap()
+                .read(cx)
+                .thread
+                .read(cx)
+                .session_id()
+                .clone()
+        });
+
+        cx.run_until_parked();
+
+        cx.update(|_, _cx| {
+            connection.end_turn(session_id, acp::StopReason::MaxTokens);
+        });
+
+        cx.run_until_parked();
+
+        thread_view.read_with(cx, |thread_view, cx| {
+            let state = thread_view.active_thread().unwrap();
+            let error = &state.read(cx).thread_error;
+            match error {
+                Some(ThreadError::Other { message, .. }) => {
+                    assert!(
+                        message.contains("Max tokens reached"),
+                        "Expected 'Max tokens reached' error, got: {}",
+                        message
+                    );
+                }
+                other => panic!(
+                    "Expected ThreadError::Other with 'Max tokens reached', got: {:?}",
+                    other.is_some()
+                ),
+            }
+        });
+    }
 }