Use RMCP for StreamableHTTP OAuth support (#3845)

crates/goose/src/agents/extension_manager.rs

@@ -27,9 +27,11 @@ use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, Extension
 use super::tool_execution::ToolCallResult;
 use crate::agents::extension::{Envs, ProcessExit};
 use crate::config::{Config, ExtensionConfigManager};
+use crate::oauth::oauth_flow;
 use crate::prompt_template;
 use mcp_client::client::{McpClient, McpClientTrait};
 use rmcp::model::{Content, GetPromptResult, Prompt, ResourceContents, Tool};
+use rmcp::transport::auth::AuthClient;
 use serde_json::Value;
 
 type McpClientBox = Arc<Mutex<Box<dyn McpClientTrait>>>;
@@ -205,6 +207,7 @@ impl ExtensionManager {
                 uri,
                 timeout,
                 headers,
+                name,
                 ..
             } => {
                 let mut default_headers = HeaderMap::new();
@@ -231,13 +234,38 @@ impl ExtensionManager {
                         ..Default::default()
                     },
                 );
-                let client = McpClient::connect(
+                let client_res = McpClient::connect(
                     transport,
                     Duration::from_secs(
                         timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
                     ),
                 )
-                .await?;
+                .await;
+                let client = if let Err(e) = client_res {
+                    // make an attempt at oauth, but failing that, return the original error,
+                    // because this might not have been an auth error at all
+                    let am = match oauth_flow(uri, name).await {
+                        Ok(am) => am,
+                        Err(_) => return Err(e.into()),
+                    };
+                    let client = AuthClient::new(reqwest::Client::default(), am);
+                    let transport = StreamableHttpClientTransport::with_client(
+                        client,
+                        StreamableHttpClientTransportConfig {
+                            uri: uri.clone().into(),
+                            ..Default::default()
+                        },
+                    );
+                    McpClient::connect(
+                        transport,
+                        Duration::from_secs(
+                            timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
+                        ),
+                    )
+                    .await?
+                } else {
+                    client_res?
+                };
                 Box::new(client)
             }
             ExtensionConfig::Stdio {
@@ -463,6 +491,7 @@ impl ExtensionManager {
                             description: tool.description,
                             input_schema: tool.input_schema,
                             annotations: tool.annotations,
+                            output_schema: tool.output_schema,
                         });
                     }
 
@@ -719,7 +748,7 @@ impl ExtensionManager {
             client_guard
                 .call_tool(&tool_name, arguments, cancellation_token)
                 .await
-                .map(|call| call.content)
+                .map(|call| call.content.unwrap_or_default())
                 .map_err(|e| ToolError::ExecutionError(e.to_string()))
         };
 
@@ -947,8 +976,9 @@ mod tests {
         ) -> Result<CallToolResult, Error> {
             match name {
                 "tool" | "test__tool" => Ok(CallToolResult {
-                    content: vec![],
+                    content: Some(vec![]),
                     is_error: None,
+                    structured_content: None,
                 }),
                 _ => Err(Error::TransportClosed),
             }

crates/goose/src/lib.rs

@@ -4,6 +4,7 @@ pub mod context_mgmt;
 mod conversation_fixer;
 pub mod message;
 pub mod model;
+pub mod oauth;
 pub mod permission;
 pub mod project;
 pub mod prompt_template;

crates/goose/src/oauth.rs

@@ -0,0 +1,81 @@
+use axum::extract::{Query, State};
+use axum::response::Html;
+use axum::routing::get;
+use axum::Router;
+use minijinja::render;
+use rmcp::transport::auth::OAuthState;
+use rmcp::transport::AuthorizationManager;
+use serde::Deserialize;
+use std::net::SocketAddr;
+use std::sync::Arc;
+use tokio::sync::{oneshot, Mutex};
+
+const CALLBACK_TEMPLATE: &str = include_str!("oauth_callback.html");
+
+#[derive(Clone)]
+struct AppState {
+    code_receiver: Arc<Mutex<Option<oneshot::Sender<String>>>>,
+}
+
+#[derive(Debug, Deserialize)]
+struct CallbackParams {
+    code: String,
+    #[allow(dead_code)]
+    state: Option<String>,
+}
+
+pub async fn oauth_flow(
+    mcp_server_url: &String,
+    name: &String,
+) -> Result<AuthorizationManager, anyhow::Error> {
+    let (code_sender, code_receiver) = oneshot::channel::<String>();
+    let app_state = AppState {
+        code_receiver: Arc::new(Mutex::new(Some(code_sender))),
+    };
+
+    let rendered = render!(CALLBACK_TEMPLATE, name => name);
+    let handler = move |Query(params): Query<CallbackParams>, State(state): State<AppState>| {
+        let rendered = rendered.clone();
+        async move {
+            if let Some(sender) = state.code_receiver.lock().await.take() {
+                let _ = sender.send(params.code);
+            }
+            Html(rendered)
+        }
+    };
+    let app = Router::new()
+        .route("/oauth_callback", get(handler))
+        .with_state(app_state);
+
+    let addr = SocketAddr::from(([127, 0, 0, 1], 0));
+    let listener = tokio::net::TcpListener::bind(addr).await?;
+    let used_addr = listener.local_addr()?;
+    tokio::spawn(async move {
+        let result = axum::serve(listener, app).await;
+
+        if let Err(e) = result {
+            eprintln!("Callback server error: {}", e);
+        }
+    });
+
+    let mut oauth_state = OAuthState::new(mcp_server_url, None).await?;
+    let redirect_uri = format!("http://localhost:{}/oauth_callback", used_addr.port());
+    oauth_state
+        .start_authorization(&[], redirect_uri.as_str())
+        .await?;
+
+    let authorization_url = oauth_state.get_authorization_url().await?;
+    if webbrowser::open(authorization_url.as_str()).is_err() {
+        eprintln!("Open the following URL to authorize {}:", name);
+        eprintln!("  {}", authorization_url);
+    }
+
+    let auth_code = code_receiver.await?;
+    oauth_state.handle_callback(&auth_code).await?;
+
+    let am = oauth_state
+        .into_authorization_manager()
+        .ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
+
+    Ok(am)
+}

crates/goose/src/oauth_callback.html

@@ -0,0 +1,73 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>{{ name }} OAuth Success</title>
+    <style>
+      body {
+        font-family: "Cash Sans", -apple-system, BlinkMacSystemFont, "Segoe UI",
+          Roboto, sans-serif;
+        display: flex;
+        flex-direction: column;
+        justify-content: center;
+        align-items: center;
+        min-height: 100vh;
+        margin: 0;
+        background-color: #f4f6f7;
+        color: #3f434b;
+      }
+
+      .container {
+        text-align: center;
+        padding: 2rem;
+        background: #ffffff;
+        border-radius: 8px;
+        box-shadow: 0px 12px 32px 0px rgba(0, 0, 0, 0.04),
+          0px 8px 16px 0px rgba(0, 0, 0, 0.02),
+          0px 2px 4px 0px rgba(0, 0, 0, 0.04),
+          0px 0px 1px 0px rgba(0, 0, 0, 0.2);
+        max-width: 400px;
+      }
+
+      h1 {
+        color: #32353b;
+        margin-bottom: 1rem;
+        font-weight: 500;
+      }
+
+      .client-name {
+        font-weight: 700;
+        color: #22252a;
+      }
+
+      button {
+        background-color: #32353b;
+        color: #ffffff;
+        border: none;
+        padding: 0.75rem 1.5rem;
+        border-radius: 4px;
+        cursor: pointer;
+        font-size: 1rem;
+        font-family: "Cash Sans", sans-serif;
+        font-weight: 500;
+        margin-top: 1rem;
+        transition: background-color 0.2s ease;
+      }
+
+      button:hover {
+        background-color: #22252a;
+      }
+    </style>
+  </head>
+  <body>
+    <div class="container">
+      <h1>Authorization Success</h1>
+      <p>
+        You have successfully authorized
+        <span class="client-name">{{ name }}</span>. You can now close this
+        window and return to Goose.
+      </p>
+    </div>
+  </body>
+</html>