>,
+)> {
+ let server = tiny_http::Server::http("127.0.0.1:0")
+ .map_err(|e| anyhow!(e).context("Failed to bind loopback listener for OAuth callback"))?;
+ let port = server
+ .server_addr()
+ .to_ip()
+ .context("server not bound to a TCP address")?
+ .port();
+
+ let redirect_uri = format!("http://127.0.0.1:{}/callback", port);
+
+ let (tx, rx) = futures::channel::oneshot::channel();
+
+ // `tiny_http` is blocking, so we run it on a background thread.
+ // The `recv_timeout` loop lets us check for cancellation (the receiver
+ // being dropped) and enforce an overall timeout.
+ std::thread::spawn(move || {
+ let deadline = std::time::Instant::now() + CALLBACK_TIMEOUT;
+
+ loop {
+ if tx.is_canceled() {
+ return;
+ }
+ let remaining = deadline.saturating_duration_since(std::time::Instant::now());
+ if remaining.is_zero() {
+ return;
+ }
+
+ let timeout = remaining.min(Duration::from_millis(500));
+ let Some(request) = (match server.recv_timeout(timeout) {
+ Ok(req) => req,
+ Err(_) => {
+ let _ = tx.send(Err(anyhow!("OAuth callback server I/O error")));
+ return;
+ }
+ }) else {
+ // Timeout with no request — loop back and check cancellation.
+ continue;
+ };
+
+ let result = handle_callback_request(&request);
+
+ let (status_code, body) = match &result {
+ Ok(_) => (
+ 200,
+ "Authorization successful
\
+ You can close this tab and return to Zed.
",
+ ),
+ Err(err) => {
+ log::error!("OAuth callback error: {}", err);
+ (
+ 400,
+ "Authorization failed
\
+ Something went wrong. Please try again from Zed.
",
+ )
+ }
+ };
+
+ let response = tiny_http::Response::from_string(body)
+ .with_status_code(status_code)
+ .with_header(
+ tiny_http::Header::from_str("Content-Type: text/html")
+ .expect("failed to construct response header"),
+ )
+ .with_header(
+ tiny_http::Header::from_str("Keep-Alive: timeout=0,max=0")
+ .expect("failed to construct response header"),
+ );
+ request.respond(response).log_err();
+
+ let _ = tx.send(result);
+ return;
+ }
+ });
+
+ Ok((redirect_uri, rx))
+}
+
+/// Extract the `code` and `state` query parameters from an OAuth callback
+/// request to `/callback`.
+fn handle_callback_request(request: &tiny_http::Request) -> Result {
+ let url = Url::parse(&format!("http://localhost{}", request.url()))
+ .context("malformed callback request URL")?;
+
+ if url.path() != "/callback" {
+ bail!("unexpected path in OAuth callback: {}", url.path());
+ }
+
+ let query = url
+ .query()
+ .ok_or_else(|| anyhow!("OAuth callback has no query string"))?;
+ OAuthCallback::parse_query(query)
+}
+
+// -- JSON fetch helper -------------------------------------------------------
+
+async fn fetch_json(
+ http_client: &Arc,
+ url: &Url,
+) -> Result {
+ validate_oauth_url(url)?;
+
+ let request = Request::builder()
+ .method(http_client::http::Method::GET)
+ .uri(url.as_str())
+ .header("Accept", "application/json")
+ .body(AsyncBody::default())?;
+
+ let mut response = http_client.send(request).await?;
+
+ if !response.status().is_success() {
+ bail!("HTTP {} fetching {}", response.status(), url);
+ }
+
+ let mut body = String::new();
+ response.body_mut().read_to_string(&mut body).await?;
+ serde_json::from_str(&body).with_context(|| format!("failed to parse JSON from {}", url))
+}
+
+// -- Serde response types for discovery --------------------------------------
+
+#[derive(Debug, Deserialize)]
+struct ProtectedResourceMetadataResponse {
+ #[serde(default)]
+ resource: Option,
+ #[serde(default)]
+ authorization_servers: Vec,
+ #[serde(default)]
+ scopes_supported: Option>,
+}
+
+#[derive(Debug, Deserialize)]
+struct AuthServerMetadataResponse {
+ #[serde(default)]
+ issuer: Option,
+ #[serde(default)]
+ authorization_endpoint: Option,
+ #[serde(default)]
+ token_endpoint: Option,
+ #[serde(default)]
+ registration_endpoint: Option,
+ #[serde(default)]
+ scopes_supported: Option>,
+ #[serde(default)]
+ code_challenge_methods_supported: Option>,
+ #[serde(default)]
+ client_id_metadata_document_supported: Option,
+}
+
+#[derive(Debug, Deserialize)]
+struct DcrResponse {
+ client_id: String,
+ #[serde(default)]
+ client_secret: Option,
+}
+
+/// Provides OAuth tokens to the HTTP transport layer.
+///
+/// The transport calls `access_token()` before each request. On a 401 response
+/// it calls `try_refresh()` and retries once if the refresh succeeds.
+#[async_trait]
+pub trait OAuthTokenProvider: Send + Sync {
+ /// Returns the current access token, if one is available.
+ fn access_token(&self) -> Option;
+
+ /// Attempts to refresh the access token. Returns `true` if a new token was
+ /// obtained and the request should be retried.
+ async fn try_refresh(&self) -> Result;
+}
+
+/// Concrete `OAuthTokenProvider` backed by a full persisted OAuth session and
+/// an HTTP client for token refresh. The same provider type is used both after
+/// an interactive authentication flow and when restoring a saved session from
+/// the keychain on startup.
+pub struct McpOAuthTokenProvider {
+ session: SyncMutex,
+ http_client: Arc,
+ token_refresh_tx: Option>,
+}
+
+impl McpOAuthTokenProvider {
+ pub fn new(
+ session: OAuthSession,
+ http_client: Arc,
+ token_refresh_tx: Option>,
+ ) -> Self {
+ Self {
+ session: SyncMutex::new(session),
+ http_client,
+ token_refresh_tx,
+ }
+ }
+
+ fn access_token_is_expired(tokens: &OAuthTokens) -> bool {
+ tokens.expires_at.is_some_and(|expires_at| {
+ SystemTime::now()
+ .checked_add(Duration::from_secs(30))
+ .is_some_and(|now_with_buffer| expires_at <= now_with_buffer)
+ })
+ }
+}
+
+#[async_trait]
+impl OAuthTokenProvider for McpOAuthTokenProvider {
+ fn access_token(&self) -> Option {
+ let session = self.session.lock();
+ if Self::access_token_is_expired(&session.tokens) {
+ return None;
+ }
+ Some(session.tokens.access_token.clone())
+ }
+
+ async fn try_refresh(&self) -> Result {
+ let (refresh_token, token_endpoint, resource, client_id) = {
+ let session = self.session.lock();
+ match session.tokens.refresh_token.clone() {
+ Some(refresh_token) => (
+ refresh_token,
+ session.token_endpoint.clone(),
+ session.resource.clone(),
+ session.client_registration.client_id.clone(),
+ ),
+ None => return Ok(false),
+ }
+ };
+
+ let resource_str = canonical_server_uri(&resource);
+
+ match refresh_tokens(
+ &self.http_client,
+ &token_endpoint,
+ &refresh_token,
+ &client_id,
+ &resource_str,
+ )
+ .await
+ {
+ Ok(mut new_tokens) => {
+ if new_tokens.refresh_token.is_none() {
+ new_tokens.refresh_token = Some(refresh_token);
+ }
+
+ {
+ let mut session = self.session.lock();
+ session.tokens = new_tokens;
+
+ if let Some(ref tx) = self.token_refresh_tx {
+ tx.unbounded_send(session.clone()).ok();
+ }
+ }
+
+ Ok(true)
+ }
+ Err(err) => {
+ log::warn!("OAuth token refresh failed: {}", err);
+ Ok(false)
+ }
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use http_client::Response;
+
+ // -- require_https_or_loopback tests ------------------------------------
+
+ #[test]
+ fn test_require_https_or_loopback_accepts_https() {
+ let url = Url::parse("https://auth.example.com/token").unwrap();
+ assert!(require_https_or_loopback(&url).is_ok());
+ }
+
+ #[test]
+ fn test_require_https_or_loopback_rejects_http_remote() {
+ let url = Url::parse("http://auth.example.com/token").unwrap();
+ assert!(require_https_or_loopback(&url).is_err());
+ }
+
+ #[test]
+ fn test_require_https_or_loopback_accepts_http_127_0_0_1() {
+ let url = Url::parse("http://127.0.0.1:8080/callback").unwrap();
+ assert!(require_https_or_loopback(&url).is_ok());
+ }
+
+ #[test]
+ fn test_require_https_or_loopback_accepts_http_ipv6_loopback() {
+ let url = Url::parse("http://[::1]:8080/callback").unwrap();
+ assert!(require_https_or_loopback(&url).is_ok());
+ }
+
+ #[test]
+ fn test_require_https_or_loopback_accepts_http_localhost() {
+ let url = Url::parse("http://localhost:8080/callback").unwrap();
+ assert!(require_https_or_loopback(&url).is_ok());
+ }
+
+ #[test]
+ fn test_require_https_or_loopback_accepts_http_localhost_case_insensitive() {
+ let url = Url::parse("http://LOCALHOST:8080/callback").unwrap();
+ assert!(require_https_or_loopback(&url).is_ok());
+ }
+
+ #[test]
+ fn test_require_https_or_loopback_rejects_http_non_loopback_ip() {
+ let url = Url::parse("http://192.168.1.1:8080/token").unwrap();
+ assert!(require_https_or_loopback(&url).is_err());
+ }
+
+ #[test]
+ fn test_require_https_or_loopback_rejects_ftp() {
+ let url = Url::parse("ftp://auth.example.com/token").unwrap();
+ assert!(require_https_or_loopback(&url).is_err());
+ }
+
+ // -- validate_oauth_url (SSRF) tests ------------------------------------
+
+ #[test]
+ fn test_validate_oauth_url_accepts_https_public() {
+ let url = Url::parse("https://auth.example.com/token").unwrap();
+ assert!(validate_oauth_url(&url).is_ok());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_private_ipv4_10() {
+ let url = Url::parse("https://10.0.0.1/token").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_private_ipv4_172() {
+ let url = Url::parse("https://172.16.0.1/token").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_private_ipv4_192() {
+ let url = Url::parse("https://192.168.1.1/token").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_link_local() {
+ let url = Url::parse("https://169.254.169.254/latest/meta-data/").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_ipv6_ula() {
+ let url = Url::parse("https://[fd12:3456:789a::1]/token").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_ipv6_unspecified() {
+ let url = Url::parse("https://[::]/token").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_ipv4_mapped_ipv6_private() {
+ let url = Url::parse("https://[::ffff:10.0.0.1]/token").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_rejects_ipv4_mapped_ipv6_link_local() {
+ let url = Url::parse("https://[::ffff:169.254.169.254]/token").unwrap();
+ assert!(validate_oauth_url(&url).is_err());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_allows_http_loopback() {
+ // Loopback is permitted (it's our callback server).
+ let url = Url::parse("http://127.0.0.1:8080/callback").unwrap();
+ assert!(validate_oauth_url(&url).is_ok());
+ }
+
+ #[test]
+ fn test_validate_oauth_url_allows_https_public_ip() {
+ let url = Url::parse("https://93.184.216.34/token").unwrap();
+ assert!(validate_oauth_url(&url).is_ok());
+ }
+
+ // -- parse_www_authenticate tests ----------------------------------------
+
+ #[test]
+ fn test_parse_www_authenticate_with_resource_metadata_and_scope() {
+ let header = r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource", scope="files:read user:profile""#;
+ let result = parse_www_authenticate(header).unwrap();
+
+ assert_eq!(
+ result.resource_metadata.as_ref().map(|u| u.as_str()),
+ Some("https://mcp.example.com/.well-known/oauth-protected-resource")
+ );
+ assert_eq!(
+ result.scope,
+ Some(vec!["files:read".to_string(), "user:profile".to_string()])
+ );
+ assert_eq!(result.error, None);
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_resource_metadata_only() {
+ let header = r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#;
+ let result = parse_www_authenticate(header).unwrap();
+
+ assert_eq!(
+ result.resource_metadata.as_ref().map(|u| u.as_str()),
+ Some("https://mcp.example.com/.well-known/oauth-protected-resource")
+ );
+ assert_eq!(result.scope, None);
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_bare_bearer() {
+ let result = parse_www_authenticate("Bearer").unwrap();
+ assert_eq!(result.resource_metadata, None);
+ assert_eq!(result.scope, None);
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_with_error() {
+ let header = r#"Bearer error="insufficient_scope", scope="files:read files:write", resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource", error_description="Additional file write permission required""#;
+ let result = parse_www_authenticate(header).unwrap();
+
+ assert_eq!(result.error, Some(BearerError::InsufficientScope));
+ assert_eq!(
+ result.error_description.as_deref(),
+ Some("Additional file write permission required")
+ );
+ assert_eq!(
+ result.scope,
+ Some(vec!["files:read".to_string(), "files:write".to_string()])
+ );
+ assert!(result.resource_metadata.is_some());
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_invalid_token_error() {
+ let header =
+ r#"Bearer error="invalid_token", error_description="The access token expired""#;
+ let result = parse_www_authenticate(header).unwrap();
+ assert_eq!(result.error, Some(BearerError::InvalidToken));
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_invalid_request_error() {
+ let header = r#"Bearer error="invalid_request""#;
+ let result = parse_www_authenticate(header).unwrap();
+ assert_eq!(result.error, Some(BearerError::InvalidRequest));
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_unknown_error() {
+ let header = r#"Bearer error="some_future_error""#;
+ let result = parse_www_authenticate(header).unwrap();
+ assert_eq!(result.error, Some(BearerError::Other));
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_rejects_non_bearer() {
+ let result = parse_www_authenticate("Basic realm=\"example\"");
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_case_insensitive_scheme() {
+ let header = r#"bearer resource_metadata="https://example.com/.well-known/oauth-protected-resource""#;
+ let result = parse_www_authenticate(header).unwrap();
+ assert!(result.resource_metadata.is_some());
+ }
+
+ #[test]
+ fn test_parse_www_authenticate_multiline_style() {
+ // Some servers emit the header spread across multiple lines joined by
+ // whitespace, as shown in the spec examples.
+ let header = "Bearer resource_metadata=\"https://mcp.example.com/.well-known/oauth-protected-resource\",\n scope=\"files:read\"";
+ let result = parse_www_authenticate(header).unwrap();
+ assert!(result.resource_metadata.is_some());
+ assert_eq!(result.scope, Some(vec!["files:read".to_string()]));
+ }
+
+ #[test]
+ fn test_protected_resource_metadata_urls_with_path() {
+ let server_url = Url::parse("https://api.example.com/v1/mcp").unwrap();
+ let urls = protected_resource_metadata_urls(&server_url);
+
+ assert_eq!(urls.len(), 2);
+ assert_eq!(
+ urls[0].as_str(),
+ "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp"
+ );
+ assert_eq!(
+ urls[1].as_str(),
+ "https://api.example.com/.well-known/oauth-protected-resource"
+ );
+ }
+
+ #[test]
+ fn test_protected_resource_metadata_urls_without_path() {
+ let server_url = Url::parse("https://mcp.example.com").unwrap();
+ let urls = protected_resource_metadata_urls(&server_url);
+
+ assert_eq!(urls.len(), 1);
+ assert_eq!(
+ urls[0].as_str(),
+ "https://mcp.example.com/.well-known/oauth-protected-resource"
+ );
+ }
+
+ #[test]
+ fn test_auth_server_metadata_urls_with_path() {
+ let issuer = Url::parse("https://auth.example.com/tenant1").unwrap();
+ let urls = auth_server_metadata_urls(&issuer);
+
+ assert_eq!(urls.len(), 3);
+ assert_eq!(
+ urls[0].as_str(),
+ "https://auth.example.com/.well-known/oauth-authorization-server/tenant1"
+ );
+ assert_eq!(
+ urls[1].as_str(),
+ "https://auth.example.com/.well-known/openid-configuration/tenant1"
+ );
+ assert_eq!(
+ urls[2].as_str(),
+ "https://auth.example.com/tenant1/.well-known/openid-configuration"
+ );
+ }
+
+ #[test]
+ fn test_auth_server_metadata_urls_without_path() {
+ let issuer = Url::parse("https://auth.example.com").unwrap();
+ let urls = auth_server_metadata_urls(&issuer);
+
+ assert_eq!(urls.len(), 2);
+ assert_eq!(
+ urls[0].as_str(),
+ "https://auth.example.com/.well-known/oauth-authorization-server"
+ );
+ assert_eq!(
+ urls[1].as_str(),
+ "https://auth.example.com/.well-known/openid-configuration"
+ );
+ }
+
+ // -- Canonical server URI tests ------------------------------------------
+
+ #[test]
+ fn test_canonical_server_uri_simple() {
+ let url = Url::parse("https://mcp.example.com").unwrap();
+ assert_eq!(canonical_server_uri(&url), "https://mcp.example.com");
+ }
+
+ #[test]
+ fn test_canonical_server_uri_with_path() {
+ let url = Url::parse("https://mcp.example.com/v1/mcp").unwrap();
+ assert_eq!(canonical_server_uri(&url), "https://mcp.example.com/v1/mcp");
+ }
+
+ #[test]
+ fn test_canonical_server_uri_strips_trailing_slash() {
+ let url = Url::parse("https://mcp.example.com/").unwrap();
+ assert_eq!(canonical_server_uri(&url), "https://mcp.example.com");
+ }
+
+ #[test]
+ fn test_canonical_server_uri_preserves_port() {
+ let url = Url::parse("https://mcp.example.com:8443").unwrap();
+ assert_eq!(canonical_server_uri(&url), "https://mcp.example.com:8443");
+ }
+
+ #[test]
+ fn test_canonical_server_uri_lowercases() {
+ let url = Url::parse("HTTPS://MCP.Example.COM/Server/MCP").unwrap();
+ assert_eq!(
+ canonical_server_uri(&url),
+ "https://mcp.example.com/Server/MCP"
+ );
+ }
+
+ // -- Scope selection tests -----------------------------------------------
+
+ #[test]
+ fn test_select_scopes_prefers_www_authenticate() {
+ let www_auth = WwwAuthenticate {
+ resource_metadata: None,
+ scope: Some(vec!["files:read".into()]),
+ error: None,
+ error_description: None,
+ };
+ let resource_meta = ProtectedResourceMetadata {
+ resource: Url::parse("https://example.com").unwrap(),
+ authorization_servers: vec![],
+ scopes_supported: Some(vec!["files:read".into(), "files:write".into()]),
+ };
+ assert_eq!(select_scopes(&www_auth, &resource_meta), vec!["files:read"]);
+ }
+
+ #[test]
+ fn test_select_scopes_falls_back_to_resource_metadata() {
+ let www_auth = WwwAuthenticate {
+ resource_metadata: None,
+ scope: None,
+ error: None,
+ error_description: None,
+ };
+ let resource_meta = ProtectedResourceMetadata {
+ resource: Url::parse("https://example.com").unwrap(),
+ authorization_servers: vec![],
+ scopes_supported: Some(vec!["admin".into()]),
+ };
+ assert_eq!(select_scopes(&www_auth, &resource_meta), vec!["admin"]);
+ }
+
+ #[test]
+ fn test_select_scopes_empty_when_nothing_available() {
+ let www_auth = WwwAuthenticate {
+ resource_metadata: None,
+ scope: None,
+ error: None,
+ error_description: None,
+ };
+ let resource_meta = ProtectedResourceMetadata {
+ resource: Url::parse("https://example.com").unwrap(),
+ authorization_servers: vec![],
+ scopes_supported: None,
+ };
+ assert!(select_scopes(&www_auth, &resource_meta).is_empty());
+ }
+
+ // -- Client registration strategy tests ----------------------------------
+
+ #[test]
+ fn test_registration_strategy_prefers_cimd() {
+ let metadata = AuthServerMetadata {
+ issuer: Url::parse("https://auth.example.com").unwrap(),
+ authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ registration_endpoint: Some(Url::parse("https://auth.example.com/register").unwrap()),
+ scopes_supported: None,
+ code_challenge_methods_supported: Some(vec!["S256".into()]),
+ client_id_metadata_document_supported: true,
+ };
+ assert_eq!(
+ determine_registration_strategy(&metadata),
+ ClientRegistrationStrategy::Cimd {
+ client_id: CIMD_URL.to_string(),
+ }
+ );
+ }
+
+ #[test]
+ fn test_registration_strategy_falls_back_to_dcr() {
+ let reg_endpoint = Url::parse("https://auth.example.com/register").unwrap();
+ let metadata = AuthServerMetadata {
+ issuer: Url::parse("https://auth.example.com").unwrap(),
+ authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ registration_endpoint: Some(reg_endpoint.clone()),
+ scopes_supported: None,
+ code_challenge_methods_supported: Some(vec!["S256".into()]),
+ client_id_metadata_document_supported: false,
+ };
+ assert_eq!(
+ determine_registration_strategy(&metadata),
+ ClientRegistrationStrategy::Dcr {
+ registration_endpoint: reg_endpoint,
+ }
+ );
+ }
+
+ #[test]
+ fn test_registration_strategy_unavailable() {
+ let metadata = AuthServerMetadata {
+ issuer: Url::parse("https://auth.example.com").unwrap(),
+ authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ registration_endpoint: None,
+ scopes_supported: None,
+ code_challenge_methods_supported: Some(vec!["S256".into()]),
+ client_id_metadata_document_supported: false,
+ };
+ assert_eq!(
+ determine_registration_strategy(&metadata),
+ ClientRegistrationStrategy::Unavailable,
+ );
+ }
+
+ // -- PKCE tests ----------------------------------------------------------
+
+ #[test]
+ fn test_pkce_challenge_verifier_length() {
+ let pkce = generate_pkce_challenge();
+ // 32 random bytes → 43 base64url chars (no padding).
+ assert_eq!(pkce.verifier.len(), 43);
+ }
+
+ #[test]
+ fn test_pkce_challenge_is_valid_base64url() {
+ let pkce = generate_pkce_challenge();
+ for c in pkce.verifier.chars().chain(pkce.challenge.chars()) {
+ assert!(
+ c.is_ascii_alphanumeric() || c == '-' || c == '_',
+ "invalid base64url character: {}",
+ c
+ );
+ }
+ }
+
+ #[test]
+ fn test_pkce_challenge_is_s256_of_verifier() {
+ let pkce = generate_pkce_challenge();
+ let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD;
+ let expected_digest = Sha256::digest(pkce.verifier.as_bytes());
+ let expected_challenge = engine.encode(expected_digest);
+ assert_eq!(pkce.challenge, expected_challenge);
+ }
+
+ #[test]
+ fn test_pkce_challenges_are_unique() {
+ let a = generate_pkce_challenge();
+ let b = generate_pkce_challenge();
+ assert_ne!(a.verifier, b.verifier);
+ }
+
+ // -- Authorization URL tests ---------------------------------------------
+
+ #[test]
+ fn test_build_authorization_url() {
+ let metadata = AuthServerMetadata {
+ issuer: Url::parse("https://auth.example.com").unwrap(),
+ authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ registration_endpoint: None,
+ scopes_supported: None,
+ code_challenge_methods_supported: Some(vec!["S256".into()]),
+ client_id_metadata_document_supported: true,
+ };
+ let pkce = PkceChallenge {
+ verifier: "test_verifier".into(),
+ challenge: "test_challenge".into(),
+ };
+ let url = build_authorization_url(
+ &metadata,
+ "https://zed.dev/oauth/client-metadata.json",
+ "http://127.0.0.1:12345/callback",
+ &["files:read".into(), "files:write".into()],
+ "https://mcp.example.com",
+ &pkce,
+ "random_state_123",
+ );
+
+ let pairs: std::collections::HashMap<_, _> = url.query_pairs().collect();
+ assert_eq!(pairs.get("response_type").unwrap(), "code");
+ assert_eq!(
+ pairs.get("client_id").unwrap(),
+ "https://zed.dev/oauth/client-metadata.json"
+ );
+ assert_eq!(
+ pairs.get("redirect_uri").unwrap(),
+ "http://127.0.0.1:12345/callback"
+ );
+ assert_eq!(pairs.get("scope").unwrap(), "files:read files:write");
+ assert_eq!(pairs.get("resource").unwrap(), "https://mcp.example.com");
+ assert_eq!(pairs.get("code_challenge").unwrap(), "test_challenge");
+ assert_eq!(pairs.get("code_challenge_method").unwrap(), "S256");
+ assert_eq!(pairs.get("state").unwrap(), "random_state_123");
+ }
+
+ #[test]
+ fn test_build_authorization_url_omits_empty_scope() {
+ let metadata = AuthServerMetadata {
+ issuer: Url::parse("https://auth.example.com").unwrap(),
+ authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ registration_endpoint: None,
+ scopes_supported: None,
+ code_challenge_methods_supported: Some(vec!["S256".into()]),
+ client_id_metadata_document_supported: false,
+ };
+ let pkce = PkceChallenge {
+ verifier: "v".into(),
+ challenge: "c".into(),
+ };
+ let url = build_authorization_url(
+ &metadata,
+ "client_123",
+ "http://127.0.0.1:9999/callback",
+ &[],
+ "https://mcp.example.com",
+ &pkce,
+ "state",
+ );
+
+ let pairs: std::collections::HashMap<_, _> = url.query_pairs().collect();
+ assert!(!pairs.contains_key("scope"));
+ }
+
+ // -- Token exchange / refresh param tests --------------------------------
+
+ #[test]
+ fn test_token_exchange_params() {
+ let params = token_exchange_params(
+ "auth_code_abc",
+ "client_xyz",
+ "http://127.0.0.1:5555/callback",
+ "verifier_123",
+ "https://mcp.example.com",
+ );
+ let map: std::collections::HashMap<&str, &str> =
+ params.iter().map(|(k, v)| (*k, v.as_str())).collect();
+
+ assert_eq!(map["grant_type"], "authorization_code");
+ assert_eq!(map["code"], "auth_code_abc");
+ assert_eq!(map["redirect_uri"], "http://127.0.0.1:5555/callback");
+ assert_eq!(map["client_id"], "client_xyz");
+ assert_eq!(map["code_verifier"], "verifier_123");
+ assert_eq!(map["resource"], "https://mcp.example.com");
+ }
+
+ #[test]
+ fn test_token_refresh_params() {
+ let params =
+ token_refresh_params("refresh_token_abc", "client_xyz", "https://mcp.example.com");
+ let map: std::collections::HashMap<&str, &str> =
+ params.iter().map(|(k, v)| (*k, v.as_str())).collect();
+
+ assert_eq!(map["grant_type"], "refresh_token");
+ assert_eq!(map["refresh_token"], "refresh_token_abc");
+ assert_eq!(map["client_id"], "client_xyz");
+ assert_eq!(map["resource"], "https://mcp.example.com");
+ }
+
+ // -- Token response tests ------------------------------------------------
+
+ #[test]
+ fn test_token_response_into_tokens_with_expiry() {
+ let response: TokenResponse = serde_json::from_str(
+ r#"{"access_token": "at_123", "refresh_token": "rt_456", "expires_in": 3600, "token_type": "Bearer"}"#,
+ )
+ .unwrap();
+
+ let tokens = response.into_tokens();
+ assert_eq!(tokens.access_token, "at_123");
+ assert_eq!(tokens.refresh_token.as_deref(), Some("rt_456"));
+ assert!(tokens.expires_at.is_some());
+ }
+
+ #[test]
+ fn test_token_response_into_tokens_minimal() {
+ let response: TokenResponse =
+ serde_json::from_str(r#"{"access_token": "at_789"}"#).unwrap();
+
+ let tokens = response.into_tokens();
+ assert_eq!(tokens.access_token, "at_789");
+ assert_eq!(tokens.refresh_token, None);
+ assert_eq!(tokens.expires_at, None);
+ }
+
+ // -- DCR body test -------------------------------------------------------
+
+ #[test]
+ fn test_dcr_registration_body_shape() {
+ let body = dcr_registration_body("http://127.0.0.1:12345/callback");
+ assert_eq!(body["client_name"], "Zed");
+ assert_eq!(body["redirect_uris"][0], "http://127.0.0.1:12345/callback");
+ assert_eq!(body["grant_types"][0], "authorization_code");
+ assert_eq!(body["response_types"][0], "code");
+ assert_eq!(body["token_endpoint_auth_method"], "none");
+ }
+
+ // -- Test helpers for async/HTTP tests -----------------------------------
+
+ fn make_fake_http_client(
+ handler: impl Fn(
+ http_client::Request,
+ ) -> std::pin::Pin<
+ Box>> + Send>,
+ > + Send
+ + Sync
+ + 'static,
+ ) -> Arc {
+ http_client::FakeHttpClient::create(handler) as Arc
+ }
+
+ fn json_response(status: u16, body: &str) -> anyhow::Result> {
+ Ok(Response::builder()
+ .status(status)
+ .header("Content-Type", "application/json")
+ .body(AsyncBody::from(body.as_bytes().to_vec()))
+ .unwrap())
+ }
+
+ // -- Discovery integration tests -----------------------------------------
+
+ #[test]
+ fn test_fetch_protected_resource_metadata() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains(".well-known/oauth-protected-resource") {
+ json_response(
+ 200,
+ r#"{
+ "resource": "https://mcp.example.com",
+ "authorization_servers": ["https://auth.example.com"],
+ "scopes_supported": ["read", "write"]
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let server_url = Url::parse("https://mcp.example.com").unwrap();
+ let www_auth = WwwAuthenticate {
+ resource_metadata: None,
+ scope: None,
+ error: None,
+ error_description: None,
+ };
+
+ let metadata = fetch_protected_resource_metadata(&client, &server_url, &www_auth)
+ .await
+ .unwrap();
+
+ assert_eq!(metadata.resource.as_str(), "https://mcp.example.com/");
+ assert_eq!(metadata.authorization_servers.len(), 1);
+ assert_eq!(
+ metadata.authorization_servers[0].as_str(),
+ "https://auth.example.com/"
+ );
+ assert_eq!(
+ metadata.scopes_supported,
+ Some(vec!["read".to_string(), "write".to_string()])
+ );
+ });
+ }
+
+ #[test]
+ fn test_fetch_protected_resource_metadata_prefers_www_authenticate_url() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri == "https://mcp.example.com/custom-resource-metadata" {
+ json_response(
+ 200,
+ r#"{
+ "resource": "https://mcp.example.com",
+ "authorization_servers": ["https://auth.example.com"]
+ }"#,
+ )
+ } else {
+ json_response(500, r#"{"error": "should not be called"}"#)
+ }
+ })
+ });
+
+ let server_url = Url::parse("https://mcp.example.com").unwrap();
+ let www_auth = WwwAuthenticate {
+ resource_metadata: Some(
+ Url::parse("https://mcp.example.com/custom-resource-metadata").unwrap(),
+ ),
+ scope: None,
+ error: None,
+ error_description: None,
+ };
+
+ let metadata = fetch_protected_resource_metadata(&client, &server_url, &www_auth)
+ .await
+ .unwrap();
+
+ assert_eq!(metadata.authorization_servers.len(), 1);
+ });
+ }
+
+ #[test]
+ fn test_fetch_protected_resource_metadata_rejects_cross_origin_url() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ // The cross-origin URL should NOT be fetched; only the
+ // well-known fallback at the server's own origin should be.
+ if uri.contains("attacker.example.com") {
+ panic!("should not fetch cross-origin resource_metadata URL");
+ } else if uri.contains(".well-known/oauth-protected-resource") {
+ json_response(
+ 200,
+ r#"{
+ "resource": "https://mcp.example.com",
+ "authorization_servers": ["https://auth.example.com"]
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let server_url = Url::parse("https://mcp.example.com").unwrap();
+ let www_auth = WwwAuthenticate {
+ resource_metadata: Some(
+ Url::parse("https://attacker.example.com/fake-metadata").unwrap(),
+ ),
+ scope: None,
+ error: None,
+ error_description: None,
+ };
+
+ let metadata = fetch_protected_resource_metadata(&client, &server_url, &www_auth)
+ .await
+ .unwrap();
+
+ // Should have used the fallback well-known URL, not the attacker's.
+ assert_eq!(metadata.resource.as_str(), "https://mcp.example.com/");
+ });
+ }
+
+ #[test]
+ fn test_fetch_auth_server_metadata() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains(".well-known/oauth-authorization-server") {
+ json_response(
+ 200,
+ r#"{
+ "issuer": "https://auth.example.com",
+ "authorization_endpoint": "https://auth.example.com/authorize",
+ "token_endpoint": "https://auth.example.com/token",
+ "registration_endpoint": "https://auth.example.com/register",
+ "code_challenge_methods_supported": ["S256"],
+ "client_id_metadata_document_supported": true
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let issuer = Url::parse("https://auth.example.com").unwrap();
+ let metadata = fetch_auth_server_metadata(&client, &issuer).await.unwrap();
+
+ assert_eq!(metadata.issuer.as_str(), "https://auth.example.com/");
+ assert_eq!(
+ metadata.authorization_endpoint.as_str(),
+ "https://auth.example.com/authorize"
+ );
+ assert_eq!(
+ metadata.token_endpoint.as_str(),
+ "https://auth.example.com/token"
+ );
+ assert!(metadata.registration_endpoint.is_some());
+ assert!(metadata.client_id_metadata_document_supported);
+ assert_eq!(
+ metadata.code_challenge_methods_supported,
+ Some(vec!["S256".to_string()])
+ );
+ });
+ }
+
+ #[test]
+ fn test_fetch_auth_server_metadata_falls_back_to_oidc() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains("openid-configuration") {
+ json_response(
+ 200,
+ r#"{
+ "issuer": "https://auth.example.com",
+ "authorization_endpoint": "https://auth.example.com/authorize",
+ "token_endpoint": "https://auth.example.com/token",
+ "code_challenge_methods_supported": ["S256"]
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let issuer = Url::parse("https://auth.example.com").unwrap();
+ let metadata = fetch_auth_server_metadata(&client, &issuer).await.unwrap();
+
+ assert_eq!(
+ metadata.authorization_endpoint.as_str(),
+ "https://auth.example.com/authorize"
+ );
+ assert!(!metadata.client_id_metadata_document_supported);
+ });
+ }
+
+ #[test]
+ fn test_fetch_auth_server_metadata_rejects_issuer_mismatch() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains(".well-known/oauth-authorization-server") {
+ // Response claims to be a different issuer.
+ json_response(
+ 200,
+ r#"{
+ "issuer": "https://evil.example.com",
+ "authorization_endpoint": "https://evil.example.com/authorize",
+ "token_endpoint": "https://evil.example.com/token",
+ "code_challenge_methods_supported": ["S256"]
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let issuer = Url::parse("https://auth.example.com").unwrap();
+ let result = fetch_auth_server_metadata(&client, &issuer).await;
+
+ assert!(result.is_err());
+ let err_msg = result.unwrap_err().to_string();
+ assert!(
+ err_msg.contains("issuer mismatch"),
+ "unexpected error: {}",
+ err_msg
+ );
+ });
+ }
+
+ // -- Full discover integration tests -------------------------------------
+
+ #[test]
+ fn test_full_discover_with_cimd() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains("oauth-protected-resource") {
+ json_response(
+ 200,
+ r#"{
+ "resource": "https://mcp.example.com",
+ "authorization_servers": ["https://auth.example.com"],
+ "scopes_supported": ["mcp:read"]
+ }"#,
+ )
+ } else if uri.contains("oauth-authorization-server") {
+ json_response(
+ 200,
+ r#"{
+ "issuer": "https://auth.example.com",
+ "authorization_endpoint": "https://auth.example.com/authorize",
+ "token_endpoint": "https://auth.example.com/token",
+ "code_challenge_methods_supported": ["S256"],
+ "client_id_metadata_document_supported": true
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let server_url = Url::parse("https://mcp.example.com").unwrap();
+ let www_auth = WwwAuthenticate {
+ resource_metadata: None,
+ scope: None,
+ error: None,
+ error_description: None,
+ };
+
+ let discovery = discover(&client, &server_url, &www_auth).await.unwrap();
+ let registration =
+ resolve_client_registration(&client, &discovery, "http://127.0.0.1:12345/callback")
+ .await
+ .unwrap();
+
+ assert_eq!(registration.client_id, CIMD_URL);
+ assert_eq!(registration.client_secret, None);
+ assert_eq!(discovery.scopes, vec!["mcp:read"]);
+ });
+ }
+
+ #[test]
+ fn test_full_discover_with_dcr_fallback() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains("oauth-protected-resource") {
+ json_response(
+ 200,
+ r#"{
+ "resource": "https://mcp.example.com",
+ "authorization_servers": ["https://auth.example.com"]
+ }"#,
+ )
+ } else if uri.contains("oauth-authorization-server") {
+ json_response(
+ 200,
+ r#"{
+ "issuer": "https://auth.example.com",
+ "authorization_endpoint": "https://auth.example.com/authorize",
+ "token_endpoint": "https://auth.example.com/token",
+ "registration_endpoint": "https://auth.example.com/register",
+ "code_challenge_methods_supported": ["S256"],
+ "client_id_metadata_document_supported": false
+ }"#,
+ )
+ } else if uri.contains("/register") {
+ json_response(
+ 201,
+ r#"{
+ "client_id": "dcr-minted-id-123",
+ "client_secret": "dcr-secret-456"
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let server_url = Url::parse("https://mcp.example.com").unwrap();
+ let www_auth = WwwAuthenticate {
+ resource_metadata: None,
+ scope: Some(vec!["files:read".into()]),
+ error: None,
+ error_description: None,
+ };
+
+ let discovery = discover(&client, &server_url, &www_auth).await.unwrap();
+ let registration =
+ resolve_client_registration(&client, &discovery, "http://127.0.0.1:9999/callback")
+ .await
+ .unwrap();
+
+ assert_eq!(registration.client_id, "dcr-minted-id-123");
+ assert_eq!(
+ registration.client_secret.as_deref(),
+ Some("dcr-secret-456")
+ );
+ assert_eq!(discovery.scopes, vec!["files:read"]);
+ });
+ }
+
+ #[test]
+ fn test_discover_fails_without_pkce_support() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains("oauth-protected-resource") {
+ json_response(
+ 200,
+ r#"{
+ "resource": "https://mcp.example.com",
+ "authorization_servers": ["https://auth.example.com"]
+ }"#,
+ )
+ } else if uri.contains("oauth-authorization-server") {
+ json_response(
+ 200,
+ r#"{
+ "issuer": "https://auth.example.com",
+ "authorization_endpoint": "https://auth.example.com/authorize",
+ "token_endpoint": "https://auth.example.com/token"
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let server_url = Url::parse("https://mcp.example.com").unwrap();
+ let www_auth = WwwAuthenticate {
+ resource_metadata: None,
+ scope: None,
+ error: None,
+ error_description: None,
+ };
+
+ let result = discover(&client, &server_url, &www_auth).await;
+ assert!(result.is_err());
+ let err_msg = result.unwrap_err().to_string();
+ assert!(
+ err_msg.contains("code_challenge_methods_supported"),
+ "unexpected error: {}",
+ err_msg
+ );
+ });
+ }
+
+ // -- Token exchange integration tests ------------------------------------
+
+ #[test]
+ fn test_exchange_code_success() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains("/token") {
+ json_response(
+ 200,
+ r#"{
+ "access_token": "new_access_token",
+ "refresh_token": "new_refresh_token",
+ "expires_in": 3600,
+ "token_type": "Bearer"
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let metadata = AuthServerMetadata {
+ issuer: Url::parse("https://auth.example.com").unwrap(),
+ authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ registration_endpoint: None,
+ scopes_supported: None,
+ code_challenge_methods_supported: Some(vec!["S256".into()]),
+ client_id_metadata_document_supported: true,
+ };
+
+ let tokens = exchange_code(
+ &client,
+ &metadata,
+ "auth_code_123",
+ CIMD_URL,
+ "http://127.0.0.1:9999/callback",
+ "verifier_abc",
+ "https://mcp.example.com",
+ )
+ .await
+ .unwrap();
+
+ assert_eq!(tokens.access_token, "new_access_token");
+ assert_eq!(tokens.refresh_token.as_deref(), Some("new_refresh_token"));
+ assert!(tokens.expires_at.is_some());
+ });
+ }
+
+ #[test]
+ fn test_refresh_tokens_success() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|req| {
+ Box::pin(async move {
+ let uri = req.uri().to_string();
+ if uri.contains("/token") {
+ json_response(
+ 200,
+ r#"{
+ "access_token": "refreshed_token",
+ "expires_in": 1800,
+ "token_type": "Bearer"
+ }"#,
+ )
+ } else {
+ json_response(404, "{}")
+ }
+ })
+ });
+
+ let token_endpoint = Url::parse("https://auth.example.com/token").unwrap();
+
+ let tokens = refresh_tokens(
+ &client,
+ &token_endpoint,
+ "old_refresh_token",
+ CIMD_URL,
+ "https://mcp.example.com",
+ )
+ .await
+ .unwrap();
+
+ assert_eq!(tokens.access_token, "refreshed_token");
+ assert_eq!(tokens.refresh_token, None);
+ assert!(tokens.expires_at.is_some());
+ });
+ }
+
+ #[test]
+ fn test_exchange_code_failure() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|_req| {
+ Box::pin(async move { json_response(400, r#"{"error": "invalid_grant"}"#) })
+ });
+
+ let metadata = AuthServerMetadata {
+ issuer: Url::parse("https://auth.example.com").unwrap(),
+ authorization_endpoint: Url::parse("https://auth.example.com/authorize").unwrap(),
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ registration_endpoint: None,
+ scopes_supported: None,
+ code_challenge_methods_supported: Some(vec!["S256".into()]),
+ client_id_metadata_document_supported: true,
+ };
+
+ let result = exchange_code(
+ &client,
+ &metadata,
+ "bad_code",
+ "client",
+ "http://127.0.0.1:1/callback",
+ "verifier",
+ "https://mcp.example.com",
+ )
+ .await;
+
+ assert!(result.is_err());
+ assert!(result.unwrap_err().to_string().contains("400"));
+ });
+ }
+
+ // -- DCR integration tests -----------------------------------------------
+
+ #[test]
+ fn test_perform_dcr() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|_req| {
+ Box::pin(async move {
+ json_response(
+ 201,
+ r#"{
+ "client_id": "dynamic-client-001",
+ "client_secret": "dynamic-secret-001"
+ }"#,
+ )
+ })
+ });
+
+ let endpoint = Url::parse("https://auth.example.com/register").unwrap();
+ let registration = perform_dcr(&client, &endpoint, "http://127.0.0.1:9999/callback")
+ .await
+ .unwrap();
+
+ assert_eq!(registration.client_id, "dynamic-client-001");
+ assert_eq!(
+ registration.client_secret.as_deref(),
+ Some("dynamic-secret-001")
+ );
+ });
+ }
+
+ #[test]
+ fn test_perform_dcr_failure() {
+ smol::block_on(async {
+ let client = make_fake_http_client(|_req| {
+ Box::pin(
+ async move { json_response(403, r#"{"error": "registration_not_allowed"}"#) },
+ )
+ });
+
+ let endpoint = Url::parse("https://auth.example.com/register").unwrap();
+ let result = perform_dcr(&client, &endpoint, "http://127.0.0.1:9999/callback").await;
+
+ assert!(result.is_err());
+ assert!(result.unwrap_err().to_string().contains("403"));
+ });
+ }
+
+ // -- OAuthCallback parse tests -------------------------------------------
+
+ #[test]
+ fn test_oauth_callback_parse_query() {
+ let callback = OAuthCallback::parse_query("code=test_auth_code&state=test_state").unwrap();
+ assert_eq!(callback.code, "test_auth_code");
+ assert_eq!(callback.state, "test_state");
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_reversed_order() {
+ let callback = OAuthCallback::parse_query("state=test_state&code=test_auth_code").unwrap();
+ assert_eq!(callback.code, "test_auth_code");
+ assert_eq!(callback.state, "test_state");
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_with_extra_params() {
+ let callback =
+ OAuthCallback::parse_query("code=test_auth_code&state=test_state&extra=ignored")
+ .unwrap();
+ assert_eq!(callback.code, "test_auth_code");
+ assert_eq!(callback.state, "test_state");
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_missing_code() {
+ let result = OAuthCallback::parse_query("state=test_state");
+ assert!(result.is_err());
+ assert!(result.unwrap_err().to_string().contains("code"));
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_missing_state() {
+ let result = OAuthCallback::parse_query("code=test_auth_code");
+ assert!(result.is_err());
+ assert!(result.unwrap_err().to_string().contains("state"));
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_empty_code() {
+ let result = OAuthCallback::parse_query("code=&state=test_state");
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_empty_state() {
+ let result = OAuthCallback::parse_query("code=test_auth_code&state=");
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_url_encoded_values() {
+ let callback = OAuthCallback::parse_query("code=abc%20def&state=test%3Dstate").unwrap();
+ assert_eq!(callback.code, "abc def");
+ assert_eq!(callback.state, "test=state");
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_error_response() {
+ let result = OAuthCallback::parse_query(
+ "error=access_denied&error_description=User%20denied%20access&state=abc",
+ );
+ assert!(result.is_err());
+ let err_msg = result.unwrap_err().to_string();
+ assert!(
+ err_msg.contains("access_denied"),
+ "unexpected error: {}",
+ err_msg
+ );
+ assert!(
+ err_msg.contains("User denied access"),
+ "unexpected error: {}",
+ err_msg
+ );
+ }
+
+ #[test]
+ fn test_oauth_callback_parse_query_error_without_description() {
+ let result = OAuthCallback::parse_query("error=server_error&state=abc");
+ assert!(result.is_err());
+ let err_msg = result.unwrap_err().to_string();
+ assert!(
+ err_msg.contains("server_error"),
+ "unexpected error: {}",
+ err_msg
+ );
+ assert!(
+ err_msg.contains("no description"),
+ "unexpected error: {}",
+ err_msg
+ );
+ }
+
+ // -- McpOAuthTokenProvider tests -----------------------------------------
+
+ fn make_test_session(
+ access_token: &str,
+ refresh_token: Option<&str>,
+ expires_at: Option,
+ ) -> OAuthSession {
+ OAuthSession {
+ token_endpoint: Url::parse("https://auth.example.com/token").unwrap(),
+ resource: Url::parse("https://mcp.example.com").unwrap(),
+ client_registration: OAuthClientRegistration {
+ client_id: "test-client".into(),
+ client_secret: None,
+ },
+ tokens: OAuthTokens {
+ access_token: access_token.into(),
+ refresh_token: refresh_token.map(String::from),
+ expires_at,
+ },
+ }
+ }
+
+ #[test]
+ fn test_mcp_oauth_provider_returns_none_when_token_expired() {
+ let expired = SystemTime::now() - Duration::from_secs(60);
+ let session = make_test_session("stale-token", Some("rt"), Some(expired));
+ let provider = McpOAuthTokenProvider::new(
+ session,
+ make_fake_http_client(|_| Box::pin(async { unreachable!() })),
+ None,
+ );
+
+ assert_eq!(provider.access_token(), None);
+ }
+
+ #[test]
+ fn test_mcp_oauth_provider_returns_token_when_not_expired() {
+ let far_future = SystemTime::now() + Duration::from_secs(3600);
+ let session = make_test_session("valid-token", Some("rt"), Some(far_future));
+ let provider = McpOAuthTokenProvider::new(
+ session,
+ make_fake_http_client(|_| Box::pin(async { unreachable!() })),
+ None,
+ );
+
+ assert_eq!(provider.access_token().as_deref(), Some("valid-token"));
+ }
+
+ #[test]
+ fn test_mcp_oauth_provider_returns_token_when_no_expiry() {
+ let session = make_test_session("no-expiry-token", Some("rt"), None);
+ let provider = McpOAuthTokenProvider::new(
+ session,
+ make_fake_http_client(|_| Box::pin(async { unreachable!() })),
+ None,
+ );
+
+ assert_eq!(provider.access_token().as_deref(), Some("no-expiry-token"));
+ }
+
+ #[test]
+ fn test_mcp_oauth_provider_refresh_without_refresh_token_returns_false() {
+ smol::block_on(async {
+ let session = make_test_session("token", None, None);
+ let provider = McpOAuthTokenProvider::new(
+ session,
+ make_fake_http_client(|_| {
+ Box::pin(async { unreachable!("no HTTP call expected") })
+ }),
+ None,
+ );
+
+ let refreshed = provider.try_refresh().await.unwrap();
+ assert!(!refreshed);
+ });
+ }
+
+ #[test]
+ fn test_mcp_oauth_provider_refresh_updates_session_and_notifies_channel() {
+ smol::block_on(async {
+ let session = make_test_session("old-access", Some("my-refresh-token"), None);
+ let (tx, mut rx) = futures::channel::mpsc::unbounded();
+
+ let http_client = make_fake_http_client(|_req| {
+ Box::pin(async {
+ json_response(
+ 200,
+ r#"{
+ "access_token": "new-access",
+ "refresh_token": "new-refresh",
+ "expires_in": 1800
+ }"#,
+ )
+ })
+ });
+
+ let provider = McpOAuthTokenProvider::new(session, http_client, Some(tx));
+
+ let refreshed = provider.try_refresh().await.unwrap();
+ assert!(refreshed);
+ assert_eq!(provider.access_token().as_deref(), Some("new-access"));
+
+ let notified_session = rx
+ .try_next()
+ .unwrap()
+ .expect("channel should have a session");
+ assert_eq!(notified_session.tokens.access_token, "new-access");
+ assert_eq!(
+ notified_session.tokens.refresh_token.as_deref(),
+ Some("new-refresh")
+ );
+ });
+ }
+
+ #[test]
+ fn test_mcp_oauth_provider_refresh_preserves_old_refresh_token_when_server_omits_it() {
+ smol::block_on(async {
+ let session = make_test_session("old-access", Some("original-refresh"), None);
+ let (tx, mut rx) = futures::channel::mpsc::unbounded();
+
+ let http_client = make_fake_http_client(|_req| {
+ Box::pin(async {
+ json_response(
+ 200,
+ r#"{
+ "access_token": "new-access",
+ "expires_in": 900
+ }"#,
+ )
+ })
+ });
+
+ let provider = McpOAuthTokenProvider::new(session, http_client, Some(tx));
+
+ let refreshed = provider.try_refresh().await.unwrap();
+ assert!(refreshed);
+
+ let notified_session = rx
+ .try_next()
+ .unwrap()
+ .expect("channel should have a session");
+ assert_eq!(notified_session.tokens.access_token, "new-access");
+ assert_eq!(
+ notified_session.tokens.refresh_token.as_deref(),
+ Some("original-refresh"),
+ );
+ });
+ }
+
+ #[test]
+ fn test_mcp_oauth_provider_refresh_returns_false_on_http_error() {
+ smol::block_on(async {
+ let session = make_test_session("old-access", Some("my-refresh"), None);
+
+ let http_client = make_fake_http_client(|_req| {
+ Box::pin(async { json_response(401, r#"{"error": "invalid_grant"}"#) })
+ });
+
+ let provider = McpOAuthTokenProvider::new(session, http_client, None);
+
+ let refreshed = provider.try_refresh().await.unwrap();
+ assert!(!refreshed);
+ // The old token should still be in place.
+ assert_eq!(provider.access_token().as_deref(), Some("old-access"));
+ });
+ }
+}
diff --git a/crates/context_server/src/transport/http.rs b/crates/context_server/src/transport/http.rs
index 70248f0278fcf80024d75d7f78cae5c29f26cc43..3e002983b5e49026d668c8baabfe8f856e4c5fe7 100644
--- a/crates/context_server/src/transport/http.rs
+++ b/crates/context_server/src/transport/http.rs
@@ -8,8 +8,30 @@ use parking_lot::Mutex as SyncMutex;
use smol::channel;
use std::{pin::Pin, sync::Arc};
+use crate::oauth::{self, OAuthTokenProvider, WwwAuthenticate};
use crate::transport::Transport;
+/// Typed errors returned by the HTTP transport that callers can downcast from
+/// `anyhow::Error` to handle specific failure modes.
+#[derive(Debug)]
+pub enum TransportError {
+ /// The server returned 401 and token refresh either wasn't possible or
+ /// failed. The caller should initiate the OAuth authorization flow.
+ AuthRequired { www_authenticate: WwwAuthenticate },
+}
+
+impl std::fmt::Display for TransportError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ TransportError::AuthRequired { .. } => {
+ write!(f, "OAuth authorization required")
+ }
+ }
+ }
+}
+
+impl std::error::Error for TransportError {}
+
// Constants from MCP spec
const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
@@ -25,8 +47,11 @@ pub struct HttpTransport {
response_rx: channel::Receiver,
error_tx: channel::Sender,
error_rx: channel::Receiver,
- // Authentication headers to include in requests
+ /// Static headers to include in every request (e.g. from server config).
headers: HashMap,
+ /// When set, the transport attaches `Authorization: Bearer` headers and
+ /// handles 401 responses with token refresh + retry.
+ token_provider: Option>,
}
impl HttpTransport {
@@ -35,6 +60,16 @@ impl HttpTransport {
endpoint: String,
headers: HashMap,
executor: BackgroundExecutor,
+ ) -> Self {
+ Self::new_with_token_provider(http_client, endpoint, headers, executor, None)
+ }
+
+ pub fn new_with_token_provider(
+ http_client: Arc,
+ endpoint: String,
+ headers: HashMap,
+ executor: BackgroundExecutor,
+ token_provider: Option>,
) -> Self {
let (response_tx, response_rx) = channel::unbounded();
let (error_tx, error_rx) = channel::unbounded();
@@ -49,14 +84,14 @@ impl HttpTransport {
error_tx,
error_rx,
headers,
+ token_provider,
}
}
- /// Send a message and handle the response based on content type
- async fn send_message(&self, message: String) -> Result<()> {
- let is_notification =
- !message.contains("\"id\":") || message.contains("notifications/initialized");
-
+ /// Build a POST request for the given message body, attaching all standard
+ /// headers (content-type, accept, session ID, static headers, and bearer
+ /// token if available).
+ fn build_request(&self, message: &[u8]) -> Result> {
let mut request_builder = Request::builder()
.method(Method::POST)
.uri(&self.endpoint)
@@ -70,15 +105,71 @@ impl HttpTransport {
request_builder = request_builder.header(key.as_str(), value.as_str());
}
- // Add session ID if we have one (except for initialize)
+ // Attach bearer token when a token provider is present.
+ if let Some(token) = self.token_provider.as_ref().and_then(|p| p.access_token()) {
+ request_builder = request_builder.header("Authorization", format!("Bearer {}", token));
+ }
+
+ // Add session ID if we have one (except for initialize).
if let Some(ref session_id) = *self.session_id.lock() {
request_builder = request_builder.header(HEADER_SESSION_ID, session_id.as_str());
}
- let request = request_builder.body(AsyncBody::from(message.into_bytes()))?;
+ Ok(request_builder.body(AsyncBody::from(message.to_vec()))?)
+ }
+
+ /// Send a message and handle the response based on content type.
+ async fn send_message(&self, message: String) -> Result<()> {
+ let is_notification =
+ !message.contains("\"id\":") || message.contains("notifications/initialized");
+
+ // If we currently have no access token, try refreshing before sending
+ // the request so restored but expired sessions do not need an initial
+ // 401 round-trip before they can recover.
+ if let Some(ref provider) = self.token_provider {
+ if provider.access_token().is_none() {
+ provider.try_refresh().await.unwrap_or(false);
+ }
+ }
+
+ let request = self.build_request(message.as_bytes())?;
let mut response = self.http_client.send(request).await?;
- // Handle different response types based on status and content-type
+ // On 401, try refreshing the token and retry once.
+ if response.status().as_u16() == 401 {
+ let www_auth_header = response
+ .headers()
+ .get("www-authenticate")
+ .and_then(|v| v.to_str().ok())
+ .unwrap_or("Bearer");
+
+ let www_authenticate =
+ oauth::parse_www_authenticate(www_auth_header).unwrap_or(WwwAuthenticate {
+ resource_metadata: None,
+ scope: None,
+ error: None,
+ error_description: None,
+ });
+
+ if let Some(ref provider) = self.token_provider {
+ if provider.try_refresh().await.unwrap_or(false) {
+ // Retry with the refreshed token.
+ let retry_request = self.build_request(message.as_bytes())?;
+ response = self.http_client.send(retry_request).await?;
+
+ // If still 401 after refresh, give up.
+ if response.status().as_u16() == 401 {
+ return Err(TransportError::AuthRequired { www_authenticate }.into());
+ }
+ } else {
+ return Err(TransportError::AuthRequired { www_authenticate }.into());
+ }
+ } else {
+ return Err(TransportError::AuthRequired { www_authenticate }.into());
+ }
+ }
+
+ // Handle different response types based on status and content-type.
match response.status() {
status if status.is_success() => {
// Check content type
@@ -233,6 +324,7 @@ impl Drop for HttpTransport {
let endpoint = self.endpoint.clone();
let session_id = self.session_id.lock().clone();
let headers = self.headers.clone();
+ let access_token = self.token_provider.as_ref().and_then(|p| p.access_token());
if let Some(session_id) = session_id {
self.executor
@@ -242,11 +334,17 @@ impl Drop for HttpTransport {
.uri(&endpoint)
.header(HEADER_SESSION_ID, &session_id);
- // Add authentication headers if present
+ // Add static authentication headers.
for (key, value) in headers {
request_builder = request_builder.header(key.as_str(), value.as_str());
}
+ // Attach bearer token if available.
+ if let Some(token) = access_token {
+ request_builder =
+ request_builder.header("Authorization", format!("Bearer {}", token));
+ }
+
let request = request_builder.body(AsyncBody::empty());
if let Ok(request) = request {
@@ -257,3 +355,402 @@ impl Drop for HttpTransport {
}
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use async_trait::async_trait;
+ use gpui::TestAppContext;
+ use parking_lot::Mutex as SyncMutex;
+ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+ /// A mock token provider that returns a configurable token and tracks
+ /// refresh attempts.
+ struct FakeTokenProvider {
+ token: SyncMutex>,
+ refreshed_token: SyncMutex>,
+ refresh_succeeds: AtomicBool,
+ refresh_count: AtomicUsize,
+ }
+
+ impl FakeTokenProvider {
+ fn new(token: Option<&str>, refresh_succeeds: bool) -> Arc {
+ Self::with_refreshed_token(token, None, refresh_succeeds)
+ }
+
+ fn with_refreshed_token(
+ token: Option<&str>,
+ refreshed_token: Option<&str>,
+ refresh_succeeds: bool,
+ ) -> Arc {
+ Arc::new(Self {
+ token: SyncMutex::new(token.map(String::from)),
+ refreshed_token: SyncMutex::new(refreshed_token.map(String::from)),
+ refresh_succeeds: AtomicBool::new(refresh_succeeds),
+ refresh_count: AtomicUsize::new(0),
+ })
+ }
+
+ fn set_token(&self, token: &str) {
+ *self.token.lock() = Some(token.to_string());
+ }
+
+ fn refresh_count(&self) -> usize {
+ self.refresh_count.load(Ordering::SeqCst)
+ }
+ }
+
+ #[async_trait]
+ impl OAuthTokenProvider for FakeTokenProvider {
+ fn access_token(&self) -> Option {
+ self.token.lock().clone()
+ }
+
+ async fn try_refresh(&self) -> Result {
+ self.refresh_count.fetch_add(1, Ordering::SeqCst);
+
+ let refresh_succeeds = self.refresh_succeeds.load(Ordering::SeqCst);
+ if refresh_succeeds {
+ if let Some(token) = self.refreshed_token.lock().clone() {
+ *self.token.lock() = Some(token);
+ }
+ }
+
+ Ok(refresh_succeeds)
+ }
+ }
+
+ fn make_fake_http_client(
+ handler: impl Fn(
+ http_client::Request,
+ ) -> std::pin::Pin<
+ Box>> + Send>,
+ > + Send
+ + Sync
+ + 'static,
+ ) -> Arc {
+ http_client::FakeHttpClient::create(handler) as Arc
+ }
+
+ fn json_response(status: u16, body: &str) -> anyhow::Result> {
+ Ok(Response::builder()
+ .status(status)
+ .header("Content-Type", "application/json")
+ .body(AsyncBody::from(body.as_bytes().to_vec()))
+ .unwrap())
+ }
+
+ #[gpui::test]
+ async fn test_bearer_token_attached_to_requests(cx: &mut TestAppContext) {
+ // Capture the Authorization header from the request.
+ let captured_auth = Arc::new(SyncMutex::new(None::));
+ let captured_auth_clone = captured_auth.clone();
+
+ let client = make_fake_http_client(move |req| {
+ let auth = req
+ .headers()
+ .get("Authorization")
+ .map(|v| v.to_str().unwrap().to_string());
+ *captured_auth_clone.lock() = auth;
+ Box::pin(async { json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#) })
+ });
+
+ let provider = FakeTokenProvider::new(Some("test-access-token"), false);
+ let transport = HttpTransport::new_with_token_provider(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ Some(provider),
+ );
+
+ transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .expect("send should succeed");
+
+ assert_eq!(
+ captured_auth.lock().as_deref(),
+ Some("Bearer test-access-token"),
+ );
+ }
+
+ #[gpui::test]
+ async fn test_no_bearer_token_without_provider(cx: &mut TestAppContext) {
+ let captured_auth = Arc::new(SyncMutex::new(None::));
+ let captured_auth_clone = captured_auth.clone();
+
+ let client = make_fake_http_client(move |req| {
+ let auth = req
+ .headers()
+ .get("Authorization")
+ .map(|v| v.to_str().unwrap().to_string());
+ *captured_auth_clone.lock() = auth;
+ Box::pin(async { json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#) })
+ });
+
+ let transport = HttpTransport::new(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ );
+
+ transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .expect("send should succeed");
+
+ assert!(captured_auth.lock().is_none());
+ }
+
+ #[gpui::test]
+ async fn test_missing_token_triggers_refresh_before_first_request(cx: &mut TestAppContext) {
+ let captured_auth = Arc::new(SyncMutex::new(None::));
+ let captured_auth_clone = captured_auth.clone();
+
+ let client = make_fake_http_client(move |req| {
+ let auth = req
+ .headers()
+ .get("Authorization")
+ .map(|v| v.to_str().unwrap().to_string());
+ *captured_auth_clone.lock() = auth;
+ Box::pin(async { json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#) })
+ });
+
+ let provider = FakeTokenProvider::with_refreshed_token(None, Some("refreshed-token"), true);
+ let transport = HttpTransport::new_with_token_provider(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ Some(provider.clone()),
+ );
+
+ transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .expect("send should succeed after proactive refresh");
+
+ assert_eq!(provider.refresh_count(), 1);
+ assert_eq!(
+ captured_auth.lock().as_deref(),
+ Some("Bearer refreshed-token"),
+ );
+ }
+
+ #[gpui::test]
+ async fn test_invalid_token_still_triggers_refresh_and_retry(cx: &mut TestAppContext) {
+ let request_count = Arc::new(AtomicUsize::new(0));
+ let request_count_clone = request_count.clone();
+
+ let client = make_fake_http_client(move |_req| {
+ let count = request_count_clone.fetch_add(1, Ordering::SeqCst);
+ Box::pin(async move {
+ if count == 0 {
+ Ok(Response::builder()
+ .status(401)
+ .header(
+ "WWW-Authenticate",
+ r#"Bearer error="invalid_token", resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#,
+ )
+ .body(AsyncBody::from(b"Unauthorized".to_vec()))
+ .unwrap())
+ } else {
+ json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#)
+ }
+ })
+ });
+
+ let provider = FakeTokenProvider::with_refreshed_token(
+ Some("old-token"),
+ Some("refreshed-token"),
+ true,
+ );
+ let transport = HttpTransport::new_with_token_provider(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ Some(provider.clone()),
+ );
+
+ transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .expect("send should succeed after refresh");
+
+ assert_eq!(provider.refresh_count(), 1);
+ assert_eq!(request_count.load(Ordering::SeqCst), 2);
+ }
+
+ #[gpui::test]
+ async fn test_401_triggers_refresh_and_retry(cx: &mut TestAppContext) {
+ let request_count = Arc::new(AtomicUsize::new(0));
+ let request_count_clone = request_count.clone();
+
+ let client = make_fake_http_client(move |_req| {
+ let count = request_count_clone.fetch_add(1, Ordering::SeqCst);
+ Box::pin(async move {
+ if count == 0 {
+ // First request: 401.
+ Ok(Response::builder()
+ .status(401)
+ .header(
+ "WWW-Authenticate",
+ r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource""#,
+ )
+ .body(AsyncBody::from(b"Unauthorized".to_vec()))
+ .unwrap())
+ } else {
+ // Retry after refresh: 200.
+ json_response(200, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#)
+ }
+ })
+ });
+
+ let provider = FakeTokenProvider::new(Some("old-token"), true);
+ // Simulate the refresh updating the token.
+ let provider_ref = provider.clone();
+ let transport = HttpTransport::new_with_token_provider(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ Some(provider.clone()),
+ );
+
+ // Set the new token that will be used on retry.
+ provider_ref.set_token("refreshed-token");
+
+ transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .expect("send should succeed after refresh");
+
+ assert_eq!(provider_ref.refresh_count(), 1);
+ assert_eq!(request_count.load(Ordering::SeqCst), 2);
+ }
+
+ #[gpui::test]
+ async fn test_401_returns_auth_required_when_refresh_fails(cx: &mut TestAppContext) {
+ let client = make_fake_http_client(|_req| {
+ Box::pin(async {
+ Ok(Response::builder()
+ .status(401)
+ .header(
+ "WWW-Authenticate",
+ r#"Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource", scope="read write""#,
+ )
+ .body(AsyncBody::from(b"Unauthorized".to_vec()))
+ .unwrap())
+ })
+ });
+
+ // Refresh returns false — no new token available.
+ let provider = FakeTokenProvider::new(Some("stale-token"), false);
+ let transport = HttpTransport::new_with_token_provider(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ Some(provider.clone()),
+ );
+
+ let err = transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .unwrap_err();
+
+ let transport_err = err
+ .downcast_ref::()
+ .expect("error should be TransportError");
+ match transport_err {
+ TransportError::AuthRequired { www_authenticate } => {
+ assert_eq!(
+ www_authenticate
+ .resource_metadata
+ .as_ref()
+ .map(|u| u.as_str()),
+ Some("https://mcp.example.com/.well-known/oauth-protected-resource"),
+ );
+ assert_eq!(
+ www_authenticate.scope,
+ Some(vec!["read".to_string(), "write".to_string()]),
+ );
+ }
+ }
+ assert_eq!(provider.refresh_count(), 1);
+ }
+
+ #[gpui::test]
+ async fn test_401_returns_auth_required_without_provider(cx: &mut TestAppContext) {
+ let client = make_fake_http_client(|_req| {
+ Box::pin(async {
+ Ok(Response::builder()
+ .status(401)
+ .header("WWW-Authenticate", "Bearer")
+ .body(AsyncBody::from(b"Unauthorized".to_vec()))
+ .unwrap())
+ })
+ });
+
+ // No token provider at all.
+ let transport = HttpTransport::new(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ );
+
+ let err = transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .unwrap_err();
+
+ let transport_err = err
+ .downcast_ref::()
+ .expect("error should be TransportError");
+ match transport_err {
+ TransportError::AuthRequired { www_authenticate } => {
+ assert!(www_authenticate.resource_metadata.is_none());
+ assert!(www_authenticate.scope.is_none());
+ }
+ }
+ }
+
+ #[gpui::test]
+ async fn test_401_after_successful_refresh_still_returns_auth_required(
+ cx: &mut TestAppContext,
+ ) {
+ // Both requests return 401 — the server rejects the refreshed token too.
+ let client = make_fake_http_client(|_req| {
+ Box::pin(async {
+ Ok(Response::builder()
+ .status(401)
+ .header("WWW-Authenticate", "Bearer")
+ .body(AsyncBody::from(b"Unauthorized".to_vec()))
+ .unwrap())
+ })
+ });
+
+ let provider = FakeTokenProvider::new(Some("token"), true);
+ let transport = HttpTransport::new_with_token_provider(
+ client,
+ "http://mcp.example.com/mcp".to_string(),
+ HashMap::default(),
+ cx.background_executor.clone(),
+ Some(provider.clone()),
+ );
+
+ let err = transport
+ .send(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#.to_string())
+ .await
+ .unwrap_err();
+
+ err.downcast_ref::()
+ .expect("error should be TransportError");
+ // Refresh was attempted exactly once.
+ assert_eq!(provider.refresh_count(), 1);
+ }
+}
diff --git a/crates/edit_prediction/Cargo.toml b/crates/edit_prediction/Cargo.toml
index d2a23b8b4ec3425072ffbe9d042ff89d26a56778..a6a7d8777cbf0d52575489e91a5ae03be2d031ea 100644
--- a/crates/edit_prediction/Cargo.toml
+++ b/crates/edit_prediction/Cargo.toml
@@ -17,7 +17,7 @@ cli-support = []
[dependencies]
ai_onboarding.workspace = true
anyhow.workspace = true
-arrayvec.workspace = true
+heapless.workspace = true
brotli.workspace = true
buffer_diff.workspace = true
client.workspace = true
diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs
index bd1bc7a3303a6f80094fd8261e90a2c5e113803d..421a51b055693617a915e622b617298f5f8a01c5 100644
--- a/crates/edit_prediction/src/edit_prediction.rs
+++ b/crates/edit_prediction/src/edit_prediction.rs
@@ -1,5 +1,4 @@
use anyhow::Result;
-use arrayvec::ArrayVec;
use client::{Client, EditPredictionUsage, UserStore};
use cloud_api_types::{OrganizationId, SubmitEditPredictionFeedbackBody};
use cloud_llm_client::predict_edits_v3::{
@@ -27,6 +26,7 @@ use gpui::{
http_client::{self, AsyncBody, Method},
prelude::*,
};
+use heapless::Vec as ArrayVec;
use language::language_settings::all_language_settings;
use language::{Anchor, Buffer, File, Point, TextBufferSnapshot, ToOffset, ToPoint};
use language::{BufferSnapshot, OffsetRangeExt};
@@ -332,7 +332,7 @@ struct ProjectState {
registered_buffers: HashMap,
current_prediction: Option,
next_pending_prediction_id: usize,
- pending_predictions: ArrayVec,
+ pending_predictions: ArrayVec,
debug_tx: Option>,
last_edit_prediction_refresh: Option<(EntityId, Instant)>,
last_jump_prediction_refresh: Option<(EntityId, Instant)>,
@@ -2311,18 +2311,24 @@ impl EditPredictionStore {
});
if project_state.pending_predictions.len() < max_pending_predictions {
- project_state.pending_predictions.push(PendingPrediction {
- id: pending_prediction_id,
- task,
- drop_on_cancel,
- });
+ project_state
+ .pending_predictions
+ .push(PendingPrediction {
+ id: pending_prediction_id,
+ task,
+ drop_on_cancel,
+ })
+ .unwrap();
} else {
let pending_prediction = project_state.pending_predictions.pop().unwrap();
- project_state.pending_predictions.push(PendingPrediction {
- id: pending_prediction_id,
- task,
- drop_on_cancel,
- });
+ project_state
+ .pending_predictions
+ .push(PendingPrediction {
+ id: pending_prediction_id,
+ task,
+ drop_on_cancel,
+ })
+ .unwrap();
project_state.cancel_pending_prediction(pending_prediction, cx);
}
}
diff --git a/crates/edit_prediction_cli/src/example.rs b/crates/edit_prediction_cli/src/example.rs
index 495ca26f97af5f2c2c1dc50ea339881853d9ebbc..196f4f96d99b64aed2ff3ae2d7a9897295a60b29 100644
--- a/crates/edit_prediction_cli/src/example.rs
+++ b/crates/edit_prediction_cli/src/example.rs
@@ -82,6 +82,10 @@ pub struct ExamplePrediction {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option,
pub provider: PredictionProvider,
+ #[serde(default, skip_serializing_if = "Option::is_none")]
+ pub cumulative_logprob: Option,
+ #[serde(default, skip_serializing_if = "Option::is_none")]
+ pub avg_logprob: Option,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
@@ -166,6 +170,10 @@ pub struct ExampleScore {
pub inserted_tokens: usize,
#[serde(default)]
pub deleted_tokens: usize,
+ #[serde(default, skip_serializing_if = "Option::is_none")]
+ pub cumulative_logprob: Option,
+ #[serde(default, skip_serializing_if = "Option::is_none")]
+ pub avg_logprob: Option,
}
impl Example {
diff --git a/crates/edit_prediction_cli/src/format_prompt.rs b/crates/edit_prediction_cli/src/format_prompt.rs
index b0cfbd77ee543d4271cda0fb952f5ba48fc4a998..2a1b49007bd19e721a6d95ebddda3758c86aaaef 100644
--- a/crates/edit_prediction_cli/src/format_prompt.rs
+++ b/crates/edit_prediction_cli/src/format_prompt.rs
@@ -150,6 +150,20 @@ pub fn zeta2_output_for_patch(
);
}
+ if version == ZetaFormat::V0318SeedMultiRegions {
+ let cursor_in_new = cursor_offset.map(|cursor_offset| {
+ let hunk_start = first_hunk_offset.unwrap_or(0);
+ result.floor_char_boundary((hunk_start + cursor_offset).min(result.len()))
+ });
+ return multi_region::encode_from_old_and_new_v0318(
+ &old_editable_region,
+ &result,
+ cursor_in_new,
+ zeta_prompt::CURSOR_MARKER,
+ multi_region::V0318_END_MARKER,
+ );
+ }
+
if version == ZetaFormat::V0316SeedMultiRegions {
let cursor_in_new = cursor_offset.map(|cursor_offset| {
let hunk_start = first_hunk_offset.unwrap_or(0);
@@ -237,7 +251,10 @@ impl TeacherPrompt {
}
}
- if response.trim().ends_with(Self::NO_EDITS) {
+ if response
+ .trim_end_matches(&[' ', '\n', '`'])
+ .ends_with(Self::NO_EDITS)
+ {
return Ok(no_edits);
}
@@ -872,4 +889,42 @@ mod tests {
let result = extract_last_codeblock(text).unwrap();
assert_eq!(result, "content here\n");
}
+
+ #[test]
+ fn test_parse_no_edits_response_with_trailing_backticks() {
+ let response = "NO_EDITS```";
+
+ let parsed = TeacherPrompt::parse(
+ &Example {
+ spec: edit_prediction::example_spec::ExampleSpec {
+ name: "test".to_string(),
+ repository_url: "https://github.com/zed-industries/zed.git".to_string(),
+ revision: "HEAD".to_string(),
+ tags: Vec::new(),
+ reasoning: None,
+ uncommitted_diff: String::new(),
+ cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
+ cursor_position: "0:0".to_string(),
+ edit_history: String::new(),
+ expected_patches: Vec::new(),
+ rejected_patch: None,
+ telemetry: None,
+ human_feedback: Vec::new(),
+ rating: None,
+ },
+ prompt_inputs: None,
+ prompt: None,
+ predictions: Vec::new(),
+ score: Vec::new(),
+ qa: Vec::new(),
+ zed_version: None,
+ state: None,
+ },
+ response,
+ )
+ .unwrap();
+
+ assert!(parsed.0.is_empty());
+ assert!(parsed.1.is_none());
+ }
}
diff --git a/crates/edit_prediction_cli/src/predict.rs b/crates/edit_prediction_cli/src/predict.rs
index 9f70861b5ef7298141441ec09606fa77e341cbfd..df797b0abaa4933e73e40b746797ffb5581d7f79 100644
--- a/crates/edit_prediction_cli/src/predict.rs
+++ b/crates/edit_prediction_cli/src/predict.rs
@@ -263,6 +263,8 @@ pub async fn run_prediction(
actual_cursor: None,
error: None,
provider,
+ cumulative_logprob: None,
+ avg_logprob: None,
});
step_progress.set_substatus("requesting prediction");
@@ -455,6 +457,8 @@ async fn predict_anthropic(
_ => PredictionProvider::TeacherNonBatching(backend),
}
},
+ cumulative_logprob: None,
+ avg_logprob: None,
};
example.predictions.push(prediction);
@@ -572,6 +576,8 @@ async fn predict_openai(
_ => PredictionProvider::TeacherNonBatching(backend),
}
},
+ cumulative_logprob: None,
+ avg_logprob: None,
};
example.predictions.push(prediction);
@@ -656,6 +662,8 @@ pub async fn predict_baseten(
actual_cursor,
error: None,
provider: PredictionProvider::Baseten(format),
+ cumulative_logprob: None,
+ avg_logprob: None,
};
example.predictions.push(prediction);
diff --git a/crates/edit_prediction_cli/src/repair.rs b/crates/edit_prediction_cli/src/repair.rs
index 9d891314bc62a44e730b584cea3423df665dc381..58da6c47e91491cc785804c7f4c2aab30887a741 100644
--- a/crates/edit_prediction_cli/src/repair.rs
+++ b/crates/edit_prediction_cli/src/repair.rs
@@ -10,7 +10,7 @@ use crate::{
BatchProvider, PredictionProvider,
anthropic_client::AnthropicClient,
example::{ActualCursor, Example, ExamplePrediction},
- format_prompt::{TeacherPrompt, extract_last_codeblock},
+ format_prompt::TeacherPrompt,
metrics::count_patch_token_changes,
openai_client::OpenAiClient,
parse_output::run_parse_output,
@@ -227,10 +227,7 @@ pub fn needs_repair(example: &Example, confidence_threshold: u8) -> bool {
/// Handles the `KEEP_PREVIOUS` sentinel by copying the teacher's prediction,
/// and delegates normal output to `TeacherPrompt::parse`.
pub fn parse(example: &Example, actual_output: &str) -> Result<(String, Option)> {
- let last_codeblock =
- extract_last_codeblock(actual_output).unwrap_or_else(|| actual_output.to_string());
-
- if last_codeblock.contains(KEEP_PREVIOUS) {
+ if actual_output.contains(KEEP_PREVIOUS) {
let original = example
.predictions
.first()
@@ -426,6 +423,8 @@ pub async fn run_repair(
actual_cursor,
error: err,
provider: PredictionProvider::Repair,
+ cumulative_logprob: None,
+ avg_logprob: None,
});
Ok(())
@@ -454,3 +453,71 @@ pub async fn sync_batches(args: &RepairArgs) -> Result<()> {
Ok(())
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::{PredictionProvider, TeacherBackend};
+ use edit_prediction::example_spec::ExampleSpec;
+ use std::{path::Path, sync::Arc};
+
+ fn example_with_previous_prediction() -> Example {
+ Example {
+ spec: ExampleSpec {
+ name: "example".to_string(),
+ repository_url: "https://github.com/zed-industries/zed.git".to_string(),
+ revision: "HEAD".to_string(),
+ tags: Vec::new(),
+ reasoning: None,
+ uncommitted_diff: String::new(),
+ cursor_path: Arc::from(Path::new("src/main.rs")),
+ cursor_position: "0:0".to_string(),
+ edit_history: String::new(),
+ expected_patches: Vec::new(),
+ rejected_patch: None,
+ telemetry: None,
+ human_feedback: Vec::new(),
+ rating: None,
+ },
+ prompt_inputs: None,
+ prompt: None,
+ predictions: vec![ExamplePrediction {
+ actual_patch: Some("previous patch".to_string()),
+ actual_output: String::new(),
+ actual_cursor: Some(ActualCursor {
+ path: "src/main.rs".to_string(),
+ row: 1,
+ column: 2,
+ offset: 3,
+ editable_region_offset: Some(4),
+ }),
+ error: None,
+ provider: PredictionProvider::Teacher(TeacherBackend::Sonnet45),
+ cumulative_logprob: None,
+ avg_logprob: None,
+ }],
+ score: Vec::new(),
+ qa: Vec::new(),
+ zed_version: None,
+ state: None,
+ }
+ }
+
+ #[test]
+ fn test_parse_keeps_previous_when_sentinel_appears_outside_last_codeblock() {
+ let example = example_with_previous_prediction();
+ let actual_output = indoc::indoc! {"
+ After reviewing the feedback, the previous prediction is still correct.
+ Use `KEEP_PREVIOUS`.
+
+ ```
+ unrelated trailing code block
+ ```
+ "};
+
+ let (patch, cursor) = parse(&example, actual_output).unwrap();
+
+ assert_eq!(patch, "previous patch");
+ assert_eq!(cursor.unwrap().offset, 3);
+ }
+}
diff --git a/crates/edit_prediction_cli/src/score.rs b/crates/edit_prediction_cli/src/score.rs
index b6f745114f6dd2a091b95b724ee53869a04a8c4e..d75cf55e85b198bc28469e83d8f9209a8a59a83f 100644
--- a/crates/edit_prediction_cli/src/score.rs
+++ b/crates/edit_prediction_cli/src/score.rs
@@ -78,6 +78,8 @@ pub async fn run_scoring(
has_isolated_whitespace_changes: false,
inserted_tokens: 0,
deleted_tokens: 0,
+ cumulative_logprob: None,
+ avg_logprob: None,
};
let cursor_path = example.spec.cursor_path.as_ref();
@@ -189,6 +191,8 @@ pub async fn run_scoring(
has_isolated_whitespace_changes,
inserted_tokens: token_changes.inserted_tokens,
deleted_tokens: token_changes.deleted_tokens,
+ cumulative_logprob: prediction.cumulative_logprob,
+ avg_logprob: prediction.avg_logprob,
});
}
diff --git a/crates/edit_prediction_context/src/edit_prediction_context_tests.rs b/crates/edit_prediction_context/src/edit_prediction_context_tests.rs
index 01c4c76e82eb0851b7552b3d9117af1212a8b3da..78ded78b7eb558c9bb5d1839a8c8c82290a13d9a 100644
--- a/crates/edit_prediction_context/src/edit_prediction_context_tests.rs
+++ b/crates/edit_prediction_context/src/edit_prediction_context_tests.rs
@@ -1028,6 +1028,7 @@ fn assert_related_files_impl(
pretty_assertions::assert_eq!(actual, expected)
}
+#[track_caller]
fn assert_definitions(definitions: &[LocationLink], first_lines: &[&str], cx: &mut TestAppContext) {
let actual_first_lines = definitions
.iter()
diff --git a/crates/edit_prediction_context/src/fake_definition_lsp.rs b/crates/edit_prediction_context/src/fake_definition_lsp.rs
index 6b6d93469b9a1fbeb856e189f4fe79da06135045..5b9e528b63f6709ce2966d2fd54d39eeeb195a36 100644
--- a/crates/edit_prediction_context/src/fake_definition_lsp.rs
+++ b/crates/edit_prediction_context/src/fake_definition_lsp.rs
@@ -174,7 +174,7 @@ pub fn register_fake_definition_server(
struct DefinitionIndex {
language: Arc,
definitions: HashMap>,
- type_annotations: HashMap,
+ type_annotations_by_file: HashMap>,
files: HashMap,
}
@@ -189,7 +189,7 @@ impl DefinitionIndex {
Self {
language,
definitions: HashMap::default(),
- type_annotations: HashMap::default(),
+ type_annotations_by_file: HashMap::default(),
files: HashMap::default(),
}
}
@@ -199,6 +199,7 @@ impl DefinitionIndex {
locations.retain(|loc| &loc.uri != uri);
!locations.is_empty()
});
+ self.type_annotations_by_file.remove(uri);
self.files.remove(uri);
}
@@ -243,11 +244,11 @@ impl DefinitionIndex {
.push(location);
}
- for (identifier_name, type_name) in extract_type_annotations(content) {
- self.type_annotations
- .entry(identifier_name)
- .or_insert(type_name);
- }
+ let type_annotations = extract_type_annotations(content)
+ .into_iter()
+ .collect::>();
+ self.type_annotations_by_file
+ .insert(uri.clone(), type_annotations);
self.files.insert(
uri,
@@ -279,7 +280,11 @@ impl DefinitionIndex {
let entry = self.files.get(&uri)?;
let name = word_at_position(&entry.contents, position)?;
- if let Some(type_name) = self.type_annotations.get(name) {
+ if let Some(type_name) = self
+ .type_annotations_by_file
+ .get(&uri)
+ .and_then(|annotations| annotations.get(name))
+ {
if let Some(locations) = self.definitions.get(type_name) {
return Some(lsp::GotoDefinitionResponse::Array(locations.clone()));
}
@@ -367,6 +372,20 @@ fn extract_base_type_name(type_str: &str) -> String {
return outer.to_string();
}
+ if let Some(call_start) = trimmed.find("::") {
+ let outer = &trimmed[..call_start];
+ if matches!(outer, "Arc" | "Box" | "Rc" | "Option" | "Vec" | "Cow") {
+ let rest = trimmed[call_start + 2..].trim_start();
+ if let Some(paren_start) = rest.find('(') {
+ let inner = &rest[paren_start + 1..];
+ let inner = inner.trim();
+ if !inner.is_empty() {
+ return extract_base_type_name(inner);
+ }
+ }
+ }
+ }
+
trimmed
.split(|c: char| !c.is_alphanumeric() && c != '_')
.next()
diff --git a/crates/editor/src/bracket_colorization.rs b/crates/editor/src/bracket_colorization.rs
index 657f1e1b23d91ca421da6a38fbeaa382a65863db..ad2fc1bd8b9666dfa5e2c4b0367984c6398c98f8 100644
--- a/crates/editor/src/bracket_colorization.rs
+++ b/crates/editor/src/bracket_colorization.rs
@@ -392,6 +392,20 @@ where
&bracket_colors_markup(&mut cx),
"All markdown brackets should be colored based on their depth, again"
);
+
+ cx.set_state(indoc! {r#"ˇ('')('')
+
+((''))('')
+
+('')((''))"#});
+ cx.executor().advance_clock(Duration::from_millis(100));
+ cx.executor().run_until_parked();
+
+ assert_eq!(
+ "«1('')1»«1('')1»\n\n«1(«2('')2»)1»«1('')1»\n\n«1('')1»«1(«2('')2»)1»\n1 hsla(207.80, 16.20%, 69.19%, 1.00)\n2 hsla(29.00, 54.00%, 65.88%, 1.00)\n",
+ &bracket_colors_markup(&mut cx),
+ "Markdown quote pairs should not interfere with parenthesis pairing"
+ );
}
#[gpui::test]
diff --git a/crates/editor/src/edit_prediction_tests.rs b/crates/editor/src/edit_prediction_tests.rs
index 40d915b841c1cf4a53025f7dce654dd129cb8de5..684213e481762d7fb09a0bd6d8b7a0b9fc6d4a36 100644
--- a/crates/editor/src/edit_prediction_tests.rs
+++ b/crates/editor/src/edit_prediction_tests.rs
@@ -1,13 +1,17 @@
use edit_prediction_types::{
EditPredictionDelegate, EditPredictionIconSet, PredictedCursorPosition,
};
-use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
+use gpui::{
+ Entity, KeyBinding, KeybindingKeystroke, Keystroke, Modifiers, NoAction, Task, prelude::*,
+};
use indoc::indoc;
-use language::Buffer;
use language::EditPredictionsMode;
-use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
+use language::{Buffer, CodeLabel};
+use multi_buffer::{Anchor, ExcerptId, MultiBufferSnapshot, ToPoint};
+use project::{Completion, CompletionResponse, CompletionSource};
use std::{
ops::Range,
+ rc::Rc,
sync::{
Arc,
atomic::{self, AtomicUsize},
@@ -17,8 +21,9 @@ use text::{Point, ToOffset};
use ui::prelude::*;
use crate::{
- AcceptEditPrediction, EditPrediction, EditPredictionKeybindAction,
- EditPredictionKeybindSurface, MenuEditPredictionsPolicy,
+ AcceptEditPrediction, CompletionContext, CompletionProvider, EditPrediction,
+ EditPredictionKeybindAction, EditPredictionKeybindSurface, MenuEditPredictionsPolicy,
+ ShowCompletions,
editor_tests::{init_test, update_test_language_settings},
test::editor_test_context::EditorTestContext,
};
@@ -482,57 +487,10 @@ async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestA
});
}
-fn load_default_keymap(cx: &mut gpui::TestAppContext) {
- cx.update(|cx| {
- cx.bind_keys(
- settings::KeymapFile::load_asset_allow_partial_failure(
- settings::DEFAULT_KEYMAP_PATH,
- cx,
- )
- .expect("failed to load default keymap"),
- );
- });
-}
-
#[gpui::test]
-async fn test_tab_is_preferred_accept_binding_over_alt_tab(cx: &mut gpui::TestAppContext) {
- init_test(cx, |_| {});
- load_default_keymap(cx);
-
- let mut cx = EditorTestContext::new(cx).await;
- let provider = cx.new(|_| FakeEditPredictionDelegate::default());
- assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.set_state("let x = ˇ;");
-
- propose_edits(&provider, vec![(8..8, "42")], &mut cx);
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
-
- cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
- let keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::Inline,
- window,
- cx,
- );
- let keystroke = keybind_display
- .accept_keystroke
- .as_ref()
- .expect("should have an accept binding");
- assert!(
- !keystroke.modifiers().modified(),
- "preferred accept binding should be unmodified (tab), got modifiers: {:?}",
- keystroke.modifiers()
- );
- assert_eq!(
- keystroke.key(),
- "tab",
- "preferred accept binding should be tab"
- );
- });
-}
-
-#[gpui::test]
-async fn test_subtle_in_code_indicator_prefers_preview_binding(cx: &mut gpui::TestAppContext) {
+async fn test_edit_prediction_preview_activates_when_prediction_arrives_with_modifier_held(
+ cx: &mut gpui::TestAppContext,
+) {
init_test(cx, |_| {});
load_default_keymap(cx);
update_test_language_settings(cx, &|settings| {
@@ -544,227 +502,324 @@ async fn test_subtle_in_code_indicator_prefers_preview_binding(cx: &mut gpui::Te
assign_editor_completion_provider(provider.clone(), &mut cx);
cx.set_state("let x = ˇ;");
- propose_edits(&provider, vec![(8..8, "42")], &mut cx);
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
-
- cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
- assert!(
- editor.edit_prediction_requires_modifier(),
- "subtle mode should require a modifier"
- );
-
- let inline_keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::Inline,
- window,
- cx,
- );
- let compact_keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::CursorPopoverCompact,
- window,
- cx,
- );
-
- let accept_keystroke = inline_keybind_display
- .accept_keystroke
- .as_ref()
- .expect("should have an accept binding");
- let preview_keystroke = inline_keybind_display
- .preview_keystroke
- .as_ref()
- .expect("should have a preview binding");
- let in_code_keystroke = inline_keybind_display
- .displayed_keystroke
- .as_ref()
- .expect("should have an in-code binding");
- let compact_cursor_popover_keystroke = compact_keybind_display
- .displayed_keystroke
- .as_ref()
- .expect("should have a compact cursor popover binding");
-
- assert_eq!(accept_keystroke.key(), "tab");
- assert!(
- !editor.has_visible_completions_menu(),
- "compact cursor-popover branch should be used without a completions menu"
- );
- assert!(
- preview_keystroke.modifiers().modified(),
- "preview binding should use modifiers in subtle mode"
- );
- assert_eq!(
- compact_cursor_popover_keystroke.key(),
- preview_keystroke.key(),
- "subtle compact cursor popover should prefer the preview binding"
- );
- assert_eq!(
- compact_cursor_popover_keystroke.modifiers(),
- preview_keystroke.modifiers(),
- "subtle compact cursor popover should use the preview binding modifiers"
- );
- assert_eq!(
- in_code_keystroke.key(),
- preview_keystroke.key(),
- "subtle in-code indicator should prefer the preview binding"
- );
- assert_eq!(
- in_code_keystroke.modifiers(),
- preview_keystroke.modifiers(),
- "subtle in-code indicator should use the preview binding modifiers"
- );
+ cx.editor(|editor, _, _| {
+ assert!(!editor.has_active_edit_prediction());
+ assert!(!editor.edit_prediction_preview_is_active());
});
-}
-#[gpui::test]
-async fn test_tab_accepts_edit_prediction_over_completion(cx: &mut gpui::TestAppContext) {
- init_test(cx, |_| {});
- load_default_keymap(cx);
-
- let mut cx = EditorTestContext::new(cx).await;
- let provider = cx.new(|_| FakeEditPredictionDelegate::default());
- assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.set_state("let x = ˇ;");
-
- propose_edits(&provider, vec![(8..8, "42")], &mut cx);
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
-
- assert_editor_active_edit_completion(&mut cx, |_, edits| {
- assert_eq!(edits.len(), 1);
- assert_eq!(edits[0].1.as_ref(), "42");
+ let preview_modifiers = cx.update_editor(|editor, window, cx| {
+ *editor
+ .preview_edit_prediction_keystroke(window, cx)
+ .unwrap()
+ .modifiers()
});
- cx.simulate_keystroke("tab");
+ cx.simulate_modifiers_change(preview_modifiers);
cx.run_until_parked();
- cx.assert_editor_state("let x = 42ˇ;");
-}
-
-#[gpui::test]
-async fn test_single_line_prediction_uses_accept_cursor_popover_action(
- cx: &mut gpui::TestAppContext,
-) {
- init_test(cx, |_| {});
- load_default_keymap(cx);
-
- let mut cx = EditorTestContext::new(cx).await;
- let provider = cx.new(|_| FakeEditPredictionDelegate::default());
- assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.set_state("let x = ˇ;");
+ cx.editor(|editor, _, _| {
+ assert!(!editor.has_active_edit_prediction());
+ assert!(editor.edit_prediction_preview_is_active());
+ });
propose_edits(&provider, vec![(8..8, "42")], &mut cx);
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
-
cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
+ editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
+ editor.update_visible_edit_prediction(window, cx)
+ });
- let keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::CursorPopoverExpanded,
- window,
- cx,
+ cx.editor(|editor, _, _| {
+ assert!(editor.has_active_edit_prediction());
+ assert!(
+ editor.edit_prediction_preview_is_active(),
+ "prediction preview should activate immediately when the prediction arrives while the preview modifier is still held",
);
+ });
+}
- let accept_keystroke = keybind_display
- .accept_keystroke
- .as_ref()
- .expect("should have an accept binding");
- let preview_keystroke = keybind_display
- .preview_keystroke
- .as_ref()
- .expect("should have a preview binding");
-
- assert_eq!(
- keybind_display.action,
- EditPredictionKeybindAction::Accept,
- "single-line prediction should show the accept action"
+fn load_default_keymap(cx: &mut gpui::TestAppContext) {
+ cx.update(|cx| {
+ cx.bind_keys(
+ settings::KeymapFile::load_asset_allow_partial_failure(
+ settings::DEFAULT_KEYMAP_PATH,
+ cx,
+ )
+ .expect("failed to load default keymap"),
);
- assert_eq!(accept_keystroke.key(), "tab");
- assert!(preview_keystroke.modifiers().modified());
});
}
#[gpui::test]
-async fn test_multi_line_prediction_uses_preview_cursor_popover_action(
- cx: &mut gpui::TestAppContext,
-) {
- init_test(cx, |_| {});
- load_default_keymap(cx);
-
- let mut cx = EditorTestContext::new(cx).await;
- let provider = cx.new(|_| FakeEditPredictionDelegate::default());
- assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.set_state("let x = ˇ;");
-
- propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
-
- cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
+async fn test_inline_edit_prediction_keybind_selection_cases(cx: &mut gpui::TestAppContext) {
+ enum InlineKeybindState {
+ Normal,
+ ShowingCompletions,
+ InLeadingWhitespace,
+ ShowingCompletionsAndLeadingWhitespace,
+ }
- let keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::CursorPopoverExpanded,
- window,
- cx,
- );
- let preview_keystroke = keybind_display
- .preview_keystroke
- .as_ref()
- .expect("should have a preview binding");
+ enum ExpectedKeystroke {
+ DefaultAccept,
+ DefaultPreview,
+ Literal(&'static str),
+ }
- assert_eq!(
- keybind_display.action,
- EditPredictionKeybindAction::Preview,
- "multi-line prediction should show the preview action"
- );
- assert!(preview_keystroke.modifiers().modified());
- });
-}
+ struct InlineKeybindCase {
+ name: &'static str,
+ use_default_keymap: bool,
+ mode: EditPredictionsMode,
+ extra_bindings: Vec,
+ state: InlineKeybindState,
+ expected_accept_keystroke: ExpectedKeystroke,
+ expected_preview_keystroke: ExpectedKeystroke,
+ expected_displayed_keystroke: ExpectedKeystroke,
+ }
-#[gpui::test]
-async fn test_single_line_prediction_with_preview_uses_accept_cursor_popover_action(
- cx: &mut gpui::TestAppContext,
-) {
init_test(cx, |_| {});
load_default_keymap(cx);
+ let mut default_cx = EditorTestContext::new(cx).await;
+ let provider = default_cx.new(|_| FakeEditPredictionDelegate::default());
+ assign_editor_completion_provider(provider.clone(), &mut default_cx);
+ default_cx.set_state("let x = ˇ;");
+ propose_edits(&provider, vec![(8..8, "42")], &mut default_cx);
+ default_cx
+ .update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
+
+ let (default_accept_keystroke, default_preview_keystroke) =
+ default_cx.update_editor(|editor, window, cx| {
+ let keybind_display = editor.edit_prediction_keybind_display(
+ EditPredictionKeybindSurface::Inline,
+ window,
+ cx,
+ );
+ let accept_keystroke = keybind_display
+ .accept_keystroke
+ .as_ref()
+ .expect("default inline edit prediction should have an accept binding")
+ .clone();
+ let preview_keystroke = keybind_display
+ .preview_keystroke
+ .as_ref()
+ .expect("default inline edit prediction should have a preview binding")
+ .clone();
+ (accept_keystroke, preview_keystroke)
+ });
+
+ let cases = [
+ InlineKeybindCase {
+ name: "default setup prefers tab over alt-tab for accept",
+ use_default_keymap: true,
+ mode: EditPredictionsMode::Eager,
+ extra_bindings: Vec::new(),
+ state: InlineKeybindState::Normal,
+ expected_accept_keystroke: ExpectedKeystroke::DefaultAccept,
+ expected_preview_keystroke: ExpectedKeystroke::DefaultPreview,
+ expected_displayed_keystroke: ExpectedKeystroke::DefaultAccept,
+ },
+ InlineKeybindCase {
+ name: "subtle mode displays preview binding inline",
+ use_default_keymap: true,
+ mode: EditPredictionsMode::Subtle,
+ extra_bindings: Vec::new(),
+ state: InlineKeybindState::Normal,
+ expected_accept_keystroke: ExpectedKeystroke::DefaultPreview,
+ expected_preview_keystroke: ExpectedKeystroke::DefaultPreview,
+ expected_displayed_keystroke: ExpectedKeystroke::DefaultPreview,
+ },
+ InlineKeybindCase {
+ name: "removing default tab binding still displays tab",
+ use_default_keymap: true,
+ mode: EditPredictionsMode::Eager,
+ extra_bindings: vec![KeyBinding::new(
+ "tab",
+ NoAction,
+ Some("Editor && edit_prediction && edit_prediction_mode == eager"),
+ )],
+ state: InlineKeybindState::Normal,
+ expected_accept_keystroke: ExpectedKeystroke::DefaultPreview,
+ expected_preview_keystroke: ExpectedKeystroke::DefaultPreview,
+ expected_displayed_keystroke: ExpectedKeystroke::DefaultPreview,
+ },
+ InlineKeybindCase {
+ name: "custom-only rebound accept key uses replacement key",
+ use_default_keymap: true,
+ mode: EditPredictionsMode::Eager,
+ extra_bindings: vec![KeyBinding::new(
+ "ctrl-enter",
+ AcceptEditPrediction,
+ Some("Editor && edit_prediction"),
+ )],
+ state: InlineKeybindState::Normal,
+ expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ },
+ InlineKeybindCase {
+ name: "showing completions restores conflict-context binding",
+ use_default_keymap: true,
+ mode: EditPredictionsMode::Eager,
+ extra_bindings: vec![KeyBinding::new(
+ "ctrl-enter",
+ AcceptEditPrediction,
+ Some("Editor && edit_prediction && showing_completions"),
+ )],
+ state: InlineKeybindState::ShowingCompletions,
+ expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ },
+ InlineKeybindCase {
+ name: "leading whitespace restores conflict-context binding",
+ use_default_keymap: false,
+ mode: EditPredictionsMode::Eager,
+ extra_bindings: vec![KeyBinding::new(
+ "ctrl-enter",
+ AcceptEditPrediction,
+ Some("Editor && edit_prediction && in_leading_whitespace"),
+ )],
+ state: InlineKeybindState::InLeadingWhitespace,
+ expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ },
+ InlineKeybindCase {
+ name: "showing completions and leading whitespace restore combined conflict binding",
+ use_default_keymap: false,
+ mode: EditPredictionsMode::Eager,
+ extra_bindings: vec![KeyBinding::new(
+ "ctrl-enter",
+ AcceptEditPrediction,
+ Some("Editor && edit_prediction && showing_completions && in_leading_whitespace"),
+ )],
+ state: InlineKeybindState::ShowingCompletionsAndLeadingWhitespace,
+ expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
+ },
+ ];
+
+ for case in cases {
+ init_test(cx, |_| {});
+ if case.use_default_keymap {
+ load_default_keymap(cx);
+ }
+ update_test_language_settings(cx, &|settings| {
+ settings.edit_predictions.get_or_insert_default().mode = Some(case.mode);
+ });
- let mut cx = EditorTestContext::new(cx).await;
- let provider = cx.new(|_| FakeEditPredictionDelegate::default());
- assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.set_state("let x = ˇ;");
+ if !case.extra_bindings.is_empty() {
+ cx.update(|cx| cx.bind_keys(case.extra_bindings.clone()));
+ }
- propose_edits_with_preview(&provider, vec![(8..8, "42")], &mut cx).await;
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
+ let mut cx = EditorTestContext::new(cx).await;
+ let provider = cx.new(|_| FakeEditPredictionDelegate::default());
+ assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
+ match case.state {
+ InlineKeybindState::Normal | InlineKeybindState::ShowingCompletions => {
+ cx.set_state("let x = ˇ;");
+ }
+ InlineKeybindState::InLeadingWhitespace
+ | InlineKeybindState::ShowingCompletionsAndLeadingWhitespace => {
+ cx.set_state(indoc! {"
+ fn main() {
+ ˇ
+ }
+ "});
+ }
+ }
- let keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::CursorPopoverExpanded,
- window,
- cx,
- );
+ propose_edits(&provider, vec![(8..8, "42")], &mut cx);
+ cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
+
+ if matches!(
+ case.state,
+ InlineKeybindState::ShowingCompletions
+ | InlineKeybindState::ShowingCompletionsAndLeadingWhitespace
+ ) {
+ assign_editor_completion_menu_provider(&mut cx);
+ cx.update_editor(|editor, window, cx| {
+ editor.show_completions(&ShowCompletions, window, cx);
+ });
+ cx.run_until_parked();
+ }
- let accept_keystroke = keybind_display
- .accept_keystroke
- .as_ref()
- .expect("should have an accept binding");
- let preview_keystroke = keybind_display
- .preview_keystroke
- .as_ref()
- .expect("should have a preview binding");
+ cx.update_editor(|editor, window, cx| {
+ assert!(
+ editor.has_active_edit_prediction(),
+ "case '{}' should have an active edit prediction",
+ case.name
+ );
- assert_eq!(
- keybind_display.action,
- EditPredictionKeybindAction::Accept,
- "single-line prediction should show the accept action even with edit_preview"
- );
- assert_eq!(accept_keystroke.key(), "tab");
- assert!(preview_keystroke.modifiers().modified());
- });
+ let keybind_display = editor.edit_prediction_keybind_display(
+ EditPredictionKeybindSurface::Inline,
+ window,
+ cx,
+ );
+ let accept_keystroke = keybind_display
+ .accept_keystroke
+ .as_ref()
+ .unwrap_or_else(|| panic!("case '{}' should have an accept binding", case.name));
+ let preview_keystroke = keybind_display
+ .preview_keystroke
+ .as_ref()
+ .unwrap_or_else(|| panic!("case '{}' should have a preview binding", case.name));
+ let displayed_keystroke = keybind_display
+ .displayed_keystroke
+ .as_ref()
+ .unwrap_or_else(|| panic!("case '{}' should have a displayed binding", case.name));
+
+ let expected_accept_keystroke = match case.expected_accept_keystroke {
+ ExpectedKeystroke::DefaultAccept => default_accept_keystroke.clone(),
+ ExpectedKeystroke::DefaultPreview => default_preview_keystroke.clone(),
+ ExpectedKeystroke::Literal(keystroke) => KeybindingKeystroke::from_keystroke(
+ Keystroke::parse(keystroke).expect("expected test keystroke to parse"),
+ ),
+ };
+ let expected_preview_keystroke = match case.expected_preview_keystroke {
+ ExpectedKeystroke::DefaultAccept => default_accept_keystroke.clone(),
+ ExpectedKeystroke::DefaultPreview => default_preview_keystroke.clone(),
+ ExpectedKeystroke::Literal(keystroke) => KeybindingKeystroke::from_keystroke(
+ Keystroke::parse(keystroke).expect("expected test keystroke to parse"),
+ ),
+ };
+ let expected_displayed_keystroke = match case.expected_displayed_keystroke {
+ ExpectedKeystroke::DefaultAccept => default_accept_keystroke.clone(),
+ ExpectedKeystroke::DefaultPreview => default_preview_keystroke.clone(),
+ ExpectedKeystroke::Literal(keystroke) => KeybindingKeystroke::from_keystroke(
+ Keystroke::parse(keystroke).expect("expected test keystroke to parse"),
+ ),
+ };
+
+ assert_eq!(
+ accept_keystroke, &expected_accept_keystroke,
+ "case '{}' selected the wrong accept binding",
+ case.name
+ );
+ assert_eq!(
+ preview_keystroke, &expected_preview_keystroke,
+ "case '{}' selected the wrong preview binding",
+ case.name
+ );
+ assert_eq!(
+ displayed_keystroke, &expected_displayed_keystroke,
+ "case '{}' selected the wrong displayed binding",
+ case.name
+ );
+
+ if matches!(case.mode, EditPredictionsMode::Subtle) {
+ assert!(
+ editor.edit_prediction_requires_modifier(),
+ "case '{}' should require a modifier",
+ case.name
+ );
+ }
+ });
+ }
}
#[gpui::test]
-async fn test_multi_line_prediction_with_preview_uses_preview_cursor_popover_action(
- cx: &mut gpui::TestAppContext,
-) {
+async fn test_tab_accepts_edit_prediction_over_completion(cx: &mut gpui::TestAppContext) {
init_test(cx, |_| {});
load_default_keymap(cx);
@@ -773,131 +828,194 @@ async fn test_multi_line_prediction_with_preview_uses_preview_cursor_popover_act
assign_editor_completion_provider(provider.clone(), &mut cx);
cx.set_state("let x = ˇ;");
- propose_edits_with_preview(&provider, vec![(8..8, "42\n43")], &mut cx).await;
+ propose_edits(&provider, vec![(8..8, "42")], &mut cx);
cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
- cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
- let keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::CursorPopoverExpanded,
- window,
- cx,
- );
- let preview_keystroke = keybind_display
- .preview_keystroke
- .as_ref()
- .expect("should have a preview binding");
-
- assert_eq!(
- keybind_display.action,
- EditPredictionKeybindAction::Preview,
- "multi-line prediction should show the preview action with edit_preview"
- );
- assert!(preview_keystroke.modifiers().modified());
+ assert_editor_active_edit_completion(&mut cx, |_, edits| {
+ assert_eq!(edits.len(), 1);
+ assert_eq!(edits[0].1.as_ref(), "42");
});
-}
-
-#[gpui::test]
-async fn test_single_line_deletion_of_newline_uses_accept_cursor_popover_action(
- cx: &mut gpui::TestAppContext,
-) {
- init_test(cx, |_| {});
- load_default_keymap(cx);
-
- let mut cx = EditorTestContext::new(cx).await;
- let provider = cx.new(|_| FakeEditPredictionDelegate::default());
- assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.set_state(indoc! {"
- fn main() {
- let value = 1;
- ˇprintln!(\"done\");
- }
- "});
-
- propose_edits(
- &provider,
- vec![(Point::new(1, 18)..Point::new(2, 17), "")],
- &mut cx,
- );
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
-
- cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
- let keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::CursorPopoverExpanded,
- window,
- cx,
- );
-
- let accept_keystroke = keybind_display
- .accept_keystroke
- .as_ref()
- .expect("should have an accept binding");
- let preview_keystroke = keybind_display
- .preview_keystroke
- .as_ref()
- .expect("should have a preview binding");
+ cx.simulate_keystroke("tab");
+ cx.run_until_parked();
- assert_eq!(
- keybind_display.action,
- EditPredictionKeybindAction::Accept,
- "deleting one newline plus adjacent text should show the accept action"
- );
- assert_eq!(accept_keystroke.key(), "tab");
- assert!(preview_keystroke.modifiers().modified());
- });
+ cx.assert_editor_state("let x = 42ˇ;");
}
#[gpui::test]
-async fn test_stale_single_line_prediction_does_not_force_preview_cursor_popover_action(
- cx: &mut gpui::TestAppContext,
-) {
- init_test(cx, |_| {});
- load_default_keymap(cx);
-
- let mut cx = EditorTestContext::new(cx).await;
- let provider = cx.new(|_| FakeEditPredictionDelegate::default());
- assign_editor_completion_provider(provider.clone(), &mut cx);
- cx.set_state("let x = ˇ;");
-
- propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
- cx.update_editor(|editor, _window, cx| {
- assert!(editor.active_edit_prediction.is_some());
- assert!(editor.stale_edit_prediction_in_menu.is_none());
- editor.take_active_edit_prediction(cx);
- assert!(editor.active_edit_prediction.is_none());
- assert!(editor.stale_edit_prediction_in_menu.is_some());
- });
+async fn test_cursor_popover_edit_prediction_keybind_cases(cx: &mut gpui::TestAppContext) {
+ enum CursorPopoverPredictionKind {
+ SingleLine,
+ MultiLine,
+ SingleLineWithPreview,
+ MultiLineWithPreview,
+ DeleteSingleNewline,
+ StaleSingleLineAfterMultiLine,
+ }
- propose_edits(&provider, vec![(8..8, "42")], &mut cx);
- cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
+ struct CursorPopoverCase {
+ name: &'static str,
+ prediction_kind: CursorPopoverPredictionKind,
+ expected_action: EditPredictionKeybindAction,
+ }
- cx.update_editor(|editor, window, cx| {
- assert!(editor.has_active_edit_prediction());
+ let cases = [
+ CursorPopoverCase {
+ name: "single line prediction uses accept action",
+ prediction_kind: CursorPopoverPredictionKind::SingleLine,
+ expected_action: EditPredictionKeybindAction::Accept,
+ },
+ CursorPopoverCase {
+ name: "multi line prediction uses preview action",
+ prediction_kind: CursorPopoverPredictionKind::MultiLine,
+ expected_action: EditPredictionKeybindAction::Preview,
+ },
+ CursorPopoverCase {
+ name: "single line prediction with preview still uses accept action",
+ prediction_kind: CursorPopoverPredictionKind::SingleLineWithPreview,
+ expected_action: EditPredictionKeybindAction::Accept,
+ },
+ CursorPopoverCase {
+ name: "multi line prediction with preview uses preview action",
+ prediction_kind: CursorPopoverPredictionKind::MultiLineWithPreview,
+ expected_action: EditPredictionKeybindAction::Preview,
+ },
+ CursorPopoverCase {
+ name: "single line newline deletion uses accept action",
+ prediction_kind: CursorPopoverPredictionKind::DeleteSingleNewline,
+ expected_action: EditPredictionKeybindAction::Accept,
+ },
+ CursorPopoverCase {
+ name: "stale multi line prediction does not force preview action",
+ prediction_kind: CursorPopoverPredictionKind::StaleSingleLineAfterMultiLine,
+ expected_action: EditPredictionKeybindAction::Accept,
+ },
+ ];
+
+ for case in cases {
+ init_test(cx, |_| {});
+ load_default_keymap(cx);
+
+ let mut cx = EditorTestContext::new(cx).await;
+ let provider = cx.new(|_| FakeEditPredictionDelegate::default());
+ assign_editor_completion_provider(provider.clone(), &mut cx);
+
+ match case.prediction_kind {
+ CursorPopoverPredictionKind::SingleLine => {
+ cx.set_state("let x = ˇ;");
+ propose_edits(&provider, vec![(8..8, "42")], &mut cx);
+ cx.update_editor(|editor, window, cx| {
+ editor.update_visible_edit_prediction(window, cx)
+ });
+ }
+ CursorPopoverPredictionKind::MultiLine => {
+ cx.set_state("let x = ˇ;");
+ propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
+ cx.update_editor(|editor, window, cx| {
+ editor.update_visible_edit_prediction(window, cx)
+ });
+ }
+ CursorPopoverPredictionKind::SingleLineWithPreview => {
+ cx.set_state("let x = ˇ;");
+ propose_edits_with_preview(&provider, vec![(8..8, "42")], &mut cx).await;
+ cx.update_editor(|editor, window, cx| {
+ editor.update_visible_edit_prediction(window, cx)
+ });
+ }
+ CursorPopoverPredictionKind::MultiLineWithPreview => {
+ cx.set_state("let x = ˇ;");
+ propose_edits_with_preview(&provider, vec![(8..8, "42\n43")], &mut cx).await;
+ cx.update_editor(|editor, window, cx| {
+ editor.update_visible_edit_prediction(window, cx)
+ });
+ }
+ CursorPopoverPredictionKind::DeleteSingleNewline => {
+ cx.set_state(indoc! {"
+ fn main() {
+ let value = 1;
+ ˇprintln!(\"done\");
+ }
+ "});
+ propose_edits(
+ &provider,
+ vec![(Point::new(1, 18)..Point::new(2, 17), "")],
+ &mut cx,
+ );
+ cx.update_editor(|editor, window, cx| {
+ editor.update_visible_edit_prediction(window, cx)
+ });
+ }
+ CursorPopoverPredictionKind::StaleSingleLineAfterMultiLine => {
+ cx.set_state("let x = ˇ;");
+ propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
+ cx.update_editor(|editor, window, cx| {
+ editor.update_visible_edit_prediction(window, cx)
+ });
+ cx.update_editor(|editor, _window, cx| {
+ assert!(editor.active_edit_prediction.is_some());
+ assert!(editor.stale_edit_prediction_in_menu.is_none());
+ editor.take_active_edit_prediction(cx);
+ assert!(editor.active_edit_prediction.is_none());
+ assert!(editor.stale_edit_prediction_in_menu.is_some());
+ });
+
+ propose_edits(&provider, vec![(8..8, "42")], &mut cx);
+ cx.update_editor(|editor, window, cx| {
+ editor.update_visible_edit_prediction(window, cx)
+ });
+ }
+ }
- let keybind_display = editor.edit_prediction_keybind_display(
- EditPredictionKeybindSurface::CursorPopoverExpanded,
- window,
- cx,
- );
- let accept_keystroke = keybind_display
- .accept_keystroke
- .as_ref()
- .expect("should have an accept binding");
+ cx.update_editor(|editor, window, cx| {
+ assert!(
+ editor.has_active_edit_prediction(),
+ "case '{}' should have an active edit prediction",
+ case.name
+ );
- assert_eq!(
- keybind_display.action,
- EditPredictionKeybindAction::Accept,
- "single-line active prediction should show the accept action"
- );
- assert!(
- editor.stale_edit_prediction_in_menu.is_none(),
- "refreshing the visible prediction should clear stale menu state"
- );
- assert_eq!(accept_keystroke.key(), "tab");
- });
+ let keybind_display = editor.edit_prediction_keybind_display(
+ EditPredictionKeybindSurface::CursorPopoverExpanded,
+ window,
+ cx,
+ );
+ let accept_keystroke = keybind_display
+ .accept_keystroke
+ .as_ref()
+ .unwrap_or_else(|| panic!("case '{}' should have an accept binding", case.name));
+ let preview_keystroke = keybind_display
+ .preview_keystroke
+ .as_ref()
+ .unwrap_or_else(|| panic!("case '{}' should have a preview binding", case.name));
+
+ assert_eq!(
+ keybind_display.action, case.expected_action,
+ "case '{}' selected the wrong cursor popover action",
+ case.name
+ );
+ assert_eq!(
+ accept_keystroke.key(),
+ "tab",
+ "case '{}' selected the wrong accept binding",
+ case.name
+ );
+ assert!(
+ preview_keystroke.modifiers().modified(),
+ "case '{}' should use a modified preview binding",
+ case.name
+ );
+
+ if matches!(
+ case.prediction_kind,
+ CursorPopoverPredictionKind::StaleSingleLineAfterMultiLine
+ ) {
+ assert!(
+ editor.stale_edit_prediction_in_menu.is_none(),
+ "case '{}' should clear stale menu state",
+ case.name
+ );
+ }
+ });
+ }
}
fn assert_editor_active_edit_completion(
@@ -1054,6 +1172,12 @@ fn assign_editor_completion_provider(
})
}
+fn assign_editor_completion_menu_provider(cx: &mut EditorTestContext) {
+ cx.update_editor(|editor, _, _| {
+ editor.set_completion_provider(Some(Rc::new(FakeCompletionMenuProvider)));
+ });
+}
+
fn propose_edits_non_zed(
provider: &Entity,
edits: Vec<(Range, &str)>,
@@ -1086,6 +1210,54 @@ fn assign_editor_completion_provider_non_zed(
})
}
+struct FakeCompletionMenuProvider;
+
+impl CompletionProvider for FakeCompletionMenuProvider {
+ fn completions(
+ &self,
+ _excerpt_id: ExcerptId,
+ _buffer: &Entity,
+ _buffer_position: text::Anchor,
+ _trigger: CompletionContext,
+ _window: &mut Window,
+ _cx: &mut Context,
+ ) -> Task>> {
+ let completion = Completion {
+ replace_range: text::Anchor::MIN..text::Anchor::MAX,
+ new_text: "fake_completion".to_string(),
+ label: CodeLabel::plain("fake_completion".to_string(), None),
+ documentation: None,
+ source: CompletionSource::Custom,
+ icon_path: None,
+ match_start: None,
+ snippet_deduplication_key: None,
+ insert_text_mode: None,
+ confirm: None,
+ };
+
+ Task::ready(Ok(vec![CompletionResponse {
+ completions: vec![completion],
+ display_options: Default::default(),
+ is_incomplete: false,
+ }]))
+ }
+
+ fn is_completion_trigger(
+ &self,
+ _buffer: &Entity,
+ _position: language::Anchor,
+ _text: &str,
+ _trigger_in_words: bool,
+ _cx: &mut Context,
+ ) -> bool {
+ false
+ }
+
+ fn filter_completions(&self) -> bool {
+ false
+ }
+}
+
#[derive(Default, Clone)]
pub struct FakeEditPredictionDelegate {
pub completion: Option,
diff --git a/crates/editor/src/editor.rs b/crates/editor/src/editor.rs
index ba7870efcb43a36442060654118bd228c0c0055c..3fb832b3ac7b2ebb6aa3f47b2dcd590d7eb081df 100644
--- a/crates/editor/src/editor.rs
+++ b/crates/editor/src/editor.rs
@@ -1869,6 +1869,7 @@ pub enum MultibufferSelectionMode {
pub struct RewrapOptions {
pub override_language_settings: bool,
pub preserve_existing_whitespace: bool,
+ pub line_length: Option,
}
impl Editor {
@@ -2885,6 +2886,11 @@ impl Editor {
if self.in_leading_whitespace {
key_context.add("in_leading_whitespace");
}
+ if self.edit_prediction_requires_modifier() {
+ key_context.set("edit_prediction_mode", "subtle")
+ } else {
+ key_context.set("edit_prediction_mode", "eager");
+ }
if self.selection_mark_mode {
key_context.add("selection_mode");
@@ -2952,7 +2958,7 @@ impl Editor {
window: &mut Window,
cx: &mut App,
) -> Option {
- let key_context = self.key_context_internal(self.has_active_edit_prediction(), window, cx);
+ let key_context = self.key_context_internal(true, window, cx);
let bindings =
match granularity {
@@ -2979,7 +2985,7 @@ impl Editor {
window: &mut Window,
cx: &mut App,
) -> Option