Skip to content

fix(rust): request MCP OAuth during session startup #2511

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
chuwik wants to merge 1 commit into
base: main
Choose a base branch
from
+18 −2
Draft

fix(rust): request MCP OAuth during session startup #2511

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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: 10 additions & 2 deletions rust/src/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2219,8 +2219,9 @@ pub struct SessionConfig {
/// Optional elicitation-request handler. When `None`,
/// `requestElicitation: false` goes on the wire.
pub elicitation_handler: Option<Arc<dyn ElicitationHandler>>,
/// Optional MCP OAuth request handler. When set, the SDK can satisfy MCP
/// server OAuth requests with host-acquired token data or cancellation.
/// Optional MCP OAuth request handler. When set, the SDK sends
/// `requestMcpOauth: true` during session creation so OAuth requests can
/// be observed from the start of runtime initialization.
pub mcp_auth_handler: Option<Arc<dyn McpAuthHandler>>,
/// Optional handler for the legacy question-and-answer `ask_user` variant.
/// When `None`, `requestUserInput: false` goes on the wire, so this client
Expand Down Expand Up @@ -2552,6 +2553,7 @@ impl SessionConfig {
let request_exit_plan_mode = self.exit_plan_mode_handler.is_some();
let request_auto_mode_switch = self.auto_mode_switch_handler.is_some();
let request_elicitation = self.elicitation_handler.is_some();
let request_mcp_oauth = self.mcp_auth_handler.is_some();
let hooks_flag = self.hooks_handler.is_some();

let mut tool_handlers: HashMap<String, Arc<dyn crate::tool::ToolHandler>> = HashMap::new();
Expand Down Expand Up @@ -2619,6 +2621,7 @@ impl SessionConfig {
request_exit_plan_mode,
request_auto_mode_switch,
request_elicitation,
request_mcp_oauth,
request_mcp_apps: self.enable_mcp_apps.unwrap_or(false),
github_mcp_tool_config: self.github_mcp_tool_config,
hooks: hooks_flag,
Expand Down Expand Up @@ -2698,6 +2701,7 @@ impl SessionConfig {
}

/// Install an [`McpAuthHandler`] for host-provided MCP OAuth tokens.
/// Enables create-time MCP OAuth event interest in the runtime.
pub fn with_mcp_auth_handler(mut self, handler: Arc<dyn McpAuthHandler>) -> Self {
self.mcp_auth_handler = Some(handler);
self
Expand Down Expand Up @@ -3573,6 +3577,7 @@ pub struct ResumeSessionConfig {
/// [`SessionConfig::elicitation_handler`].
pub elicitation_handler: Option<Arc<dyn ElicitationHandler>>,
/// Optional MCP OAuth handler. See [`SessionConfig::mcp_auth_handler`].
/// When set, the SDK sends `requestMcpOauth: true` during resume.
pub mcp_auth_handler: Option<Arc<dyn McpAuthHandler>>,
/// Optional user-input handler. See
/// [`SessionConfig::user_input_handler`].
Expand Down Expand Up @@ -3762,6 +3767,7 @@ impl ResumeSessionConfig {
let request_exit_plan_mode = self.exit_plan_mode_handler.is_some();
let request_auto_mode_switch = self.auto_mode_switch_handler.is_some();
let request_elicitation = self.elicitation_handler.is_some();
let request_mcp_oauth = self.mcp_auth_handler.is_some();
let hooks_flag = self.hooks_handler.is_some();

let mut tool_handlers: HashMap<String, Arc<dyn crate::tool::ToolHandler>> = HashMap::new();
Expand Down Expand Up @@ -3830,6 +3836,7 @@ impl ResumeSessionConfig {
request_exit_plan_mode,
request_auto_mode_switch,
request_elicitation,
request_mcp_oauth,
request_mcp_apps: self.enable_mcp_apps.unwrap_or(false),
github_mcp_tool_config: self.github_mcp_tool_config,
hooks: hooks_flag,
Expand Down Expand Up @@ -4002,6 +4009,7 @@ impl ResumeSessionConfig {
}

/// Install an [`McpAuthHandler`] for host-provided MCP OAuth tokens.
/// Enables resume-time MCP OAuth event interest in the runtime.
pub fn with_mcp_auth_handler(mut self, handler: Arc<dyn McpAuthHandler>) -> Self {
self.mcp_auth_handler = Some(handler);
self
Expand Down
2 changes: 2 additions & 0 deletions rust/src/wire.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,7 @@ pub(crate) struct SessionCreateWire {
pub request_exit_plan_mode: bool,
pub request_auto_mode_switch: bool,
pub request_elicitation: bool,
pub request_mcp_oauth: bool,
pub request_mcp_apps: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub github_mcp_tool_config: Option<GitHubMcpToolConfig>,
Expand Down Expand Up @@ -276,6 +277,7 @@ pub(crate) struct SessionResumeWire {
pub request_exit_plan_mode: bool,
pub request_auto_mode_switch: bool,
pub request_elicitation: bool,
pub request_mcp_oauth: bool,
pub request_mcp_apps: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub github_mcp_tool_config: Option<GitHubMcpToolConfig>,
Expand Down
6 changes: 6 additions & 0 deletions rust/tests/session_test.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -724,6 +724,7 @@ async fn create_session_registers_mcp_auth_interest_only_with_handler() {
let create_req = read_framed(&mut server_read).await;
assert_eq!(create_req["method"], "session.create");
assert_eq!(create_req["params"]["requestPermission"], true);
assert_eq!(create_req["params"]["requestMcpOauth"], false);
let session_id = requested_session_id(&create_req).to_string();
server_respond_create(&mut server_write, &create_req, &session_id).await;
let session = timeout(TIMEOUT, create_handle).await.unwrap().unwrap();
Expand All @@ -750,6 +751,7 @@ async fn create_session_registers_mcp_auth_interest_only_with_handler() {
let create_req = read_framed(&mut server_read).await;
assert_eq!(create_req["method"], "session.create");
assert_eq!(create_req["params"]["requestPermission"], true);
assert_eq!(create_req["params"]["requestMcpOauth"], true);
let session_id = requested_session_id(&create_req).to_string();
server_respond_create(&mut server_write, &create_req, &session_id).await;

Expand Down Expand Up @@ -798,6 +800,7 @@ async fn cloud_create_session_registers_mcp_auth_interest_after_create_only_with
assert_eq!(create_req["method"], "session.create");
assert!(create_req["params"].get("sessionId").is_none());
assert_eq!(create_req["params"]["requestPermission"], true);
assert_eq!(create_req["params"]["requestMcpOauth"], false);
server_respond_create(&mut server_write, &create_req, "server-assigned-session-1").await;
let session = timeout(TIMEOUT, create_handle).await.unwrap().unwrap();
let no_extra_request = timeout(Duration::from_millis(50), read_framed(&mut server_read)).await;
Expand All @@ -824,6 +827,7 @@ async fn cloud_create_session_registers_mcp_auth_interest_after_create_only_with
assert_eq!(create_req["method"], "session.create");
assert!(create_req["params"].get("sessionId").is_none());
assert_eq!(create_req["params"]["requestPermission"], true);
assert_eq!(create_req["params"]["requestMcpOauth"], true);
server_respond_create(&mut server_write, &create_req, "server-assigned-session-2").await;

let interest_req = read_framed(&mut server_read).await;
Expand Down Expand Up @@ -868,6 +872,7 @@ async fn resume_session_registers_mcp_auth_interest_only_with_handler() {
let resume_req = read_framed(&mut server_read).await;
assert_eq!(resume_req["method"], "session.resume");
assert_eq!(resume_req["params"]["requestPermission"], true);
assert_eq!(resume_req["params"]["requestMcpOauth"], false);
server_respond_create(&mut server_write, &resume_req, "session-without-auth").await;
respond_to_reload(&mut server_read, &mut server_write).await;
let session = timeout(TIMEOUT, resume_handle).await.unwrap().unwrap();
Expand All @@ -893,6 +898,7 @@ async fn resume_session_registers_mcp_auth_interest_only_with_handler() {
let resume_req = read_framed(&mut server_read).await;
assert_eq!(resume_req["method"], "session.resume");
assert_eq!(resume_req["params"]["requestPermission"], true);
assert_eq!(resume_req["params"]["requestMcpOauth"], true);
server_respond_create(&mut server_write, &resume_req, "session-with-auth").await;

let interest_req = read_framed(&mut server_read).await;
Expand Down
Loading