feat: google sheets support (in google drive builtin MCP server) (#1601)

Cargo.lock

@@ -2265,6 +2265,26 @@ dependencies = [
  "yup-oauth2",
 ]
 
+[[package]]
+name = "google-sheets4"
+version = "6.0.0+20240621"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4f8ccfc6418e81d1e2ed66fad49d0487526281505b8a0ed8ee770dc7d6bb1e5"
+dependencies = [
+ "chrono",
+ "google-apis-common",
+ "hyper 1.6.0",
+ "hyper-rustls 0.27.5",
+ "hyper-util",
+ "mime",
+ "serde",
+ "serde_json",
+ "serde_with",
+ "tokio",
+ "url",
+ "yup-oauth2",
+]
+
 [[package]]
 name = "goose"
 version = "1.0.13"
@@ -2388,6 +2408,7 @@ dependencies = [
  "docx-rs",
  "etcetera",
  "google-drive3",
+ "google-sheets4",
  "http-body-util",
  "ignore",
  "image 0.24.9",

crates/goose-mcp/Cargo.toml

@@ -33,6 +33,7 @@ etcetera = "0.8.0"
 tempfile = "3.8"
 include_dir = "0.7.4"
 google-drive3 = "6.0.0"
+google-sheets4 = "6.0.0"
 webbrowser = "0.8"
 http-body-util = "0.1.2"
 regex = "1.11.1"

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

@@ -29,6 +29,8 @@ use google_drive3::{
     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
@@ -67,10 +69,14 @@ pub struct GoogleDriveRouter {
     tools: Vec<Tool>,
     instructions: String,
     drive: DriveHub<HttpsConnector<HttpConnector>>,
+    sheets: Sheets<HttpsConnector<HttpConnector>>,
 }
 
 impl GoogleDriveRouter {
-    async fn google_auth() -> DriveHub<HttpsConnector<HttpConnector>> {
+    async fn google_auth() -> (
+        DriveHub<HttpsConnector<HttpConnector>>,
+        Sheets<HttpsConnector<HttpConnector>>,
+    ) {
         let oauth_config = env::var("GOOGLE_DRIVE_OAUTH_CONFIG");
         let keyfile_path_str = env::var("GOOGLE_DRIVE_OAUTH_PATH")
             .unwrap_or_else(|_| "./gcp-oauth.keys.json".to_string());
@@ -129,11 +135,14 @@ impl GoogleDriveRouter {
                         .build(),
                 );
 
-        DriveHub::new(client, auth)
+        let drive_hub = DriveHub::new(client.clone(), auth.clone());
+        let sheets_hub = Sheets::new(client, auth);
+
+        (drive_hub, sheets_hub)
     }
 
     pub async fn new() -> Self {
-        let drive = Self::google_auth().await;
+        let (drive, sheets) = Self::google_auth().await;
 
         // handle auth
         let search_tool = Tool::new(
@@ -185,13 +194,49 @@ impl GoogleDriveRouter {
             }),
         );
 
+        let sheets_tool = Tool::new(
+            "sheets_tool".to_string(),
+            indoc! {r#"
+                Work with Google Sheets data using various operations.
+                Supports operations:
+                - list_sheets: List all sheets in a spreadsheet
+                - get_columns: Get column headers from a specific sheet
+                - get_values: Get values from a range
+            "#}
+            .to_string(),
+            json!({
+              "type": "object",
+              "properties": {
+                  "spreadsheetId": {
+                      "type": "string",
+                      "description": "The ID of the spreadsheet to work with",
+                  },
+                  "operation": {
+                      "type": "string",
+                      "enum": ["list_sheets", "get_columns", "get_values"],
+                      "description": "The operation to perform on the spreadsheet",
+                  },
+                  "sheetName": {
+                      "type": "string",
+                      "description": "The name of the sheet to work with (optional for some operations)",
+                  },
+                  "range": {
+                      "type": "string",
+                      "description": "The A1 notation of the range to retrieve values (e.g., 'Sheet1!A1:D10')",
+                  }
+              },
+              "required": ["spreadsheetId", "operation"],
+            }),
+        );
+
         let instructions = indoc::formatdoc! {r#"
             Google Drive MCP Server Instructions
 
             ## Overview
-            The Google Drive MCP server provides two main tools for interacting with Google Drive files:
+            The Google Drive MCP server provides tools for interacting with Google Drive files and Google Sheets:
             1. search - Find files in your Google Drive
             2. read - Read file contents directly using a uri in the `gdrive:///uri` format
+            3. sheets_tool - Work with Google Sheets data using various operations
 
             ## Available Tools
 
@@ -210,6 +255,18 @@ impl GoogleDriveRouter {
             Limitations: Google Sheets exporting only supports reading the first sheet. This is an important limitation that should
             be communicated to the user whenever dealing with a Google Sheet (mimeType: application/vnd.google-apps.spreadsheet).
 
+            ### 3. Sheets Tool
+            Work with Google Sheets data using various operations:
+            - list_sheets: List all sheets in a spreadsheet
+            - get_columns: Get column headers from a specific sheet
+            - get_values: Get values from a range
+
+            Parameters:
+            - spreadsheetId: The ID of the spreadsheet (can be obtained from search results)
+            - operation: The operation to perform (one of the operations listed above)
+            - sheetName: The name of the sheet to work with (optional for some operations)
+            - range: The A1 notation of the range to retrieve values (e.g., 'Sheet1!A1:D10')
+
             ## File Format Handling
             The server automatically handles different file types:
             - Google Docs → Markdown
@@ -222,6 +279,7 @@ impl GoogleDriveRouter {
 
             1. First, search for the file you want to read, searching by name.
             2. Then, use the file URI from the search results to read its contents.
+            3. For Google Sheets, use the sheets_tool with the appropriate operation.
 
             ## Best Practices
             1. Always use search first to find the correct file URI
@@ -240,9 +298,10 @@ impl GoogleDriveRouter {
         "#};
 
         Self {
-            tools: vec![search_tool, read_tool],
+            tools: vec![search_tool, read_tool, sheets_tool],
             instructions,
             drive,
+            sheets,
         }
     }
 
@@ -506,6 +565,152 @@ impl GoogleDriveRouter {
         }
     }
 
+    // Implement sheets_tool functionality
+    async fn sheets_tool(&self, params: Value) -> Result<Vec<Content>, ToolError> {
+        let spreadsheet_id = params.get("spreadsheetId").and_then(|q| q.as_str()).ok_or(
+            ToolError::InvalidParameters("The spreadsheetId is required".to_string()),
+        )?;
+
+        let operation = params.get("operation").and_then(|q| q.as_str()).ok_or(
+            ToolError::InvalidParameters("The operation is required".to_string()),
+        )?;
+
+        match operation {
+            "list_sheets" => {
+                // Get spreadsheet metadata to list all sheets
+                let result = self
+                    .sheets
+                    .spreadsheets()
+                    .get(spreadsheet_id)
+                    .clear_scopes()
+                    .add_scope(Scope::Readonly)
+                    .doit()
+                    .await;
+
+                match result {
+                    Err(e) => Err(ToolError::ExecutionError(format!(
+                        "Failed to execute Google Sheets get query, {}.",
+                        e
+                    ))),
+                    Ok(r) => {
+                        let spreadsheet = r.1;
+                        let sheets = spreadsheet.sheets.unwrap_or_default();
+                        let sheets_info = sheets
+                            .into_iter()
+                            .filter_map(|sheet| {
+                                let properties = sheet.properties?;
+                                let title = properties.title?;
+                                let sheet_id = properties.sheet_id?;
+                                let grid_properties = properties.grid_properties?;
+                                Some(format!(
+                                    "Sheet: {} (ID: {}, Rows: {}, Columns: {})",
+                                    title,
+                                    sheet_id,
+                                    grid_properties.row_count.unwrap_or(0),
+                                    grid_properties.column_count.unwrap_or(0)
+                                ))
+                            })
+                            .collect::<Vec<String>>()
+                            .join("\n");
+
+                        Ok(vec![Content::text(sheets_info).with_priority(0.1)])
+                    }
+                }
+            },
+            "get_columns" => {
+                // Get the sheet name if provided, otherwise we'll use the first sheet
+                let sheet_name = params
+                    .get("sheetName")
+                    .and_then(|q| q.as_str())
+                    .map(|s| format!("{}!1:1", s))
+                    .unwrap_or_else(|| "1:1".to_string()); // Default to first row of first sheet
+
+                let result = self
+                    .sheets
+                    .spreadsheets()
+                    .values_get(spreadsheet_id, &sheet_name)
+                    .clear_scopes()
+                    .add_scope(Scope::Readonly)
+                    .doit()
+                    .await;
+
+                match result {
+                    Err(e) => Err(ToolError::ExecutionError(format!(
+                        "Failed to execute Google Sheets get_columns query, {}.",
+                        e
+                    ))),
+                    Ok(r) => {
+                        let value_range = r.1;
+                        // Extract just the headers (first row)
+                        let headers = match value_range.values {
+                            Some(mut values) if !values.is_empty() => {
+                                // Take the first row only
+                                let headers = values.remove(0);
+                                let header_values: Vec<String> = headers
+                                    .into_iter()
+                                    .map(|cell| cell.as_str().unwrap_or_default().to_string())
+                                    .collect();
+                                header_values.join(", ")
+                            }
+                            _ => "No headers found".to_string(),
+                        };
+
+                        Ok(vec![Content::text(headers).with_priority(0.1)])
+                    }
+                }
+            },
+            "get_values" => {
+                let range = params
+                    .get("range")
+                    .and_then(|q| q.as_str())
+                    .ok_or(ToolError::InvalidParameters(
+                        "The range is required for get_values operation".to_string(),
+                    ))?;
+
+                let result = self
+                    .sheets
+                    .spreadsheets()
+                    .values_get(spreadsheet_id, range)
+                    .clear_scopes()
+                    .add_scope(Scope::Readonly)
+                    .doit()
+                    .await;
+
+                match result {
+                    Err(e) => Err(ToolError::ExecutionError(format!(
+                        "Failed to execute Google Sheets values_get query, {}.",
+                        e
+                    ))),
+                    Ok(r) => {
+                        let value_range = r.1;
+                        // Convert the values to a CSV string
+                        let csv_content = match value_range.values {
+                            Some(values) => {
+                                let mut csv_string = String::new();
+                                for row in values {
+                                    let row_values: Vec<String> = row
+                                        .into_iter()
+                                        .map(|cell| cell.as_str().unwrap_or_default().to_string())
+                                        .collect();
+                                    csv_string.push_str(&row_values.join(","));
+                                    csv_string.push('\n');
+                                }
+                                csv_string
+                            }
+                            None => "No data found".to_string(),
+                        };
+
+                        Ok(vec![Content::text(csv_content).with_priority(0.1)])
+                    }
+                }
+            },
+            _ => Err(ToolError::InvalidParameters(format!(
+                "Invalid operation: {}. Supported operations are: list_sheets, get_columns, get_values",
+                operation
+            ))),
+        }
+    }
+
     async fn read_google_resource(&self, uri: String) -> Result<String, ResourceError> {
         self.read(json!({"uri": uri}))
             .await
@@ -599,6 +804,7 @@ impl Router for GoogleDriveRouter {
             match tool_name.as_str() {
                 "search" => this.search(arguments).await,
                 "read" => this.read(arguments).await,
+                "sheets_tool" => this.sheets_tool(arguments).await,
                 _ => Err(ToolError::NotFound(format!("Tool {} not found", tool_name))),
             }
         })
@@ -644,6 +850,7 @@ impl Clone for GoogleDriveRouter {
             tools: self.tools.clone(),
             instructions: self.instructions.clone(),
             drive: self.drive.clone(),
+            sheets: self.sheets.clone(),
         }
     }
 }