Merge pull request #3 from cfdude/dev

feat(drive): Add update_drive_file function with OAuth fixes and comprehensive file management

Author: cfdude

.gitignore

@@ -29,3 +29,4 @@ mcp_server_debug.log
 # ---- Local development files -------------------------------------------
 /.credentials
 /.claude
+.serena/

README.md

@@ -698,8 +698,9 @@ cp .env.oauth21 .env
 |------|------|-------------|
 | `search_drive_files` | **Core** | Search files with query syntax |
 | `get_drive_file_content` | **Core** | Read file content (Office formats) |
-| `list_drive_items` | Extended | List folder contents |
 | `create_drive_file` | **Core** | Create files or fetch from URLs |
+| `list_drive_items` | Extended | List folder contents |
+| `update_drive_file` | Extended | Update file metadata, move between folders |
 
 </td>
 </tr>

auth/__init__.py

@@ -1 +1 @@
-# Make the auth directory a Python package
+# Make the auth directory a Python package

auth/auth_info_middleware.py

@@ -27,116 +27,126 @@
 class GoogleWorkspaceAuthProvider(AuthProvider):
     """
     Authentication provider for Google Workspace integration.
-    
+
     This provider implements the Remote Authentication pattern where:
     - Google acts as the Authorization Server (AS)
     - This MCP server acts as a Resource Server (RS)
     - Tokens are verified using Google's public keys
     """
-    
+
     def __init__(self):
         """Initialize the Google Workspace auth provider."""
         super().__init__()
-        
+
         # Get configuration from OAuth config
         from auth.oauth_config import get_oauth_config
+
         config = get_oauth_config()
-        
+
         self.client_id = config.client_id
         self.client_secret = config.client_secret
         self.base_url = config.get_oauth_base_url()
         self.port = config.port
-        
+
         if not self.client_id:
-            logger.warning("GOOGLE_OAUTH_CLIENT_ID not set - OAuth 2.1 authentication will not work")
+            logger.warning(
+                "GOOGLE_OAUTH_CLIENT_ID not set - OAuth 2.1 authentication will not work"
+            )
             return
-            
+
         # Initialize JWT verifier for Google tokens
         self.jwt_verifier = JWTVerifier(
             jwks_uri="https://www.googleapis.com/oauth2/v3/certs",
             issuer="https://accounts.google.com",
             audience=self.client_id,
-            algorithm="RS256"
+            algorithm="RS256",
         )
-        
+
         # Session bridging now handled by OAuth21SessionStore
-        
+
     async def verify_token(self, token: str) -> Optional[AccessToken]:
         """
         Verify a bearer token issued by Google.
-        
+
         Args:
             token: The bearer token to verify
-            
+
         Returns:
             AccessToken object if valid, None otherwise
         """
         if not self.client_id:
             return None
-            
+
         try:
             # Use FastMCP's JWT verifier
             access_token = await self.jwt_verifier.verify_token(token)
-            
+
             if access_token:
                 # Store session info in OAuth21SessionStore for credential bridging
                 user_email = access_token.claims.get("email")
                 if user_email:
                     from auth.oauth21_session_store import get_oauth21_session_store
+
                     store = get_oauth21_session_store()
                     session_id = f"google_{access_token.claims.get('sub', 'unknown')}"
-                    
+
                     # Try to get FastMCP session ID for binding
                     mcp_session_id = None
                     try:
                         from fastmcp.server.dependencies import get_context
+
                         ctx = get_context()
-                        if ctx and hasattr(ctx, 'session_id'):
+                        if ctx and hasattr(ctx, "session_id"):
                             mcp_session_id = ctx.session_id
-                            logger.debug(f"Binding MCP session {mcp_session_id} to user {user_email}")
+                            logger.debug(
+                                f"Binding MCP session {mcp_session_id} to user {user_email}"
+                            )
                     except Exception:
                         pass
-                    
+
                     store.store_session(
                         user_email=user_email,
                         access_token=token,
                         scopes=access_token.scopes or [],
                         session_id=session_id,
-                        mcp_session_id=mcp_session_id
+                        mcp_session_id=mcp_session_id,
                     )
-                
-                logger.debug(f"Successfully verified Google token for user: {user_email}")
-                
+
+                logger.debug(
+                    f"Successfully verified Google token for user: {user_email}"
+                )
+
             return access_token
-            
+
         except Exception as e:
             logger.error(f"Failed to verify Google token: {e}")
             return None
-    
+
     def customize_auth_routes(self, routes: List[Route]) -> List[Route]:
         """
-        NOTE: This method is not currently used. All OAuth 2.1 routes are implemented 
+        NOTE: This method is not currently used. All OAuth 2.1 routes are implemented
         directly in core/server.py using @server.custom_route decorators.
-        
+
         This method exists for compatibility with FastMCP's AuthProvider interface
         but the routes it would define are handled elsewhere.
         """
         # Routes are implemented directly in core/server.py
         return routes
-    
+
     def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
         """
         Get session information for credential bridging from OAuth21SessionStore.
-        
+
         Args:
             session_id: The session identifier
-            
+
         Returns:
             Session information if found
         """
         from auth.oauth21_session_store import get_oauth21_session_store
+
         store = get_oauth21_session_store()
-        
+
         # Try to get user by session_id (assuming it's the MCP session ID)
         user_email = store.get_user_by_mcp_session(session_id)
         if user_email:
@@ -145,28 +155,27 @@ def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
                 return {
                     "access_token": credentials.token,
                     "user_email": user_email,
-                    "scopes": credentials.scopes or []
+                    "scopes": credentials.scopes or [],
                 }
         return None
-    
+
     def create_session_from_token(self, token: str, user_email: str) -> str:
         """
         Create a session from an access token for credential bridging using OAuth21SessionStore.
-        
+
         Args:
             token: The access token
             user_email: The user's email address
-            
+
         Returns:
             Session ID
         """
         from auth.oauth21_session_store import get_oauth21_session_store
+
         store = get_oauth21_session_store()
         session_id = f"google_{user_email}"
-        
+
         store.store_session(
-            user_email=user_email,
-            access_token=token,
-            session_id=session_id
+            user_email=user_email, access_token=token, session_id=session_id
         )
-        return session_id
+        return session_id

auth/fastmcp_google_auth.py

@@ -14,7 +14,7 @@
 from google.auth.exceptions import RefreshError
 from googleapiclient.discovery import build
 from googleapiclient.errors import HttpError
-from auth.scopes import SCOPES, get_current_scopes # noqa
+from auth.scopes import SCOPES, get_current_scopes  # noqa
 from auth.oauth21_session_store import get_oauth21_session_store
 from auth.credential_store import get_credential_store
 from auth.oauth_config import get_oauth_config, is_stateless_mode
@@ -135,11 +135,15 @@ def save_credentials_to_session(session_id: str, credentials: Credentials):
             client_secret=credentials.client_secret,
             scopes=credentials.scopes,
             expiry=credentials.expiry,
-            mcp_session_id=session_id
+            mcp_session_id=session_id,
+        )
+        logger.debug(
+            f"Credentials saved to OAuth21SessionStore for session_id: {session_id}, user: {user_email}"
         )
-        logger.debug(f"Credentials saved to OAuth21SessionStore for session_id: {session_id}, user: {user_email}")
     else:
-        logger.warning(f"Could not save credentials to session store - no user email found for session: {session_id}")
+        logger.warning(
+            f"Could not save credentials to session store - no user email found for session: {session_id}"
+        )
 
 
 def load_credentials_from_session(session_id: str) -> Optional[Credentials]:
@@ -477,7 +481,7 @@ def handle_auth_callback(
             scopes=credentials.scopes,
             expiry=credentials.expiry,
             mcp_session_id=session_id,
-            issuer="https://accounts.google.com"  # Add issuer for Google tokens
+            issuer="https://accounts.google.com",  # Add issuer for Google tokens
         )
 
         # If session_id is provided, also save to session cache for compatibility
@@ -521,7 +525,9 @@ def get_credentials(
             # Try to get credentials by MCP session
             credentials = store.get_credentials_by_mcp_session(session_id)
             if credentials:
-                logger.info(f"[get_credentials] Found OAuth 2.1 credentials for MCP session {session_id}")
+                logger.info(
+                    f"[get_credentials] Found OAuth 2.1 credentials for MCP session {session_id}"
+                )
 
                 # Check scopes
                 if not all(scope in credentials.scopes for scope in required_scopes):
@@ -537,7 +543,9 @@ def get_credentials(
                     # Try to refresh
                     try:
                         credentials.refresh(Request())
-                        logger.info(f"[get_credentials] Refreshed OAuth 2.1 credentials for session {session_id}")
+                        logger.info(
+                            f"[get_credentials] Refreshed OAuth 2.1 credentials for session {session_id}"
+                        )
                         # Update stored credentials
                         user_email = store.get_user_by_mcp_session(session_id)
                         if user_email:
@@ -547,11 +555,13 @@ def get_credentials(
                                 refresh_token=credentials.refresh_token,
                                 scopes=credentials.scopes,
                                 expiry=credentials.expiry,
-                                mcp_session_id=session_id
+                                mcp_session_id=session_id,
                             )
                         return credentials
                     except Exception as e:
-                        logger.error(f"[get_credentials] Failed to refresh OAuth 2.1 credentials: {e}")
+                        logger.error(
+                            f"[get_credentials] Failed to refresh OAuth 2.1 credentials: {e}"
+                        )
                         return None
         except ImportError:
             pass  # OAuth 2.1 store not available
@@ -672,7 +682,9 @@ def get_credentials(
                     credential_store = get_credential_store()
                     credential_store.store_credential(user_google_email, credentials)
                 else:
-                    logger.info(f"Skipping credential file save in stateless mode for {user_google_email}")
+                    logger.info(
+                        f"Skipping credential file save in stateless mode for {user_google_email}"
+                    )
 
                 # Also update OAuth21SessionStore
                 store = get_oauth21_session_store()
@@ -686,7 +698,7 @@ def get_credentials(
                     scopes=credentials.scopes,
                     expiry=credentials.expiry,
                     mcp_session_id=session_id,
-                    issuer="https://accounts.google.com"  # Add issuer for Google tokens
+                    issuer="https://accounts.google.com",  # Add issuer for Google tokens
                 )
 
             if session_id:  # Update session cache if it was the source or is active
@@ -775,9 +787,13 @@ async def get_authenticated_google_service(
             # First try context variable (works in async context)
             session_id = get_fastmcp_session_id()
             if session_id:
-                logger.debug(f"[{tool_name}] Got FastMCP session ID from context: {session_id}")
+                logger.debug(
+                    f"[{tool_name}] Got FastMCP session ID from context: {session_id}"
+                )
             else:
-                logger.debug(f"[{tool_name}] Context variable returned None/empty session ID")
+                logger.debug(
+                    f"[{tool_name}] Context variable returned None/empty session ID"
+                )
         except Exception as e:
             logger.debug(
                 f"[{tool_name}] Could not get FastMCP session from context: {e}"
@@ -787,17 +803,25 @@ async def get_authenticated_google_service(
         if not session_id and get_fastmcp_context:
             try:
                 fastmcp_ctx = get_fastmcp_context()
-                if fastmcp_ctx and hasattr(fastmcp_ctx, 'session_id'):
+                if fastmcp_ctx and hasattr(fastmcp_ctx, "session_id"):
                     session_id = fastmcp_ctx.session_id
-                    logger.debug(f"[{tool_name}] Got FastMCP session ID directly: {session_id}")
+                    logger.debug(
+                        f"[{tool_name}] Got FastMCP session ID directly: {session_id}"
+                    )
                 else:
-                    logger.debug(f"[{tool_name}] FastMCP context exists but no session_id attribute")
+                    logger.debug(
+                        f"[{tool_name}] FastMCP context exists but no session_id attribute"
+                    )
             except Exception as e:
-                logger.debug(f"[{tool_name}] Could not get FastMCP context directly: {e}")
+                logger.debug(
+                    f"[{tool_name}] Could not get FastMCP context directly: {e}"
+                )
 
         # Final fallback: log if we still don't have session_id
         if not session_id:
-            logger.warning(f"[{tool_name}] Unable to obtain FastMCP session ID from any source")
+            logger.warning(
+                f"[{tool_name}] Unable to obtain FastMCP session ID from any source"
+            )
 
     logger.info(
         f"[{tool_name}] Attempting to get authenticated {service_name} service. Email: '{user_google_email}', Session: '{session_id}'"
@@ -818,8 +842,12 @@ async def get_authenticated_google_service(
     )
 
     if not credentials or not credentials.valid:
-        logger.warning(f"[{tool_name}] No valid credentials. Email: '{user_google_email}'.")
-        logger.info(f"[{tool_name}] Valid email '{user_google_email}' provided, initiating auth flow.")
+        logger.warning(
+            f"[{tool_name}] No valid credentials. Email: '{user_google_email}'."
+        )
+        logger.info(
+            f"[{tool_name}] Valid email '{user_google_email}' provided, initiating auth flow."
+        )
 
         # Ensure OAuth callback is available
         from auth.oauth_callback_server import ensure_oauth_callback_available

auth/google_auth.py

@@ -61,8 +61,9 @@ def __init__(self):
 
         # Get configuration from OAuth config
         from auth.oauth_config import get_oauth_config
+
         config = get_oauth_config()
-        
+
         self.client_id = config.client_id
         self.client_secret = config.client_secret
         self.base_url = config.get_oauth_base_url()