feat(google_drive): use oauth2 crate for PKCE support, make token storage generic over Serializable

.github/workflows/ci.yml

@@ -70,6 +70,28 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-cargo-build-
 
+      # Add disk space cleanup before linting
+      - name: Check disk space before build
+        run: df -h
+
+      #https://github.com/actions/runner-images/issues/2840
+      - name: Clean up disk space
+        run: |
+          echo "Cleaning up disk space..."
+          sudo rm -rf \
+            /opt/google/chrome \
+            /opt/microsoft/msedge \
+            /opt/microsoft/powershell \
+            /usr/lib/mono \
+            /usr/local/lib/android \
+            /usr/local/lib/node_modules \
+            /usr/local/share/chromium \
+            /usr/local/share/powershell \
+            /usr/share/dotnet \
+            /usr/share/swift
+
+          df -h
+
       - name: Build and Test
         run: |
           gnome-keyring-daemon --components=secrets --daemonize --unlock <<< 'foobar'
@@ -129,4 +151,4 @@ jobs:
     uses: ./.github/workflows/bundle-desktop.yml
     if: github.event_name == 'pull_request'
     with:
-      signing: false
+      signing: false

crates/goose-mcp/Cargo.toml

@@ -47,6 +47,7 @@ docx-rs = "0.4.7"
 image = "0.24.9"
 umya-spreadsheet = "2.2.3"
 keyring = { version = "3.6.1", features = ["apple-native", "windows-native", "sync-secret-service"] }
+oauth2 = { version = "5.0.0", features = ["reqwest"] }
 
 [dev-dependencies]
 serial_test = "3.0.0"

crates/goose-mcp/src/google_drive/mod.rs

@@ -1,13 +1,13 @@
-mod token_storage;
+mod oauth_pkce;
+pub mod storage;
 
 use indoc::indoc;
+use oauth_pkce::PkceOAuth2Client;
 use regex::Regex;
 use serde_json::{json, Value};
-use token_storage::{CredentialsManager, KeychainTokenStorage};
-
 use std::io::Cursor;
-use std::sync::Arc;
-use std::{env, fs, future::Future, path::Path, pin::Pin};
+use std::{env, fs, future::Future, path::Path, pin::Pin, sync::Arc};
+use storage::CredentialsManager;
 
 use mcp_core::content::Content;
 use mcp_core::{
@@ -26,47 +26,15 @@ use google_drive3::{
     api::{File, Scope},
     hyper_rustls::{self, HttpsConnector},
     hyper_util::{self, client::legacy::connect::HttpConnector},
-    yup_oauth2::{
-        self,
-        authenticator_delegate::{DefaultInstalledFlowDelegate, InstalledFlowDelegate},
-        InstalledFlowAuthenticator,
-    },
     DriveHub,
 };
 use google_sheets4::{self, Sheets};
 use http_body_util::BodyExt;
 
-/// async function to be pinned by the `present_user_url` method of the trait
-/// we use the existing `DefaultInstalledFlowDelegate::present_user_url` method as a fallback for
-/// when the browser did not open for example, the user still see's the URL.
-async fn browser_user_url(url: &str, need_code: bool) -> Result<String, String> {
-    tracing::info!(oauth_url = url, "Attempting OAuth login flow");
-    if let Err(e) = webbrowser::open(url) {
-        tracing::debug!(oauth_url = url, error = ?e, "Failed to open OAuth flow");
-        println!("Please open this URL in your browser:\n{}", url);
-    }
-    let def_delegate = DefaultInstalledFlowDelegate;
-    def_delegate.present_user_url(url, need_code).await
-}
-
-/// our custom delegate struct we will implement a flow delegate trait for:
-/// in this case we will implement the `InstalledFlowDelegated` trait
-#[derive(Copy, Clone)]
-struct LocalhostBrowserDelegate;
-
-/// here we implement only the present_user_url method with the added webbrowser opening
-/// the other behaviour of the trait does not need to be changed.
-impl InstalledFlowDelegate for LocalhostBrowserDelegate {
-    /// the actual presenting of URL and browser opening happens in the function defined above here
-    /// we only pin it
-    fn present_user_url<'a>(
-        &'a self,
-        url: &'a str,
-        need_code: bool,
-    ) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send + 'a>> {
-        Box::pin(browser_user_url(url, need_code))
-    }
-}
+// Constants for credential storage
+pub const KEYCHAIN_SERVICE: &str = "mcp_google_drive";
+pub const KEYCHAIN_USERNAME: &str = "oauth_credentials";
+pub const KEYCHAIN_DISK_FALLBACK_ENV: &str = "GOOGLE_DRIVE_DISK_FALLBACK";
 
 #[derive(Debug)]
 enum FileOperation {
@@ -141,38 +109,31 @@ impl GoogleDriveRouter {
             }
         }
 
-        // Create a credentials manager for storing tokens securely
-        let credentials_manager = Arc::new(CredentialsManager::new(credentials_path.clone()));
-
-        // Read the application secret from the OAuth keyfile
-        let secret = yup_oauth2::read_application_secret(keyfile_path)
-            .await
-            .expect("expected keyfile for google auth");
-
-        // Create custom token storage using our credentials manager
-        let token_storage = KeychainTokenStorage::new(
-            secret
-                .project_id
-                .clone()
-                .unwrap_or("unknown-project-id".to_string())
-                .to_string(),
-            credentials_manager.clone(),
-        );
-
-        // Create the authenticator with the installed flow
-        let auth = InstalledFlowAuthenticator::builder(
-            secret,
-            yup_oauth2::InstalledFlowReturnMethod::HTTPRedirect,
-        )
-        .with_storage(Box::new(token_storage)) // Use our custom storage
-        .flow_delegate(Box::new(LocalhostBrowserDelegate))
-        .build()
-        .await
-        .expect("expected successful authentication");
+        // Check if we should fall back to disk, must be explicitly enabled
+        let fallback_to_disk = match env::var(KEYCHAIN_DISK_FALLBACK_ENV) {
+            Ok(value) => value.to_lowercase() == "true",
+            Err(_) => false,
+        };
 
-        // Create the HTTP client
-        let client =
-            hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new())
+        // Create a credentials manager for storing tokens securely
+        let credentials_manager = Arc::new(CredentialsManager::new(
+            credentials_path.clone(),
+            fallback_to_disk,
+            KEYCHAIN_SERVICE.to_string(),
+            KEYCHAIN_USERNAME.to_string(),
+        ));
+
+        // Read the OAuth credentials from the keyfile
+        match fs::read_to_string(keyfile_path) {
+            Ok(_) => {
+                // Create the PKCE OAuth2 client
+                let auth = PkceOAuth2Client::new(keyfile_path, credentials_manager.clone())
+                    .expect("Failed to create OAuth2 client");
+
+                // Create the HTTP client
+                let client = hyper_util::client::legacy::Client::builder(
+                    hyper_util::rt::TokioExecutor::new(),
+                )
                 .build(
                     hyper_rustls::HttpsConnectorBuilder::new()
                         .with_native_roots()
@@ -182,11 +143,21 @@ impl GoogleDriveRouter {
                         .build(),
                 );
 
-        let drive_hub = DriveHub::new(client.clone(), auth.clone());
-        let sheets_hub = Sheets::new(client, auth);
+                let drive_hub = DriveHub::new(client.clone(), auth.clone());
+                let sheets_hub = Sheets::new(client, auth);
 
-        // Create and return the DriveHub
-        (drive_hub, sheets_hub, credentials_manager)
+                // Create and return the DriveHub, Sheets and our PKCE OAuth2 client
+                (drive_hub, sheets_hub, credentials_manager)
+            }
+            Err(e) => {
+                tracing::error!(
+                    "Failed to read OAuth config from {}: {}",
+                    keyfile_path.display(),
+                    e
+                );
+                panic!("Failed to read OAuth config: {}", e);
+            }
+        }
     }
 
     pub async fn new() -> Self {
@@ -715,7 +686,7 @@ impl GoogleDriveRouter {
                         .collect::<Vec<_>>()
                         .join("\n");
 
-                Ok(vec![Content::text(content.to_string())])
+                Ok(vec![Content::text(content.to_string()).with_priority(0.3)])
             }
         }
     }