fix(mcp): oauth issues (#2925)

* fix incorrect scope for mcp oauth

* reverts custom tool config enum change

* fixes display task overriding sign in notice

* updates schema

Author: dingfeli

crates/chat-cli/src/cli/chat/prompt.rs

@@ -745,10 +745,7 @@ mod tests {
             let key_event = KeyEvent(KeyCode::Char(key), Modifiers::CTRL);
 
             // Try to bind and get the previous handler
-            let previous_handler = test_editor.bind_sequence(
-                key_event,
-                EventHandler::Simple(Cmd::Noop)
-            );
+            let previous_handler = test_editor.bind_sequence(key_event, EventHandler::Simple(Cmd::Noop));
 
             // If there was a previous handler, it means the key was already bound
             // (which could be our custom binding overriding Emacs)

crates/chat-cli/src/cli/chat/tool_manager.rs

@@ -1210,6 +1210,7 @@ fn spawn_display_task(
                                     terminal::Clear(terminal::ClearType::CurrentLine),
                                 )?;
                                 queue_oauth_message(&name, &mut output)?;
+                                queue_init_message(spinner_logo_idx, complete, failed, total, &mut output)?;
                             },
                         },
                         Err(_e) => {

crates/chat-cli/src/cli/chat/tools/custom_tool.rs

@@ -22,7 +22,10 @@ use crate::cli::agent::{
 };
 use crate::cli::chat::CONTINUATION_LINE;
 use crate::cli::chat::token_counter::TokenCounter;
-use crate::mcp_client::RunningService;
+use crate::mcp_client::{
+    RunningService,
+    oauth_util,
+};
 use crate::os::Os;
 use crate::util::MCP_SERVER_TOOL_DELIMITER;
 use crate::util::pattern_matching::matches_any_pattern;
@@ -43,17 +46,20 @@ impl Default for TransportType {
 }
 
 #[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, JsonSchema)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
 pub struct CustomToolConfig {
-    /// The type of transport the mcp server is expecting. For http transport, only url (for now)
-    /// is taken into account.
+    /// The transport type to use for communication with the MCP server
     #[serde(default)]
     pub r#type: TransportType,
-    /// The URL endpoint for HTTP-based MCP servers
+    /// The URL for HTTP-based MCP server communication
     #[serde(default)]
     pub url: String,
     /// HTTP headers to include when communicating with HTTP-based MCP servers
     #[serde(default)]
     pub headers: HashMap<String, String>,
+    /// Scopes with which oauth is done
+    #[serde(default = "get_default_scopes")]
+    pub oauth_scopes: Vec<String>,
     /// The command string used to initialize the mcp server
     #[serde(default)]
     pub command: String,
@@ -74,6 +80,13 @@ pub struct CustomToolConfig {
     pub is_from_legacy_mcp_json: bool,
 }
 
+pub fn get_default_scopes() -> Vec<String> {
+    oauth_util::get_default_scopes()
+        .iter()
+        .map(|s| (*s).to_string())
+        .collect::<Vec<_>>()
+}
+
 pub fn default_timeout() -> u64 {
     120 * 1000
 }

crates/chat-cli/src/cli/mcp.rs

@@ -406,7 +406,7 @@ impl StatusArgs {
                         style::Print(format!("Disabled: {}\n", cfg.disabled)),
                         style::Print(format!(
                             "Env Vars: {}\n",
-                            cfg.env.as_ref().map_or_else(
+                            cfg.env.map_or_else(
                                 || "(none)".into(),
                                 |e| e
                                     .iter()

crates/chat-cli/src/mcp_client/client.rs

@@ -151,6 +151,8 @@ pub enum McpClientError {
     Parse(#[from] url::ParseError),
     #[error(transparent)]
     Auth(#[from] crate::auth::AuthError),
+    #[error("{0}")]
+    MalformedConfig(&'static str),
 }
 
 /// Decorates the method passed in with retry logic, but only if the [RunningService] has an
@@ -336,20 +338,23 @@ impl McpClientService {
                             HttpTransport::WithAuth((transport, mut auth_client)) => {
                                 // The crate does not automatically refresh tokens when they expire. We
                                 // would need to handle that here
-                                let url = self.config.url.clone();
+                                let url = &backup_config.url;
                                 let service = match self.into_dyn().serve(transport).await.map_err(Box::new) {
                                     Ok(service) => service,
                                     Err(e) if matches!(*e, ClientInitializeError::ConnectionClosed(_)) => {
                                         debug!("## mcp: first hand shake attempt failed: {:?}", e);
                                         let refresh_res = auth_client.refresh_token().await;
                                         let new_self = McpClientService::new(
                                             server_name.clone(),
-                                            backup_config,
+                                            backup_config.clone(),
                                             messenger_clone.clone(),
                                         );
 
+                                        let scopes = &backup_config.oauth_scopes;
+                                        let timeout = backup_config.timeout;
+                                        let headers = &backup_config.headers;
                                         let new_transport =
-                                            get_http_transport(&os_clone, &url, Some(auth_client.auth_client.clone()), &*messenger_dup).await?;
+                                            get_http_transport(&os_clone, url, timeout, scopes, headers,Some(auth_client.auth_client.clone()), &*messenger_dup).await?;
 
                                         match new_transport {
                                             HttpTransport::WithAuth((new_transport, new_auth_client)) => {
@@ -367,8 +372,8 @@ impl McpClientService {
                                                         // again. We do this by deleting the cred
                                                         // and discarding the client to trigger a full auth flow
                                                         tokio::fs::remove_file(&auth_client.cred_full_path).await?;
-                                                        let new_transport  =
-                                                            get_http_transport(&os_clone, &url, None, &*messenger_dup).await?;
+                                                        let new_transport =
+                                                            get_http_transport(&os_clone, url, timeout, scopes,headers,None, &*messenger_dup).await?;
 
                                                         match new_transport {
                                                             HttpTransport::WithAuth((new_transport, new_auth_client)) => {
@@ -495,17 +500,32 @@ impl McpClientService {
     }
 
     async fn get_transport(&mut self, os: &Os, messenger: &dyn Messenger) -> Result<Transport, McpClientError> {
-        // TODO: figure out what to do with headers
         let CustomToolConfig {
-            r#type: transport_type,
+            r#type,
             url,
+            headers,
+            oauth_scopes: scopes,
             command: command_as_str,
             args,
             env: config_envs,
+            timeout,
             ..
         } = &mut self.config;
 
-        match transport_type {
+        let is_malformed_http = matches!(r#type, TransportType::Http) && url.is_empty();
+        let is_malformed_stdio = matches!(r#type, TransportType::Stdio) && command_as_str.is_empty();
+
+        if is_malformed_http {
+            return Err(McpClientError::MalformedConfig(
+                "MCP config is malformed: transport type is specified to be http but url is empty",
+            ));
+        } else if is_malformed_stdio {
+            return Err(McpClientError::MalformedConfig(
+                "MCP config is malformed: transport type is specified to be stdio but command is empty",
+            ));
+        }
+
+        match r#type {
             TransportType::Stdio => {
                 let expanded_cmd = canonicalizes_path(os, command_as_str)?;
                 let command = Command::new(expanded_cmd).configure(|cmd| {
@@ -525,7 +545,7 @@ impl McpClientService {
                 Ok(Transport::Stdio((tokio_child_process, child_stderr)))
             },
             TransportType::Http => {
-                let http_transport = get_http_transport(os, url, None, messenger).await?;
+                let http_transport = get_http_transport(os, url, *timeout, scopes, headers, None, messenger).await?;
 
                 Ok(Transport::Http(http_transport))
             },
@@ -562,7 +582,6 @@ impl McpClientService {
 
     async fn on_tool_list_changed(&self, context: NotificationContext<RoleClient>) {
         let NotificationContext { peer, .. } = context;
-        let _timeout = self.config.timeout;
 
         paginated_fetch! {
             final_result_type: ListToolsResult,
@@ -578,7 +597,6 @@ impl McpClientService {
 
     async fn on_prompt_list_changed(&self, context: NotificationContext<RoleClient>) {
         let NotificationContext { peer, .. } = context;
-        let _timeout = self.config.timeout;
 
         paginated_fetch! {
             final_result_type: ListPromptsResult,

crates/chat-cli/src/mcp_client/oauth_util.rs

@@ -1,10 +1,14 @@
+use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::str::FromStr;
 use std::sync::Arc;
 
-use http::StatusCode;
+use http::{
+    HeaderMap,
+    StatusCode,
+};
 use http_body_util::Full;
 use hyper::Response;
 use hyper::body::Bytes;
@@ -69,6 +73,8 @@ pub enum OauthUtilError {
     Directory(#[from] DirectoryError),
     #[error(transparent)]
     Reqwest(#[from] reqwest::Error),
+    #[error("{0}")]
+    Http(String),
     #[error("Malformed directory")]
     MalformDirectory,
     #[error("Missing credential")]
@@ -162,13 +168,16 @@ pub enum HttpTransport {
     WithoutAuth(WorkerTransport<StreamableHttpClientWorker<Client>>),
 }
 
-fn get_scopes() -> &'static [&'static str] {
-    &["openid", "mcp", "email", "profile"]
+pub fn get_default_scopes() -> &'static [&'static str] {
+    &["openid", "email", "profile", "offline_access"]
 }
 
 pub async fn get_http_transport(
     os: &Os,
     url: &str,
+    timeout: u64,
+    scopes: &[String],
+    headers: &HashMap<String, String>,
     auth_client: Option<AuthClient<Client>>,
     messenger: &dyn Messenger,
 ) -> Result<HttpTransport, OauthUtilError> {
@@ -178,16 +187,28 @@ pub async fn get_http_transport(
     let cred_full_path = cred_dir.join(format!("{key}.token.json"));
     let reg_full_path = cred_dir.join(format!("{key}.registration.json"));
 
-    let reqwest_client = reqwest::Client::default();
+    let mut client_builder = reqwest::ClientBuilder::new().timeout(std::time::Duration::from_millis(timeout));
+    if !headers.is_empty() {
+        let headers = HeaderMap::try_from(headers).map_err(|e| OauthUtilError::Http(e.to_string()))?;
+        client_builder = client_builder.default_headers(headers);
+    };
+    let reqwest_client = client_builder.build()?;
+
     let probe_resp = reqwest_client.get(url.clone()).send().await?;
     match probe_resp.status() {
         StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
             debug!("## mcp: requires auth, auth client passed in is {:?}", auth_client);
             let auth_client = match auth_client {
                 Some(auth_client) => auth_client,
                 None => {
-                    let am =
-                        get_auth_manager(url.clone(), cred_full_path.clone(), reg_full_path.clone(), messenger).await?;
+                    let am = get_auth_manager(
+                        url.clone(),
+                        cred_full_path.clone(),
+                        reg_full_path.clone(),
+                        scopes,
+                        messenger,
+                    )
+                    .await?;
                     AuthClient::new(reqwest_client, am)
                 },
             };
@@ -204,7 +225,12 @@ pub async fn get_http_transport(
             Ok(HttpTransport::WithAuth((transport, auth_dg)))
         },
         _ => {
-            let transport = StreamableHttpClientTransport::from_uri(url.as_str());
+            let transport =
+                StreamableHttpClientTransport::with_client(reqwest_client, StreamableHttpClientTransportConfig {
+                    uri: url.as_str().into(),
+                    allow_stateless: false,
+                    ..Default::default()
+                });
 
             Ok(HttpTransport::WithoutAuth(transport))
         },
@@ -215,6 +241,7 @@ async fn get_auth_manager(
     url: Url,
     cred_full_path: PathBuf,
     reg_full_path: PathBuf,
+    scopes: &[String],
     messenger: &dyn Messenger,
 ) -> Result<AuthorizationManager, OauthUtilError> {
     let cred_as_bytes = tokio::fs::read(&cred_full_path).await;
@@ -237,7 +264,7 @@ async fn get_auth_manager(
         _ => {
             info!("Error reading cached credentials");
             debug!("## mcp: cache read failed. constructing auth manager from scratch");
-            let (am, redirect_uri) = get_auth_manager_impl(oauth_state, messenger).await?;
+            let (am, redirect_uri) = get_auth_manager_impl(oauth_state, scopes, messenger).await?;
 
             // Client registration is done in [start_authorization]
             // If we have gotten past that point that means we have the info to persist the
@@ -246,7 +273,10 @@ async fn get_auth_manager(
             let reg = Registration {
                 client_id,
                 client_secret: None,
-                scopes: get_scopes().iter().map(|s| (*s).to_string()).collect::<Vec<_>>(),
+                scopes: get_default_scopes()
+                    .iter()
+                    .map(|s| (*s).to_string())
+                    .collect::<Vec<_>>(),
                 redirect_uri,
             };
             let reg_as_str = serde_json::to_string_pretty(&reg)?;
@@ -268,6 +298,7 @@ async fn get_auth_manager(
 
 async fn get_auth_manager_impl(
     mut oauth_state: OAuthState,
+    scopes: &[String],
     messenger: &dyn Messenger,
 ) -> Result<(AuthorizationManager, String), OauthUtilError> {
     let socket_addr = SocketAddr::from(([127, 0, 0, 1], 0));
@@ -278,7 +309,9 @@ async fn get_auth_manager_impl(
     info!("Listening on local host port {:?} for oauth", actual_addr);
 
     let redirect_uri = format!("http://{}", actual_addr);
-    oauth_state.start_authorization(get_scopes(), &redirect_uri).await?;
+    let scopes_as_str = scopes.iter().map(String::as_str).collect::<Vec<_>>();
+    let scopes_as_slice = scopes_as_str.as_slice();
+    oauth_state.start_authorization(scopes_as_slice, &redirect_uri).await?;
 
     let auth_url = oauth_state.get_authorization_url().await?;
     _ = messenger.send_oauth_link(auth_url).await;
@@ -333,9 +366,19 @@ async fn make_svc(
             let query = uri.query().unwrap_or("");
             let params: std::collections::HashMap<String, String> =
                 url::form_urlencoded::parse(query.as_bytes()).into_owned().collect();
+            debug!("## mcp: uri: {}, query: {}, params: {:?}", uri, query, params);
 
             let self_clone = self.clone();
             Box::pin(async move {
+                let error = params.get("error");
+                let resp = if let Some(err) = error {
+                    mk_response(format!(
+                        "Oauth failed. Check url for precise reasons. Possible reasons: {err}.\nIf this is scope related. You can try configuring the server scopes to be an empty array via adding oauth_scopes: []"
+                    ))
+                } else {
+                    mk_response("You can close this page now".to_string())
+                };
+
                 let code = params.get("code").cloned().unwrap_or_default();
                 if let Some(sender) = self_clone
                     .one_shot_sender
@@ -345,7 +388,8 @@ async fn make_svc(
                 {
                     sender.send(code).map_err(LoopBackError::Send)?;
                 }
-                mk_response("You can close this page now".to_string())
+
+                resp
             })
         }
     }

schemas/agent-v1.json

@@ -73,6 +73,27 @@
             "type": "string",
             "default": ""
           },
+          "headers": {
+            "description": "HTTP headers to include when communicating with HTTP-based MCP servers",
+            "type": "object",
+            "additionalProperties": {
+              "type": "string"
+            },
+            "default": {}
+          },
+          "oauthScopes": {
+            "description": "Scopes with which oauth is done",
+            "type": "array",
+            "items": {
+              "type": "string"
+            },
+            "default": [
+              "openid",
+              "email",
+              "profile",
+              "offline_access"
+            ]
+          },
           "command": {
             "description": "The command string used to initialize the mcp server",
             "type": "string",