diff --git a/crates/copilot_chat/src/copilot_chat.rs b/crates/copilot_chat/src/copilot_chat.rs index a690333e0c..92073f5c0b 100644 --- a/crates/copilot_chat/src/copilot_chat.rs +++ b/crates/copilot_chat/src/copilot_chat.rs @@ -41,12 +41,16 @@ impl CopilotChatConfiguration { pub fn graphql_url(&self) -> String { if let Some(enterprise_uri) = &self.enterprise_uri { let domain = Self::parse_domain(enterprise_uri); - format!("https://{}/api/graphql", domain) + format!("https://api.{domain}/graphql") } else { "https://api.github.com/graphql".to_string() } } + fn credentials_url(&self) -> String { + format!("https://{}/copilot-agent", self.oauth_domain()) + } + pub fn chat_completions_url(&self, api_endpoint: &str) -> String { format!("{}/chat/completions", api_endpoint) } @@ -505,11 +509,6 @@ struct GlobalCopilotChat(gpui::Entity); impl Global for GlobalCopilotChat {} -/// The keychain URL under which the Copilot agent's OAuth token is stored. This -/// is intentionally distinct from the edit-prediction token so the two Copilot -/// providers are authenticated entirely separately. -const COPILOT_AGENT_CREDENTIALS_URL: &str = "https://github.com/copilot-agent"; - /// Authentication state for the Copilot agent (chat) provider. #[derive(Clone, Debug)] pub enum CopilotChatStatus { @@ -553,10 +552,11 @@ fn oauth_token_from_env() -> Option { async fn load_stored_token( credentials_provider: &Arc, + configuration: &CopilotChatConfiguration, cx: &AsyncApp, ) -> Option { let (_, token) = credentials_provider - .read_credentials(COPILOT_AGENT_CREDENTIALS_URL, cx) + .read_credentials(&configuration.credentials_url(), cx) .await .ok() .flatten()?; @@ -580,16 +580,24 @@ impl CopilotChat { // Load a previously-stored token (or the one from the environment) and // fetch models if we end up authenticated. cx.spawn(async move |this, cx| { - let (env_token, credentials_provider) = this.read_with(cx, |this, _| { - (this.oauth_token.clone(), this.credentials_provider.clone()) - })?; + let (env_token, credentials_provider, configuration) = + this.read_with(cx, |this, _| { + ( + this.oauth_token.clone(), + this.credentials_provider.clone(), + this.configuration.clone(), + ) + })?; let token = match env_token { Some(token) => Some(token), - None => load_stored_token(&credentials_provider, cx).await, + None => load_stored_token(&credentials_provider, &configuration, cx).await, }; - this.update(cx, |this, cx| { + let configuration_is_current = this.update(cx, |this, cx| { + if this.configuration != configuration { + return false; + } this.oauth_token = token.clone(); this.status = if token.is_some() { CopilotChatStatus::Authorized @@ -597,9 +605,10 @@ impl CopilotChat { CopilotChatStatus::SignedOut }; cx.notify(); + true })?; - if token.is_some() { + if configuration_is_current && token.is_some() { Self::update_models(&this, cx).await?; } anyhow::Ok(()) @@ -634,6 +643,7 @@ impl CopilotChat { let client = self.client.clone(); let configuration = self.configuration.clone(); + let credentials_url = configuration.credentials_url(); let credentials_provider = self.credentials_provider.clone(); let executor = cx.background_executor().clone(); @@ -657,12 +667,7 @@ impl CopilotChat { .await?; credentials_provider - .write_credentials( - COPILOT_AGENT_CREDENTIALS_URL, - "Bearer", - token.as_bytes(), - cx, - ) + .write_credentials(&credentials_url, "Bearer", token.as_bytes(), cx) .await .context("writing Copilot agent credentials to the keychain")?; @@ -694,6 +699,7 @@ impl CopilotChat { pub fn sign_out(&mut self, cx: &mut Context) -> Task> { let credentials_provider = self.credentials_provider.clone(); + let credentials_url = self.configuration.credentials_url(); self.oauth_token = None; self.api_endpoint = None; self.models = None; @@ -703,7 +709,7 @@ impl CopilotChat { cx.spawn(async move |_, cx| { credentials_provider - .delete_credentials(COPILOT_AGENT_CREDENTIALS_URL, cx) + .delete_credentials(&credentials_url, cx) .await .context("deleting Copilot agent credentials from the keychain")?; anyhow::Ok(()) @@ -852,6 +858,7 @@ impl CopilotChat { ) -> Result { let api_endpoint = match discover_api_endpoint(oauth_token, configuration, client).await { Ok(endpoint) => endpoint, + Err(error) if configuration.enterprise_uri.is_some() => return Err(error), Err(error) => { log::warn!( "Failed to discover Copilot API endpoint via GraphQL, \ @@ -877,12 +884,42 @@ impl CopilotChat { let same_configuration = self.configuration == configuration; self.configuration = configuration; if !same_configuration { + self.oauth_token = None; self.api_endpoint = None; + self.models = None; + self.sign_in_task = None; + self.status = CopilotChatStatus::Starting; + cx.notify(); + + let credentials_provider = self.credentials_provider.clone(); + let configuration = self.configuration.clone(); cx.spawn(async move |this, cx| { - Self::update_models(&this, cx).await?; + let token = load_stored_token(&credentials_provider, &configuration, cx).await; + let configuration_is_current = this.update(cx, |this, cx| { + if this.configuration != configuration { + return false; + } + this.oauth_token = token.clone(); + this.status = if token.is_some() { + CopilotChatStatus::Authorized + } else { + CopilotChatStatus::SignedOut + }; + cx.notify(); + true + })?; + + if configuration_is_current && token.is_some() { + if let Err(error) = Self::update_models(&this, cx).await { + this.update(cx, |this, cx| { + this.status = CopilotChatStatus::Error(error.to_string().into()); + cx.notify(); + })?; + } + } Ok::<_, anyhow::Error>(()) }) - .detach(); + .detach_and_log_err(cx); } } } @@ -1184,6 +1221,34 @@ async fn stream_messages( mod tests { use super::*; + #[test] + fn test_github_urls() { + let configuration = CopilotChatConfiguration::default(); + assert_eq!( + configuration.graphql_url(), + "https://api.github.com/graphql" + ); + assert_eq!( + configuration.credentials_url(), + "https://github.com/copilot-agent" + ); + } + + #[test] + fn test_ghe_data_residency_urls() { + let configuration = CopilotChatConfiguration { + enterprise_uri: Some("https://acme.ghe.com/".to_string()), + }; + assert_eq!( + configuration.graphql_url(), + "https://api.acme.ghe.com/graphql" + ); + assert_eq!( + configuration.credentials_url(), + "https://acme.ghe.com/copilot-agent" + ); + } + #[test] fn test_resilient_model_schema_deserialize() { let json = r#"{