agent_ui: Fix panic when multiple rewrite tool uses share a char offset counter (#52458)

## Context

The process_tool_use closure used a single chars_read_so_far counter
shared across all REWRITE_SECTION_TOOL_NAME tool use events. When the
LLM produced multiple separate rewrite tool uses (different IDs), the
counter from the first tool use carried over into the second. If the
second tool use had a shorter replacement_text, the slice indexing
panicked with index out of bounds.

(This caused my first ever Zed panic today.)

Fix by tracking bytes-read per tool use ID using a HashMap keyed on
LanguageModelToolUseId. Also replace the direct slice index with .get()
so any remaining out-of-bounds case is handled gracefully instead of
panicking.

Looks like it hasn't been reported yet. At least I couldn't find any
related issues.

## How to Review

I hope the fix is pretty self-explanatory. We keep track of multiple
rewrite tools, which explains the extra ceremony around getting the
right character count, but apart from that it's pretty straightforward.

The performance impact should be neglible because:

1. Once the completion stream ends (or the task is dropped/cancelled),
the entire Arc<Mutex<HashMap>> is freed.
2. The number of entries is tiny. In a single completion request the LLM
will produce at most a handful of
REWRITE_SECTION_TOOL_NAME tool uses. Each entry is just an Arc<str> key
+ usize value.

So the HashMap grows for the duration of one handle_completion call and
is then dropped wholesale.

## Self-Review Checklist

<!-- Check before requesting review: -->
- [X] I've reviewed my own diff for quality, security, and reliability
- [X] Unsafe blocks (if any) have justifying comments
- [X] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [X] Tests cover the new/changed behavior
- [X] Performance impact has been considered and is acceptable

Release Notes:

- Fixed an out-of-bounds panic when the AI produced multiple inline
rewrites in a single completion.

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::HashSet;
+use collections::{HashMap, HashSet};
 use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};
 use futures::FutureExt;
 use futures::{
@@ -17,7 +17,7 @@ use language_model::{
     CompletionIntent, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
     LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
     LanguageModelRequestTool, LanguageModelTextStream, LanguageModelToolChoice,
-    LanguageModelToolUse, Role, TokenUsage,
+    LanguageModelToolUse, LanguageModelToolUseId, Role, TokenUsage,
 };
 use language_models::provider::anthropic::telemetry::{
     AnthropicCompletionType, AnthropicEventData, AnthropicEventReporter, AnthropicEventType,
@@ -1169,9 +1169,10 @@ impl CodegenAlternative {
                 Failure(String),
             }
 
-            let chars_read_so_far = Arc::new(Mutex::new(0usize));
+            let chars_read_by_tool_id: Arc<Mutex<HashMap<LanguageModelToolUseId, usize>>> =
+                Arc::new(Mutex::new(HashMap::default()));
             let process_tool_use = move |tool_use: LanguageModelToolUse| -> Option<ToolUseOutput> {
-                let mut chars_read_so_far = chars_read_so_far.lock();
+                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) =
@@ -1179,7 +1180,13 @@ impl CodegenAlternative {
                         else {
                             return None;
                         };
-                        let text = input.replacement_text[*chars_read_so_far..].to_string();
+                        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,
@@ -1845,7 +1852,7 @@ mod tests {
             .unbounded_send(rewrite_tool_use("tool_1", &text[..chunk_len], false))
             .unwrap();
         events_tx
-            .unbounded_send(rewrite_tool_use("tool_2", &text, true))
+            .unbounded_send(rewrite_tool_use("tool_1", &text, true))
             .unwrap();
         events_tx
             .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
@@ -1859,6 +1866,52 @@ mod tests {
         );
     }
 
+    // Regression test: a second rewrite tool use with a *shorter* replacement_text
+    // 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,
+            )
+        });
+
+        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();
+        events_tx
+            .unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
+            .unwrap();
+        drop(events_tx);
+        cx.run_until_parked();
+
+        assert_eq!(
+            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
+            "longer replacement textshort"
+        );
+    }
+
     #[gpui::test]
     async fn test_strip_invalid_spans_from_codeblock() {
         assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;