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

Release Notes:

- N/A

---------

Co-authored-by: Richard Feldman <oss@rtfeldman.com>

crates/agent2/src/agent.rs

@@ -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);