Shea/gdrive labels (#2537)

Co-authored-by: Michael Neale <michael.neale@gmail.com>
Co-authored-by: Kalvin C <kalvinnchau@users.noreply.github.com>

crates/goose-mcp/Cargo.toml

@@ -35,6 +35,7 @@ chrono = { version = "0.4.38", features = ["serde"] }
 etcetera = "0.8.0"
 tempfile = "3.8"
 include_dir = "0.7.4"
+google-apis-common = "7.0.0"
 google-drive3 = "6.0.0"
 google-sheets4 = "6.0.0"
 google-docs1 = "6.0.0"
@@ -47,9 +48,21 @@ lopdf = "0.35.0"
 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", "vendored"] }
+keyring = { version = "3.6.1", features = [
+    "apple-native",
+    "windows-native",
+    "sync-secret-service",
+    "vendored",
+] }
 oauth2 = { version = "5.0.0", features = ["reqwest"] }
+utoipa = { version = "4.1", optional = true }
+hyper = "1"
+serde_with = "3"
+
 
 [dev-dependencies]
 serial_test = "3.0.0"
 sysinfo = "0.32.1"
+
+[features]
+utoipa = ["dep:utoipa"]

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

@@ -0,0 +1,478 @@
+#![allow(clippy::ptr_arg, dead_code, clippy::enum_variant_names)]
+
+use std::collections::{BTreeSet, HashMap};
+
+use google_apis_common as common;
+use tokio::time::sleep;
+
+/// A scope is needed when requesting an
+/// [authorization token](https://developers.google.com/workspace/drive/labels/guides/authorize).
+#[derive(PartialEq, Eq, Ord, PartialOrd, Hash, Debug, Clone, Copy)]
+pub enum Scope {
+    /// View, use, and manage Drive labels.
+    DriveLabels,
+
+    /// View and use Drive labels.
+    DriveLabelsReadonly,
+
+    /// View, edit, create, and delete all Drive labels in your organization,
+    /// and view your organization's label-related administration policies.
+    DriveLabelsAdmin,
+
+    /// View all Drive labels and label-related administration policies in your
+    /// organization.
+    DriveLabelsAdminReadonly,
+}
+
+impl AsRef<str> for Scope {
+    fn as_ref(&self) -> &str {
+        match *self {
+            Scope::DriveLabels => "https://www.googleapis.com/auth/drive.labels",
+            Scope::DriveLabelsReadonly => "https://www.googleapis.com/auth/drive.labels.readonly",
+            Scope::DriveLabelsAdmin => "https://www.googleapis.com/auth/drive.admin.labels",
+            Scope::DriveLabelsAdminReadonly => {
+                "https://www.googleapis.com/auth/drive.admin.labels.readonly"
+            }
+        }
+    }
+}
+
+#[allow(clippy::derivable_impls)]
+impl Default for Scope {
+    fn default() -> Scope {
+        Scope::DriveLabelsReadonly
+    }
+}
+
+#[derive(Clone)]
+pub struct DriveLabelsHub<C> {
+    pub client: common::Client<C>,
+    pub auth: Box<dyn common::GetToken>,
+    _user_agent: String,
+    _base_url: String,
+}
+
+impl<C> common::Hub for DriveLabelsHub<C> {}
+
+impl<'a, C> DriveLabelsHub<C> {
+    pub fn new<A: 'static + common::GetToken>(
+        client: common::Client<C>,
+        auth: A,
+    ) -> DriveLabelsHub<C> {
+        DriveLabelsHub {
+            client,
+            auth: Box::new(auth),
+            _user_agent: "google-api-rust-client/6.0.0".to_string(),
+            _base_url: "https://drivelabels.googleapis.com/".to_string(),
+        }
+    }
+
+    pub fn labels(&'a self) -> LabelMethods<'a, C> {
+        LabelMethods { hub: self }
+    }
+
+    /// Set the user-agent header field to use in all requests to the server.
+    /// It defaults to `google-api-rust-client/6.0.0`.
+    ///
+    /// Returns the previously set user-agent.
+    pub fn user_agent(&mut self, agent_name: String) -> String {
+        std::mem::replace(&mut self._user_agent, agent_name)
+    }
+
+    /// Set the base url to use in all requests to the server.
+    /// It defaults to `https://www.googleapis.com/drive/v3/`.
+    ///
+    /// Returns the previously set base url.
+    pub fn base_url(&mut self, new_base_url: String) -> String {
+        std::mem::replace(&mut self._base_url, new_base_url)
+    }
+}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct Label {
+    #[serde(rename = "name")]
+    pub name: Option<String>,
+    #[serde(rename = "id")]
+    pub id: Option<String>,
+    #[serde(rename = "revisionId")]
+    pub revision_id: Option<String>,
+    #[serde(rename = "labelType")]
+    pub label_type: Option<String>,
+    #[serde(rename = "creator")]
+    pub creator: Option<User>,
+    #[serde(rename = "createTime")]
+    pub create_time: Option<String>,
+    #[serde(rename = "revisionCreator")]
+    pub revision_creator: Option<User>,
+    #[serde(rename = "revisionCreateTime")]
+    pub revision_create_time: Option<String>,
+    #[serde(rename = "publisher")]
+    pub publisher: Option<User>,
+    #[serde(rename = "publishTime")]
+    pub publish_time: Option<String>,
+    #[serde(rename = "disabler")]
+    pub disabler: Option<User>,
+    #[serde(rename = "disableTime")]
+    pub disable_time: Option<String>,
+    #[serde(rename = "customer")]
+    pub customer: Option<String>,
+    pub properties: Option<LabelProperty>,
+    pub fields: Option<Vec<Field>>,
+    // We ignore the remaining fields.
+}
+
+impl common::Part for Label {}
+
+impl common::ResponseResult for Label {}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct LabelProperty {
+    pub title: Option<String>,
+    pub description: Option<String>,
+}
+
+impl common::Part for LabelProperty {}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct Field {
+    id: Option<String>,
+    #[serde(rename = "queryKey")]
+    query_key: Option<String>,
+    properties: Option<FieldProperty>,
+    #[serde(rename = "selectionOptions")]
+    selection_options: Option<SelectionOption>,
+}
+
+impl common::Part for Field {}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct FieldProperty {
+    #[serde(rename = "displayName")]
+    pub display_name: Option<String>,
+    pub required: Option<bool>,
+}
+
+impl common::Part for FieldProperty {}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct SelectionOption {
+    #[serde(rename = "listOptions")]
+    pub list_options: Option<String>,
+    pub choices: Option<Vec<Choice>>,
+}
+
+impl common::Part for SelectionOption {}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct Choice {
+    id: Option<String>,
+    properties: Option<ChoiceProperties>,
+    // We ignore the remaining fields.
+}
+
+impl common::Part for Choice {}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct ChoiceProperties {
+    #[serde(rename = "displayName")]
+    display_name: Option<String>,
+    description: Option<String>,
+}
+
+impl common::Part for ChoiceProperties {}
+
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct LabelList {
+    pub labels: Option<Vec<Label>>,
+    #[serde(rename = "nextPageToken")]
+    pub next_page_token: Option<String>,
+}
+
+impl common::ResponseResult for LabelList {}
+
+/// Information about a Drive user.
+///
+/// This type is not used in any activity, and only used as *part* of another schema.
+///
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+#[serde_with::serde_as]
+#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct User {
+    /// Output only. A plain text displayable name for this user.
+    #[serde(rename = "displayName")]
+    pub display_name: Option<String>,
+    /// Output only. The email address of the user. This may not be present in certain contexts if the user has not made their email address visible to the requester.
+    #[serde(rename = "emailAddress")]
+    pub email_address: Option<String>,
+    /// Output only. Identifies what kind of resource this is. Value: the fixed string `"drive#user"`.
+    pub kind: Option<String>,
+    /// Output only. Whether this user is the requesting user.
+    pub me: Option<bool>,
+    /// Output only. The user's ID as visible in Permission resources.
+    #[serde(rename = "permissionId")]
+    pub permission_id: Option<String>,
+    /// Output only. A link to the user's profile photo, if available.
+    #[serde(rename = "photoLink")]
+    pub photo_link: Option<String>,
+}
+
+impl common::Part for User {}
+
+pub struct LabelMethods<'a, C>
+where
+    C: 'a,
+{
+    hub: &'a DriveLabelsHub<C>,
+}
+
+impl<C> common::MethodsBuilder for LabelMethods<'_, C> {}
+
+impl<'a, C> LabelMethods<'a, C> {
+    /// Create a builder to help you perform the following tasks:
+    ///
+    /// List labels
+    pub fn list(&self) -> LabelListCall<'a, C> {
+        LabelListCall {
+            hub: self.hub,
+            _delegate: Default::default(),
+            _additional_params: Default::default(),
+            _scopes: Default::default(),
+        }
+    }
+}
+
+/// Lists the workspace's labels.
+pub struct LabelListCall<'a, C>
+where
+    C: 'a,
+{
+    hub: &'a DriveLabelsHub<C>,
+    _delegate: Option<&'a mut dyn common::Delegate>,
+    _additional_params: HashMap<String, String>,
+    _scopes: BTreeSet<String>,
+}
+
+impl<C> common::CallBuilder for LabelListCall<'_, C> {}
+
+impl<'a, C> LabelListCall<'a, C>
+where
+    C: common::Connector,
+{
+    /// Perform the operation you have built so far.
+    pub async fn doit(mut self) -> common::Result<(common::Response, LabelList)> {
+        use common::url::Params;
+        use hyper::header::{AUTHORIZATION, CONTENT_LENGTH, USER_AGENT};
+
+        let mut dd = common::DefaultDelegate;
+        let dlg: &mut dyn common::Delegate = self._delegate.unwrap_or(&mut dd);
+        dlg.begin(common::MethodInfo {
+            id: "drivelabels.labels.list",
+            http_method: hyper::Method::GET,
+        });
+
+        for &field in ["alt"].iter() {
+            if self._additional_params.contains_key(field) {
+                dlg.finished(false);
+                return Err(common::Error::FieldClash(field));
+            }
+        }
+
+        // TODO: We don't handle any of the query params.
+        let mut params = Params::with_capacity(2 + self._additional_params.len());
+
+        params.extend(self._additional_params.iter());
+
+        params.push("alt", "json");
+        let url = self.hub._base_url.clone() + "v2/labels";
+
+        if self._scopes.is_empty() {
+            self._scopes
+                .insert(Scope::DriveLabelsReadonly.as_ref().to_string());
+        }
+
+        let url = params.parse_with_url(&url);
+
+        loop {
+            let token = match self
+                .hub
+                .auth
+                .get_token(&self._scopes.iter().map(String::as_str).collect::<Vec<_>>()[..])
+                .await
+            {
+                Ok(token) => token,
+                Err(e) => match dlg.token(e) {
+                    Ok(token) => token,
+                    Err(e) => {
+                        dlg.finished(false);
+                        return Err(common::Error::MissingToken(e));
+                    }
+                },
+            };
+            let req_result = {
+                let client = &self.hub.client;
+                dlg.pre_request();
+                let mut req_builder = hyper::Request::builder()
+                    .method(hyper::Method::GET)
+                    .uri(url.as_str())
+                    .header(USER_AGENT, self.hub._user_agent.clone());
+
+                if let Some(token) = token.as_ref() {
+                    req_builder = req_builder.header(AUTHORIZATION, format!("Bearer {}", token));
+                }
+
+                let request = req_builder
+                    .header(CONTENT_LENGTH, 0_u64)
+                    .body(common::to_body::<String>(None));
+                client.request(request.unwrap()).await
+            };
+
+            match req_result {
+                Err(err) => {
+                    if let common::Retry::After(d) = dlg.http_error(&err) {
+                        sleep(d).await;
+                        continue;
+                    }
+                    dlg.finished(false);
+                    return Err(common::Error::HttpError(err));
+                }
+                Ok(res) => {
+                    let (parts, body) = res.into_parts();
+                    let body = common::Body::new(body);
+                    if !parts.status.is_success() {
+                        let bytes = common::to_bytes(body).await.unwrap_or_default();
+                        let error = serde_json::from_str(&common::to_string(&bytes));
+                        let response = common::to_response(parts, bytes.into());
+
+                        if let common::Retry::After(d) =
+                            dlg.http_failure(&response, error.as_ref().ok())
+                        {
+                            sleep(d).await;
+                            continue;
+                        }
+
+                        dlg.finished(false);
+
+                        return Err(match error {
+                            Ok(value) => common::Error::BadRequest(value),
+                            _ => common::Error::Failure(response),
+                        });
+                    }
+                    let response = {
+                        let bytes = common::to_bytes(body).await.unwrap_or_default();
+                        let encoded = common::to_string(&bytes);
+                        match serde_json::from_str(&encoded) {
+                            Ok(decoded) => (common::to_response(parts, bytes.into()), decoded),
+                            Err(error) => {
+                                dlg.response_json_decode_error(&encoded, &error);
+                                return Err(common::Error::JsonDecodeError(
+                                    encoded.to_string(),
+                                    error,
+                                ));
+                            }
+                        }
+                    };
+
+                    dlg.finished(true);
+                    return Ok(response);
+                }
+            }
+        }
+    }
+
+    /// The delegate implementation is consulted whenever there is an intermediate result, or if something goes wrong
+    /// while executing the actual API request.
+    ///
+    /// ````text
+    ///                   It should be used to handle progress information, and to implement a certain level of resilience.
+    /// ````
+    ///
+    /// Sets the *delegate* property to the given value.
+    pub fn delegate(mut self, new_value: &'a mut dyn common::Delegate) -> LabelListCall<'a, C> {
+        self._delegate = Some(new_value);
+        self
+    }
+
+    /// Set any additional parameter of the query string used in the request.
+    /// It should be used to set parameters which are not yet available through their own
+    /// setters.
+    ///
+    /// Please note that this method must not be used to set any of the known parameters
+    /// which have their own setter method. If done anyway, the request will fail.
+    ///
+    /// # Additional Parameters
+    ///
+    /// * *$.xgafv* (query-string) - V1 error format.
+    /// * *access_token* (query-string) - OAuth access token.
+    /// * *alt* (query-string) - Data format for response.
+    /// * *callback* (query-string) - JSONP
+    /// * *fields* (query-string) - Selector specifying which fields to include in a partial response.
+    /// * *key* (query-string) - API key. Your API key identifies your project and provides you with API access, quota, and reports. Required unless you provide an OAuth 2.0 token.
+    /// * *oauth_token* (query-string) - OAuth 2.0 token for the current user.
+    /// * *prettyPrint* (query-boolean) - Returns response with indentations and line breaks.
+    /// * *quotaUser* (query-string) - Available to use for quota purposes for server-side applications. Can be any arbitrary string assigned to a user, but should not exceed 40 characters.
+    /// * *uploadType* (query-string) - Legacy upload protocol for media (e.g. "media", "multipart").
+    /// * *upload_protocol* (query-string) - Upload protocol for media (e.g. "raw", "multipart").
+    pub fn param<T>(mut self, name: T, value: T) -> LabelListCall<'a, C>
+    where
+        T: AsRef<str>,
+    {
+        self._additional_params
+            .insert(name.as_ref().to_string(), value.as_ref().to_string());
+        self
+    }
+
+    /// Identifies the authorization scope for the method you are building.
+    ///
+    /// Use this method to actively specify which scope should be used, instead of the default [`Scope`] variant
+    /// [`Scope::DriveLabelsReadonly`].
+    ///
+    /// The `scope` will be added to a set of scopes. This is important as one can maintain access
+    /// tokens for more than one scope.
+    ///
+    /// Usually there is more than one suitable scope to authorize an operation, some of which may
+    /// encompass more rights than others. For example, for listing resources, a *read-only* scope will be
+    /// sufficient, a read-write scope will do as well.
+    pub fn add_scope<St>(mut self, scope: St) -> LabelListCall<'a, C>
+    where
+        St: AsRef<str>,
+    {
+        self._scopes.insert(String::from(scope.as_ref()));
+        self
+    }
+    /// Identifies the authorization scope(s) for the method you are building.
+    ///
+    /// See [`Self::add_scope()`] for details.
+    pub fn add_scopes<I, St>(mut self, scopes: I) -> LabelListCall<'a, C>
+    where
+        I: IntoIterator<Item = St>,
+        St: AsRef<str>,
+    {
+        self._scopes
+            .extend(scopes.into_iter().map(|s| String::from(s.as_ref())));
+        self
+    }
+
+    /// Removes all scopes, and no default scope will be used either.
+    /// In this case, you have to specify your API-key using the `key` parameter (see [`Self::param()`]
+    /// for details).
+    pub fn clear_scopes(mut self) -> LabelListCall<'a, C> {
+        self._scopes.clear();
+        self
+    }
+}

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

@@ -1,8 +1,10 @@
+mod google_labels;
 mod oauth_pkce;
 pub mod storage;
 
 use anyhow::{Context, Error};
 use base64::Engine;
+use chrono::NaiveDate;
 use indoc::indoc;
 use lazy_static::lazy_static;
 use mcp_core::tool::ToolAnnotations;
@@ -28,11 +30,15 @@ use google_docs1::{self, Docs};
 use google_drive3::common::ReadSeek;
 use google_drive3::{
     self,
-    api::{Comment, File, FileShortcutDetails, Permission, Reply, Scope},
+    api::{
+        Comment, File, FileShortcutDetails, LabelFieldModification, LabelModification,
+        ModifyLabelsRequest, Permission, Reply, Scope,
+    },
     hyper_rustls::{self, HttpsConnector},
     hyper_util::{self, client::legacy::connect::HttpConnector},
     DriveHub,
 };
+use google_labels::DriveLabelsHub;
 use google_sheets4::{self, Sheets};
 use http_body_util::BodyExt;
 
@@ -80,6 +86,7 @@ pub struct GoogleDriveRouter {
     tools: Vec<Tool>,
     instructions: String,
     drive: DriveHub<HttpsConnector<HttpConnector>>,
+    drive_labels: DriveLabelsHub<HttpsConnector<HttpConnector>>,
     sheets: Sheets<HttpsConnector<HttpConnector>>,
     docs: Docs<HttpsConnector<HttpConnector>>,
     credentials_manager: Arc<CredentialsManager>,
@@ -88,6 +95,7 @@ pub struct GoogleDriveRouter {
 impl GoogleDriveRouter {
     async fn google_auth() -> (
         DriveHub<HttpsConnector<HttpConnector>>,
+        DriveLabelsHub<HttpsConnector<HttpConnector>>,
         Sheets<HttpsConnector<HttpConnector>>,
         Docs<HttpsConnector<HttpConnector>>,
         Arc<CredentialsManager>,
@@ -162,7 +170,7 @@ impl GoogleDriveRouter {
         // Read the OAuth credentials from the keyfile
         match fs::read_to_string(keyfile_path) {
             Ok(_) => {
-                // Create the PKCE OAuth2 clien
+                // Create the PKCE OAuth2 client
                 let auth = PkceOAuth2Client::new(keyfile_path, credentials_manager.clone())
                     .expect("Failed to create OAuth2 client");
 
@@ -180,11 +188,18 @@ impl GoogleDriveRouter {
                 );
 
                 let drive_hub = DriveHub::new(client.clone(), auth.clone());
+                let drive_labels_hub = DriveLabelsHub::new(client.clone(), auth.clone());
                 let sheets_hub = Sheets::new(client.clone(), auth.clone());
                 let docs_hub = Docs::new(client, auth);
 
                 // Create and return the DriveHub, Sheets and our PKCE OAuth2 client
-                (drive_hub, sheets_hub, docs_hub, credentials_manager)
+                (
+                    drive_hub,
+                    drive_labels_hub,
+                    sheets_hub,
+                    docs_hub,
+                    credentials_manager,
+                )
             }
             Err(e) => {
                 tracing::error!(
@@ -199,24 +214,28 @@ impl GoogleDriveRouter {
 
     pub async fn new() -> Self {
         // handle auth
-        let (drive, sheets, docs, credentials_manager) = Self::google_auth().await;
+        let (drive, drive_labels, sheets, docs, credentials_manager) = Self::google_auth().await;
 
         let search_tool = Tool::new(
             "search".to_string(),
             indoc! {r#"
-                Search for files in google drive by name, given an input search query. At least one of ('name', 'mimeType', or 'parent') are required.
+                List or search for files or labels in google drive by name, given an input search query. At least one of ('name', 'mimeType', or 'parent') are required for file searches.
             "#}
             .to_string(),
             json!({
               "type": "object",
               "properties": {
+                "driveType": {
+                    "type": "string",
+                    "description": "Required type of object to list or search (file, label)."
+                },
                 "name": {
                     "type": "string",
                     "description": "String to search for in the file's name.",
                 },
                 "mimeType": {
                     "type": "string",
-                    "description": "MIME type to constrain the search to.",
+                    "description": "Use when searching for a file to constrain the results to just this MIME type.",
                 },
                 "parent": {
                     "type": "string",
@@ -233,8 +252,13 @@ impl GoogleDriveRouter {
                 "pageSize": {
                     "type": "number",
                     "description": "How many items to return from the search query, default 10, max 100",
+                },
+                "includeLabels": {
+                    "type": "boolean",
+                    "description": "When searching or listing files, also get any applied labels.",
                 }
               },
+              "required": ["driveType"],
             }),
             Some(ToolAnnotations {
                     title: Some("Search GDrive".to_string()),
@@ -370,7 +394,7 @@ impl GoogleDriveRouter {
         let update_file_tool = Tool::new(
             "update_file".to_string(),
             indoc! {r#"
-                Update an existing file in Google Drive with new content.
+                Update an existing file in Google Drive with new content or edit the file's labels.
             "#}
             .to_string(),
             json!({
@@ -380,6 +404,10 @@ impl GoogleDriveRouter {
                       "type": "string",
                       "description": "The ID of the file to update.",
                   },
+                  "allowSharedDrives": {
+                      "type": "boolean",
+                      "description": "Whether to allow access to shared drives or just your personal drive (default: false)",
+                  },
                   "mimeType": {
                       "type": "string",
                       "description": "The MIME type of the file.",
@@ -392,15 +420,77 @@ impl GoogleDriveRouter {
                       "type": "string",
                       "description": "Path to a local file to use to update the Google Drive file. Mutually exclusive with body (required for Google Slides type)",
                   },
-                  "allowSharedDrives": {
-                      "type": "boolean",
-                      "description": "Whether to allow access to shared drives or just your personal drive (default: false)",
-                  }
+                  "updateLabels": {
+                      "type": "array",
+                      "description": "Array of label operations to perform on the file. Each operation may remove one label, unset one field, or update one field.",
+                      "items": {
+                          "type": "object",
+                          "properties": {
+                              "labelId": {
+                                  "type": "string",
+                                  "description": "The ID of the label to be operated upon."
+                              },
+                              "operation": {
+                                  "type": "string",
+                                  "enum":  ["removeLabel", "unsetField", "addOrUpdateLabel"],
+                                  "description": "The operation to perform. You may 'removeLabel' to completely remove the label from the file, 'unsetField' to remove a field from an applied label, or 'addOrUpdateLabel' to add a new label (with or without fields), or change the value of a field on an applied label."
+                              },
+                              "fieldId": {
+                                  "type": "string",
+                                  "description": "The ID of the field to be operated upon."
+                              },
+                              "dateValue": {
+                                  "type": "array",
+                                  "description": "If updating a date field, an array of RFC 3339 dates (format YYYY-MM-DD) to update to.",
+                                  "items": {
+                                      "type": "string",
+                                      "description": "An RFC 3339 full-date format YYYY-MM-DD.",
+                                  }
+                              },
+                              "textValue": {
+                                  "type": "array",
+                                  "description": "If updating a text field, the string values to update to.",
+                                  "items": {
+                                      "type": "string",
+                                      "description": "Text field values.",
+                                  }
+                              },
+                              "choiceValue": {
+                                  "type": "array",
+                                  "description": "If updating a Choice field, the ID(s) of the desired choice field(s).",
+                                  "items": {
+                                      "type": "string",
+                                      "description": "Choice ID as a string",
+                                  }
+                              },
+                              "integerValue": {
+                                  "type": "array",
+                                  "description": "If updating an integer field, the integer values to use.",
+                                  "items": {
+                                      "type": "integer",
+                                      "description": "The integer value.",
+                                  }
+                              },
+                              "userValue": {
+                                  "type": "array",
+                                  "description": "If updating a user field, an array of the email address(es) of the user(s) to set as the field value.",
+                                  "items": {
+                                      "type": "string",
+                                      "description": "Email address as a string",
+                                  }
+                              }
+                          }
+                      }
+                  },
               },
-              "required": ["fileId", "mimeType"],
+              "required": ["fileId"],
+              "dependentRequired": {
+                  "body": ["mimeType"],
+                  "path": ["mimeType"]
+              }
             }),
             Some(ToolAnnotations {
-                title: Some("Update a file".to_string()),
+                title: Some("Update a file's contents or labels".to_string()),
                 read_only_hint: false,
                 destructive_hint: true,
                 idempotent_hint: false,
@@ -466,7 +556,13 @@ impl GoogleDriveRouter {
               },
               "required": ["spreadsheetId", "operation"],
             }),
-            None,
+            Some(ToolAnnotations {
+                title: Some("Work with Google Sheets data using various operations.".to_string()),
+                read_only_hint: false,
+                destructive_hint: true,
+                idempotent_hint: false,
+                open_world_hint: false,
+            }),
         );
 
         let docs_tool = Tool::new(
@@ -517,7 +613,13 @@ impl GoogleDriveRouter {
               },
               "required": ["documentId", "operation"],
             }),
-            None,
+            Some(ToolAnnotations {
+                title: Some("Work with Google Docs data using various operations.".to_string()),
+                read_only_hint: false,
+                destructive_hint: true,
+                idempotent_hint: false,
+                open_world_hint: false,
+            }),
         );
 
         let get_comments_tool = Tool::new(
@@ -549,7 +651,7 @@ impl GoogleDriveRouter {
             "manage_comment".to_string(),
             indoc! {r#"
                 Manage comment for a Google Drive file.
-                
+
                 Supports the operations:
                 - create: Create a comment for the latest revision of a Google Drive file. The Google Drive API only supports unanchored comments (they don't refer to a specific location in the file).
                 - reply: Add a reply to a comment thread, or resolve a comment.
@@ -702,7 +804,7 @@ impl GoogleDriveRouter {
 
             ## Overview
             The Google Drive MCP server provides tools for interacting with Google Drive files, Google Sheets, and Google Docs:
-            1. search - Find files in your Google Drive
+            1. search - List or search for files or labels in your Google Drive
             2. read - Read file contents directly using a uri in the `gdrive:///uri` format
             3. move_file - Move a file to a new location in Google Drive
             4. list_drives - List the shared drives to which you have access
@@ -711,16 +813,18 @@ impl GoogleDriveRouter {
             7. get_comments - List a file or folder's comments
             8. manage_comment - Manage comment for a Google Drive file.
             9. create_file - Create a new file
-            10. update_file - Update a existing file
+            10. update_file - Update an existing file's contents or labels
             11. sheets_tool - Work with Google Sheets data using various operations
             12. docs_tool - Work with Google Docs data using various operations
 
             ## Available Tools
 
             ### 1. Search Tool
-            Search for files in Google Drive, by name and ordered by most recently viewedByMeTime.
+            Search for or list files or labels in Google Drive. Files are
+            searched by name and ordered by most recently viewedByMeTime.
             A corpora parameter controls which corpus is searched.
-            Returns: List of files with their names, MIME types, and IDs
+            Returns: List of files with their names, MIME types, and IDs or a
+            list of labels and their fields.
 
             ### 2. Read File Tool
             Read a file's contents using its ID, and optionally include images as base64 encoded data.
@@ -775,7 +879,7 @@ impl GoogleDriveRouter {
 
             ### 8. Manage Comment Tool
             Create or reply comment for a Google Drive file.
-            
+
             ### 9. Create File Tool
             Create any kind of file, including Google Workspace files (Docs, Sheets, or Slides) directly in Google Drive.
             - For Google Docs: Converts Markdown text to a Google Document
@@ -788,12 +892,18 @@ impl GoogleDriveRouter {
             include the changes as part of the entire document.
 
             ### 10. Update File Tool
-            Replace the entire contents of an existing file with new content, including Google Workspace files (Docs, Sheets, or Slides).
+            Replace the entire contents of an existing file with new content,
+            including Google Workspace files (Docs, Sheets, or Slides), or
+            update the labels applied to a file.
             - For Google Docs: Updates with new Markdown text
             - For Google Sheets: Updates with new CSV text
             - For Google Slides: Updates with a new PowerPoint file (requires a path to the powerpoint file)
             - Other: No file conversion.
 
+            Label operations include adding a new label, unsetting a field for
+            an already-applied label, removing a label, or changing the field
+            value for an applied label.
+
             ### 11. Sheets Tool
             Work with Google Sheets data using various operations:
             - list_sheets: List all sheets in a spreadsheet
@@ -879,6 +989,7 @@ impl GoogleDriveRouter {
             ],
             instructions,
             drive,
+            drive_labels,
             sheets,
             docs,
             credentials_manager,
@@ -887,6 +998,22 @@ impl GoogleDriveRouter {
 
     // Implement search tool functionality
     async fn search(&self, params: Value) -> Result<Vec<Content>, ToolError> {
+        // To minimize tool growth, we search/list for a number of different
+        // objects in Gdrive with sub-funcs.
+        let drive_type = params.get("driveType").and_then(|q| q.as_str()).ok_or(
+            ToolError::InvalidParameters("The type is required".to_string()),
+        )?;
+        match drive_type {
+            "file" => return self.search_files(params).await,
+            "label" => return self.list_labels(params).await,
+            t => Err(ToolError::InvalidParameters(format!(
+                "type must be one of ('file', 'label'), got {}",
+                t
+            ))),
+        }
+    }
+
+    async fn search_files(&self, params: Value) -> Result<Vec<Content>, ToolError> {
         let name = params.get("name").and_then(|q| q.as_str());
         let mime_type = params.get("mimeType").and_then(|q| q.as_str());
         let drive_id = params.get("driveId").and_then(|q| q.as_str());
@@ -928,6 +1055,11 @@ impl GoogleDriveRouter {
             })
             .unwrap_or(Ok(10))?;
 
+        let include_labels = params
+            .get("includeLabels")
+            .and_then(|b| b.as_bool())
+            .unwrap_or(false);
+
         let mut query = Vec::new();
         if let Some(n) = name {
             query.push(
@@ -958,7 +1090,13 @@ impl GoogleDriveRouter {
             .corpora(corpus)
             .q(query_string.as_str())
             .order_by("viewedByMeTime desc")
-            .param("fields", "files(id, name, mimeType, modifiedTime, size)")
+            .param(
+                "fields",
+                &format!(
+                    "files(id, name, mimeType, modifiedTime, size{})",
+                    if include_labels { ", labelInfo" } else { "" }
+                ),
+            )
             .page_size(page_size)
             .supports_all_drives(true)
             .include_items_from_all_drives(true)
@@ -968,8 +1106,34 @@ impl GoogleDriveRouter {
         if let (Some(d), "drive") = (drive_id, corpus) {
             builder = builder.drive_id(d);
         }
-        let result = builder.doit().await;
+        // If we want labels, we have to go look up the IDs first.
+        // let mut label_results: Vec<Label> = Vec::new();
+        if include_labels {
+            let label_builder = self
+                .drive_labels
+                .labels()
+                .list()
+                .param("view", "LABEL_VIEW_BASIC");
+            // .param("view", "LABEL_VIEW_FULL");
 
+            let label_results = match label_builder.doit().await {
+                Ok(r) => r.1.labels.unwrap_or_default(),
+                Err(e) => {
+                    return Err(ToolError::ExecutionError(format!(
+                        "Failed to execute google drive label list '{}'.",
+                        e
+                    )))
+                }
+            };
+            let label_ids = label_results
+                .iter()
+                .filter_map(|l| l.id.clone())
+                .collect::<Vec<_>>()
+                .join(",");
+            builder = builder.include_labels(&label_ids);
+        }
+
+        let result = builder.doit().await;
         match result {
             Err(e) => Err(ToolError::ExecutionError(format!(
                 "Failed to execute google drive search query '{}', {}.",
@@ -982,10 +1146,15 @@ impl GoogleDriveRouter {
                         .map(|files| {
                             files.into_iter().map(|f| {
                                 format!(
-                                    "{} ({}) (uri: {})",
+                                    "{} ({}) (uri: {}){}",
                                     f.name.unwrap_or_default(),
                                     f.mime_type.unwrap_or_default(),
-                                    f.id.unwrap_or_default()
+                                    f.id.unwrap_or_default(),
+                                    if include_labels {
+                                        format!(" (labels: {:?})", f.label_info.unwrap_or_default())
+                                    } else {
+                                        "".to_string()
+                                    }
                                 )
                             })
                         })
@@ -1985,17 +2154,173 @@ impl GoogleDriveRouter {
                     "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 allow_shared_drives = params
+            .get("allowSharedDrives")
+            .and_then(|q| q.as_bool())
+            .unwrap_or_default();
 
+        let mime_type = params.get("mimeType").and_then(|q| q.as_str());
         let body = params.get("body").and_then(|q| q.as_str());
         let path = params.get("path").and_then(|q| q.as_str());
 
+        let mut final_result = vec![];
+
+        if mime_type.is_some() && (body.is_some() || path.is_some()) {
+            let update_result = self
+                .update_file_contents(file_id, mime_type.unwrap(), body, path, allow_shared_drives)
+                .await?;
+            final_result.extend(update_result);
+        };
+
+        if let Some(label_ops) = params.get("updateLabels").and_then(|q| q.as_array()) {
+            let label_result = self.update_label(file_id, label_ops).await?;
+            final_result.extend(label_result);
+        };
+        Ok(final_result)
+    }
+
+    async fn update_label(
+        &self,
+        file_id: &str,
+        label_ops: &Vec<Value>,
+    ) -> Result<Vec<Content>, ToolError> {
+        let mut req = ModifyLabelsRequest::default();
+        let mut label_mods = vec![];
+
+        for op in label_ops {
+            if let Some(op) = op.as_object() {
+                let label_id = op.get("labelId").and_then(|o| o.as_str()).ok_or(
+                    ToolError::InvalidParameters(
+                        "The labelId param is required for label changes".to_string(),
+                    ),
+                )?;
+                match op.get("operation").and_then(|o| o.as_str()) {
+                    Some("removeLabel") => {
+                        let removal = LabelModification {
+                            label_id: Some(label_id.to_string()),
+                            remove_label: Some(true),
+                            ..Default::default()
+                        };
+                        label_mods.push(removal);
+                    }
+                    Some("unsetField") => {
+                        let field_id = op.get("fieldId").and_then(|o| o.as_str()).ok_or(
+                            ToolError::InvalidParameters(
+                                "The fieldId param is required for unsetting a field.".to_string(),
+                            ),
+                        )?;
+                        let field_mods = LabelFieldModification {
+                            field_id: Some(field_id.to_string()),
+                            unset_values: Some(true),
+                            ..Default::default()
+                        };
+                        let unset_mod = LabelModification {
+                            label_id: Some(label_id.to_string()),
+                            field_modifications: Some(vec![field_mods]),
+                            ..Default::default()
+                        };
+                        label_mods.push(unset_mod);
+                    }
+                    Some("addOrUpdateLabel") => {
+                        let mut field_mods = LabelFieldModification::default();
+                        // Not all labels _have_ fields.
+                        if let Some(field_id) = op.get("fieldId").and_then(|o| o.as_str()) {
+                            field_mods.field_id = Some(field_id.to_string());
+                        }
+
+                        if let Some(date_value) = op.get("dateValue").and_then(|o| o.as_array()) {
+                            let parsed_dates: Result<Vec<NaiveDate>, ToolError> = date_value
+                                .iter()
+                                .filter_map(|d| d.as_str())
+                                .map(|d| {
+                                    NaiveDate::parse_from_str(d, "%Y-%m-%d").map_err(|e| {
+                                        ToolError::InvalidParameters(format!(
+                                            "Error parsing field date: {}",
+                                            e
+                                        ))
+                                    })
+                                })
+                                .collect();
+
+                            field_mods.set_date_values = Some(parsed_dates?);
+                        } else if let Some(text_value) =
+                            op.get("textValue").and_then(|o| o.as_array())
+                        {
+                            field_mods.set_text_values = Some(
+                                text_value
+                                    .iter()
+                                    .map(|s| s.as_str().unwrap_or_default().to_string())
+                                    .collect(),
+                            );
+                        } else if let Some(choice_value) =
+                            op.get("choiceValue").and_then(|o| o.as_array())
+                        {
+                            field_mods.set_selection_values = Some(
+                                choice_value
+                                    .iter()
+                                    .map(|s| s.as_str().unwrap_or_default().to_string())
+                                    .collect(),
+                            );
+                        } else if let Some(int_value) =
+                            op.get("integerValue").and_then(|o| o.as_array())
+                        {
+                            field_mods.set_integer_values = Some(
+                                int_value
+                                    .iter()
+                                    .map(|s| s.as_i64().unwrap_or_default())
+                                    .collect(),
+                            );
+                        } else if let Some(user_value) =
+                            op.get("userValue").and_then(|o| o.as_array())
+                        {
+                            field_mods.set_user_values = Some(
+                                user_value
+                                    .iter()
+                                    .map(|s| s.as_str().unwrap_or_default().to_string())
+                                    .collect(),
+                            );
+                        }
+
+                        let update_mod = LabelModification {
+                            label_id: Some(label_id.to_string()),
+                            field_modifications: Some(vec![field_mods]),
+                            ..Default::default()
+                        };
+                        label_mods.push(update_mod);
+                    }
+                    _ => {
+                        return Err(ToolError::InvalidParameters(format!(
+                            "Label operation invalid: {:?}",
+                            op.get("operation")
+                        )))
+                    }
+                }
+            };
+        }
+        req.label_modifications = Some(label_mods);
+
+        let result = self.drive.files().modify_labels(req, file_id).doit().await;
+        match result {
+            Err(e) => Err(ToolError::ExecutionError(format!(
+                "Failed to update label for google drive file {}, {}.",
+                file_id, e
+            ))),
+            Ok(r) => Ok(vec![Content::text(format!(
+                "file URI: {}, labels modified: {:?}",
+                file_id,
+                r.1.modified_labels.unwrap_or_default()
+            ))]),
+        }
+    }
+
+    async fn update_file_contents(
+        &self,
+        file_id: &str,
+        mime_type: &str,
+        body: Option<&str>,
+        path: Option<&str>,
+        allow_shared_drives: bool,
+    ) -> Result<Vec<Content>, ToolError> {
         // Determine source and target MIME types based on file_type
         let (source_mime_type, target_mime_type, reader): (String, String, Box<dyn ReadSeek>) =
             match mime_type {
@@ -2064,11 +2389,6 @@ impl GoogleDriveRouter {
                 }
             };
 
-        let allow_shared_drives = params
-            .get("allowSharedDrives")
-            .and_then(|q| q.as_bool())
-            .unwrap_or_default();
-
         self.upload_to_drive(
             FileOperation::Update {
                 file_id: file_id.to_string(),
@@ -2898,6 +3218,43 @@ impl GoogleDriveRouter {
             )),
         }
     }
+
+    async fn list_labels(&self, _params: Value) -> Result<Vec<Content>, ToolError> {
+        let builder = self
+            .drive_labels
+            .labels()
+            .list()
+            .param("view", "LABEL_VIEW_FULL");
+
+        let result = builder.doit().await;
+        match result {
+            Err(e) => Err(ToolError::ExecutionError(format!(
+                "Failed to list labels for Google Drive {}",
+                e
+            ))),
+            Ok(r) => {
+                let content =
+                    r.1.labels
+                        .map(|labels| {
+                            labels.into_iter().map(|l| {
+                                format!(
+                                    "name: {} label_type: {} properties: {:?} uri: {} fields: {:?}",
+                                    l.name.unwrap_or_default(),
+                                    l.label_type.unwrap_or_default(),
+                                    l.properties.unwrap_or_default(),
+                                    l.id.unwrap_or_default(),
+                                    l.fields.unwrap_or_default()
+                                )
+                            })
+                        })
+                        .into_iter()
+                        .flatten()
+                        .collect::<Vec<_>>()
+                        .join("\n");
+                Ok(vec![Content::text(content.to_string()).with_priority(0.3)])
+            }
+        }
+    }
 }
 
 impl Router for GoogleDriveRouter {
@@ -2986,6 +3343,7 @@ impl Clone for GoogleDriveRouter {
             tools: self.tools.clone(),
             instructions: self.instructions.clone(),
             drive: self.drive.clone(),
+            drive_labels: self.drive_labels.clone(),
             sheets: self.sheets.clone(),
             docs: self.docs.clone(),
             credentials_manager: self.credentials_manager.clone(),

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

@@ -46,6 +46,7 @@ struct TokenData {
     #[serde(skip_serializing_if = "Option::is_none")]
     expires_at: Option<u64>,
     project_id: String,
+    scopes: Vec<String>,
 }
 
 /// PkceOAuth2Client implements the GetToken trait required by DriveHub
@@ -56,6 +57,7 @@ pub struct PkceOAuth2Client {
     credentials_manager: Arc<CredentialsManager>,
     http_client: reqwest::Client,
     project_id: String,
+    scopes: Vec<String>,
 }
 
 impl PkceOAuth2Client {
@@ -69,6 +71,7 @@ impl PkceOAuth2Client {
 
         // Extract the project_id from the config
         let project_id = config.installed.project_id.clone();
+        let scopes = vec![];
 
         // Create OAuth URLs
         let auth_url =
@@ -97,6 +100,7 @@ impl PkceOAuth2Client {
             credentials_manager,
             http_client,
             project_id,
+            scopes,
         })
     }
 
@@ -185,6 +189,7 @@ impl PkceOAuth2Client {
                 refresh_token: refresh_token_str.clone(),
                 expires_at,
                 project_id: self.project_id.clone(),
+                scopes: scopes.iter().map(|s| s.to_string()).collect(),
             };
 
             // Store updated token data
@@ -239,6 +244,7 @@ impl PkceOAuth2Client {
             refresh_token: new_refresh_token.clone(),
             expires_at,
             project_id: self.project_id.clone(),
+            scopes: self.scopes.clone(),
         };
 
         // Store updated token data
@@ -315,37 +321,66 @@ impl GetToken for PkceOAuth2Client {
             if let Ok(token_data) = self.credentials_manager.read_credentials::<TokenData>() {
                 // Verify the project_id matches
                 if token_data.project_id == self.project_id {
-                    // Check if the token is expired or expiring within a 5-min buffer
-                    if !self.is_token_expired(token_data.expires_at, 300) {
-                        return Ok(Some(token_data.access_token));
-                    }
+                    // Convert stored scopes to &str slices for comparison
+                    let stored_scope_refs: Vec<&str> =
+                        token_data.scopes.iter().map(|s| s.as_str()).collect();
 
-                    // Token is expired or will expire soon, try to refresh it
-                    debug!("Token is expired or will expire soon, refreshing...");
+                    // Check if we need additional scopes
+                    let needs_additional_scopes = scopes.iter().any(|&scope| {
+                        !stored_scope_refs
+                            .iter()
+                            .any(|&stored| stored.contains(scope))
+                    });
 
-                    // Try to refresh the token
-                    if let Ok(access_token) = self.refresh_token(&token_data.refresh_token).await {
-                        debug!("Successfully refreshed access token");
-                        return Ok(Some(access_token));
+                    if !needs_additional_scopes {
+                        // Check if the token is expired or expiring within a 5-min buffer
+                        if !self.is_token_expired(token_data.expires_at, 300) {
+                            return Ok(Some(token_data.access_token));
+                        }
+
+                        // Token is expired or will expire soon, try to refresh it
+                        debug!("Token is expired or will expire soon, refreshing...");
+
+                        // Try to refresh the token
+                        if let Ok(access_token) =
+                            self.refresh_token(&token_data.refresh_token).await
+                        {
+                            debug!("Successfully refreshed access token");
+                            return Ok(Some(access_token));
+                        }
+                    } else {
+                        // Only allocate new strings when we need to combine scopes
+                        let mut combined_scopes: Vec<&str> =
+                            Vec::with_capacity(scopes.len() + stored_scope_refs.len());
+                        combined_scopes.extend(scopes);
+                        combined_scopes.extend(stored_scope_refs.iter().filter(|&&stored| {
+                            !scopes.iter().any(|&scope| stored.contains(scope))
+                        }));
+
+                        return self
+                            .perform_oauth_flow(&combined_scopes)
+                            .await
+                            .map(Some)
+                            .map_err(|e| {
+                                error!("OAuth flow failed: {}", e);
+                                e
+                            });
                     }
                 }
             }
-
             // If we get here, either:
             // 1. The project ID didn't match
             // 2. Token refresh failed
             // 3. There are no valid tokens yet
+            // 4. We didn't have to change the scopes of an existing token
             // Fallback: perform interactive OAuth flow
-            match self.perform_oauth_flow(scopes).await {
-                Ok(token) => {
-                    debug!("Successfully obtained new access token through OAuth flow");
-                    Ok(Some(token))
-                }
-                Err(e) => {
+            self.perform_oauth_flow(scopes)
+                .await
+                .map(Some)
+                .map_err(|e| {
                     error!("OAuth flow failed: {}", e);
-                    Err(e)
-                }
-            }
+                    e
+                })
         })
     }
 }