diff --git a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs index 0e45205678..2a26c1c12b 100644 --- a/crates/chat-cli/src/cli/chat/tools/custom_tool.rs +++ b/crates/chat-cli/src/cli/chat/tools/custom_tool.rs @@ -44,6 +44,15 @@ impl Default for TransportType { } } +#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, JsonSchema)] +#[serde(rename_all = "camelCase")] +pub struct OAuthConfig { + /// Custom redirect URI for OAuth flow (e.g., "127.0.0.1:7778") + /// If not specified, a random available port will be assigned by the OS + #[serde(skip_serializing_if = "Option::is_none")] + pub redirect_uri: Option, +} + #[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, JsonSchema)] #[serde(rename_all = "camelCase")] pub struct CustomToolConfig { @@ -59,6 +68,9 @@ pub struct CustomToolConfig { /// Scopes with which oauth is done #[serde(default = "get_default_scopes")] pub oauth_scopes: Vec, + /// OAuth configuration for this server + #[serde(skip_serializing_if = "Option::is_none")] + pub oauth: Option, /// The command string used to initialize the mcp server #[serde(default)] pub command: String, diff --git a/crates/chat-cli/src/mcp_client/client.rs b/crates/chat-cli/src/mcp_client/client.rs index fefaa1a9cb..6598fc71db 100644 --- a/crates/chat-cli/src/mcp_client/client.rs +++ b/crates/chat-cli/src/mcp_client/client.rs @@ -453,11 +453,12 @@ impl McpClientService { url, headers, oauth_scopes: scopes, + oauth, timeout, .. } = &self.config; - let http_service_builder = HttpServiceBuilder::new(url, os, url, *timeout, scopes, headers, messenger); + let http_service_builder = HttpServiceBuilder::new(url, os, url, *timeout, scopes, headers, oauth, messenger); let (service, auth_client_wrapper) = http_service_builder.try_build(&self).await?; diff --git a/crates/chat-cli/src/mcp_client/oauth_util.rs b/crates/chat-cli/src/mcp_client/oauth_util.rs index dfd1e9e0e7..5a2199defe 100644 --- a/crates/chat-cli/src/mcp_client/oauth_util.rs +++ b/crates/chat-cli/src/mcp_client/oauth_util.rs @@ -196,6 +196,7 @@ pub struct HttpServiceBuilder<'a> { pub timeout: u64, pub scopes: &'a [String], pub headers: &'a HashMap, + pub oauth_config: &'a Option, pub messenger: &'a dyn Messenger, } @@ -207,6 +208,7 @@ impl<'a> HttpServiceBuilder<'a> { timeout: u64, scopes: &'a [String], headers: &'a HashMap, + oauth_config: &'a Option, messenger: &'a dyn Messenger, ) -> Self { Self { @@ -216,6 +218,7 @@ impl<'a> HttpServiceBuilder<'a> { timeout, scopes, headers, + oauth_config, messenger, } } @@ -231,6 +234,7 @@ impl<'a> HttpServiceBuilder<'a> { timeout, scopes, headers, + oauth_config, messenger, } = self; @@ -292,7 +296,9 @@ impl<'a> HttpServiceBuilder<'a> { cred_full_path.clone(), reg_full_path.clone(), scopes, + oauth_config, messenger, + os, ) .await?; @@ -452,7 +458,9 @@ async fn get_auth_manager( cred_full_path: PathBuf, reg_full_path: PathBuf, scopes: &[String], + oauth_config: &Option, messenger: &dyn Messenger, + os: &Os, ) -> Result { let cred_as_bytes = tokio::fs::read(&cred_full_path).await; let reg_as_bytes = tokio::fs::read(®_full_path).await; @@ -474,7 +482,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, scopes, messenger).await?; + let (am, redirect_uri) = get_auth_manager_impl(oauth_state, scopes, oauth_config, messenger, os).await?; // Client registration is done in [start_authorization] // If we have gotten past that point that means we have the info to persist the @@ -509,9 +517,21 @@ async fn get_auth_manager( async fn get_auth_manager_impl( mut oauth_state: OAuthState, scopes: &[String], + oauth_config: &Option, messenger: &dyn Messenger, + _os: &Os, ) -> Result<(AuthorizationManager, String), OauthUtilError> { - let socket_addr = SocketAddr::from(([127, 0, 0, 1], 0)); + // Get port from per-server oauth config, or use 0 for random port assignment + let port = oauth_config + .as_ref() + .and_then(|cfg| cfg.redirect_uri.as_ref()) + .and_then(|uri| { + // Parse port from redirect_uri like "127.0.0.1:7778" or ":7778" + uri.split(':').last().and_then(|p| p.parse::().ok()) + }) + .unwrap_or(0); // Port 0 = OS assigns random available port + + let socket_addr = SocketAddr::from(([127, 0, 0, 1], port)); let cancellation_token = tokio_util::sync::CancellationToken::new(); let (tx, rx) = tokio::sync::oneshot::channel::<(String, String)>(); diff --git a/schemas/agent-v1.json b/schemas/agent-v1.json index 5e72b08476..fb4c5e8964 100644 --- a/schemas/agent-v1.json +++ b/schemas/agent-v1.json @@ -94,6 +94,22 @@ "offline_access" ] }, + "oauth": { + "description": "OAuth configuration for this server", + "type": [ + "object", + "null" + ], + "properties": { + "redirectUri": { + "description": "Custom redirect URI for OAuth flow (e.g., \"127.0.0.1:7778\"). If not specified, a random available port will be assigned by the OS", + "type": [ + "string", + "null" + ] + } + } + }, "command": { "description": "The command string used to initialize the mcp server", "type": "string",