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

acp: Allow editing of thread titles in agent2 (#36706)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Richard Feldman <oss@rtfeldman.com>
This commit is contained in:
Antonio Scandurra 2025-08-21 22:24:13 +02:00 • committed by GitHub
parent 555692fac6
commit 731b5d0def
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 254 additions and 105 deletions
Download patch file Download diff file Expand all files Collapse all files

71
crates/agent2/src/agent.rs
View file

@ -2,7 +2,7 @@ use crate::{
ContextServerRegistry, Thread, ThreadEvent, ThreadsDatabase, ToolCallAuthorization,
UserMessageContent, templates::Templates,
};
use crate::{HistoryStore, TokenUsageUpdated};
use crate::{HistoryStore, TitleUpdated, TokenUsageUpdated};
use acp_thread::{AcpThread, AgentModelSelector};
use action_log::ActionLog;
use agent_client_protocol as acp;
@ -253,6 +253,7 @@ impl NativeAgent {
cx.observe_release(&acp_thread, |this, acp_thread, _cx| {
this.sessions.remove(acp_thread.session_id());
}),
cx.subscribe(&thread_handle, Self::handle_thread_title_updated),
cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
cx.observe(&thread_handle, move |this, thread, cx| {
this.save_thread(thread, cx)
@ -441,6 +442,26 @@ impl NativeAgent {
})
}
fn handle_thread_title_updated(
&mut self,
thread: Entity<Thread>,
_: &TitleUpdated,
cx: &mut Context<Self>,
) {
let session_id = thread.read(cx).id();
let Some(session) = self.sessions.get(session_id) else {
return;
};
let thread = thread.downgrade();
let acp_thread = session.acp_thread.clone();
cx.spawn(async move |_, cx| {
let title = thread.read_with(cx, |thread, _| thread.title())?;
let task = acp_thread.update(cx, |acp_thread, cx| acp_thread.set_title(title, cx))?;
task.await
})
.detach_and_log_err(cx);
}
fn handle_thread_token_usage_updated(
&mut self,
thread: Entity<Thread>,
@ -717,10 +738,6 @@ impl NativeAgentConnection {
thread.update_tool_call(update, cx)
})??;
}
ThreadEvent::TitleUpdate(title) => {
acp_thread
.update(cx, |thread, cx| thread.update_title(title, cx))??;
}
ThreadEvent::Retry(status) => {
acp_thread.update(cx, |thread, cx| {
thread.update_retry_status(status, cx)
@ -856,8 +873,7 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
.models
.model_from_id(&LanguageModels::model_id(&default_model.model))
});
let thread = cx.new(|cx| {
Ok(cx.new(|cx| {
Thread::new(
project.clone(),
agent.project_context.clone(),
@ -867,9 +883,7 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
default_model,
cx,
)
});
Ok(thread)
}))
},
)??;
agent.update(cx, |agent, cx| agent.register_session(thread, cx))
@ -941,11 +955,11 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
});
}
fn session_editor(
fn truncate(
&self,
session_id: &agent_client_protocol::SessionId,
cx: &mut App,
) -> Option<Rc<dyn acp_thread::AgentSessionEditor>> {
) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
self.0.update(cx, |agent, _cx| {
agent.sessions.get(session_id).map(|session| {
Rc::new(NativeAgentSessionEditor {
@ -956,6 +970,17 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
})
}
fn set_title(
&self,
session_id: &acp::SessionId,
_cx: &mut App,
) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
Some(Rc::new(NativeAgentSessionSetTitle {
connection: self.clone(),
session_id: session_id.clone(),
}) as _)
}
fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
}
@ -991,8 +1016,8 @@ struct NativeAgentSessionEditor {
acp_thread: WeakEntity<AcpThread>,
}
impl acp_thread::AgentSessionEditor for NativeAgentSessionEditor {
fn truncate(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
impl acp_thread::AgentSessionTruncate for NativeAgentSessionEditor {
fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
match self.thread.update(cx, |thread, cx| {
thread.truncate(message_id.clone(), cx)?;
Ok(thread.latest_token_usage())
@ -1024,6 +1049,22 @@ impl acp_thread::AgentSessionResume for NativeAgentSessionResume {
}
}
struct NativeAgentSessionSetTitle {
connection: NativeAgentConnection,
session_id: acp::SessionId,
}
impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
let Some(session) = self.connection.0.read(cx).sessions.get(&self.session_id) else {
return Task::ready(Err(anyhow!("session not found")));
};
let thread = session.thread.clone();
thread.update(cx, |thread, cx| thread.set_title(title, cx));
Task::ready(Ok(()))
}
}
#[cfg(test)]
mod tests {
use crate::HistoryEntryId;
@ -1323,6 +1364,8 @@ mod tests {
)
});
cx.run_until_parked();
// Drop the ACP thread, which should cause the session to be dropped as well.
cx.update(|_| {
drop(thread);

1
crates/agent2/src/tests/mod.rs
View file

@ -1383,6 +1383,7 @@ async fn test_title_generation(cx: &mut TestAppContext) {
summary_model.send_last_completion_stream_text_chunk("oodnight Moon");
summary_model.end_last_completion_stream();
send.collect::<Vec<_>>().await;
cx.run_until_parked();
thread.read_with(cx, |thread, _| assert_eq!(thread.title(), "Hello world"));
// Send another message, ensuring no title is generated this time.

109
crates/agent2/src/thread.rs
View file

@ -487,7 +487,6 @@ pub enum ThreadEvent {
ToolCall(acp::ToolCall),
ToolCallUpdate(acp_thread::ToolCallUpdate),
ToolCallAuthorization(ToolCallAuthorization),
TitleUpdate(SharedString),
Retry(acp_thread::RetryStatus),
Stop(acp::StopReason),
}
@ -514,6 +513,7 @@ pub struct Thread {
prompt_id: PromptId,
updated_at: DateTime<Utc>,
title: Option<SharedString>,
pending_title_generation: Option<Task<()>>,
summary: Option<SharedString>,
messages: Vec<Message>,
completion_mode: CompletionMode,
@ -555,6 +555,7 @@ impl Thread {
prompt_id: PromptId::new(),
updated_at: Utc::now(),
title: None,
pending_title_generation: None,
summary: None,
messages: Vec::new(),
completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
@ -705,6 +706,7 @@ impl Thread {
} else {
Some(db_thread.title.clone())
},
pending_title_generation: None,
summary: db_thread.detailed_summary,
messages: db_thread.messages,
completion_mode: db_thread.completion_mode.unwrap_or_default(),
@ -1086,7 +1088,7 @@ impl Thread {
event_stream: event_stream.clone(),
_task: cx.spawn(async move |this, cx| {
log::info!("Starting agent turn execution");
let mut update_title = None;
let turn_result: Result<()> = async {
let mut intent = CompletionIntent::UserPrompt;
loop {
@ -1095,8 +1097,8 @@ impl Thread {
let mut end_turn = true;
this.update(cx, |this, cx| {
// Generate title if needed.
if this.title.is_none() && update_title.is_none() {
update_title = Some(this.update_title(&event_stream, cx));
if this.title.is_none() && this.pending_title_generation.is_none() {
this.generate_title(cx);
}
// End the turn if the model didn't use tools.
@ -1120,10 +1122,6 @@ impl Thread {
.await;
_ = this.update(cx, |this, cx| this.flush_pending_message(cx));
if let Some(update_title) = update_title {
update_title.await.context("update title failed").log_err();
}
match turn_result {
Ok(()) => {
log::info!("Turn execution completed");
@ -1607,19 +1605,15 @@ impl Thread {
})
}
fn update_title(
&mut self,
event_stream: &ThreadEventStream,
cx: &mut Context<Self>,
) -> Task<Result<()>> {
fn generate_title(&mut self, cx: &mut Context<Self>) {
let Some(model) = self.summarization_model.clone() else {
return;
};
log::info!(
"Generating title with model: {:?}",
self.summarization_model.as_ref().map(|model| model.name())
);
let Some(model) = self.summarization_model.clone() else {
return Task::ready(Ok(()));
};
let event_stream = event_stream.clone();
let mut request = LanguageModelRequest {
intent: Some(CompletionIntent::ThreadSummarization),
temperature: AgentSettings::temperature_for_model(&model, cx),
@ -1635,42 +1629,51 @@ impl Thread {
content: vec![SUMMARIZE_THREAD_PROMPT.into()],
cache: false,
});
cx.spawn(async move |this, cx| {
self.pending_title_generation = Some(cx.spawn(async move |this, cx| {
let mut title = String::new();
let mut messages = model.stream_completion(request, cx).await?;
while let Some(event) = messages.next().await {
let event = event?;
let text = match event {
LanguageModelCompletionEvent::Text(text) => text,
LanguageModelCompletionEvent::StatusUpdate(
CompletionRequestStatus::UsageUpdated { amount, limit },
) => {
this.update(cx, |thread, cx| {
thread.update_model_request_usage(amount, limit, cx);
})?;
continue;
let generate = async {
let mut messages = model.stream_completion(request, cx).await?;
while let Some(event) = messages.next().await {
let event = event?;
let text = match event {
LanguageModelCompletionEvent::Text(text) => text,
LanguageModelCompletionEvent::StatusUpdate(
CompletionRequestStatus::UsageUpdated { amount, limit },
) => {
this.update(cx, |thread, cx| {
thread.update_model_request_usage(amount, limit, cx);
})?;
continue;
}
_ => continue,
};
let mut lines = text.lines();
title.extend(lines.next());
// Stop if the LLM generated multiple lines.
if lines.next().is_some() {
break;
}
_ => continue,
};
let mut lines = text.lines();
title.extend(lines.next());
// Stop if the LLM generated multiple lines.
if lines.next().is_some() {
break;
}
anyhow::Ok(())
};
if generate.await.context("failed to generate title").is_ok() {
_ = this.update(cx, |this, cx| this.set_title(title.into(), cx));
}
_ = this.update(cx, |this, _| this.pending_title_generation = None);
}));
}
log::info!("Setting title: {}", title);
this.update(cx, |this, cx| {
let title = SharedString::from(title);
event_stream.send_title_update(title.clone());
this.title = Some(title);
cx.notify();
})
})
pub fn set_title(&mut self, title: SharedString, cx: &mut Context<Self>) {
self.pending_title_generation = None;
if Some(&title) != self.title.as_ref() {
self.title = Some(title);
cx.emit(TitleUpdated);
cx.notify();
}
}
fn last_user_message(&self) -> Option<&UserMessage> {
@ -1975,6 +1978,10 @@ pub struct TokenUsageUpdated(pub Option<acp_thread::TokenUsage>);
impl EventEmitter<TokenUsageUpdated> for Thread {}
pub struct TitleUpdated;
impl EventEmitter<TitleUpdated> for Thread {}
pub trait AgentTool
where
Self: 'static + Sized,
@ -2132,12 +2139,6 @@ where
struct ThreadEventStream(mpsc::UnboundedSender<Result<ThreadEvent>>);
impl ThreadEventStream {
fn send_title_update(&self, text: SharedString) {
self.0
.unbounded_send(Ok(ThreadEvent::TitleUpdate(text)))
.ok();
}
fn send_user_message(&self, message: &UserMessage) {
self.0
.unbounded_send(Ok(ThreadEvent::UserMessage(message.clone())))