Fix warnings

Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>

.github/workflows/eval.yml

@@ -0,0 +1,71 @@
+name: Run Agent Eval
+
+on:
+  schedule:
+    - cron: "0 * * * *"
+
+  pull_request:
+    branches:
+      - "**"
+    types: [opened, synchronize, reopened, labeled]
+
+  workflow_dispatch:
+
+concurrency:
+  # Allow only one workflow per any non-`main` branch.
+  group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }}
+  cancel-in-progress: true
+
+env:
+  CARGO_TERM_COLOR: always
+  CARGO_INCREMENTAL: 0
+  RUST_BACKTRACE: 1
+  ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+  ZED_CLIENT_CHECKSUM_SEED: ${{ secrets.ZED_CLIENT_CHECKSUM_SEED }}
+  ZED_EVAL_TELEMETRY: 1
+
+jobs:
+  run_eval:
+    timeout-minutes: 60
+    name: Run Agent Eval
+    if: >
+      github.repository_owner == 'zed-industries' &&
+      (github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'run-eval'))
+    runs-on:
+      - buildjet-16vcpu-ubuntu-2204
+    steps:
+      - name: Add Rust to the PATH
+        run: echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+
+      - name: Checkout repo
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4
+        with:
+          clean: false
+
+      - name: Cache dependencies
+        uses: swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6 # v2
+        with:
+          save-if: ${{ github.ref == 'refs/heads/main' }}
+          cache-provider: "buildjet"
+
+      - name: Install Linux dependencies
+        run: ./script/linux
+
+      - name: Configure CI
+        run: |
+          mkdir -p ./../.cargo
+          cp ./.cargo/ci-config.toml ./../.cargo/config.toml
+
+      - name: Compile eval
+        run: cargo build --package=eval
+
+      - name: Run eval
+        run: cargo run --package=eval -- --repetitions=3 --concurrency=1
+
+      # Even the Linux runner is not stateful, in theory there is no need to do this cleanup.
+      # But, to avoid potential issues in the future if we choose to use a stateful Linux runner and forget to add code
+      # to clean up the config file, I’ve included the cleanup code here as a precaution.
+      # While it’s not strictly necessary at this moment, I believe it’s better to err on the side of caution.
+      - name: Clean CI config file
+        if: always()
+        run: rm -rf ./../.cargo

.github/workflows/run_agent_eval_daily.yml

@@ -1,28 +0,0 @@
-name: Run Eval Daily
-
-on:
-  schedule:
-    - cron: "0 2 * * *"
-  workflow_dispatch:
-
-env:
-  CARGO_TERM_COLOR: always
-  CARGO_INCREMENTAL: 0
-  RUST_BACKTRACE: 1
-
-jobs:
-  run_eval:
-    name: Run Eval
-    if: github.repository_owner == 'zed-industries'
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout repo
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4
-        with:
-          clean: false
-
-      - name: Setup Rust
-        uses: dtolnay/rust-toolchain@stable
-
-      - name: Run cargo eval
-        run: cargo run -p eval

.rules

@@ -5,6 +5,7 @@
 * Prefer implementing functionality in existing files unless it is a new logical component. Avoid creating many small files.
 * Avoid using functions that panic like `unwrap()`, instead use mechanisms like `?` to propagate errors.
 * Be careful with operations like indexing which may panic if the indexes are out of bounds.
+* Never create files with `mod.rs` paths - prefer `src/some_module.rs` instead of `src/some_module/mod.rs`.
 
 # GPUI
 
@@ -108,3 +109,13 @@ When a view's state has changed in a way that may affect its rendering, it shoul
 While updating an entity (`cx: Context<T>`), it can emit an event using `cx.emit(event)`. Entities register which events they can emit by declaring `impl EventEmittor<EventType> for EntityType {}`.
 
 Other entities can then register a callback to handle these events by doing `cx.subscribe(other_entity, |this, other_entity, event, cx| ...)`. This will return a `Subscription` which deregisters the callback when dropped.  Typically `cx.subscribe` happens when creating a new entity and the subscriptions are stored in a `_subscriptions: Vec<Subscription>` field.
+
+## Recent API changes
+
+GPUI has had some changes to its APIs. Always write code using the new APIs:
+
+* `spawn` methods now take async closures (`AsyncFn`), and so should be called like `cx.spawn(async move |cx| ...)`.
+* Use `Entity<T>`. This replaces `Model<T>` and `View<T>` which longer exists and should NEVER be used.
+* Use `App` references. This replaces `AppContext` which no longer exists and should NEVER be used.
+* Use `Context<T>` references. This replaces `ModelContext<T>` which no longer exists and should NEVER be used.
+* `Window` is now passed around explicitly. The new interface adds a `Window` reference parameter to some methods, and adds some new "*_in" methods for plumbing `Window`. The old types `WindowContext` and `ViewContext<T>` should NEVER be used.

Cargo.toml

@@ -480,6 +480,7 @@ num-format = "0.4.4"
 ordered-float = "2.1.1"
 palette = { version = "0.7.5", default-features = false, features = ["std"] }
 parking_lot = "0.12.1"
+partial-json-fixer = "0.5.3"
 pathdiff = "0.2"
 pet = { git = "https://github.com/microsoft/python-environment-tools.git", rev = "845945b830297a50de0e24020b980a65e4820559" }
 pet-fs = { git = "https://github.com/microsoft/python-environment-tools.git", rev = "845945b830297a50de0e24020b980a65e4820559" }
@@ -604,7 +605,7 @@ wasmtime-wasi = "29"
 which = "6.0.0"
 wit-component = "0.221"
 workspace-hack = "0.1.0"
-zed_llm_client = "0.6.1"
+zed_llm_client = "0.7.0"
 zstd = "0.11"
 metal = "0.29"
 

assets/icons/image.svg

@@ -0,0 +1 @@
+<svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-image-icon lucide-image"><rect width="18" height="18" x="3" y="3" rx="2" ry="2"/><circle cx="9" cy="9" r="2"/><path d="m21 15-3.086-3.086a2 2 0 0 0-2.828 0L6 21"/></svg>

assets/keymaps/default-macos.json

@@ -1028,10 +1028,10 @@
       // Using `ctrl-shift-space` in Zed requires disabling the macOS global shortcut.
       // System Preferences->Keyboard->Keyboard Shortcuts->Input Sources->Select the previous input source (uncheck)
       "ctrl-shift-space": "terminal::ToggleViMode",
-      "ctrl-k up": "pane::SplitUp",
-      "ctrl-k down": "pane::SplitDown",
-      "ctrl-k left": "pane::SplitLeft",
-      "ctrl-k right": "pane::SplitRight"
+      "ctrl-alt-up": "pane::SplitUp",
+      "ctrl-alt-down": "pane::SplitDown",
+      "ctrl-alt-left": "pane::SplitLeft",
+      "ctrl-alt-right": "pane::SplitRight"
     }
   },
   {

assets/prompts/assistant_system_prompt.hbs

@@ -73,9 +73,9 @@ There are project rules that apply to these root directories:
 {{/each}}
 {{/if}}
 
-{{#if has_default_user_rules}}
+{{#if has_user_rules}}
 The user has specified the following rules that should be applied:
-{{#each default_user_rules}}
+{{#each user_rules}}
 
 {{#if title}}
 Rules title: {{title}}

assets/settings/default.json

@@ -1489,7 +1489,12 @@
     "use_multiline_find": false,
     "use_smartcase_find": false,
     "highlight_on_yank_duration": 200,
-    "custom_digraphs": {}
+    "custom_digraphs": {},
+    // Cursor shape for the each mode.
+    // Specify the mode as the key and the shape as the value.
+    // The mode can be one of the following: "normal", "replace", "insert", "visual".
+    // The shape can be one of the following: "block", "bar", "underline", "hollow".
+    "cursor_shape": {}
   },
   // The server to connect to. If the environment variable
   // ZED_SERVER_URL is set, it will override this setting.

crates/agent/src/active_thread.rs

@@ -1,4 +1,4 @@
-use crate::context::{AssistantContext, ContextId, format_context_as_string};
+use crate::context::{AssistantContext, ContextId, RULES_ICON, format_context_as_string};
 use crate::context_picker::MentionLink;
 use crate::thread::{
     LastRestoreCheckpoint, MessageId, MessageSegment, Thread, ThreadError, ThreadEvent,
@@ -6,7 +6,9 @@ use crate::thread::{
 };
 use crate::thread_store::{RulesLoadingError, ThreadStore};
 use crate::tool_use::{PendingToolUseStatus, ToolUse};
-use crate::ui::{AddedContext, AgentNotification, AgentNotificationEvent, ContextPill};
+use crate::ui::{
+    AddedContext, AgentNotification, AgentNotificationEvent, AnimatedLabel, ContextPill,
+};
 use crate::{AssistantPanel, OpenActiveThreadAsMarkdown};
 use anyhow::Context as _;
 use assistant_settings::{AssistantSettings, NotifyWhenAgentWaiting};
@@ -266,14 +268,6 @@ fn default_markdown_style(window: &Window, cx: &App) -> MarkdownStyle {
     }
 }
 
-fn render_tool_use_markdown(
-    text: SharedString,
-    language_registry: Arc<LanguageRegistry>,
-    cx: &mut App,
-) -> Entity<Markdown> {
-    cx.new(|cx| Markdown::new(text, Some(language_registry), None, cx))
-}
-
 fn tool_use_markdown_style(window: &Window, cx: &mut App) -> MarkdownStyle {
     let theme_settings = ThemeSettings::get_global(cx);
     let colors = cx.theme().colors();
@@ -439,47 +433,39 @@ fn render_markdown_code_block(
                         workspace
                             .update(cx, {
                                 |workspace, cx| {
-                                    if let Some(project_path) = workspace
+                                    let Some(project_path) = workspace
                                         .project()
                                         .read(cx)
                                         .find_project_path(&path_range.path, cx)
-                                    {
-                                        let target = path_range.range.as_ref().map(|range| {
-                                            Point::new(
-                                                // Line number is 1-based
-                                                range.start.line.saturating_sub(1),
-                                                range.start.col.unwrap_or(0),
-                                            )
-                                        });
-                                        let open_task = workspace.open_path(
-                                            project_path,
-                                            None,
-                                            true,
-                                            window,
-                                            cx,
-                                        );
-                                        window
-                                            .spawn(cx, async move |cx| {
-                                                let item = open_task.await?;
-                                                if let Some(target) = target {
-                                                    if let Some(active_editor) =
-                                                        item.downcast::<Editor>()
-                                                    {
-                                                        active_editor
-                                                            .downgrade()
-                                                            .update_in(cx, |editor, window, cx| {
-                                                                editor
-                                                                    .go_to_singleton_buffer_point(
-                                                                        target, window, cx,
-                                                                    );
-                                                            })
-                                                            .log_err();
-                                                    }
-                                                }
-                                                anyhow::Ok(())
-                                            })
-                                            .detach_and_log_err(cx);
-                                    }
+                                    else {
+                                        return;
+                                    };
+                                    let Some(target) = path_range.range.as_ref().map(|range| {
+                                        Point::new(
+                                            // Line number is 1-based
+                                            range.start.line.saturating_sub(1),
+                                            range.start.col.unwrap_or(0),
+                                        )
+                                    }) else {
+                                        return;
+                                    };
+                                    let open_task =
+                                        workspace.open_path(project_path, None, true, window, cx);
+                                    window
+                                        .spawn(cx, async move |cx| {
+                                            let item = open_task.await?;
+                                            if let Some(active_editor) = item.downcast::<Editor>() {
+                                                active_editor
+                                                    .update_in(cx, |editor, window, cx| {
+                                                        editor.go_to_singleton_buffer_point(
+                                                            target, window, cx,
+                                                        );
+                                                    })
+                                                    .ok();
+                                            }
+                                            anyhow::Ok(())
+                                        })
+                                        .detach_and_log_err(cx);
                                 }
                             })
                             .ok();
@@ -678,6 +664,26 @@ fn open_markdown_link(
                 })
                 .detach_and_log_err(cx);
         }
+        Some(MentionLink::Selection(path, line_range)) => {
+            let open_task = workspace.update(cx, |workspace, cx| {
+                workspace.open_path(path, None, true, window, cx)
+            });
+            window
+                .spawn(cx, async move |cx| {
+                    let active_editor = open_task
+                        .await?
+                        .downcast::<Editor>()
+                        .context("Item is not an editor")?;
+                    active_editor.update_in(cx, |editor, window, cx| {
+                        editor.change_selections(Some(Autoscroll::center()), window, cx, |s| {
+                            s.select_ranges([Point::new(line_range.start as u32, 0)
+                                ..Point::new(line_range.start as u32, 0)])
+                        });
+                        anyhow::Ok(())
+                    })
+                })
+                .detach_and_log_err(cx);
+        }
         Some(MentionLink::Thread(thread_id)) => workspace.update(cx, |workspace, cx| {
             if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {
                 panel.update(cx, |panel, cx| {
@@ -688,6 +694,12 @@ fn open_markdown_link(
             }
         }),
         Some(MentionLink::Fetch(url)) => cx.open_url(&url),
+        Some(MentionLink::Rules(prompt_id)) => window.dispatch_action(
+            Box::new(OpenPromptLibrary {
+                prompt_to_select: Some(prompt_id.0),
+            }),
+            cx,
+        ),
         None => cx.open_url(&text),
     }
 }
@@ -787,10 +799,11 @@ impl ActiveThread {
         self.thread.read(cx).summary_or_default()
     }
 
-    pub fn cancel_last_completion(&mut self, cx: &mut App) -> bool {
+    pub fn cancel_last_completion(&mut self, window: &mut Window, cx: &mut App) -> bool {
         self.last_error.take();
-        self.thread
-            .update(cx, |thread, cx| thread.cancel_last_completion(cx))
+        self.thread.update(cx, |thread, cx| {
+            thread.cancel_last_completion(Some(window.window_handle()), cx)
+        })
     }
 
     pub fn last_error(&self) -> Option<ThreadError> {
@@ -861,21 +874,34 @@ impl ActiveThread {
         tool_output: SharedString,
         cx: &mut Context<Self>,
     ) {
-        let rendered = RenderedToolUse {
-            label: render_tool_use_markdown(tool_label.into(), self.language_registry.clone(), cx),
-            input: render_tool_use_markdown(
-                format!(
-                    "```json\n{}\n```",
-                    serde_json::to_string_pretty(tool_input).unwrap_or_default()
-                )
-                .into(),
-                self.language_registry.clone(),
-                cx,
-            ),
-            output: render_tool_use_markdown(tool_output, self.language_registry.clone(), cx),
-        };
-        self.rendered_tool_uses
-            .insert(tool_use_id.clone(), rendered);
+        let rendered = self
+            .rendered_tool_uses
+            .entry(tool_use_id.clone())
+            .or_insert_with(|| RenderedToolUse {
+                label: cx.new(|cx| {
+                    Markdown::new("".into(), Some(self.language_registry.clone()), None, cx)
+                }),
+                input: cx.new(|cx| {
+                    Markdown::new("".into(), Some(self.language_registry.clone()), None, cx)
+                }),
+                output: cx.new(|cx| {
+                    Markdown::new("".into(), Some(self.language_registry.clone()), None, cx)
+                }),
+            });
+
+        rendered.label.update(cx, |this, cx| {
+            this.replace(tool_label, cx);
+        });
+        rendered.input.update(cx, |this, cx| {
+            let input = format!(
+                "```json\n{}\n```",
+                serde_json::to_string_pretty(tool_input).unwrap_or_default()
+            );
+            this.replace(input, cx);
+        });
+        rendered.output.update(cx, |this, cx| {
+            this.replace(tool_output, cx);
+        });
     }
 
     fn handle_thread_event(
@@ -968,6 +994,19 @@ impl ActiveThread {
                     );
                 }
             }
+            ThreadEvent::StreamedToolUse {
+                tool_use_id,
+                ui_text,
+                input,
+            } => {
+                self.render_tool_use_markdown(
+                    tool_use_id.clone(),
+                    ui_text.clone(),
+                    input,
+                    "".into(),
+                    cx,
+                );
+            }
             ThreadEvent::ToolFinished {
                 pending_tool_use, ..
             } => {
@@ -986,6 +1025,7 @@ impl ActiveThread {
                 }
             }
             ThreadEvent::CheckpointChanged => cx.notify(),
+            ThreadEvent::ReceivedTextChunk => {}
         }
     }
 
@@ -1048,9 +1088,21 @@ impl ActiveThread {
     ) {
         let options = AgentNotification::window_options(screen, cx);
 
+        let project_name = self.workspace.upgrade().and_then(|workspace| {
+            workspace
+                .read(cx)
+                .project()
+                .read(cx)
+                .visible_worktrees(cx)
+                .next()
+                .map(|worktree| worktree.read(cx).root_name().to_string())
+        });
+
         if let Some(screen_window) = cx
             .open_window(options, |_, cx| {
-                cx.new(|_| AgentNotification::new(title.clone(), caption.clone(), icon))
+                cx.new(|_| {
+                    AgentNotification::new(title.clone(), caption.clone(), icon, project_name)
+                })
             })
             .log_err()
         {
@@ -1255,7 +1307,7 @@ impl ActiveThread {
     fn confirm_editing_message(
         &mut self,
         _: &menu::Confirm,
-        _: &mut Window,
+        window: &mut Window,
         cx: &mut Context<Self>,
     ) {
         let Some((message_id, state)) = self.editing_message.take() else {
@@ -1285,7 +1337,7 @@ impl ActiveThread {
 
         self.thread.update(cx, |thread, cx| {
             thread.advance_prompt_id();
-            thread.send_to_model(model.model, cx)
+            thread.send_to_model(model.model, Some(window.window_handle()), cx);
         });
         cx.notify();
     }
@@ -1447,45 +1499,8 @@ impl ActiveThread {
 
         let needs_confirmation = tool_uses.iter().any(|tool_use| tool_use.needs_confirmation);
 
-        let generating_label = (is_generating && is_last_message).then(|| {
-            Label::new("Generating")
-                .color(Color::Muted)
-                .size(LabelSize::Small)
-                .with_animations(
-                    "generating-label",
-                    vec![
-                        Animation::new(Duration::from_secs(1)),
-                        Animation::new(Duration::from_secs(1)).repeat(),
-                    ],
-                    |mut label, animation_ix, delta| {
-                        match animation_ix {
-                            0 => {
-                                let chars_to_show = (delta * 10.).ceil() as usize;
-                                let text = &"Generating"[0..chars_to_show];
-                                label.set_text(text);
-                            }
-                            1 => {
-                                let text = match delta {
-                                    d if d < 0.25 => "Generating",
-                                    d if d < 0.5 => "Generating.",
-                                    d if d < 0.75 => "Generating..",
-                                    _ => "Generating...",
-                                };
-                                label.set_text(text);
-                            }
-                            _ => {}
-                        }
-                        label
-                    },
-                )
-                .with_animation(
-                    "pulsating-label",
-                    Animation::new(Duration::from_secs(2))
-                        .repeat()
-                        .with_easing(pulsating_between(0.6, 1.)),
-                    |label, delta| label.map_element(|label| label.alpha(delta)),
-                )
-        });
+        let generating_label = (is_generating && is_last_message)
+            .then(|| AnimatedLabel::new("Generating").size(LabelSize::Small));
 
         // Don't render user messages that are just there for returning tool results.
         if message.role == Role::User && thread.message_has_tool_results(message_id) {
@@ -1512,9 +1527,7 @@ impl ActiveThread {
             .map(|(_, state)| state.editor.clone());
 
         let colors = cx.theme().colors();
-        let active_color = colors.element_active;
         let editor_bg_color = colors.editor_background;
-        let bg_user_message_header = editor_bg_color.blend(active_color.opacity(0.25));
 
         let open_as_markdown = IconButton::new(("open-as-markdown", ix), IconName::FileCode)
             .shape(ui::IconButtonShape::Square)
@@ -1641,7 +1654,7 @@ impl ActiveThread {
         let message_content =
             has_content.then(|| {
                 v_flex()
-                    .gap_1p5()
+                    .gap_1()
                     .when(!message_is_empty, |parent| {
                         parent.child(
                             if let Some(edit_message_editor) = edit_message_editor.clone() {
@@ -1664,7 +1677,6 @@ impl ActiveThread {
                                     .on_action(cx.listener(Self::cancel_editing_message))
                                     .on_action(cx.listener(Self::confirm_editing_message))
                                     .min_h_6()
-                                    .pt_1()
                                     .child(EditorElement::new(
                                         &edit_message_editor,
                                         EditorStyle {
@@ -1679,7 +1691,6 @@ impl ActiveThread {
                             } else {
                                 div()
                                     .min_h_6()
-                                    .text_ui(cx)
                                     .child(self.render_message_content(
                                         message_id,
                                         rendered_message,
@@ -1738,35 +1749,18 @@ impl ActiveThread {
                 .pb_4()
                 .child(
                     v_flex()
-                        .bg(colors.editor_background)
+                        .bg(editor_bg_color)
                         .rounded_lg()
                         .border_1()
                         .border_color(colors.border)
                         .shadow_md()
+                        .child(div().p_2().children(message_content))
                         .child(
                             h_flex()
-                                .py_1()
-                                .pl_2()
-                                .pr_1()
-                                .bg(bg_user_message_header)
-                                .border_b_1()
-                                .border_color(colors.border)
-                                .justify_between()
-                                .rounded_t_md()
-                                .child(
-                                    h_flex()
-                                        .gap_1p5()
-                                        .child(
-                                            Icon::new(IconName::PersonCircle)
-                                                .size(IconSize::XSmall)
-                                                .color(Color::Muted),
-                                        )
-                                        .child(
-                                            Label::new("You")
-                                                .size(LabelSize::Small)
-                                                .color(Color::Muted),
-                                        ),
-                                )
+                                .p_1()
+                                .border_t_1()
+                                .border_color(colors.border_variant)
+                                .justify_end()
                                 .child(
                                     h_flex()
                                         .gap_1()
@@ -1819,8 +1813,12 @@ impl ActiveThread {
                                             edit_message_editor.is_none() && allow_editing_message,
                                             |this| {
                                                 this.child(
-                                                    Button::new("edit-message", "Edit")
+                                                    Button::new("edit-message", "Edit Message")
                                                         .label_size(LabelSize::Small)
+                                                        .icon(IconName::Pencil)
+                                                        .icon_size(IconSize::XSmall)
+                                                        .icon_color(Color::Muted)
+                                                        .icon_position(IconPosition::Start)
                                                         .on_click(cx.listener({
                                                             let message_segments =
                                                                 message.segments.clone();
@@ -1837,8 +1835,7 @@ impl ActiveThread {
                                             },
                                         ),
                                 ),
-                        )
-                        .child(div().p_2().children(message_content)),
+                        ),
                 ),
             Role::Assistant => v_flex()
                 .id(("message-container", ix))
@@ -1846,11 +1843,9 @@ impl ActiveThread {
                 .gap_2()
                 .children(message_content)
                 .when(has_tool_uses, |parent| {
-                    parent.children(
-                        tool_uses
-                            .into_iter()
-                            .map(|tool_use| self.render_tool_use(tool_use, window, cx)),
-                    )
+                    parent.children(tool_uses.into_iter().map(|tool_use| {
+                        self.render_tool_use(tool_use, window, workspace.clone(), cx)
+                    }))
                 }),
             Role::System => div().id(("message-container", ix)).py_1().px_2().child(
                 v_flex()
@@ -2077,11 +2072,13 @@ impl ActiveThread {
             .map(|m| m.role)
             .unwrap_or(Role::User);
 
-        let is_assistant = message_role == Role::Assistant;
+        let is_assistant_message = message_role == Role::Assistant;
+        let is_user_message = message_role == Role::User;
 
         v_flex()
             .text_ui(cx)
             .gap_2()
+            .when(is_user_message, |this| this.text_xs())
             .children(
                 rendered_message.segments.iter().enumerate().map(
                     |(index, segment)| match segment {
@@ -2102,10 +2099,28 @@ impl ActiveThread {
                         RenderedMessageSegment::Text(markdown) => {
                             let markdown_element = MarkdownElement::new(
                                 markdown.clone(),
-                                default_markdown_style(window, cx),
+                                if is_user_message {
+                                    let mut style = default_markdown_style(window, cx);
+                                    let mut text_style = window.text_style();
+                                    let theme_settings = ThemeSettings::get_global(cx);
+
+                                    let buffer_font = theme_settings.buffer_font.family.clone();
+                                    let buffer_font_size = TextSize::Small.rems(cx);
+
+                                    text_style.refine(&TextStyleRefinement {
+                                        font_family: Some(buffer_font),
+                                        font_size: Some(buffer_font_size.into()),
+                                        ..Default::default()
+                                    });
+
+                                    style.base_text_style = text_style;
+                                    style
+                                } else {
+                                    default_markdown_style(window, cx)
+                                },
                             );
 
-                            let markdown_element = if is_assistant {
+                            let markdown_element = if is_assistant_message {
                                 markdown_element.code_block_renderer(
                                     markdown::CodeBlockRenderer::Custom {
                                         render: Arc::new({
@@ -2244,34 +2259,7 @@ impl ActiveThread {
                                             .size(IconSize::XSmall)
                                             .color(Color::Muted),
                                     )
-                                    .child({
-                                        Label::new("Thinking")
-                                            .color(Color::Muted)
-                                            .size(LabelSize::Small)
-                                            .with_animation(
-                                                "generating-label",
-                                                Animation::new(Duration::from_secs(1)).repeat(),
-                                                |mut label, delta| {
-                                                    let text = match delta {
-                                                        d if d < 0.25 => "Thinking",
-                                                        d if d < 0.5 => "Thinking.",
-                                                        d if d < 0.75 => "Thinking..",
-                                                        _ => "Thinking...",
-                                                    };
-                                                    label.set_text(text);
-                                                    label
-                                                },
-                                            )
-                                            .with_animation(
-                                                "pulsating-label",
-                                                Animation::new(Duration::from_secs(2))
-                                                    .repeat()
-                                                    .with_easing(pulsating_between(0.6, 1.)),
-                                                |label, delta| {
-                                                    label.map_element(|label| label.alpha(delta))
-                                                },
-                                            )
-                                    }),
+                                    .child(AnimatedLabel::new("Thinking").size(LabelSize::Small)),
                             )
                             .child(
                                 h_flex()
@@ -2450,10 +2438,11 @@ impl ActiveThread {
         &self,
         tool_use: ToolUse,
         window: &mut Window,
+        workspace: WeakEntity<Workspace>,
         cx: &mut Context<Self>,
     ) -> impl IntoElement + use<> {
         if let Some(card) = self.thread.read(cx).card_for_tool(&tool_use.id) {
-            return card.render(&tool_use.status, window, cx);
+            return card.render(&tool_use.status, window, workspace, cx);
         }
 
         let is_open = self
@@ -2469,16 +2458,18 @@ impl ActiveThread {
             .upgrade()
             .map(|workspace| workspace.read(cx).app_state().fs.clone());
         let needs_confirmation = matches!(&tool_use.status, ToolUseStatus::NeedsConfirmation);
-        let edit_tools = tool_use.needs_confirmation;
+        let needs_confirmation_tools = tool_use.needs_confirmation;
 
         let status_icons = div().child(match &tool_use.status {
-            ToolUseStatus::Pending | ToolUseStatus::NeedsConfirmation => {
+            ToolUseStatus::NeedsConfirmation => {
                 let icon = Icon::new(IconName::Warning)
                     .color(Color::Warning)
                     .size(IconSize::Small);
                 icon.into_any_element()
             }
-            ToolUseStatus::Running => {
+            ToolUseStatus::Pending
+            | ToolUseStatus::InputStillStreaming
+            | ToolUseStatus::Running => {
                 let icon = Icon::new(IconName::ArrowCircle)
                     .color(Color::Accent)
                     .size(IconSize::Small);
@@ -2564,34 +2555,34 @@ impl ActiveThread {
                             }),
                         )),
                 ),
-                ToolUseStatus::Running => container.child(
-                    results_content_container().child(
-                        h_flex()
-                            .gap_1()
-                            .pb_1()
-                            .border_t_1()
-                            .border_color(self.tool_card_border_color(cx))
-                            .child(
-                                Icon::new(IconName::ArrowCircle)
-                                    .size(IconSize::Small)
-                                    .color(Color::Accent)
-                                    .with_animation(
-                                        "arrow-circle",
-                                        Animation::new(Duration::from_secs(2)).repeat(),
-                                        |icon, delta| {
-                                            icon.transform(Transformation::rotate(percentage(
-                                                delta,
-                                            )))
-                                        },
-                                    ),
-                            )
-                            .child(
-                                Label::new("Running…")
-                                    .size(LabelSize::XSmall)
-                                    .color(Color::Muted)
-                                    .buffer_font(cx),
-                            ),
-                    ),
+                ToolUseStatus::InputStillStreaming | ToolUseStatus::Running => container.child(
+                    results_content_container()
+                        .border_t_1()
+                        .border_color(self.tool_card_border_color(cx))
+                        .child(
+                            h_flex()
+                                .gap_1()
+                                .child(
+                                    Icon::new(IconName::ArrowCircle)
+                                        .size(IconSize::Small)
+                                        .color(Color::Accent)
+                                        .with_animation(
+                                            "arrow-circle",
+                                            Animation::new(Duration::from_secs(2)).repeat(),
+                                            |icon, delta| {
+                                                icon.transform(Transformation::rotate(percentage(
+                                                    delta,
+                                                )))
+                                            },
+                                        ),
+                                )
+                                .child(
+                                    Label::new("Running…")
+                                        .size(LabelSize::XSmall)
+                                        .color(Color::Muted)
+                                        .buffer_font(cx),
+                                ),
+                        ),
                 ),
                 ToolUseStatus::Error(_) => container.child(
                     results_content_container()
@@ -2655,8 +2646,8 @@ impl ActiveThread {
                 ))
         };
 
-        v_flex().gap_1().mb_3().map(|element| {
-            if !edit_tools {
+        v_flex().gap_1().mb_2().map(|element| {
+            if !needs_confirmation_tools {
                 element.child(
                     v_flex()
                         .child(
@@ -2834,30 +2825,7 @@ impl ActiveThread {
                                 .border_color(self.tool_card_border_color(cx))
                                 .rounded_b_lg()
                                 .child(
-                                    Label::new("Waiting for Confirmation…")
-                                        .color(Color::Muted)
-                                        .size(LabelSize::Small)
-                                        .with_animation(
-                                            "generating-label",
-                                            Animation::new(Duration::from_secs(1)).repeat(),
-                                            |mut label, delta| {
-                                                let text = match delta {
-                                                    d if d < 0.25 => "Waiting for Confirmation",
-                                                    d if d < 0.5 => "Waiting for Confirmation.",
-                                                    d if d < 0.75 => "Waiting for Confirmation..",
-                                                    _ => "Waiting for Confirmation...",
-                                                };
-                                                label.set_text(text);
-                                                label
-                                            },
-                                        )
-                                        .with_animation(
-                                            "pulsating-label",
-                                            Animation::new(Duration::from_secs(2))
-                                                .repeat()
-                                                .with_easing(pulsating_between(0.6, 1.)),
-                                            |label, delta| label.map_element(|label| label.alpha(delta)),
-                                        ),
+                                    AnimatedLabel::new("Waiting for Confirmation").size(LabelSize::Small)
                                 )
                                 .child(
                                     h_flex()
@@ -2957,10 +2925,10 @@ impl ActiveThread {
             return div().into_any();
         };
 
-        let default_user_rules_text = if project_context.default_user_rules.is_empty() {
+        let user_rules_text = if project_context.user_rules.is_empty() {
             None
-        } else if project_context.default_user_rules.len() == 1 {
-            let user_rules = &project_context.default_user_rules[0];
+        } else if project_context.user_rules.len() == 1 {
+            let user_rules = &project_context.user_rules[0];
 
             match user_rules.title.as_ref() {
                 Some(title) => Some(format!("Using \"{title}\" user rule")),
@@ -2969,14 +2937,14 @@ impl ActiveThread {
         } else {
             Some(format!(
                 "Using {} user rules",
-                project_context.default_user_rules.len()
+                project_context.user_rules.len()
             ))
         };
 
-        let first_default_user_rules_id = project_context
-            .default_user_rules
+        let first_user_rules_id = project_context
+            .user_rules
             .first()
-            .map(|user_rules| user_rules.uuid);
+            .map(|user_rules| user_rules.uuid.0);
 
         let rules_files = project_context
             .worktrees
@@ -2993,7 +2961,7 @@ impl ActiveThread {
             rules_files => Some(format!("Using {} project rules files", rules_files.len())),
         };
 
-        if default_user_rules_text.is_none() && rules_file_text.is_none() {
+        if user_rules_text.is_none() && rules_file_text.is_none() {
             return div().into_any();
         }
 
@@ -3001,45 +2969,42 @@ impl ActiveThread {
             .pt_2()
             .px_2p5()
             .gap_1()
-            .when_some(
-                default_user_rules_text,
-                |parent, default_user_rules_text| {
-                    parent.child(
-                        h_flex()
-                            .w_full()
-                            .child(
-                                Icon::new(IconName::File)
-                                    .size(IconSize::XSmall)
-                                    .color(Color::Disabled),
-                            )
-                            .child(
-                                Label::new(default_user_rules_text)
-                                    .size(LabelSize::XSmall)
-                                    .color(Color::Muted)
-                                    .truncate()
-                                    .buffer_font(cx)
-                                    .ml_1p5()
-                                    .mr_0p5(),
-                            )
-                            .child(
-                                IconButton::new("open-prompt-library", IconName::ArrowUpRightAlt)
-                                    .shape(ui::IconButtonShape::Square)
-                                    .icon_size(IconSize::XSmall)
-                                    .icon_color(Color::Ignored)
-                                    // TODO: Figure out a way to pass focus handle here so we can display the `OpenPromptLibrary`  keybinding
-                                    .tooltip(Tooltip::text("View User Rules"))
-                                    .on_click(move |_event, window, cx| {
-                                        window.dispatch_action(
-                                            Box::new(OpenPromptLibrary {
-                                                prompt_to_focus: first_default_user_rules_id,
-                                            }),
-                                            cx,
-                                        )
-                                    }),
-                            ),
-                    )
-                },
-            )
+            .when_some(user_rules_text, |parent, user_rules_text| {
+                parent.child(
+                    h_flex()
+                        .w_full()
+                        .child(
+                            Icon::new(RULES_ICON)
+                                .size(IconSize::XSmall)
+                                .color(Color::Disabled),
+                        )
+                        .child(
+                            Label::new(user_rules_text)
+                                .size(LabelSize::XSmall)
+                                .color(Color::Muted)
+                                .truncate()
+                                .buffer_font(cx)
+                                .ml_1p5()
+                                .mr_0p5(),
+                        )
+                        .child(
+                            IconButton::new("open-prompt-library", IconName::ArrowUpRightAlt)
+                                .shape(ui::IconButtonShape::Square)
+                                .icon_size(IconSize::XSmall)
+                                .icon_color(Color::Ignored)
+                                // TODO: Figure out a way to pass focus handle here so we can display the `OpenPromptLibrary`  keybinding
+                                .tooltip(Tooltip::text("View User Rules"))
+                                .on_click(move |_event, window, cx| {
+                                    window.dispatch_action(
+                                        Box::new(OpenPromptLibrary {
+                                            prompt_to_select: first_user_rules_id,
+                                        }),
+                                        cx,
+                                    )
+                                }),
+                        ),
+                )
+            })
             .when_some(rules_file_text, |parent, rules_file_text| {
                 parent.child(
                     h_flex()
@@ -3074,7 +3039,7 @@ impl ActiveThread {
         &mut self,
         tool_use_id: LanguageModelToolUseId,
         _: &ClickEvent,
-        _window: &mut Window,
+        window: &mut Window,
         cx: &mut Context<Self>,
     ) {
         if let Some(PendingToolUseStatus::NeedsConfirmation(c)) = self
@@ -3090,6 +3055,7 @@ impl ActiveThread {
                     c.input.clone(),
                     &c.messages,
                     c.tool.clone(),
+                    Some(window.window_handle()),
                     cx,
                 );
             });
@@ -3101,11 +3067,12 @@ impl ActiveThread {
         tool_use_id: LanguageModelToolUseId,
         tool_name: Arc<str>,
         _: &ClickEvent,
-        _window: &mut Window,
+        window: &mut Window,
         cx: &mut Context<Self>,
     ) {
+        let window_handle = window.window_handle();
         self.thread.update(cx, |thread, cx| {
-            thread.deny_tool_use(tool_use_id, tool_name, cx);
+            thread.deny_tool_use(tool_use_id, tool_name, Some(window_handle), cx);
         });
     }
 
@@ -3259,12 +3226,10 @@ pub(crate) fn open_context(
             }
         }
         AssistantContext::Directory(directory_context) => {
-            let project_path = directory_context.project_path(cx);
+            let entry_id = directory_context.entry_id;
             workspace.update(cx, |workspace, cx| {
-                workspace.project().update(cx, |project, cx| {
-                    if let Some(entry) = project.entry_for_path(&project_path, cx) {
-                        cx.emit(project::Event::RevealInProjectPanel(entry.id));
-                    }
+                workspace.project().update(cx, |_project, cx| {
+                    cx.emit(project::Event::RevealInProjectPanel(entry_id));
                 })
             })
         }
@@ -3287,15 +3252,15 @@ pub(crate) fn open_context(
                     .detach();
             }
         }
-        AssistantContext::Excerpt(excerpt_context) => {
-            if let Some(project_path) = excerpt_context
+        AssistantContext::Selection(selection_context) => {
+            if let Some(project_path) = selection_context
                 .context_buffer
                 .buffer
                 .read(cx)
                 .project_path(cx)
             {
-                let snapshot = excerpt_context.context_buffer.buffer.read(cx).snapshot();
-                let target_position = excerpt_context.range.start.to_point(&snapshot);
+                let snapshot = selection_context.context_buffer.buffer.read(cx).snapshot();
+                let target_position = selection_context.range.start.to_point(&snapshot);
 
                 open_editor_at_position(project_path, target_position, &workspace, window, cx)
                     .detach();
@@ -3316,6 +3281,13 @@ pub(crate) fn open_context(
                 }
             })
         }
+        AssistantContext::Rules(rules_context) => window.dispatch_action(
+            Box::new(OpenPromptLibrary {
+                prompt_to_select: Some(rules_context.prompt_id.0),
+            }),
+            cx,
+        ),
+        AssistantContext::Image(_) => {}
     }
 }
 

crates/agent/src/agent_diff.rs

@@ -1,4 +1,4 @@
-use crate::{Keep, KeepAll, Reject, RejectAll, Thread, ThreadEvent};
+use crate::{Keep, KeepAll, Reject, RejectAll, Thread, ThreadEvent, ui::AnimatedLabel};
 use anyhow::Result;
 use buffer_diff::DiffHunkStatus;
 use collections::{HashMap, HashSet};
@@ -8,8 +8,8 @@ use editor::{
     scroll::Autoscroll,
 };
 use gpui::{
-    Action, AnyElement, AnyView, App, Entity, EventEmitter, FocusHandle, Focusable, SharedString,
-    Subscription, Task, WeakEntity, Window, prelude::*,
+    Action, AnyElement, AnyView, App, Empty, Entity, EventEmitter, FocusHandle, Focusable,
+    SharedString, Subscription, Task, WeakEntity, Window, prelude::*,
 };
 use language::{Capability, DiskState, OffsetRangeExt, Point};
 use multi_buffer::PathKey;
@@ -307,6 +307,10 @@ impl AgentDiff {
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
+        if self.thread.read(cx).is_generating() {
+            return;
+        }
+
         let snapshot = self.multibuffer.read(cx).snapshot(cx);
         let diff_hunks_in_ranges = self
             .editor
@@ -339,6 +343,10 @@ impl AgentDiff {
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
+        if self.thread.read(cx).is_generating() {
+            return;
+        }
+
         let snapshot = self.multibuffer.read(cx).snapshot(cx);
         let diff_hunks_in_ranges = self
             .editor
@@ -650,6 +658,11 @@ fn render_diff_hunk_controls(
     cx: &mut App,
 ) -> AnyElement {
     let editor = editor.clone();
+
+    if agent_diff.read(cx).thread.read(cx).is_generating() {
+        return Empty.into_any();
+    }
+
     h_flex()
         .h(line_height)
         .mr_0p5()
@@ -857,8 +870,14 @@ impl Render for AgentDiffToolbar {
             None => return div(),
         };
 
-        let is_empty = agent_diff.read(cx).multibuffer.read(cx).is_empty();
+        let is_generating = agent_diff.read(cx).thread.read(cx).is_generating();
+        if is_generating {
+            return div()
+                .w(rems(6.5625)) // Arbitrary 105px size—so the label doesn't dance around
+                .child(AnimatedLabel::new("Generating"));
+        }
 
+        let is_empty = agent_diff.read(cx).multibuffer.read(cx).is_empty();
         if is_empty {
             return div();
         }
@@ -969,7 +988,7 @@ mod tests {
             .await
             .unwrap();
         cx.update(|_, cx| {
-            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit(

crates/agent/src/assistant_panel.rs

@@ -25,7 +25,7 @@ use language_model::{LanguageModelProviderTosView, LanguageModelRegistry};
 use language_model_selector::ToggleModelSelector;
 use project::Project;
 use prompt_library::{PromptLibrary, open_prompt_library};
-use prompt_store::{PromptBuilder, PromptId};
+use prompt_store::{PromptBuilder, PromptId, UserPromptId};
 use proto::Plan;
 use settings::{Settings, update_settings_file};
 use time::UtcOffset;
@@ -79,11 +79,11 @@ pub fn init(cx: &mut App) {
                         panel.update(cx, |panel, cx| panel.new_prompt_editor(window, cx));
                     }
                 })
-                .register_action(|workspace, _: &OpenPromptLibrary, window, cx| {
+                .register_action(|workspace, action: &OpenPromptLibrary, window, cx| {
                     if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {
                         workspace.focus_panel::<AssistantPanel>(window, cx);
                         panel.update(cx, |panel, cx| {
-                            panel.deploy_prompt_library(&OpenPromptLibrary::default(), window, cx)
+                            panel.deploy_prompt_library(action, window, cx)
                         });
                     }
                 })
@@ -356,14 +356,9 @@ impl AssistantPanel {
         &self.thread_store
     }
 
-    fn cancel(
-        &mut self,
-        _: &editor::actions::Cancel,
-        _window: &mut Window,
-        cx: &mut Context<Self>,
-    ) {
+    fn cancel(&mut self, _: &editor::actions::Cancel, window: &mut Window, cx: &mut Context<Self>) {
         self.thread
-            .update(cx, |thread, cx| thread.cancel_last_completion(cx));
+            .update(cx, |thread, cx| thread.cancel_last_completion(window, cx));
     }
 
     fn new_thread(&mut self, action: &NewThread, window: &mut Window, cx: &mut Context<Self>) {
@@ -502,7 +497,9 @@ impl AssistantPanel {
                     None,
                 ))
             }),
-            action.prompt_to_focus.map(|uuid| PromptId::User { uuid }),
+            action.prompt_to_select.map(|uuid| PromptId::User {
+                uuid: UserPromptId(uuid),
+            }),
             cx,
         )
         .detach_and_log_err(cx);
@@ -1548,7 +1545,7 @@ impl AssistantPanel {
     fn render_usage_banner(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
         let usage = self.thread.read(cx).last_usage()?;
 
-        Some(UsageBanner::new(zed_llm_client::Plan::ZedProTrial, usage.amount).into_any_element())
+        Some(UsageBanner::new(zed_llm_client::Plan::ZedProTrial, usage).into_any_element())
     }
 
     fn render_last_error(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
@@ -1949,7 +1946,9 @@ impl AssistantPanelDelegate for ConcreteAssistantPanelDelegate {
                                 .collect::<Vec<_>>();
 
                             for (buffer, range) in selection_ranges {
-                                store.add_excerpt(range, buffer, cx).detach_and_log_err(cx);
+                                store
+                                    .add_selection(buffer, range, cx)
+                                    .detach_and_log_err(cx);
                             }
                         })
                     })

crates/agent/src/context.rs

@@ -1,9 +1,15 @@
-use std::{ops::Range, path::Path, sync::Arc};
+use std::{
+    ops::Range,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
-use gpui::{App, Entity, SharedString};
-use language::{Buffer, File};
-use language_model::LanguageModelRequestMessage;
-use project::{ProjectPath, Worktree};
+use futures::{FutureExt, future::Shared};
+use gpui::{App, Entity, SharedString, Task};
+use language::Buffer;
+use language_model::{LanguageModelImage, LanguageModelRequestMessage};
+use project::{ProjectEntryId, ProjectPath, Worktree};
+use prompt_store::UserPromptId;
 use rope::Point;
 use serde::{Deserialize, Serialize};
 use text::{Anchor, BufferId};
@@ -12,6 +18,8 @@ use util::post_inc;
 
 use crate::thread::Thread;
 
+pub const RULES_ICON: IconName = IconName::Context;
+
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
 pub struct ContextId(pub(crate) usize);
 
@@ -20,13 +28,16 @@ impl ContextId {
         Self(post_inc(&mut self.0))
     }
 }
+
 pub enum ContextKind {
     File,
     Directory,
     Symbol,
-    Excerpt,
+    Selection,
     FetchedUrl,
     Thread,
+    Rules,
+    Image,
 }
 
 impl ContextKind {
@@ -35,9 +46,11 @@ impl ContextKind {
             ContextKind::File => IconName::File,
             ContextKind::Directory => IconName::Folder,
             ContextKind::Symbol => IconName::Code,
-            ContextKind::Excerpt => IconName::Code,
+            ContextKind::Selection => IconName::Context,
             ContextKind::FetchedUrl => IconName::Globe,
             ContextKind::Thread => IconName::MessageBubbles,
+            ContextKind::Rules => RULES_ICON,
+            ContextKind::Image => IconName::Image,
         }
     }
 }
@@ -49,7 +62,9 @@ pub enum AssistantContext {
     Symbol(SymbolContext),
     FetchedUrl(FetchedUrlContext),
     Thread(ThreadContext),
-    Excerpt(ExcerptContext),
+    Selection(SelectionContext),
+    Rules(RulesContext),
+    Image(ImageContext),
 }
 
 impl AssistantContext {
@@ -60,7 +75,9 @@ impl AssistantContext {
             Self::Symbol(symbol) => symbol.id,
             Self::FetchedUrl(url) => url.id,
             Self::Thread(thread) => thread.id,
-            Self::Excerpt(excerpt) => excerpt.id,
+            Self::Selection(selection) => selection.id,
+            Self::Rules(rules) => rules.id,
+            Self::Image(image) => image.id,
         }
     }
 }
@@ -75,17 +92,25 @@ pub struct FileContext {
 pub struct DirectoryContext {
     pub id: ContextId,
     pub worktree: Entity<Worktree>,
-    pub path: Arc<Path>,
+    pub entry_id: ProjectEntryId,
+    pub last_path: Arc<Path>,
     /// Buffers of the files within the directory.
     pub context_buffers: Vec<ContextBuffer>,
 }
 
 impl DirectoryContext {
-    pub fn project_path(&self, cx: &App) -> ProjectPath {
-        ProjectPath {
-            worktree_id: self.worktree.read(cx).id(),
-            path: self.path.clone(),
-        }
+    pub fn entry<'a>(&self, cx: &'a App) -> Option<&'a project::Entry> {
+        self.worktree.read(cx).entry_for_id(self.entry_id)
+    }
+
+    pub fn project_path(&self, cx: &App) -> Option<ProjectPath> {
+        let worktree = self.worktree.read(cx);
+        worktree
+            .entry_for_id(self.entry_id)
+            .map(|entry| ProjectPath {
+                worktree_id: worktree.id(),
+                path: entry.path.clone(),
+            })
     }
 }
 
@@ -120,17 +145,51 @@ impl ThreadContext {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct ImageContext {
+    pub id: ContextId,
+    pub original_image: Arc<gpui::Image>,
+    pub image_task: Shared<Task<Option<LanguageModelImage>>>,
+}
+
+impl ImageContext {
+    pub fn image(&self) -> Option<LanguageModelImage> {
+        self.image_task.clone().now_or_never().flatten()
+    }
+
+    pub fn is_loading(&self) -> bool {
+        self.image_task.clone().now_or_never().is_none()
+    }
+
+    pub fn is_error(&self) -> bool {
+        self.image_task
+            .clone()
+            .now_or_never()
+            .map(|result| result.is_none())
+            .unwrap_or(false)
+    }
+}
+
 #[derive(Clone)]
 pub struct ContextBuffer {
     pub id: BufferId,
-    // TODO: Entity<Buffer> holds onto the thread even if the thread is deleted. Should probably be
+    // TODO: Entity<Buffer> holds onto the buffer even if the buffer is deleted. Should probably be
     // a WeakEntity and handle removal from the UI when it has dropped.
     pub buffer: Entity<Buffer>,
-    pub file: Arc<dyn File>,
+    pub last_full_path: Arc<Path>,
     pub version: clock::Global,
     pub text: SharedString,
 }
 
+impl ContextBuffer {
+    pub fn full_path(&self, cx: &App) -> PathBuf {
+        let file = self.buffer.read(cx).file();
+        // Note that in practice file can't be `None` because it is present when this is created and
+        // there's no way for buffers to go from having a file to not.
+        file.map_or(self.last_full_path.to_path_buf(), |file| file.full_path(cx))
+    }
+}
+
 impl std::fmt::Debug for ContextBuffer {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("ContextBuffer")
@@ -161,13 +220,21 @@ pub struct ContextSymbolId {
 }
 
 #[derive(Debug, Clone)]
-pub struct ExcerptContext {
+pub struct SelectionContext {
     pub id: ContextId,
     pub range: Range<Anchor>,
     pub line_range: Range<Point>,
     pub context_buffer: ContextBuffer,
 }
 
+#[derive(Debug, Clone)]
+pub struct RulesContext {
+    pub id: ContextId,
+    pub prompt_id: UserPromptId,
+    pub title: SharedString,
+    pub text: SharedString,
+}
+
 /// Formats a collection of contexts into a string representation
 pub fn format_context_as_string<'a>(
     contexts: impl Iterator<Item = &'a AssistantContext>,
@@ -176,27 +243,31 @@ pub fn format_context_as_string<'a>(
     let mut file_context = Vec::new();
     let mut directory_context = Vec::new();
     let mut symbol_context = Vec::new();
-    let mut excerpt_context = Vec::new();
+    let mut selection_context = Vec::new();
     let mut fetch_context = Vec::new();
     let mut thread_context = Vec::new();
+    let mut rules_context = Vec::new();
 
     for context in contexts {
         match context {
             AssistantContext::File(context) => file_context.push(context),
             AssistantContext::Directory(context) => directory_context.push(context),
             AssistantContext::Symbol(context) => symbol_context.push(context),
-            AssistantContext::Excerpt(context) => excerpt_context.push(context),
+            AssistantContext::Selection(context) => selection_context.push(context),
             AssistantContext::FetchedUrl(context) => fetch_context.push(context),
             AssistantContext::Thread(context) => thread_context.push(context),
+            AssistantContext::Rules(context) => rules_context.push(context),
+            AssistantContext::Image(_) => {}
         }
     }
 
     if file_context.is_empty()
         && directory_context.is_empty()
         && symbol_context.is_empty()
-        && excerpt_context.is_empty()
+        && selection_context.is_empty()
         && fetch_context.is_empty()
         && thread_context.is_empty()
+        && rules_context.is_empty()
     {
         return None;
     }
@@ -232,13 +303,13 @@ pub fn format_context_as_string<'a>(
         result.push_str("</symbols>\n");
     }
 
-    if !excerpt_context.is_empty() {
-        result.push_str("<excerpts>\n");
-        for context in excerpt_context {
+    if !selection_context.is_empty() {
+        result.push_str("<selections>\n");
+        for context in selection_context {
             result.push_str(&context.context_buffer.text);
             result.push('\n');
         }
-        result.push_str("</excerpts>\n");
+        result.push_str("</selections>\n");
     }
 
     if !fetch_context.is_empty() {
@@ -263,6 +334,18 @@ pub fn format_context_as_string<'a>(
         result.push_str("</conversation_threads>\n");
     }
 
+    if !rules_context.is_empty() {
+        result.push_str(
+            "<user_rules>\n\
+            The user has specified the following rules that should be applied:\n\n",
+        );
+        for context in &rules_context {
+            result.push_str(&context.text);
+            result.push('\n');
+        }
+        result.push_str("</user_rules>\n");
+    }
+
     result.push_str("</context>\n");
     Some(result)
 }

crates/agent/src/context_picker.rs

@@ -1,6 +1,7 @@
 mod completion_provider;
 mod fetch_context_picker;
 mod file_context_picker;
+mod rules_context_picker;
 mod symbol_context_picker;
 mod thread_context_picker;
 
@@ -16,30 +17,91 @@ use gpui::{
     App, DismissEvent, Empty, Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task,
     WeakEntity,
 };
+use language::Buffer;
 use multi_buffer::MultiBufferRow;
 use project::{Entry, ProjectPath};
+use prompt_store::UserPromptId;
+use rules_context_picker::RulesContextEntry;
 use symbol_context_picker::SymbolContextPicker;
 use thread_context_picker::{ThreadContextEntry, render_thread_context_entry};
 use ui::{
     ButtonLike, ContextMenu, ContextMenuEntry, ContextMenuItem, Disclosure, TintColor, prelude::*,
 };
+use uuid::Uuid;
 use workspace::{Workspace, notifications::NotifyResultExt};
 
 use crate::AssistantPanel;
+use crate::context::RULES_ICON;
 pub use crate::context_picker::completion_provider::ContextPickerCompletionProvider;
 use crate::context_picker::fetch_context_picker::FetchContextPicker;
 use crate::context_picker::file_context_picker::FileContextPicker;
+use crate::context_picker::rules_context_picker::RulesContextPicker;
 use crate::context_picker::thread_context_picker::ThreadContextPicker;
 use crate::context_store::ContextStore;
 use crate::thread::ThreadId;
 use crate::thread_store::ThreadStore;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum ContextPickerEntry {
+    Mode(ContextPickerMode),
+    Action(ContextPickerAction),
+}
+
+impl ContextPickerEntry {
+    pub fn keyword(&self) -> &'static str {
+        match self {
+            Self::Mode(mode) => mode.keyword(),
+            Self::Action(action) => action.keyword(),
+        }
+    }
+
+    pub fn label(&self) -> &'static str {
+        match self {
+            Self::Mode(mode) => mode.label(),
+            Self::Action(action) => action.label(),
+        }
+    }
+
+    pub fn icon(&self) -> IconName {
+        match self {
+            Self::Mode(mode) => mode.icon(),
+            Self::Action(action) => action.icon(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum ContextPickerMode {
     File,
     Symbol,
     Fetch,
     Thread,
+    Rules,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum ContextPickerAction {
+    AddSelections,
+}
+
+impl ContextPickerAction {
+    pub fn keyword(&self) -> &'static str {
+        match self {
+            Self::AddSelections => "selection",
+        }
+    }
+
+    pub fn label(&self) -> &'static str {
+        match self {
+            Self::AddSelections => "Selection",
+        }
+    }
+
+    pub fn icon(&self) -> IconName {
+        match self {
+            Self::AddSelections => IconName::Context,
+        }
+    }
 }
 
 impl TryFrom<&str> for ContextPickerMode {
@@ -51,18 +113,20 @@ impl TryFrom<&str> for ContextPickerMode {
             "symbol" => Ok(Self::Symbol),
             "fetch" => Ok(Self::Fetch),
             "thread" => Ok(Self::Thread),
+            "rules" => Ok(Self::Rules),
             _ => Err(format!("Invalid context picker mode: {}", value)),
         }
     }
 }
 
 impl ContextPickerMode {
-    pub fn mention_prefix(&self) -> &'static str {
+    pub fn keyword(&self) -> &'static str {
         match self {
             Self::File => "file",
             Self::Symbol => "symbol",
             Self::Fetch => "fetch",
             Self::Thread => "thread",
+            Self::Rules => "rules",
         }
     }
 
@@ -72,6 +136,7 @@ impl ContextPickerMode {
             Self::Symbol => "Symbols",
             Self::Fetch => "Fetch",
             Self::Thread => "Threads",
+            Self::Rules => "Rules",
         }
     }
 
@@ -81,6 +146,7 @@ impl ContextPickerMode {
             Self::Symbol => IconName::Code,
             Self::Fetch => IconName::Globe,
             Self::Thread => IconName::MessageBubbles,
+            Self::Rules => RULES_ICON,
         }
     }
 }
@@ -92,6 +158,7 @@ enum ContextPickerState {
     Symbol(Entity<SymbolContextPicker>),
     Fetch(Entity<FetchContextPicker>),
     Thread(Entity<ThreadContextPicker>),
+    Rules(Entity<RulesContextPicker>),
 }
 
 pub(super) struct ContextPicker {
@@ -155,7 +222,13 @@ impl ContextPicker {
                 .enumerate()
                 .map(|(ix, entry)| self.recent_menu_item(context_picker.clone(), ix, entry));
 
-            let modes = supported_context_picker_modes(&self.thread_store);
+            let entries = self
+                .workspace
+                .upgrade()
+                .map(|workspace| {
+                    available_context_picker_entries(&self.thread_store, &workspace, cx)
+                })
+                .unwrap_or_default();
 
             menu.when(has_recent, |menu| {
                 menu.custom_row(|_, _| {
@@ -171,15 +244,15 @@ impl ContextPicker {
             })
             .extend(recent_entries)
             .when(has_recent, |menu| menu.separator())
-            .extend(modes.into_iter().map(|mode| {
+            .extend(entries.into_iter().map(|entry| {
                 let context_picker = context_picker.clone();
 
-                ContextMenuEntry::new(mode.label())
-                    .icon(mode.icon())
+                ContextMenuEntry::new(entry.label())
+                    .icon(entry.icon())
                     .icon_size(IconSize::XSmall)
                     .icon_color(Color::Muted)
                     .handler(move |window, cx| {
-                        context_picker.update(cx, |this, cx| this.select_mode(mode, window, cx))
+                        context_picker.update(cx, |this, cx| this.select_entry(entry, window, cx))
                     })
             }))
             .keep_open_on_confirm()
@@ -198,61 +271,87 @@ impl ContextPicker {
         self.thread_store.is_some()
     }
 
-    fn select_mode(
+    fn select_entry(
         &mut self,
-        mode: ContextPickerMode,
+        entry: ContextPickerEntry,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
         let context_picker = cx.entity().downgrade();
 
-        match mode {
-            ContextPickerMode::File => {
-                self.mode = ContextPickerState::File(cx.new(|cx| {
-                    FileContextPicker::new(
-                        context_picker.clone(),
-                        self.workspace.clone(),
-                        self.context_store.clone(),
-                        window,
-                        cx,
-                    )
-                }));
-            }
-            ContextPickerMode::Symbol => {
-                self.mode = ContextPickerState::Symbol(cx.new(|cx| {
-                    SymbolContextPicker::new(
-                        context_picker.clone(),
-                        self.workspace.clone(),
-                        self.context_store.clone(),
-                        window,
-                        cx,
-                    )
-                }));
-            }
-            ContextPickerMode::Fetch => {
-                self.mode = ContextPickerState::Fetch(cx.new(|cx| {
-                    FetchContextPicker::new(
-                        context_picker.clone(),
-                        self.workspace.clone(),
-                        self.context_store.clone(),
-                        window,
-                        cx,
-                    )
-                }));
-            }
-            ContextPickerMode::Thread => {
-                if let Some(thread_store) = self.thread_store.as_ref() {
-                    self.mode = ContextPickerState::Thread(cx.new(|cx| {
-                        ThreadContextPicker::new(
-                            thread_store.clone(),
+        match entry {
+            ContextPickerEntry::Mode(mode) => match mode {
+                ContextPickerMode::File => {
+                    self.mode = ContextPickerState::File(cx.new(|cx| {
+                        FileContextPicker::new(
                             context_picker.clone(),
+                            self.workspace.clone(),
                             self.context_store.clone(),
                             window,
                             cx,
                         )
                     }));
                 }
-            }
+                ContextPickerMode::Symbol => {
+                    self.mode = ContextPickerState::Symbol(cx.new(|cx| {
+                        SymbolContextPicker::new(
+                            context_picker.clone(),
+                            self.workspace.clone(),
+                            self.context_store.clone(),
+                            window,
+                            cx,
+                        )
+                    }));
+                }
+                ContextPickerMode::Rules => {
+                    if let Some(thread_store) = self.thread_store.as_ref() {
+                        self.mode = ContextPickerState::Rules(cx.new(|cx| {
+                            RulesContextPicker::new(
+                                thread_store.clone(),
+                                context_picker.clone(),
+                                self.context_store.clone(),
+                                window,
+                                cx,
+                            )
+                        }));
+                    }
+                }
+                ContextPickerMode::Fetch => {
+                    self.mode = ContextPickerState::Fetch(cx.new(|cx| {
+                        FetchContextPicker::new(
+                            context_picker.clone(),
+                            self.workspace.clone(),
+                            self.context_store.clone(),
+                            window,
+                            cx,
+                        )
+                    }));
+                }
+                ContextPickerMode::Thread => {
+                    if let Some(thread_store) = self.thread_store.as_ref() {
+                        self.mode = ContextPickerState::Thread(cx.new(|cx| {
+                            ThreadContextPicker::new(
+                                thread_store.clone(),
+                                context_picker.clone(),
+                                self.context_store.clone(),
+                                window,
+                                cx,
+                            )
+                        }));
+                    }
+                }
+            },
+            ContextPickerEntry::Action(action) => match action {
+                ContextPickerAction::AddSelections => {
+                    if let Some((context_store, workspace)) =
+                        self.context_store.upgrade().zip(self.workspace.upgrade())
+                    {
+                        add_selections_as_context(&context_store, &workspace, cx);
+                    }
+
+                    cx.emit(DismissEvent);
+                }
+            },
         }
 
         cx.notify();
@@ -381,6 +480,7 @@ impl ContextPicker {
             ContextPickerState::Symbol(entity) => entity.update(cx, |_, cx| cx.notify()),
             ContextPickerState::Fetch(entity) => entity.update(cx, |_, cx| cx.notify()),
             ContextPickerState::Thread(entity) => entity.update(cx, |_, cx| cx.notify()),
+            ContextPickerState::Rules(entity) => entity.update(cx, |_, cx| cx.notify()),
         }
     }
 }
@@ -395,6 +495,7 @@ impl Focusable for ContextPicker {
             ContextPickerState::Symbol(symbol_picker) => symbol_picker.focus_handle(cx),
             ContextPickerState::Fetch(fetch_picker) => fetch_picker.focus_handle(cx),
             ContextPickerState::Thread(thread_picker) => thread_picker.focus_handle(cx),
+            ContextPickerState::Rules(user_rules_picker) => user_rules_picker.focus_handle(cx),
         }
     }
 }
@@ -410,6 +511,9 @@ impl Render for ContextPicker {
                 ContextPickerState::Symbol(symbol_picker) => parent.child(symbol_picker.clone()),
                 ContextPickerState::Fetch(fetch_picker) => parent.child(fetch_picker.clone()),
                 ContextPickerState::Thread(thread_picker) => parent.child(thread_picker.clone()),
+                ContextPickerState::Rules(user_rules_picker) => {
+                    parent.child(user_rules_picker.clone())
+                }
             })
     }
 }
@@ -421,18 +525,37 @@ enum RecentEntry {
     Thread(ThreadContextEntry),
 }
 
-fn supported_context_picker_modes(
+fn available_context_picker_entries(
     thread_store: &Option<WeakEntity<ThreadStore>>,
-) -> Vec<ContextPickerMode> {
-    let mut modes = vec![
-        ContextPickerMode::File,
-        ContextPickerMode::Symbol,
-        ContextPickerMode::Fetch,
+    workspace: &Entity<Workspace>,
+    cx: &mut App,
+) -> Vec<ContextPickerEntry> {
+    let mut entries = vec![
+        ContextPickerEntry::Mode(ContextPickerMode::File),
+        ContextPickerEntry::Mode(ContextPickerMode::Symbol),
     ];
-    if thread_store.is_some() {
-        modes.push(ContextPickerMode::Thread);
+
+    let has_selection = workspace
+        .read(cx)
+        .active_item(cx)
+        .and_then(|item| item.downcast::<Editor>())
+        .map_or(false, |editor| {
+            editor.update(cx, |editor, cx| editor.has_non_empty_selection(cx))
+        });
+    if has_selection {
+        entries.push(ContextPickerEntry::Action(
+            ContextPickerAction::AddSelections,
+        ));
     }
-    modes
+
+    if thread_store.is_some() {
+        entries.push(ContextPickerEntry::Mode(ContextPickerMode::Thread));
+        entries.push(ContextPickerEntry::Mode(ContextPickerMode::Rules));
+    }
+
+    entries.push(ContextPickerEntry::Mode(ContextPickerMode::Fetch));
+
+    entries
 }
 
 fn recent_context_picker_entries(
@@ -491,6 +614,54 @@ fn recent_context_picker_entries(
     recent
 }
 
+fn add_selections_as_context(
+    context_store: &Entity<ContextStore>,
+    workspace: &Entity<Workspace>,
+    cx: &mut App,
+) {
+    let selection_ranges = selection_ranges(workspace, cx);
+    context_store.update(cx, |context_store, cx| {
+        for (buffer, range) in selection_ranges {
+            context_store
+                .add_selection(buffer, range, cx)
+                .detach_and_log_err(cx);
+        }
+    })
+}
+
+fn selection_ranges(
+    workspace: &Entity<Workspace>,
+    cx: &mut App,
+) -> Vec<(Entity<Buffer>, Range<text::Anchor>)> {
+    let Some(editor) = workspace
+        .read(cx)
+        .active_item(cx)
+        .and_then(|item| item.act_as::<Editor>(cx))
+    else {
+        return Vec::new();
+    };
+
+    editor.update(cx, |editor, cx| {
+        let selections = editor.selections.all_adjusted(cx);
+
+        let buffer = editor.buffer().clone().read(cx);
+        let snapshot = buffer.snapshot(cx);
+
+        selections
+            .into_iter()
+            .map(|s| snapshot.anchor_after(s.start)..snapshot.anchor_before(s.end))
+            .flat_map(|range| {
+                let (start_buffer, start) = buffer.text_anchor_for_position(range.start, cx)?;
+                let (end_buffer, end) = buffer.text_anchor_for_position(range.end, cx)?;
+                if start_buffer != end_buffer {
+                    return None;
+                }
+                Some((start_buffer, start..end))
+            })
+            .collect::<Vec<_>>()
+    })
+}
+
 pub(crate) fn insert_fold_for_mention(
     excerpt_id: ExcerptId,
     crease_start: text::Anchor,
@@ -510,24 +681,11 @@ pub(crate) fn insert_fold_for_mention(
         let start = start.bias_right(&snapshot);
         let end = snapshot.anchor_before(start.to_offset(&snapshot) + content_len);
 
-        let placeholder = FoldPlaceholder {
-            render: render_fold_icon_button(
-                crease_icon_path,
-                crease_label,
-                editor_entity.downgrade(),
-            ),
-            merge_adjacent: false,
-            ..Default::default()
-        };
-
-        let render_trailer =
-            move |_row, _unfold, _window: &mut Window, _cx: &mut App| Empty.into_any();
-
-        let crease = Crease::inline(
+        let crease = crease_for_mention(
+            crease_label,
+            crease_icon_path,
             start..end,
-            placeholder.clone(),
-            fold_toggle("mention"),
-            render_trailer,
+            editor_entity.downgrade(),
         );
 
         editor.display_map.update(cx, |display_map, cx| {
@@ -536,6 +694,29 @@ pub(crate) fn insert_fold_for_mention(
     });
 }
 
+pub fn crease_for_mention(
+    label: SharedString,
+    icon_path: SharedString,
+    range: Range<Anchor>,
+    editor_entity: WeakEntity<Editor>,
+) -> Crease<Anchor> {
+    let placeholder = FoldPlaceholder {
+        render: render_fold_icon_button(icon_path, label, editor_entity),
+        merge_adjacent: false,
+        ..Default::default()
+    };
+
+    let render_trailer = move |_row, _unfold, _window: &mut Window, _cx: &mut App| Empty.into_any();
+
+    let crease = Crease::inline(
+        range,
+        placeholder.clone(),
+        fold_toggle("mention"),
+        render_trailer,
+    );
+    crease
+}
+
 fn render_fold_icon_button(
     icon_path: SharedString,
     label: SharedString,
@@ -624,15 +805,19 @@ fn fold_toggle(
 pub enum MentionLink {
     File(ProjectPath, Entry),
     Symbol(ProjectPath, String),
+    Selection(ProjectPath, Range<usize>),
     Fetch(String),
     Thread(ThreadId),
+    Rules(UserPromptId),
 }
 
 impl MentionLink {
     const FILE: &str = "@file";
     const SYMBOL: &str = "@symbol";
+    const SELECTION: &str = "@selection";
     const THREAD: &str = "@thread";
     const FETCH: &str = "@fetch";
+    const RULES: &str = "@rules";
 
     const SEPARATOR: &str = ":";
 
@@ -640,7 +825,9 @@ impl MentionLink {
         url.starts_with(Self::FILE)
             || url.starts_with(Self::SYMBOL)
             || url.starts_with(Self::FETCH)
+            || url.starts_with(Self::SELECTION)
             || url.starts_with(Self::THREAD)
+            || url.starts_with(Self::RULES)
     }
 
     pub fn for_file(file_name: &str, full_path: &str) -> String {
@@ -657,14 +844,31 @@ impl MentionLink {
         )
     }
 
-    pub fn for_fetch(url: &str) -> String {
-        format!("[@{}]({}:{})", url, Self::FETCH, url)
+    pub fn for_selection(file_name: &str, full_path: &str, line_range: Range<usize>) -> String {
+        format!(
+            "[@{} ({}-{})]({}:{}:{}-{})",
+            file_name,
+            line_range.start,
+            line_range.end,
+            Self::SELECTION,
+            full_path,
+            line_range.start,
+            line_range.end
+        )
     }
 
     pub fn for_thread(thread: &ThreadContextEntry) -> String {
         format!("[@{}]({}:{})", thread.summary, Self::THREAD, thread.id)
     }
 
+    pub fn for_fetch(url: &str) -> String {
+        format!("[@{}]({}:{})", url, Self::FETCH, url)
+    }
+
+    pub fn for_rules(rules: &RulesContextEntry) -> String {
+        format!("[@{}]({}:{})", rules.title, Self::RULES, rules.prompt_id.0)
+    }
+
     pub fn try_parse(link: &str, workspace: &Entity<Workspace>, cx: &App) -> Option<Self> {
         fn extract_project_path_from_link(
             path: &str,
@@ -701,11 +905,29 @@ impl MentionLink {
                 let project_path = extract_project_path_from_link(path, workspace, cx)?;
                 Some(MentionLink::Symbol(project_path, symbol.to_string()))
             }
+            Self::SELECTION => {
+                let (path, line_args) = argument.split_once(Self::SEPARATOR)?;
+                let project_path = extract_project_path_from_link(path, workspace, cx)?;
+
+                let line_range = {
+                    let (start, end) = line_args
+                        .trim_start_matches('(')
+                        .trim_end_matches(')')
+                        .split_once('-')?;
+                    start.parse::<usize>().ok()?..end.parse::<usize>().ok()?
+                };
+
+                Some(MentionLink::Selection(project_path, line_range))
+            }
             Self::THREAD => {
                 let thread_id = ThreadId::from(argument);
                 Some(MentionLink::Thread(thread_id))
             }
             Self::FETCH => Some(MentionLink::Fetch(argument.to_string())),
+            Self::RULES => {
+                let prompt_id = UserPromptId(Uuid::try_parse(argument).ok()?);
+                Some(MentionLink::Rules(prompt_id))
+            }
             _ => None,
         }
     }

crates/agent/src/context_picker/completion_provider.rs

@@ -1,24 +1,27 @@
 use std::cell::RefCell;
 use std::ops::Range;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::rc::Rc;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 
 use anyhow::Result;
-use editor::{CompletionProvider, Editor, ExcerptId};
+use editor::{CompletionProvider, Editor, ExcerptId, ToOffset as _};
 use file_icons::FileIcons;
 use fuzzy::{StringMatch, StringMatchCandidate};
 use gpui::{App, Entity, Task, WeakEntity};
 use http_client::HttpClientWithUrl;
+use itertools::Itertools;
 use language::{Buffer, CodeLabel, HighlightId};
 use lsp::CompletionContext;
 use project::{Completion, CompletionIntent, ProjectPath, Symbol, WorktreeId};
+use prompt_store::PromptId;
 use rope::Point;
-use text::{Anchor, ToPoint};
+use text::{Anchor, OffsetRangeExt, ToPoint};
 use ui::prelude::*;
 use workspace::Workspace;
 
+use crate::context::RULES_ICON;
 use crate::context_picker::file_context_picker::search_files;
 use crate::context_picker::symbol_context_picker::search_symbols;
 use crate::context_store::ContextStore;
@@ -26,11 +29,12 @@ use crate::thread_store::ThreadStore;
 
 use super::fetch_context_picker::fetch_url_content;
 use super::file_context_picker::FileMatch;
+use super::rules_context_picker::{RulesContextEntry, search_rules};
 use super::symbol_context_picker::SymbolMatch;
 use super::thread_context_picker::{ThreadContextEntry, ThreadMatch, search_threads};
 use super::{
-    ContextPickerMode, MentionLink, RecentEntry, recent_context_picker_entries,
-    supported_context_picker_modes,
+    ContextPickerAction, ContextPickerEntry, ContextPickerMode, MentionLink, RecentEntry,
+    available_context_picker_entries, recent_context_picker_entries, selection_ranges,
 };
 
 pub(crate) enum Match {
@@ -38,22 +42,24 @@ pub(crate) enum Match {
     File(FileMatch),
     Thread(ThreadMatch),
     Fetch(SharedString),
-    Mode(ModeMatch),
+    Rules(RulesContextEntry),
+    Entry(EntryMatch),
 }
 
-pub struct ModeMatch {
+pub struct EntryMatch {
     mat: Option<StringMatch>,
-    mode: ContextPickerMode,
+    entry: ContextPickerEntry,
 }
 
 impl Match {
     pub fn score(&self) -> f64 {
         match self {
             Match::File(file) => file.mat.score,
-            Match::Mode(mode) => mode.mat.as_ref().map(|mat| mat.score).unwrap_or(1.),
+            Match::Entry(mode) => mode.mat.as_ref().map(|mat| mat.score).unwrap_or(1.),
             Match::Thread(_) => 1.,
             Match::Symbol(_) => 1.,
             Match::Fetch(_) => 1.,
+            Match::Rules(_) => 1.,
         }
     }
 }
@@ -112,6 +118,21 @@ fn search(
                 Task::ready(Vec::new())
             }
         }
+        Some(ContextPickerMode::Rules) => {
+            if let Some(thread_store) = thread_store.as_ref().and_then(|t| t.upgrade()) {
+                let search_rules_task =
+                    search_rules(query.clone(), cancellation_flag.clone(), thread_store, cx);
+                cx.background_spawn(async move {
+                    search_rules_task
+                        .await
+                        .into_iter()
+                        .map(Match::Rules)
+                        .collect::<Vec<_>>()
+                })
+            } else {
+                Task::ready(Vec::new())
+            }
+        }
         None => {
             if query.is_empty() {
                 let mut matches = recent_entries
@@ -142,9 +163,14 @@ fn search(
                     .collect::<Vec<_>>();
 
                 matches.extend(
-                    supported_context_picker_modes(&thread_store)
+                    available_context_picker_entries(&thread_store, &workspace, cx)
                         .into_iter()
-                        .map(|mode| Match::Mode(ModeMatch { mode, mat: None })),
+                        .map(|mode| {
+                            Match::Entry(EntryMatch {
+                                entry: mode,
+                                mat: None,
+                            })
+                        }),
                 );
 
                 Task::ready(matches)
@@ -154,11 +180,11 @@ fn search(
                 let search_files_task =
                     search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);
 
-                let modes = supported_context_picker_modes(&thread_store);
-                let mode_candidates = modes
+                let entries = available_context_picker_entries(&thread_store, &workspace, cx);
+                let entry_candidates = entries
                     .iter()
                     .enumerate()
-                    .map(|(ix, mode)| StringMatchCandidate::new(ix, mode.mention_prefix()))
+                    .map(|(ix, entry)| StringMatchCandidate::new(ix, entry.keyword()))
                     .collect::<Vec<_>>();
 
                 cx.background_spawn(async move {
@@ -168,8 +194,8 @@ fn search(
                         .map(Match::File)
                         .collect::<Vec<_>>();
 
-                    let mode_matches = fuzzy::match_strings(
-                        &mode_candidates,
+                    let entry_matches = fuzzy::match_strings(
+                        &entry_candidates,
                         &query,
                         false,
                         100,
@@ -178,9 +204,9 @@ fn search(
                     )
                     .await;
 
-                    matches.extend(mode_matches.into_iter().map(|mat| {
-                        Match::Mode(ModeMatch {
-                            mode: modes[mat.candidate_id],
+                    matches.extend(entry_matches.into_iter().map(|mat| {
+                        Match::Entry(EntryMatch {
+                            entry: entries[mat.candidate_id],
                             mat: Some(mat),
                         })
                     }));
@@ -220,19 +246,137 @@ impl ContextPickerCompletionProvider {
         }
     }
 
-    fn completion_for_mode(source_range: Range<Anchor>, mode: ContextPickerMode) -> Completion {
-        Completion {
-            replace_range: source_range.clone(),
-            new_text: format!("@{} ", mode.mention_prefix()),
-            label: CodeLabel::plain(mode.label().to_string(), None),
-            icon_path: Some(mode.icon().path().into()),
-            documentation: None,
-            source: project::CompletionSource::Custom,
-            insert_text_mode: None,
-            // This ensures that when a user accepts this completion, the
-            // completion menu will still be shown after "@category " is
-            // inserted
-            confirm: Some(Arc::new(|_, _, _| true)),
+    fn completion_for_entry(
+        entry: ContextPickerEntry,
+        excerpt_id: ExcerptId,
+        source_range: Range<Anchor>,
+        editor: Entity<Editor>,
+        context_store: Entity<ContextStore>,
+        workspace: &Entity<Workspace>,
+        cx: &mut App,
+    ) -> Option<Completion> {
+        match entry {
+            ContextPickerEntry::Mode(mode) => Some(Completion {
+                replace_range: source_range.clone(),
+                new_text: format!("@{} ", mode.keyword()),
+                label: CodeLabel::plain(mode.label().to_string(), None),
+                icon_path: Some(mode.icon().path().into()),
+                documentation: None,
+                source: project::CompletionSource::Custom,
+                insert_text_mode: None,
+                // This ensures that when a user accepts this completion, the
+                // completion menu will still be shown after "@category " is
+                // inserted
+                confirm: Some(Arc::new(|_, _, _| true)),
+            }),
+            ContextPickerEntry::Action(action) => {
+                let (new_text, on_action) = match action {
+                    ContextPickerAction::AddSelections => {
+                        let selections = selection_ranges(workspace, cx);
+
+                        let selection_infos = selections
+                            .iter()
+                            .map(|(buffer, range)| {
+                                let full_path = buffer
+                                    .read(cx)
+                                    .file()
+                                    .map(|file| file.full_path(cx))
+                                    .unwrap_or_else(|| PathBuf::from("untitled"));
+                                let file_name = full_path
+                                    .file_name()
+                                    .unwrap_or_default()
+                                    .to_string_lossy()
+                                    .to_string();
+                                let line_range = range.to_point(&buffer.read(cx).snapshot());
+
+                                let link = MentionLink::for_selection(
+                                    &file_name,
+                                    &full_path.to_string_lossy(),
+                                    line_range.start.row as usize..line_range.end.row as usize,
+                                );
+                                (file_name, link, line_range)
+                            })
+                            .collect::<Vec<_>>();
+
+                        let new_text = selection_infos.iter().map(|(_, link, _)| link).join(" ");
+
+                        let callback = Arc::new({
+                            let context_store = context_store.clone();
+                            let selections = selections.clone();
+                            let selection_infos = selection_infos.clone();
+                            move |_, _: &mut Window, cx: &mut App| {
+                                context_store.update(cx, |context_store, cx| {
+                                    for (buffer, range) in &selections {
+                                        context_store
+                                            .add_selection(buffer.clone(), range.clone(), cx)
+                                            .detach_and_log_err(cx)
+                                    }
+                                });
+
+                                let editor = editor.clone();
+                                let selection_infos = selection_infos.clone();
+                                cx.defer(move |cx| {
+                                    let mut current_offset = 0;
+                                    for (file_name, link, line_range) in selection_infos.iter() {
+                                        let snapshot =
+                                            editor.read(cx).buffer().read(cx).snapshot(cx);
+                                        let Some(start) = snapshot
+                                            .anchor_in_excerpt(excerpt_id, source_range.start)
+                                        else {
+                                            return;
+                                        };
+
+                                        let offset = start.to_offset(&snapshot) + current_offset;
+                                        let text_len = link.len();
+
+                                        let range = snapshot.anchor_after(offset)
+                                            ..snapshot.anchor_after(offset + text_len);
+
+                                        let crease = super::crease_for_mention(
+                                            format!(
+                                                "{} ({}-{})",
+                                                file_name,
+                                                line_range.start.row + 1,
+                                                line_range.end.row + 1
+                                            )
+                                            .into(),
+                                            IconName::Context.path().into(),
+                                            range,
+                                            editor.downgrade(),
+                                        );
+
+                                        editor.update(cx, |editor, cx| {
+                                            editor.display_map.update(cx, |display_map, cx| {
+                                                display_map.fold(vec![crease], cx);
+                                            });
+                                        });
+
+                                        current_offset += text_len + 1;
+                                    }
+                                });
+
+                                false
+                            }
+                        });
+
+                        (new_text, callback)
+                    }
+                };
+
+                Some(Completion {
+                    replace_range: source_range.clone(),
+                    new_text,
+                    label: CodeLabel::plain(action.label().to_string(), None),
+                    icon_path: Some(action.icon().path().into()),
+                    documentation: None,
+                    source: project::CompletionSource::Custom,
+                    insert_text_mode: None,
+                    // This ensures that when a user accepts this completion, the
+                    // completion menu will still be shown after "@category " is
+                    // inserted
+                    confirm: Some(on_action),
+                })
+            }
         }
     }
 
@@ -287,6 +431,60 @@ impl ContextPickerCompletionProvider {
         }
     }
 
+    fn completion_for_rules(
+        rules: RulesContextEntry,
+        excerpt_id: ExcerptId,
+        source_range: Range<Anchor>,
+        editor: Entity<Editor>,
+        context_store: Entity<ContextStore>,
+        thread_store: Entity<ThreadStore>,
+    ) -> Completion {
+        let new_text = MentionLink::for_rules(&rules);
+        let new_text_len = new_text.len();
+        Completion {
+            replace_range: source_range.clone(),
+            new_text,
+            label: CodeLabel::plain(rules.title.to_string(), None),
+            documentation: None,
+            insert_text_mode: None,
+            source: project::CompletionSource::Custom,
+            icon_path: Some(RULES_ICON.path().into()),
+            confirm: Some(confirm_completion_callback(
+                RULES_ICON.path().into(),
+                rules.title.clone(),
+                excerpt_id,
+                source_range.start,
+                new_text_len,
+                editor.clone(),
+                move |cx| {
+                    let prompt_uuid = rules.prompt_id;
+                    let prompt_id = PromptId::User { uuid: prompt_uuid };
+                    let context_store = context_store.clone();
+                    let Some(prompt_store) = thread_store.read(cx).prompt_store() else {
+                        log::error!("Can't add user rules as prompt store is missing.");
+                        return;
+                    };
+                    let prompt_store = prompt_store.read(cx);
+                    let Some(metadata) = prompt_store.metadata(prompt_id) else {
+                        return;
+                    };
+                    let Some(title) = metadata.title else {
+                        return;
+                    };
+                    let text_task = prompt_store.load(prompt_id, cx);
+
+                    cx.spawn(async move |cx| {
+                        let text = text_task.await?;
+                        context_store.update(cx, |context_store, cx| {
+                            context_store.add_rules(prompt_uuid, title, text, false, cx)
+                        })
+                    })
+                    .detach_and_log_err(cx);
+                },
+            )),
+        }
+    }
+
     fn completion_for_fetch(
         source_range: Range<Anchor>,
         url_to_fetch: SharedString,
@@ -593,6 +791,17 @@ impl CompletionProvider for ContextPickerCompletionProvider {
                                 thread_store,
                             ))
                         }
+                        Match::Rules(user_rules) => {
+                            let thread_store = thread_store.as_ref().and_then(|t| t.upgrade())?;
+                            Some(Self::completion_for_rules(
+                                user_rules,
+                                excerpt_id,
+                                source_range.clone(),
+                                editor.clone(),
+                                context_store.clone(),
+                                thread_store,
+                            ))
+                        }
                         Match::Fetch(url) => Some(Self::completion_for_fetch(
                             source_range.clone(),
                             url,
@@ -601,9 +810,15 @@ impl CompletionProvider for ContextPickerCompletionProvider {
                             context_store.clone(),
                             http_client.clone(),
                         )),
-                        Match::Mode(ModeMatch { mode, .. }) => {
-                            Some(Self::completion_for_mode(source_range.clone(), mode))
-                        }
+                        Match::Entry(EntryMatch { entry, .. }) => Self::completion_for_entry(
+                            entry,
+                            excerpt_id,
+                            source_range.clone(),
+                            editor.clone(),
+                            context_store.clone(),
+                            &workspace,
+                            cx,
+                        ),
                     })
                     .collect()
             })?))

crates/agent/src/context_picker/rules_context_picker.rs

@@ -0,0 +1,248 @@
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+
+use anyhow::anyhow;
+use gpui::{App, DismissEvent, Entity, FocusHandle, Focusable, Task, WeakEntity};
+use picker::{Picker, PickerDelegate};
+use prompt_store::{PromptId, UserPromptId};
+use ui::{ListItem, prelude::*};
+
+use crate::context::RULES_ICON;
+use crate::context_picker::ContextPicker;
+use crate::context_store::{self, ContextStore};
+use crate::thread_store::ThreadStore;
+
+pub struct RulesContextPicker {
+    picker: Entity<Picker<RulesContextPickerDelegate>>,
+}
+
+impl RulesContextPicker {
+    pub fn new(
+        thread_store: WeakEntity<ThreadStore>,
+        context_picker: WeakEntity<ContextPicker>,
+        context_store: WeakEntity<context_store::ContextStore>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Self {
+        let delegate = RulesContextPickerDelegate::new(thread_store, context_picker, context_store);
+        let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx));
+
+        RulesContextPicker { picker }
+    }
+}
+
+impl Focusable for RulesContextPicker {
+    fn focus_handle(&self, cx: &App) -> FocusHandle {
+        self.picker.focus_handle(cx)
+    }
+}
+
+impl Render for RulesContextPicker {
+    fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
+        self.picker.clone()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RulesContextEntry {
+    pub prompt_id: UserPromptId,
+    pub title: SharedString,
+}
+
+pub struct RulesContextPickerDelegate {
+    thread_store: WeakEntity<ThreadStore>,
+    context_picker: WeakEntity<ContextPicker>,
+    context_store: WeakEntity<context_store::ContextStore>,
+    matches: Vec<RulesContextEntry>,
+    selected_index: usize,
+}
+
+impl RulesContextPickerDelegate {
+    pub fn new(
+        thread_store: WeakEntity<ThreadStore>,
+        context_picker: WeakEntity<ContextPicker>,
+        context_store: WeakEntity<context_store::ContextStore>,
+    ) -> Self {
+        RulesContextPickerDelegate {
+            thread_store,
+            context_picker,
+            context_store,
+            matches: Vec::new(),
+            selected_index: 0,
+        }
+    }
+}
+
+impl PickerDelegate for RulesContextPickerDelegate {
+    type ListItem = ListItem;
+
+    fn match_count(&self) -> usize {
+        self.matches.len()
+    }
+
+    fn selected_index(&self) -> usize {
+        self.selected_index
+    }
+
+    fn set_selected_index(
+        &mut self,
+        ix: usize,
+        _window: &mut Window,
+        _cx: &mut Context<Picker<Self>>,
+    ) {
+        self.selected_index = ix;
+    }
+
+    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
+        "Search available rules…".into()
+    }
+
+    fn update_matches(
+        &mut self,
+        query: String,
+        window: &mut Window,
+        cx: &mut Context<Picker<Self>>,
+    ) -> Task<()> {
+        let Some(thread_store) = self.thread_store.upgrade() else {
+            return Task::ready(());
+        };
+
+        let search_task = search_rules(query, Arc::new(AtomicBool::default()), thread_store, cx);
+        cx.spawn_in(window, async move |this, cx| {
+            let matches = search_task.await;
+            this.update(cx, |this, cx| {
+                this.delegate.matches = matches;
+                this.delegate.selected_index = 0;
+                cx.notify();
+            })
+            .ok();
+        })
+    }
+
+    fn confirm(&mut self, _secondary: bool, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
+        let Some(entry) = self.matches.get(self.selected_index) else {
+            return;
+        };
+
+        let Some(thread_store) = self.thread_store.upgrade() else {
+            return;
+        };
+
+        let prompt_id = entry.prompt_id;
+
+        let load_rules_task = thread_store.update(cx, |thread_store, cx| {
+            thread_store.load_rules(prompt_id, cx)
+        });
+
+        cx.spawn(async move |this, cx| {
+            let (metadata, text) = load_rules_task.await?;
+            let Some(title) = metadata.title else {
+                return Err(anyhow!("Encountered user rule with no title when attempting to add it to agent context."));
+            };
+            this.update(cx, |this, cx| {
+                this.delegate
+                    .context_store
+                    .update(cx, |context_store, cx| {
+                        context_store.add_rules(prompt_id, title, text, true, cx)
+                    })
+                    .ok();
+            })
+        })
+        .detach_and_log_err(cx);
+    }
+
+    fn dismissed(&mut self, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
+        self.context_picker
+            .update(cx, |_, cx| {
+                cx.emit(DismissEvent);
+            })
+            .ok();
+    }
+
+    fn render_match(
+        &self,
+        ix: usize,
+        selected: bool,
+        _window: &mut Window,
+        cx: &mut Context<Picker<Self>>,
+    ) -> Option<Self::ListItem> {
+        let thread = &self.matches[ix];
+
+        Some(ListItem::new(ix).inset(true).toggle_state(selected).child(
+            render_thread_context_entry(thread, self.context_store.clone(), cx),
+        ))
+    }
+}
+
+pub fn render_thread_context_entry(
+    user_rules: &RulesContextEntry,
+    context_store: WeakEntity<ContextStore>,
+    cx: &mut App,
+) -> Div {
+    let added = context_store.upgrade().map_or(false, |ctx_store| {
+        ctx_store
+            .read(cx)
+            .includes_user_rules(&user_rules.prompt_id)
+            .is_some()
+    });
+
+    h_flex()
+        .gap_1p5()
+        .w_full()
+        .justify_between()
+        .child(
+            h_flex()
+                .gap_1p5()
+                .max_w_72()
+                .child(
+                    Icon::new(RULES_ICON)
+                        .size(IconSize::XSmall)
+                        .color(Color::Muted),
+                )
+                .child(Label::new(user_rules.title.clone()).truncate()),
+        )
+        .when(added, |el| {
+            el.child(
+                h_flex()
+                    .gap_1()
+                    .child(
+                        Icon::new(IconName::Check)
+                            .size(IconSize::Small)
+                            .color(Color::Success),
+                    )
+                    .child(Label::new("Added").size(LabelSize::Small)),
+            )
+        })
+}
+
+pub(crate) fn search_rules(
+    query: String,
+    cancellation_flag: Arc<AtomicBool>,
+    thread_store: Entity<ThreadStore>,
+    cx: &mut App,
+) -> Task<Vec<RulesContextEntry>> {
+    let Some(prompt_store) = thread_store.read(cx).prompt_store() else {
+        return Task::ready(vec![]);
+    };
+    let search_task = prompt_store.read(cx).search(query, cancellation_flag, cx);
+    cx.background_spawn(async move {
+        search_task
+            .await
+            .into_iter()
+            .flat_map(|metadata| {
+                // Default prompts are filtered out as they are automatically included.
+                if metadata.default {
+                    None
+                } else {
+                    match metadata.id {
+                        PromptId::EditWorkflow => None,
+                        PromptId::User { uuid } => Some(RulesContextEntry {
+                            prompt_id: uuid,
+                            title: metadata.title?,
+                        }),
+                    }
+                }
+            })
+            .collect::<Vec<_>>()
+    })
+}

crates/agent/src/context_picker/thread_context_picker.rs

@@ -103,11 +103,11 @@ impl PickerDelegate for ThreadContextPickerDelegate {
         window: &mut Window,
         cx: &mut Context<Picker<Self>>,
     ) -> Task<()> {
-        let Some(threads) = self.thread_store.upgrade() else {
+        let Some(thread_store) = self.thread_store.upgrade() else {
             return Task::ready(());
         };
 
-        let search_task = search_threads(query, Arc::new(AtomicBool::default()), threads, cx);
+        let search_task = search_threads(query, Arc::new(AtomicBool::default()), thread_store, cx);
         cx.spawn_in(window, async move |this, cx| {
             let matches = search_task.await;
             this.update(cx, |this, cx| {
@@ -217,15 +217,15 @@ pub(crate) fn search_threads(
     thread_store: Entity<ThreadStore>,
     cx: &mut App,
 ) -> Task<Vec<ThreadMatch>> {
-    let threads = thread_store.update(cx, |this, _cx| {
-        this.threads()
-            .into_iter()
-            .map(|thread| ThreadContextEntry {
-                id: thread.id,
-                summary: thread.summary,
-            })
-            .collect::<Vec<_>>()
-    });
+    let threads = thread_store
+        .read(cx)
+        .threads()
+        .into_iter()
+        .map(|thread| ThreadContextEntry {
+            id: thread.id,
+            summary: thread.summary,
+        })
+        .collect::<Vec<_>>();
 
     let executor = cx.background_executor().clone();
     cx.background_spawn(async move {

crates/agent/src/context_store.rs

@@ -6,9 +6,11 @@ use anyhow::{Context as _, Result, anyhow};
 use collections::{BTreeMap, HashMap, HashSet};
 use futures::future::join_all;
 use futures::{self, Future, FutureExt, future};
-use gpui::{App, AppContext as _, Context, Entity, SharedString, Task, WeakEntity};
-use language::{Buffer, File};
-use project::{Project, ProjectItem, ProjectPath, Worktree};
+use gpui::{App, AppContext as _, Context, Entity, Image, SharedString, Task, WeakEntity};
+use language::Buffer;
+use language_model::LanguageModelImage;
+use project::{Project, ProjectEntryId, ProjectItem, ProjectPath, Worktree};
+use prompt_store::UserPromptId;
 use rope::{Point, Rope};
 use text::{Anchor, BufferId, OffsetRangeExt};
 use util::{ResultExt as _, maybe};
@@ -16,7 +18,8 @@ use util::{ResultExt as _, maybe};
 use crate::ThreadStore;
 use crate::context::{
     AssistantContext, ContextBuffer, ContextId, ContextSymbol, ContextSymbolId, DirectoryContext,
-    ExcerptContext, FetchedUrlContext, FileContext, SymbolContext, ThreadContext,
+    FetchedUrlContext, FileContext, ImageContext, RulesContext, SelectionContext, SymbolContext,
+    ThreadContext,
 };
 use crate::context_strip::SuggestedContext;
 use crate::thread::{Thread, ThreadId};
@@ -25,7 +28,6 @@ pub struct ContextStore {
     project: WeakEntity<Project>,
     context: Vec<AssistantContext>,
     thread_store: Option<WeakEntity<ThreadStore>>,
-    // TODO: If an EntityId is used for all context types (like BufferId), can remove ContextId.
     next_context_id: ContextId,
     files: BTreeMap<BufferId, ContextId>,
     directories: HashMap<ProjectPath, ContextId>,
@@ -35,6 +37,7 @@ pub struct ContextStore {
     threads: HashMap<ThreadId, ContextId>,
     thread_summary_tasks: Vec<Task<()>>,
     fetched_urls: HashMap<String, ContextId>,
+    user_rules: HashMap<UserPromptId, ContextId>,
 }
 
 impl ContextStore {
@@ -55,6 +58,7 @@ impl ContextStore {
             threads: HashMap::default(),
             thread_summary_tasks: Vec::new(),
             fetched_urls: HashMap::default(),
+            user_rules: HashMap::default(),
         }
     }
 
@@ -72,6 +76,7 @@ impl ContextStore {
         self.directories.clear();
         self.threads.clear();
         self.fetched_urls.clear();
+        self.user_rules.clear();
     }
 
     pub fn add_file_from_path(
@@ -109,13 +114,12 @@ impl ContextStore {
                 return anyhow::Ok(());
             }
 
-            let (buffer_info, text_task) =
-                this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, cx))??;
-
-            let text = text_task.await;
+            let context_buffer = this
+                .update(cx, |_, cx| load_context_buffer(buffer, cx))??
+                .await;
 
             this.update(cx, |this, cx| {
-                this.insert_file(make_context_buffer(buffer_info, text), cx);
+                this.insert_file(context_buffer, cx);
             })?;
 
             anyhow::Ok(())
@@ -128,14 +132,11 @@ impl ContextStore {
         cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
         cx.spawn(async move |this, cx| {
-            let (buffer_info, text_task) =
-                this.update(cx, |_, cx| collect_buffer_info_and_text(buffer, cx))??;
+            let context_buffer = this
+                .update(cx, |_, cx| load_context_buffer(buffer, cx))??
+                .await;
 
-            let text = text_task.await;
-
-            this.update(cx, |this, cx| {
-                this.insert_file(make_context_buffer(buffer_info, text), cx)
-            })?;
+            this.update(cx, |this, cx| this.insert_file(context_buffer, cx))?;
 
             anyhow::Ok(())
         })
@@ -159,6 +160,14 @@ impl ContextStore {
             return Task::ready(Err(anyhow!("failed to read project")));
         };
 
+        let Some(entry_id) = project
+            .read(cx)
+            .entry_for_path(&project_path, cx)
+            .map(|entry| entry.id)
+        else {
+            return Task::ready(Err(anyhow!("no entry found for directory context")));
+        };
+
         let already_included = match self.includes_directory(&project_path) {
             Some(FileInclusion::Direct(context_id)) => {
                 if remove_if_exists {
@@ -200,27 +209,15 @@ impl ContextStore {
 
             let buffers = open_buffers_task.await;
 
-            let mut buffer_infos = Vec::new();
-            let mut text_tasks = Vec::new();
-            this.update(cx, |_, cx| {
-                // Skip all binary files and other non-UTF8 files
-                for buffer in buffers.into_iter().flatten() {
-                    if let Some((buffer_info, text_task)) =
-                        collect_buffer_info_and_text(buffer, cx).log_err()
-                    {
-                        buffer_infos.push(buffer_info);
-                        text_tasks.push(text_task);
-                    }
-                }
-                anyhow::Ok(())
-            })??;
+            let context_buffer_tasks = this.update(cx, |_, cx| {
+                buffers
+                    .into_iter()
+                    .flatten()
+                    .flat_map(move |buffer| load_context_buffer(buffer, cx).log_err())
+                    .collect::<Vec<_>>()
+            })?;
 
-            let buffer_texts = future::join_all(text_tasks).await;
-            let context_buffers = buffer_infos
-                .into_iter()
-                .zip(buffer_texts)
-                .map(|(info, text)| make_context_buffer(info, text))
-                .collect::<Vec<_>>();
+            let context_buffers = future::join_all(context_buffer_tasks).await;
 
             if context_buffers.is_empty() {
                 let full_path = cx.update(|cx| worktree.read(cx).full_path(&project_path.path))?;
@@ -228,7 +225,7 @@ impl ContextStore {
             }
 
             this.update(cx, |this, cx| {
-                this.insert_directory(worktree, project_path, context_buffers, cx);
+                this.insert_directory(worktree, entry_id, project_path, context_buffers, cx);
             })?;
 
             anyhow::Ok(())
@@ -238,19 +235,21 @@ impl ContextStore {
     fn insert_directory(
         &mut self,
         worktree: Entity<Worktree>,
+        entry_id: ProjectEntryId,
         project_path: ProjectPath,
         context_buffers: Vec<ContextBuffer>,
         cx: &mut Context<Self>,
     ) {
         let id = self.next_context_id.post_inc();
-        let path = project_path.path.clone();
+        let last_path = project_path.path.clone();
         self.directories.insert(project_path, id);
 
         self.context
             .push(AssistantContext::Directory(DirectoryContext {
                 id,
                 worktree,
-                path,
+                entry_id,
+                last_path,
                 context_buffers,
             }));
         cx.notify();
@@ -290,27 +289,23 @@ impl ContextStore {
             }
         }
 
-        let (buffer_info, collect_content_task) = match collect_buffer_info_and_text_for_range(
-            buffer,
-            symbol_enclosing_range.clone(),
-            cx,
-        ) {
-            Ok((_, buffer_info, collect_context_task)) => (buffer_info, collect_context_task),
-            Err(err) => return Task::ready(Err(err)),
-        };
+        let context_buffer_task =
+            match load_context_buffer_range(buffer, symbol_enclosing_range.clone(), cx) {
+                Ok((_line_range, context_buffer_task)) => context_buffer_task,
+                Err(err) => return Task::ready(Err(err)),
+            };
 
         cx.spawn(async move |this, cx| {
-            let content = collect_content_task.await;
+            let context_buffer = context_buffer_task.await;
 
             this.update(cx, |this, cx| {
                 this.insert_symbol(
                     make_context_symbol(
-                        buffer_info,
+                        context_buffer,
                         project_path,
                         symbol_name,
                         symbol_range,
                         symbol_enclosing_range,
-                        content,
                     ),
                     cx,
                 )
@@ -390,6 +385,42 @@ impl ContextStore {
         cx.notify();
     }
 
+    pub fn add_rules(
+        &mut self,
+        prompt_id: UserPromptId,
+        title: impl Into<SharedString>,
+        text: impl Into<SharedString>,
+        remove_if_exists: bool,
+        cx: &mut Context<ContextStore>,
+    ) {
+        if let Some(context_id) = self.includes_user_rules(&prompt_id) {
+            if remove_if_exists {
+                self.remove_context(context_id, cx);
+            }
+        } else {
+            self.insert_user_rules(prompt_id, title, text, cx);
+        }
+    }
+
+    pub fn insert_user_rules(
+        &mut self,
+        prompt_id: UserPromptId,
+        title: impl Into<SharedString>,
+        text: impl Into<SharedString>,
+        cx: &mut Context<ContextStore>,
+    ) {
+        let id = self.next_context_id.post_inc();
+
+        self.user_rules.insert(prompt_id, id);
+        self.context.push(AssistantContext::Rules(RulesContext {
+            id,
+            prompt_id,
+            title: title.into(),
+            text: text.into(),
+        }));
+        cx.notify();
+    }
+
     pub fn add_fetched_url(
         &mut self,
         url: String,
@@ -419,33 +450,54 @@ impl ContextStore {
         cx.notify();
     }
 
-    pub fn add_excerpt(
+    pub fn add_image(&mut self, image: Arc<Image>, cx: &mut Context<ContextStore>) {
+        let image_task = LanguageModelImage::from_image(image.clone(), cx).shared();
+        let id = self.next_context_id.post_inc();
+        self.context.push(AssistantContext::Image(ImageContext {
+            id,
+            original_image: image,
+            image_task,
+        }));
+        cx.notify();
+    }
+
+    pub fn wait_for_images(&self, cx: &App) -> Task<()> {
+        let tasks = self
+            .context
+            .iter()
+            .filter_map(|ctx| match ctx {
+                AssistantContext::Image(ctx) => Some(ctx.image_task.clone()),
+                _ => None,
+            })
+            .collect::<Vec<_>>();
+
+        cx.spawn(async move |_cx| {
+            join_all(tasks).await;
+        })
+    }
+
+    pub fn add_selection(
         &mut self,
-        range: Range<Anchor>,
         buffer: Entity<Buffer>,
+        range: Range<Anchor>,
         cx: &mut Context<ContextStore>,
     ) -> Task<Result<()>> {
         cx.spawn(async move |this, cx| {
-            let (line_range, buffer_info, text_task) = this.update(cx, |_, cx| {
-                collect_buffer_info_and_text_for_range(buffer, range.clone(), cx)
+            let (line_range, context_buffer_task) = this.update(cx, |_, cx| {
+                load_context_buffer_range(buffer, range.clone(), cx)
             })??;
 
-            let text = text_task.await;
+            let context_buffer = context_buffer_task.await;
 
             this.update(cx, |this, cx| {
-                this.insert_excerpt(
-                    make_context_buffer(buffer_info, text),
-                    range,
-                    line_range,
-                    cx,
-                )
+                this.insert_selection(context_buffer, range, line_range, cx)
             })?;
 
             anyhow::Ok(())
         })
     }
 
-    fn insert_excerpt(
+    fn insert_selection(
         &mut self,
         context_buffer: ContextBuffer,
         range: Range<Anchor>,
@@ -453,12 +505,13 @@ impl ContextStore {
         cx: &mut Context<Self>,
     ) {
         let id = self.next_context_id.post_inc();
-        self.context.push(AssistantContext::Excerpt(ExcerptContext {
-            id,
-            range,
-            line_range,
-            context_buffer,
-        }));
+        self.context
+            .push(AssistantContext::Selection(SelectionContext {
+                id,
+                range,
+                line_range,
+                context_buffer,
+            }));
         cx.notify();
     }
 
@@ -511,13 +564,17 @@ impl ContextStore {
                 self.symbol_buffers.remove(&symbol.context_symbol.id);
                 self.symbols.retain(|_, context_id| *context_id != id);
             }
-            AssistantContext::Excerpt(_) => {}
+            AssistantContext::Selection(_) => {}
             AssistantContext::FetchedUrl(_) => {
                 self.fetched_urls.retain(|_, context_id| *context_id != id);
             }
             AssistantContext::Thread(_) => {
                 self.threads.retain(|_, context_id| *context_id != id);
             }
+            AssistantContext::Rules(RulesContext { prompt_id, .. }) => {
+                self.user_rules.remove(&prompt_id);
+            }
+            AssistantContext::Image(_) => {}
         }
 
         cx.notify();
@@ -614,6 +671,10 @@ impl ContextStore {
         self.threads.get(thread_id).copied()
     }
 
+    pub fn includes_user_rules(&self, prompt_id: &UserPromptId) -> Option<ContextId> {
+        self.user_rules.get(prompt_id).copied()
+    }
+
     pub fn includes_url(&self, url: &str) -> Option<ContextId> {
         self.fetched_urls.get(url).copied()
     }
@@ -639,9 +700,11 @@ impl ContextStore {
                 }
                 AssistantContext::Directory(_)
                 | AssistantContext::Symbol(_)
-                | AssistantContext::Excerpt(_)
+                | AssistantContext::Selection(_)
                 | AssistantContext::FetchedUrl(_)
-                | AssistantContext::Thread(_) => None,
+                | AssistantContext::Thread(_)
+                | AssistantContext::Rules(_)
+                | AssistantContext::Image(_) => None,
             })
             .collect()
     }
@@ -656,92 +719,78 @@ pub enum FileInclusion {
     InDirectory(ProjectPath),
 }
 
-// ContextBuffer without text.
-struct BufferInfo {
-    id: BufferId,
-    buffer: Entity<Buffer>,
-    file: Arc<dyn File>,
-    version: clock::Global,
-}
-
-fn make_context_buffer(info: BufferInfo, text: SharedString) -> ContextBuffer {
-    ContextBuffer {
-        id: info.id,
-        buffer: info.buffer,
-        file: info.file,
-        version: info.version,
-        text,
-    }
-}
-
 fn make_context_symbol(
-    info: BufferInfo,
+    context_buffer: ContextBuffer,
     path: ProjectPath,
     name: SharedString,
     range: Range<Anchor>,
     enclosing_range: Range<Anchor>,
-    text: SharedString,
 ) -> ContextSymbol {
     ContextSymbol {
         id: ContextSymbolId { name, range, path },
-        buffer_version: info.version,
+        buffer_version: context_buffer.version,
         enclosing_range,
-        buffer: info.buffer,
-        text,
+        buffer: context_buffer.buffer,
+        text: context_buffer.text,
     }
 }
 
-fn collect_buffer_info_and_text_for_range(
+fn load_context_buffer_range(
     buffer: Entity<Buffer>,
     range: Range<Anchor>,
     cx: &App,
-) -> Result<(Range<Point>, BufferInfo, Task<SharedString>)> {
-    let content = buffer
-        .read(cx)
-        .text_for_range(range.clone())
-        .collect::<Rope>();
-
-    let line_range = range.to_point(&buffer.read(cx).snapshot());
-
-    let buffer_info = collect_buffer_info(buffer, cx)?;
-    let full_path = buffer_info.file.full_path(cx);
-
-    let text_task = cx.background_spawn({
-        let line_range = line_range.clone();
-        async move { to_fenced_codeblock(&full_path, content, Some(line_range)) }
-    });
-
-    Ok((line_range, buffer_info, text_task))
-}
-
-fn collect_buffer_info_and_text(
-    buffer: Entity<Buffer>,
-    cx: &App,
-) -> Result<(BufferInfo, Task<SharedString>)> {
-    let content = buffer.read(cx).as_rope().clone();
-
-    let buffer_info = collect_buffer_info(buffer, cx)?;
-    let full_path = buffer_info.file.full_path(cx);
-
-    let text_task =
-        cx.background_spawn(async move { to_fenced_codeblock(&full_path, content, None) });
-
-    Ok((buffer_info, text_task))
-}
-
-fn collect_buffer_info(buffer: Entity<Buffer>, cx: &App) -> Result<BufferInfo> {
+) -> Result<(Range<Point>, Task<ContextBuffer>)> {
     let buffer_ref = buffer.read(cx);
-    let file = buffer_ref.file().context("file context must have a path")?;
+    let id = buffer_ref.remote_id();
+
+    let file = buffer_ref.file().context("context buffer missing path")?;
+    let full_path = file.full_path(cx);
 
     // Important to collect version at the same time as content so that staleness logic is correct.
     let version = buffer_ref.version();
+    let content = buffer_ref.text_for_range(range.clone()).collect::<Rope>();
+    let line_range = range.to_point(&buffer_ref.snapshot());
 
-    Ok(BufferInfo {
-        buffer,
-        id: buffer_ref.remote_id(),
-        file: file.clone(),
-        version,
-    })
+    // Build the text on a background thread.
+    let task = cx.background_spawn({
+        let line_range = line_range.clone();
+        async move {
+            let text = to_fenced_codeblock(&full_path, content, Some(line_range));
+            ContextBuffer {
+                id,
+                buffer,
+                last_full_path: full_path.into(),
+                version,
+                text,
+            }
+        }
+    });
+
+    Ok((line_range, task))
+}
+
+fn load_context_buffer(buffer: Entity<Buffer>, cx: &App) -> Result<Task<ContextBuffer>> {
+    let buffer_ref = buffer.read(cx);
+    let id = buffer_ref.remote_id();
+
+    let file = buffer_ref.file().context("context buffer missing path")?;
+    let full_path = file.full_path(cx);
+
+    // Important to collect version at the same time as content so that staleness logic is correct.
+    let version = buffer_ref.version();
+    let content = buffer_ref.as_rope().clone();
+
+    // Build the text on a background thread.
+    Ok(cx.background_spawn(async move {
+        let text = to_fenced_codeblock(&full_path, content, None);
+        ContextBuffer {
+            id,
+            buffer,
+            last_full_path: full_path.into(),
+            version,
+            text,
+        }
+    }))
 }
 
 fn to_fenced_codeblock(
@@ -828,6 +877,7 @@ pub fn refresh_context_store_text(
         let task = maybe!({
             match context {
                 AssistantContext::File(file_context) => {
+                    // TODO: Should refresh if the path has changed, as it's in the text.
                     if changed_buffers.is_empty()
                         || changed_buffers.contains(&file_context.context_buffer.buffer)
                     {
@@ -836,8 +886,9 @@ pub fn refresh_context_store_text(
                     }
                 }
                 AssistantContext::Directory(directory_context) => {
-                    let directory_path = directory_context.project_path(cx);
-                    let should_refresh = changed_buffers.is_empty()
+                    let directory_path = directory_context.project_path(cx)?;
+                    let should_refresh = directory_path.path != directory_context.last_path
+                        || changed_buffers.is_empty()
                         || changed_buffers.iter().any(|buffer| {
                             let Some(buffer_path) = buffer.read(cx).project_path(cx) else {
                                 return false;
@@ -847,10 +898,16 @@ pub fn refresh_context_store_text(
 
                     if should_refresh {
                         let context_store = context_store.clone();
-                        return refresh_directory_text(context_store, directory_context, cx);
+                        return refresh_directory_text(
+                            context_store,
+                            directory_context,
+                            directory_path,
+                            cx,
+                        );
                     }
                 }
                 AssistantContext::Symbol(symbol_context) => {
+                    // TODO: Should refresh if the path has changed, as it's in the text.
                     if changed_buffers.is_empty()
                         || changed_buffers.contains(&symbol_context.context_symbol.buffer)
                     {
@@ -858,12 +915,13 @@ pub fn refresh_context_store_text(
                         return refresh_symbol_text(context_store, symbol_context, cx);
                     }
                 }
-                AssistantContext::Excerpt(excerpt_context) => {
+                AssistantContext::Selection(selection_context) => {
+                    // TODO: Should refresh if the path has changed, as it's in the text.
                     if changed_buffers.is_empty()
-                        || changed_buffers.contains(&excerpt_context.context_buffer.buffer)
+                        || changed_buffers.contains(&selection_context.context_buffer.buffer)
                     {
                         let context_store = context_store.clone();
-                        return refresh_excerpt_text(context_store, excerpt_context, cx);
+                        return refresh_selection_text(context_store, selection_context, cx);
                     }
                 }
                 AssistantContext::Thread(thread_context) => {
@@ -876,6 +934,11 @@ pub fn refresh_context_store_text(
                 // and doing the caching properly could be tricky (unless it's already handled by
                 // the HttpClient?).
                 AssistantContext::FetchedUrl(_) => {}
+                AssistantContext::Rules(user_rules_context) => {
+                    let context_store = context_store.clone();
+                    return Some(refresh_user_rules(context_store, user_rules_context, cx));
+                }
+                AssistantContext::Image(_) => {}
             }
 
             None
@@ -914,6 +977,7 @@ fn refresh_file_text(
 fn refresh_directory_text(
     context_store: Entity<ContextStore>,
     directory_context: &DirectoryContext,
+    directory_path: ProjectPath,
     cx: &App,
 ) -> Option<Task<()>> {
     let mut stale = false;
@@ -938,7 +1002,8 @@ fn refresh_directory_text(
 
     let id = directory_context.id;
     let worktree = directory_context.worktree.clone();
-    let path = directory_context.path.clone();
+    let entry_id = directory_context.entry_id;
+    let last_path = directory_path.path;
     Some(cx.spawn(async move |cx| {
         let context_buffers = context_buffers.await;
         context_store
@@ -946,7 +1011,8 @@ fn refresh_directory_text(
                 let new_directory_context = DirectoryContext {
                     id,
                     worktree,
-                    path,
+                    entry_id,
+                    last_path,
                     context_buffers,
                 };
                 context_store.replace_context(AssistantContext::Directory(new_directory_context));
@@ -977,26 +1043,27 @@ fn refresh_symbol_text(
     }
 }
 
-fn refresh_excerpt_text(
+fn refresh_selection_text(
     context_store: Entity<ContextStore>,
-    excerpt_context: &ExcerptContext,
+    selection_context: &SelectionContext,
     cx: &App,
 ) -> Option<Task<()>> {
-    let id = excerpt_context.id;
-    let range = excerpt_context.range.clone();
-    let task = refresh_context_excerpt(&excerpt_context.context_buffer, range.clone(), cx);
+    let id = selection_context.id;
+    let range = selection_context.range.clone();
+    let task = refresh_context_excerpt(&selection_context.context_buffer, range.clone(), cx);
     if let Some(task) = task {
         Some(cx.spawn(async move |cx| {
             let (line_range, context_buffer) = task.await;
             context_store
                 .update(cx, |context_store, _| {
-                    let new_excerpt_context = ExcerptContext {
+                    let new_selection_context = SelectionContext {
                         id,
                         range,
                         line_range,
                         context_buffer,
                     };
-                    context_store.replace_context(AssistantContext::Excerpt(new_excerpt_context));
+                    context_store
+                        .replace_context(AssistantContext::Selection(new_selection_context));
                 })
                 .ok();
         }))
@@ -1026,15 +1093,49 @@ fn refresh_thread_text(
     })
 }
 
-fn refresh_context_buffer(
-    context_buffer: &ContextBuffer,
+fn refresh_user_rules(
+    context_store: Entity<ContextStore>,
+    user_rules_context: &RulesContext,
     cx: &App,
-) -> Option<impl Future<Output = ContextBuffer> + use<>> {
+) -> Task<()> {
+    let id = user_rules_context.id;
+    let prompt_id = user_rules_context.prompt_id;
+    let Some(thread_store) = context_store.read(cx).thread_store.as_ref() else {
+        return Task::ready(());
+    };
+    let Ok(load_task) = thread_store.read_with(cx, |thread_store, cx| {
+        thread_store.load_rules(prompt_id, cx)
+    }) else {
+        return Task::ready(());
+    };
+    cx.spawn(async move |cx| {
+        if let Ok((metadata, text)) = load_task.await {
+            if let Some(title) = metadata.title.clone() {
+                context_store
+                    .update(cx, |context_store, _cx| {
+                        context_store.replace_context(AssistantContext::Rules(RulesContext {
+                            id,
+                            prompt_id,
+                            title,
+                            text: text.into(),
+                        }));
+                    })
+                    .ok();
+                return;
+            }
+        }
+        context_store
+            .update(cx, |context_store, cx| {
+                context_store.remove_context(id, cx);
+            })
+            .ok();
+    })
+}
+
+fn refresh_context_buffer(context_buffer: &ContextBuffer, cx: &App) -> Option<Task<ContextBuffer>> {
     let buffer = context_buffer.buffer.read(cx);
     if buffer.version.changed_since(&context_buffer.version) {
-        let (buffer_info, text_task) =
-            collect_buffer_info_and_text(context_buffer.buffer.clone(), cx).log_err()?;
-        Some(text_task.map(move |text| make_context_buffer(buffer_info, text)))
+        load_context_buffer(context_buffer.buffer.clone(), cx).log_err()
     } else {
         None
     }
@@ -1047,10 +1148,9 @@ fn refresh_context_excerpt(
 ) -> Option<impl Future<Output = (Range<Point>, ContextBuffer)> + use<>> {
     let buffer = context_buffer.buffer.read(cx);
     if buffer.version.changed_since(&context_buffer.version) {
-        let (line_range, buffer_info, text_task) =
-            collect_buffer_info_and_text_for_range(context_buffer.buffer.clone(), range, cx)
-                .log_err()?;
-        Some(text_task.map(move |text| (line_range, make_context_buffer(buffer_info, text))))
+        let (line_range, context_buffer_task) =
+            load_context_buffer_range(context_buffer.buffer.clone(), range, cx).log_err()?;
+        Some(context_buffer_task.map(move |context_buffer| (line_range, context_buffer)))
     } else {
         None
     }
@@ -1063,7 +1163,7 @@ fn refresh_context_symbol(
     let buffer = context_symbol.buffer.read(cx);
     let project_path = buffer.project_path(cx)?;
     if buffer.version.changed_since(&context_symbol.buffer_version) {
-        let (_, buffer_info, text_task) = collect_buffer_info_and_text_for_range(
+        let (_line_range, context_buffer_task) = load_context_buffer_range(
             context_symbol.buffer.clone(),
             context_symbol.enclosing_range.clone(),
             cx,
@@ -1072,15 +1172,8 @@ fn refresh_context_symbol(
         let name = context_symbol.id.name.clone();
         let range = context_symbol.id.range.clone();
         let enclosing_range = context_symbol.enclosing_range.clone();
-        Some(text_task.map(move |text| {
-            make_context_symbol(
-                buffer_info,
-                project_path,
-                name,
-                range,
-                enclosing_range,
-                text,
-            )
+        Some(context_buffer_task.map(move |context_buffer| {
+            make_context_symbol(context_buffer, project_path, name, range, enclosing_range)
         }))
     } else {
         None

crates/agent/src/inline_assistant.rs

@@ -1328,7 +1328,7 @@ impl InlineAssistant {
                 editor.highlight_rows::<InlineAssist>(
                     row_range,
                     cx.theme().status().info_background,
-                    false,
+                    Default::default(),
                     cx,
                 );
             }
@@ -1393,7 +1393,7 @@ impl InlineAssistant {
                     editor.highlight_rows::<DeletedLines>(
                         Anchor::min()..Anchor::max(),
                         cx.theme().status().deleted_background,
-                        false,
+                        Default::default(),
                         cx,
                     );
                     editor

crates/agent/src/message_editor.rs

@@ -6,7 +6,7 @@ use crate::context::{AssistantContext, format_context_as_string};
 use crate::tool_compatibility::{IncompatibleToolsState, IncompatibleToolsTooltip};
 use buffer_diff::BufferDiff;
 use collections::HashSet;
-use editor::actions::MoveUp;
+use editor::actions::{MoveUp, Paste};
 use editor::{
     ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorEvent, EditorMode,
     EditorStyle, MultiBuffer,
@@ -14,8 +14,8 @@ use editor::{
 use file_icons::FileIcons;
 use fs::Fs;
 use gpui::{
-    Animation, AnimationExt, App, Entity, EventEmitter, Focusable, Subscription, Task, TextStyle,
-    WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,
+    Animation, AnimationExt, App, ClipboardEntry, Entity, EventEmitter, Focusable, Subscription,
+    Task, TextStyle, WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,
 };
 use language::{Buffer, Language};
 use language_model::{ConfiguredModel, LanguageModelRegistry, LanguageModelRequestMessage};
@@ -195,6 +195,7 @@ impl MessageEditor {
                 editor.set_mode(EditorMode::Full {
                     scale_ui_elements_with_buffer_font_size: false,
                     show_active_line_background: false,
+                    sized_by_content: false,
                 })
             } else {
                 editor.set_mode(EditorMode::AutoHeight {
@@ -271,15 +272,18 @@ impl MessageEditor {
 
         let refresh_task =
             refresh_context_store_text(self.context_store.clone(), &HashSet::default(), cx);
+        let wait_for_images = self.context_store.read(cx).wait_for_images(cx);
 
         let thread = self.thread.clone();
         let context_store = self.context_store.clone();
         let git_store = self.project.read(cx).git_store().clone();
         let checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
+        let window_handle = window.window_handle();
 
         cx.spawn(async move |this, cx| {
             let checkpoint = checkpoint.await.ok();
             refresh_task.await;
+            wait_for_images.await;
 
             thread
                 .update(cx, |thread, cx| {
@@ -293,7 +297,12 @@ impl MessageEditor {
                     let excerpt_ids = context_store
                         .context()
                         .iter()
-                        .filter(|ctx| matches!(ctx, AssistantContext::Excerpt(_)))
+                        .filter(|ctx| {
+                            matches!(
+                                ctx,
+                                AssistantContext::Selection(_) | AssistantContext::Image(_)
+                            )
+                        })
                         .map(|ctx| ctx.id())
                         .collect::<Vec<_>>();
 
@@ -326,7 +335,7 @@ impl MessageEditor {
             thread
                 .update(cx, |thread, cx| {
                     thread.advance_prompt_id();
-                    thread.send_to_model(model, cx);
+                    thread.send_to_model(model, Some(window_handle), cx);
                 })
                 .log_err();
         })
@@ -334,9 +343,9 @@ impl MessageEditor {
     }
 
     fn stop_current_and_send_new_message(&mut self, window: &mut Window, cx: &mut Context<Self>) {
-        let cancelled = self
-            .thread
-            .update(cx, |thread, cx| thread.cancel_last_completion(cx));
+        let cancelled = self.thread.update(cx, |thread, cx| {
+            thread.cancel_last_completion(Some(window.window_handle()), cx)
+        });
 
         if cancelled {
             self.set_editor_is_expanded(false, cx);
@@ -370,8 +379,38 @@ impl MessageEditor {
         }
     }
 
-    fn handle_review_click(&self, window: &mut Window, cx: &mut Context<Self>) {
+    fn paste(&mut self, _: &Paste, _: &mut Window, cx: &mut Context<Self>) {
+        let images = cx
+            .read_from_clipboard()
+            .map(|item| {
+                item.into_entries()
+                    .filter_map(|entry| {
+                        if let ClipboardEntry::Image(image) = entry {
+                            Some(image)
+                        } else {
+                            None
+                        }
+                    })
+                    .collect::<Vec<_>>()
+            })
+            .unwrap_or_default();
+
+        if images.is_empty() {
+            return;
+        }
+        cx.stop_propagation();
+
+        self.context_store.update(cx, |store, cx| {
+            for image in images {
+                store.add_image(Arc::new(image), cx);
+            }
+        });
+    }
+
+    fn handle_review_click(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        self.edits_expanded = true;
         AgentDiff::deploy(self.thread.clone(), self.workspace.clone(), window, cx).log_err();
+        cx.notify();
     }
 
     fn handle_file_click(
@@ -445,6 +484,7 @@ impl MessageEditor {
             .on_action(cx.listener(Self::move_up))
             .on_action(cx.listener(Self::toggle_chat_mode))
             .on_action(cx.listener(Self::expand_message_editor))
+            .capture_action(cx.listener(Self::paste))
             .gap_2()
             .p_2()
             .bg(editor_bg_color)

crates/agent/src/thread.rs

@@ -13,10 +13,12 @@ use feature_flags::{self, FeatureFlagAppExt};
 use futures::future::Shared;
 use futures::{FutureExt, StreamExt as _};
 use git::repository::DiffType;
-use gpui::{App, AppContext, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
+use gpui::{
+    AnyWindowHandle, App, AppContext, Context, Entity, EventEmitter, SharedString, Task, WeakEntity,
+};
 use language_model::{
     ConfiguredModel, LanguageModel, LanguageModelCompletionEvent, LanguageModelId,
-    LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest,
+    LanguageModelImage, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest,
     LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
     LanguageModelToolUseId, MaxMonthlySpendReachedError, MessageContent,
     ModelRequestLimitReachedError, PaymentRequiredError, RequestUsage, Role, StopReason,
@@ -38,7 +40,7 @@ use crate::thread_store::{
     SerializedMessage, SerializedMessageSegment, SerializedThread, SerializedToolResult,
     SerializedToolUse, SharedProjectContext,
 };
-use crate::tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState, USING_TOOL_MARKER};
+use crate::tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState};
 
 #[derive(
     Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
@@ -97,6 +99,7 @@ pub struct Message {
     pub role: Role,
     pub segments: Vec<MessageSegment>,
     pub context: String,
+    pub images: Vec<LanguageModelImage>,
 }
 
 impl Message {
@@ -167,12 +170,9 @@ pub enum MessageSegment {
 
 impl MessageSegment {
     pub fn should_display(&self) -> bool {
-        // We add USING_TOOL_MARKER when making a request that includes tool uses
-        // without non-whitespace text around them, and this can cause the model
-        // to mimic the pattern, so we consider those segments not displayable.
         match self {
-            Self::Text(text) => text.is_empty() || text.trim() == USING_TOOL_MARKER,
-            Self::Thinking { text, .. } => text.is_empty() || text.trim() == USING_TOOL_MARKER,
+            Self::Text(text) => text.is_empty(),
+            Self::Thinking { text, .. } => text.is_empty(),
             Self::RedactedThinking(_) => false,
         }
     }
@@ -317,6 +317,7 @@ pub struct Thread {
     request_callback: Option<
         Box<dyn FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>])>,
     >,
+    remaining_turns: u32,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -370,6 +371,7 @@ impl Thread {
             message_feedback: HashMap::default(),
             last_auto_capture_at: None,
             request_callback: None,
+            remaining_turns: u32::MAX,
         }
     }
 
@@ -418,6 +420,7 @@ impl Thread {
                         })
                         .collect(),
                     context: message.context,
+                    images: Vec::new(),
                 })
                 .collect(),
             next_message_id,
@@ -443,6 +446,7 @@ impl Thread {
             message_feedback: HashMap::default(),
             last_auto_capture_at: None,
             request_callback: None,
+            remaining_turns: u32::MAX,
         }
     }
 
@@ -523,7 +527,7 @@ impl Thread {
         self.messages.iter().find(|message| message.id == id)
     }
 
-    pub fn messages(&self) -> impl Iterator<Item = &Message> {
+    pub fn messages(&self) -> impl ExactSizeIterator<Item = &Message> {
         self.messages.iter()
     }
 
@@ -617,24 +621,12 @@ impl Thread {
                     .await
                     .unwrap_or(false);
 
-                if equal {
-                    git_store
-                        .update(cx, |store, cx| {
-                            store.delete_checkpoint(pending_checkpoint.git_checkpoint, cx)
-                        })?
-                        .detach();
-                } else {
+                if !equal {
                     this.update(cx, |this, cx| {
                         this.insert_checkpoint(pending_checkpoint, cx)
                     })?;
                 }
 
-                git_store
-                    .update(cx, |store, cx| {
-                        store.delete_checkpoint(final_checkpoint, cx)
-                    })?
-                    .detach();
-
                 Ok(())
             }
             Err(_) => this.update(cx, |this, cx| {
@@ -750,31 +742,41 @@ impl Thread {
                 }
             }
 
+            if let Some(message) = self.messages.iter_mut().find(|m| m.id == message_id) {
+                message.images = new_context
+                    .iter()
+                    .filter_map(|context| {
+                        if let AssistantContext::Image(image_context) = context {
+                            image_context.image_task.clone().now_or_never().flatten()
+                        } else {
+                            None
+                        }
+                    })
+                    .collect::<Vec<_>>();
+            }
+
             self.action_log.update(cx, |log, cx| {
                 // Track all buffers added as context
                 for ctx in &new_context {
                     match ctx {
                         AssistantContext::File(file_ctx) => {
-                            log.buffer_added_as_context(file_ctx.context_buffer.buffer.clone(), cx);
+                            log.track_buffer(file_ctx.context_buffer.buffer.clone(), cx);
                         }
                         AssistantContext::Directory(dir_ctx) => {
                             for context_buffer in &dir_ctx.context_buffers {
-                                log.buffer_added_as_context(context_buffer.buffer.clone(), cx);
+                                log.track_buffer(context_buffer.buffer.clone(), cx);
                             }
                         }
                         AssistantContext::Symbol(symbol_ctx) => {
-                            log.buffer_added_as_context(
-                                symbol_ctx.context_symbol.buffer.clone(),
-                                cx,
-                            );
+                            log.track_buffer(symbol_ctx.context_symbol.buffer.clone(), cx);
                         }
-                        AssistantContext::Excerpt(excerpt_context) => {
-                            log.buffer_added_as_context(
-                                excerpt_context.context_buffer.buffer.clone(),
-                                cx,
-                            );
+                        AssistantContext::Selection(selection_context) => {
+                            log.track_buffer(selection_context.context_buffer.buffer.clone(), cx);
                         }
-                        AssistantContext::FetchedUrl(_) | AssistantContext::Thread(_) => {}
+                        AssistantContext::FetchedUrl(_)
+                        | AssistantContext::Thread(_)
+                        | AssistantContext::Rules(_)
+                        | AssistantContext::Image(_) => {}
                     }
                 }
             });
@@ -815,6 +817,7 @@ impl Thread {
             role,
             segments,
             context: String::new(),
+            images: Vec::new(),
         });
         self.touch_updated_at();
         cx.emit(ThreadEvent::MessageAdded(id));
@@ -942,7 +945,26 @@ impl Thread {
         })
     }
 
-    pub fn send_to_model(&mut self, model: Arc<dyn LanguageModel>, cx: &mut Context<Self>) {
+    pub fn remaining_turns(&self) -> u32 {
+        self.remaining_turns
+    }
+
+    pub fn set_remaining_turns(&mut self, remaining_turns: u32) {
+        self.remaining_turns = remaining_turns;
+    }
+
+    pub fn send_to_model(
+        &mut self,
+        model: Arc<dyn LanguageModel>,
+        window: Option<AnyWindowHandle>,
+        cx: &mut Context<Self>,
+    ) {
+        if self.remaining_turns == 0 {
+            return;
+        }
+
+        self.remaining_turns -= 1;
+
         let mut request = self.to_completion_request(cx);
         if model.supports_tools() {
             request.tools = {
@@ -967,7 +989,7 @@ impl Thread {
             };
         }
 
-        self.stream_completion(request, model, cx);
+        self.stream_completion(request, model, window, cx);
     }
 
     pub fn used_tools_since_last_user_message(&self) -> bool {
@@ -1038,6 +1060,21 @@ impl Thread {
                     .push(MessageContent::Text(message.context.to_string()));
             }
 
+            if !message.images.is_empty() {
+                // Some providers only support image parts after an initial text part
+                if request_message.content.is_empty() {
+                    request_message
+                        .content
+                        .push(MessageContent::Text("Images attached by user:".to_string()));
+                }
+
+                for image in &message.images {
+                    request_message
+                        .content
+                        .push(MessageContent::Image(image.clone()))
+                }
+            }
+
             for segment in &message.segments {
                 match segment {
                     MessageSegment::Text(text) => {
@@ -1171,6 +1208,7 @@ impl Thread {
         &mut self,
         request: LanguageModelRequest,
         model: Arc<dyn LanguageModel>,
+        window: Option<AnyWindowHandle>,
         cx: &mut Context<Self>,
     ) {
         let pending_completion_id = post_inc(&mut self.completion_count);
@@ -1232,6 +1270,7 @@ impl Thread {
                                 current_token_usage = token_usage;
                             }
                             LanguageModelCompletionEvent::Text(chunk) => {
+                                cx.emit(ThreadEvent::ReceivedTextChunk);
                                 if let Some(last_message) = thread.messages.last_mut() {
                                     if last_message.role == Role::Assistant {
                                         last_message.push_text(&chunk);
@@ -1291,12 +1330,27 @@ impl Thread {
                                         thread.insert_message(Role::Assistant, vec![], cx)
                                     });
 
-                                thread.tool_use.request_tool_use(
+                                let tool_use_id = tool_use.id.clone();
+                                let streamed_input = if tool_use.is_input_complete {
+                                    None
+                                } else {
+                                    Some((&tool_use.input).clone())
+                                };
+
+                                let ui_text = thread.tool_use.request_tool_use(
                                     last_assistant_message_id,
                                     tool_use,
                                     tool_use_metadata.clone(),
                                     cx,
                                 );
+
+                                if let Some(input) = streamed_input {
+                                    cx.emit(ThreadEvent::StreamedToolUse {
+                                        tool_use_id,
+                                        ui_text,
+                                        input,
+                                    });
+                                }
                             }
                         }
 
@@ -1336,7 +1390,7 @@ impl Thread {
                     match result.as_ref() {
                         Ok(stop_reason) => match stop_reason {
                             StopReason::ToolUse => {
-                                let tool_uses = thread.use_pending_tools(cx);
+                                let tool_uses = thread.use_pending_tools(window, cx);
                                 cx.emit(ThreadEvent::UsePendingTools { tool_uses });
                             }
                             StopReason::EndTurn => {}
@@ -1381,7 +1435,7 @@ impl Thread {
                                 }));
                             }
 
-                            thread.cancel_last_completion(cx);
+                            thread.cancel_last_completion(window, cx);
                         }
                     }
                     cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new)));
@@ -1550,7 +1604,11 @@ impl Thread {
         )
     }
 
-    pub fn use_pending_tools(&mut self, cx: &mut Context<Self>) -> Vec<PendingToolUse> {
+    pub fn use_pending_tools(
+        &mut self,
+        window: Option<AnyWindowHandle>,
+        cx: &mut Context<Self>,
+    ) -> Vec<PendingToolUse> {
         self.auto_capture_telemetry(cx);
         let request = self.to_completion_request(cx);
         let messages = Arc::new(request.messages);
@@ -1582,6 +1640,7 @@ impl Thread {
                         tool_use.input.clone(),
                         &messages,
                         tool,
+                        window,
                         cx,
                     );
                 }
@@ -1598,9 +1657,10 @@ impl Thread {
         input: serde_json::Value,
         messages: &[LanguageModelRequestMessage],
         tool: Arc<dyn Tool>,
+        window: Option<AnyWindowHandle>,
         cx: &mut Context<Thread>,
     ) {
-        let task = self.spawn_tool_use(tool_use_id.clone(), messages, input, tool, cx);
+        let task = self.spawn_tool_use(tool_use_id.clone(), messages, input, tool, window, cx);
         self.tool_use
             .run_pending_tool(tool_use_id, ui_text.into(), task);
     }
@@ -1611,6 +1671,7 @@ impl Thread {
         messages: &[LanguageModelRequestMessage],
         input: serde_json::Value,
         tool: Arc<dyn Tool>,
+        window: Option<AnyWindowHandle>,
         cx: &mut Context<Thread>,
     ) -> Task<()> {
         let tool_name: Arc<str> = tool.name().into();
@@ -1623,6 +1684,7 @@ impl Thread {
                 messages,
                 self.project.clone(),
                 self.action_log.clone(),
+                window,
                 cx,
             )
         };
@@ -1645,7 +1707,7 @@ impl Thread {
                             output,
                             cx,
                         );
-                        thread.tool_finished(tool_use_id, pending_tool_use, false, cx);
+                        thread.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
                     })
                     .ok();
             }
@@ -1657,6 +1719,7 @@ impl Thread {
         tool_use_id: LanguageModelToolUseId,
         pending_tool_use: Option<PendingToolUse>,
         canceled: bool,
+        window: Option<AnyWindowHandle>,
         cx: &mut Context<Self>,
     ) {
         if self.all_tools_finished() {
@@ -1664,7 +1727,7 @@ impl Thread {
             if let Some(ConfiguredModel { model, .. }) = model_registry.default_model() {
                 self.attach_tool_results(cx);
                 if !canceled {
-                    self.send_to_model(model, cx);
+                    self.send_to_model(model, window, cx);
                 }
             }
         }
@@ -1686,7 +1749,11 @@ impl Thread {
     /// Cancels the last pending completion, if there are any pending.
     ///
     /// Returns whether a completion was canceled.
-    pub fn cancel_last_completion(&mut self, cx: &mut Context<Self>) -> bool {
+    pub fn cancel_last_completion(
+        &mut self,
+        window: Option<AnyWindowHandle>,
+        cx: &mut Context<Self>,
+    ) -> bool {
         let canceled = if self.pending_completions.pop().is_some() {
             true
         } else {
@@ -1697,6 +1764,7 @@ impl Thread {
                     pending_tool_use.id.clone(),
                     Some(pending_tool_use),
                     true,
+                    window,
                     cx,
                 );
             }
@@ -1766,7 +1834,7 @@ impl Thread {
                 thread_data,
                 final_project_snapshot
             );
-            client.telemetry().flush_events();
+            client.telemetry().flush_events().await;
 
             Ok(())
         })
@@ -1811,7 +1879,7 @@ impl Thread {
                     thread_data,
                     final_project_snapshot
                 );
-                client.telemetry().flush_events();
+                client.telemetry().flush_events().await;
 
                 Ok(())
             })
@@ -2067,7 +2135,7 @@ impl Thread {
                             github_login = github_login
                         );
 
-                        client.telemetry().flush_events();
+                        client.telemetry().flush_events().await;
                     }
                 }
             })
@@ -2153,6 +2221,7 @@ impl Thread {
         &mut self,
         tool_use_id: LanguageModelToolUseId,
         tool_name: Arc<str>,
+        window: Option<AnyWindowHandle>,
         cx: &mut Context<Self>,
     ) {
         let err = Err(anyhow::anyhow!(
@@ -2161,7 +2230,7 @@ impl Thread {
 
         self.tool_use
             .insert_tool_output(tool_use_id.clone(), tool_name, err, cx);
-        self.tool_finished(tool_use_id.clone(), None, true, cx);
+        self.tool_finished(tool_use_id.clone(), None, true, window, cx);
     }
 }
 
@@ -2185,8 +2254,14 @@ pub enum ThreadEvent {
     ShowError(ThreadError),
     UsageUpdated(RequestUsage),
     StreamedCompletion,
+    ReceivedTextChunk,
     StreamedAssistantText(MessageId, String),
     StreamedAssistantThinking(MessageId, String),
+    StreamedToolUse {
+        tool_use_id: LanguageModelToolUseId,
+        ui_text: Arc<str>,
+        input: serde_json::Value,
+    },
     Stopped(Result<StopReason, Arc<anyhow::Error>>),
     MessageAdded(MessageId),
     MessageEdited(MessageId),

crates/agent/src/thread_store.rs

@@ -24,8 +24,8 @@ use heed::types::SerdeBincode;
 use language_model::{LanguageModelToolUseId, Role, TokenUsage};
 use project::{Project, Worktree};
 use prompt_store::{
-    DefaultUserRulesContext, ProjectContext, PromptBuilder, PromptId, PromptStore,
-    PromptsUpdatedEvent, RulesFileContext, WorktreeContext,
+    ProjectContext, PromptBuilder, PromptId, PromptMetadata, PromptStore, PromptsUpdatedEvent,
+    RulesFileContext, UserPromptId, UserRulesContext, WorktreeContext,
 };
 use serde::{Deserialize, Serialize};
 use settings::{Settings as _, SettingsStore};
@@ -62,6 +62,7 @@ pub struct ThreadStore {
     project: Entity<Project>,
     tools: Entity<ToolWorkingSet>,
     prompt_builder: Arc<PromptBuilder>,
+    prompt_store: Option<Entity<PromptStore>>,
     context_server_manager: Entity<ContextServerManager>,
     context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,
     threads: Vec<SerializedThreadMetadata>,
@@ -135,6 +136,7 @@ impl ThreadStore {
         let (ready_tx, ready_rx) = oneshot::channel();
         let mut ready_tx = Some(ready_tx);
         let reload_system_prompt_task = cx.spawn({
+            let prompt_store = prompt_store.clone();
             async move |thread_store, cx| {
                 loop {
                     let Some(reload_task) = thread_store
@@ -158,6 +160,7 @@ impl ThreadStore {
             project,
             tools,
             prompt_builder,
+            prompt_store,
             context_server_manager,
             context_server_tool_ids: HashMap::default(),
             threads: Vec::new(),
@@ -245,7 +248,7 @@ impl ThreadStore {
             let default_user_rules = default_user_rules
                 .into_iter()
                 .flat_map(|(contents, prompt_metadata)| match contents {
-                    Ok(contents) => Some(DefaultUserRulesContext {
+                    Ok(contents) => Some(UserRulesContext {
                         uuid: match prompt_metadata.id {
                             PromptId::User { uuid } => uuid,
                             PromptId::EditWorkflow => return None,
@@ -346,6 +349,27 @@ impl ThreadStore {
         self.context_server_manager.clone()
     }
 
+    pub fn prompt_store(&self) -> Option<Entity<PromptStore>> {
+        self.prompt_store.clone()
+    }
+
+    pub fn load_rules(
+        &self,
+        prompt_id: UserPromptId,
+        cx: &App,
+    ) -> Task<Result<(PromptMetadata, String)>> {
+        let prompt_id = PromptId::User { uuid: prompt_id };
+        let Some(prompt_store) = self.prompt_store.as_ref() else {
+            return Task::ready(Err(anyhow!("Prompt store unexpectedly missing.")));
+        };
+        let prompt_store = prompt_store.read(cx);
+        let Some(metadata) = prompt_store.metadata(prompt_id) else {
+            return Task::ready(Err(anyhow!("User rules not found in library.")));
+        };
+        let text_task = prompt_store.load(prompt_id, cx);
+        cx.background_spawn(async move { Ok((metadata, text_task.await?)) })
+    }
+
     pub fn tools(&self) -> Entity<ToolWorkingSet> {
         self.tools.clone()
     }

crates/agent/src/tool_use.rs

@@ -27,8 +27,6 @@ pub struct ToolUse {
     pub needs_confirmation: bool,
 }
 
-pub const USING_TOOL_MARKER: &str = "<using_tool>";
-
 pub struct ToolUseState {
     tools: Entity<ToolWorkingSet>,
     tool_uses_by_assistant_message: HashMap<MessageId, Vec<LanguageModelToolUse>>,
@@ -75,6 +73,7 @@ impl ToolUseState {
                                 id: tool_use.id.clone(),
                                 name: tool_use.name.clone().into(),
                                 input: tool_use.input.clone(),
+                                is_input_complete: true,
                             })
                             .collect::<Vec<_>>();
 
@@ -176,6 +175,9 @@ impl ToolUseState {
                         PendingToolUseStatus::Error(ref err) => {
                             ToolUseStatus::Error(err.clone().into())
                         }
+                        PendingToolUseStatus::InputStillStreaming => {
+                            ToolUseStatus::InputStillStreaming
+                        }
                     }
                 } else {
                     ToolUseStatus::Pending
@@ -192,7 +194,12 @@ impl ToolUseState {
             tool_uses.push(ToolUse {
                 id: tool_use.id.clone(),
                 name: tool_use.name.clone().into(),
-                ui_text: self.tool_ui_label(&tool_use.name, &tool_use.input, cx),
+                ui_text: self.tool_ui_label(
+                    &tool_use.name,
+                    &tool_use.input,
+                    tool_use.is_input_complete,
+                    cx,
+                ),
                 input: tool_use.input.clone(),
                 status,
                 icon,
@@ -207,10 +214,15 @@ impl ToolUseState {
         &self,
         tool_name: &str,
         input: &serde_json::Value,
+        is_input_complete: bool,
         cx: &App,
     ) -> SharedString {
         if let Some(tool) = self.tools.read(cx).tool(tool_name, cx) {
-            tool.ui_text(input).into()
+            if is_input_complete {
+                tool.ui_text(input).into()
+            } else {
+                tool.still_streaming_ui_text(input).into()
+            }
         } else {
             format!("Unknown tool {tool_name:?}").into()
         }
@@ -258,22 +270,50 @@ impl ToolUseState {
         tool_use: LanguageModelToolUse,
         metadata: ToolUseMetadata,
         cx: &App,
-    ) {
-        self.tool_uses_by_assistant_message
+    ) -> Arc<str> {
+        let tool_uses = self
+            .tool_uses_by_assistant_message
             .entry(assistant_message_id)
-            .or_default()
-            .push(tool_use.clone());
+            .or_default();
 
-        self.tool_use_metadata_by_id
-            .insert(tool_use.id.clone(), metadata);
+        let mut existing_tool_use_found = false;
 
-        // The tool use is being requested by the Assistant, so we want to
-        // attach the tool results to the next user message.
-        let next_user_message_id = MessageId(assistant_message_id.0 + 1);
-        self.tool_uses_by_user_message
-            .entry(next_user_message_id)
-            .or_default()
-            .push(tool_use.id.clone());
+        for existing_tool_use in tool_uses.iter_mut() {
+            if existing_tool_use.id == tool_use.id {
+                *existing_tool_use = tool_use.clone();
+                existing_tool_use_found = true;
+            }
+        }
+
+        if !existing_tool_use_found {
+            tool_uses.push(tool_use.clone());
+        }
+
+        let status = if tool_use.is_input_complete {
+            self.tool_use_metadata_by_id
+                .insert(tool_use.id.clone(), metadata);
+
+            // The tool use is being requested by the Assistant, so we want to
+            // attach the tool results to the next user message.
+            let next_user_message_id = MessageId(assistant_message_id.0 + 1);
+            self.tool_uses_by_user_message
+                .entry(next_user_message_id)
+                .or_default()
+                .push(tool_use.id.clone());
+
+            PendingToolUseStatus::Idle
+        } else {
+            PendingToolUseStatus::InputStillStreaming
+        };
+
+        let ui_text: Arc<str> = self
+            .tool_ui_label(
+                &tool_use.name,
+                &tool_use.input,
+                tool_use.is_input_complete,
+                cx,
+            )
+            .into();
 
         self.pending_tool_uses_by_id.insert(
             tool_use.id.clone(),
@@ -281,13 +321,13 @@ impl ToolUseState {
                 assistant_message_id,
                 id: tool_use.id,
                 name: tool_use.name.clone(),
-                ui_text: self
-                    .tool_ui_label(&tool_use.name, &tool_use.input, cx)
-                    .into(),
+                ui_text: ui_text.clone(),
                 input: tool_use.input,
-                status: PendingToolUseStatus::Idle,
+                status,
             },
         );
+
+        ui_text
     }
 
     pub fn run_pending_tool(
@@ -410,28 +450,8 @@ impl ToolUseState {
         request_message: &mut LanguageModelRequestMessage,
     ) {
         if let Some(tool_uses) = self.tool_uses_by_assistant_message.get(&message_id) {
-            let mut found_tool_use = false;
-
             for tool_use in tool_uses {
                 if self.tool_results.contains_key(&tool_use.id) {
-                    if !found_tool_use {
-                        // The API fails if a message contains a tool use without any (non-whitespace) text around it
-                        match request_message.content.last_mut() {
-                            Some(MessageContent::Text(txt)) => {
-                                if txt.is_empty() {
-                                    txt.push_str(USING_TOOL_MARKER);
-                                }
-                            }
-                            None | Some(_) => {
-                                request_message
-                                    .content
-                                    .push(MessageContent::Text(USING_TOOL_MARKER.into()));
-                            }
-                        };
-                    }
-
-                    found_tool_use = true;
-
                     // Do not send tool uses until they are completed
                     request_message
                         .content
@@ -497,6 +517,7 @@ pub struct Confirmation {
 
 #[derive(Debug, Clone)]
 pub enum PendingToolUseStatus {
+    InputStillStreaming,
     Idle,
     NeedsConfirmation(Arc<Confirmation>),
     Running { _task: Shared<Task<()>> },

crates/agent/src/ui.rs

@@ -1,7 +1,9 @@
 mod agent_notification;
+mod animated_label;
 mod context_pill;
 mod usage_banner;
 
 pub use agent_notification::*;
+pub use animated_label::*;
 pub use context_pill::*;
 pub use usage_banner::*;

crates/agent/src/ui/agent_notification.rs

@@ -12,6 +12,7 @@ pub struct AgentNotification {
     title: SharedString,
     caption: SharedString,
     icon: IconName,
+    project_name: Option<SharedString>,
 }
 
 impl AgentNotification {
@@ -19,11 +20,13 @@ impl AgentNotification {
         title: impl Into<SharedString>,
         caption: impl Into<SharedString>,
         icon: IconName,
+        project_name: Option<impl Into<SharedString>>,
     ) -> Self {
         Self {
             title: title.into(),
             caption: caption.into(),
             icon,
+            project_name: project_name.map(|name| name.into()),
         }
     }
 
@@ -130,11 +133,34 @@ impl Render for AgentNotification {
                                     .child(gradient_overflow()),
                             )
                             .child(
-                                div()
+                                h_flex()
                                     .relative()
+                                    .gap_1p5()
                                     .text_size(px(12.))
                                     .text_color(cx.theme().colors().text_muted)
                                     .truncate()
+                                    .when_some(
+                                        self.project_name.clone(),
+                                        |description, project_name| {
+                                            description.child(
+                                                h_flex()
+                                                    .gap_1p5()
+                                                    .child(
+                                                        div()
+                                                            .max_w_16()
+                                                            .truncate()
+                                                            .child(project_name),
+                                                    )
+                                                    .child(
+                                                        div().size(px(3.)).rounded_full().bg(cx
+                                                            .theme()
+                                                            .colors()
+                                                            .text
+                                                            .opacity(0.5)),
+                                                    ),
+                                            )
+                                        },
+                                    )
                                     .child(self.caption.clone())
                                     .child(gradient_overflow()),
                             ),

crates/agent/src/ui/animated_label.rs

@@ -0,0 +1,116 @@
+use gpui::{Animation, AnimationExt, FontWeight, pulsating_between};
+use std::time::Duration;
+use ui::prelude::*;
+
+#[derive(IntoElement)]
+pub struct AnimatedLabel {
+    base: Label,
+    text: SharedString,
+}
+
+impl AnimatedLabel {
+    pub fn new(text: impl Into<SharedString>) -> Self {
+        let text = text.into();
+        AnimatedLabel {
+            base: Label::new(text.clone()),
+            text,
+        }
+    }
+}
+
+impl LabelCommon for AnimatedLabel {
+    fn size(mut self, size: LabelSize) -> Self {
+        self.base = self.base.size(size);
+        self
+    }
+
+    fn weight(mut self, weight: FontWeight) -> Self {
+        self.base = self.base.weight(weight);
+        self
+    }
+
+    fn line_height_style(mut self, line_height_style: LineHeightStyle) -> Self {
+        self.base = self.base.line_height_style(line_height_style);
+        self
+    }
+
+    fn color(mut self, color: Color) -> Self {
+        self.base = self.base.color(color);
+        self
+    }
+
+    fn strikethrough(mut self) -> Self {
+        self.base = self.base.strikethrough();
+        self
+    }
+
+    fn italic(mut self) -> Self {
+        self.base = self.base.italic();
+        self
+    }
+
+    fn alpha(mut self, alpha: f32) -> Self {
+        self.base = self.base.alpha(alpha);
+        self
+    }
+
+    fn underline(mut self) -> Self {
+        self.base = self.base.underline();
+        self
+    }
+
+    fn truncate(mut self) -> Self {
+        self.base = self.base.truncate();
+        self
+    }
+
+    fn single_line(mut self) -> Self {
+        self.base = self.base.single_line();
+        self
+    }
+
+    fn buffer_font(mut self, cx: &App) -> Self {
+        self.base = self.base.buffer_font(cx);
+        self
+    }
+}
+
+impl RenderOnce for AnimatedLabel {
+    fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement {
+        let text = self.text.clone();
+
+        self.base
+            .color(Color::Muted)
+            .with_animations(
+                "animated-label",
+                vec![
+                    Animation::new(Duration::from_secs(1)),
+                    Animation::new(Duration::from_secs(1)).repeat(),
+                ],
+                move |mut label, animation_ix, delta| {
+                    match animation_ix {
+                        0 => {
+                            let chars_to_show = (delta * text.len() as f32).ceil() as usize;
+                            let text = SharedString::from(text[0..chars_to_show].to_string());
+                            label.set_text(text);
+                        }
+                        1 => match delta {
+                            d if d < 0.25 => label.set_text(text.clone()),
+                            d if d < 0.5 => label.set_text(format!("{}.", text)),
+                            d if d < 0.75 => label.set_text(format!("{}..", text)),
+                            _ => label.set_text(format!("{}...", text)),
+                        },
+                        _ => {}
+                    }
+                    label
+                },
+            )
+            .with_animation(
+                "pulsating-label",
+                Animation::new(Duration::from_secs(2))
+                    .repeat()
+                    .with_easing(pulsating_between(0.6, 1.)),
+                |label, delta| label.map_element(|label| label.alpha(delta)),
+            )
+    }
+}

crates/agent/src/ui/context_pill.rs

@@ -1,11 +1,14 @@
+use std::sync::Arc;
 use std::{rc::Rc, time::Duration};
 
 use file_icons::FileIcons;
-use gpui::ClickEvent;
-use gpui::{Animation, AnimationExt as _, pulsating_between};
-use ui::{IconButtonShape, Tooltip, prelude::*};
+use futures::FutureExt;
+use gpui::{Animation, AnimationExt as _, Image, MouseButton, pulsating_between};
+use gpui::{ClickEvent, Task};
+use language_model::LanguageModelImage;
+use ui::{IconButtonShape, Tooltip, prelude::*, tooltip_container};
 
-use crate::context::{AssistantContext, ContextId, ContextKind};
+use crate::context::{AssistantContext, ContextId, ContextKind, ImageContext};
 
 #[derive(IntoElement)]
 pub enum ContextPill {
@@ -120,74 +123,100 @@ impl RenderOnce for ContextPill {
                 on_remove,
                 focused,
                 on_click,
-            } => base_pill
-                .bg(color.element_background)
-                .border_color(if *focused {
-                    color.border_focused
-                } else {
-                    color.border.opacity(0.5)
-                })
-                .pr(if on_remove.is_some() { px(2.) } else { px(4.) })
-                .child(
-                    h_flex()
-                        .id("context-data")
-                        .gap_1()
-                        .child(
-                            div().max_w_64().child(
-                                Label::new(context.name.clone())
-                                    .size(LabelSize::Small)
-                                    .truncate(),
-                            ),
-                        )
-                        .when_some(context.parent.as_ref(), |element, parent_name| {
-                            if *dupe_name {
-                                element.child(
-                                    Label::new(parent_name.clone())
-                                        .size(LabelSize::XSmall)
-                                        .color(Color::Muted),
-                                )
-                            } else {
-                                element
-                            }
-                        })
-                        .when_some(context.tooltip.as_ref(), |element, tooltip| {
-                            element.tooltip(Tooltip::text(tooltip.clone()))
-                        }),
-                )
-                .when_some(on_remove.as_ref(), |element, on_remove| {
-                    element.child(
-                        IconButton::new(("remove", context.id.0), IconName::Close)
-                            .shape(IconButtonShape::Square)
-                            .icon_size(IconSize::XSmall)
-                            .tooltip(Tooltip::text("Remove Context"))
-                            .on_click({
-                                let on_remove = on_remove.clone();
-                                move |event, window, cx| on_remove(event, window, cx)
+            } => {
+                let status_is_error = matches!(context.status, ContextStatus::Error { .. });
+
+                base_pill
+                    .pr(if on_remove.is_some() { px(2.) } else { px(4.) })
+                    .map(|pill| {
+                        if status_is_error {
+                            pill.bg(cx.theme().status().error_background)
+                                .border_color(cx.theme().status().error_border)
+                        } else if *focused {
+                            pill.bg(color.element_background)
+                                .border_color(color.border_focused)
+                        } else {
+                            pill.bg(color.element_background)
+                                .border_color(color.border.opacity(0.5))
+                        }
+                    })
+                    .child(
+                        h_flex()
+                            .id("context-data")
+                            .gap_1()
+                            .child(
+                                div().max_w_64().child(
+                                    Label::new(context.name.clone())
+                                        .size(LabelSize::Small)
+                                        .truncate(),
+                                ),
+                            )
+                            .when_some(context.parent.as_ref(), |element, parent_name| {
+                                if *dupe_name {
+                                    element.child(
+                                        Label::new(parent_name.clone())
+                                            .size(LabelSize::XSmall)
+                                            .color(Color::Muted),
+                                    )
+                                } else {
+                                    element
+                                }
+                            })
+                            .when_some(context.tooltip.as_ref(), |element, tooltip| {
+                                element.tooltip(Tooltip::text(tooltip.clone()))
+                            })
+                            .map(|element| match &context.status {
+                                ContextStatus::Ready => element
+                                    .when_some(
+                                        context.render_preview.as_ref(),
+                                        |element, render_preview| {
+                                            element.hoverable_tooltip({
+                                                let render_preview = render_preview.clone();
+                                                move |_, cx| {
+                                                    cx.new(|_| ContextPillPreview {
+                                                        render_preview: render_preview.clone(),
+                                                    })
+                                                    .into()
+                                                }
+                                            })
+                                        },
+                                    )
+                                    .into_any(),
+                                ContextStatus::Loading { message } => element
+                                    .tooltip(ui::Tooltip::text(message.clone()))
+                                    .with_animation(
+                                        "pulsating-ctx-pill",
+                                        Animation::new(Duration::from_secs(2))
+                                            .repeat()
+                                            .with_easing(pulsating_between(0.4, 0.8)),
+                                        |label, delta| label.opacity(delta),
+                                    )
+                                    .into_any_element(),
+                                ContextStatus::Error { message } => element
+                                    .tooltip(ui::Tooltip::text(message.clone()))
+                                    .into_any_element(),
                             }),
                     )
-                })
-                .when_some(on_click.as_ref(), |element, on_click| {
-                    let on_click = on_click.clone();
-                    element
-                        .cursor_pointer()
-                        .on_click(move |event, window, cx| on_click(event, window, cx))
-                })
-                .map(|element| {
-                    if context.summarizing {
+                    .when_some(on_remove.as_ref(), |element, on_remove| {
+                        element.child(
+                            IconButton::new(("remove", context.id.0), IconName::Close)
+                                .shape(IconButtonShape::Square)
+                                .icon_size(IconSize::XSmall)
+                                .tooltip(Tooltip::text("Remove Context"))
+                                .on_click({
+                                    let on_remove = on_remove.clone();
+                                    move |event, window, cx| on_remove(event, window, cx)
+                                }),
+                        )
+                    })
+                    .when_some(on_click.as_ref(), |element, on_click| {
+                        let on_click = on_click.clone();
                         element
-                            .tooltip(ui::Tooltip::text("Summarizing..."))
-                            .with_animation(
-                                "pulsating-ctx-pill",
-                                Animation::new(Duration::from_secs(2))
-                                    .repeat()
-                                    .with_easing(pulsating_between(0.4, 0.8)),
-                                |label, delta| label.opacity(delta),
-                            )
-                            .into_any_element()
-                    } else {
-                        element.into_any()
-                    }
-                }),
+                            .cursor_pointer()
+                            .on_click(move |event, window, cx| on_click(event, window, cx))
+                    })
+                    .into_any_element()
+            }
             ContextPill::Suggested {
                 name,
                 icon_path: _,
@@ -198,15 +227,15 @@ impl RenderOnce for ContextPill {
                 .cursor_pointer()
                 .pr_1()
                 .border_dashed()
-                .border_color(if *focused {
-                    color.border_focused
-                } else {
-                    color.border
+                .map(|pill| {
+                    if *focused {
+                        pill.border_color(color.border_focused)
+                            .bg(color.element_background.opacity(0.5))
+                    } else {
+                        pill.border_color(color.border)
+                    }
                 })
                 .hover(|style| style.bg(color.element_hover.opacity(0.5)))
-                .when(*focused, |this| {
-                    this.bg(color.element_background.opacity(0.5))
-                })
                 .child(
                     div().max_w_64().child(
                         Label::new(name.clone())
@@ -227,6 +256,13 @@ impl RenderOnce for ContextPill {
     }
 }
 
+pub enum ContextStatus {
+    Ready,
+    Loading { message: SharedString },
+    Error { message: SharedString },
+}
+
+#[derive(RegisterComponent)]
 pub struct AddedContext {
     pub id: ContextId,
     pub kind: ContextKind,
@@ -234,14 +270,15 @@ pub struct AddedContext {
     pub parent: Option<SharedString>,
     pub tooltip: Option<SharedString>,
     pub icon_path: Option<SharedString>,
-    pub summarizing: bool,
+    pub status: ContextStatus,
+    pub render_preview: Option<Rc<dyn Fn(&mut Window, &mut App) -> AnyElement + 'static>>,
 }
 
 impl AddedContext {
     pub fn new(context: &AssistantContext, cx: &App) -> AddedContext {
         match context {
             AssistantContext::File(file_context) => {
-                let full_path = file_context.context_buffer.file.full_path(cx);
+                let full_path = file_context.context_buffer.full_path(cx);
                 let full_path_string: SharedString =
                     full_path.to_string_lossy().into_owned().into();
                 let name = full_path
@@ -259,15 +296,20 @@ impl AddedContext {
                     parent,
                     tooltip: Some(full_path_string),
                     icon_path: FileIcons::get_icon(&full_path, cx),
-                    summarizing: false,
+                    status: ContextStatus::Ready,
+                    render_preview: None,
                 }
             }
 
             AssistantContext::Directory(directory_context) => {
-                let full_path = directory_context
-                    .worktree
-                    .read(cx)
-                    .full_path(&directory_context.path);
+                let worktree = directory_context.worktree.read(cx);
+                // If the directory no longer exists, use its last known path.
+                let full_path = worktree
+                    .entry_for_id(directory_context.entry_id)
+                    .map_or_else(
+                        || directory_context.last_path.clone(),
+                        |entry| worktree.full_path(&entry.path).into(),
+                    );
                 let full_path_string: SharedString =
                     full_path.to_string_lossy().into_owned().into();
                 let name = full_path
@@ -285,7 +327,8 @@ impl AddedContext {
                     parent,
                     tooltip: Some(full_path_string),
                     icon_path: None,
-                    summarizing: false,
+                    status: ContextStatus::Ready,
+                    render_preview: None,
                 }
             }
 
@@ -296,11 +339,12 @@ impl AddedContext {
                 parent: None,
                 tooltip: None,
                 icon_path: None,
-                summarizing: false,
+                status: ContextStatus::Ready,
+                render_preview: None,
             },
 
-            AssistantContext::Excerpt(excerpt_context) => {
-                let full_path = excerpt_context.context_buffer.file.full_path(cx);
+            AssistantContext::Selection(selection_context) => {
+                let full_path = selection_context.context_buffer.full_path(cx);
                 let mut full_path_string = full_path.to_string_lossy().into_owned();
                 let mut name = full_path
                     .file_name()
@@ -309,8 +353,8 @@ impl AddedContext {
 
                 let line_range_text = format!(
                     " ({}-{})",
-                    excerpt_context.line_range.start.row + 1,
-                    excerpt_context.line_range.end.row + 1
+                    selection_context.line_range.start.row + 1,
+                    selection_context.line_range.end.row + 1
                 );
 
                 full_path_string.push_str(&line_range_text);
@@ -322,13 +366,25 @@ impl AddedContext {
                     .map(|n| n.to_string_lossy().into_owned().into());
 
                 AddedContext {
-                    id: excerpt_context.id,
-                    kind: ContextKind::File, // Use File icon for excerpts
+                    id: selection_context.id,
+                    kind: ContextKind::Selection,
                     name: name.into(),
                     parent,
-                    tooltip: Some(full_path_string.into()),
+                    tooltip: None,
                     icon_path: FileIcons::get_icon(&full_path, cx),
-                    summarizing: false,
+                    status: ContextStatus::Ready,
+                    render_preview: Some(Rc::new({
+                        let content = selection_context.context_buffer.text.clone();
+                        move |_, cx| {
+                            div()
+                                .id("context-pill-selection-preview")
+                                .overflow_scroll()
+                                .max_w_128()
+                                .max_h_96()
+                                .child(Label::new(content.clone()).buffer_font(cx))
+                                .into_any_element()
+                        }
+                    })),
                 }
             }
 
@@ -339,7 +395,8 @@ impl AddedContext {
                 parent: None,
                 tooltip: None,
                 icon_path: None,
-                summarizing: false,
+                status: ContextStatus::Ready,
+                render_preview: None,
             },
 
             AssistantContext::Thread(thread_context) => AddedContext {
@@ -349,11 +406,143 @@ impl AddedContext {
                 parent: None,
                 tooltip: None,
                 icon_path: None,
-                summarizing: thread_context
+                status: if thread_context
                     .thread
                     .read(cx)
-                    .is_generating_detailed_summary(),
+                    .is_generating_detailed_summary()
+                {
+                    ContextStatus::Loading {
+                        message: "Summarizing…".into(),
+                    }
+                } else {
+                    ContextStatus::Ready
+                },
+                render_preview: None,
+            },
+
+            AssistantContext::Rules(user_rules_context) => AddedContext {
+                id: user_rules_context.id,
+                kind: ContextKind::Rules,
+                name: user_rules_context.title.clone(),
+                parent: None,
+                tooltip: None,
+                icon_path: None,
+                status: ContextStatus::Ready,
+                render_preview: None,
+            },
+
+            AssistantContext::Image(image_context) => AddedContext {
+                id: image_context.id,
+                kind: ContextKind::Image,
+                name: "Image".into(),
+                parent: None,
+                tooltip: None,
+                icon_path: None,
+                status: if image_context.is_loading() {
+                    ContextStatus::Loading {
+                        message: "Loading…".into(),
+                    }
+                } else if image_context.is_error() {
+                    ContextStatus::Error {
+                        message: "Failed to load image".into(),
+                    }
+                } else {
+                    ContextStatus::Ready
+                },
+                render_preview: Some(Rc::new({
+                    let image = image_context.original_image.clone();
+                    move |_, _| {
+                        gpui::img(image.clone())
+                            .max_w_96()
+                            .max_h_96()
+                            .into_any_element()
+                    }
+                })),
             },
         }
     }
 }
+
+struct ContextPillPreview {
+    render_preview: Rc<dyn Fn(&mut Window, &mut App) -> AnyElement>,
+}
+
+impl Render for ContextPillPreview {
+    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        tooltip_container(window, cx, move |this, window, cx| {
+            this.occlude()
+                .on_mouse_move(|_, _, cx| cx.stop_propagation())
+                .on_mouse_down(MouseButton::Left, |_, _, cx| cx.stop_propagation())
+                .child((self.render_preview)(window, cx))
+        })
+    }
+}
+
+impl Component for AddedContext {
+    fn scope() -> ComponentScope {
+        ComponentScope::Agent
+    }
+
+    fn sort_name() -> &'static str {
+        "AddedContext"
+    }
+
+    fn preview(_window: &mut Window, cx: &mut App) -> Option<AnyElement> {
+        let image_ready = (
+            "Ready",
+            AddedContext::new(
+                &AssistantContext::Image(ImageContext {
+                    id: ContextId(0),
+                    original_image: Arc::new(Image::empty()),
+                    image_task: Task::ready(Some(LanguageModelImage::empty())).shared(),
+                }),
+                cx,
+            ),
+        );
+
+        let image_loading = (
+            "Loading",
+            AddedContext::new(
+                &AssistantContext::Image(ImageContext {
+                    id: ContextId(1),
+                    original_image: Arc::new(Image::empty()),
+                    image_task: cx
+                        .background_spawn(async move {
+                            smol::Timer::after(Duration::from_secs(60 * 5)).await;
+                            Some(LanguageModelImage::empty())
+                        })
+                        .shared(),
+                }),
+                cx,
+            ),
+        );
+
+        let image_error = (
+            "Error",
+            AddedContext::new(
+                &AssistantContext::Image(ImageContext {
+                    id: ContextId(2),
+                    original_image: Arc::new(Image::empty()),
+                    image_task: Task::ready(None).shared(),
+                }),
+                cx,
+            ),
+        );
+
+        Some(
+            v_flex()
+                .gap_6()
+                .children(
+                    vec![image_ready, image_loading, image_error]
+                        .into_iter()
+                        .map(|(text, context)| {
+                            single_example(
+                                text,
+                                ContextPill::added(context, false, false, None).into_any_element(),
+                            )
+                        }),
+                )
+                .into_any(),
+        )
+    }
+}

crates/agent/src/ui/usage_banner.rs

@@ -1,31 +1,30 @@
 use client::zed_urls;
+use language_model::RequestUsage;
 use ui::{Banner, ProgressBar, Severity, prelude::*};
 use zed_llm_client::{Plan, UsageLimit};
 
 #[derive(IntoElement, RegisterComponent)]
 pub struct UsageBanner {
     plan: Plan,
-    requests: i32,
+    usage: RequestUsage,
 }
 
 impl UsageBanner {
-    pub fn new(plan: Plan, requests: i32) -> Self {
-        Self { plan, requests }
+    pub fn new(plan: Plan, usage: RequestUsage) -> Self {
+        Self { plan, usage }
     }
 }
 
 impl RenderOnce for UsageBanner {
     fn render(self, _window: &mut Window, cx: &mut App) -> impl IntoElement {
-        let request_limit = self.plan.model_requests_limit();
-
-        let used_percentage = match request_limit {
-            UsageLimit::Limited(limit) => Some((self.requests as f32 / limit as f32) * 100.),
+        let used_percentage = match self.usage.limit {
+            UsageLimit::Limited(limit) => Some((self.usage.amount as f32 / limit as f32) * 100.),
             UsageLimit::Unlimited => None,
         };
 
-        let (severity, message) = match request_limit {
+        let (severity, message) = match self.usage.limit {
             UsageLimit::Limited(limit) => {
-                if self.requests >= limit {
+                if self.usage.amount >= limit {
                     let message = match self.plan {
                         Plan::ZedPro => "Monthly request limit reached",
                         Plan::ZedProTrial => "Trial request limit reached",
@@ -33,7 +32,7 @@ impl RenderOnce for UsageBanner {
                     };
 
                     (Severity::Error, message)
-                } else if (self.requests as f32 / limit as f32) >= 0.9 {
+                } else if (self.usage.amount as f32 / limit as f32) >= 0.9 {
                     (Severity::Warning, "Approaching request limit")
                 } else {
                     let message = match self.plan {
@@ -81,11 +80,11 @@ impl RenderOnce for UsageBanner {
                             .child(ProgressBar::new("usage", percent, 100., cx))
                     }))
                     .child(
-                        Label::new(match request_limit {
+                        Label::new(match self.usage.limit {
                             UsageLimit::Limited(limit) => {
-                                format!("{} / {limit}", self.requests)
+                                format!("{} / {limit}", self.usage.amount)
                             }
-                            UsageLimit::Unlimited => format!("{} / ∞", self.requests),
+                            UsageLimit::Unlimited => format!("{} / ∞", self.usage.amount),
                         })
                         .size(LabelSize::Small)
                         .color(Color::Muted),
@@ -104,74 +103,131 @@ impl Component for UsageBanner {
     }
 
     fn preview(_window: &mut Window, _cx: &mut App) -> Option<AnyElement> {
+        let trial_limit = Plan::ZedProTrial.model_requests_limit();
         let trial_examples = vec![
             single_example(
                 "Zed Pro Trial - New User",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedProTrial, 10))
+                    .child(UsageBanner::new(
+                        Plan::ZedProTrial,
+                        RequestUsage {
+                            limit: trial_limit,
+                            amount: 10,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro Trial - Approaching Limit",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedProTrial, 135))
+                    .child(UsageBanner::new(
+                        Plan::ZedProTrial,
+                        RequestUsage {
+                            limit: trial_limit,
+                            amount: 135,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro Trial - Request Limit Reached",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedProTrial, 150))
+                    .child(UsageBanner::new(
+                        Plan::ZedProTrial,
+                        RequestUsage {
+                            limit: trial_limit,
+                            amount: 150,
+                        },
+                    ))
                     .into_any_element(),
             ),
         ];
 
+        let free_limit = Plan::Free.model_requests_limit();
         let free_examples = vec![
             single_example(
                 "Free - Normal Usage",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::Free, 25))
+                    .child(UsageBanner::new(
+                        Plan::Free,
+                        RequestUsage {
+                            limit: free_limit,
+                            amount: 25,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Free - Approaching Limit",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::Free, 45))
+                    .child(UsageBanner::new(
+                        Plan::Free,
+                        RequestUsage {
+                            limit: free_limit,
+                            amount: 45,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Free - Request Limit Reached",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::Free, 50))
+                    .child(UsageBanner::new(
+                        Plan::Free,
+                        RequestUsage {
+                            limit: free_limit,
+                            amount: 50,
+                        },
+                    ))
                     .into_any_element(),
             ),
         ];
 
+        let zed_pro_limit = Plan::ZedPro.model_requests_limit();
         let zed_pro_examples = vec![
             single_example(
                 "Zed Pro - Normal Usage",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedPro, 250))
+                    .child(UsageBanner::new(
+                        Plan::ZedPro,
+                        RequestUsage {
+                            limit: zed_pro_limit,
+                            amount: 250,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro - Approaching Limit",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedPro, 450))
+                    .child(UsageBanner::new(
+                        Plan::ZedPro,
+                        RequestUsage {
+                            limit: zed_pro_limit,
+                            amount: 450,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro - Request Limit Reached",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedPro, 500))
+                    .child(UsageBanner::new(
+                        Plan::ZedPro,
+                        RequestUsage {
+                            limit: zed_pro_limit,
+                            amount: 500,
+                        },
+                    ))
                     .into_any_element(),
             ),
         ];

crates/assistant/src/assistant_panel.rs

@@ -23,11 +23,10 @@ use gpui::{
 use language::LanguageRegistry;
 use language_model::{
     AuthenticateError, ConfiguredModel, LanguageModelProviderId, LanguageModelRegistry,
-    ZED_CLOUD_PROVIDER_ID,
 };
 use project::Project;
 use prompt_library::{PromptLibrary, open_prompt_library};
-use prompt_store::{PromptBuilder, PromptId};
+use prompt_store::{PromptBuilder, PromptId, UserPromptId};
 
 use search::{BufferSearchBar, buffer_search::DivRegistrar};
 use settings::{Settings, update_settings_file};
@@ -58,11 +57,11 @@ pub fn init(cx: &mut App) {
                 .register_action(AssistantPanel::show_configuration)
                 .register_action(AssistantPanel::create_new_context)
                 .register_action(AssistantPanel::restart_context_servers)
-                .register_action(|workspace, _: &OpenPromptLibrary, window, cx| {
+                .register_action(|workspace, action: &OpenPromptLibrary, window, cx| {
                     if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {
                         workspace.focus_panel::<AssistantPanel>(window, cx);
                         panel.update(cx, |panel, cx| {
-                            panel.deploy_prompt_library(&OpenPromptLibrary::default(), window, cx)
+                            panel.deploy_prompt_library(action, window, cx)
                         });
                     }
                 });
@@ -489,8 +488,8 @@ impl AssistantPanel {
 
         // If we're signed out and don't have a provider configured, or we're signed-out AND Zed.dev is
         // the provider, we want to show a nudge to sign in.
-        let show_zed_ai_notice = client_status.is_signed_out()
-            && model.map_or(true, |model| model.provider.id().0 == ZED_CLOUD_PROVIDER_ID);
+        let show_zed_ai_notice =
+            client_status.is_signed_out() && model.map_or(true, |model| model.is_provided_by_zed());
 
         self.show_zed_ai_notice = show_zed_ai_notice;
         cx.notify();
@@ -1060,7 +1059,9 @@ impl AssistantPanel {
                     None,
                 ))
             }),
-            action.prompt_to_focus.map(|uuid| PromptId::User { uuid }),
+            action.prompt_to_select.map(|uuid| PromptId::User {
+                uuid: UserPromptId(uuid),
+            }),
             cx,
         )
         .detach_and_log_err(cx);

crates/assistant/src/inline_assistant.rs

@@ -1226,7 +1226,7 @@ impl InlineAssistant {
                 editor.highlight_rows::<InlineAssist>(
                     row_range,
                     cx.theme().status().info_background,
-                    false,
+                    Default::default(),
                     cx,
                 );
             }
@@ -1291,7 +1291,7 @@ impl InlineAssistant {
                     editor.highlight_rows::<DeletedLines>(
                         Anchor::min()..Anchor::max(),
                         cx.theme().status().deleted_background,
-                        false,
+                        Default::default(),
                         cx,
                     );
                     editor

crates/assistant_context_editor/src/context_editor.rs

@@ -2089,7 +2089,7 @@ impl ContextEditor {
                         continue;
                     };
                     let image_id = image.id();
-                    let image_task = LanguageModelImage::from_image(image, cx).shared();
+                    let image_task = LanguageModelImage::from_image(Arc::new(image), cx).shared();
 
                     for image_position in image_positions.iter() {
                         context.insert_content(

crates/assistant_slash_commands/src/prompt_command.rs

@@ -44,9 +44,10 @@ impl SlashCommand for PromptSlashCommand {
         let store = PromptStore::global(cx);
         let query = arguments.to_owned().join(" ");
         cx.spawn(async move |cx| {
+            let cancellation_flag = Arc::new(AtomicBool::default());
             let prompts: Vec<PromptMetadata> = store
                 .await?
-                .read_with(cx, |store, cx| store.search(query, cx))?
+                .read_with(cx, |store, cx| store.search(query, cancellation_flag, cx))?
                 .await;
             Ok(prompts
                 .into_iter()

crates/assistant_tool/Cargo.toml

@@ -28,6 +28,7 @@ serde.workspace = true
 serde_json.workspace = true
 text.workspace = true
 util.workspace = true
+workspace.workspace = true
 workspace-hack.workspace = true
 
 [dev-dependencies]

crates/assistant_tool/src/action_log.rs

@@ -39,10 +39,9 @@ impl ActionLog {
         self.edited_since_project_diagnostics_check
     }
 
-    fn track_buffer(
+    fn track_buffer_internal(
         &mut self,
         buffer: Entity<Buffer>,
-        created: bool,
         cx: &mut Context<Self>,
     ) -> &mut TrackedBuffer {
         let tracked_buffer = self
@@ -59,7 +58,11 @@ impl ActionLog {
                 let base_text;
                 let status;
                 let unreviewed_changes;
-                if created {
+                if buffer
+                    .read(cx)
+                    .file()
+                    .map_or(true, |file| !file.disk_state().exists())
+                {
                     base_text = Rope::default();
                     status = TrackedBufferStatus::Created;
                     unreviewed_changes = Patch::new(vec![Edit {
@@ -146,7 +149,7 @@ impl ActionLog {
                     // resurrected externally, we want to clear the changes we
                     // were tracking and reset the buffer's state.
                     self.tracked_buffers.remove(&buffer);
-                    self.track_buffer(buffer, false, cx);
+                    self.track_buffer_internal(buffer, cx);
                 }
                 cx.notify();
             }
@@ -260,26 +263,15 @@ impl ActionLog {
     }
 
     /// Track a buffer as read, so we can notify the model about user edits.
-    pub fn buffer_read(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
-        self.track_buffer(buffer, false, cx);
-    }
-
-    /// Track a buffer that was added as context, so we can notify the model about user edits.
-    pub fn buffer_added_as_context(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
-        self.track_buffer(buffer, false, cx);
-    }
-
-    /// Track a buffer as read, so we can notify the model about user edits.
-    pub fn will_create_buffer(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
-        self.track_buffer(buffer.clone(), true, cx);
-        self.buffer_edited(buffer, cx)
+    pub fn track_buffer(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
+        self.track_buffer_internal(buffer, cx);
     }
 
     /// Mark a buffer as edited, so we can refresh it in the context
     pub fn buffer_edited(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
         self.edited_since_project_diagnostics_check = true;
 
-        let tracked_buffer = self.track_buffer(buffer.clone(), false, cx);
+        let tracked_buffer = self.track_buffer_internal(buffer.clone(), cx);
         if let TrackedBufferStatus::Deleted = tracked_buffer.status {
             tracked_buffer.status = TrackedBufferStatus::Modified;
         }
@@ -287,7 +279,7 @@ impl ActionLog {
     }
 
     pub fn will_delete_buffer(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
-        let tracked_buffer = self.track_buffer(buffer.clone(), false, cx);
+        let tracked_buffer = self.track_buffer_internal(buffer.clone(), cx);
         match tracked_buffer.status {
             TrackedBufferStatus::Created => {
                 self.tracked_buffers.remove(&buffer);
@@ -397,7 +389,7 @@ impl ActionLog {
 
                 // Clear all tracked changes for this buffer and start over as if we just read it.
                 self.tracked_buffers.remove(&buffer);
-                self.track_buffer(buffer.clone(), false, cx);
+                self.track_buffer_internal(buffer.clone(), cx);
                 cx.notify();
                 save
             }
@@ -695,12 +687,20 @@ mod tests {
         init_test(cx);
 
         let fs = FakeFs::new(cx.executor());
-        let project = Project::test(fs.clone(), [], cx).await;
+        fs.insert_tree(path!("/dir"), json!({"file": "abc\ndef\nghi\njkl\nmno"}))
+            .await;
+        let project = Project::test(fs.clone(), [path!("/dir").as_ref()], cx).await;
         let action_log = cx.new(|_| ActionLog::new(project.clone()));
-        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl\nmno", cx));
+        let file_path = project
+            .read_with(cx, |project, cx| project.find_project_path("dir/file", cx))
+            .unwrap();
+        let buffer = project
+            .update(cx, |project, cx| project.open_buffer(file_path, cx))
+            .await
+            .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 1)..Point::new(1, 2), "E")], None, cx)
@@ -765,12 +765,23 @@ mod tests {
         init_test(cx);
 
         let fs = FakeFs::new(cx.executor());
-        let project = Project::test(fs.clone(), [], cx).await;
+        fs.insert_tree(
+            path!("/dir"),
+            json!({"file": "abc\ndef\nghi\njkl\nmno\npqr"}),
+        )
+        .await;
+        let project = Project::test(fs.clone(), [path!("/dir").as_ref()], cx).await;
         let action_log = cx.new(|_| ActionLog::new(project.clone()));
-        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl\nmno\npqr", cx));
+        let file_path = project
+            .read_with(cx, |project, cx| project.find_project_path("dir/file", cx))
+            .unwrap();
+        let buffer = project
+            .update(cx, |project, cx| project.open_buffer(file_path, cx))
+            .await
+            .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 0)..Point::new(2, 0), "")], None, cx)
@@ -839,12 +850,20 @@ mod tests {
         init_test(cx);
 
         let fs = FakeFs::new(cx.executor());
-        let project = Project::test(fs.clone(), [], cx).await;
+        fs.insert_tree(path!("/dir"), json!({"file": "abc\ndef\nghi\njkl\nmno"}))
+            .await;
+        let project = Project::test(fs.clone(), [path!("/dir").as_ref()], cx).await;
         let action_log = cx.new(|_| ActionLog::new(project.clone()));
-        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl\nmno", cx));
+        let file_path = project
+            .read_with(cx, |project, cx| project.find_project_path("dir/file", cx))
+            .unwrap();
+        let buffer = project
+            .update(cx, |project, cx| project.open_buffer(file_path, cx))
+            .await
+            .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 2)..Point::new(2, 3), "F\nGHI")], None, cx)
@@ -928,25 +947,21 @@ mod tests {
         init_test(cx);
 
         let fs = FakeFs::new(cx.executor());
+        fs.insert_tree(path!("/dir"), json!({})).await;
         let project = Project::test(fs.clone(), [path!("/dir").as_ref()], cx).await;
         let action_log = cx.new(|_| ActionLog::new(project.clone()));
-
-        let fs = FakeFs::new(cx.executor());
-        fs.insert_tree(path!("/dir"), json!({})).await;
-
-        let project = Project::test(fs.clone(), [path!("/dir").as_ref()], cx).await;
         let file_path = project
             .read_with(cx, |project, cx| project.find_project_path("dir/file1", cx))
             .unwrap();
 
-        // Simulate file2 being recreated by a tool.
         let buffer = project
             .update(cx, |project, cx| project.open_buffer(file_path, cx))
             .await
             .unwrap();
         cx.update(|cx| {
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| buffer.set_text("lorem", cx));
-            action_log.update(cx, |log, cx| log.will_create_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
         });
         project
             .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))
@@ -1067,8 +1082,9 @@ mod tests {
             .update(cx, |project, cx| project.open_buffer(file2_path, cx))
             .await
             .unwrap();
+        action_log.update(cx, |log, cx| log.track_buffer(buffer2.clone(), cx));
         buffer2.update(cx, |buffer, cx| buffer.set_text("IPSUM", cx));
-        action_log.update(cx, |log, cx| log.will_create_buffer(buffer2.clone(), cx));
+        action_log.update(cx, |log, cx| log.buffer_edited(buffer2.clone(), cx));
         project
             .update(cx, |project, cx| project.save_buffer(buffer2.clone(), cx))
             .await
@@ -1113,7 +1129,7 @@ mod tests {
             .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 1)..Point::new(1, 2), "E\nXYZ")], None, cx)
@@ -1248,7 +1264,7 @@ mod tests {
             .unwrap();
 
         cx.update(|cx| {
-            action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| {
                 buffer
                     .edit([(Point::new(1, 1)..Point::new(1, 2), "E\nXYZ")], None, cx)
@@ -1381,8 +1397,9 @@ mod tests {
             .await
             .unwrap();
         cx.update(|cx| {
+            action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
             buffer.update(cx, |buffer, cx| buffer.set_text("content", cx));
-            action_log.update(cx, |log, cx| log.will_create_buffer(buffer.clone(), cx));
+            action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
         });
         project
             .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))
@@ -1438,7 +1455,7 @@ mod tests {
             .await
             .unwrap();
 
-        action_log.update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));
+        action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
 
         for _ in 0..operations {
             match rng.gen_range(0..100) {
@@ -1490,7 +1507,7 @@ mod tests {
             log::info!("quiescing...");
             cx.run_until_parked();
             action_log.update(cx, |log, cx| {
-                let tracked_buffer = log.track_buffer(buffer.clone(), false, cx);
+                let tracked_buffer = log.track_buffer_internal(buffer.clone(), cx);
                 let mut old_text = tracked_buffer.base_text.clone();
                 let new_text = buffer.read(cx).as_rope();
                 for edit in tracked_buffer.unreviewed_changes.edits() {

crates/assistant_tool/src/assistant_tool.rs

@@ -10,14 +10,16 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use gpui::AnyElement;
+use gpui::AnyWindowHandle;
 use gpui::Context;
 use gpui::IntoElement;
 use gpui::Window;
-use gpui::{App, Entity, SharedString, Task};
+use gpui::{App, Entity, SharedString, Task, WeakEntity};
 use icons::IconName;
 use language_model::LanguageModelRequestMessage;
 use language_model::LanguageModelToolSchemaFormat;
 use project::Project;
+use workspace::Workspace;
 
 pub use crate::action_log::*;
 pub use crate::tool_registry::*;
@@ -30,6 +32,7 @@ pub fn init(cx: &mut App) {
 
 #[derive(Debug, Clone)]
 pub enum ToolUseStatus {
+    InputStillStreaming,
     NeedsConfirmation,
     Pending,
     Running,
@@ -41,6 +44,7 @@ impl ToolUseStatus {
     pub fn text(&self) -> SharedString {
         match self {
             ToolUseStatus::NeedsConfirmation => "".into(),
+            ToolUseStatus::InputStillStreaming => "".into(),
             ToolUseStatus::Pending => "".into(),
             ToolUseStatus::Running => "".into(),
             ToolUseStatus::Finished(out) => out.clone(),
@@ -63,6 +67,7 @@ pub trait ToolCard: 'static + Sized {
         &mut self,
         status: &ToolUseStatus,
         window: &mut Window,
+        workspace: WeakEntity<Workspace>,
         cx: &mut Context<Self>,
     ) -> impl IntoElement;
 }
@@ -74,6 +79,7 @@ pub struct AnyToolCard {
         entity: gpui::AnyEntity,
         status: &ToolUseStatus,
         window: &mut Window,
+        workspace: WeakEntity<Workspace>,
         cx: &mut App,
     ) -> AnyElement,
 }
@@ -84,11 +90,14 @@ impl<T: ToolCard> From<Entity<T>> for AnyToolCard {
             entity: gpui::AnyEntity,
             status: &ToolUseStatus,
             window: &mut Window,
+            workspace: WeakEntity<Workspace>,
             cx: &mut App,
         ) -> AnyElement {
             let entity = entity.downcast::<T>().unwrap();
             entity.update(cx, |entity, cx| {
-                entity.render(status, window, cx).into_any_element()
+                entity
+                    .render(status, window, workspace, cx)
+                    .into_any_element()
             })
         }
 
@@ -100,8 +109,14 @@ impl<T: ToolCard> From<Entity<T>> for AnyToolCard {
 }
 
 impl AnyToolCard {
-    pub fn render(&self, status: &ToolUseStatus, window: &mut Window, cx: &mut App) -> AnyElement {
-        (self.render)(self.entity.clone(), status, window, cx)
+    pub fn render(
+        &self,
+        status: &ToolUseStatus,
+        window: &mut Window,
+        workspace: WeakEntity<Workspace>,
+        cx: &mut App,
+    ) -> AnyElement {
+        (self.render)(self.entity.clone(), status, window, workspace, cx)
     }
 }
 
@@ -148,6 +163,12 @@ pub trait Tool: 'static + Send + Sync {
     /// Returns markdown to be displayed in the UI for this tool.
     fn ui_text(&self, input: &serde_json::Value) -> String;
 
+    /// Returns markdown to be displayed in the UI for this tool, while the input JSON is still streaming
+    /// (so information may be missing).
+    fn still_streaming_ui_text(&self, input: &serde_json::Value) -> String {
+        self.ui_text(input)
+    }
+
     /// Runs the tool with the provided input.
     fn run(
         self: Arc<Self>,
@@ -155,6 +176,7 @@ pub trait Tool: 'static + Send + Sync {
         messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult;
 }

crates/assistant_tools/Cargo.toml

@@ -14,10 +14,11 @@ path = "src/assistant_tools.rs"
 [dependencies]
 anyhow.workspace = true
 assistant_tool.workspace = true
+buffer_diff.workspace = true
 chrono.workspace = true
 collections.workspace = true
 component.workspace = true
-feature_flags.workspace = true
+editor.workspace = true
 futures.workspace = true
 gpui.workspace = true
 html_to_markdown.workspace = true
@@ -36,11 +37,14 @@ serde_json.workspace = true
 ui.workspace = true
 util.workspace = true
 web_search.workspace = true
+workspace.workspace = true
 workspace-hack.workspace = true
 worktree.workspace = true
 zed_llm_client.workspace = true
 
 [dev-dependencies]
+client = { workspace = true, features = ["test-support"] }
+clock = { workspace = true, features = ["test-support"] }
 collections = { workspace = true, features = ["test-support"] }
 gpui = { workspace = true, features = ["test-support"] }
 language = { workspace = true, features = ["test-support"] }

crates/assistant_tools/src/assistant_tools.rs

@@ -22,15 +22,16 @@ mod schema;
 mod symbol_info_tool;
 mod terminal_tool;
 mod thinking_tool;
+mod ui;
 mod web_search_tool;
 
 use std::sync::Arc;
 
 use assistant_tool::ToolRegistry;
 use copy_path_tool::CopyPathTool;
-use feature_flags::FeatureFlagAppExt;
 use gpui::App;
 use http_client::HttpClientWithUrl;
+use language_model::LanguageModelRegistry;
 use move_path_tool::MovePathTool;
 use web_search_tool::WebSearchTool;
 
@@ -55,6 +56,11 @@ use crate::symbol_info_tool::SymbolInfoTool;
 use crate::terminal_tool::TerminalTool;
 use crate::thinking_tool::ThinkingTool;
 
+pub use create_file_tool::CreateFileToolInput;
+pub use edit_file_tool::EditFileToolInput;
+pub use path_search_tool::PathSearchToolInput;
+pub use read_file_tool::ReadFileToolInput;
+
 pub fn init(http_client: Arc<HttpClientWithUrl>, cx: &mut App) {
     assistant_tool::init(cx);
 
@@ -82,34 +88,45 @@ pub fn init(http_client: Arc<HttpClientWithUrl>, cx: &mut App) {
     registry.register_tool(ThinkingTool);
     registry.register_tool(FetchTool::new(http_client));
 
-    cx.observe_flag::<feature_flags::ZedProWebSearchTool, _>({
-        move |is_enabled, cx| {
-            if is_enabled {
-                ToolRegistry::global(cx).register_tool(WebSearchTool);
-            } else {
-                ToolRegistry::global(cx).unregister_tool(WebSearchTool);
+    cx.subscribe(
+        &LanguageModelRegistry::global(cx),
+        move |registry, event, cx| match event {
+            language_model::Event::DefaultModelChanged => {
+                let using_zed_provider = registry
+                    .read(cx)
+                    .default_model()
+                    .map_or(false, |default| default.is_provided_by_zed());
+                if using_zed_provider {
+                    ToolRegistry::global(cx).register_tool(WebSearchTool);
+                } else {
+                    ToolRegistry::global(cx).unregister_tool(WebSearchTool);
+                }
             }
-        }
-    })
+            _ => {}
+        },
+    )
     .detach();
 }
 
 #[cfg(test)]
 mod tests {
+    use client::Client;
+    use clock::FakeSystemClock;
     use http_client::FakeHttpClient;
 
     use super::*;
 
     #[gpui::test]
     fn test_builtin_tool_schema_compatibility(cx: &mut App) {
-        crate::init(
-            Arc::new(http_client::HttpClientWithUrl::new(
-                FakeHttpClient::with_200_response(),
-                "https://zed.dev",
-                None,
-            )),
+        settings::init(cx);
+
+        let client = Client::new(
+            Arc::new(FakeSystemClock::new()),
+            FakeHttpClient::with_200_response(),
             cx,
         );
+        language_model::init(client.clone(), cx);
+        crate::init(client.http_client(), cx);
 
         for tool in ToolRegistry::global(cx).tools() {
             let actual_schema = tool

crates/assistant_tools/src/batch_tool.rs

@@ -2,7 +2,7 @@ use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult, ToolWorkingSet};
 use futures::future::join_all;
-use gpui::{App, AppContext, Entity, Task};
+use gpui::{AnyWindowHandle, App, AppContext, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -218,6 +218,7 @@ impl Tool for BatchTool {
         messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<BatchToolInput>(input) {
@@ -258,7 +259,9 @@ impl Tool for BatchTool {
                     let action_log = action_log.clone();
                     let messages = messages.clone();
                     let tool_result = cx
-                        .update(|cx| tool.run(invocation.input, &messages, project, action_log, cx))
+                        .update(|cx| {
+                            tool.run(invocation.input, &messages, project, action_log, window, cx)
+                        })
                         .map_err(|err| anyhow!("Failed to start tool '{}': {}", tool_name, err))?;
 
                     tasks.push(tool_result.output);

crates/assistant_tools/src/code_action_tool.rs

@@ -1,6 +1,6 @@
 use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use language::{self, Anchor, Buffer, ToPointUtf16};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::{self, LspAction, Project};
@@ -140,6 +140,7 @@ impl Tool for CodeActionTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<CodeActionToolInput>(input) {
@@ -159,7 +160,7 @@ impl Tool for CodeActionTool {
             };
 
             action_log.update(cx, |action_log, cx| {
-                action_log.buffer_read(buffer.clone(), cx);
+                action_log.track_buffer(buffer.clone(), cx);
             })?;
 
             let range = {

crates/assistant_tools/src/code_symbols_tool.rs

@@ -6,7 +6,7 @@ use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
 use collections::IndexMap;
-use gpui::{App, AsyncApp, Entity, Task};
+use gpui::{AnyWindowHandle, App, AsyncApp, Entity, Task};
 use language::{OutlineItem, ParseStatus, Point};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::{Project, Symbol};
@@ -128,6 +128,7 @@ impl Tool for CodeSymbolsTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<CodeSymbolsInput>(input) {
@@ -174,7 +175,7 @@ pub async fn file_outline(
     };
 
     action_log.update(cx, |action_log, cx| {
-        action_log.buffer_read(buffer.clone(), cx);
+        action_log.track_buffer(buffer.clone(), cx);
     })?;
 
     // Wait until the buffer has been fully parsed, so that we can read its outline.

crates/assistant_tools/src/contents_tool.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use crate::{code_symbols_tool::file_outline, schema::json_schema_for};
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use itertools::Itertools;
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
@@ -102,6 +102,7 @@ impl Tool for ContentsTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<ContentsToolInput>(input) {
@@ -209,7 +210,7 @@ impl Tool for ContentsTool {
                     })?;
 
                     action_log.update(cx, |log, cx| {
-                        log.buffer_read(buffer, cx);
+                        log.track_buffer(buffer, cx);
                     })?;
 
                     Ok(result)
@@ -221,7 +222,7 @@ impl Tool for ContentsTool {
                         let result = buffer.read_with(cx, |buffer, _cx| buffer.text())?;
 
                         action_log.update(cx, |log, cx| {
-                            log.buffer_read(buffer, cx);
+                            log.track_buffer(buffer, cx);
                         })?;
 
                         Ok(result)

crates/assistant_tools/src/copy_path_tool.rs

@@ -1,6 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
+use gpui::AnyWindowHandle;
 use gpui::{App, AppContext, Entity, Task};
 use language_model::LanguageModelRequestMessage;
 use language_model::LanguageModelToolSchemaFormat;
@@ -76,6 +77,7 @@ impl Tool for CopyPathTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<CopyPathToolInput>(input) {

crates/assistant_tools/src/create_directory_tool.rs

@@ -1,6 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
+use gpui::AnyWindowHandle;
 use gpui::{App, Entity, Task};
 use language_model::LanguageModelRequestMessage;
 use language_model::LanguageModelToolSchemaFormat;
@@ -67,6 +68,7 @@ impl Tool for CreateDirectoryTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<CreateDirectoryToolInput>(input) {

crates/assistant_tools/src/create_file_tool.rs

@@ -1,6 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
+use gpui::AnyWindowHandle;
 use gpui::{App, Entity, Task};
 use language_model::LanguageModelRequestMessage;
 use language_model::LanguageModelToolSchemaFormat;
@@ -33,8 +34,18 @@ pub struct CreateFileToolInput {
     pub contents: String,
 }
 
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+struct PartialInput {
+    #[serde(default)]
+    path: String,
+    #[serde(default)]
+    contents: String,
+}
+
 pub struct CreateFileTool;
 
+const DEFAULT_UI_TEXT: &str = "Create file";
+
 impl Tool for CreateFileTool {
     fn name(&self) -> String {
         "create_file".into()
@@ -62,7 +73,14 @@ impl Tool for CreateFileTool {
                 let path = MarkdownString::inline_code(&input.path);
                 format!("Create file {path}")
             }
-            Err(_) => "Create file".to_string(),
+            Err(_) => DEFAULT_UI_TEXT.to_string(),
+        }
+    }
+
+    fn still_streaming_ui_text(&self, input: &serde_json::Value) -> String {
+        match serde_json::from_value::<PartialInput>(input.clone()).ok() {
+            Some(input) if !input.path.is_empty() => input.path,
+            _ => DEFAULT_UI_TEXT.to_string(),
         }
     }
 
@@ -72,6 +90,7 @@ impl Tool for CreateFileTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<CreateFileToolInput>(input) {
@@ -95,9 +114,12 @@ impl Tool for CreateFileTool {
                 .await
                 .map_err(|err| anyhow!("Unable to open buffer for {destination_path}: {err}"))?;
             cx.update(|cx| {
+                action_log.update(cx, |action_log, cx| {
+                    action_log.track_buffer(buffer.clone(), cx)
+                });
                 buffer.update(cx, |buffer, cx| buffer.set_text(contents, cx));
                 action_log.update(cx, |action_log, cx| {
-                    action_log.will_create_buffer(buffer.clone(), cx)
+                    action_log.buffer_edited(buffer.clone(), cx)
                 });
             })?;
 
@@ -111,3 +133,60 @@ impl Tool for CreateFileTool {
         .into()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn still_streaming_ui_text_with_path() {
+        let tool = CreateFileTool;
+        let input = json!({
+            "path": "src/main.rs",
+            "contents": "fn main() {\n    println!(\"Hello, world!\");\n}"
+        });
+
+        assert_eq!(tool.still_streaming_ui_text(&input), "src/main.rs");
+    }
+
+    #[test]
+    fn still_streaming_ui_text_without_path() {
+        let tool = CreateFileTool;
+        let input = json!({
+            "path": "",
+            "contents": "fn main() {\n    println!(\"Hello, world!\");\n}"
+        });
+
+        assert_eq!(tool.still_streaming_ui_text(&input), DEFAULT_UI_TEXT);
+    }
+
+    #[test]
+    fn still_streaming_ui_text_with_null() {
+        let tool = CreateFileTool;
+        let input = serde_json::Value::Null;
+
+        assert_eq!(tool.still_streaming_ui_text(&input), DEFAULT_UI_TEXT);
+    }
+
+    #[test]
+    fn ui_text_with_valid_input() {
+        let tool = CreateFileTool;
+        let input = json!({
+            "path": "src/main.rs",
+            "contents": "fn main() {\n    println!(\"Hello, world!\");\n}"
+        });
+
+        assert_eq!(tool.ui_text(&input), "Create file `src/main.rs`");
+    }
+
+    #[test]
+    fn ui_text_with_invalid_input() {
+        let tool = CreateFileTool;
+        let input = json!({
+            "invalid": "field"
+        });
+
+        assert_eq!(tool.ui_text(&input), DEFAULT_UI_TEXT);
+    }
+}

crates/assistant_tools/src/delete_path_tool.rs

@@ -2,7 +2,7 @@ use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
 use futures::{SinkExt, StreamExt, channel::mpsc};
-use gpui::{App, AppContext, Entity, Task};
+use gpui::{AnyWindowHandle, App, AppContext, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::{Project, ProjectPath};
 use schemars::JsonSchema;
@@ -62,6 +62,7 @@ impl Tool for DeletePathTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let path_str = match serde_json::from_value::<DeletePathToolInput>(input) {

crates/assistant_tools/src/diagnostics_tool.rs

@@ -1,7 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use language::{DiagnosticSeverity, OffsetRangeExt};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
@@ -82,6 +82,7 @@ impl Tool for DiagnosticsTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         match serde_json::from_value::<DiagnosticsToolInput>(input)

crates/assistant_tools/src/edit_file_tool.rs

@@ -1,15 +1,28 @@
-use crate::{replace::replace_with_flexible_indent, schema::json_schema_for};
+use crate::{
+    replace::{replace_exact, replace_with_flexible_indent},
+    schema::json_schema_for,
+};
 use anyhow::{Context as _, Result, anyhow};
-use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, AppContext, AsyncApp, Entity, Task};
+use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolCard, ToolResult, ToolUseStatus};
+use buffer_diff::{BufferDiff, BufferDiffSnapshot};
+use editor::{Editor, EditorMode, MultiBuffer, PathKey};
+use gpui::{
+    AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EntityId, Task, WeakEntity,
+};
+use language::{
+    Anchor, Buffer, Capability, LanguageRegistry, LineEnding, OffsetRangeExt, Rope, TextBuffer,
+};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
-use std::{path::PathBuf, sync::Arc};
-use ui::IconName;
-
-use crate::replace::replace_exact;
+use std::{
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+use ui::{Disclosure, Tooltip, Window, prelude::*};
+use util::ResultExt;
+use workspace::Workspace;
 
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
 pub struct EditFileToolInput {
@@ -47,8 +60,22 @@ pub struct EditFileToolInput {
     pub new_string: String,
 }
 
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+struct PartialInput {
+    #[serde(default)]
+    path: String,
+    #[serde(default)]
+    display_description: String,
+    #[serde(default)]
+    old_string: String,
+    #[serde(default)]
+    new_string: String,
+}
+
 pub struct EditFileTool;
 
+const DEFAULT_UI_TEXT: &str = "Editing file";
+
 impl Tool for EditFileTool {
     fn name(&self) -> String {
         "edit_file".into()
@@ -73,16 +100,33 @@ impl Tool for EditFileTool {
     fn ui_text(&self, input: &serde_json::Value) -> String {
         match serde_json::from_value::<EditFileToolInput>(input.clone()) {
             Ok(input) => input.display_description,
-            Err(_) => "Edit file".to_string(),
+            Err(_) => "Editing file".to_string(),
         }
     }
 
+    fn still_streaming_ui_text(&self, input: &serde_json::Value) -> String {
+        if let Some(input) = serde_json::from_value::<PartialInput>(input.clone()).ok() {
+            let description = input.display_description.trim();
+            if !description.is_empty() {
+                return description.to_string();
+            }
+
+            let path = input.path.trim();
+            if !path.is_empty() {
+                return path.to_string();
+            }
+        }
+
+        DEFAULT_UI_TEXT.to_string()
+    }
+
     fn run(
         self: Arc<Self>,
         input: serde_json::Value,
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<EditFileToolInput>(input) {
@@ -90,7 +134,18 @@ impl Tool for EditFileTool {
             Err(err) => return Task::ready(Err(anyhow!(err))).into(),
         };
 
-        cx.spawn(async move |cx: &mut AsyncApp| {
+        let card = window.and_then(|window| {
+            window
+                .update(cx, |_, window, cx| {
+                    cx.new(|cx| {
+                        EditFileToolCard::new(input.path.clone(), project.clone(), window, cx)
+                    })
+                })
+                .ok()
+        });
+
+        let card_clone = card.clone();
+        let task = cx.spawn(async move |cx: &mut AsyncApp| {
             let project_path = project.read_with(cx, |project, cx| {
                 project
                     .find_project_path(&input.path, cx)
@@ -98,26 +153,38 @@ impl Tool for EditFileTool {
             })??;
 
             let buffer = project
-                .update(cx, |project, cx| project.open_buffer(project_path, cx))?
+                .update(cx, |project, cx| {
+                    project.open_buffer(project_path.clone(), cx)
+                })?
                 .await?;
 
             let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
 
             if input.old_string.is_empty() {
-                return Err(anyhow!("`old_string` cannot be empty. Use a different tool if you want to create a file."));
+                return Err(anyhow!(
+                    "`old_string` can't be empty, use another tool if you want to create a file."
+                ));
             }
 
             if input.old_string == input.new_string {
-                return Err(anyhow!("The `old_string` and `new_string` are identical, so no changes would be made."));
+                return Err(anyhow!(
+                    "The `old_string` and `new_string` are identical, so no changes would be made."
+                ));
             }
 
             let result = cx
                 .background_spawn(async move {
                     // Try to match exactly
                     let diff = replace_exact(&input.old_string, &input.new_string, &snapshot)
-                    .await
-                    // If that fails, try being flexible about indentation
-                    .or_else(|| replace_with_flexible_indent(&input.old_string, &input.new_string, &snapshot))?;
+                        .await
+                        // If that fails, try being flexible about indentation
+                        .or_else(|| {
+                            replace_with_flexible_indent(
+                                &input.old_string,
+                                &input.new_string,
+                                &snapshot,
+                            )
+                        })?;
 
                     if diff.edits.is_empty() {
                         return None;
@@ -147,37 +214,471 @@ impl Tool for EditFileTool {
                     }
                 })?;
 
-                return Err(err)
+                return Err(err);
             };
 
             let snapshot = cx.update(|cx| {
-                action_log.update(cx, |log, cx| {
-                    log.buffer_read(buffer.clone(), cx)
-                });
+                action_log.update(cx, |log, cx| log.track_buffer(buffer.clone(), cx));
+
                 let snapshot = buffer.update(cx, |buffer, cx| {
                     buffer.finalize_last_transaction();
                     buffer.apply_diff(diff, cx);
                     buffer.finalize_last_transaction();
                     buffer.snapshot()
                 });
-                action_log.update(cx, |log, cx| {
-                    log.buffer_edited(buffer.clone(), cx)
-                });
+                action_log.update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
                 snapshot
             })?;
 
-            project.update( cx, |project, cx| {
-                project.save_buffer(buffer, cx)
-            })?.await?;
+            project
+                .update(cx, |project, cx| project.save_buffer(buffer, cx))?
+                .await?;
 
-            let diff_str = cx.background_spawn(async move {
-                let new_text = snapshot.text();
-                language::unified_diff(&old_text, &new_text)
-            }).await;
+            let new_text = snapshot.text();
+            let diff_str = cx
+                .background_spawn({
+                    let old_text = old_text.clone();
+                    let new_text = new_text.clone();
+                    async move { language::unified_diff(&old_text, &new_text) }
+                })
+                .await;
 
+            if let Some(card) = card_clone {
+                card.update(cx, |card, cx| {
+                    card.set_diff(project_path.path.clone(), old_text, new_text, cx);
+                })
+                .log_err();
+            }
 
-            Ok(format!("Edited {}:\n\n```diff\n{}\n```", input.path.display(), diff_str))
+            Ok(format!(
+                "Edited {}:\n\n```diff\n{}\n```",
+                input.path.display(),
+                diff_str
+            ))
+        });
 
-        }).into()
+        ToolResult {
+            output: task,
+            card: card.map(AnyToolCard::from),
+        }
+    }
+}
+
+pub struct EditFileToolCard {
+    path: PathBuf,
+    editor: Entity<Editor>,
+    multibuffer: Entity<MultiBuffer>,
+    project: Entity<Project>,
+    diff_task: Option<Task<Result<()>>>,
+    preview_expanded: bool,
+    full_height_expanded: bool,
+    editor_unique_id: EntityId,
+}
+
+impl EditFileToolCard {
+    fn new(path: PathBuf, project: Entity<Project>, window: &mut Window, cx: &mut App) -> Self {
+        let multibuffer = cx.new(|_| MultiBuffer::without_headers(Capability::ReadOnly));
+        let editor = cx.new(|cx| {
+            let mut editor = Editor::new(
+                EditorMode::Full {
+                    scale_ui_elements_with_buffer_font_size: false,
+                    show_active_line_background: false,
+                    sized_by_content: true,
+                },
+                multibuffer.clone(),
+                Some(project.clone()),
+                window,
+                cx,
+            );
+            editor.set_show_scrollbars(false, cx);
+            editor.set_show_gutter(false, cx);
+            editor.disable_inline_diagnostics();
+            editor.disable_scrolling(cx);
+            editor.disable_expand_excerpt_buttons(cx);
+            editor.set_show_breakpoints(false, cx);
+            editor.set_show_code_actions(false, cx);
+            editor.set_show_git_diff_gutter(false, cx);
+            editor.set_expand_all_diff_hunks(cx);
+            editor
+        });
+        Self {
+            editor_unique_id: editor.entity_id(),
+            path,
+            project,
+            editor,
+            multibuffer,
+            diff_task: None,
+            preview_expanded: true,
+            full_height_expanded: false,
+        }
+    }
+
+    fn set_diff(
+        &mut self,
+        path: Arc<Path>,
+        old_text: String,
+        new_text: String,
+        cx: &mut Context<Self>,
+    ) {
+        let language_registry = self.project.read(cx).languages().clone();
+        self.diff_task = Some(cx.spawn(async move |this, cx| {
+            let buffer = build_buffer(new_text, path.clone(), &language_registry, cx).await?;
+            let buffer_diff = build_buffer_diff(old_text, &buffer, &language_registry, cx).await?;
+
+            this.update(cx, |this, cx| {
+                this.multibuffer.update(cx, |multibuffer, cx| {
+                    let snapshot = buffer.read(cx).snapshot();
+                    let diff = buffer_diff.read(cx);
+                    let diff_hunk_ranges = diff
+                        .hunks_intersecting_range(Anchor::MIN..Anchor::MAX, &snapshot, cx)
+                        .map(|diff_hunk| diff_hunk.buffer_range.to_point(&snapshot))
+                        .collect::<Vec<_>>();
+                    let (_, is_newly_added) = multibuffer.set_excerpts_for_path(
+                        PathKey::for_buffer(&buffer, cx),
+                        buffer,
+                        diff_hunk_ranges,
+                        editor::DEFAULT_MULTIBUFFER_CONTEXT,
+                        cx,
+                    );
+                    debug_assert!(is_newly_added);
+                    multibuffer.add_diff(buffer_diff, cx);
+                });
+                cx.notify();
+            })
+        }));
+    }
+}
+
+impl ToolCard for EditFileToolCard {
+    fn render(
+        &mut self,
+        status: &ToolUseStatus,
+        window: &mut Window,
+        workspace: WeakEntity<Workspace>,
+        cx: &mut Context<Self>,
+    ) -> impl IntoElement {
+        let failed = matches!(status, ToolUseStatus::Error(_));
+
+        let path_label_button = h_flex()
+            .id(("edit-tool-path-label-button", self.editor_unique_id))
+            .w_full()
+            .max_w_full()
+            .px_1()
+            .gap_0p5()
+            .cursor_pointer()
+            .rounded_sm()
+            .opacity(0.8)
+            .hover(|label| {
+                label
+                    .opacity(1.)
+                    .bg(cx.theme().colors().element_hover.opacity(0.5))
+            })
+            .tooltip(Tooltip::text("Jump to File"))
+            .child(
+                h_flex()
+                    .child(
+                        Icon::new(IconName::Pencil)
+                            .size(IconSize::XSmall)
+                            .color(Color::Muted),
+                    )
+                    .child(
+                        div()
+                            .text_size(rems(0.8125))
+                            .child(self.path.display().to_string())
+                            .ml_1p5()
+                            .mr_0p5(),
+                    )
+                    .child(
+                        Icon::new(IconName::ArrowUpRight)
+                            .size(IconSize::XSmall)
+                            .color(Color::Ignored),
+                    ),
+            )
+            .on_click({
+                let path = self.path.clone();
+                let workspace = workspace.clone();
+                move |_, window, cx| {
+                    workspace
+                        .update(cx, {
+                            |workspace, cx| {
+                                let Some(project_path) =
+                                    workspace.project().read(cx).find_project_path(&path, cx)
+                                else {
+                                    return;
+                                };
+                                let open_task =
+                                    workspace.open_path(project_path, None, true, window, cx);
+                                window
+                                    .spawn(cx, async move |cx| {
+                                        let item = open_task.await?;
+                                        if let Some(active_editor) = item.downcast::<Editor>() {
+                                            active_editor
+                                                .update_in(cx, |editor, window, cx| {
+                                                    editor.go_to_singleton_buffer_point(
+                                                        language::Point::new(0, 0),
+                                                        window,
+                                                        cx,
+                                                    );
+                                                })
+                                                .log_err();
+                                        }
+                                        anyhow::Ok(())
+                                    })
+                                    .detach_and_log_err(cx);
+                            }
+                        })
+                        .ok();
+                }
+            })
+            .into_any_element();
+
+        let codeblock_header_bg = cx
+            .theme()
+            .colors()
+            .element_background
+            .blend(cx.theme().colors().editor_foreground.opacity(0.025));
+
+        let codeblock_header = h_flex()
+            .flex_none()
+            .p_1()
+            .gap_1()
+            .justify_between()
+            .rounded_t_md()
+            .when(!failed, |header| header.bg(codeblock_header_bg))
+            .child(path_label_button)
+            .map(|container| {
+                if failed {
+                    container.child(
+                        Icon::new(IconName::Close)
+                            .size(IconSize::Small)
+                            .color(Color::Error),
+                    )
+                } else {
+                    container.child(
+                        Disclosure::new(
+                            ("edit-file-disclosure", self.editor_unique_id),
+                            self.preview_expanded,
+                        )
+                        .opened_icon(IconName::ChevronUp)
+                        .closed_icon(IconName::ChevronDown)
+                        .on_click(cx.listener(
+                            move |this, _event, _window, _cx| {
+                                this.preview_expanded = !this.preview_expanded;
+                            },
+                        )),
+                    )
+                }
+            });
+
+        let editor = self.editor.update(cx, |editor, cx| {
+            editor.render(window, cx).into_any_element()
+        });
+
+        let (full_height_icon, full_height_tooltip_label) = if self.full_height_expanded {
+            (IconName::ChevronUp, "Collapse Code Block")
+        } else {
+            (IconName::ChevronDown, "Expand Code Block")
+        };
+
+        let gradient_overlay = div()
+            .absolute()
+            .bottom_0()
+            .left_0()
+            .w_full()
+            .h_2_5()
+            .rounded_b_lg()
+            .bg(gpui::linear_gradient(
+                0.,
+                gpui::linear_color_stop(cx.theme().colors().editor_background, 0.),
+                gpui::linear_color_stop(cx.theme().colors().editor_background.opacity(0.), 1.),
+            ));
+
+        let border_color = cx.theme().colors().border.opacity(0.6);
+
+        v_flex()
+            .mb_2()
+            .border_1()
+            .when(failed, |card| card.border_dashed())
+            .border_color(border_color)
+            .rounded_lg()
+            .overflow_hidden()
+            .child(codeblock_header)
+            .when(!failed && self.preview_expanded, |card| {
+                card.child(
+                    v_flex()
+                        .relative()
+                        .overflow_hidden()
+                        .border_t_1()
+                        .border_color(border_color)
+                        .bg(cx.theme().colors().editor_background)
+                        .map(|editor_container| {
+                            if self.full_height_expanded {
+                                editor_container.h_full()
+                            } else {
+                                editor_container.max_h_64()
+                            }
+                        })
+                        .child(div().pl_1().child(editor))
+                        .when(!self.full_height_expanded, |editor_container| {
+                            editor_container.child(gradient_overlay)
+                        }),
+                )
+            })
+            .when(!failed && self.preview_expanded, |card| {
+                card.child(
+                    h_flex()
+                        .id(("edit-tool-card-inner-hflex", self.editor_unique_id))
+                        .flex_none()
+                        .cursor_pointer()
+                        .h_5()
+                        .justify_center()
+                        .rounded_b_md()
+                        .border_t_1()
+                        .border_color(border_color)
+                        .bg(cx.theme().colors().editor_background)
+                        .hover(|style| style.bg(cx.theme().colors().element_hover.opacity(0.1)))
+                        .child(
+                            Icon::new(full_height_icon)
+                                .size(IconSize::Small)
+                                .color(Color::Muted),
+                        )
+                        .tooltip(Tooltip::text(full_height_tooltip_label))
+                        .on_click(cx.listener(move |this, _event, _window, _cx| {
+                            this.full_height_expanded = !this.full_height_expanded;
+                        })),
+                )
+            })
+    }
+}
+
+async fn build_buffer(
+    mut text: String,
+    path: Arc<Path>,
+    language_registry: &Arc<language::LanguageRegistry>,
+    cx: &mut AsyncApp,
+) -> Result<Entity<Buffer>> {
+    let line_ending = LineEnding::detect(&text);
+    LineEnding::normalize(&mut text);
+    let text = Rope::from(text);
+    let language = cx
+        .update(|_cx| language_registry.language_for_file_path(&path))?
+        .await
+        .ok();
+    let buffer = cx.new(|cx| {
+        let buffer = TextBuffer::new_normalized(
+            0,
+            cx.entity_id().as_non_zero_u64().into(),
+            line_ending,
+            text,
+        );
+        let mut buffer = Buffer::build(buffer, None, Capability::ReadWrite);
+        buffer.set_language(language, cx);
+        buffer
+    })?;
+    Ok(buffer)
+}
+
+async fn build_buffer_diff(
+    mut old_text: String,
+    buffer: &Entity<Buffer>,
+    language_registry: &Arc<LanguageRegistry>,
+    cx: &mut AsyncApp,
+) -> Result<Entity<BufferDiff>> {
+    LineEnding::normalize(&mut old_text);
+
+    let buffer = cx.update(|cx| buffer.read(cx).snapshot())?;
+
+    let base_buffer = cx
+        .update(|cx| {
+            Buffer::build_snapshot(
+                old_text.clone().into(),
+                buffer.language().cloned(),
+                Some(language_registry.clone()),
+                cx,
+            )
+        })?
+        .await;
+
+    let diff_snapshot = cx
+        .update(|cx| {
+            BufferDiffSnapshot::new_with_base_buffer(
+                buffer.text.clone(),
+                Some(old_text.into()),
+                base_buffer,
+                cx,
+            )
+        })?
+        .await;
+
+    cx.new(|cx| {
+        let mut diff = BufferDiff::new(&buffer.text, cx);
+        diff.set_snapshot(diff_snapshot, &buffer.text, cx);
+        diff
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn still_streaming_ui_text_with_path() {
+        let tool = EditFileTool;
+        let input = json!({
+            "path": "src/main.rs",
+            "display_description": "",
+            "old_string": "old code",
+            "new_string": "new code"
+        });
+
+        assert_eq!(tool.still_streaming_ui_text(&input), "src/main.rs");
+    }
+
+    #[test]
+    fn still_streaming_ui_text_with_description() {
+        let tool = EditFileTool;
+        let input = json!({
+            "path": "",
+            "display_description": "Fix error handling",
+            "old_string": "old code",
+            "new_string": "new code"
+        });
+
+        assert_eq!(tool.still_streaming_ui_text(&input), "Fix error handling");
+    }
+
+    #[test]
+    fn still_streaming_ui_text_with_path_and_description() {
+        let tool = EditFileTool;
+        let input = json!({
+            "path": "src/main.rs",
+            "display_description": "Fix error handling",
+            "old_string": "old code",
+            "new_string": "new code"
+        });
+
+        assert_eq!(tool.still_streaming_ui_text(&input), "Fix error handling");
+    }
+
+    #[test]
+    fn still_streaming_ui_text_no_path_or_description() {
+        let tool = EditFileTool;
+        let input = json!({
+            "path": "",
+            "display_description": "",
+            "old_string": "old code",
+            "new_string": "new code"
+        });
+
+        assert_eq!(tool.still_streaming_ui_text(&input), DEFAULT_UI_TEXT);
+    }
+
+    #[test]
+    fn still_streaming_ui_text_with_null() {
+        let tool = EditFileTool;
+        let input = serde_json::Value::Null;
+
+        assert_eq!(tool.still_streaming_ui_text(&input), DEFAULT_UI_TEXT);
     }
 }

crates/assistant_tools/src/fetch_tool.rs

@@ -6,7 +6,7 @@ use crate::schema::json_schema_for;
 use anyhow::{Context as _, Result, anyhow, bail};
 use assistant_tool::{ActionLog, Tool, ToolResult};
 use futures::AsyncReadExt as _;
-use gpui::{App, AppContext as _, Entity, Task};
+use gpui::{AnyWindowHandle, App, AppContext as _, Entity, Task};
 use html_to_markdown::{TagHandler, convert_html_to_markdown, markdown};
 use http_client::{AsyncBody, HttpClientWithUrl};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
@@ -145,6 +145,7 @@ impl Tool for FetchTool {
         _messages: &[LanguageModelRequestMessage],
         _project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<FetchToolInput>(input) {

crates/assistant_tools/src/grep_tool.rs

@@ -2,7 +2,7 @@ use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
 use futures::StreamExt;
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use language::OffsetRangeExt;
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::{
@@ -96,6 +96,7 @@ impl Tool for GrepTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         const CONTEXT_LINES: u32 = 2;
@@ -405,7 +406,7 @@ mod tests {
     ) -> String {
         let tool = Arc::new(GrepTool);
         let action_log = cx.new(|_cx| ActionLog::new(project.clone()));
-        let task = cx.update(|cx| tool.run(input, &[], project, action_log, cx));
+        let task = cx.update(|cx| tool.run(input, &[], project, action_log, None, cx));
 
         match task.output.await {
             Ok(result) => result,

crates/assistant_tools/src/list_directory_tool.rs

@@ -1,7 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -76,6 +76,7 @@ impl Tool for ListDirectoryTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<ListDirectoryToolInput>(input) {

crates/assistant_tools/src/move_path_tool.rs

@@ -1,7 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, AppContext, Entity, Task};
+use gpui::{AnyWindowHandle, App, AppContext, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -89,6 +89,7 @@ impl Tool for MovePathTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<MovePathToolInput>(input) {

crates/assistant_tools/src/now_tool.rs

@@ -4,7 +4,7 @@ use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
 use chrono::{Local, Utc};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -59,6 +59,7 @@ impl Tool for NowTool {
         _messages: &[LanguageModelRequestMessage],
         _project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         _cx: &mut App,
     ) -> ToolResult {
         let input: NowToolInput = match serde_json::from_value(input) {

crates/assistant_tools/src/open_tool.rs

@@ -1,7 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, AppContext, Entity, Task};
+use gpui::{AnyWindowHandle, App, AppContext, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -52,6 +52,7 @@ impl Tool for OpenTool {
         _messages: &[LanguageModelRequestMessage],
         _project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input: OpenToolInput = match serde_json::from_value(input) {

crates/assistant_tools/src/path_search_tool.rs

@@ -1,7 +1,7 @@
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, AppContext, Entity, Task};
+use gpui::{AnyWindowHandle, App, AppContext, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -70,6 +70,7 @@ impl Tool for PathSearchTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let (offset, glob) = match serde_json::from_value::<PathSearchToolInput>(input) {

crates/assistant_tools/src/read_file_tool.rs

@@ -1,7 +1,8 @@
 use crate::{code_symbols_tool::file_outline, schema::json_schema_for};
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
+
 use indoc::formatdoc;
 use itertools::Itertools;
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
@@ -87,6 +88,7 @@ impl Tool for ReadFileTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<ReadFileToolInput>(input) {
@@ -134,7 +136,7 @@ impl Tool for ReadFileTool {
                 })?;
 
                 action_log.update(cx, |log, cx| {
-                    log.buffer_read(buffer, cx);
+                    log.track_buffer(buffer, cx);
                 })?;
 
                 Ok(result)
@@ -147,7 +149,7 @@ impl Tool for ReadFileTool {
                     let result = buffer.read_with(cx, |buffer, _cx| buffer.text())?;
 
                     action_log.update(cx, |log, cx| {
-                        log.buffer_read(buffer, cx);
+                        log.track_buffer(buffer, cx);
                     })?;
 
                     Ok(result)
@@ -193,7 +195,7 @@ mod test {
                     "path": "root/nonexistent_file.txt"
                 });
                 Arc::new(ReadFileTool)
-                    .run(input, &[], project.clone(), action_log, cx)
+                    .run(input, &[], project.clone(), action_log, None, cx)
                     .output
             })
             .await;
@@ -223,7 +225,7 @@ mod test {
                     "path": "root/small_file.txt"
                 });
                 Arc::new(ReadFileTool)
-                    .run(input, &[], project.clone(), action_log, cx)
+                    .run(input, &[], project.clone(), action_log, None, cx)
                     .output
             })
             .await;
@@ -253,7 +255,7 @@ mod test {
                     "path": "root/large_file.rs"
                 });
                 Arc::new(ReadFileTool)
-                    .run(input, &[], project.clone(), action_log.clone(), cx)
+                    .run(input, &[], project.clone(), action_log.clone(), None, cx)
                     .output
             })
             .await;
@@ -277,7 +279,7 @@ mod test {
                     "offset": 1
                 });
                 Arc::new(ReadFileTool)
-                    .run(input, &[], project.clone(), action_log, cx)
+                    .run(input, &[], project.clone(), action_log, None, cx)
                     .output
             })
             .await;
@@ -323,7 +325,7 @@ mod test {
                     "end_line": 4
                 });
                 Arc::new(ReadFileTool)
-                    .run(input, &[], project.clone(), action_log, cx)
+                    .run(input, &[], project.clone(), action_log, None, cx)
                     .output
             })
             .await;

crates/assistant_tools/src/rename_tool.rs

@@ -1,6 +1,6 @@
 use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use language::{self, Buffer, ToPointUtf16};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
@@ -87,6 +87,7 @@ impl Tool for RenameTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<RenameToolInput>(input) {
@@ -106,7 +107,7 @@ impl Tool for RenameTool {
             };
 
             action_log.update(cx, |action_log, cx| {
-                action_log.buffer_read(buffer.clone(), cx);
+                action_log.track_buffer(buffer.clone(), cx);
             })?;
 
             let position = {

crates/assistant_tools/src/symbol_info_tool.rs

@@ -1,6 +1,6 @@
 use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, AsyncApp, Entity, Task};
+use gpui::{AnyWindowHandle, App, AsyncApp, Entity, Task};
 use language::{self, Anchor, Buffer, BufferSnapshot, Location, Point, ToPoint, ToPointUtf16};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
@@ -121,6 +121,7 @@ impl Tool for SymbolInfoTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<SymbolInfoToolInput>(input) {
@@ -140,7 +141,7 @@ impl Tool for SymbolInfoTool {
             };
 
             action_log.update(cx, |action_log, cx| {
-                action_log.buffer_read(buffer.clone(), cx);
+                action_log.track_buffer(buffer.clone(), cx);
             })?;
 
             let position = {

crates/assistant_tools/src/terminal_tool.rs

@@ -3,7 +3,7 @@ use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
 use futures::io::BufReader;
 use futures::{AsyncBufReadExt, AsyncReadExt, FutureExt};
-use gpui::{App, AppContext, Entity, Task};
+use gpui::{AnyWindowHandle, App, AppContext, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -78,6 +78,7 @@ impl Tool for TerminalTool {
         _messages: &[LanguageModelRequestMessage],
         project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input: TerminalToolInput = match serde_json::from_value(input) {

crates/assistant_tools/src/thinking_tool.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use crate::schema::json_schema_for;
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolResult};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
@@ -50,6 +50,7 @@ impl Tool for ThinkingTool {
         _messages: &[LanguageModelRequestMessage],
         _project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         _cx: &mut App,
     ) -> ToolResult {
         // This tool just "thinks out loud" and doesn't perform any actions.

crates/assistant_tools/src/ui.rs

@@ -0,0 +1,3 @@
+mod tool_call_card_header;
+
+pub use tool_call_card_header::*;

crates/assistant_tools/src/ui/tool_call_card_header.rs

@@ -0,0 +1,102 @@
+use gpui::{Animation, AnimationExt, App, IntoElement, pulsating_between};
+use std::time::Duration;
+use ui::{Tooltip, prelude::*};
+
+/// A reusable header component for tool call cards.
+#[derive(IntoElement)]
+pub struct ToolCallCardHeader {
+    icon: IconName,
+    primary_text: SharedString,
+    secondary_text: Option<SharedString>,
+    is_loading: bool,
+    error: Option<String>,
+}
+
+impl ToolCallCardHeader {
+    pub fn new(icon: IconName, primary_text: impl Into<SharedString>) -> Self {
+        Self {
+            icon,
+            primary_text: primary_text.into(),
+            secondary_text: None,
+            is_loading: false,
+            error: None,
+        }
+    }
+
+    pub fn with_secondary_text(mut self, text: impl Into<SharedString>) -> Self {
+        self.secondary_text = Some(text.into());
+        self
+    }
+
+    pub fn loading(mut self) -> Self {
+        self.is_loading = true;
+        self
+    }
+
+    pub fn with_error(mut self, error: impl Into<String>) -> Self {
+        self.error = Some(error.into());
+        self
+    }
+}
+
+impl RenderOnce for ToolCallCardHeader {
+    fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement {
+        let font_size = rems(0.8125);
+        let secondary_text = self.secondary_text;
+
+        h_flex()
+            .id("tool-label-container")
+            .gap_1p5()
+            .max_w_full()
+            .overflow_x_scroll()
+            .opacity(0.8)
+            .child(
+                h_flex().h(window.line_height()).justify_center().child(
+                    Icon::new(self.icon)
+                        .size(IconSize::XSmall)
+                        .color(Color::Muted),
+                ),
+            )
+            .child(
+                h_flex()
+                    .h(window.line_height())
+                    .gap_1p5()
+                    .text_size(font_size)
+                    .map(|this| {
+                        if let Some(error) = &self.error {
+                            this.child(format!("{} failed", self.primary_text)).child(
+                                IconButton::new("error_info", IconName::Warning)
+                                    .shape(ui::IconButtonShape::Square)
+                                    .icon_size(IconSize::XSmall)
+                                    .icon_color(Color::Warning)
+                                    .tooltip(Tooltip::text(error.clone())),
+                            )
+                        } else {
+                            this.child(self.primary_text.clone())
+                        }
+                    })
+                    .when_some(secondary_text, |this, secondary_text| {
+                        this.child(
+                            div()
+                                .size(px(3.))
+                                .rounded_full()
+                                .bg(cx.theme().colors().text),
+                        )
+                        .child(div().text_size(font_size).child(secondary_text.clone()))
+                    })
+                    .with_animation(
+                        "loading-label",
+                        Animation::new(Duration::from_secs(2))
+                            .repeat()
+                            .with_easing(pulsating_between(0.6, 1.)),
+                        move |this, delta| {
+                            if self.is_loading {
+                                this.opacity(delta)
+                            } else {
+                                this
+                            }
+                        },
+                    ),
+            )
+    }
+}

crates/assistant_tools/src/web_search_tool.rs

@@ -1,12 +1,12 @@
 use std::{sync::Arc, time::Duration};
 
 use crate::schema::json_schema_for;
+use crate::ui::ToolCallCardHeader;
 use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool, ToolCard, ToolResult, ToolUseStatus};
-use futures::{FutureExt, TryFutureExt};
+use futures::{Future, FutureExt, TryFutureExt};
 use gpui::{
-    Animation, AnimationExt, App, AppContext, Context, Entity, IntoElement, Task, Window,
-    pulsating_between,
+    AnyWindowHandle, App, AppContext, Context, Entity, IntoElement, Task, WeakEntity, Window,
 };
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
@@ -14,6 +14,7 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use ui::{IconName, Tooltip, prelude::*};
 use web_search::WebSearchRegistry;
+use workspace::Workspace;
 use zed_llm_client::{WebSearchCitation, WebSearchResponse};
 
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
@@ -47,7 +48,7 @@ impl Tool for WebSearchTool {
     }
 
     fn ui_text(&self, _input: &serde_json::Value) -> String {
-        "Web Search".to_string()
+        "Searching the Web".to_string()
     }
 
     fn run(
@@ -56,6 +57,7 @@ impl Tool for WebSearchTool {
         _messages: &[LanguageModelRequestMessage],
         _project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         let input = match serde_json::from_value::<WebSearchToolInput>(input) {
@@ -113,63 +115,33 @@ impl ToolCard for WebSearchToolCard {
         &mut self,
         _status: &ToolUseStatus,
         _window: &mut Window,
+        _workspace: WeakEntity<Workspace>,
         cx: &mut Context<Self>,
     ) -> impl IntoElement {
-        let header = h_flex()
-            .id("tool-label-container")
-            .gap_1p5()
-            .max_w_full()
-            .overflow_x_scroll()
-            .child(
-                Icon::new(IconName::Globe)
-                    .size(IconSize::XSmall)
-                    .color(Color::Muted),
-            )
-            .child(match self.response.as_ref() {
-                Some(Ok(response)) => {
-                    let text: SharedString = if response.citations.len() == 1 {
-                        "1 result".into()
-                    } else {
-                        format!("{} results", response.citations.len()).into()
-                    };
-                    h_flex()
-                        .gap_1p5()
-                        .child(Label::new("Searched the Web").size(LabelSize::Small))
-                        .child(
-                            div()
-                                .size(px(3.))
-                                .rounded_full()
-                                .bg(cx.theme().colors().text),
-                        )
-                        .child(Label::new(text).size(LabelSize::Small))
-                        .into_any_element()
-                }
-                Some(Err(error)) => div()
-                    .id("web-search-error")
-                    .child(Label::new("Web Search failed").size(LabelSize::Small))
-                    .tooltip(Tooltip::text(error.to_string()))
-                    .into_any_element(),
-
-                None => Label::new("Searching the Web…")
-                    .size(LabelSize::Small)
-                    .with_animation(
-                        "web-search-label",
-                        Animation::new(Duration::from_secs(2))
-                            .repeat()
-                            .with_easing(pulsating_between(0.6, 1.)),
-                        |label, delta| label.alpha(delta),
-                    )
-                    .into_any_element(),
-            })
-            .into_any();
+        let header = match self.response.as_ref() {
+            Some(Ok(response)) => {
+                let text: SharedString = if response.citations.len() == 1 {
+                    "1 result".into()
+                } else {
+                    format!("{} results", response.citations.len()).into()
+                };
+                ToolCallCardHeader::new(IconName::Globe, "Searched the Web")
+                    .with_secondary_text(text)
+            }
+            Some(Err(error)) => {
+                ToolCallCardHeader::new(IconName::Globe, "Web Search").with_error(error.to_string())
+            }
+            None => ToolCallCardHeader::new(IconName::Globe, "Searching the Web").loading(),
+        };
 
         let content =
             self.response.as_ref().and_then(|response| match response {
                 Ok(response) => {
                     Some(
                         v_flex()
+                            .overflow_hidden()
                             .ml_1p5()
-                            .pl_1p5()
+                            .pl(px(5.))
                             .border_l_1()
                             .border_color(cx.theme().colors().border_variant)
                             .gap_1()
@@ -209,7 +181,7 @@ impl ToolCard for WebSearchToolCard {
                 Err(_) => None,
             });
 
-        v_flex().my_2().gap_1().child(header).children(content)
+        v_flex().mb_3().gap_1().child(header).children(content)
     }
 }
 
@@ -253,8 +225,13 @@ impl Component for WebSearchTool {
                         div()
                             .size_full()
                             .child(in_progress_search.update(cx, |tool, cx| {
-                                tool.render(&ToolUseStatus::Pending, window, cx)
-                                    .into_any_element()
+                                tool.render(
+                                    &ToolUseStatus::Pending,
+                                    window,
+                                    WeakEntity::new_invalid(),
+                                    cx,
+                                )
+                                .into_any_element()
                             }))
                             .into_any_element(),
                     ),
@@ -263,8 +240,13 @@ impl Component for WebSearchTool {
                         div()
                             .size_full()
                             .child(successful_search.update(cx, |tool, cx| {
-                                tool.render(&ToolUseStatus::Finished("".into()), window, cx)
-                                    .into_any_element()
+                                tool.render(
+                                    &ToolUseStatus::Finished("".into()),
+                                    window,
+                                    WeakEntity::new_invalid(),
+                                    cx,
+                                )
+                                .into_any_element()
                             }))
                             .into_any_element(),
                     ),
@@ -273,8 +255,13 @@ impl Component for WebSearchTool {
                         div()
                             .size_full()
                             .child(error_search.update(cx, |tool, cx| {
-                                tool.render(&ToolUseStatus::Error("".into()), window, cx)
-                                    .into_any_element()
+                                tool.render(
+                                    &ToolUseStatus::Error("".into()),
+                                    window,
+                                    WeakEntity::new_invalid(),
+                                    cx,
+                                )
+                                .into_any_element()
                             }))
                             .into_any_element(),
                     ),

crates/client/src/telemetry.rs

@@ -4,7 +4,7 @@ use crate::TelemetrySettings;
 use anyhow::Result;
 use clock::SystemClock;
 use futures::channel::mpsc;
-use futures::{Future, StreamExt};
+use futures::{Future, FutureExt, StreamExt};
 use gpui::{App, AppContext as _, BackgroundExecutor, Task};
 use http_client::{self, AsyncBody, HttpClient, HttpClientWithUrl, Method, Request};
 use parking_lot::Mutex;
@@ -290,6 +290,10 @@ impl Telemetry {
         paths::logs_dir().join("telemetry.log")
     }
 
+    pub fn has_checksum_seed(&self) -> bool {
+        ZED_CLIENT_CHECKSUM_SEED.is_some()
+    }
+
     pub fn start(
         self: &Arc<Self>,
         system_id: Option<String>,
@@ -430,7 +434,7 @@ impl Telemetry {
             let executor = self.executor.clone();
             state.flush_events_task = Some(self.executor.spawn(async move {
                 executor.timer(FLUSH_INTERVAL).await;
-                this.flush_events();
+                this.flush_events().detach();
             }));
         }
 
@@ -456,7 +460,7 @@ impl Telemetry {
 
         if state.installation_id.is_some() && state.events_queue.len() >= state.max_queue_size {
             drop(state);
-            self.flush_events();
+            self.flush_events().detach();
         }
     }
 
@@ -499,60 +503,59 @@ impl Telemetry {
             .body(json_bytes.into())?)
     }
 
-    pub fn flush_events(self: &Arc<Self>) {
+    pub fn flush_events(self: &Arc<Self>) -> Task<()> {
         let mut state = self.state.lock();
         state.first_event_date_time = None;
         let mut events = mem::take(&mut state.events_queue);
         state.flush_events_task.take();
         drop(state);
         if events.is_empty() {
-            return;
+            return Task::ready(());
         }
 
         let this = self.clone();
-        self.executor
-            .spawn(
-                async move {
-                    let mut json_bytes = Vec::new();
+        self.executor.spawn(
+            async move {
+                let mut json_bytes = Vec::new();
 
-                    if let Some(file) = &mut this.state.lock().log_file {
-                        for event in &mut events {
-                            json_bytes.clear();
-                            serde_json::to_writer(&mut json_bytes, event)?;
-                            file.write_all(&json_bytes)?;
-                            file.write_all(b"\n")?;
-                        }
+                if let Some(file) = &mut this.state.lock().log_file {
+                    for event in &mut events {
+                        json_bytes.clear();
+                        serde_json::to_writer(&mut json_bytes, event)?;
+                        file.write_all(&json_bytes)?;
+                        file.write_all(b"\n")?;
                     }
-
-                    let request_body = {
-                        let state = this.state.lock();
-
-                        EventRequestBody {
-                            system_id: state.system_id.as_deref().map(Into::into),
-                            installation_id: state.installation_id.as_deref().map(Into::into),
-                            session_id: state.session_id.clone(),
-                            metrics_id: state.metrics_id.as_deref().map(Into::into),
-                            is_staff: state.is_staff,
-                            app_version: state.app_version.clone(),
-                            os_name: state.os_name.clone(),
-                            os_version: state.os_version.clone(),
-                            architecture: state.architecture.to_string(),
-
-                            release_channel: state.release_channel.map(Into::into),
-                            events,
-                        }
-                    };
-
-                    let request = this.build_request(json_bytes, request_body)?;
-                    let response = this.http_client.send(request).await?;
-                    if response.status() != 200 {
-                        log::error!("Failed to send events: HTTP {:?}", response.status());
-                    }
-                    anyhow::Ok(())
                 }
-                .log_err(),
-            )
-            .detach();
+
+                let request_body = {
+                    let state = this.state.lock();
+
+                    EventRequestBody {
+                        system_id: state.system_id.as_deref().map(Into::into),
+                        installation_id: state.installation_id.as_deref().map(Into::into),
+                        session_id: state.session_id.clone(),
+                        metrics_id: state.metrics_id.as_deref().map(Into::into),
+                        is_staff: state.is_staff,
+                        app_version: state.app_version.clone(),
+                        os_name: state.os_name.clone(),
+                        os_version: state.os_version.clone(),
+                        architecture: state.architecture.to_string(),
+
+                        release_channel: state.release_channel.map(Into::into),
+                        events,
+                    }
+                };
+
+                let request = this.build_request(json_bytes, request_body)?;
+                let response = this.http_client.send(request).await?;
+                if response.status() != 200 {
+                    log::error!("Failed to send events: HTTP {:?}", response.status());
+                }
+                anyhow::Ok(())
+            }
+            .log_err()
+            .map(|_| ()),
+        )
     }
 }
 

crates/collab/Cargo.toml

@@ -128,6 +128,7 @@ serde_json.workspace = true
 session = { workspace = true, features = ["test-support"] }
 settings = { workspace = true, features = ["test-support"] }
 sqlx = { version = "0.8", features = ["sqlite"] }
+task.workspace = true
 theme.workspace = true
 unindent.workspace = true
 util.workspace = true

crates/collab/migrations.sqlite/20221109000000_test_schema.sql

@@ -492,7 +492,8 @@ CREATE TABLE IF NOT EXISTS billing_customers (
     created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
     user_id INTEGER NOT NULL REFERENCES users (id),
     has_overdue_invoices BOOLEAN NOT NULL DEFAULT FALSE,
-    stripe_customer_id TEXT NOT NULL
+    stripe_customer_id TEXT NOT NULL,
+    trial_started_at TIMESTAMP
 );
 
 CREATE UNIQUE INDEX "uix_billing_customers_on_user_id" ON billing_customers (user_id);

crates/collab/migrations/20250422194500_add_trial_started_at_to_billing_customers.sql

@@ -0,0 +1,2 @@
+alter table billing_customers
+    add column trial_started_at timestamp without time zone;

crates/collab/migrations/20250423150129_add_head_commit_details_to_project_repositories.sql

@@ -0,0 +1,2 @@
+alter table project_repositories
+    add column head_commit_details varchar;

crates/collab/src/api/billing.rs

@@ -287,7 +287,7 @@ async fn create_billing_subscription(
         }
     }
 
-    let customer_id = if let Some(existing_customer) = existing_billing_customer {
+    let customer_id = if let Some(existing_customer) = &existing_billing_customer {
         CustomerId::from_str(&existing_customer.stripe_customer_id)
             .context("failed to parse customer ID")?
     } else {
@@ -320,6 +320,15 @@ async fn create_billing_subscription(
                 .await?
         }
         Some(ProductCode::ZedProTrial) => {
+            if let Some(existing_billing_customer) = &existing_billing_customer {
+                if existing_billing_customer.trial_started_at.is_some() {
+                    return Err(Error::http(
+                        StatusCode::FORBIDDEN,
+                        "user already used free trial".into(),
+                    ));
+                }
+            }
+
             stripe_billing
                 .checkout_with_zed_pro_trial(
                     app.config.zed_pro_price_id()?,
@@ -817,6 +826,24 @@ async fn handle_customer_subscription_event(
             .await?
             .ok_or_else(|| anyhow!("billing customer not found"))?;
 
+    if let Some(SubscriptionKind::ZedProTrial) = subscription_kind {
+        if subscription.status == SubscriptionStatus::Trialing {
+            let current_period_start =
+                DateTime::from_timestamp(subscription.current_period_start, 0)
+                    .ok_or_else(|| anyhow!("No trial subscription period start"))?;
+
+            app.db
+                .update_billing_customer(
+                    billing_customer.id,
+                    &UpdateBillingCustomerParams {
+                        trial_started_at: ActiveValue::set(Some(current_period_start.naive_utc())),
+                        ..Default::default()
+                    },
+                )
+                .await?;
+        }
+    }
+
     let was_canceled_due_to_payment_failure = subscription.status == SubscriptionStatus::Canceled
         && subscription
             .cancellation_details
@@ -843,6 +870,28 @@ async fn handle_customer_subscription_event(
         .get_billing_subscription_by_stripe_subscription_id(&subscription.id)
         .await?
     {
+        let llm_db = app
+            .llm_db
+            .clone()
+            .ok_or_else(|| anyhow!("LLM DB not initialized"))?;
+
+        let new_period_start_at =
+            chrono::DateTime::from_timestamp(subscription.current_period_start, 0)
+                .ok_or_else(|| anyhow!("No subscription period start"))?;
+        let new_period_end_at =
+            chrono::DateTime::from_timestamp(subscription.current_period_end, 0)
+                .ok_or_else(|| anyhow!("No subscription period end"))?;
+
+        llm_db
+            .transfer_existing_subscription_usage(
+                billing_customer.user_id,
+                &existing_subscription,
+                subscription_kind,
+                new_period_start_at,
+                new_period_end_at,
+            )
+            .await?;
+
         app.db
             .update_billing_subscription(
                 existing_subscription.id,

crates/collab/src/api/events.rs

@@ -516,6 +516,7 @@ pub async fn post_events(
     if let Some(kinesis_client) = app.kinesis_client.clone() {
         if let Some(stream) = app.config.kinesis_stream.clone() {
             let mut request = kinesis_client.put_records().stream_name(stream);
+            let mut has_records = false;
             for row in for_snowflake(
                 request_body.clone(),
                 first_event_at,
@@ -530,9 +531,12 @@ pub async fn post_events(
                             .build()
                             .unwrap(),
                     );
+                    has_records = true;
                 }
             }
-            request.send().await.log_err();
+            if has_records {
+                request.send().await.log_err();
+            }
         }
     };
 
@@ -555,7 +559,7 @@ fn for_snowflake(
     country_code: Option<String>,
     checksum_matched: bool,
 ) -> impl Iterator<Item = SnowflakeRow> {
-    body.events.into_iter().flat_map(move |event| {
+    body.events.into_iter().filter_map(move |event| {
         let timestamp =
             first_event_at + Duration::milliseconds(event.milliseconds_since_first_event);
         // We will need to double check, but I believe all of the events that
@@ -744,9 +748,11 @@ fn for_snowflake(
         // NOTE: most amplitude user properties are read out of our event_properties
         // dictionary. See https://app.amplitude.com/data/zed/Zed/sources/detail/production/falcon%3A159998
         // for how that is configured.
-        let user_properties = Some(serde_json::json!({
-            "is_staff": body.is_staff,
-        }));
+        let user_properties = body.is_staff.map(|is_staff| {
+            serde_json::json!({
+                "is_staff": is_staff,
+            })
+        });
 
         Some(SnowflakeRow {
             time: timestamp,

crates/collab/src/db/queries/billing_customers.rs

@@ -11,6 +11,7 @@ pub struct UpdateBillingCustomerParams {
     pub user_id: ActiveValue<UserId>,
     pub stripe_customer_id: ActiveValue<String>,
     pub has_overdue_invoices: ActiveValue<bool>,
+    pub trial_started_at: ActiveValue<Option<DateTime>>,
 }
 
 impl Database {
@@ -45,7 +46,8 @@ impl Database {
                 user_id: params.user_id.clone(),
                 stripe_customer_id: params.stripe_customer_id.clone(),
                 has_overdue_invoices: params.has_overdue_invoices.clone(),
-                ..Default::default()
+                trial_started_at: params.trial_started_at.clone(),
+                created_at: ActiveValue::not_set(),
             })
             .exec(&*tx)
             .await?;

crates/collab/src/db/tables/billing_customer.rs

@@ -10,6 +10,7 @@ pub struct Model {
     pub user_id: UserId,
     pub stripe_customer_id: String,
     pub has_overdue_invoices: bool,
+    pub trial_started_at: Option<DateTime>,
     pub created_at: DateTime,
 }
 

crates/collab/src/llm/db/queries/subscription_usages.rs

@@ -1,8 +1,87 @@
-use crate::db::UserId;
+use chrono::Timelike;
+use time::PrimitiveDateTime;
+
+use crate::db::billing_subscription::SubscriptionKind;
+use crate::db::{UserId, billing_subscription};
 
 use super::*;
 
+fn convert_chrono_to_time(datetime: DateTimeUtc) -> anyhow::Result<PrimitiveDateTime> {
+    use chrono::{Datelike as _, Timelike as _};
+
+    let date = time::Date::from_calendar_date(
+        datetime.year(),
+        time::Month::try_from(datetime.month() as u8).unwrap(),
+        datetime.day() as u8,
+    )?;
+
+    let time = time::Time::from_hms_nano(
+        datetime.hour() as u8,
+        datetime.minute() as u8,
+        datetime.second() as u8,
+        datetime.nanosecond(),
+    )?;
+
+    Ok(PrimitiveDateTime::new(date, time))
+}
+
 impl LlmDatabase {
+    pub async fn create_subscription_usage(
+        &self,
+        user_id: UserId,
+        period_start_at: DateTimeUtc,
+        period_end_at: DateTimeUtc,
+        plan: SubscriptionKind,
+        model_requests: i32,
+        edit_predictions: i32,
+    ) -> Result<subscription_usage::Model> {
+        self.transaction(|tx| async move {
+            self.create_subscription_usage_in_tx(
+                user_id,
+                period_start_at,
+                period_end_at,
+                plan,
+                model_requests,
+                edit_predictions,
+                &tx,
+            )
+            .await
+        })
+        .await
+    }
+
+    async fn create_subscription_usage_in_tx(
+        &self,
+        user_id: UserId,
+        period_start_at: DateTimeUtc,
+        period_end_at: DateTimeUtc,
+        plan: SubscriptionKind,
+        model_requests: i32,
+        edit_predictions: i32,
+        tx: &DatabaseTransaction,
+    ) -> Result<subscription_usage::Model> {
+        // Clear out the nanoseconds so that these timestamps are comparable with Unix timestamps.
+        let period_start_at = period_start_at.with_nanosecond(0).unwrap();
+        let period_end_at = period_end_at.with_nanosecond(0).unwrap();
+
+        let period_start_at = convert_chrono_to_time(period_start_at)?;
+        let period_end_at = convert_chrono_to_time(period_end_at)?;
+
+        Ok(
+            subscription_usage::Entity::insert(subscription_usage::ActiveModel {
+                id: ActiveValue::not_set(),
+                user_id: ActiveValue::set(user_id),
+                period_start_at: ActiveValue::set(period_start_at),
+                period_end_at: ActiveValue::set(period_end_at),
+                plan: ActiveValue::set(plan),
+                model_requests: ActiveValue::set(model_requests),
+                edit_predictions: ActiveValue::set(edit_predictions),
+            })
+            .exec_with_returning(tx)
+            .await?,
+        )
+    }
+
     pub async fn get_subscription_usage_for_period(
         &self,
         user_id: UserId,
@@ -10,12 +89,77 @@ impl LlmDatabase {
         period_end_at: DateTimeUtc,
     ) -> Result<Option<subscription_usage::Model>> {
         self.transaction(|tx| async move {
-            Ok(subscription_usage::Entity::find()
-                .filter(subscription_usage::Column::UserId.eq(user_id))
-                .filter(subscription_usage::Column::PeriodStartAt.eq(period_start_at))
-                .filter(subscription_usage::Column::PeriodEndAt.eq(period_end_at))
-                .one(&*tx)
-                .await?)
+            self.get_subscription_usage_for_period_in_tx(
+                user_id,
+                period_start_at,
+                period_end_at,
+                &tx,
+            )
+            .await
+        })
+        .await
+    }
+
+    async fn get_subscription_usage_for_period_in_tx(
+        &self,
+        user_id: UserId,
+        period_start_at: DateTimeUtc,
+        period_end_at: DateTimeUtc,
+        tx: &DatabaseTransaction,
+    ) -> Result<Option<subscription_usage::Model>> {
+        Ok(subscription_usage::Entity::find()
+            .filter(subscription_usage::Column::UserId.eq(user_id))
+            .filter(subscription_usage::Column::PeriodStartAt.eq(period_start_at))
+            .filter(subscription_usage::Column::PeriodEndAt.eq(period_end_at))
+            .one(tx)
+            .await?)
+    }
+
+    pub async fn transfer_existing_subscription_usage(
+        &self,
+        user_id: UserId,
+        existing_subscription: &billing_subscription::Model,
+        new_subscription_kind: Option<SubscriptionKind>,
+        new_period_start_at: DateTimeUtc,
+        new_period_end_at: DateTimeUtc,
+    ) -> Result<Option<subscription_usage::Model>> {
+        self.transaction(|tx| async move {
+            match existing_subscription.kind {
+                Some(SubscriptionKind::ZedProTrial) => {
+                    let trial_period_start_at = existing_subscription
+                        .current_period_start_at()
+                        .ok_or_else(|| anyhow!("No trial subscription period start"))?;
+                    let trial_period_end_at = existing_subscription
+                        .current_period_end_at()
+                        .ok_or_else(|| anyhow!("No trial subscription period end"))?;
+
+                    let existing_usage = self
+                        .get_subscription_usage_for_period_in_tx(
+                            user_id,
+                            trial_period_start_at,
+                            trial_period_end_at,
+                            &tx,
+                        )
+                        .await?;
+                    if let Some(existing_usage) = existing_usage {
+                        return Ok(Some(
+                            self.create_subscription_usage_in_tx(
+                                user_id,
+                                new_period_start_at,
+                                new_period_end_at,
+                                new_subscription_kind.unwrap_or(existing_usage.plan),
+                                existing_usage.model_requests,
+                                existing_usage.edit_predictions,
+                                &tx,
+                            )
+                            .await?,
+                        ));
+                    }
+                }
+                _ => {}
+            }
+
+            Ok(None)
         })
         .await
     }

crates/collab/src/llm/db/tests.rs

@@ -1,4 +1,5 @@
 mod provider_tests;
+mod subscription_usage_tests;
 
 use gpui::BackgroundExecutor;
 use parking_lot::Mutex;

crates/collab/src/llm/db/tests/subscription_usage_tests.rs

@@ -0,0 +1,69 @@
+use chrono::{Duration, Utc};
+use pretty_assertions::assert_eq;
+
+use crate::db::billing_subscription::SubscriptionKind;
+use crate::db::{UserId, billing_subscription};
+use crate::llm::db::LlmDatabase;
+use crate::test_llm_db;
+
+test_llm_db!(
+    test_transfer_existing_subscription_usage,
+    test_transfer_existing_subscription_usage_postgres
+);
+
+async fn test_transfer_existing_subscription_usage(db: &mut LlmDatabase) {
+    let user_id = UserId(1);
+
+    let now = Utc::now();
+
+    let trial_period_start_at = now - Duration::days(14);
+    let trial_period_end_at = now;
+
+    let new_period_start_at = now;
+    let new_period_end_at = now + Duration::days(30);
+
+    let existing_subscription = billing_subscription::Model {
+        kind: Some(SubscriptionKind::ZedProTrial),
+        stripe_current_period_start: Some(trial_period_start_at.timestamp()),
+        stripe_current_period_end: Some(trial_period_end_at.timestamp()),
+        ..Default::default()
+    };
+
+    let existing_usage = db
+        .create_subscription_usage(
+            user_id,
+            trial_period_start_at,
+            trial_period_end_at,
+            SubscriptionKind::ZedProTrial,
+            25,
+            1_000,
+        )
+        .await
+        .unwrap();
+
+    let transferred_usage = db
+        .transfer_existing_subscription_usage(
+            user_id,
+            &existing_subscription,
+            Some(SubscriptionKind::ZedPro),
+            new_period_start_at,
+            new_period_end_at,
+        )
+        .await
+        .unwrap();
+
+    assert!(
+        transferred_usage.is_some(),
+        "subscription usage not transferred successfully"
+    );
+    let transferred_usage = transferred_usage.unwrap();
+
+    assert_eq!(
+        transferred_usage.model_requests,
+        existing_usage.model_requests
+    );
+    assert_eq!(
+        transferred_usage.edit_predictions,
+        existing_usage.edit_predictions
+    );
+}

crates/collab/src/llm/token.rs

@@ -1,4 +1,5 @@
 use crate::Cents;
+use crate::db::billing_subscription::SubscriptionKind;
 use crate::db::{billing_subscription, user};
 use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
 use crate::{Config, db::billing_preference};
@@ -32,6 +33,8 @@ pub struct LlmTokenClaims {
     pub plan: Plan,
     #[serde(default)]
     pub subscription_period: Option<(NaiveDateTime, NaiveDateTime)>,
+    #[serde(default)]
+    pub can_use_web_search_tool: bool,
 }
 
 const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
@@ -43,7 +46,6 @@ impl LlmTokenClaims {
         billing_preferences: Option<billing_preference::Model>,
         feature_flags: &Vec<String>,
         has_legacy_llm_subscription: bool,
-        plan: rpc::proto::Plan,
         subscription: Option<billing_subscription::Model>,
         system_id: Option<String>,
         config: &Config,
@@ -70,6 +72,7 @@ impl LlmTokenClaims {
             bypass_account_age_check: feature_flags
                 .iter()
                 .any(|flag| flag == "bypass-account-age-check"),
+            can_use_web_search_tool: feature_flags.iter().any(|flag| flag == "assistant2"),
             has_llm_subscription: has_legacy_llm_subscription,
             max_monthly_spend_in_cents: billing_preferences
                 .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
@@ -78,11 +81,14 @@ impl LlmTokenClaims {
             custom_llm_monthly_allowance_in_cents: user
                 .custom_llm_monthly_allowance_in_cents
                 .map(|allowance| allowance as u32),
-            plan: match plan {
-                rpc::proto::Plan::Free => Plan::Free,
-                rpc::proto::Plan::ZedPro => Plan::ZedPro,
-                rpc::proto::Plan::ZedProTrial => Plan::ZedProTrial,
-            },
+            plan: subscription
+                .as_ref()
+                .and_then(|subscription| subscription.kind)
+                .map_or(Plan::Free, |kind| match kind {
+                    SubscriptionKind::ZedFree => Plan::Free,
+                    SubscriptionKind::ZedPro => Plan::ZedPro,
+                    SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
+                }),
             subscription_period: maybe!({
                 let subscription = subscription?;
                 let period_start_at = subscription.current_period_start_at()?;

crates/collab/src/rpc.rs

@@ -4147,7 +4147,6 @@ async fn get_llm_api_token(
         billing_preferences,
         &flags,
         has_legacy_llm_subscription,
-        session.current_plan(&db).await?,
         billing_subscription,
         session.system_id.clone(),
         &session.app_state.config,

crates/collab/src/stripe_billing.rs

@@ -416,9 +416,16 @@ impl StripeBilling {
         let mut params = stripe::CreateCheckoutSession::new();
         params.subscription_data = Some(stripe::CreateCheckoutSessionSubscriptionData {
             trial_period_days: Some(14),
+            trial_settings: Some(stripe::CreateCheckoutSessionSubscriptionDataTrialSettings {
+                end_behavior: stripe::CreateCheckoutSessionSubscriptionDataTrialSettingsEndBehavior {
+                    missing_payment_method: stripe::CreateCheckoutSessionSubscriptionDataTrialSettingsEndBehaviorMissingPaymentMethod::Pause,
+                }
+            }),
             ..Default::default()
         });
         params.mode = Some(stripe::CheckoutSessionMode::Subscription);
+        params.payment_method_collection =
+            Some(stripe::CheckoutSessionPaymentMethodCollection::IfRequired);
         params.customer = Some(customer_id);
         params.client_reference_id = Some(github_login);
         params.line_items = Some(vec![stripe::CreateCheckoutSessionLineItems {

crates/collab/src/tests/debug_panel_tests.rs

@@ -1,7 +1,7 @@
 use call::ActiveCall;
-use dap::requests::{Initialize, Launch, StackTrace};
 use dap::DebugRequestType;
-use dap::{requests::SetBreakpoints, SourceBreakpoint};
+use dap::requests::{Initialize, Launch, StackTrace};
+use dap::{SourceBreakpoint, requests::SetBreakpoints};
 use debugger_ui::debugger_panel::DebugPanel;
 use debugger_ui::session::DebugSession;
 use editor::Editor;
@@ -13,7 +13,7 @@ use std::{
     path::Path,
     sync::atomic::{AtomicBool, Ordering},
 };
-use workspace::{dock::Panel, Workspace};
+use workspace::{Workspace, dock::Panel};
 
 use super::{TestClient, TestServer};
 

crates/collab/src/tests/following_tests.rs

@@ -6,17 +6,18 @@ use collab_ui::{
     channel_view::ChannelView,
     notifications::project_shared_notification::ProjectSharedNotification,
 };
-use editor::{Editor, ExcerptRange, MultiBuffer};
+use editor::{Editor, MultiBuffer, PathKey};
 use gpui::{
     AppContext as _, BackgroundExecutor, BorrowAppContext, Entity, SharedString, TestAppContext,
-    VisualTestContext, point,
+    VisualContext, VisualTestContext, point,
 };
 use language::Capability;
 use project::WorktreeSettings;
 use rpc::proto::PeerId;
 use serde_json::json;
 use settings::SettingsStore;
-use util::path;
+use text::{Point, ToPoint};
+use util::{path, test::sample_text};
 use workspace::{SplitDirection, Workspace, item::ItemHandle as _};
 
 use super::TestClient;
@@ -295,8 +296,20 @@ async fn test_basic_following(
                 .unwrap()
         });
         let mut result = MultiBuffer::new(Capability::ReadWrite);
-        result.push_excerpts(buffer_a1, [ExcerptRange::new(0..3)], cx);
-        result.push_excerpts(buffer_a2, [ExcerptRange::new(4..7)], cx);
+        result.set_excerpts_for_path(
+            PathKey::for_buffer(&buffer_a1, cx),
+            buffer_a1,
+            [Point::row_range(1..2)],
+            1,
+            cx,
+        );
+        result.set_excerpts_for_path(
+            PathKey::for_buffer(&buffer_a2, cx),
+            buffer_a2,
+            [Point::row_range(5..6)],
+            1,
+            cx,
+        );
         result
     });
     let multibuffer_editor_a = workspace_a.update_in(cx_a, |workspace, window, cx| {
@@ -2070,6 +2083,83 @@ async fn share_workspace(
         .await
 }
 
+#[gpui::test]
+async fn test_following_after_replacement(cx_a: &mut TestAppContext, cx_b: &mut TestAppContext) {
+    let (_server, client_a, client_b, channel) = TestServer::start2(cx_a, cx_b).await;
+
+    let (workspace, cx_a) = client_a.build_test_workspace(cx_a).await;
+    join_channel(channel, &client_a, cx_a).await.unwrap();
+    share_workspace(&workspace, cx_a).await.unwrap();
+    let buffer = workspace.update(cx_a, |workspace, cx| {
+        workspace.project().update(cx, |project, cx| {
+            project.create_local_buffer(&sample_text(26, 5, 'a'), None, cx)
+        })
+    });
+    let multibuffer = cx_a.new(|cx| {
+        let mut mb = MultiBuffer::new(Capability::ReadWrite);
+        mb.set_excerpts_for_path(
+            PathKey::for_buffer(&buffer, cx),
+            buffer.clone(),
+            [Point::row_range(1..1), Point::row_range(5..5)],
+            1,
+            cx,
+        );
+        mb
+    });
+    let snapshot = buffer.update(cx_a, |buffer, _| buffer.snapshot());
+    let editor: Entity<Editor> = cx_a.new_window_entity(|window, cx| {
+        Editor::for_multibuffer(
+            multibuffer.clone(),
+            Some(workspace.read(cx).project().clone()),
+            window,
+            cx,
+        )
+    });
+    workspace.update_in(cx_a, |workspace, window, cx| {
+        workspace.add_item_to_center(Box::new(editor.clone()) as _, window, cx)
+    });
+    editor.update_in(cx_a, |editor, window, cx| {
+        editor.change_selections(None, window, cx, |s| {
+            s.select_ranges([Point::row_range(4..4)]);
+        })
+    });
+    let positions = editor.update(cx_a, |editor, _| {
+        editor
+            .selections
+            .disjoint_anchor_ranges()
+            .map(|range| range.start.text_anchor.to_point(&snapshot))
+            .collect::<Vec<_>>()
+    });
+    multibuffer.update(cx_a, |multibuffer, cx| {
+        multibuffer.set_excerpts_for_path(
+            PathKey::for_buffer(&buffer, cx),
+            buffer,
+            [Point::row_range(1..5)],
+            1,
+            cx,
+        );
+    });
+
+    let (workspace_b, cx_b) = client_b.join_workspace(channel, cx_b).await;
+    cx_b.run_until_parked();
+    let editor_b = workspace_b
+        .update(cx_b, |workspace, cx| {
+            workspace
+                .active_item(cx)
+                .and_then(|item| item.downcast::<Editor>())
+        })
+        .unwrap();
+
+    let new_positions = editor_b.update(cx_b, |editor, _| {
+        editor
+            .selections
+            .disjoint_anchor_ranges()
+            .map(|range| range.start.text_anchor.to_point(&snapshot))
+            .collect::<Vec<_>>()
+    });
+    assert_eq!(positions, new_positions);
+}
+
 #[gpui::test]
 async fn test_following_to_channel_notes_other_workspace(
     cx_a: &mut TestAppContext,

crates/collab/src/tests/remote_editing_collaboration_tests.rs

@@ -2,11 +2,13 @@ use crate::tests::TestServer;
 use call::ActiveCall;
 use collections::{HashMap, HashSet};
 
+use debugger_ui::debugger_panel::DebugPanel;
 use extension::ExtensionHostProxy;
 use fs::{FakeFs, Fs as _, RemoveOptions};
 use futures::StreamExt as _;
 use gpui::{
     AppContext as _, BackgroundExecutor, SemanticVersion, TestAppContext, UpdateGlobal as _,
+    VisualContext,
 };
 use http_client::BlockedHttpClient;
 use language::{
@@ -24,6 +26,7 @@ use project::{
 };
 use remote::SshRemoteClient;
 use remote_server::{HeadlessAppState, HeadlessProject};
+use rpc::proto;
 use serde_json::json;
 use settings::SettingsStore;
 use std::{path::Path, sync::Arc};
@@ -576,3 +579,108 @@ async fn test_ssh_collaboration_formatting_with_prettier(
         "Prettier formatting was not applied to client buffer after host's request"
     );
 }
+
+#[gpui::test]
+async fn test_remote_server_debugger(cx_a: &mut TestAppContext, server_cx: &mut TestAppContext) {
+    cx_a.update(|cx| {
+        release_channel::init(SemanticVersion::default(), cx);
+        command_palette_hooks::init(cx);
+        if std::env::var("RUST_LOG").is_ok() {
+            env_logger::try_init().ok();
+        }
+    });
+    server_cx.update(|cx| {
+        release_channel::init(SemanticVersion::default(), cx);
+    });
+    let (opts, server_ssh) = SshRemoteClient::fake_server(cx_a, server_cx);
+    let remote_fs = FakeFs::new(server_cx.executor());
+    remote_fs
+        .insert_tree(
+            path!("/code"),
+            json!({
+                "lib.rs": "fn one() -> usize { 1 }"
+            }),
+        )
+        .await;
+
+    // User A connects to the remote project via SSH.
+    server_cx.update(HeadlessProject::init);
+    let remote_http_client = Arc::new(BlockedHttpClient);
+    let node = NodeRuntime::unavailable();
+    let languages = Arc::new(LanguageRegistry::new(server_cx.executor()));
+    let _headless_project = server_cx.new(|cx| {
+        client::init_settings(cx);
+        HeadlessProject::new(
+            HeadlessAppState {
+                session: server_ssh,
+                fs: remote_fs.clone(),
+                http_client: remote_http_client,
+                node_runtime: node,
+                languages,
+                extension_host_proxy: Arc::new(ExtensionHostProxy::new()),
+            },
+            cx,
+        )
+    });
+
+    let client_ssh = SshRemoteClient::fake_client(opts, cx_a).await;
+    let mut server = TestServer::start(server_cx.executor()).await;
+    let client_a = server.create_client(cx_a, "user_a").await;
+    cx_a.update(|cx| {
+        debugger_ui::init(cx);
+        command_palette_hooks::init(cx);
+    });
+    let (project_a, _) = client_a
+        .build_ssh_project(path!("/code"), client_ssh.clone(), cx_a)
+        .await;
+
+    let (workspace, cx_a) = client_a.build_workspace(&project_a, cx_a);
+
+    let debugger_panel = workspace
+        .update_in(cx_a, |_workspace, window, cx| {
+            cx.spawn_in(window, DebugPanel::load)
+        })
+        .await
+        .unwrap();
+
+    workspace.update_in(cx_a, |workspace, window, cx| {
+        workspace.add_panel(debugger_panel, window, cx);
+    });
+
+    cx_a.run_until_parked();
+    let debug_panel = workspace
+        .update(cx_a, |workspace, cx| workspace.panel::<DebugPanel>(cx))
+        .unwrap();
+
+    let workspace_window = cx_a
+        .window_handle()
+        .downcast::<workspace::Workspace>()
+        .unwrap();
+
+    let session = debugger_ui::tests::start_debug_session(&workspace_window, cx_a, |_| {}).unwrap();
+    cx_a.run_until_parked();
+    debug_panel.update(cx_a, |debug_panel, cx| {
+        assert_eq!(
+            debug_panel.active_session().unwrap().read(cx).session(cx),
+            session
+        )
+    });
+
+    session.update(cx_a, |session, _| {
+        assert_eq!(session.binary().command, "ssh");
+    });
+
+    let shutdown_session = workspace.update(cx_a, |workspace, cx| {
+        workspace.project().update(cx, |project, cx| {
+            project.dap_store().update(cx, |dap_store, cx| {
+                dap_store.shutdown_session(session.read(cx).session_id(), cx)
+            })
+        })
+    });
+
+    client_ssh.update(cx_a, |a, _| {
+        a.shutdown_processes(Some(proto::ShutdownRemoteServer {}))
+    });
+
+    shutdown_session.await.unwrap();
+}

crates/context_server/src/context_server_tool.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use anyhow::{Result, anyhow, bail};
 use assistant_tool::{ActionLog, Tool, ToolResult, ToolSource};
-use gpui::{App, Entity, Task};
+use gpui::{AnyWindowHandle, App, Entity, Task};
 use icons::IconName;
 use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
@@ -77,6 +77,7 @@ impl Tool for ContextServerTool {
         _messages: &[LanguageModelRequestMessage],
         _project: Entity<Project>,
         _action_log: Entity<ActionLog>,
+        _window: Option<AnyWindowHandle>,
         cx: &mut App,
     ) -> ToolResult {
         if let Some(server) = self.server_manager.read(cx).get_server(&self.server_id) {

crates/dap/src/transport.rs

@@ -22,7 +22,7 @@ use std::{
     time::Duration,
 };
 use task::TcpArgumentsTemplate;
-use util::ResultExt as _;
+use util::{ResultExt as _, TryFutureExt};
 
 use crate::{adapters::DebugAdapterBinary, debugger_settings::DebuggerSettings};
 
@@ -126,6 +126,7 @@ pub(crate) struct TransportDelegate {
     pending_requests: Requests,
     transport: Transport,
     server_tx: Arc<Mutex<Option<Sender<Message>>>>,
+    _tasks: Vec<gpui::Task<Option<()>>>,
 }
 
 impl TransportDelegate {
@@ -140,6 +141,7 @@ impl TransportDelegate {
             log_handlers: Default::default(),
             current_requests: Default::default(),
             pending_requests: Default::default(),
+            _tasks: Default::default(),
         };
         let messages = this.start_handlers(transport_pipes, cx).await?;
         Ok((messages, this))
@@ -166,35 +168,43 @@ impl TransportDelegate {
 
         cx.update(|cx| {
             if let Some(stdout) = params.stdout.take() {
-                cx.background_executor()
-                    .spawn(Self::handle_adapter_log(stdout, log_handler.clone()))
-                    .detach_and_log_err(cx);
+                self._tasks.push(
+                    cx.background_executor()
+                        .spawn(Self::handle_adapter_log(stdout, log_handler.clone()).log_err()),
+                );
             }
 
-            cx.background_executor()
-                .spawn(Self::handle_output(
-                    params.output,
-                    client_tx,
-                    self.pending_requests.clone(),
-                    log_handler.clone(),
-                ))
-                .detach_and_log_err(cx);
+            self._tasks.push(
+                cx.background_executor().spawn(
+                    Self::handle_output(
+                        params.output,
+                        client_tx,
+                        self.pending_requests.clone(),
+                        log_handler.clone(),
+                    )
+                    .log_err(),
+                ),
+            );
 
             if let Some(stderr) = params.stderr.take() {
-                cx.background_executor()
-                    .spawn(Self::handle_error(stderr, self.log_handlers.clone()))
-                    .detach_and_log_err(cx);
+                self._tasks.push(
+                    cx.background_executor()
+                        .spawn(Self::handle_error(stderr, self.log_handlers.clone()).log_err()),
+                );
             }
 
-            cx.background_executor()
-                .spawn(Self::handle_input(
-                    params.input,
-                    client_rx,
-                    self.current_requests.clone(),
-                    self.pending_requests.clone(),
-                    log_handler.clone(),
-                ))
-                .detach_and_log_err(cx);
+            self._tasks.push(
+                cx.background_executor().spawn(
+                    Self::handle_input(
+                        params.input,
+                        client_rx,
+                        self.current_requests.clone(),
+                        self.pending_requests.clone(),
+                        log_handler.clone(),
+                    )
+                    .log_err(),
+                ),
+            );
         })?;
 
         {
@@ -367,6 +377,7 @@ impl TransportDelegate {
     where
         Stderr: AsyncRead + Unpin + Send + 'static,
     {
+        log::debug!("Handle error started");
         let mut buffer = String::new();
 
         let mut reader = BufReader::new(stderr);

crates/debugger_tools/Cargo.toml

@@ -12,6 +12,9 @@ workspace = true
 path = "src/debugger_tools.rs"
 doctest = false
 
+[features]
+test-support = []
+
 [dependencies]
 anyhow.workspace = true
 dap.workspace = true

crates/debugger_tools/src/dap_log.rs

@@ -41,7 +41,7 @@ struct DapLogView {
     _subscriptions: Vec<Subscription>,
 }
 
-struct LogStore {
+pub struct LogStore {
     projects: HashMap<WeakEntity<Project>, ProjectState>,
     debug_clients: HashMap<SessionId, DebugAdapterState>,
     rpc_tx: UnboundedSender<(SessionId, IoKind, String)>,
@@ -101,7 +101,7 @@ impl DebugAdapterState {
 }
 
 impl LogStore {
-    fn new(cx: &Context<Self>) -> Self {
+    pub fn new(cx: &Context<Self>) -> Self {
         let (rpc_tx, mut rpc_rx) = unbounded::<(SessionId, IoKind, String)>();
         cx.spawn(async move |this, cx| {
             while let Some((client_id, io_kind, message)) = rpc_rx.next().await {
@@ -845,3 +845,29 @@ impl EventEmitter<Event> for LogStore {}
 impl EventEmitter<Event> for DapLogView {}
 impl EventEmitter<EditorEvent> for DapLogView {}
 impl EventEmitter<SearchEvent> for DapLogView {}
+
+#[cfg(any(test, feature = "test-support"))]
+impl LogStore {
+    pub fn contained_session_ids(&self) -> Vec<SessionId> {
+        self.debug_clients.keys().cloned().collect()
+    }
+
+    pub fn rpc_messages_for_session_id(&self, session_id: SessionId) -> Vec<String> {
+        self.debug_clients
+            .get(&session_id)
+            .expect("This session should exist if a test is calling")
+            .rpc_messages
+            .messages
+            .clone()
+            .into()
+    }
+
+    pub fn log_messages_for_session_id(&self, session_id: SessionId) -> Vec<String> {
+        self.debug_clients
+            .get(&session_id)
+            .expect("This session should exist if a test is calling")
+            .log_messages
+            .clone()
+            .into()
+    }
+}

crates/debugger_ui/Cargo.toml

@@ -20,6 +20,9 @@ test-support = [
     "project/test-support",
     "util/test-support",
     "workspace/test-support",
+    "env_logger",
+    "unindent",
+    "debugger_tools"
 ]
 
 [dependencies]
@@ -37,6 +40,7 @@ gpui.workspace = true
 language.workspace = true
 log.workspace = true
 menu.workspace = true
+parking_lot.workspace = true
 picker.workspace = true
 pretty_assertions.workspace = true
 project.workspace = true
@@ -53,9 +57,13 @@ ui.workspace = true
 util.workspace = true
 workspace.workspace = true
 workspace-hack.workspace = true
+env_logger = { workspace = true, optional = true }
+debugger_tools = { workspace = true, optional = true }
+unindent = { workspace = true, optional = true }
 
 [dev-dependencies]
 dap = { workspace = true, features = ["test-support"] }
+debugger_tools = { workspace = true, features = ["test-support"] }
 editor = { workspace = true, features = ["test-support"] }
 env_logger.workspace = true
 gpui = { workspace = true, features = ["test-support"] }

crates/debugger_ui/src/attach_modal.rs

@@ -1,7 +1,7 @@
 use dap::DebugRequest;
 use fuzzy::{StringMatch, StringMatchCandidate};
-use gpui::Subscription;
 use gpui::{DismissEvent, Entity, EventEmitter, Focusable, Render};
+use gpui::{Subscription, WeakEntity};
 use picker::{Picker, PickerDelegate};
 
 use std::sync::Arc;
@@ -9,7 +9,9 @@ use sysinfo::System;
 use ui::{Context, Tooltip, prelude::*};
 use ui::{ListItem, ListItemSpacing};
 use util::debug_panic;
-use workspace::ModalView;
+use workspace::{ModalView, Workspace};
+
+use crate::debugger_panel::DebugPanel;
 
 #[derive(Debug, Clone)]
 pub(super) struct Candidate {
@@ -22,19 +24,19 @@ pub(crate) struct AttachModalDelegate {
     selected_index: usize,
     matches: Vec<StringMatch>,
     placeholder_text: Arc<str>,
-    project: Entity<project::Project>,
+    workspace: WeakEntity<Workspace>,
     pub(crate) debug_config: task::DebugTaskDefinition,
     candidates: Arc<[Candidate]>,
 }
 
 impl AttachModalDelegate {
     fn new(
-        project: Entity<project::Project>,
+        workspace: Entity<Workspace>,
         debug_config: task::DebugTaskDefinition,
         candidates: Arc<[Candidate]>,
     ) -> Self {
         Self {
-            project,
+            workspace: workspace.downgrade(),
             debug_config,
             candidates,
             selected_index: 0,
@@ -51,7 +53,7 @@ pub struct AttachModal {
 
 impl AttachModal {
     pub fn new(
-        project: Entity<project::Project>,
+        workspace: Entity<Workspace>,
         debug_config: task::DebugTaskDefinition,
         modal: bool,
         window: &mut Window,
@@ -75,11 +77,11 @@ impl AttachModal {
             .collect();
         processes.sort_by_key(|k| k.name.clone());
         let processes = processes.into_iter().collect();
-        Self::with_processes(project, debug_config, processes, modal, window, cx)
+        Self::with_processes(workspace, debug_config, processes, modal, window, cx)
     }
 
     pub(super) fn with_processes(
-        project: Entity<project::Project>,
+        workspace: Entity<Workspace>,
         debug_config: task::DebugTaskDefinition,
         processes: Arc<[Candidate]>,
         modal: bool,
@@ -88,7 +90,7 @@ impl AttachModal {
     ) -> Self {
         let picker = cx.new(|cx| {
             Picker::uniform_list(
-                AttachModalDelegate::new(project, debug_config, processes),
+                AttachModalDelegate::new(workspace, debug_config, processes),
                 window,
                 cx,
             )
@@ -202,7 +204,7 @@ impl PickerDelegate for AttachModalDelegate {
         })
     }
 
-    fn confirm(&mut self, _: bool, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
+    fn confirm(&mut self, _: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
         let candidate = self
             .matches
             .get(self.selected_index())
@@ -225,14 +227,17 @@ impl PickerDelegate for AttachModalDelegate {
             }
         }
 
-        let config = self.debug_config.clone();
-        self.project
-            .update(cx, |project, cx| {
-                let ret = project.start_debug_session(config, cx);
-                ret
-            })
-            .detach_and_log_err(cx);
-
+        let definition = self.debug_config.clone();
+        let panel = self
+            .workspace
+            .update(cx, |workspace, cx| workspace.panel::<DebugPanel>(cx))
+            .ok()
+            .flatten();
+        if let Some(panel) = panel {
+            panel.update(cx, |panel, cx| {
+                panel.start_session(definition, window, cx);
+            });
+        }
         cx.emit(DismissEvent);
     }
 

crates/debugger_ui/src/debugger_panel.rs

@@ -6,6 +6,7 @@ use crate::{new_session_modal::NewSessionModal, session::DebugSession};
 use anyhow::{Result, anyhow};
 use collections::HashMap;
 use command_palette_hooks::CommandPaletteFilter;
+use dap::StartDebuggingRequestArguments;
 use dap::{
     ContinuedEvent, LoadedSourceEvent, ModuleEvent, OutputEvent, StoppedEvent, ThreadEvent,
     client::SessionId, debugger_settings::DebuggerSettings,
@@ -17,6 +18,7 @@ use gpui::{
     actions, anchored, deferred,
 };
 
+use project::debugger::session::{Session, SessionStateEvent};
 use project::{
     Project,
     debugger::{
@@ -30,10 +32,9 @@ use settings::Settings;
 use std::any::TypeId;
 use std::path::Path;
 use std::sync::Arc;
-use task::DebugTaskDefinition;
+use task::{DebugTaskDefinition, DebugTaskTemplate};
 use terminal_view::terminal_panel::TerminalPanel;
 use ui::{ContextMenu, Divider, DropdownMenu, Tooltip, prelude::*};
-use util::debug_panic;
 use workspace::{
     Workspace,
     dock::{DockPosition, Panel, PanelEvent},
@@ -63,7 +64,7 @@ pub struct DebugPanel {
     active_session: Option<Entity<DebugSession>>,
     /// This represents the last debug definition that was created in the new session modal
     pub(crate) past_debug_definition: Option<DebugTaskDefinition>,
-    project: WeakEntity<Project>,
+    project: Entity<Project>,
     workspace: WeakEntity<Workspace>,
     focus_handle: FocusHandle,
     context_menu: Option<(Entity<ContextMenu>, Point<Pixels>, Subscription)>,
@@ -97,10 +98,10 @@ impl DebugPanel {
                             window,
                             |panel, _, event: &tasks_ui::ShowAttachModal, window, cx| {
                                 panel.workspace.update(cx, |workspace, cx| {
-                                    let project = workspace.project().clone();
+                                    let workspace_handle = cx.entity().clone();
                                     workspace.toggle_modal(window, cx, |window, cx| {
                                         crate::attach_modal::AttachModal::new(
-                                            project,
+                                            workspace_handle,
                                             event.debug_config.clone(),
                                             true,
                                             window,
@@ -127,7 +128,7 @@ impl DebugPanel {
                 _subscriptions,
                 past_debug_definition: None,
                 focus_handle: cx.focus_handle(),
-                project: project.downgrade(),
+                project,
                 workspace: workspace.weak_handle(),
                 context_menu: None,
             };
@@ -219,7 +220,7 @@ impl DebugPanel {
 
     pub fn load(
         workspace: WeakEntity<Workspace>,
-        cx: AsyncWindowContext,
+        cx: &mut AsyncWindowContext,
     ) -> Task<Result<Entity<Self>>> {
         cx.spawn(async move |cx| {
             workspace.update_in(cx, |workspace, window, cx| {
@@ -245,114 +246,226 @@ impl DebugPanel {
                     });
                 })
                 .detach();
+                workspace.set_debugger_provider(DebuggerProvider(debug_panel.clone()));
 
                 debug_panel
             })
         })
     }
 
+    pub fn start_session(
+        &mut self,
+        definition: DebugTaskDefinition,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let task_contexts = self
+            .workspace
+            .update(cx, |workspace, cx| {
+                tasks_ui::task_contexts(workspace, window, cx)
+            })
+            .ok();
+        let dap_store = self.project.read(cx).dap_store().clone();
+
+        cx.spawn_in(window, async move |this, cx| {
+            let task_context = if let Some(task) = task_contexts {
+                task.await
+                    .active_worktree_context
+                    .map_or(task::TaskContext::default(), |context| context.1)
+            } else {
+                task::TaskContext::default()
+            };
+
+            let (session, task) = dap_store.update(cx, |dap_store, cx| {
+                let template = DebugTaskTemplate {
+                    locator: None,
+                    definition: definition.clone(),
+                };
+                let session = if let Some(debug_config) = template
+                    .to_zed_format()
+                    .resolve_task("debug_task", &task_context)
+                    .and_then(|resolved_task| resolved_task.resolved_debug_adapter_config())
+                {
+                    dap_store.new_session(debug_config.definition, None, cx)
+                } else {
+                    dap_store.new_session(definition.clone(), None, cx)
+                };
+
+                (session.clone(), dap_store.boot_session(session, cx))
+            })?;
+
+            match task.await {
+                Err(e) => {
+                    this.update(cx, |this, cx| {
+                        this.workspace
+                            .update(cx, |workspace, cx| {
+                                workspace.show_error(&e, cx);
+                            })
+                            .ok();
+                    })
+                    .ok();
+
+                    session
+                        .update(cx, |session, cx| session.shutdown(cx))?
+                        .await;
+                }
+                Ok(_) => Self::register_session(this, session, cx).await?,
+            }
+
+            anyhow::Ok(())
+        })
+        .detach_and_log_err(cx);
+    }
+
+    async fn register_session(
+        this: WeakEntity<Self>,
+        session: Entity<Session>,
+        cx: &mut AsyncWindowContext,
+    ) -> Result<()> {
+        let adapter_name = session.update(cx, |session, _| session.adapter_name())?;
+        this.update_in(cx, |_, window, cx| {
+            cx.subscribe_in(
+                &session,
+                window,
+                move |_, session, event: &SessionStateEvent, window, cx| match event {
+                    SessionStateEvent::Restart => {
+                        let mut curr_session = session.clone();
+                        while let Some(parent_session) = curr_session
+                            .read_with(cx, |session, _| session.parent_session().cloned())
+                        {
+                            curr_session = parent_session;
+                        }
+
+                        let definition = curr_session.update(cx, |session, _| session.definition());
+                        let task = curr_session.update(cx, |session, cx| session.shutdown(cx));
+
+                        let definition = definition.clone();
+                        cx.spawn_in(window, async move |this, cx| {
+                            task.await;
+
+                            this.update_in(cx, |this, window, cx| {
+                                this.start_session(definition, window, cx)
+                            })
+                        })
+                        .detach_and_log_err(cx);
+                    }
+                    _ => {}
+                },
+            )
+            .detach();
+        })
+        .ok();
+
+        let serialized_layout = persistence::get_serialized_pane_layout(adapter_name).await;
+
+        let workspace = this.update_in(cx, |this, window, cx| {
+            this.sessions.retain(|session| {
+                session
+                    .read(cx)
+                    .mode()
+                    .as_running()
+                    .map_or(false, |running_state| {
+                        !running_state.read(cx).session().read(cx).is_terminated()
+                    })
+            });
+
+            let session_item = DebugSession::running(
+                this.project.clone(),
+                this.workspace.clone(),
+                session,
+                cx.weak_entity(),
+                serialized_layout,
+                window,
+                cx,
+            );
+
+            if let Some(running) = session_item.read(cx).mode().as_running().cloned() {
+                // We might want to make this an event subscription and only notify when a new thread is selected
+                // This is used to filter the command menu correctly
+                cx.observe(&running, |_, _, cx| cx.notify()).detach();
+            }
+
+            this.sessions.push(session_item.clone());
+            this.activate_session(session_item, window, cx);
+            this.workspace.clone()
+        })?;
+
+        workspace.update_in(cx, |workspace, window, cx| {
+            workspace.focus_panel::<Self>(window, cx);
+        })?;
+        Ok(())
+    }
+
+    pub fn start_child_session(
+        &mut self,
+        request: &StartDebuggingRequestArguments,
+        parent_session: Entity<Session>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let Some(worktree) = parent_session.read(cx).worktree() else {
+            log::error!("Attempted to start a child session from non local debug session");
+            return;
+        };
+
+        let dap_store_handle = self.project.read(cx).dap_store().clone();
+        let breakpoint_store = self.project.read(cx).breakpoint_store();
+        let definition = parent_session.read(cx).definition().clone();
+        let mut binary = parent_session.read(cx).binary().clone();
+        binary.request_args = request.clone();
+
+        cx.spawn_in(window, async move |this, cx| {
+            let (session, task) = dap_store_handle.update(cx, |dap_store, cx| {
+                let session =
+                    dap_store.new_session(definition.clone(), Some(parent_session.clone()), cx);
+
+                let task = session.update(cx, |session, cx| {
+                    session.boot(
+                        binary,
+                        worktree,
+                        breakpoint_store,
+                        dap_store_handle.downgrade(),
+                        cx,
+                    )
+                });
+                (session, task)
+            })?;
+
+            match task.await {
+                Err(e) => {
+                    this.update(cx, |this, cx| {
+                        this.workspace
+                            .update(cx, |workspace, cx| {
+                                workspace.show_error(&e, cx);
+                            })
+                            .ok();
+                    })
+                    .ok();
+
+                    session
+                        .update(cx, |session, cx| session.shutdown(cx))?
+                        .await;
+                }
+                Ok(_) => Self::register_session(this, session, cx).await?,
+            }
+
+            anyhow::Ok(())
+        })
+        .detach_and_log_err(cx);
+    }
+
     pub fn active_session(&self) -> Option<Entity<DebugSession>> {
         self.active_session.clone()
     }
 
-    pub fn debug_panel_items_by_client(
-        &self,
-        client_id: &SessionId,
-        cx: &Context<Self>,
-    ) -> Vec<Entity<DebugSession>> {
-        self.sessions
-            .iter()
-            .filter(|item| item.read(cx).session_id(cx) == *client_id)
-            .map(|item| item.clone())
-            .collect()
-    }
-
-    pub fn debug_panel_item_by_client(
-        &self,
-        client_id: SessionId,
-        cx: &mut Context<Self>,
-    ) -> Option<Entity<DebugSession>> {
-        self.sessions
-            .iter()
-            .find(|item| {
-                let item = item.read(cx);
-
-                item.session_id(cx) == client_id
-            })
-            .cloned()
-    }
-
     fn handle_dap_store_event(
         &mut self,
-        dap_store: &Entity<DapStore>,
+        _dap_store: &Entity<DapStore>,
         event: &dap_store::DapStoreEvent,
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
         match event {
-            dap_store::DapStoreEvent::DebugSessionInitialized(session_id) => {
-                let Some(session) = dap_store.read(cx).session_by_id(session_id) else {
-                    return log::error!(
-                        "Couldn't get session with id: {session_id:?} from DebugClientStarted event"
-                    );
-                };
-
-                let adapter_name = session.read(cx).adapter_name();
-
-                let session_id = *session_id;
-                cx.spawn_in(window, async move |this, cx| {
-                    let serialized_layout =
-                        persistence::get_serialized_pane_layout(adapter_name).await;
-
-                    this.update_in(cx, |this, window, cx| {
-                        let Some(project) = this.project.upgrade() else {
-                            return log::error!(
-                                "Debug Panel out lived it's weak reference to Project"
-                            );
-                        };
-
-                        if this
-                            .sessions
-                            .iter()
-                            .any(|item| item.read(cx).session_id(cx) == session_id)
-                        {
-                            // We already have an item for this session.
-                            debug_panic!("We should never reuse session ids");
-                            return;
-                        }
-
-                        this.sessions.retain(|session| {
-                            session
-                                .read(cx)
-                                .mode()
-                                .as_running()
-                                .map_or(false, |running_state| {
-                                    !running_state.read(cx).session().read(cx).is_terminated()
-                                })
-                        });
-
-                        let session_item = DebugSession::running(
-                            project,
-                            this.workspace.clone(),
-                            session,
-                            cx.weak_entity(),
-                            serialized_layout,
-                            window,
-                            cx,
-                        );
-
-                        if let Some(running) = session_item.read(cx).mode().as_running().cloned() {
-                            // We might want to make this an event subscription and only notify when a new thread is selected
-                            // This is used to filter the command menu correctly
-                            cx.observe(&running, |_, _, cx| cx.notify()).detach();
-                        }
-
-                        this.sessions.push(session_item.clone());
-                        this.activate_session(session_item, window, cx);
-                    })
-                })
-                .detach();
-            }
             dap_store::DapStoreEvent::RunInTerminal {
                 title,
                 cwd,
@@ -374,6 +487,12 @@ impl DebugPanel {
                 )
                 .detach_and_log_err(cx);
             }
+            dap_store::DapStoreEvent::SpawnChildSession {
+                request,
+                parent_session,
+            } => {
+                self.start_child_session(request, parent_session.clone(), window, cx);
+            }
             _ => {}
         }
     }
@@ -408,7 +527,7 @@ impl DebugPanel {
                         cwd,
                         title,
                     },
-                    task::RevealStrategy::Always,
+                    task::RevealStrategy::Never,
                     window,
                     cx,
                 );
@@ -468,8 +587,6 @@ impl DebugPanel {
                 let session = this.dap_store().read(cx).session_by_id(session_id);
                 session.map(|session| !session.read(cx).is_terminated())
             })
-            .ok()
-            .flatten()
             .unwrap_or_default();
 
         cx.spawn_in(window, async move |this, cx| {
@@ -893,7 +1010,6 @@ impl DebugPanel {
 
 impl EventEmitter<PanelEvent> for DebugPanel {}
 impl EventEmitter<DebugPanelEvent> for DebugPanel {}
-impl EventEmitter<project::Event> for DebugPanel {}
 
 impl Focusable for DebugPanel {
     fn focus_handle(&self, _: &App) -> FocusHandle {
@@ -1039,3 +1155,15 @@ impl Render for DebugPanel {
             .into_any()
     }
 }
+
+struct DebuggerProvider(Entity<DebugPanel>);
+
+impl workspace::DebuggerProvider for DebuggerProvider {
+    fn start_session(&self, definition: DebugTaskDefinition, window: &mut Window, cx: &mut App) {
+        self.0.update(cx, |_, cx| {
+            cx.defer_in(window, |this, window, cx| {
+                this.start_session(definition, window, cx);
+            })
+        })
+    }
+}

crates/debugger_ui/src/debugger_ui.rs

@@ -16,7 +16,7 @@ mod new_session_modal;
 mod persistence;
 pub(crate) mod session;
 
-#[cfg(test)]
+#[cfg(any(test, feature = "test-support"))]
 pub mod tests;
 
 actions!(

crates/debugger_ui/src/new_session_modal.rs

@@ -4,14 +4,12 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use anyhow::{Result, anyhow};
 use dap::{DapRegistry, DebugRequest};
 use editor::{Editor, EditorElement, EditorStyle};
 use gpui::{
     App, AppContext, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, Render, TextStyle,
     WeakEntity,
 };
-use project::Project;
 use settings::Settings;
 use task::{DebugTaskDefinition, DebugTaskTemplate, LaunchRequest};
 use theme::ThemeSettings;
@@ -21,7 +19,6 @@ use ui::{
     LabelCommon as _, ParentElement, RenderOnce, SharedString, Styled, StyledExt, ToggleButton,
     ToggleState, Toggleable, Window, div, h_flex, relative, rems, v_flex,
 };
-use util::ResultExt;
 use workspace::{ModalView, Workspace};
 
 use crate::{attach_modal::AttachModal, debugger_panel::DebugPanel};
@@ -88,11 +85,11 @@ impl NewSessionModal {
         }
     }
 
-    fn debug_config(&self, cx: &App) -> Option<DebugTaskDefinition> {
+    fn debug_config(&self, cx: &App, debugger: &str) -> DebugTaskDefinition {
         let request = self.mode.debug_task(cx);
-        Some(DebugTaskDefinition {
-            adapter: self.debugger.clone()?.to_string(),
-            label: suggested_label(&request, self.debugger.as_deref()?),
+        DebugTaskDefinition {
+            adapter: debugger.to_owned(),
+            label: suggested_label(&request, debugger),
             request,
             initialize_args: self.initialize_args.clone(),
             tcp_connection: None,
@@ -100,26 +97,26 @@ impl NewSessionModal {
                 ToggleState::Selected => Some(true),
                 _ => None,
             },
-        })
+        }
     }
 
-    fn start_new_session(&self, window: &mut Window, cx: &mut Context<Self>) -> Result<()> {
-        let workspace = self.workspace.clone();
-        let config = self
-            .debug_config(cx)
-            .ok_or_else(|| anyhow!("Failed to create a debug config"))?;
+    fn start_new_session(&self, window: &mut Window, cx: &mut Context<Self>) {
+        let Some(debugger) = self.debugger.as_ref() else {
+            // todo: show in UI.
+            log::error!("No debugger selected");
+            return;
+        };
+        let config = self.debug_config(cx, debugger);
+        let debug_panel = self.debug_panel.clone();
 
-        let _ = self.debug_panel.update(cx, |panel, _| {
-            panel.past_debug_definition = Some(config.clone());
-        });
-
-        let task_contexts = workspace
+        let task_contexts = self
+            .workspace
             .update(cx, |workspace, cx| {
                 tasks_ui::task_contexts(workspace, window, cx)
             })
             .ok();
 
-        cx.spawn(async move |this, cx| {
+        cx.spawn_in(window, async move |this, cx| {
             let task_context = if let Some(task) = task_contexts {
                 task.await
                     .active_worktree_context
@@ -127,9 +124,8 @@ impl NewSessionModal {
             } else {
                 task::TaskContext::default()
             };
-            let project = workspace.update(cx, |workspace, _| workspace.project().clone())?;
 
-            let task = project.update(cx, |this, cx| {
+            debug_panel.update_in(cx, |debug_panel, window, cx| {
                 let template = DebugTaskTemplate {
                     locator: None,
                     definition: config.clone(),
@@ -139,23 +135,18 @@ impl NewSessionModal {
                     .resolve_task("debug_task", &task_context)
                     .and_then(|resolved_task| resolved_task.resolved_debug_adapter_config())
                 {
-                    this.start_debug_session(debug_config.definition, cx)
+                    debug_panel.start_session(debug_config.definition, window, cx)
                 } else {
-                    this.start_debug_session(config, cx)
+                    debug_panel.start_session(config, window, cx)
                 }
             })?;
-            let spawn_result = task.await;
-            if spawn_result.is_ok() {
-                this.update(cx, |_, cx| {
-                    cx.emit(DismissEvent);
-                })
-                .ok();
-            }
-            spawn_result?;
+            this.update(cx, |_, cx| {
+                cx.emit(DismissEvent);
+            })
+            .ok();
             anyhow::Result::<_, anyhow::Error>::Ok(())
         })
         .detach_and_log_err(cx);
-        Ok(())
     }
 
     fn update_attach_picker(
@@ -249,15 +240,12 @@ impl NewSessionModal {
                                     );
                                 }
                                 DebugRequest::Attach(_) => {
-                                    let Ok(project) = this
-                                        .workspace
-                                        .read_with(cx, |this, _| this.project().clone())
-                                    else {
+                                    let Some(workspace) = this.workspace.upgrade() else {
                                         return;
                                     };
                                     this.mode = NewSessionMode::attach(
                                         this.debugger.clone(),
-                                        project,
+                                        workspace,
                                         window,
                                         cx,
                                     );
@@ -357,7 +345,7 @@ struct AttachMode {
 impl AttachMode {
     fn new(
         debugger: Option<SharedString>,
-        project: Entity<Project>,
+        workspace: Entity<Workspace>,
         window: &mut Window,
         cx: &mut Context<NewSessionModal>,
     ) -> Entity<Self> {
@@ -370,7 +358,7 @@ impl AttachMode {
             stop_on_entry: Some(false),
         };
         let attach_picker = cx.new(|cx| {
-            let modal = AttachModal::new(project, debug_definition.clone(), false, window, cx);
+            let modal = AttachModal::new(workspace, debug_definition.clone(), false, window, cx);
             window.focus(&modal.focus_handle(cx));
 
             modal
@@ -470,11 +458,11 @@ impl RenderOnce for NewSessionMode {
 impl NewSessionMode {
     fn attach(
         debugger: Option<SharedString>,
-        project: Entity<Project>,
+        workspace: Entity<Workspace>,
         window: &mut Window,
         cx: &mut Context<NewSessionModal>,
     ) -> Self {
-        Self::Attach(AttachMode::new(debugger, project, window, cx))
+        Self::Attach(AttachMode::new(debugger, workspace, window, cx))
     }
     fn launch(
         past_launch_config: Option<LaunchRequest>,
@@ -569,15 +557,12 @@ impl Render for NewSessionModal {
                                 .toggle_state(matches!(self.mode, NewSessionMode::Attach(_)))
                                 .style(ui::ButtonStyle::Subtle)
                                 .on_click(cx.listener(|this, _, window, cx| {
-                                    let Ok(project) = this
-                                        .workspace
-                                        .read_with(cx, |this, _| this.project().clone())
-                                    else {
+                                    let Some(workspace) = this.workspace.upgrade() else {
                                         return;
                                     };
                                     this.mode = NewSessionMode::attach(
                                         this.debugger.clone(),
-                                        project,
+                                        workspace,
                                         window,
                                         cx,
                                     );
@@ -631,7 +616,7 @@ impl Render for NewSessionModal {
                             .child(
                                 Button::new("debugger-spawn", "Start")
                                     .on_click(cx.listener(|this, _, window, cx| {
-                                        this.start_new_session(window, cx).log_err();
+                                        this.start_new_session(window, cx);
                                     }))
                                     .disabled(self.debugger.is_none()),
                             ),

crates/debugger_ui/src/session.rs

@@ -88,6 +88,12 @@ impl DebugSession {
         }
     }
 
+    pub fn session(&self, cx: &App) -> Entity<Session> {
+        match &self.mode {
+            DebugSessionState::Running(entity) => entity.read(cx).session().clone(),
+        }
+    }
+
     pub(crate) fn shutdown(&mut self, cx: &mut Context<Self>) {
         match &self.mode {
             DebugSessionState::Running(state) => state.update(cx, |state, cx| state.shutdown(cx)),
@@ -115,13 +121,7 @@ impl DebugSession {
         };
 
         self.label
-            .get_or_init(|| {
-                session
-                    .read(cx)
-                    .as_local()
-                    .expect("Remote Debug Sessions are not implemented yet")
-                    .label()
-            })
+            .get_or_init(|| session.read(cx).label())
             .to_owned()
     }
 

crates/debugger_ui/src/session/running.rs

@@ -411,13 +411,26 @@ impl RunningState {
                             .log_err();
 
                         if let Some(thread_id) = thread_id {
-                            this.select_thread(*thread_id, cx);
+                            this.select_thread(*thread_id, window, cx);
                         }
                     }
                     SessionEvent::Threads => {
                         let threads = this.session.update(cx, |this, cx| this.threads(cx));
-                        this.select_current_thread(&threads, cx);
+                        this.select_current_thread(&threads, window, cx);
                     }
+                    SessionEvent::CapabilitiesLoaded => {
+                        let capabilities = this.capabilities(cx);
+                        if !capabilities.supports_modules_request.unwrap_or(false) {
+                            this.remove_pane_item(DebuggerPaneItem::Modules, window, cx);
+                        }
+                        if !capabilities
+                            .supports_loaded_sources_request
+                            .unwrap_or(false)
+                        {
+                            this.remove_pane_item(DebuggerPaneItem::LoadedSources, window, cx);
+                        }
+                    }
+
                     _ => {}
                 }
                 cx.notify()
@@ -447,35 +460,14 @@ impl RunningState {
             workspace::PaneGroup::with_root(root)
         } else {
             pane_close_subscriptions.clear();
-            let module_list = if session
-                .read(cx)
-                .capabilities()
-                .supports_modules_request
-                .unwrap_or(false)
-            {
-                Some(&module_list)
-            } else {
-                None
-            };
-
-            let loaded_source_list = if session
-                .read(cx)
-                .capabilities()
-                .supports_loaded_sources_request
-                .unwrap_or(false)
-            {
-                Some(&loaded_source_list)
-            } else {
-                None
-            };
 
             let root = Self::default_pane_layout(
                 project,
                 &workspace,
                 &stack_frame_list,
                 &variable_list,
-                module_list,
-                loaded_source_list,
+                &module_list,
+                &loaded_source_list,
                 &console,
                 &breakpoint_list,
                 &mut pane_close_subscriptions,
@@ -512,11 +504,6 @@ impl RunningState {
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
-        debug_assert!(
-            item_kind.is_supported(self.session.read(cx).capabilities()),
-            "We should only allow removing supported item kinds"
-        );
-
         if let Some((pane, item_id)) = self.panes.panes().iter().find_map(|pane| {
             Some(pane).zip(
                 pane.read(cx)
@@ -744,6 +731,7 @@ impl RunningState {
     pub fn select_current_thread(
         &mut self,
         threads: &Vec<(Thread, ThreadStatus)>,
+        window: &mut Window,
         cx: &mut Context<Self>,
     ) {
         let selected_thread = self
@@ -756,7 +744,7 @@ impl RunningState {
         };
 
         if Some(ThreadId(selected_thread.id)) != self.thread_id {
-            self.select_thread(ThreadId(selected_thread.id), cx);
+            self.select_thread(ThreadId(selected_thread.id), window, cx);
         }
     }
 
@@ -769,7 +757,7 @@ impl RunningState {
             .map(|id| self.session().read(cx).thread_status(id))
     }
 
-    fn select_thread(&mut self, thread_id: ThreadId, cx: &mut Context<Self>) {
+    fn select_thread(&mut self, thread_id: ThreadId, window: &mut Window, cx: &mut Context<Self>) {
         if self.thread_id.is_some_and(|id| id == thread_id) {
             return;
         }
@@ -777,8 +765,7 @@ impl RunningState {
         self.thread_id = Some(thread_id);
 
         self.stack_frame_list
-            .update(cx, |list, cx| list.refresh(cx));
-        cx.notify();
+            .update(cx, |list, cx| list.schedule_refresh(true, window, cx));
     }
 
     pub fn continue_thread(&mut self, cx: &mut Context<Self>) {
@@ -930,9 +917,9 @@ impl RunningState {
                 for (thread, _) in threads {
                     let state = state.clone();
                     let thread_id = thread.id;
-                    this = this.entry(thread.name, None, move |_, cx| {
+                    this = this.entry(thread.name, None, move |window, cx| {
                         state.update(cx, |state, cx| {
-                            state.select_thread(ThreadId(thread_id), cx);
+                            state.select_thread(ThreadId(thread_id), window, cx);
                         });
                     });
                 }
@@ -946,8 +933,8 @@ impl RunningState {
         workspace: &WeakEntity<Workspace>,
         stack_frame_list: &Entity<StackFrameList>,
         variable_list: &Entity<VariableList>,
-        module_list: Option<&Entity<ModuleList>>,
-        loaded_source_list: Option<&Entity<LoadedSourceList>>,
+        module_list: &Entity<ModuleList>,
+        loaded_source_list: &Entity<LoadedSourceList>,
         console: &Entity<Console>,
         breakpoints: &Entity<BreakpointList>,
         subscriptions: &mut HashMap<EntityId, Subscription>,
@@ -1003,41 +990,36 @@ impl RunningState {
                 window,
                 cx,
             );
-            if let Some(module_list) = module_list {
-                this.add_item(
-                    Box::new(SubView::new(
-                        module_list.focus_handle(cx),
-                        module_list.clone().into(),
-                        DebuggerPaneItem::Modules,
-                        None,
-                        cx,
-                    )),
-                    false,
-                    false,
+            this.add_item(
+                Box::new(SubView::new(
+                    module_list.focus_handle(cx),
+                    module_list.clone().into(),
+                    DebuggerPaneItem::Modules,
                     None,
-                    window,
                     cx,
-                );
-                this.activate_item(0, false, false, window, cx);
-            }
+                )),
+                false,
+                false,
+                None,
+                window,
+                cx,
+            );
 
-            if let Some(loaded_source_list) = loaded_source_list {
-                this.add_item(
-                    Box::new(SubView::new(
-                        loaded_source_list.focus_handle(cx),
-                        loaded_source_list.clone().into(),
-                        DebuggerPaneItem::LoadedSources,
-                        None,
-                        cx,
-                    )),
-                    false,
-                    false,
+            this.add_item(
+                Box::new(SubView::new(
+                    loaded_source_list.focus_handle(cx),
+                    loaded_source_list.clone().into(),
+                    DebuggerPaneItem::LoadedSources,
                     None,
-                    window,
                     cx,
-                );
-                this.activate_item(1, false, false, window, cx);
-            }
+                )),
+                false,
+                false,
+                None,
+                window,
+                cx,
+            );
+            this.activate_item(0, false, false, window, cx);
         });
 
         let rightmost_pane = new_debugger_pane(workspace.clone(), project.clone(), window, cx);

crates/debugger_ui/src/session/running/stack_frame_list.rs

@@ -1,5 +1,6 @@
 use std::path::Path;
 use std::sync::Arc;
+use std::time::Duration;
 
 use anyhow::{Result, anyhow};
 use dap::StackFrameId;
@@ -28,11 +29,11 @@ pub struct StackFrameList {
     _subscription: Subscription,
     session: Entity<Session>,
     state: WeakEntity<RunningState>,
-    invalidate: bool,
     entries: Vec<StackFrameEntry>,
     workspace: WeakEntity<Workspace>,
     selected_stack_frame_id: Option<StackFrameId>,
     scrollbar_state: ScrollbarState,
+    _refresh_task: Task<()>,
 }
 
 #[allow(clippy::large_enum_variant)]
@@ -68,14 +69,17 @@ impl StackFrameList {
         );
 
         let _subscription =
-            cx.subscribe_in(&session, window, |this, _, event, _, cx| match event {
-                SessionEvent::Stopped(_) | SessionEvent::StackTrace | SessionEvent::Threads => {
-                    this.refresh(cx);
+            cx.subscribe_in(&session, window, |this, _, event, window, cx| match event {
+                SessionEvent::Threads => {
+                    this.schedule_refresh(false, window, cx);
+                }
+                SessionEvent::Stopped(..) | SessionEvent::StackTrace => {
+                    this.schedule_refresh(true, window, cx);
                 }
                 _ => {}
             });
 
-        Self {
+        let mut this = Self {
             scrollbar_state: ScrollbarState::new(list.clone()),
             list,
             session,
@@ -83,10 +87,12 @@ impl StackFrameList {
             focus_handle,
             state,
             _subscription,
-            invalidate: true,
             entries: Default::default(),
             selected_stack_frame_id: None,
-        }
+            _refresh_task: Task::ready(()),
+        };
+        this.schedule_refresh(true, window, cx);
+        this
     }
 
     #[cfg(test)]
@@ -136,10 +142,32 @@ impl StackFrameList {
         self.selected_stack_frame_id
     }
 
-    pub(super) fn refresh(&mut self, cx: &mut Context<Self>) {
-        self.invalidate = true;
-        self.entries.clear();
-        cx.notify();
+    pub(super) fn schedule_refresh(
+        &mut self,
+        select_first: bool,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        const REFRESH_DEBOUNCE: Duration = Duration::from_millis(20);
+
+        self._refresh_task = cx.spawn_in(window, async move |this, cx| {
+            let debounce = this
+                .update(cx, |this, cx| {
+                    let new_stack_frames = this.stack_frames(cx);
+                    new_stack_frames.is_empty() && !this.entries.is_empty()
+                })
+                .ok()
+                .unwrap_or_default();
+
+            if debounce {
+                cx.background_executor().timer(REFRESH_DEBOUNCE).await;
+            }
+            this.update_in(cx, |this, window, cx| {
+                this.build_entries(select_first, window, cx);
+                cx.notify();
+            })
+            .ok();
+        })
     }
 
     pub fn build_entries(
@@ -515,13 +543,7 @@ impl StackFrameList {
 }
 
 impl Render for StackFrameList {
-    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
-        if self.invalidate {
-            self.build_entries(self.entries.is_empty(), window, cx);
-            self.invalidate = false;
-            cx.notify();
-        }
-
+    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
         div()
             .size_full()
             .p_1()

crates/debugger_ui/src/tests.rs

@@ -1,16 +1,29 @@
+use std::sync::Arc;
+
+use anyhow::{Result, anyhow};
+use dap::{DebugRequest, client::DebugAdapterClient};
 use gpui::{Entity, TestAppContext, WindowHandle};
-use project::Project;
+use project::{Project, debugger::session::Session};
 use settings::SettingsStore;
+use task::DebugTaskDefinition;
 use terminal_view::terminal_panel::TerminalPanel;
 use workspace::Workspace;
 
 use crate::{debugger_panel::DebugPanel, session::DebugSession};
 
+#[cfg(test)]
 mod attach_modal;
+#[cfg(test)]
 mod console;
+#[cfg(test)]
+mod dap_logger;
+#[cfg(test)]
 mod debugger_panel;
+#[cfg(test)]
 mod module_list;
+#[cfg(test)]
 mod stack_frame_list;
+#[cfg(test)]
 mod variable_list;
 
 pub fn init_test(cx: &mut gpui::TestAppContext) {
@@ -42,7 +55,7 @@ pub async fn init_test_workspace(
     let debugger_panel = workspace_handle
         .update(cx, |_, window, cx| {
             cx.spawn_in(window, async move |this, cx| {
-                DebugPanel::load(this, cx.clone()).await
+                DebugPanel::load(this, cx).await
             })
         })
         .unwrap()
@@ -82,3 +95,46 @@ pub fn active_debug_session_panel(
         })
         .unwrap()
 }
+
+pub fn start_debug_session_with<T: Fn(&Arc<DebugAdapterClient>) + 'static>(
+    workspace: &WindowHandle<Workspace>,
+    cx: &mut gpui::TestAppContext,
+    config: DebugTaskDefinition,
+    configure: T,
+) -> Result<Entity<Session>> {
+    let _subscription = project::debugger::test::intercept_debug_sessions(cx, configure);
+    workspace.update(cx, |workspace, window, cx| {
+        workspace.start_debug_session(config, window, cx)
+    })?;
+    cx.run_until_parked();
+    let session = workspace.read_with(cx, |workspace, cx| {
+        workspace
+            .panel::<DebugPanel>(cx)
+            .and_then(|panel| panel.read(cx).active_session())
+            .and_then(|session| session.read(cx).mode().as_running().cloned())
+            .map(|running| running.read(cx).session().clone())
+            .ok_or_else(|| anyhow!("Failed to get active session"))
+    })??;
+
+    Ok(session)
+}
+
+pub fn start_debug_session<T: Fn(&Arc<DebugAdapterClient>) + 'static>(
+    workspace: &WindowHandle<Workspace>,
+    cx: &mut gpui::TestAppContext,
+    configure: T,
+) -> Result<Entity<Session>> {
+    start_debug_session_with(
+        workspace,
+        cx,
+        DebugTaskDefinition {
+            adapter: "fake-adapter".to_string(),
+            request: DebugRequest::Launch(Default::default()),
+            label: "test".to_string(),
+            initialize_args: None,
+            tcp_connection: None,
+            stop_on_entry: None,
+        },
+        configure,
+    )
+}

crates/debugger_ui/src/tests/attach_modal.rs

@@ -1,4 +1,4 @@
-use crate::{attach_modal::Candidate, *};
+use crate::{attach_modal::Candidate, tests::start_debug_session_with, *};
 use attach_modal::AttachModal;
 use dap::{FakeAdapter, client::SessionId};
 use gpui::{BackgroundExecutor, TestAppContext, VisualTestContext};
@@ -26,8 +26,8 @@ async fn test_direct_attach_to_process(executor: BackgroundExecutor, cx: &mut Te
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session_with(
-        &project,
+    let session = start_debug_session_with(
+        &workspace,
         cx,
         DebugTaskDefinition {
             adapter: "fake-adapter".to_string(),
@@ -47,7 +47,6 @@ async fn test_direct_attach_to_process(executor: BackgroundExecutor, cx: &mut Te
             });
         },
     )
-    .await
     .unwrap();
 
     cx.run_until_parked();
@@ -99,9 +98,10 @@ async fn test_show_attach_modal_and_select_process(
         });
     let attach_modal = workspace
         .update(cx, |workspace, window, cx| {
+            let workspace_handle = cx.entity();
             workspace.toggle_modal(window, cx, |window, cx| {
                 AttachModal::with_processes(
-                    project.clone(),
+                    workspace_handle,
                     DebugTaskDefinition {
                         adapter: FakeAdapter::ADAPTER_NAME.into(),
                         request: dap::DebugRequest::Attach(AttachRequest::default()),

crates/debugger_ui/src/tests/console.rs

@@ -1,4 +1,7 @@
-use crate::{tests::active_debug_session_panel, *};
+use crate::{
+    tests::{active_debug_session_panel, start_debug_session},
+    *,
+};
 use dap::requests::StackTrace;
 use gpui::{BackgroundExecutor, TestAppContext, VisualTestContext};
 use project::{FakeFs, Project};
@@ -28,9 +31,7 @@ async fn test_handle_output_event(executor: BackgroundExecutor, cx: &mut TestApp
         })
         .unwrap();
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<StackTrace, _>(move |_, _| {

crates/debugger_ui/src/tests/dap_logger.rs

@@ -0,0 +1,118 @@
+use crate::tests::{init_test, init_test_workspace, start_debug_session};
+use dap::requests::{StackTrace, Threads};
+use debugger_tools::LogStore;
+use gpui::{BackgroundExecutor, TestAppContext, VisualTestContext};
+use project::Project;
+use serde_json::json;
+use std::cell::OnceCell;
+
+#[gpui::test]
+async fn test_dap_logger_captures_all_session_rpc_messages(
+    executor: BackgroundExecutor,
+    cx: &mut TestAppContext,
+) {
+    let log_store_cell = std::rc::Rc::new(OnceCell::new());
+
+    cx.update(|cx| {
+        let log_store_cell = log_store_cell.clone();
+        cx.observe_new::<LogStore>(move |_, _, cx| {
+            log_store_cell.set(cx.entity()).unwrap();
+        })
+        .detach();
+        debugger_tools::init(cx);
+    });
+    init_test(cx);
+
+    let log_store = log_store_cell.get().unwrap().clone();
+
+    // Create a filesystem with a simple project
+    let fs = project::FakeFs::new(executor.clone());
+    fs.insert_tree(
+        "/project",
+        json!({
+            "main.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"
+        }),
+    )
+    .await;
+
+    assert!(
+        log_store.read_with(cx, |log_store, _| log_store
+            .contained_session_ids()
+            .is_empty()),
+        "log_store shouldn't contain any session IDs before any sessions were created"
+    );
+
+    let project = Project::test(fs, ["/project".as_ref()], cx).await;
+
+    let workspace = init_test_workspace(&project, cx).await;
+    let cx = &mut VisualTestContext::from_window(*workspace, cx);
+
+    // Start a debug session
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
+    let session_id = session.read_with(cx, |session, _| session.session_id());
+    let client = session.update(cx, |session, _| session.adapter_client().unwrap());
+
+    assert_eq!(
+        log_store.read_with(cx, |log_store, _| log_store.contained_session_ids().len()),
+        1,
+    );
+
+    assert!(
+        log_store.read_with(cx, |log_store, _| log_store
+            .contained_session_ids()
+            .contains(&session_id)),
+        "log_store should contain the session IDs of the started session"
+    );
+
+    assert!(
+        !log_store.read_with(cx, |log_store, _| log_store
+            .rpc_messages_for_session_id(session_id)
+            .is_empty()),
+        "We should have the initialization sequence in the log store"
+    );
+
+    // Set up basic responses for common requests
+    client.on_request::<Threads, _>(move |_, _| {
+        Ok(dap::ThreadsResponse {
+            threads: vec![dap::Thread {
+                id: 1,
+                name: "Thread 1".into(),
+            }],
+        })
+    });
+
+    client.on_request::<StackTrace, _>(move |_, _| {
+        Ok(dap::StackTraceResponse {
+            stack_frames: Vec::default(),
+            total_frames: None,
+        })
+    });
+
+    // Run until all pending tasks are executed
+    cx.run_until_parked();
+
+    // Simulate a stopped event to generate more DAP messages
+    client
+        .fake_event(dap::messages::Events::Stopped(dap::StoppedEvent {
+            reason: dap::StoppedEventReason::Pause,
+            description: None,
+            thread_id: Some(1),
+            preserve_focus_hint: None,
+            text: None,
+            all_threads_stopped: None,
+            hit_breakpoint_ids: None,
+        }))
+        .await;
+
+    cx.run_until_parked();
+
+    // Shutdown the debug session
+    let shutdown_session = project.update(cx, |project, cx| {
+        project.dap_store().update(cx, |dap_store, cx| {
+            dap_store.shutdown_session(session.read(cx).session_id(), cx)
+        })
+    });
+
+    shutdown_session.await.unwrap();
+    cx.run_until_parked();
+}

crates/debugger_ui/src/tests/debugger_panel.rs

@@ -1,4 +1,4 @@
-use crate::*;
+use crate::{tests::start_debug_session, *};
 use dap::{
     ErrorResponse, Message, RunInTerminalRequestArguments, SourceBreakpoint,
     StartDebuggingRequestArguments, StartDebuggingRequestArgumentsRequest,
@@ -48,9 +48,7 @@ async fn test_basic_show_debug_panel(executor: BackgroundExecutor, cx: &mut Test
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<Threads, _>(move |_, _| {
@@ -187,9 +185,7 @@ async fn test_we_can_only_have_one_panel_per_debug_session(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<Threads, _>(move |_, _| {
@@ -354,9 +350,7 @@ async fn test_handle_successful_run_in_terminal_reverse_request(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client
@@ -419,6 +413,86 @@ async fn test_handle_successful_run_in_terminal_reverse_request(
     shutdown_session.await.unwrap();
 }
 
+#[gpui::test]
+async fn test_handle_start_debugging_request(
+    executor: BackgroundExecutor,
+    cx: &mut TestAppContext,
+) {
+    init_test(cx);
+
+    let fs = FakeFs::new(executor.clone());
+
+    fs.insert_tree(
+        "/project",
+        json!({
+            "main.rs": "First line\nSecond line\nThird line\nFourth line",
+        }),
+    )
+    .await;
+
+    let project = Project::test(fs, ["/project".as_ref()], cx).await;
+    let workspace = init_test_workspace(&project, cx).await;
+    let cx = &mut VisualTestContext::from_window(*workspace, cx);
+
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
+    let client = session.update(cx, |session, _| session.adapter_client().unwrap());
+
+    let fake_config = json!({"one": "two"});
+    let launched_with = Arc::new(parking_lot::Mutex::new(None));
+
+    let _subscription = project::debugger::test::intercept_debug_sessions(cx, {
+        let launched_with = launched_with.clone();
+        move |client| {
+            let launched_with = launched_with.clone();
+            client.on_request::<dap::requests::Launch, _>(move |_, args| {
+                launched_with.lock().replace(args.raw);
+                Ok(())
+            });
+            client.on_request::<dap::requests::Attach, _>(move |_, _| {
+                assert!(false, "should not get attach request");
+                Ok(())
+            });
+        }
+    });
+
+    client
+        .fake_reverse_request::<StartDebugging>(StartDebuggingRequestArguments {
+            request: StartDebuggingRequestArgumentsRequest::Launch,
+            configuration: fake_config.clone(),
+        })
+        .await;
+
+    cx.run_until_parked();
+
+    workspace
+        .update(cx, |workspace, _window, cx| {
+            let debug_panel = workspace.panel::<DebugPanel>(cx).unwrap();
+            let active_session = debug_panel
+                .read(cx)
+                .active_session()
+                .unwrap()
+                .read(cx)
+                .session(cx);
+            let parent_session = active_session.read(cx).parent_session().unwrap();
+
+            assert_eq!(
+                active_session.read(cx).definition(),
+                parent_session.read(cx).definition()
+            );
+        })
+        .unwrap();
+
+    assert_eq!(&fake_config, launched_with.lock().as_ref().unwrap());
+
+    let shutdown_session = project.update(cx, |project, cx| {
+        project.dap_store().update(cx, |dap_store, cx| {
+            dap_store.shutdown_session(session.read(cx).session_id(), cx)
+        })
+    });
+
+    shutdown_session.await.unwrap();
+}
+
 // // covers that we always send a response back, if something when wrong,
 // // while spawning the terminal
 #[gpui::test]
@@ -444,9 +518,7 @@ async fn test_handle_error_run_in_terminal_reverse_request(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client
@@ -522,9 +594,7 @@ async fn test_handle_start_debugging_reverse_request(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<dap::requests::Threads, _>(move |_, _| {
@@ -629,9 +699,7 @@ async fn test_shutdown_children_when_parent_session_shutdown(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let parent_session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let parent_session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = parent_session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<dap::requests::Threads, _>(move |_, _| {
@@ -737,9 +805,7 @@ async fn test_shutdown_parent_session_if_all_children_are_shutdown(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let parent_session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let parent_session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = parent_session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_response::<StartDebugging, _>(move |_| {}).await;
@@ -858,7 +924,7 @@ async fn test_debug_panel_item_thread_status_reset_on_failure(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |client| {
+    let session = start_debug_session(&workspace, cx, |client| {
         client.on_request::<dap::requests::Initialize, _>(move |_, _| {
             Ok(dap::Capabilities {
                 supports_step_back: Some(true),
@@ -866,7 +932,6 @@ async fn test_debug_panel_item_thread_status_reset_on_failure(
             })
         });
     })
-    .await
     .unwrap();
 
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
@@ -1073,9 +1138,7 @@ async fn test_send_breakpoints_when_editor_has_been_saved(
         .update(cx, |_, _, cx| worktree.read(cx).id())
         .unwrap();
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     let buffer = project
@@ -1290,9 +1353,7 @@ async fn test_unsetting_breakpoints_on_clear_breakpoint_action(
         editor.toggle_breakpoint(&actions::ToggleBreakpoint, window, cx);
     });
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     let called_set_breakpoints = Arc::new(AtomicBool::new(false));
@@ -1358,7 +1419,7 @@ async fn test_debug_session_is_shutdown_when_attach_and_launch_request_fails(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let task = project::debugger::test::start_debug_session(&project, cx, |client| {
+    start_debug_session(&workspace, cx, |client| {
         client.on_request::<dap::requests::Initialize, _>(|_, _| {
             Err(ErrorResponse {
                 error: Some(Message {
@@ -1372,12 +1433,8 @@ async fn test_debug_session_is_shutdown_when_attach_and_launch_request_fails(
                 }),
             })
         });
-    });
-
-    assert!(
-        task.await.is_err(),
-        "Session should failed to start if launch request fails"
-    );
+    })
+    .ok();
 
     cx.run_until_parked();
 

crates/debugger_ui/src/tests/module_list.rs

@@ -1,16 +1,13 @@
 use crate::{
     debugger_panel::DebugPanel,
-    tests::{active_debug_session_panel, init_test, init_test_workspace},
+    tests::{active_debug_session_panel, init_test, init_test_workspace, start_debug_session},
 };
 use dap::{
     StoppedEvent,
     requests::{Initialize, Modules},
 };
 use gpui::{BackgroundExecutor, TestAppContext, VisualTestContext};
-use project::{
-    FakeFs, Project,
-    debugger::{self},
-};
+use project::{FakeFs, Project};
 use std::sync::{
     Arc,
     atomic::{AtomicBool, AtomicI32, Ordering},
@@ -31,7 +28,7 @@ async fn test_module_list(executor: BackgroundExecutor, cx: &mut TestAppContext)
         .unwrap();
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |client| {
+    let session = start_debug_session(&workspace, cx, |client| {
         client.on_request::<Initialize, _>(move |_, _| {
             Ok(dap::Capabilities {
                 supports_modules_request: Some(true),
@@ -39,7 +36,6 @@ async fn test_module_list(executor: BackgroundExecutor, cx: &mut TestAppContext)
             })
         });
     })
-    .await
     .unwrap();
 
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());

crates/debugger_ui/src/tests/stack_frame_list.rs

@@ -1,7 +1,7 @@
 use crate::{
     debugger_panel::DebugPanel,
     session::running::stack_frame_list::StackFrameEntry,
-    tests::{active_debug_session_panel, init_test, init_test_workspace},
+    tests::{active_debug_session_panel, init_test, init_test_workspace, start_debug_session},
 };
 use dap::{
     StackFrame,
@@ -9,7 +9,7 @@ use dap::{
 };
 use editor::{Editor, ToPoint as _};
 use gpui::{BackgroundExecutor, TestAppContext, VisualTestContext};
-use project::{FakeFs, Project, debugger};
+use project::{FakeFs, Project};
 use serde_json::json;
 use std::sync::Arc;
 use unindent::Unindent as _;
@@ -50,9 +50,7 @@ async fn test_fetch_initial_stack_frames_and_go_to_stack_frame(
     let project = Project::test(fs, [path!("/project").as_ref()], cx).await;
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
     client.on_request::<Scopes, _>(move |_, _| Ok(dap::ScopesResponse { scopes: vec![] }));
 
@@ -154,7 +152,7 @@ async fn test_fetch_initial_stack_frames_and_go_to_stack_frame(
     cx.run_until_parked();
 
     // select first thread
-    active_debug_session_panel(workspace, cx).update_in(cx, |session, _, cx| {
+    active_debug_session_panel(workspace, cx).update_in(cx, |session, window, cx| {
         session
             .mode()
             .as_running()
@@ -164,6 +162,7 @@ async fn test_fetch_initial_stack_frames_and_go_to_stack_frame(
                     &running_state
                         .session()
                         .update(cx, |session, cx| session.threads(cx)),
+                    window,
                     cx,
                 );
             });
@@ -229,9 +228,7 @@ async fn test_select_stack_frame(executor: BackgroundExecutor, cx: &mut TestAppC
     });
 
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<Threads, _>(move |_, _| {
@@ -334,7 +331,7 @@ async fn test_select_stack_frame(executor: BackgroundExecutor, cx: &mut TestAppC
     cx.run_until_parked();
 
     // select first thread
-    active_debug_session_panel(workspace, cx).update_in(cx, |session, _, cx| {
+    active_debug_session_panel(workspace, cx).update_in(cx, |session, window, cx| {
         session
             .mode()
             .as_running()
@@ -344,6 +341,7 @@ async fn test_select_stack_frame(executor: BackgroundExecutor, cx: &mut TestAppC
                     &running_state
                         .session()
                         .update(cx, |session, cx| session.threads(cx)),
+                    window,
                     cx,
                 );
             });
@@ -495,9 +493,7 @@ async fn test_collapsed_entries(executor: BackgroundExecutor, cx: &mut TestAppCo
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<Threads, _>(move |_, _| {
@@ -710,7 +706,7 @@ async fn test_collapsed_entries(executor: BackgroundExecutor, cx: &mut TestAppCo
     cx.run_until_parked();
 
     // select first thread
-    active_debug_session_panel(workspace, cx).update_in(cx, |session, _, cx| {
+    active_debug_session_panel(workspace, cx).update_in(cx, |session, window, cx| {
         session
             .mode()
             .as_running()
@@ -720,6 +716,7 @@ async fn test_collapsed_entries(executor: BackgroundExecutor, cx: &mut TestAppCo
                     &running_state
                         .session()
                         .update(cx, |session, cx| session.threads(cx)),
+                    window,
                     cx,
                 );
             });

crates/debugger_ui/src/tests/variable_list.rs

@@ -6,7 +6,7 @@ use std::sync::{
 use crate::{
     DebugPanel,
     session::running::variable_list::{CollapseSelectedEntry, ExpandSelectedEntry},
-    tests::{active_debug_session_panel, init_test, init_test_workspace},
+    tests::{active_debug_session_panel, init_test, init_test_workspace, start_debug_session},
 };
 use collections::HashMap;
 use dap::{
@@ -15,7 +15,7 @@ use dap::{
 };
 use gpui::{BackgroundExecutor, TestAppContext, VisualTestContext};
 use menu::{SelectFirst, SelectNext, SelectPrevious};
-use project::{FakeFs, Project, debugger};
+use project::{FakeFs, Project};
 use serde_json::json;
 use unindent::Unindent as _;
 use util::path;
@@ -54,9 +54,7 @@ async fn test_basic_fetch_initial_scope_and_variables(
         })
         .unwrap();
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<dap::requests::Threads, _>(move |_, _| {
@@ -266,9 +264,7 @@ async fn test_fetch_variables_for_multiple_scopes(
         .unwrap();
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<dap::requests::Threads, _>(move |_, _| {
@@ -528,9 +524,7 @@ async fn test_keyboard_navigation(executor: BackgroundExecutor, cx: &mut TestApp
         })
         .unwrap();
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<dap::requests::Threads, _>(move |_, _| {
@@ -1313,9 +1307,7 @@ async fn test_variable_list_only_sends_requests_when_rendering(
     let workspace = init_test_workspace(&project, cx).await;
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<dap::requests::Threads, _>(move |_, _| {
@@ -1560,9 +1552,7 @@ async fn test_it_fetches_scopes_variables_when_you_select_a_stack_frame(
         .unwrap();
     let cx = &mut VisualTestContext::from_window(*workspace, cx);
 
-    let session = debugger::test::start_debug_session(&project, cx, |_| {})
-        .await
-        .unwrap();
+    let session = start_debug_session(&workspace, cx, |_| {}).unwrap();
     let client = session.update(cx, |session, _| session.adapter_client().unwrap());
 
     client.on_request::<dap::requests::Threads, _>(move |_, _| {