Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions crates/chat-cli/src/cli/chat/tools/custom_tool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,15 @@ impl Default for TransportType {
}
}

#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct OAuthConfig {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Assuming we want this to more or less match what you linked w/ gemini - https://github.com/google-gemini/gemini-cli/blob/main/docs/tools/mcp-server.md#oauth-configuration-properties

There's already a top-level oauth_scopes unfortunately which could've been included in here, but I think it does make sense to add any and all future OAuth related config items here going forward

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If i have some time this week, I can refactor oauth_scopes to be inside this config to consolidate this going forward

/// 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<String>,
}

#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct CustomToolConfig {
Expand All @@ -59,6 +68,9 @@ pub struct CustomToolConfig {
/// Scopes with which oauth is done
#[serde(default = "get_default_scopes")]
pub oauth_scopes: Vec<String>,
/// OAuth configuration for this server
#[serde(skip_serializing_if = "Option::is_none")]
pub oauth: Option<OAuthConfig>,
/// The command string used to initialize the mcp server
#[serde(default)]
pub command: String,
Expand Down
3 changes: 2 additions & 1 deletion crates/chat-cli/src/mcp_client/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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?;

Expand Down
24 changes: 22 additions & 2 deletions crates/chat-cli/src/mcp_client/oauth_util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,7 @@ pub struct HttpServiceBuilder<'a> {
pub timeout: u64,
pub scopes: &'a [String],
pub headers: &'a HashMap<String, String>,
pub oauth_config: &'a Option<crate::cli::chat::tools::custom_tool::OAuthConfig>,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit - import OAuthConfig

pub messenger: &'a dyn Messenger,
}

Expand All @@ -207,6 +208,7 @@ impl<'a> HttpServiceBuilder<'a> {
timeout: u64,
scopes: &'a [String],
headers: &'a HashMap<String, String>,
oauth_config: &'a Option<crate::cli::chat::tools::custom_tool::OAuthConfig>,
messenger: &'a dyn Messenger,
) -> Self {
Self {
Expand All @@ -216,6 +218,7 @@ impl<'a> HttpServiceBuilder<'a> {
timeout,
scopes,
headers,
oauth_config,
messenger,
}
}
Expand All @@ -231,6 +234,7 @@ impl<'a> HttpServiceBuilder<'a> {
timeout,
scopes,
headers,
oauth_config,
messenger,
} = self;

Expand Down Expand Up @@ -292,7 +296,9 @@ impl<'a> HttpServiceBuilder<'a> {
cred_full_path.clone(),
reg_full_path.clone(),
scopes,
oauth_config,
messenger,
os,
)
.await?;

Expand Down Expand Up @@ -452,7 +458,9 @@ async fn get_auth_manager(
cred_full_path: PathBuf,
reg_full_path: PathBuf,
scopes: &[String],
oauth_config: &Option<crate::cli::chat::tools::custom_tool::OAuthConfig>,
messenger: &dyn Messenger,
os: &Os,
) -> Result<AuthorizationManager, OauthUtilError> {
let cred_as_bytes = tokio::fs::read(&cred_full_path).await;
let reg_as_bytes = tokio::fs::read(&reg_full_path).await;
Expand All @@ -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
Expand Down Expand Up @@ -509,9 +517,21 @@ async fn get_auth_manager(
async fn get_auth_manager_impl(
mut oauth_state: OAuthState,
scopes: &[String],
oauth_config: &Option<crate::cli::chat::tools::custom_tool::OAuthConfig>,
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::<u16>().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)>();

Expand Down
16 changes: 16 additions & 0 deletions schemas/agent-v1.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down