From 8bda119d80584e445cb85b1505b41e4f8d0bce0e Mon Sep 17 00:00:00 2001 From: Shea Craig Date: Fri, 14 Mar 2025 15:04:33 -0400 Subject: [PATCH] feat: google_drive write tools and read comment tool (#1650) --- crates/goose-mcp/Cargo.toml | 5 +- crates/goose-mcp/src/google_drive/mod.rs | 735 ++++++++++++++++++++++- 2 files changed, 737 insertions(+), 3 deletions(-) diff --git a/crates/goose-mcp/Cargo.toml b/crates/goose-mcp/Cargo.toml index cd718134..66ab11dc 100644 --- a/crates/goose-mcp/Cargo.toml +++ b/crates/goose-mcp/Cargo.toml @@ -26,7 +26,10 @@ kill_tree = "0.2.4" shellexpand = "3.1.0" indoc = "2.0.5" xcap = "0.0.14" -reqwest = { version = "0.11", features = ["json", "rustls-tls"] , default-features = false} +reqwest = { version = "0.11", features = [ + "json", + "rustls-tls", +], default-features = false } async-trait = "0.1" chrono = { version = "0.4.38", features = ["serde"] } etcetera = "0.8.0" diff --git a/crates/goose-mcp/src/google_drive/mod.rs b/crates/goose-mcp/src/google_drive/mod.rs index 1bd545ff..fa19fd47 100644 --- a/crates/goose-mcp/src/google_drive/mod.rs +++ b/crates/goose-mcp/src/google_drive/mod.rs @@ -3,9 +3,12 @@ mod token_storage; use indoc::indoc; use regex::Regex; use serde_json::{json, Value}; -use std::{env, fs, future::Future, path::Path, pin::Pin, sync::Arc}; use token_storage::{CredentialsManager, KeychainTokenStorage}; +use std::io::Cursor; +use std::sync::Arc; +use std::{env, fs, future::Future, path::Path, pin::Pin}; + use mcp_core::content::Content; use mcp_core::{ handler::{PromptError, ResourceError, ToolError}, @@ -17,6 +20,7 @@ use mcp_core::{ use mcp_server::router::CapabilitiesBuilder; use mcp_server::Router; +use google_drive3::common::ReadSeek; use google_drive3::{ self, api::{File, Scope}, @@ -64,6 +68,12 @@ impl InstalledFlowDelegate for LocalhostBrowserDelegate { } } +#[derive(Debug)] +enum FileOperation { + Create { name: String }, + Update { file_id: String }, +} + pub struct GoogleDriveRouter { tools: Vec, instructions: String, @@ -232,6 +242,246 @@ impl GoogleDriveRouter { }), ); + let upload_tool = Tool::new( + "upload".to_string(), + indoc! {r#" + Upload a file to Google Drive. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "The desired filename to use for the uploaded file.", + }, + "mimeType": { + "type": "string", + "description": "The MIME type of the file.", + }, + "body": { + "type": "string", + "description": "Plain text body of the file to upload. Mutually exclusive with path.", + }, + "path": { + "type": "string", + "description": "Path to the file to upload. Mutually exclusive with body.", + }, + "parent_id": { + "type": "string", + "description": "ID of the parent folder in which to create the file. (default: creates files in the root of 'My Drive')", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["name", "mimeType"], + }), + ); + + let create_doc_tool = Tool::new( + "create_doc".to_string(), + indoc! {r#" + Create a Google Doc from markdown text in Google Drive. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Name of the file to create", + }, + "body": { + "type": "string", + "description": "Markdown text of the file to create.", + }, + "parent_id": { + "type": "string", + "description": "ID of the parent folder in which to create the file. (default: creates files in the root of 'My Drive')", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["name", "body"], + }), + ); + + let create_sheets_tool = Tool::new( + "create_sheets".to_string(), + indoc! {r#" + Create a Google Sheets document from csv text in Google Drive. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Name of the file to create", + }, + "body": { + "type": "string", + "description": "CSV text of the file to create.", + }, + "parent_id": { + "type": "string", + "description": "ID of the parent folder in which to create the file. (default: creates files in the root of 'My Drive')", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["name", "body"], + }), + ); + + let create_slides_tool = Tool::new( + "create_slides".to_string(), + indoc! {r#" + Create a Google Slides document in Google Drive by converting a PowerPoint file. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Name of the file to create", + }, + "path": { + "type": "string", + "description": "Path to a PowerPoint file to upload.", + }, + "parent_id": { + "type": "string", + "description": "ID of the parent folder in which to create the file. (default: creates files in the root of 'My Drive')", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["name", "path"], + }), + ); + + let update_tool = Tool::new( + "update".to_string(), + indoc! {r#" + Update a Google Drive file with new content. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "fileId": { + "type": "string", + "description": "The ID of the file to update.", + }, + "mimeType": { + "type": "string", + "description": "The MIME type of the file.", + }, + "body": { + "type": "string", + "description": "Plain text body of the file to upload. Mutually exclusive with path.", + }, + "path": { + "type": "string", + "description": "Path to a local file to use to update the Google Drive file. Mutually exclusive with body.", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["fileId", "mimeType"], + }), + ); + + let update_doc_tool = Tool::new( + "update_doc".to_string(), + indoc! {r#" + Update a Google Doc from markdown text. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "fileId": { + "type": "string", + "description": "ID of the file to update", + }, + "body": { + "type": "string", + "description": "Complete markdown text of the file to update.", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["fileId", "body"], + }), + ); + + let update_sheets_tool = Tool::new( + "update_sheets".to_string(), + indoc! {r#" + Update a Google Sheets document from csv text. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "fileId": { + "type": "string", + "description": "ID of the file to update", + }, + "body": { + "type": "string", + "description": "Complete CSV text of the updated file.", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["fileId", "body"], + }), + ); + + let update_slides_tool = Tool::new( + "update_slides".to_string(), + indoc! {r#" + Updatea Google Slides document in Google Drive by converting a PowerPoint file. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "fileId": { + "type": "string", + "description": "ID of the file to update", + }, + "path": { + "type": "string", + "description": "Path to a PowerPoint file to upload to replace the existing file.", + }, + "allow_shared_drives": { + "type": "boolean", + "description": "Whether to allow access to shared drives or just your personal drive (default: false)", + } + }, + "required": ["fileId", "path"], + }), + ); + let sheets_tool = Tool::new( "sheets_tool".to_string(), indoc! {r#" @@ -267,6 +517,28 @@ impl GoogleDriveRouter { }), ); + let comment_list_tool = Tool::new( + "comment_list".to_string(), + indoc! {r#" + List comments for a file in google drive by id, given an input file id. + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "fileId": { + "type": "string", + "description": "Id of the file to list comments for.", + }, + "pageSize": { + "type": "number", + "description": "How many items to return from the search query, default 10, max 100", + } + }, + "required": ["fileId"], + }), + ); + let instructions = indoc::formatdoc! {r#" Google Drive MCP Server Instructions @@ -336,7 +608,20 @@ impl GoogleDriveRouter { "#}; Self { - tools: vec![search_tool, read_tool, sheets_tool], + tools: vec![ + search_tool, + read_tool, + upload_tool, + create_doc_tool, + create_sheets_tool, + create_slides_tool, + update_tool, + update_doc_tool, + update_sheets_tool, + update_slides_tool, + sheets_tool, + comment_list_tool, + ], instructions, drive, sheets, @@ -810,6 +1095,443 @@ impl GoogleDriveRouter { } } } + + async fn upload_to_drive( + &self, + operation: FileOperation, + content: Box, + source_mime_type: &str, + target_mime_type: &str, + parent: Option<&str>, + support_all_drives: bool, + ) -> Result, ToolError> { + let mut req = File { + mime_type: Some(target_mime_type.to_string()), + ..Default::default() + }; + if let Some(p) = parent { + req.parents = Some(vec![p.to_string()]); + } + + let builder = self.drive.files(); + let result = match operation { + FileOperation::Create { ref name } => { + req.name = Some(name.to_string()); + builder + .create(req) + .use_content_as_indexable_text(true) + .supports_all_drives(support_all_drives) + .upload(content, source_mime_type.parse().unwrap()) + .await + } + FileOperation::Update { ref file_id } => { + builder + .update(req, file_id) + .use_content_as_indexable_text(true) + .supports_all_drives(support_all_drives) + .upload(content, source_mime_type.parse().unwrap()) + .await + } + }; + match result { + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to upload google drive file {:?}, {}.", + operation, e + ))), + Ok(r) => Ok(vec![Content::text(format!( + "{} ({}) (uri: {})", + r.1.name.unwrap_or_default(), + r.1.mime_type.unwrap_or_default(), + r.1.id.unwrap_or_default() + ))]), + } + } + + async fn upload(&self, params: Value) -> Result, ToolError> { + let filename = + params + .get("name") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The name param is required".to_string(), + ))?; + let mime_type = + params + .get("mimeType") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The mimeType param is required".to_string(), + ))?; + let body = params.get("body").and_then(|q| q.as_str()); + let path = params.get("path").and_then(|q| q.as_str()); + let reader: Box = match (body, path) { + (None, None) | (Some(_), Some(_)) => { + return Err(ToolError::InvalidParameters( + "Either the body or path param is required".to_string(), + )) + } + (Some(b), None) => Box::new(Cursor::new(b.as_bytes().to_owned())), + (None, Some(p)) => Box::new(std::fs::File::open(p).map_err(|e| { + ToolError::ExecutionError(format!("Error opening {}: {}", p, e).to_string()) + })?), + }; + let parent = params.get("parent").and_then(|q| q.as_str()); + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + self.upload_to_drive( + FileOperation::Create { + name: filename.to_string(), + }, + reader, + mime_type, + mime_type, + parent, + support_all_drives, + ) + .await + } + + async fn create_doc(&self, params: Value) -> Result, ToolError> { + let filename = + params + .get("name") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The name param is required".to_string(), + ))?; + let body = + params + .get("body") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The body param is required".to_string(), + ))?; + let source_mime_type = "text/markdown"; + let target_mime_type = "application/vnd.google-apps.document"; + let parent = params.get("parent").and_then(|q| q.as_str()); + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + let cursor = Box::new(Cursor::new(body.as_bytes().to_owned())); + self.upload_to_drive( + FileOperation::Create { + name: filename.to_string(), + }, + cursor, + source_mime_type, + target_mime_type, + parent, + support_all_drives, + ) + .await + } + + async fn create_sheets(&self, params: Value) -> Result, ToolError> { + let filename = + params + .get("name") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The name param is required".to_string(), + ))?; + let body = + params + .get("body") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The body param is required".to_string(), + ))?; + let source_mime_type = "text/csv"; + let target_mime_type = "application/vnd.google-apps.spreadsheet"; + let parent = params.get("parent").and_then(|q| q.as_str()); + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + let cursor = Box::new(Cursor::new(body.as_bytes().to_owned())); + self.upload_to_drive( + FileOperation::Create { + name: filename.to_string(), + }, + cursor, + source_mime_type, + target_mime_type, + parent, + support_all_drives, + ) + .await + } + + async fn create_slides(&self, params: Value) -> Result, ToolError> { + let filename = + params + .get("name") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The name param is required".to_string(), + ))?; + let path = + params + .get("path") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The path param is required".to_string(), + ))?; + let reader = Box::new(std::fs::File::open(path).map_err(|e| { + ToolError::ExecutionError(format!("Error opening {}: {}", path, e).to_string()) + })?); + let source_mime_type = + "application/vnd.openxmlformats-officedocument.presentationml.presentation"; + let target_mime_type = "application/vnd.google-apps.presentation"; + let parent = params.get("parent").and_then(|q| q.as_str()); + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + self.upload_to_drive( + FileOperation::Create { + name: filename.to_string(), + }, + reader, + source_mime_type, + target_mime_type, + parent, + support_all_drives, + ) + .await + } + + async fn update(&self, params: Value) -> Result, ToolError> { + let file_id = + params + .get("fileId") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The fileId param is required".to_string(), + ))?; + let mime_type = + params + .get("mimeType") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The mimeType param is required".to_string(), + ))?; + let body = params.get("body").and_then(|q| q.as_str()); + let path = params.get("path").and_then(|q| q.as_str()); + let reader: Box = match (body, path) { + (None, None) | (Some(_), Some(_)) => { + return Err(ToolError::InvalidParameters( + "Either the body or path param is required".to_string(), + )) + } + (Some(b), None) => Box::new(Cursor::new(b.as_bytes().to_owned())), + (None, Some(p)) => Box::new(std::fs::File::open(p).map_err(|e| { + ToolError::ExecutionError(format!("Error opening {}: {}", p, e).to_string()) + })?), + }; + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + + self.upload_to_drive( + FileOperation::Update { + file_id: file_id.to_string(), + }, + reader, + mime_type, + mime_type, + None, + support_all_drives, + ) + .await + } + + async fn update_doc(&self, params: Value) -> Result, ToolError> { + let file_id = + params + .get("fileId") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The fileId param is required".to_string(), + ))?; + let body = + params + .get("body") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The body param is required".to_string(), + ))?; + let source_mime_type = "text/markdown"; + let target_mime_type = "application/vnd.google-apps.document"; + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + let cursor = Box::new(Cursor::new(body.as_bytes().to_owned())); + self.upload_to_drive( + FileOperation::Update { + file_id: file_id.to_string(), + }, + cursor, + source_mime_type, + target_mime_type, + None, + support_all_drives, + ) + .await + } + + async fn update_sheets(&self, params: Value) -> Result, ToolError> { + let file_id = + params + .get("fileId") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The fileId param is required".to_string(), + ))?; + let body = + params + .get("body") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The body param is required".to_string(), + ))?; + let source_mime_type = "text/csv"; + let target_mime_type = "application/vnd.google-apps.spreadsheet"; + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + let cursor = Box::new(Cursor::new(body.as_bytes().to_owned())); + self.upload_to_drive( + FileOperation::Update { + file_id: file_id.to_string(), + }, + cursor, + source_mime_type, + target_mime_type, + None, + support_all_drives, + ) + .await + } + + async fn update_slides(&self, params: Value) -> Result, ToolError> { + let file_id = + params + .get("fileId") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The fileId param is required".to_string(), + ))?; + let path = + params + .get("path") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The path param is required".to_string(), + ))?; + let reader = Box::new(std::fs::File::open(path).map_err(|e| { + ToolError::ExecutionError(format!("Error opening {}: {}", path, e).to_string()) + })?); + let source_mime_type = + "application/vnd.openxmlformats-officedocument.presentationml.presentation"; + let target_mime_type = "application/vnd.google-apps.presentation"; + let support_all_drives = params + .get("supportAllDrives") + .and_then(|q| q.as_bool()) + .unwrap_or_default(); + self.upload_to_drive( + FileOperation::Update { + file_id: file_id.to_string(), + }, + reader, + source_mime_type, + target_mime_type, + None, + support_all_drives, + ) + .await + } + + async fn comment_list(&self, params: Value) -> Result, ToolError> { + let file_id = + params + .get("fileId") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The fileId param is required".to_string(), + ))?; + + // extract pageSize, and convert it to an i32, default to 10 + let page_size: i32 = params + .get("pageSize") + .map(|s| { + s.as_i64() + .and_then(|n| i32::try_from(n).ok()) + .ok_or_else(|| ToolError::InvalidParameters(format!("Invalid pageSize: {}", s))) + .and_then(|n| { + if (0..=100).contains(&n) { + Ok(n) + } else { + Err(ToolError::InvalidParameters(format!( + "pageSize must be between 0 and 100, got {}", + n + ))) + } + }) + }) + .unwrap_or(Ok(10))?; + + let result = self + .drive + .comments() + .list(file_id) + .page_size(page_size) + .param( + "fields", + "comments(author, content, createdTime, modifiedTime, id, anchor, resolved)", + ) + .clear_scopes() + .add_scope(Scope::Readonly) + .doit() + .await; + + match result { + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to execute google drive comment list, {}.", + e + ))), + Ok(r) => { + let content = + r.1.comments + .map(|comments| { + comments.into_iter().map(|c| { + format!( + "Author:{:?} Content: {} (created time: {}) (modified time: {})(anchor: {}) (resolved: {}) (id: {})", + c.author.unwrap_or_default(), + c.content.unwrap_or_default(), + c.created_time.unwrap_or_default(), + c.modified_time.unwrap_or_default(), + c.anchor.unwrap_or_default(), + c.resolved.unwrap_or_default(), + c.id.unwrap_or_default() + ) + }) + }) + .into_iter() + .flatten() + .collect::>() + .join("\n"); + + Ok(vec![Content::text(content.to_string())]) + } + } + } } impl Router for GoogleDriveRouter { @@ -843,7 +1565,16 @@ impl Router for GoogleDriveRouter { match tool_name.as_str() { "search" => this.search(arguments).await, "read" => this.read(arguments).await, + "upload" => this.upload(arguments).await, + "create_doc" => this.create_doc(arguments).await, + "create_sheets" => this.create_sheets(arguments).await, + "create_slides" => this.create_slides(arguments).await, + "update" => this.update(arguments).await, + "update_doc" => this.update_doc(arguments).await, + "update_sheets" => this.update_sheets(arguments).await, + "update_slides" => this.update_slides(arguments).await, "sheets_tool" => this.sheets_tool(arguments).await, + "comment_list" => this.comment_list(arguments).await, _ => Err(ToolError::NotFound(format!("Tool {} not found", tool_name))), } })