Add system prompt and tool permission to agent2 (#35781)
Release Notes: - N/A --------- Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com> Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>
This commit is contained in:
Antonio Scandurra
GitHub
parent
4dbd24d75f
commit
03876d076e
21 changed files with 1111 additions and 304 deletions
|
|
@ -1,16 +1,39 @@
|
|||
use crate::ToolCallAuthorization;
|
||||
use crate::{templates::Templates, AgentResponseEvent, Thread};
|
||||
use acp_thread::ModelSelector;
|
||||
use agent_client_protocol as acp;
|
||||
use anyhow::{anyhow, Result};
|
||||
use futures::StreamExt;
|
||||
use gpui::{App, AppContext, AsyncApp, Entity, Subscription, Task, WeakEntity};
|
||||
use anyhow::{anyhow, Context as _, Result};
|
||||
use futures::{future, StreamExt};
|
||||
use gpui::{
|
||||
App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity,
|
||||
};
|
||||
use language_model::{LanguageModel, LanguageModelRegistry};
|
||||
use project::Project;
|
||||
use project::{Project, ProjectItem, ProjectPath, Worktree};
|
||||
use prompt_store::{
|
||||
ProjectContext, PromptId, PromptStore, RulesFileContext, UserRulesContext, WorktreeContext,
|
||||
};
|
||||
use std::cell::RefCell;
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::rc::Rc;
|
||||
use std::sync::Arc;
|
||||
use util::ResultExt;
|
||||
|
||||
use crate::{templates::Templates, AgentResponseEvent, Thread};
|
||||
const RULES_FILE_NAMES: [&'static str; 9] = [
|
||||
".rules",
|
||||
".cursorrules",
|
||||
".windsurfrules",
|
||||
".clinerules",
|
||||
".github/copilot-instructions.md",
|
||||
"CLAUDE.md",
|
||||
"AGENT.md",
|
||||
"AGENTS.md",
|
||||
"GEMINI.md",
|
||||
];
|
||||
|
||||
pub struct RulesLoadingError {
|
||||
pub message: SharedString,
|
||||
}
|
||||
|
||||
/// Holds both the internal Thread and the AcpThread for a session
|
||||
struct Session {
|
||||
|
|
@ -24,17 +47,247 @@ struct Session {
|
|||
pub struct NativeAgent {
|
||||
/// Session ID -> Session mapping
|
||||
sessions: HashMap<acp::SessionId, Session>,
|
||||
/// Shared project context for all threads
|
||||
project_context: Rc<RefCell<ProjectContext>>,
|
||||
project_context_needs_refresh: watch::Sender<()>,
|
||||
_maintain_project_context: Task<Result<()>>,
|
||||
/// Shared templates for all threads
|
||||
templates: Arc<Templates>,
|
||||
project: Entity<Project>,
|
||||
prompt_store: Option<Entity<PromptStore>>,
|
||||
_subscriptions: Vec<Subscription>,
|
||||
}
|
||||
|
||||
impl NativeAgent {
|
||||
pub fn new(templates: Arc<Templates>) -> Self {
|
||||
pub async fn new(
|
||||
project: Entity<Project>,
|
||||
templates: Arc<Templates>,
|
||||
prompt_store: Option<Entity<PromptStore>>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<Entity<NativeAgent>> {
|
||||
log::info!("Creating new NativeAgent");
|
||||
Self {
|
||||
sessions: HashMap::new(),
|
||||
templates,
|
||||
|
||||
let project_context = cx
|
||||
.update(|cx| Self::build_project_context(&project, prompt_store.as_ref(), cx))?
|
||||
.await;
|
||||
|
||||
cx.new(|cx| {
|
||||
let mut subscriptions = vec![cx.subscribe(&project, Self::handle_project_event)];
|
||||
if let Some(prompt_store) = prompt_store.as_ref() {
|
||||
subscriptions.push(cx.subscribe(prompt_store, Self::handle_prompts_updated_event))
|
||||
}
|
||||
|
||||
let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
|
||||
watch::channel(());
|
||||
Self {
|
||||
sessions: HashMap::new(),
|
||||
project_context: Rc::new(RefCell::new(project_context)),
|
||||
project_context_needs_refresh: project_context_needs_refresh_tx,
|
||||
_maintain_project_context: cx.spawn(async move |this, cx| {
|
||||
Self::maintain_project_context(this, project_context_needs_refresh_rx, cx).await
|
||||
}),
|
||||
templates,
|
||||
project,
|
||||
prompt_store,
|
||||
_subscriptions: subscriptions,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn maintain_project_context(
|
||||
this: WeakEntity<Self>,
|
||||
mut needs_refresh: watch::Receiver<()>,
|
||||
cx: &mut AsyncApp,
|
||||
) -> Result<()> {
|
||||
while needs_refresh.changed().await.is_ok() {
|
||||
let project_context = this
|
||||
.update(cx, |this, cx| {
|
||||
Self::build_project_context(&this.project, this.prompt_store.as_ref(), cx)
|
||||
})?
|
||||
.await;
|
||||
this.update(cx, |this, _| this.project_context.replace(project_context))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn build_project_context(
|
||||
project: &Entity<Project>,
|
||||
prompt_store: Option<&Entity<PromptStore>>,
|
||||
cx: &mut App,
|
||||
) -> Task<ProjectContext> {
|
||||
let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
|
||||
let worktree_tasks = worktrees
|
||||
.into_iter()
|
||||
.map(|worktree| {
|
||||
Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
|
||||
prompt_store.read_with(cx, |prompt_store, cx| {
|
||||
let prompts = prompt_store.default_prompt_metadata();
|
||||
let load_tasks = prompts.into_iter().map(|prompt_metadata| {
|
||||
let contents = prompt_store.load(prompt_metadata.id, cx);
|
||||
async move { (contents.await, prompt_metadata) }
|
||||
});
|
||||
cx.background_spawn(future::join_all(load_tasks))
|
||||
})
|
||||
} else {
|
||||
Task::ready(vec![])
|
||||
};
|
||||
|
||||
cx.spawn(async move |_cx| {
|
||||
let (worktrees, default_user_rules) =
|
||||
future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
|
||||
|
||||
let worktrees = worktrees
|
||||
.into_iter()
|
||||
.map(|(worktree, _rules_error)| {
|
||||
// TODO: show error message
|
||||
// if let Some(rules_error) = rules_error {
|
||||
// this.update(cx, |_, cx| cx.emit(rules_error)).ok();
|
||||
// }
|
||||
worktree
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let default_user_rules = default_user_rules
|
||||
.into_iter()
|
||||
.flat_map(|(contents, prompt_metadata)| match contents {
|
||||
Ok(contents) => Some(UserRulesContext {
|
||||
uuid: match prompt_metadata.id {
|
||||
PromptId::User { uuid } => uuid,
|
||||
PromptId::EditWorkflow => return None,
|
||||
},
|
||||
title: prompt_metadata.title.map(|title| title.to_string()),
|
||||
contents,
|
||||
}),
|
||||
Err(_err) => {
|
||||
// TODO: show error message
|
||||
// this.update(cx, |_, cx| {
|
||||
// cx.emit(RulesLoadingError {
|
||||
// message: format!("{err:?}").into(),
|
||||
// });
|
||||
// })
|
||||
// .ok();
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
ProjectContext::new(worktrees, default_user_rules)
|
||||
})
|
||||
}
|
||||
|
||||
fn load_worktree_info_for_system_prompt(
|
||||
worktree: Entity<Worktree>,
|
||||
project: Entity<Project>,
|
||||
cx: &mut App,
|
||||
) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
|
||||
let tree = worktree.read(cx);
|
||||
let root_name = tree.root_name().into();
|
||||
let abs_path = tree.abs_path();
|
||||
|
||||
let mut context = WorktreeContext {
|
||||
root_name,
|
||||
abs_path,
|
||||
rules_file: None,
|
||||
};
|
||||
|
||||
let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
|
||||
let Some(rules_task) = rules_task else {
|
||||
return Task::ready((context, None));
|
||||
};
|
||||
|
||||
cx.spawn(async move |_| {
|
||||
let (rules_file, rules_file_error) = match rules_task.await {
|
||||
Ok(rules_file) => (Some(rules_file), None),
|
||||
Err(err) => (
|
||||
None,
|
||||
Some(RulesLoadingError {
|
||||
message: format!("{err}").into(),
|
||||
}),
|
||||
),
|
||||
};
|
||||
context.rules_file = rules_file;
|
||||
(context, rules_file_error)
|
||||
})
|
||||
}
|
||||
|
||||
fn load_worktree_rules_file(
|
||||
worktree: Entity<Worktree>,
|
||||
project: Entity<Project>,
|
||||
cx: &mut App,
|
||||
) -> Option<Task<Result<RulesFileContext>>> {
|
||||
let worktree = worktree.read(cx);
|
||||
let worktree_id = worktree.id();
|
||||
let selected_rules_file = RULES_FILE_NAMES
|
||||
.into_iter()
|
||||
.filter_map(|name| {
|
||||
worktree
|
||||
.entry_for_path(name)
|
||||
.filter(|entry| entry.is_file())
|
||||
.map(|entry| entry.path.clone())
|
||||
})
|
||||
.next();
|
||||
|
||||
// Note that Cline supports `.clinerules` being a directory, but that is not currently
|
||||
// supported. This doesn't seem to occur often in GitHub repositories.
|
||||
selected_rules_file.map(|path_in_worktree| {
|
||||
let project_path = ProjectPath {
|
||||
worktree_id,
|
||||
path: path_in_worktree.clone(),
|
||||
};
|
||||
let buffer_task =
|
||||
project.update(cx, |project, cx| project.open_buffer(project_path, cx));
|
||||
let rope_task = cx.spawn(async move |cx| {
|
||||
buffer_task.await?.read_with(cx, |buffer, cx| {
|
||||
let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
|
||||
anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
|
||||
})?
|
||||
});
|
||||
// Build a string from the rope on a background thread.
|
||||
cx.background_spawn(async move {
|
||||
let (project_entry_id, rope) = rope_task.await?;
|
||||
anyhow::Ok(RulesFileContext {
|
||||
path_in_worktree,
|
||||
text: rope.to_string().trim().to_string(),
|
||||
project_entry_id: project_entry_id.to_usize(),
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_project_event(
|
||||
&mut self,
|
||||
_project: Entity<Project>,
|
||||
event: &project::Event,
|
||||
_cx: &mut Context<Self>,
|
||||
) {
|
||||
match event {
|
||||
project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
|
||||
self.project_context_needs_refresh.send(()).ok();
|
||||
}
|
||||
project::Event::WorktreeUpdatedEntries(_, items) => {
|
||||
if items.iter().any(|(path, _, _)| {
|
||||
RULES_FILE_NAMES
|
||||
.iter()
|
||||
.any(|name| path.as_ref() == Path::new(name))
|
||||
}) {
|
||||
self.project_context_needs_refresh.send(()).ok();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_prompts_updated_event(
|
||||
&mut self,
|
||||
_prompt_store: Entity<PromptStore>,
|
||||
_event: &prompt_store::PromptsUpdatedEvent,
|
||||
_cx: &mut Context<Self>,
|
||||
) {
|
||||
self.project_context_needs_refresh.send(()).ok();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -120,8 +373,21 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
|
|||
|
||||
cx.spawn(async move |cx| {
|
||||
log::debug!("Starting thread creation in async context");
|
||||
|
||||
// Generate session ID
|
||||
let session_id = acp::SessionId(uuid::Uuid::new_v4().to_string().into());
|
||||
log::info!("Created session with ID: {}", session_id);
|
||||
|
||||
// Create AcpThread
|
||||
let acp_thread = cx.update(|cx| {
|
||||
cx.new(|cx| {
|
||||
acp_thread::AcpThread::new("agent2", self.clone(), project.clone(), session_id.clone(), cx)
|
||||
})
|
||||
})?;
|
||||
let action_log = cx.update(|cx| acp_thread.read(cx).action_log().clone())?;
|
||||
|
||||
// Create Thread
|
||||
let (session_id, thread) = agent.update(
|
||||
let thread = agent.update(
|
||||
cx,
|
||||
|agent, cx: &mut gpui::Context<NativeAgent>| -> Result<_> {
|
||||
// Fetch default model from registry settings
|
||||
|
|
@ -146,22 +412,11 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
|
|||
anyhow!("No default model configured. Please configure a default model in settings.")
|
||||
})?;
|
||||
|
||||
let thread = cx.new(|_| Thread::new(project.clone(), agent.templates.clone(), default_model));
|
||||
|
||||
// Generate session ID
|
||||
let session_id = acp::SessionId(uuid::Uuid::new_v4().to_string().into());
|
||||
log::info!("Created session with ID: {}", session_id);
|
||||
Ok((session_id, thread))
|
||||
let thread = cx.new(|_| Thread::new(project, agent.project_context.clone(), action_log, agent.templates.clone(), default_model));
|
||||
Ok(thread)
|
||||
},
|
||||
)??;
|
||||
|
||||
// Create AcpThread
|
||||
let acp_thread = cx.update(|cx| {
|
||||
cx.new(|cx| {
|
||||
acp_thread::AcpThread::new("agent2", self.clone(), project, session_id.clone(), cx)
|
||||
})
|
||||
})?;
|
||||
|
||||
// Store the session
|
||||
agent.update(cx, |agent, cx| {
|
||||
agent.sessions.insert(
|
||||
|
|
@ -264,6 +519,28 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
|
|||
)
|
||||
})??;
|
||||
}
|
||||
AgentResponseEvent::ToolCallAuthorization(ToolCallAuthorization {
|
||||
tool_call,
|
||||
options,
|
||||
response,
|
||||
}) => {
|
||||
let recv = acp_thread.update(cx, |thread, cx| {
|
||||
thread.request_tool_call_authorization(tool_call, options, cx)
|
||||
})?;
|
||||
cx.background_spawn(async move {
|
||||
if let Some(option) = recv
|
||||
.await
|
||||
.context("authorization sender was dropped")
|
||||
.log_err()
|
||||
{
|
||||
response
|
||||
.send(option)
|
||||
.map(|_| anyhow!("authorization receiver was dropped"))
|
||||
.log_err();
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
AgentResponseEvent::ToolCall(tool_call) => {
|
||||
acp_thread.update(cx, |thread, cx| {
|
||||
thread.handle_session_update(
|
||||
|
|
@ -343,3 +620,77 @@ fn convert_prompt_to_message(blocks: Vec<acp::ContentBlock>) -> String {
|
|||
|
||||
message
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use fs::FakeFs;
|
||||
use gpui::TestAppContext;
|
||||
use serde_json::json;
|
||||
use settings::SettingsStore;
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_maintaining_project_context(cx: &mut TestAppContext) {
|
||||
init_test(cx);
|
||||
let fs = FakeFs::new(cx.executor());
|
||||
fs.insert_tree(
|
||||
"/",
|
||||
json!({
|
||||
"a": {}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let project = Project::test(fs.clone(), [], cx).await;
|
||||
let agent = NativeAgent::new(project.clone(), Templates::new(), None, &mut cx.to_async())
|
||||
.await
|
||||
.unwrap();
|
||||
agent.read_with(cx, |agent, _| {
|
||||
assert_eq!(agent.project_context.borrow().worktrees, vec![])
|
||||
});
|
||||
|
||||
let worktree = project
|
||||
.update(cx, |project, cx| project.create_worktree("/a", true, cx))
|
||||
.await
|
||||
.unwrap();
|
||||
cx.run_until_parked();
|
||||
agent.read_with(cx, |agent, _| {
|
||||
assert_eq!(
|
||||
agent.project_context.borrow().worktrees,
|
||||
vec![WorktreeContext {
|
||||
root_name: "a".into(),
|
||||
abs_path: Path::new("/a").into(),
|
||||
rules_file: None
|
||||
}]
|
||||
)
|
||||
});
|
||||
|
||||
// Creating `/a/.rules` updates the project context.
|
||||
fs.insert_file("/a/.rules", Vec::new()).await;
|
||||
cx.run_until_parked();
|
||||
agent.read_with(cx, |agent, cx| {
|
||||
let rules_entry = worktree.read(cx).entry_for_path(".rules").unwrap();
|
||||
assert_eq!(
|
||||
agent.project_context.borrow().worktrees,
|
||||
vec![WorktreeContext {
|
||||
root_name: "a".into(),
|
||||
abs_path: Path::new("/a").into(),
|
||||
rules_file: Some(RulesFileContext {
|
||||
path_in_worktree: Path::new(".rules").into(),
|
||||
text: "".into(),
|
||||
project_entry_id: rules_entry.id.to_usize()
|
||||
})
|
||||
}]
|
||||
)
|
||||
});
|
||||
}
|
||||
|
||||
fn init_test(cx: &mut TestAppContext) {
|
||||
env_logger::try_init().ok();
|
||||
cx.update(|cx| {
|
||||
let settings_store = SettingsStore::test(cx);
|
||||
cx.set_global(settings_store);
|
||||
Project::init_settings(cx);
|
||||
language::init(cx);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue