Centralize project context provided to the assistant (#11471)

This PR restructures the way that tools and attachments add information
about the current project to a conversation with the assistant. Rather
than each tool call or attachment generating a new tool or system
message containing information about the project, they can all
collectively mutate a new type called a `ProjectContext`, which stores
all of the project data that should be sent to the assistant. That data
is then formatted in a single place, and passed to the assistant in one
system message.

This prevents multiple tools/attachments from including redundant
context.

Release Notes:

- N/A

---------

Co-authored-by: Kyle <kylek@zed.dev>

crates/assistant_tooling/Cargo.toml

@@ -13,10 +13,18 @@ path = "src/assistant_tooling.rs"
 
 [dependencies]
 anyhow.workspace = true
+collections.workspace = true
+futures.workspace = true
 gpui.workspace = true
+project.workspace = true
 schemars.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+sum_tree.workspace = true
+util.workspace = true
 
 [dev-dependencies]
 gpui = { workspace = true, features = ["test-support"] }
+project = { workspace = true, features = ["test-support"] }
+settings = { workspace = true, features = ["test-support"] }
+unindent.workspace = true

crates/assistant_tooling/src/assistant_tooling.rs

@@ -1,5 +1,9 @@
-pub mod registry;
-pub mod tool;
+mod attachment_registry;
+mod project_context;
+mod tool_registry;
 
-pub use crate::registry::ToolRegistry;
-pub use crate::tool::{LanguageModelTool, ToolFunctionCall, ToolFunctionDefinition};
+pub use attachment_registry::{AttachmentRegistry, LanguageModelAttachment, UserAttachment};
+pub use project_context::ProjectContext;
+pub use tool_registry::{
+    LanguageModelTool, ToolFunctionCall, ToolFunctionDefinition, ToolOutput, ToolRegistry,
+};

crates/assistant_tooling/src/attachment_registry.rs

@@ -0,0 +1,148 @@
+use crate::{ProjectContext, ToolOutput};
+use anyhow::{anyhow, Result};
+use collections::HashMap;
+use futures::future::join_all;
+use gpui::{AnyView, Render, Task, View, WindowContext};
+use std::{
+    any::TypeId,
+    sync::{
+        atomic::{AtomicBool, Ordering::SeqCst},
+        Arc,
+    },
+};
+use util::ResultExt as _;
+
+pub struct AttachmentRegistry {
+    registered_attachments: HashMap<TypeId, RegisteredAttachment>,
+}
+
+pub trait LanguageModelAttachment {
+    type Output: 'static;
+    type View: Render + ToolOutput;
+
+    fn run(&self, cx: &mut WindowContext) -> Task<Result<Self::Output>>;
+
+    fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View>;
+}
+
+/// A collected attachment from running an attachment tool
+pub struct UserAttachment {
+    pub view: AnyView,
+    generate_fn: fn(AnyView, &mut ProjectContext, cx: &mut WindowContext) -> String,
+}
+
+/// Internal representation of an attachment tool to allow us to treat them dynamically
+struct RegisteredAttachment {
+    enabled: AtomicBool,
+    call: Box<dyn Fn(&mut WindowContext) -> Task<Result<UserAttachment>>>,
+}
+
+impl AttachmentRegistry {
+    pub fn new() -> Self {
+        Self {
+            registered_attachments: HashMap::default(),
+        }
+    }
+
+    pub fn register<A: LanguageModelAttachment + 'static>(&mut self, attachment: A) {
+        let call = Box::new(move |cx: &mut WindowContext| {
+            let result = attachment.run(cx);
+
+            cx.spawn(move |mut cx| async move {
+                let result: Result<A::Output> = result.await;
+                let view = cx.update(|cx| A::view(result, cx))?;
+
+                Ok(UserAttachment {
+                    view: view.into(),
+                    generate_fn: generate::<A>,
+                })
+            })
+        });
+
+        self.registered_attachments.insert(
+            TypeId::of::<A>(),
+            RegisteredAttachment {
+                call,
+                enabled: AtomicBool::new(true),
+            },
+        );
+        return;
+
+        fn generate<T: LanguageModelAttachment>(
+            view: AnyView,
+            project: &mut ProjectContext,
+            cx: &mut WindowContext,
+        ) -> String {
+            view.downcast::<T::View>()
+                .unwrap()
+                .update(cx, |view, cx| T::View::generate(view, project, cx))
+        }
+    }
+
+    pub fn set_attachment_tool_enabled<A: LanguageModelAttachment + 'static>(
+        &self,
+        is_enabled: bool,
+    ) {
+        if let Some(attachment) = self.registered_attachments.get(&TypeId::of::<A>()) {
+            attachment.enabled.store(is_enabled, SeqCst);
+        }
+    }
+
+    pub fn is_attachment_tool_enabled<A: LanguageModelAttachment + 'static>(&self) -> bool {
+        if let Some(attachment) = self.registered_attachments.get(&TypeId::of::<A>()) {
+            attachment.enabled.load(SeqCst)
+        } else {
+            false
+        }
+    }
+
+    pub fn call<A: LanguageModelAttachment + 'static>(
+        &self,
+        cx: &mut WindowContext,
+    ) -> Task<Result<UserAttachment>> {
+        let Some(attachment) = self.registered_attachments.get(&TypeId::of::<A>()) else {
+            return Task::ready(Err(anyhow!("no attachment tool")));
+        };
+
+        (attachment.call)(cx)
+    }
+
+    pub fn call_all_attachment_tools(
+        self: Arc<Self>,
+        cx: &mut WindowContext<'_>,
+    ) -> Task<Result<Vec<UserAttachment>>> {
+        let this = self.clone();
+        cx.spawn(|mut cx| async move {
+            let attachment_tasks = cx.update(|cx| {
+                let mut tasks = Vec::new();
+                for attachment in this
+                    .registered_attachments
+                    .values()
+                    .filter(|attachment| attachment.enabled.load(SeqCst))
+                {
+                    tasks.push((attachment.call)(cx))
+                }
+
+                tasks
+            })?;
+
+            let attachments = join_all(attachment_tasks.into_iter()).await;
+
+            Ok(attachments
+                .into_iter()
+                .filter_map(|attachment| attachment.log_err())
+                .collect())
+        })
+    }
+}
+
+impl UserAttachment {
+    pub fn generate(&self, output: &mut ProjectContext, cx: &mut WindowContext) -> Option<String> {
+        let result = (self.generate_fn)(self.view.clone(), output, cx);
+        if result.is_empty() {
+            None
+        } else {
+            Some(result)
+        }
+    }
+}

crates/assistant_tooling/src/project_context.rs

@@ -0,0 +1,296 @@
+use anyhow::{anyhow, Result};
+use gpui::{AppContext, Model, Task, WeakModel};
+use project::{Fs, Project, ProjectPath, Worktree};
+use std::{cmp::Ordering, fmt::Write as _, ops::Range, sync::Arc};
+use sum_tree::TreeMap;
+
+pub struct ProjectContext {
+    files: TreeMap<ProjectPath, PathState>,
+    project: WeakModel<Project>,
+    fs: Arc<dyn Fs>,
+}
+
+#[derive(Debug, Clone)]
+enum PathState {
+    PathOnly,
+    EntireFile,
+    Excerpts { ranges: Vec<Range<usize>> },
+}
+
+impl ProjectContext {
+    pub fn new(project: WeakModel<Project>, fs: Arc<dyn Fs>) -> Self {
+        Self {
+            files: TreeMap::default(),
+            fs,
+            project,
+        }
+    }
+
+    pub fn add_path(&mut self, project_path: ProjectPath) {
+        if self.files.get(&project_path).is_none() {
+            self.files.insert(project_path, PathState::PathOnly);
+        }
+    }
+
+    pub fn add_excerpts(&mut self, project_path: ProjectPath, new_ranges: &[Range<usize>]) {
+        let previous_state = self
+            .files
+            .get(&project_path)
+            .unwrap_or(&PathState::PathOnly);
+
+        let mut ranges = match previous_state {
+            PathState::EntireFile => return,
+            PathState::PathOnly => Vec::new(),
+            PathState::Excerpts { ranges } => ranges.to_vec(),
+        };
+
+        for new_range in new_ranges {
+            let ix = ranges.binary_search_by(|probe| {
+                if probe.end < new_range.start {
+                    Ordering::Less
+                } else if probe.start > new_range.end {
+                    Ordering::Greater
+                } else {
+                    Ordering::Equal
+                }
+            });
+
+            match ix {
+                Ok(mut ix) => {
+                    let existing = &mut ranges[ix];
+                    existing.start = existing.start.min(new_range.start);
+                    existing.end = existing.end.max(new_range.end);
+                    while ix + 1 < ranges.len() && ranges[ix + 1].start <= ranges[ix].end {
+                        ranges[ix].end = ranges[ix].end.max(ranges[ix + 1].end);
+                        ranges.remove(ix + 1);
+                    }
+                    while ix > 0 && ranges[ix - 1].end >= ranges[ix].start {
+                        ranges[ix].start = ranges[ix].start.min(ranges[ix - 1].start);
+                        ranges.remove(ix - 1);
+                        ix -= 1;
+                    }
+                }
+                Err(ix) => {
+                    ranges.insert(ix, new_range.clone());
+                }
+            }
+        }
+
+        self.files
+            .insert(project_path, PathState::Excerpts { ranges });
+    }
+
+    pub fn add_file(&mut self, project_path: ProjectPath) {
+        self.files.insert(project_path, PathState::EntireFile);
+    }
+
+    pub fn generate_system_message(&self, cx: &mut AppContext) -> Task<Result<String>> {
+        let project = self
+            .project
+            .upgrade()
+            .ok_or_else(|| anyhow!("project dropped"));
+        let files = self.files.clone();
+        let fs = self.fs.clone();
+        cx.spawn(|cx| async move {
+            let project = project?;
+            let mut result = "project structure:\n".to_string();
+
+            let mut last_worktree: Option<Model<Worktree>> = None;
+            for (project_path, path_state) in files.iter() {
+                if let Some(worktree) = &last_worktree {
+                    if worktree.read_with(&cx, |tree, _| tree.id())? != project_path.worktree_id {
+                        last_worktree = None;
+                    }
+                }
+
+                let worktree;
+                if let Some(last_worktree) = &last_worktree {
+                    worktree = last_worktree.clone();
+                } else if let Some(tree) = project.read_with(&cx, |project, cx| {
+                    project.worktree_for_id(project_path.worktree_id, cx)
+                })? {
+                    worktree = tree;
+                    last_worktree = Some(worktree.clone());
+                    let worktree_name =
+                        worktree.read_with(&cx, |tree, _cx| tree.root_name().to_string())?;
+                    writeln!(&mut result, "# {}", worktree_name).unwrap();
+                } else {
+                    continue;
+                }
+
+                let worktree_abs_path = worktree.read_with(&cx, |tree, _cx| tree.abs_path())?;
+                let path = &project_path.path;
+                writeln!(&mut result, "## {}", path.display()).unwrap();
+
+                match path_state {
+                    PathState::PathOnly => {}
+                    PathState::EntireFile => {
+                        let text = fs.load(&worktree_abs_path.join(&path)).await?;
+                        writeln!(&mut result, "~~~\n{text}\n~~~").unwrap();
+                    }
+                    PathState::Excerpts { ranges } => {
+                        let text = fs.load(&worktree_abs_path.join(&path)).await?;
+
+                        writeln!(&mut result, "~~~").unwrap();
+
+                        // Assumption: ranges are in order, not overlapping
+                        let mut prev_range_end = 0;
+                        for range in ranges {
+                            if range.start > prev_range_end {
+                                writeln!(&mut result, "...").unwrap();
+                                prev_range_end = range.end;
+                            }
+
+                            let mut start = range.start;
+                            let mut end = range.end.min(text.len());
+                            while !text.is_char_boundary(start) {
+                                start += 1;
+                            }
+                            while !text.is_char_boundary(end) {
+                                end -= 1;
+                            }
+                            result.push_str(&text[start..end]);
+                            if !result.ends_with('\n') {
+                                result.push('\n');
+                            }
+                        }
+
+                        if prev_range_end < text.len() {
+                            writeln!(&mut result, "...").unwrap();
+                        }
+
+                        writeln!(&mut result, "~~~").unwrap();
+                    }
+                }
+            }
+            Ok(result)
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::path::Path;
+
+    use super::*;
+    use gpui::TestAppContext;
+    use project::FakeFs;
+    use serde_json::json;
+    use settings::SettingsStore;
+
+    use unindent::Unindent as _;
+
+    #[gpui::test]
+    async fn test_system_message_generation(cx: &mut TestAppContext) {
+        init_test(cx);
+
+        let file_3_contents = r#"
+            fn test1() {}
+            fn test2() {}
+            fn test3() {}
+        "#
+        .unindent();
+
+        let fs = FakeFs::new(cx.executor());
+        fs.insert_tree(
+            "/code",
+            json!({
+                "root1": {
+                    "lib": {
+                        "file1.rs": "mod example;",
+                        "file2.rs": "",
+                    },
+                    "test": {
+                        "file3.rs": file_3_contents,
+                    }
+                },
+                "root2": {
+                    "src": {
+                        "main.rs": ""
+                    }
+                }
+            }),
+        )
+        .await;
+
+        let project = Project::test(
+            fs.clone(),
+            ["/code/root1".as_ref(), "/code/root2".as_ref()],
+            cx,
+        )
+        .await;
+
+        let worktree_ids = project.read_with(cx, |project, cx| {
+            project
+                .worktrees()
+                .map(|worktree| worktree.read(cx).id())
+                .collect::<Vec<_>>()
+        });
+
+        let mut ax = ProjectContext::new(project.downgrade(), fs);
+
+        ax.add_file(ProjectPath {
+            worktree_id: worktree_ids[0],
+            path: Path::new("lib/file1.rs").into(),
+        });
+
+        let message = cx
+            .update(|cx| ax.generate_system_message(cx))
+            .await
+            .unwrap();
+        assert_eq!(
+            r#"
+            project structure:
+            # root1
+            ## lib/file1.rs
+            ~~~
+            mod example;
+            ~~~
+            "#
+            .unindent(),
+            message
+        );
+
+        ax.add_excerpts(
+            ProjectPath {
+                worktree_id: worktree_ids[0],
+                path: Path::new("test/file3.rs").into(),
+            },
+            &[
+                file_3_contents.find("fn test2").unwrap()
+                    ..file_3_contents.find("fn test3").unwrap(),
+            ],
+        );
+
+        let message = cx
+            .update(|cx| ax.generate_system_message(cx))
+            .await
+            .unwrap();
+        assert_eq!(
+            r#"
+            project structure:
+            # root1
+            ## lib/file1.rs
+            ~~~
+            mod example;
+            ~~~
+            ## test/file3.rs
+            ~~~
+            ...
+            fn test2() {}
+            ...
+            ~~~
+            "#
+            .unindent(),
+            message
+        );
+    }
+
+    fn init_test(cx: &mut TestAppContext) {
+        cx.update(|cx| {
+            let settings_store = SettingsStore::test(cx);
+            cx.set_global(settings_store);
+            Project::init_settings(cx);
+        });
+    }
+}

crates/assistant_tooling/src/tool.rs

@@ -1,111 +0,0 @@
-use anyhow::Result;
-use gpui::{div, AnyElement, AnyView, IntoElement, Render, Task, View, WindowContext};
-use schemars::{schema::RootSchema, schema_for, JsonSchema};
-use serde::Deserialize;
-use std::fmt::Display;
-
-#[derive(Default, Deserialize)]
-pub struct ToolFunctionCall {
-    pub id: String,
-    pub name: String,
-    pub arguments: String,
-    #[serde(skip)]
-    pub result: Option<ToolFunctionCallResult>,
-}
-
-pub enum ToolFunctionCallResult {
-    NoSuchTool,
-    ParsingFailed,
-    Finished { for_model: String, view: AnyView },
-}
-
-impl ToolFunctionCallResult {
-    pub fn format(&self, name: &String) -> String {
-        match self {
-            ToolFunctionCallResult::NoSuchTool => format!("No tool for {name}"),
-            ToolFunctionCallResult::ParsingFailed => {
-                format!("Unable to parse arguments for {name}")
-            }
-            ToolFunctionCallResult::Finished { for_model, .. } => for_model.clone(),
-        }
-    }
-
-    pub fn into_any_element(&self, name: &String) -> AnyElement {
-        match self {
-            ToolFunctionCallResult::NoSuchTool => {
-                format!("Language Model attempted to call {name}").into_any_element()
-            }
-            ToolFunctionCallResult::ParsingFailed => {
-                format!("Language Model called {name} with bad arguments").into_any_element()
-            }
-            ToolFunctionCallResult::Finished { view, .. } => view.clone().into_any_element(),
-        }
-    }
-}
-
-#[derive(Clone)]
-pub struct ToolFunctionDefinition {
-    pub name: String,
-    pub description: String,
-    pub parameters: RootSchema,
-}
-
-impl Display for ToolFunctionDefinition {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let schema = serde_json::to_string(&self.parameters).ok();
-        let schema = schema.unwrap_or("None".to_string());
-        write!(f, "Name: {}:\n", self.name)?;
-        write!(f, "Description: {}\n", self.description)?;
-        write!(f, "Parameters: {}", schema)
-    }
-}
-
-pub trait LanguageModelTool {
-    /// The input type that will be passed in to `execute` when the tool is called
-    /// by the language model.
-    type Input: for<'de> Deserialize<'de> + JsonSchema;
-
-    /// The output returned by executing the tool.
-    type Output: 'static;
-
-    type View: Render;
-
-    /// Returns the name of the tool.
-    ///
-    /// This name is exposed to the language model to allow the model to pick
-    /// which tools to use. As this name is used to identify the tool within a
-    /// tool registry, it should be unique.
-    fn name(&self) -> String;
-
-    /// Returns the description of the tool.
-    ///
-    /// This can be used to _prompt_ the model as to what the tool does.
-    fn description(&self) -> String;
-
-    /// Returns the OpenAI Function definition for the tool, for direct use with OpenAI's API.
-    fn definition(&self) -> ToolFunctionDefinition {
-        let root_schema = schema_for!(Self::Input);
-
-        ToolFunctionDefinition {
-            name: self.name(),
-            description: self.description(),
-            parameters: root_schema,
-        }
-    }
-
-    /// Executes the tool with the given input.
-    fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>>;
-
-    fn format(input: &Self::Input, output: &Result<Self::Output>) -> String;
-
-    fn output_view(
-        tool_call_id: String,
-        input: Self::Input,
-        output: Result<Self::Output>,
-        cx: &mut WindowContext,
-    ) -> View<Self::View>;
-
-    fn render_running(_cx: &mut WindowContext) -> impl IntoElement {
-        div()
-    }
-}

crates/assistant_tooling/src/tool_registry.rs

@@ -1,54 +1,115 @@
 use anyhow::{anyhow, Result};
-use gpui::{div, AnyElement, IntoElement as _, ParentElement, Styled, Task, WindowContext};
+use gpui::{
+    div, AnyElement, AnyView, IntoElement, ParentElement, Render, Styled, Task, View, WindowContext,
+};
+use schemars::{schema::RootSchema, schema_for, JsonSchema};
+use serde::Deserialize;
 use std::{
     any::TypeId,
     collections::HashMap,
+    fmt::Display,
     sync::atomic::{AtomicBool, Ordering::SeqCst},
 };
 
-use crate::tool::{
-    LanguageModelTool, ToolFunctionCall, ToolFunctionCallResult, ToolFunctionDefinition,
-};
+use crate::ProjectContext;
 
-// Internal Tool representation for the registry
-pub struct Tool {
-    enabled: AtomicBool,
-    type_id: TypeId,
-    call: Box<dyn Fn(&ToolFunctionCall, &mut WindowContext) -> Task<Result<ToolFunctionCall>>>,
-    render_running: Box<dyn Fn(&mut WindowContext) -> gpui::AnyElement>,
-    definition: ToolFunctionDefinition,
+pub struct ToolRegistry {
+    registered_tools: HashMap<String, RegisteredTool>,
 }
 
-impl Tool {
-    fn new(
-        type_id: TypeId,
-        call: Box<dyn Fn(&ToolFunctionCall, &mut WindowContext) -> Task<Result<ToolFunctionCall>>>,
-        render_running: Box<dyn Fn(&mut WindowContext) -> gpui::AnyElement>,
-        definition: ToolFunctionDefinition,
-    ) -> Self {
-        Self {
-            enabled: AtomicBool::new(true),
-            type_id,
-            call,
-            render_running,
-            definition,
+#[derive(Default, Deserialize)]
+pub struct ToolFunctionCall {
+    pub id: String,
+    pub name: String,
+    pub arguments: String,
+    #[serde(skip)]
+    pub result: Option<ToolFunctionCallResult>,
+}
+
+pub enum ToolFunctionCallResult {
+    NoSuchTool,
+    ParsingFailed,
+    Finished {
+        view: AnyView,
+        generate_fn: fn(AnyView, &mut ProjectContext, &mut WindowContext) -> String,
+    },
+}
+
+#[derive(Clone)]
+pub struct ToolFunctionDefinition {
+    pub name: String,
+    pub description: String,
+    pub parameters: RootSchema,
+}
+
+pub trait LanguageModelTool {
+    /// The input type that will be passed in to `execute` when the tool is called
+    /// by the language model.
+    type Input: for<'de> Deserialize<'de> + JsonSchema;
+
+    /// The output returned by executing the tool.
+    type Output: 'static;
+
+    type View: Render + ToolOutput;
+
+    /// Returns the name of the tool.
+    ///
+    /// This name is exposed to the language model to allow the model to pick
+    /// which tools to use. As this name is used to identify the tool within a
+    /// tool registry, it should be unique.
+    fn name(&self) -> String;
+
+    /// Returns the description of the tool.
+    ///
+    /// This can be used to _prompt_ the model as to what the tool does.
+    fn description(&self) -> String;
+
+    /// Returns the OpenAI Function definition for the tool, for direct use with OpenAI's API.
+    fn definition(&self) -> ToolFunctionDefinition {
+        let root_schema = schema_for!(Self::Input);
+
+        ToolFunctionDefinition {
+            name: self.name(),
+            description: self.description(),
+            parameters: root_schema,
         }
     }
+
+    /// Executes the tool with the given input.
+    fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>>;
+
+    fn output_view(
+        input: Self::Input,
+        output: Result<Self::Output>,
+        cx: &mut WindowContext,
+    ) -> View<Self::View>;
+
+    fn render_running(_cx: &mut WindowContext) -> impl IntoElement {
+        div()
+    }
 }
 
-pub struct ToolRegistry {
-    tools: HashMap<String, Tool>,
+pub trait ToolOutput: Sized {
+    fn generate(&self, project: &mut ProjectContext, cx: &mut WindowContext) -> String;
+}
+
+struct RegisteredTool {
+    enabled: AtomicBool,
+    type_id: TypeId,
+    call: Box<dyn Fn(&ToolFunctionCall, &mut WindowContext) -> Task<Result<ToolFunctionCall>>>,
+    render_running: fn(&mut WindowContext) -> gpui::AnyElement,
+    definition: ToolFunctionDefinition,
 }
 
 impl ToolRegistry {
     pub fn new() -> Self {
         Self {
-            tools: HashMap::new(),
+            registered_tools: HashMap::new(),
         }
     }
 
     pub fn set_tool_enabled<T: 'static + LanguageModelTool>(&self, is_enabled: bool) {
-        for tool in self.tools.values() {
+        for tool in self.registered_tools.values() {
             if tool.type_id == TypeId::of::<T>() {
                 tool.enabled.store(is_enabled, SeqCst);
                 return;
@@ -57,7 +118,7 @@ impl ToolRegistry {
     }
 
     pub fn is_tool_enabled<T: 'static + LanguageModelTool>(&self) -> bool {
-        for tool in self.tools.values() {
+        for tool in self.registered_tools.values() {
             if tool.type_id == TypeId::of::<T>() {
                 return tool.enabled.load(SeqCst);
             }
@@ -66,7 +127,7 @@ impl ToolRegistry {
     }
 
     pub fn definitions(&self) -> Vec<ToolFunctionDefinition> {
-        self.tools
+        self.registered_tools
             .values()
             .filter(|tool| tool.enabled.load(SeqCst))
             .map(|tool| tool.definition.clone())
@@ -84,7 +145,7 @@ impl ToolRegistry {
                 .child(result.into_any_element(&tool_call.name))
                 .into_any_element(),
             None => self
-                .tools
+                .registered_tools
                 .get(&tool_call.name)
                 .map(|tool| (tool.render_running)(cx))
                 .unwrap_or_else(|| div().into_any_element()),
@@ -96,13 +157,12 @@ impl ToolRegistry {
         tool: T,
         _cx: &mut WindowContext,
     ) -> Result<()> {
-        let definition = tool.definition();
-
         let name = tool.name();
-
-        let registered_tool = Tool::new(
-            TypeId::of::<T>(),
-            Box::new(
+        let registered_tool = RegisteredTool {
+            type_id: TypeId::of::<T>(),
+            definition: tool.definition(),
+            enabled: AtomicBool::new(true),
+            call: Box::new(
                 move |tool_call: &ToolFunctionCall, cx: &mut WindowContext| {
                     let name = tool_call.name.clone();
                     let arguments = tool_call.arguments.clone();
@@ -121,8 +181,7 @@ impl ToolRegistry {
 
                     cx.spawn(move |mut cx| async move {
                         let result: Result<T::Output> = result.await;
-                        let for_model = T::format(&input, &result);
-                        let view = cx.update(|cx| T::output_view(id.clone(), input, result, cx))?;
+                        let view = cx.update(|cx| T::output_view(input, result, cx))?;
 
                         Ok(ToolFunctionCall {
                             id,
@@ -130,23 +189,35 @@ impl ToolRegistry {
                             arguments,
                             result: Some(ToolFunctionCallResult::Finished {
                                 view: view.into(),
-                                for_model,
+                                generate_fn: generate::<T>,
                             }),
                         })
                     })
                 },
             ),
-            Box::new(|cx| T::render_running(cx).into_any_element()),
-            definition,
-        );
-
-        let previous = self.tools.insert(name.clone(), registered_tool);
+            render_running: render_running::<T>,
+        };
 
+        let previous = self.registered_tools.insert(name.clone(), registered_tool);
         if previous.is_some() {
             return Err(anyhow!("already registered a tool with name {}", name));
         }
 
-        Ok(())
+        return Ok(());
+
+        fn render_running<T: LanguageModelTool>(cx: &mut WindowContext) -> AnyElement {
+            T::render_running(cx).into_any_element()
+        }
+
+        fn generate<T: LanguageModelTool>(
+            view: AnyView,
+            project: &mut ProjectContext,
+            cx: &mut WindowContext,
+        ) -> String {
+            view.downcast::<T::View>()
+                .unwrap()
+                .update(cx, |view, cx| T::View::generate(view, project, cx))
+        }
     }
 
     /// Task yields an error if the window for the given WindowContext is closed before the task completes.
@@ -159,7 +230,7 @@ impl ToolRegistry {
         let arguments = tool_call.arguments.clone();
         let id = tool_call.id.clone();
 
-        let tool = match self.tools.get(&name) {
+        let tool = match self.registered_tools.get(&name) {
             Some(tool) => tool,
             None => {
                 let name = name.clone();
@@ -176,6 +247,47 @@ impl ToolRegistry {
     }
 }
 
+impl ToolFunctionCallResult {
+    pub fn generate(
+        &self,
+        name: &String,
+        project: &mut ProjectContext,
+        cx: &mut WindowContext,
+    ) -> String {
+        match self {
+            ToolFunctionCallResult::NoSuchTool => format!("No tool for {name}"),
+            ToolFunctionCallResult::ParsingFailed => {
+                format!("Unable to parse arguments for {name}")
+            }
+            ToolFunctionCallResult::Finished { generate_fn, view } => {
+                (generate_fn)(view.clone(), project, cx)
+            }
+        }
+    }
+
+    fn into_any_element(&self, name: &String) -> AnyElement {
+        match self {
+            ToolFunctionCallResult::NoSuchTool => {
+                format!("Language Model attempted to call {name}").into_any_element()
+            }
+            ToolFunctionCallResult::ParsingFailed => {
+                format!("Language Model called {name} with bad arguments").into_any_element()
+            }
+            ToolFunctionCallResult::Finished { view, .. } => view.clone().into_any_element(),
+        }
+    }
+}
+
+impl Display for ToolFunctionDefinition {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let schema = serde_json::to_string(&self.parameters).ok();
+        let schema = schema.unwrap_or("None".to_string());
+        write!(f, "Name: {}:\n", self.name)?;
+        write!(f, "Description: {}\n", self.description)?;
+        write!(f, "Parameters: {}", schema)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -213,6 +325,12 @@ mod test {
         }
     }
 
+    impl ToolOutput for WeatherView {
+        fn generate(&self, _output: &mut ProjectContext, _cx: &mut WindowContext) -> String {
+            serde_json::to_string(&self.result).unwrap()
+        }
+    }
+
     impl LanguageModelTool for WeatherTool {
         type Input = WeatherQuery;
         type Output = WeatherResult;
@@ -240,7 +358,6 @@ mod test {
         }
 
         fn output_view(
-            _tool_call_id: String,
             _input: Self::Input,
             result: Result<Self::Output>,
             cx: &mut WindowContext,
@@ -250,10 +367,6 @@ mod test {
                 WeatherView { result }
             })
         }
-
-        fn format(_: &Self::Input, output: &Result<Self::Output>) -> String {
-            serde_json::to_string(&output.as_ref().unwrap()).unwrap()
-        }
     }
 
     #[gpui::test]