diff --git a/crates/acp_thread/src/connection.rs b/crates/acp_thread/src/connection.rs index b2bd00f61a..87c8ccf65c 100644 --- a/crates/acp_thread/src/connection.rs +++ b/crates/acp_thread/src/connection.rs @@ -127,6 +127,14 @@ pub trait AgentConnection { fn authenticate(&self, method: acp::AuthMethodId, cx: &mut App) -> Task>; + fn supports_logout(&self, _cx: &App) -> bool { + false + } + + fn logout(&self, _cx: &mut App) -> Task> { + Task::ready(Err(anyhow::Error::msg("Logout is not supported"))) + } + fn prompt( &self, user_message_id: UserMessageId, diff --git a/crates/agent_servers/src/acp.rs b/crates/agent_servers/src/acp.rs index f0a036ca4d..b3328790a5 100644 --- a/crates/agent_servers/src/acp.rs +++ b/crates/agent_servers/src/acp.rs @@ -1759,6 +1759,22 @@ impl AgentConnection for AcpConnection { }) } + fn supports_logout(&self, cx: &App) -> bool { + cx.has_flag::() && self.agent_capabilities.auth.logout.is_some() + } + + fn logout(&self, cx: &mut App) -> Task> { + if !self.supports_logout(cx) { + return Task::ready(Err(anyhow!("Logout is not supported by this agent."))); + } + + let conn = self.connection.clone(); + cx.foreground_executor().spawn(async move { + into_foreground_future(conn.send_request(acp::LogoutRequest::new())).await?; + Ok(()) + }) + } + fn prompt( &self, _id: acp_thread::UserMessageId, @@ -2026,6 +2042,7 @@ pub mod test_support { pub connection: Rc, pub load_session_count: Arc, pub close_session_count: Arc, + pub logout_count: Arc, pub keep_agent_alive: Task>, } @@ -2119,6 +2136,14 @@ pub mod test_support { self.inner.authenticate(method, cx) } + fn supports_logout(&self, cx: &App) -> bool { + self.inner.supports_logout(cx) + } + + fn logout(&self, cx: &mut App) -> Task> { + self.inner.logout(cx) + } + fn prompt( &self, user_message_id: UserMessageId, @@ -2201,6 +2226,7 @@ pub mod test_support { ) -> Result { let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let logout_count = Arc::new(AtomicUsize::new(0)); let sessions: Rc>> = Rc::new(RefCell::new(HashMap::default())); let client_session_list: Rc>>> = @@ -2269,6 +2295,16 @@ pub mod test_support { }, agent_client_protocol::on_receive_request!(), ) + .on_receive_request( + { + let logout_count = logout_count.clone(); + async move |_req: acp::LogoutRequest, responder, _cx| { + logout_count.fetch_add(1, Ordering::SeqCst); + responder.respond(acp::LogoutResponse::new()) + } + }, + agent_client_protocol::on_receive_request!(), + ) .on_receive_notification( async move |_notif: acp::CancelNotification, _cx| Ok(()), agent_client_protocol::on_receive_notification!(), @@ -2341,6 +2377,7 @@ pub mod test_support { connection: Rc::new(connection), load_session_count, close_session_count, + logout_count, keep_agent_alive, }) } @@ -2373,6 +2410,7 @@ mod tests { use feature_flags::FeatureFlag as _; use super::*; + use gpui::UpdateGlobal as _; #[test] fn terminal_auth_task_builds_spawn_from_prebuilt_command() { @@ -2641,6 +2679,76 @@ mod tests { ); } + #[gpui::test] + async fn logout_is_gated_by_beta_flag_and_agent_capability(cx: &mut gpui::TestAppContext) { + cx.update(|cx| { + let store = settings::SettingsStore::test(cx); + cx.set_global(store); + settings::SettingsStore::update_global(cx, |store, _| { + store.register_setting::(); + }); + feature_flags::FeatureFlagStore::init(cx); + }); + + let fs = fs::FakeFs::new(cx.executor()); + fs.insert_tree("/", serde_json::json!({ "a": {} })).await; + let project = project::Project::test(fs, [std::path::Path::new("/a")], cx).await; + let mut harness = test_support::connect_fake_acp_connection(project, cx).await; + cx.update(|cx| { + settings::SettingsStore::update_global(cx, |store, _| { + store.register_setting::(); + }); + feature_flags::FeatureFlagStore::init(cx); + }); + + assert!(!cx.update(|cx| harness.connection.supports_logout(cx))); + let unsupported_logout = cx.update(|cx| harness.connection.logout(cx)); + let error = unsupported_logout + .await + .expect_err("logout should be rejected when the agent does not advertise support"); + assert_eq!(error.to_string(), "Logout is not supported by this agent."); + assert_eq!(harness.logout_count.load(Ordering::SeqCst), 0); + + Rc::get_mut(&mut harness.connection) + .expect("test harness should own the only ACP connection handle") + .agent_capabilities + .auth = acp::AgentAuthCapabilities::new().logout(acp::LogoutCapabilities::new()); + + cx.update(|cx| { + settings::SettingsStore::update_global(cx, |store, cx| { + store.update_user_settings(cx, |content| { + content + .feature_flags + .get_or_insert_default() + .insert("acp-beta".to_string(), "off".to_string()); + }); + }); + }); + assert!(!cx.update(|cx| harness.connection.supports_logout(cx))); + let disabled_logout = cx.update(|cx| harness.connection.logout(cx)); + let error = disabled_logout + .await + .expect_err("logout should be rejected when acp-beta is disabled"); + assert_eq!(error.to_string(), "Logout is not supported by this agent."); + assert_eq!(harness.logout_count.load(Ordering::SeqCst), 0); + + cx.update(|cx| { + settings::SettingsStore::update_global(cx, |store, cx| { + store.update_user_settings(cx, |content| { + content + .feature_flags + .get_or_insert_default() + .insert("acp-beta".to_string(), "on".to_string()); + }); + }); + }); + assert!(cx.update(|cx| harness.connection.supports_logout(cx))); + cx.update(|cx| harness.connection.logout(cx)) + .await + .expect("logout should be sent when the agent advertises support"); + assert_eq!(harness.logout_count.load(Ordering::SeqCst), 1); + } + #[cfg(not(windows))] #[gpui::test] async fn startup_returns_error_when_agent_exits_before_initialization( diff --git a/crates/agent_ui/src/agent_panel.rs b/crates/agent_ui/src/agent_panel.rs index 88c5baa2b2..635ad042b8 100644 --- a/crates/agent_ui/src/agent_panel.rs +++ b/crates/agent_ui/src/agent_panel.rs @@ -23,9 +23,9 @@ use settings::{LanguageModelProviderSetting, LanguageModelSelection}; use zed_actions::{ DecreaseBufferFontSize, IncreaseBufferFontSize, ResetBufferFontSize, agent::{ - AddSelectionToThread, ConflictContent, OpenSettings, ReauthenticateAgent, ResetAgentZoom, - ResetOnboarding, ResolveConflictedFilesWithAgent, ResolveConflictsWithAgent, - ReviewBranchDiff, + AddSelectionToThread, ConflictContent, LogoutAgent, OpenSettings, ReauthenticateAgent, + ResetAgentZoom, ResetOnboarding, ResolveConflictedFilesWithAgent, + ResolveConflictsWithAgent, ReviewBranchDiff, }, assistant::{FocusAgent, OpenRulesLibrary, Toggle, ToggleFocus}, }; @@ -4161,6 +4161,9 @@ impl AgentPanel { } _ => false, }; + let supports_logout = self + .active_conversation_view() + .is_some_and(|conversation_view| conversation_view.read(cx).supports_logout(cx)); PopoverMenu::new("agent-options-menu") .trigger_with_tooltip( @@ -4239,6 +4242,9 @@ impl AgentPanel { if has_auth_methods { menu = menu.action("Reauthenticate", Box::new(ReauthenticateAgent)) } + if supports_logout { + menu = menu.action("Log Out", Box::new(LogoutAgent)) + } menu })) @@ -5063,6 +5069,13 @@ impl Render for AgentPanel { }) } })) + .on_action(cx.listener(|this, _: &LogoutAgent, window, cx| { + if let Some(conversation_view) = this.active_conversation_view() { + conversation_view.update(cx, |conversation_view, cx| { + conversation_view.logout(window, cx) + }) + } + })) .child(self.render_toolbar(window, cx)) .children(self.render_new_user_onboarding(window, cx)) .map(|parent| match self.visible_surface() { diff --git a/crates/agent_ui/src/conversation_view.rs b/crates/agent_ui/src/conversation_view.rs index 394281abad..6ecb85267a 100644 --- a/crates/agent_ui/src/conversation_view.rs +++ b/crates/agent_ui/src/conversation_view.rs @@ -519,6 +519,12 @@ impl ConversationView { }) } + pub fn supports_logout(&self, cx: &App) -> bool { + self.as_connected().is_some_and(|connected| { + connected.auth_state.is_ok() && connected.connection.supports_logout(cx) + }) + } + pub fn active_thread(&self) -> Option<&Entity> { match &self.server_state { ServerState::Connected(connected) => connected.active_view(), @@ -2910,6 +2916,54 @@ impl ConversationView { Self::handle_auth_required(this, AuthRequired::new(), agent_id, connection, window, cx); }) } + + pub(crate) fn logout(&mut self, window: &mut Window, cx: &mut Context) { + if !self.supports_logout(cx) { + return; + } + + if let Some(active) = self.root_thread_view() { + active.update(cx, |active, cx| active.clear_thread_error(cx)); + } + let Some(connection) = self + .as_connected() + .map(|connected| connected.connection.clone()) + else { + return; + }; + let logout = connection.logout(cx); + self.auth_task = Some(cx.spawn_in(window, { + async move |this, cx| { + let result = logout.await; + this.update_in(cx, |this, window, cx| { + if let Err(err) = result { + if let Some(active) = this.root_thread_view() { + active.update(cx, |active, cx| active.handle_thread_error(err, cx)); + } + } else if let Some(connected) = this.as_connected_mut() { + connected.auth_state = AuthState::Unauthenticated { + description: None, + configuration_view: None, + pending_auth_method: None, + _subscription: None, + }; + if let Some(view) = connected.active_view() + && view + .read(cx) + .message_editor + .focus_handle(cx) + .is_focused(window) + { + this.focus_handle.focus(window, cx) + } + cx.notify(); + } + drop(this.auth_task.take()); + }) + .ok(); + } + })); + } } fn loading_contents_spinner(size: IconSize) -> AnyElement { @@ -3788,7 +3842,7 @@ pub(crate) mod tests { // When new_session returns AuthRequired, the server should transition // to Connected + Unauthenticated rather than getting stuck in Loading. - conversation_view.read_with(cx, |view, _cx| { + conversation_view.read_with(cx, |view, cx| { let connected = view .as_connected() .expect("Should be in Connected state even though auth is required"); @@ -3796,6 +3850,10 @@ pub(crate) mod tests { !connected.auth_state.is_ok(), "Auth state should be Unauthenticated" ); + assert!( + !view.supports_logout(cx), + "Logout should be hidden while unauthenticated" + ); assert!( connected.active_id.is_none(), "There should be no active thread since no session was created" @@ -3831,6 +3889,10 @@ pub(crate) mod tests { .as_connected() .expect("Should still be in Connected state after auth"); assert!(connected.auth_state.is_ok(), "Auth state should be Ok"); + assert!( + view.supports_logout(cx), + "Logout should be available after authentication" + ); assert!( connected.active_id.is_some(), "There should be an active thread after successful auth" @@ -3849,6 +3911,23 @@ pub(crate) mod tests { "The new thread should have no errors" ); }); + + conversation_view.update_in(cx, |view, window, cx| view.logout(window, cx)); + cx.run_until_parked(); + + conversation_view.read_with(cx, |view, cx| { + let connected = view + .as_connected() + .expect("Should still be in Connected state after logout"); + assert!( + !connected.auth_state.is_ok(), + "Auth state should be Unauthenticated after logout" + ); + assert!( + !view.supports_logout(cx), + "Logout should be hidden after logout" + ); + }); } #[gpui::test] @@ -4910,6 +4989,15 @@ pub(crate) mod tests { } } + fn supports_logout(&self, _cx: &App) -> bool { + true + } + + fn logout(&self, _cx: &mut App) -> Task> { + *self.authenticated.lock() = false; + Task::ready(Ok(())) + } + fn prompt( &self, _id: acp_thread::UserMessageId, diff --git a/crates/zed_actions/src/lib.rs b/crates/zed_actions/src/lib.rs index 0803f9edd9..bcfaad5e8a 100644 --- a/crates/zed_actions/src/lib.rs +++ b/crates/zed_actions/src/lib.rs @@ -505,6 +505,8 @@ pub mod agent { ToggleModelSelector, /// Triggers re-authentication on Gemini ReauthenticateAgent, + /// Logs out of the current external agent + LogoutAgent, /// Add the current selection as context for threads in the agent panel. #[action(deprecated_aliases = ["assistant::QuoteSelection", "agent::QuoteSelection"])] AddSelectionToThread,