Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Support profiles in agent2 (#36034)

Browse source 
We still need a profile selector.

Release Notes:

- N/A

---------

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>
This commit is contained in:
Antonio Scandurra 2025-08-12 14:17:48 +02:00 • committed by GitHub
parent 13bf45dd4a
commit 2444321756
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 587 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

231
crates/agent2/src/tools/context_server_registry.rs Normal file
View file

@ -0,0 +1,231 @@
use crate::{AgentToolOutput, AnyAgentTool, ToolCallEventStream};
use agent_client_protocol::ToolKind;
use anyhow::{Result, anyhow, bail};
use collections::{BTreeMap, HashMap};
use context_server::ContextServerId;
use gpui::{App, Context, Entity, SharedString, Task};
use project::context_server_store::{ContextServerStatus, ContextServerStore};
use std::sync::Arc;
use util::ResultExt;
pub struct ContextServerRegistry {
server_store: Entity<ContextServerStore>,
registered_servers: HashMap<ContextServerId, RegisteredContextServer>,
_subscription: gpui::Subscription,
}
struct RegisteredContextServer {
tools: BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,
load_tools: Task<Result<()>>,
}
impl ContextServerRegistry {
pub fn new(server_store: Entity<ContextServerStore>, cx: &mut Context<Self>) -> Self {
let mut this = Self {
server_store: server_store.clone(),
registered_servers: HashMap::default(),
_subscription: cx.subscribe(&server_store, Self::handle_context_server_store_event),
};
for server in server_store.read(cx).running_servers() {
this.reload_tools_for_server(server.id(), cx);
}
this
}
pub fn servers(
&self,
) -> impl Iterator<
Item = (
&ContextServerId,
&BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,
),
> {
self.registered_servers
.iter()
.map(|(id, server)| (id, &server.tools))
}
fn reload_tools_for_server(&mut self, server_id: ContextServerId, cx: &mut Context<Self>) {
let Some(server) = self.server_store.read(cx).get_running_server(&server_id) else {
return;
};
let Some(client) = server.client() else {
return;
};
if !client.capable(context_server::protocol::ServerCapability::Tools) {
return;
}
let registered_server =
self.registered_servers
.entry(server_id.clone())
.or_insert(RegisteredContextServer {
tools: BTreeMap::default(),
load_tools: Task::ready(Ok(())),
});
registered_server.load_tools = cx.spawn(async move |this, cx| {
let response = client
.request::<context_server::types::requests::ListTools>(())
.await;
this.update(cx, |this, cx| {
let Some(registered_server) = this.registered_servers.get_mut(&server_id) else {
return;
};
registered_server.tools.clear();
if let Some(response) = response.log_err() {
for tool in response.tools {
let tool = Arc::new(ContextServerTool::new(
this.server_store.clone(),
server.id(),
tool,
));
registered_server.tools.insert(tool.name(), tool);
}
cx.notify();
}
})
});
}
fn handle_context_server_store_event(
&mut self,
_: Entity<ContextServerStore>,
event: &project::context_server_store::Event,
cx: &mut Context<Self>,
) {
match event {
project::context_server_store::Event::ServerStatusChanged { server_id, status } => {
match status {
ContextServerStatus::Starting => {}
ContextServerStatus::Running => {
self.reload_tools_for_server(server_id.clone(), cx);
}
ContextServerStatus::Stopped | ContextServerStatus::Error(_) => {
self.registered_servers.remove(&server_id);
cx.notify();
}
}
}
}
}
}
struct ContextServerTool {
store: Entity<ContextServerStore>,
server_id: ContextServerId,
tool: context_server::types::Tool,
}
impl ContextServerTool {
fn new(
store: Entity<ContextServerStore>,
server_id: ContextServerId,
tool: context_server::types::Tool,
) -> Self {
Self {
store,
server_id,
tool,
}
}
}
impl AnyAgentTool for ContextServerTool {
fn name(&self) -> SharedString {
self.tool.name.clone().into()
}
fn description(&self) -> SharedString {
self.tool.description.clone().unwrap_or_default().into()
}
fn kind(&self) -> ToolKind {
ToolKind::Other
}
fn initial_title(&self, _input: serde_json::Value) -> SharedString {
format!("Run MCP tool `{}`", self.tool.name).into()
}
fn input_schema(
&self,
format: language_model::LanguageModelToolSchemaFormat,
) -> Result<serde_json::Value> {
let mut schema = self.tool.input_schema.clone();
assistant_tool::adapt_schema_to_format(&mut schema, format)?;
Ok(match schema {
serde_json::Value::Null => {
serde_json::json!({ "type": "object", "properties": [] })
}
serde_json::Value::Object(map) if map.is_empty() => {
serde_json::json!({ "type": "object", "properties": [] })
}
_ => schema,
})
}
fn run(
self: Arc<Self>,
input: serde_json::Value,
_event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<AgentToolOutput>> {
let Some(server) = self.store.read(cx).get_running_server(&self.server_id) else {
return Task::ready(Err(anyhow!("Context server not found")));
};
let tool_name = self.tool.name.clone();
let server_clone = server.clone();
let input_clone = input.clone();
cx.spawn(async move |_cx| {
let Some(protocol) = server_clone.client() else {
bail!("Context server not initialized");
};
let arguments = if let serde_json::Value::Object(map) = input_clone {
Some(map.into_iter().collect())
} else {
None
};
log::trace!(
"Running tool: {} with arguments: {:?}",
tool_name,
arguments
);
let response = protocol
.request::<context_server::types::requests::CallTool>(
context_server::types::CallToolParams {
name: tool_name,
arguments,
meta: None,
},
)
.await?;
let mut result = String::new();
for content in response.content {
match content {
context_server::types::ToolResponseContent::Text { text } => {
result.push_str(&text);
}
context_server::types::ToolResponseContent::Image { .. } => {
log::warn!("Ignoring image content from tool response");
}
context_server::types::ToolResponseContent::Audio { .. } => {
log::warn!("Ignoring audio content from tool response");
}
context_server::types::ToolResponseContent::Resource { .. } => {
log::warn!("Ignoring resource content from tool response");
}
}
}
Ok(AgentToolOutput {
raw_output: result.clone().into(),
llm_output: result.into(),
})
})
}
}

18
crates/agent2/src/tools/diagnostics_tool.rs
View file

@ -85,7 +85,7 @@ impl AgentTool for DiagnosticsTool {
fn run(
self: Arc<Self>,
input: Self::Input,
event_stream: ToolCallEventStream,
_event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<Self::Output>> {
match input.path {
@ -119,11 +119,6 @@ impl AgentTool for DiagnosticsTool {
range.start.row + 1,
entry.diagnostic.message
)?;
event_stream.update_fields(acp::ToolCallUpdateFields {
content: Some(vec![output.clone().into()]),
..Default::default()
});
}
if output.is_empty() {
@ -158,18 +153,9 @@ impl AgentTool for DiagnosticsTool {
}
if has_diagnostics {
event_stream.update_fields(acp::ToolCallUpdateFields {
content: Some(vec![output.clone().into()]),
..Default::default()
});
Task::ready(Ok(output))
} else {
let text = "No errors or warnings found in the project.";
event_stream.update_fields(acp::ToolCallUpdateFields {
content: Some(vec![text.into()]),
..Default::default()
});
Task::ready(Ok(text.into()))
Task::ready(Ok("No errors or warnings found in the project.".into()))
}
}
}

66
crates/agent2/src/tools/edit_file_tool.rs
View file

@ -454,9 +454,8 @@ fn resolve_path(
#[cfg(test)]
mod tests {
use crate::Templates;
use super::*;
use crate::{ContextServerRegistry, Templates};
use action_log::ActionLog;
use client::TelemetrySettings;
use fs::Fs;
@ -475,9 +474,20 @@ mod tests {
fs.insert_tree("/root", json!({})).await;
let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread =
cx.new(|_| Thread::new(project, Rc::default(), action_log, Templates::new(), model));
let thread = cx.new(|cx| {
Thread::new(
project,
Rc::default(),
context_server_registry,
action_log,
Templates::new(),
model,
cx,
)
});
let result = cx
.update(|cx| {
let input = EditFileToolInput {
@ -661,14 +671,18 @@ mod tests {
});
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project,
Rc::default(),
context_server_registry,
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
@ -792,15 +806,19 @@ mod tests {
.unwrap();
let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project,
Rc::default(),
context_server_registry,
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
@ -914,15 +932,19 @@ mod tests {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project,
Rc::default(),
context_server_registry,
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
let tool = Arc::new(EditFileTool { thread });
@ -1041,15 +1063,19 @@ mod tests {
let fs = project::FakeFs::new(cx.executor());
fs.insert_tree("/project", json!({})).await;
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project,
Rc::default(),
context_server_registry,
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
let tool = Arc::new(EditFileTool { thread });
@ -1148,14 +1174,18 @@ mod tests {
.await;
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project.clone(),
Rc::default(),
context_server_registry.clone(),
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
let tool = Arc::new(EditFileTool { thread });
@ -1225,14 +1255,18 @@ mod tests {
.await;
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project.clone(),
Rc::default(),
context_server_registry.clone(),
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
let tool = Arc::new(EditFileTool { thread });
@ -1305,14 +1339,18 @@ mod tests {
.await;
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project.clone(),
Rc::default(),
context_server_registry.clone(),
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
let tool = Arc::new(EditFileTool { thread });
@ -1382,14 +1420,18 @@ mod tests {
let fs = project::FakeFs::new(cx.executor());
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let action_log = cx.new(|_| ActionLog::new(project.clone()));
let context_server_registry =
cx.new(|cx| ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
let model = Arc::new(FakeLanguageModel::default());
let thread = cx.new(|_| {
let thread = cx.new(|cx| {
Thread::new(
project.clone(),
Rc::default(),
context_server_registry,
action_log.clone(),
Templates::new(),
model.clone(),
cx,
)
});
let tool = Arc::new(EditFileTool { thread });

8
crates/agent2/src/tools/fetch_tool.rs
View file

@ -136,7 +136,7 @@ impl AgentTool for FetchTool {
fn run(
self: Arc<Self>,
input: Self::Input,
event_stream: ToolCallEventStream,
_event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<Self::Output>> {
let text = cx.background_spawn({
@ -149,12 +149,6 @@ impl AgentTool for FetchTool {
if text.trim().is_empty() {
bail!("no textual content found");
}
event_stream.update_fields(acp::ToolCallUpdateFields {
content: Some(vec![text.clone().into()]),
..Default::default()
});
Ok(text)
})
}

3
crates/agent2/src/tools/find_path_tool.rs
View file

@ -139,9 +139,6 @@ impl AgentTool for FindPathTool {
})
.collect(),
),
raw_output: Some(serde_json::json!({
"paths": &matches,
})),
..Default::default()
});

25
crates/agent2/src/tools/grep_tool.rs
View file

@ -101,7 +101,7 @@ impl AgentTool for GrepTool {
fn run(
self: Arc<Self>,
input: Self::Input,
event_stream: ToolCallEventStream,
_event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<Self::Output>> {
const CONTEXT_LINES: u32 = 2;
@ -282,33 +282,22 @@ impl AgentTool for GrepTool {
}
}
event_stream.update_fields(acp::ToolCallUpdateFields {
content: Some(vec![output.clone().into()]),
..Default::default()
});
matches_found += 1;
}
}
let output = if matches_found == 0 {
"No matches found".to_string()
if matches_found == 0 {
Ok("No matches found".into())
} else if has_more_matches {
format!(
Ok(format!(
"Showing matches {}-{} (there were more matches found; use offset: {} to see next page):\n{output}",
input.offset + 1,
input.offset + matches_found,
input.offset + RESULTS_PER_PAGE,
)
))
} else {
format!("Found {matches_found} matches:\n{output}")
};
event_stream.update_fields(acp::ToolCallUpdateFields {
content: Some(vec![output.clone().into()]),
..Default::default()
});
Ok(output)
Ok(format!("Found {matches_found} matches:\n{output}"))
}
})
}
}

11
crates/agent2/src/tools/now_tool.rs
View file

@ -47,20 +47,13 @@ impl AgentTool for NowTool {
fn run(
self: Arc<Self>,
input: Self::Input,
event_stream: ToolCallEventStream,
_event_stream: ToolCallEventStream,
_cx: &mut App,
) -> Task<Result<String>> {
let now = match input.timezone {
Timezone::Utc => Utc::now().to_rfc3339(),
Timezone::Local => Local::now().to_rfc3339(),
};
let content = format!("The current datetime is {now}.");
event_stream.update_fields(acp::ToolCallUpdateFields {
content: Some(vec![content.clone().into()]),
..Default::default()
});
Task::ready(Ok(content))
Task::ready(Ok(format!("The current datetime is {now}.")))
}
}