From fbda709ff083c5b92be38b1c32dab5119bbe2e08 Mon Sep 17 00:00:00 2001 From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:48:31 +0000 Subject: [PATCH] agent_ui: Fix inline assistant when reasoning precedes tool use (#61220) The inline assistant's streaming-tool path waited for the first `rewrite_section` output, but treated any other completion event as terminal. Responses API providers can emit `ReasoningDetails`, `Thinking`, and other metadata before the tool call. This stopped the stream early and marked generation as done without applying a replacement. EOF was not handled explicitly. Parse errors, provider errors, incomplete tool calls, and truncated responses could also be reported as successful completions. Continue consuming informational events until a supported tool call arrives. Track the rewrite tool call's streaming state, reject a second `rewrite_section` call instead of concatenating unrelated outputs, accept `failure_message` only after its input is complete, and reject a `rewrite_section` arriving after a failure message (and vice versa). Stop, EOF, malformed tool input, and stream errors now surface as `CodegenStatus::Error` instead of silently completing. A completed tool call followed by `Stop(MaxTokens)` or `Stop(Refusal)` is also an error: the OpenAI Responses mapper flushes truncated tool calls as complete before emitting the stop reason, so the stop reason must be checked even when the tool input parsed. Completed failure messages are stored before `Finished` is emitted, and a failure message left over from a previous generation is now cleared when a new one starts. Add regression coverage for metadata around streamed rewrites, Stop and EOF without a tool call, incomplete rewrites, interrupted responses, multiple rewrite calls, failure messages, malformed input, mixed tool calls, and stream errors before and after rewrite output begins. --- Closes #52714 Release Notes: - Fixed the inline assistant doing nothing when the model emits reasoning before its tool call. --------- Co-authored-by: Claude Fable 5 Co-authored-by: Smit Barmase --- crates/agent_ui/src/buffer_codegen.rs | 976 +++++++++++++++++++++----- 1 file changed, 814 insertions(+), 162 deletions(-) diff --git a/crates/agent_ui/src/buffer_codegen.rs b/crates/agent_ui/src/buffer_codegen.rs index 022bf5c658..5782f7a086 100644 --- a/crates/agent_ui/src/buffer_codegen.rs +++ b/crates/agent_ui/src/buffer_codegen.rs @@ -1,7 +1,7 @@ use crate::{context::LoadedContext, inline_prompt_editor::CodegenStatus}; use agent_settings::AgentSettings; -use anyhow::{Context as _, Result}; -use collections::{HashMap, HashSet}; +use anyhow::{Context as _, Result, anyhow}; +use collections::HashSet; use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint}; use futures::FutureExt; use futures::{ @@ -19,7 +19,7 @@ use language_model::{ CompletionIntent, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelTextStream, LanguageModelToolChoice, - LanguageModelToolUse, LanguageModelToolUseId, Role, TokenUsage, + LanguageModelToolUse, LanguageModelToolUseId, Role, StopReason, TokenUsage, }; use language_models::provider::anthropic::telemetry::{ AnthropicCompletionType, AnthropicEventData, AnthropicEventReporter, AnthropicEventType, @@ -415,6 +415,7 @@ impl CodegenAlternative { ) -> Result<()> { // Clear the model explanation since the user has started a new generation. self.description = None; + self.failure = None; if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() { self.buffer.update(cx, |buffer, cx| { @@ -1163,162 +1164,293 @@ impl CodegenAlternative { enum ToolUseOutput { Rewrite { text: String, - description: Option, }, - Failure(String), + Failure { + message: String, + is_input_complete: bool, + }, } - enum ModelUpdate { - Description(String), - Failure(String), + struct RewriteState { + tool_use_id: LanguageModelToolUseId, + chars_read: usize, + is_input_complete: bool, } - let chars_read_by_tool_id: Arc>> = - Arc::new(Mutex::new(HashMap::default())); - let process_tool_use = move |tool_use: LanguageModelToolUse| -> Option { - let mut chars_read_by_tool_id = chars_read_by_tool_id.lock(); - match tool_use.name.as_ref() { - REWRITE_SECTION_TOOL_NAME => { - let Ok(input) = tool_use.input.parse::() else { - return None; - }; - let chars_read_so_far = - chars_read_by_tool_id.entry(tool_use.id).or_insert(0); - let Some(text_slice) = input.replacement_text.get(*chars_read_so_far..) - else { - return None; - }; - let text = text_slice.to_string(); - *chars_read_so_far = input.replacement_text.len(); - Some(ToolUseOutput::Rewrite { - text, - description: None, - }) + let rewrite_state: Arc>> = Arc::new(Mutex::new(None)); + let process_rewrite_state = rewrite_state.clone(); + let process_tool_use = + move |tool_use: LanguageModelToolUse| -> Result> { + match tool_use.name.as_ref() { + REWRITE_SECTION_TOOL_NAME => { + let input = match tool_use.input.parse::() { + Ok(input) => input, + Err(_) if !tool_use.is_input_complete => return Ok(None), + Err(error) => { + return Err(error.context(format!( + "invalid input for `{REWRITE_SECTION_TOOL_NAME}` tool" + ))); + } + }; + let mut rewrite_state = process_rewrite_state.lock(); + if let Some(state) = rewrite_state.as_ref() + && state.tool_use_id != tool_use.id + { + return Err(anyhow!( + "Inline assistant called `{REWRITE_SECTION_TOOL_NAME}` \ + multiple times" + )); + } + let state = rewrite_state.get_or_insert_with(|| RewriteState { + tool_use_id: tool_use.id, + chars_read: 0, + is_input_complete: false, + }); + let Some(text_slice) = input.replacement_text.get(state.chars_read..) + else { + return Ok(None); + }; + let text = text_slice.to_string(); + state.chars_read = input.replacement_text.len(); + state.is_input_complete |= tool_use.is_input_complete; + Ok(Some(ToolUseOutput::Rewrite { text })) + } + FAILURE_MESSAGE_TOOL_NAME => { + let mut input = match tool_use.input.parse::() { + Ok(input) => input, + Err(_) if !tool_use.is_input_complete => return Ok(None), + Err(error) => { + return Err(error.context(format!( + "invalid input for `{FAILURE_MESSAGE_TOOL_NAME}` tool" + ))); + } + }; + Ok(Some(ToolUseOutput::Failure { + message: std::mem::take(&mut input.message), + is_input_complete: tool_use.is_input_complete, + })) + } + _ => Ok(None), } - FAILURE_MESSAGE_TOOL_NAME => { - let Ok(mut input) = tool_use.input.parse::() else { - return None; + }; + + let mut message_id = None; + let mut failure_started = false; + let mut failure_complete = false; + let last_token_usage = Arc::new(Mutex::new(TokenUsage::default())); + + let first_text = loop { + let Some(event) = completion_events.next().await else { + if failure_complete { + finish_with_status(CodegenStatus::Done, cx); + } else { + finish_with_status( + CodegenStatus::Error(anyhow!( + "Inline assistant response ended without completing \ + `{REWRITE_SECTION_TOOL_NAME}` or `{FAILURE_MESSAGE_TOOL_NAME}`" + )), + cx, + ); + } + return; + }; + + match event { + Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => { + message_id = Some(id); + } + Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => { + let output = match process_tool_use(tool_use) { + Ok(output) => output, + Err(error) => { + finish_with_status(CodegenStatus::Error(error), cx); + return; + } }; - Some(ToolUseOutput::Failure(std::mem::take(&mut input.message))) + if let Some(output) = output { + match output { + ToolUseOutput::Rewrite { text } => { + if failure_started { + finish_with_status( + CodegenStatus::Error(anyhow!( + "Inline assistant called \ + `{REWRITE_SECTION_TOOL_NAME}` after \ + `{FAILURE_MESSAGE_TOOL_NAME}`" + )), + cx, + ); + return; + } + break text; + } + ToolUseOutput::Failure { + message, + is_input_complete, + } => { + failure_started = true; + failure_complete |= is_input_complete; + if is_input_complete + && codegen + .update(cx, |this, cx| { + this.failure = Some(message); + cx.notify(); + }) + .is_err() + { + return; + } + } + } + } + } + Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => { + *last_token_usage.lock() = token_usage; + } + Ok(LanguageModelCompletionEvent::Stop(reason)) => { + if !failure_complete { + finish_with_status( + CodegenStatus::Error(anyhow!( + "Inline assistant stopped before completing \ + `{REWRITE_SECTION_TOOL_NAME}` or \ + `{FAILURE_MESSAGE_TOOL_NAME}` ({reason:?})" + )), + cx, + ); + } else if matches!(reason, StopReason::MaxTokens | StopReason::Refusal) { + finish_with_status( + CodegenStatus::Error(anyhow!( + "Inline assistant response was interrupted ({reason:?})" + )), + cx, + ); + } else { + finish_with_status(CodegenStatus::Done, cx); + } + return; + } + Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { + tool_name, + json_parse_error, + .. + }) => { + finish_with_status( + CodegenStatus::Error(anyhow!( + "Invalid input for inline assistant tool \ + `{tool_name}`: {json_parse_error}" + )), + cx, + ); + return; + } + Ok( + LanguageModelCompletionEvent::Queued { .. } + | LanguageModelCompletionEvent::Started + | LanguageModelCompletionEvent::Text(_) + | LanguageModelCompletionEvent::Thinking { .. } + | LanguageModelCompletionEvent::RedactedThinking { .. } + | LanguageModelCompletionEvent::ReasoningDetails(_) + | LanguageModelCompletionEvent::Compaction(_), + ) => {} + Err(error) => { + finish_with_status(CodegenStatus::Error(error.into()), cx); + return; } - _ => None, } }; - let (message_tx, mut message_rx) = futures::channel::mpsc::unbounded::(); + let move_last_token_usage = last_token_usage.clone(); + let move_rewrite_state = rewrite_state.clone(); - cx.spawn({ - let codegen = codegen.clone(); - async move |cx| { - while let Some(update) = message_rx.next().await { - let _ = codegen.update(cx, |this, _cx| match update { - ModelUpdate::Description(d) => this.description = Some(d), - ModelUpdate::Failure(f) => this.failure = Some(f), - }); - } - } - }) - .detach(); - - let mut message_id = None; - let mut first_text = None; - let last_token_usage = Arc::new(Mutex::new(TokenUsage::default())); - let total_text = Arc::new(Mutex::new(String::new())); - - loop { - if let Some(first_event) = completion_events.next().await { - match first_event { - Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => { - message_id = Some(id); - } + let remaining_text_stream = completion_events.filter_map(move |e| { + let process_tool_use = process_tool_use.clone(); + let last_token_usage = move_last_token_usage.clone(); + let rewrite_state = move_rewrite_state.clone(); + async move { + match e { Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => { - if let Some(output) = process_tool_use(tool_use) { - let (text, update) = match output { - ToolUseOutput::Rewrite { text, description } => { - (Some(text), description.map(ModelUpdate::Description)) - } - ToolUseOutput::Failure(message) => { - (None, Some(ModelUpdate::Failure(message))) - } - }; - if let Some(update) = update { - let _ = message_tx.unbounded_send(update); - } - first_text = text; - if first_text.is_some() { - break; - } + let output = match process_tool_use(tool_use) { + Ok(Some(output)) => output, + Ok(None) => return None, + Err(error) => return Some(Err(error.into())), + }; + match output { + ToolUseOutput::Rewrite { text } => Some(Ok(text)), + ToolUseOutput::Failure { .. } => Some(Err(anyhow!( + "Inline assistant called `{FAILURE_MESSAGE_TOOL_NAME}` \ + after starting `{REWRITE_SECTION_TOOL_NAME}`" + ) + .into())), } } Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => { *last_token_usage.lock() = token_usage; + None } - Ok(LanguageModelCompletionEvent::Text(text)) => { - let mut lock = total_text.lock(); - lock.push_str(&text); - } - Ok(e) => { - log::warn!("Unexpected event: {:?}", e); - break; - } - Err(e) => { - finish_with_status(CodegenStatus::Error(e.into()), cx); - break; + Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { + tool_name, + json_parse_error, + .. + }) => Some(Err(anyhow!( + "Invalid input for inline assistant tool \ + `{tool_name}`: {json_parse_error}" + ) + .into())), + Ok(LanguageModelCompletionEvent::Stop(reason)) => { + let rewrite_complete = rewrite_state + .lock() + .as_ref() + .is_some_and(|state| state.is_input_complete); + if !rewrite_complete { + Some(Err(anyhow!( + "Inline assistant stopped before completing \ + `{REWRITE_SECTION_TOOL_NAME}` ({reason:?})" + ) + .into())) + } else if matches!(reason, StopReason::MaxTokens | StopReason::Refusal) + { + Some(Err(anyhow!( + "Inline assistant response was interrupted ({reason:?})" + ) + .into())) + } else { + None + } } + Ok( + LanguageModelCompletionEvent::Queued { .. } + | LanguageModelCompletionEvent::Started + | LanguageModelCompletionEvent::Text(_) + | LanguageModelCompletionEvent::Thinking { .. } + | LanguageModelCompletionEvent::RedactedThinking { .. } + | LanguageModelCompletionEvent::StartMessage { .. } + | LanguageModelCompletionEvent::ReasoningDetails(_) + | LanguageModelCompletionEvent::Compaction(_), + ) => None, + Err(error) => Some(Err(error)), } } - } - - let Some(first_text) = first_text else { - finish_with_status(CodegenStatus::Done, cx); - return; - }; - - let move_last_token_usage = last_token_usage.clone(); - - let text_stream = Box::pin(futures::stream::once(async { Ok(first_text) }).chain( - completion_events.filter_map(move |e| { - let process_tool_use = process_tool_use.clone(); - let last_token_usage = move_last_token_usage.clone(); - let total_text = total_text.clone(); - let mut message_tx = message_tx.clone(); - async move { - match e { - Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => { - let Some(output) = process_tool_use(tool_use) else { - return None; - }; - let (text, update) = match output { - ToolUseOutput::Rewrite { text, description } => { - (Some(text), description.map(ModelUpdate::Description)) - } - ToolUseOutput::Failure(message) => { - (None, Some(ModelUpdate::Failure(message))) - } - }; - if let Some(update) = update { - let _ = message_tx.send(update).await; - } - text.map(Ok) - } - Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => { - *last_token_usage.lock() = token_usage; - None - } - Ok(LanguageModelCompletionEvent::Text(text)) => { - let mut lock = total_text.lock(); - lock.push_str(&text); - None - } - Ok(LanguageModelCompletionEvent::Stop(_reason)) => None, - e => { - log::error!("UNEXPECTED EVENT {:?}", e); - None - } - } - } - }), - )); + }); + let completion_status = futures::stream::once(async move { + if rewrite_state + .lock() + .as_ref() + .is_some_and(|state| state.is_input_complete) + { + None + } else { + Some(Err(anyhow!( + "Inline assistant response ended before completing \ + `{REWRITE_SECTION_TOOL_NAME}`" + ) + .into())) + } + }) + .filter_map(futures::future::ready); + let text_stream = Box::pin( + futures::stream::once(async { Ok(first_text) }) + .chain(remaining_text_stream) + .chain(completion_status), + ); let language_model_text_stream = LanguageModelTextStream { message_id: message_id, @@ -1870,36 +2002,78 @@ mod tests { // than the first would cause an index-out-of-bounds panic because the // chars_read_so_far counter was shared across all tool use IDs. #[gpui::test] - async fn test_separate_tool_uses_have_independent_char_counters(cx: &mut TestAppContext) { - init_test(cx); - let buffer = cx.new(|cx| Buffer::local("", cx)); - let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx)); - let range = buffer.read_with(cx, |buffer, cx| { - let snapshot = buffer.snapshot(cx); - snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(0, 0)) - }); - let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap()); - let codegen = cx.new(|cx| { - CodegenAlternative::new( - buffer.clone(), - range.clone(), - true, - prompt_builder, - Uuid::new_v4(), - cx, - ) - }); - + async fn test_tool_completion_rejects_multiple_rewrite_tool_calls(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); let events_tx = simulate_tool_based_completion(&codegen, cx); - // tool_1 has longer text; tool_2 has shorter text. With the old shared - // counter, processing tool_2 would attempt replacement_text[N..] where - // N > replacement_text.len(), panicking with index out of bounds. + events_tx .unbounded_send(rewrite_tool_use("tool_1", "longer replacement text", true)) .unwrap(); events_tx .unbounded_send(rewrite_tool_use("tool_2", "short", true)) .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some("Inline assistant called `rewrite_section` multiple times") + ); + } + + #[gpui::test] + async fn test_tool_completion_ignores_non_tool_events(cx: &mut TestAppContext) { + let (buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(LanguageModelCompletionEvent::Queued { position: 1 }) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Started) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::StartMessage { + message_id: "message_1".into(), + }) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::ReasoningDetails( + serde_json::json!({ "encrypted_content": "reasoning" }), + )) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Text( + "I'll update the section.".into(), + )) + .unwrap(); + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "replace", false)) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Thinking { + text: "Finishing the rewrite".into(), + signature: None, + }) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::RedactedThinking { + data: "redacted".into(), + }) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { + output_tokens: 1, + ..TokenUsage::default() + })) + .unwrap(); + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "replacement", true)) + .unwrap(); events_tx .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)) .unwrap(); @@ -1908,7 +2082,431 @@ mod tests { assert_eq!( buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()), - "longer replacement textshort" + "replacement" + ); + assert!(codegen.read_with(cx, |codegen, _cx| matches!( + &codegen.status, + CodegenStatus::Done + ))); + } + + #[gpui::test] + async fn test_tool_completion_errors_when_stopped_without_tool_call(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(LanguageModelCompletionEvent::ReasoningDetails( + serde_json::json!([]), + )) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some( + "Inline assistant stopped before completing `rewrite_section` or \ + `failure_message` (EndTurn)" + ) + ); + } + + #[gpui::test] + async fn test_tool_completion_errors_on_eof_without_tool_call(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some( + "Inline assistant response ended without completing `rewrite_section` or \ + `failure_message`" + ) + ); + } + + #[gpui::test] + async fn test_tool_completion_accepts_failure_message(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(failure_tool_use("The selection cannot be rewritten.", true)) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + assert!(codegen.read_with(cx, |codegen, _cx| matches!( + &codegen.status, + CodegenStatus::Done + ))); + assert_eq!( + codegen.read_with(cx, |codegen, _cx| codegen.current_failure()), + Some("The selection cannot be rewritten.".into()) + ); + } + + #[gpui::test] + async fn test_tool_completion_errors_on_invalid_tool_json(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(LanguageModelCompletionEvent::ToolUseJsonParseError { + id: "tool_1".into(), + tool_name: REWRITE_SECTION_TOOL_NAME.into(), + raw_input: r#"{"replacement_text":"#.into(), + json_parse_error: "EOF while parsing a string".into(), + }) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some( + "Invalid input for inline assistant tool `rewrite_section`: \ + EOF while parsing a string" + ) + ); + } + + #[gpui::test] + async fn test_tool_completion_errors_when_incomplete_rewrite_stops(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "partial", false)) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::MaxTokens)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some( + "Inline assistant stopped before completing `rewrite_section` \ + (MaxTokens)" + ) + ); + } + + #[gpui::test] + async fn test_tool_completion_errors_when_incomplete_rewrite_reaches_eof( + cx: &mut TestAppContext, + ) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "partial", false)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some("Inline assistant response ended before completing `rewrite_section`") + ); + } + + #[gpui::test] + async fn test_tool_completion_errors_when_interrupted_after_complete_rewrite( + cx: &mut TestAppContext, + ) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "replacement", true)) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::MaxTokens)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some("Inline assistant response was interrupted (MaxTokens)") + ); + } + + #[gpui::test] + async fn test_tool_completion_errors_when_interrupted_after_failure_message( + cx: &mut TestAppContext, + ) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(failure_tool_use("The selection cannot be rewritten.", true)) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::MaxTokens)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some("Inline assistant response was interrupted (MaxTokens)") + ); + } + + #[gpui::test] + async fn test_tool_completion_requires_failure_message_to_complete(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(failure_tool_use("Partial failure", false)) + .unwrap(); + events_tx + .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some( + "Inline assistant stopped before completing `rewrite_section` or \ + `failure_message` (EndTurn)" + ) + ); + assert_eq!( + codegen.read_with(cx, |codegen, _cx| codegen.current_failure()), + None + ); + } + + #[gpui::test] + async fn test_tool_completion_rejects_failure_after_rewrite(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "replacement", true)) + .unwrap(); + events_tx + .unbounded_send(failure_tool_use("Cannot rewrite selection", true)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some( + "Inline assistant called `failure_message` after starting \ + `rewrite_section`" + ) + ); + } + + #[gpui::test] + async fn test_tool_completion_rejects_rewrite_after_failure(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(failure_tool_use("Cannot rewrite selection", true)) + .unwrap(); + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "replacement", true)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some("Inline assistant called `rewrite_section` after `failure_message`") + ); + } + + #[gpui::test] + async fn test_tool_completion_rejects_rewrite_after_partial_failure(cx: &mut TestAppContext) { + let (buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(failure_tool_use("Cannot rewrite", false)) + .unwrap(); + events_tx + .unbounded_send(rewrite_tool_use("tool_1", "replacement", true)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!( + error.as_deref(), + Some("Inline assistant called `rewrite_section` after `failure_message`") + ); + assert_eq!( + buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()), + "" + ); + } + + #[gpui::test] + async fn test_tool_completion_propagates_stream_error_before_rewrite(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion_results(&codegen, cx); + + events_tx + .unbounded_send(Err(anyhow!("provider stream failed").into())) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!(error.as_deref(), Some("provider stream failed")); + } + + #[gpui::test] + async fn test_tool_completion_propagates_stream_error_after_rewrite(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion_results(&codegen, cx); + + events_tx + .unbounded_send(Ok(rewrite_tool_use("tool_1", "replacement", true))) + .unwrap(); + events_tx + .unbounded_send(Err(anyhow!("provider stream failed").into())) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert_eq!(error.as_deref(), Some("provider stream failed")); + } + + #[gpui::test] + async fn test_tool_completion_errors_on_invalid_tool_input(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(LanguageModelCompletionEvent::ToolUse( + LanguageModelToolUse { + id: "tool_1".into(), + name: REWRITE_SECTION_TOOL_NAME.into(), + raw_input: r#"{"replacement_text":42}"#.into(), + input: language_model::LanguageModelToolUseInput::Json( + serde_json::json!({ "replacement_text": 42 }), + ), + is_input_complete: true, + thought_signature: None, + }, + )) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + + let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status { + CodegenStatus::Error(error) => Some(error.to_string()), + _ => None, + }); + assert!( + error + .as_deref() + .is_some_and(|error| error.starts_with("invalid input for `rewrite_section` tool")) + ); + } + + #[gpui::test] + async fn test_start_clears_previous_failure(cx: &mut TestAppContext) { + let (_buffer, codegen) = new_tool_based_codegen(cx); + cx.update(|cx| AgentSettings::register(cx)); + let events_tx = simulate_tool_based_completion(&codegen, cx); + + events_tx + .unbounded_send(failure_tool_use("Cannot rewrite selection", true)) + .unwrap(); + drop(events_tx); + cx.run_until_parked(); + assert_eq!( + codegen.read_with(cx, |codegen, _cx| codegen.current_failure()), + Some("Cannot rewrite selection".into()) + ); + + codegen.update(cx, |codegen, cx| { + codegen + .start( + "Rewrite the selection".into(), + Task::ready(None).shared(), + Arc::new(FakeLanguageModel::default()), + cx, + ) + .unwrap(); + }); + + assert_eq!( + codegen.read_with(cx, |codegen, _cx| codegen.current_failure()), + None ); } @@ -1957,6 +2555,30 @@ mod tests { cx.set_global(cx.update(SettingsStore::test)); } + fn new_tool_based_codegen( + cx: &mut TestAppContext, + ) -> (Entity, Entity) { + init_test(cx); + let buffer = cx.new(|cx| Buffer::local("", cx)); + let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx)); + let range = buffer.read_with(cx, |buffer, cx| { + let snapshot = buffer.snapshot(cx); + snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(0, 0)) + }); + let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap()); + let codegen = cx.new(|cx| { + CodegenAlternative::new( + buffer.clone(), + range, + true, + prompt_builder, + Uuid::new_v4(), + cx, + ) + }); + (buffer, codegen) + } + fn simulate_response_stream( codegen: &Entity, cx: &mut TestAppContext, @@ -1978,6 +2600,20 @@ mod tests { chunks_tx } + fn simulate_tool_based_completion_results( + codegen: &Entity, + cx: &mut TestAppContext, + ) -> mpsc::UnboundedSender> + { + let (events_tx, events_rx) = mpsc::unbounded(); + let model = Arc::new(FakeLanguageModel::default()); + codegen.update(cx, |codegen, cx| { + codegen.generation = + codegen.handle_completion(model, Task::ready(Ok(events_rx.boxed())), cx); + }); + events_tx + } + fn simulate_tool_based_completion( codegen: &Entity, cx: &mut TestAppContext, @@ -1995,6 +2631,22 @@ mod tests { events_tx } + fn failure_tool_use(message: &str, is_input_complete: bool) -> LanguageModelCompletionEvent { + let input = FailureMessageInput { + message: message.into(), + }; + LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse { + id: "failure_tool".into(), + name: FAILURE_MESSAGE_TOOL_NAME.into(), + raw_input: serde_json::to_string(&input).unwrap(), + input: language_model::LanguageModelToolUseInput::Json( + serde_json::to_value(&input).unwrap(), + ), + is_input_complete, + thought_signature: None, + }) + } + fn rewrite_tool_use( id: &str, replacement_text: &str,