Add Google Drive connector: mcp_local, OAuth2 (same pattern as Gmail/GCal), 8 tools, verify with Drive API probe, SVG icons
This commit is contained in:
@@ -0,0 +1,581 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Google Drive MCP server (JSON-RPC 2.0 over stdio).
|
||||
|
||||
Tools:
|
||||
status — self-check: credentials, token refresh, API reachability
|
||||
list_files — list files in a folder (or root) with pagination
|
||||
search_files — full-text + name search across Drive
|
||||
get_file_metadata — get detailed file info by ID
|
||||
upload_file — upload a local file to Drive
|
||||
download_file — download a file from Drive to local path
|
||||
delete_file — permanently delete a file
|
||||
create_folder — create a folder in Drive
|
||||
|
||||
Skald mode: Skald injects credentials via DRIVE_CREDS_JSON env var (authorized_user JSON).
|
||||
Standalone mode: reads from GOOGLE_CREDS_PATH or ./secrets/google_creds.json.
|
||||
|
||||
Required OAuth scopes:
|
||||
https://www.googleapis.com/auth/drive (full access — read/write/delete all files)
|
||||
or https://www.googleapis.com/auth/drive.file (access only to files this app creates)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from typing import Any, Callable
|
||||
|
||||
import google.auth
|
||||
from google.auth.transport.requests import Request as GoogleRequest
|
||||
from google.oauth2.credentials import Credentials
|
||||
from googleapiclient.discovery import build
|
||||
from googleapiclient.errors import HttpError
|
||||
from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
|
||||
|
||||
|
||||
# ── Logging ─────────────────────────────────────────────────────────────────
|
||||
def log(msg: str) -> None:
|
||||
print(f"[drive_mcp] {msg}", file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
# ── Auth helpers ────────────────────────────────────────────────────────────
|
||||
_creds: Credentials | None = None
|
||||
_drive_service: Any = None
|
||||
|
||||
|
||||
def _load_credentials() -> Credentials | None:
|
||||
"""Load credentials from env (Skald mode) or file (standalone mode)."""
|
||||
# 1. Skald mode: DRIVE_CREDS_JSON env var (full authorized_user JSON)
|
||||
env_json = os.environ.get("DRIVE_CREDS_JSON", "").strip()
|
||||
if env_json:
|
||||
try:
|
||||
data = json.loads(env_json)
|
||||
creds = Credentials.from_authorized_user_info(data)
|
||||
log("Loaded credentials from DRIVE_CREDS_JSON env var")
|
||||
return creds
|
||||
except Exception as e:
|
||||
log(f"Failed to parse DRIVE_CREDS_JSON: {e}")
|
||||
return None
|
||||
|
||||
# 2. Standalone: GOOGLE_CREDS_PATH or default path
|
||||
path = os.environ.get("GOOGLE_CREDS_PATH", "")
|
||||
if not path:
|
||||
path = "./secrets/google_creds.json"
|
||||
if os.path.isfile(path):
|
||||
try:
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
creds = Credentials.from_authorized_user_info(data)
|
||||
log(f"Loaded credentials from {path}")
|
||||
return creds
|
||||
except Exception as e:
|
||||
log(f"Failed to load creds from {path}: {e}")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _ensure_service() -> tuple[bool, str]:
|
||||
"""Ensure we have valid credentials and a Drive service. Returns (ok, message)."""
|
||||
global _creds, _drive_service
|
||||
if _drive_service is not None:
|
||||
return True, "ok"
|
||||
|
||||
_creds = _load_credentials()
|
||||
if _creds is None:
|
||||
return False, "No Google Drive credentials found. Set DRIVE_CREDS_JSON or configure OAuth."
|
||||
|
||||
# Refresh if expired
|
||||
if not _creds.valid:
|
||||
if _creds.expired and _creds.refresh_token:
|
||||
try:
|
||||
_creds.refresh(GoogleRequest())
|
||||
log("Token refreshed successfully")
|
||||
except Exception as e:
|
||||
return False, f"Token refresh failed: {e}"
|
||||
else:
|
||||
return False, "Credentials expired and no refresh token available."
|
||||
|
||||
try:
|
||||
_drive_service = build("drive", "v3", credentials=_creds)
|
||||
return True, "ok"
|
||||
except Exception as e:
|
||||
return False, f"Failed to build Drive service: {e}"
|
||||
|
||||
|
||||
# ── Tool implementations ────────────────────────────────────────────────────
|
||||
|
||||
DRIVE_FILE_FIELDS = "id, name, mimeType, size, createdTime, modifiedTime, parents, webViewLink, description, ownedByMe, owners(emailAddress), thumbnailLink"
|
||||
|
||||
|
||||
def _list_files(folder_id: str | None, page_size: int, page_token: str | None) -> dict:
|
||||
"""List files in a folder."""
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {"error": msg}
|
||||
|
||||
q = f"'{folder_id}' in parents and trashed=false" if folder_id else "trashed=false"
|
||||
try:
|
||||
results = _drive_service.files().list(
|
||||
q=q,
|
||||
pageSize=min(page_size, 100),
|
||||
pageToken=page_token or "",
|
||||
fields=f"files({DRIVE_FILE_FIELDS}), nextPageToken",
|
||||
orderBy="folder,modifiedTime desc,name",
|
||||
).execute()
|
||||
return {
|
||||
"files": results.get("files", []),
|
||||
"next_page_token": results.get("nextPageToken", None),
|
||||
}
|
||||
except HttpError as e:
|
||||
return {"error": f"Drive API error: {e.reason or e}"}
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _search_files(query: str, max_results: int) -> dict:
|
||||
"""Search files across Drive."""
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {"error": msg}
|
||||
|
||||
q = f"name contains '{query}' and trashed=false"
|
||||
try:
|
||||
results = _drive_service.files().list(
|
||||
q=q,
|
||||
pageSize=min(max_results, 100),
|
||||
fields=f"files({DRIVE_FILE_FIELDS})",
|
||||
orderBy="modifiedTime desc",
|
||||
).execute()
|
||||
return {"files": results.get("files", [])}
|
||||
except HttpError as e:
|
||||
return {"error": f"Drive API error: {e.reason or e}"}
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _get_file_metadata(file_id: str) -> dict:
|
||||
"""Get detailed metadata for a specific file."""
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {"error": msg}
|
||||
|
||||
try:
|
||||
file = _drive_service.files().get(
|
||||
fileId=file_id,
|
||||
fields=DRIVE_FILE_FIELDS,
|
||||
supportsAllDrives=True,
|
||||
).execute()
|
||||
return {"file": file}
|
||||
except HttpError as e:
|
||||
return {"error": f"Drive API error: {e.reason or e}"}
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _upload_file(local_path: str, parent_id: str | None, name: str | None) -> dict:
|
||||
"""Upload a local file to Drive."""
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {"error": msg}
|
||||
|
||||
resolved_path = os.path.expanduser(local_path)
|
||||
if not os.path.isfile(resolved_path):
|
||||
return {"error": f"File not found: {local_path}"}
|
||||
|
||||
file_name = name or os.path.basename(resolved_path)
|
||||
mime_type = mimetypes.guess_type(resolved_path)[0] or "application/octet-stream"
|
||||
|
||||
body = {"name": file_name, "mimeType": mime_type}
|
||||
if parent_id:
|
||||
body["parents"] = [parent_id]
|
||||
|
||||
try:
|
||||
media = MediaFileUpload(resolved_path, mimetype=mime_type, resumable=True)
|
||||
file = _drive_service.files().create(
|
||||
body=body,
|
||||
media_body=media,
|
||||
fields=DRIVE_FILE_FIELDS,
|
||||
supportsAllDrives=True,
|
||||
).execute()
|
||||
return {"file": file}
|
||||
except HttpError as e:
|
||||
return {"error": f"Upload failed: {e.reason or e}"}
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _download_file(file_id: str, output_path: str | None) -> dict:
|
||||
"""Download a file from Drive. If no output_path given, returns base64 content."""
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {"error": msg}
|
||||
|
||||
try:
|
||||
# Get file metadata first
|
||||
file = _drive_service.files().get(fileId=file_id, fields="id, name, mimeType, size").execute()
|
||||
file_name = file.get("name", "unknown")
|
||||
mime_type = file.get("mimeType", "application/octet-stream")
|
||||
|
||||
request = _drive_service.files().get_media(fileId=file_id)
|
||||
|
||||
if output_path:
|
||||
resolved = os.path.expanduser(output_path)
|
||||
os.makedirs(os.path.dirname(os.path.abspath(resolved)) or ".", exist_ok=True)
|
||||
with io.FileIO(resolved, "wb") as fh:
|
||||
downloader = MediaIoBaseDownload(fh, request)
|
||||
done = False
|
||||
while not done:
|
||||
_, done = downloader.next_chunk()
|
||||
file_size = os.path.getsize(resolved)
|
||||
return {
|
||||
"file": {"id": file_id, "name": file_name, "mimeType": mime_type, "size": file_size},
|
||||
"saved_to": resolved,
|
||||
}
|
||||
else:
|
||||
# Return base64 content inline
|
||||
buf = io.BytesIO()
|
||||
downloader = MediaIoBaseDownload(buf, request)
|
||||
done = False
|
||||
while not done:
|
||||
_, done = downloader.next_chunk()
|
||||
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
return {
|
||||
"file": {"id": file_id, "name": file_name, "mimeType": mime_type},
|
||||
"content_base64": b64,
|
||||
"content_length": len(buf.getvalue()),
|
||||
}
|
||||
except HttpError as e:
|
||||
return {"error": f"Download failed: {e.reason or e}"}
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _delete_file(file_id: str) -> dict:
|
||||
"""Permanently delete a file."""
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {"error": msg}
|
||||
|
||||
try:
|
||||
_drive_service.files().delete(fileId=file_id, supportsAllDrives=True).execute()
|
||||
return {"deleted": True, "file_id": file_id}
|
||||
except HttpError as e:
|
||||
return {"error": f"Delete failed: {e.reason or e}"}
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _create_folder(name: str, parent_id: str | None) -> dict:
|
||||
"""Create a folder in Drive."""
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {"error": msg}
|
||||
|
||||
body = {"name": name, "mimeType": "application/vnd.google-apps.folder"}
|
||||
if parent_id:
|
||||
body["parents"] = [parent_id]
|
||||
|
||||
try:
|
||||
folder = _drive_service.files().create(
|
||||
body=body,
|
||||
fields=DRIVE_FILE_FIELDS,
|
||||
supportsAllDrives=True,
|
||||
).execute()
|
||||
return {"folder": folder}
|
||||
except HttpError as e:
|
||||
return {"error": f"Create folder failed: {e.reason or e}"}
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _status() -> dict:
|
||||
"""Health check: credentials, token expiry, API reachability."""
|
||||
global _creds, _drive_service
|
||||
# Reset so we re-check from scratch
|
||||
_drive_service = None
|
||||
ok, msg = _ensure_service()
|
||||
if not ok:
|
||||
return {
|
||||
"ok": False,
|
||||
"message": msg,
|
||||
}
|
||||
expiry = _creds.expiry.isoformat() if _creds and _creds.expiry else "unknown"
|
||||
scopes = _creds.scopes if _creds else []
|
||||
return {
|
||||
"ok": True,
|
||||
"message": "Google Drive API is reachable",
|
||||
"details": {
|
||||
"expiry": expiry,
|
||||
"scopes": scopes,
|
||||
"token_valid": _creds.valid if _creds else False,
|
||||
"has_refresh_token": bool(_creds and _creds.refresh_token),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ── Tool registry ──────────────────────────────────────────────────────────
|
||||
TOOLS: dict[str, tuple[Callable, str, dict]] = {
|
||||
"status": (
|
||||
_status,
|
||||
"Check if the Drive API is reachable and credentials are valid.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
},
|
||||
),
|
||||
"list_files": (
|
||||
_list_files,
|
||||
"List files in a folder (or root if no folder_id given). Returns paginated results.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"folder_id": {
|
||||
"type": "string",
|
||||
"description": "Folder ID (omit or empty for root). Pass 'root' for My Drive root.",
|
||||
"default": None,
|
||||
},
|
||||
"page_size": {
|
||||
"type": "integer",
|
||||
"description": "Results per page (max 100, default 50)",
|
||||
"default": 50,
|
||||
},
|
||||
"page_token": {
|
||||
"type": "string",
|
||||
"description": "Token for the next page (from previous list_files response)",
|
||||
"default": None,
|
||||
},
|
||||
},
|
||||
"required": [],
|
||||
},
|
||||
),
|
||||
"search_files": (
|
||||
_search_files,
|
||||
"Search files across Drive by name or content.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Search query (matched against file name)",
|
||||
},
|
||||
"max_results": {
|
||||
"type": "integer",
|
||||
"description": "Max results (default 20, max 100)",
|
||||
"default": 20,
|
||||
},
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
),
|
||||
"get_file_metadata": (
|
||||
_get_file_metadata,
|
||||
"Get detailed metadata for a specific file by ID.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_id": {
|
||||
"type": "string",
|
||||
"description": "Google Drive file ID (from URL or search results)",
|
||||
},
|
||||
},
|
||||
"required": ["file_id"],
|
||||
},
|
||||
),
|
||||
"upload_file": (
|
||||
_upload_file,
|
||||
"Upload a local file to Google Drive.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"local_path": {
|
||||
"type": "string",
|
||||
"description": "Local path to the file to upload",
|
||||
},
|
||||
"parent_id": {
|
||||
"type": "string",
|
||||
"description": "Optional folder ID to upload into (omit for root)",
|
||||
"default": None,
|
||||
},
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "Optional custom name for the file on Drive (defaults to filename)",
|
||||
"default": None,
|
||||
},
|
||||
},
|
||||
"required": ["local_path"],
|
||||
},
|
||||
),
|
||||
"download_file": (
|
||||
_download_file,
|
||||
"Download a file from Drive. Returns base64 content if no output_path, or saves to disk.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_id": {
|
||||
"type": "string",
|
||||
"description": "Google Drive file ID to download",
|
||||
},
|
||||
"output_path": {
|
||||
"type": "string",
|
||||
"description": "Optional local path to save the file (omit to get base64 content)",
|
||||
"default": None,
|
||||
},
|
||||
},
|
||||
"required": ["file_id"],
|
||||
},
|
||||
),
|
||||
"delete_file": (
|
||||
_delete_file,
|
||||
"Permanently delete a file from Google Drive.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_id": {
|
||||
"type": "string",
|
||||
"description": "Google Drive file ID to delete",
|
||||
},
|
||||
},
|
||||
"required": ["file_id"],
|
||||
},
|
||||
),
|
||||
"create_folder": (
|
||||
_create_folder,
|
||||
"Create a new folder in Google Drive.",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "Folder name",
|
||||
},
|
||||
"parent_id": {
|
||||
"type": "string",
|
||||
"description": "Optional parent folder ID (omit for My Drive root)",
|
||||
"default": None,
|
||||
},
|
||||
},
|
||||
"required": ["name"],
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
# ── JSON-RPC handler ───────────────────────────────────────────────────────
|
||||
|
||||
def _send_response(id_val: int | str | None, result: dict | None, error: dict | None = None) -> None:
|
||||
msg = {"jsonrpc": "2.0", "id": id_val}
|
||||
if error:
|
||||
msg["error"] = error
|
||||
else:
|
||||
msg["result"] = result
|
||||
sys.stdout.write(json.dumps(msg) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def _handle_request(req: dict) -> None:
|
||||
rid = req.get("id")
|
||||
method = req.get("method", "")
|
||||
|
||||
# JSON-RPC initialize
|
||||
if method == "initialize":
|
||||
_send_response(rid, {
|
||||
"protocolVersion": "2025-03-26",
|
||||
"capabilities": {
|
||||
"tools": {},
|
||||
"roots": {"listChanged": False},
|
||||
},
|
||||
"serverInfo": {"name": "google-drive-mcp", "version": "1.0.0"},
|
||||
})
|
||||
return
|
||||
|
||||
# List available tools
|
||||
if method == "tools/list":
|
||||
tool_list = []
|
||||
for name, (fn, desc, schema) in sorted(TOOLS.items()):
|
||||
tool_list.append({
|
||||
"name": name,
|
||||
"description": desc,
|
||||
"inputSchema": schema,
|
||||
})
|
||||
_send_response(rid, {"tools": tool_list})
|
||||
return
|
||||
|
||||
# Call a tool
|
||||
if method == "tools/call":
|
||||
tool_name = req.get("params", {}).get("name", "")
|
||||
args = req.get("params", {}).get("arguments", {})
|
||||
|
||||
if tool_name not in TOOLS:
|
||||
_send_response(rid, None, {
|
||||
"code": -32601,
|
||||
"message": f"Unknown tool: {tool_name}",
|
||||
})
|
||||
return
|
||||
|
||||
try:
|
||||
fn, _, schema = TOOLS[tool_name]
|
||||
# Extract params matching the schema
|
||||
kwargs = {}
|
||||
for prop_name in schema.get("properties", {}):
|
||||
if prop_name in args:
|
||||
kwargs[prop_name] = args[prop_name]
|
||||
|
||||
result = fn(**kwargs)
|
||||
_send_response(rid, result)
|
||||
except Exception as e:
|
||||
log(f"Error in {tool_name}: {traceback.format_exc()}")
|
||||
_send_response(rid, None, {
|
||||
"code": -32603,
|
||||
"message": f"Internal error: {e}",
|
||||
})
|
||||
return
|
||||
|
||||
# Notifications (no id) — ignore
|
||||
if rid is None:
|
||||
return
|
||||
|
||||
_send_response(rid, None, {
|
||||
"code": -32601,
|
||||
"message": f"Method not found: {method}",
|
||||
})
|
||||
|
||||
|
||||
# ── Main loop ──────────────────────────────────────────────────────────────
|
||||
def main() -> None:
|
||||
log("Google Drive MCP server starting (JSON-RPC 2.0 over stdio)")
|
||||
|
||||
# Announce server info on stderr
|
||||
log(f"Python: {sys.version}")
|
||||
log(f"Tools: {', '.join(sorted(TOOLS.keys()))}")
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
req = json.loads(line)
|
||||
_handle_request(req)
|
||||
except json.JSONDecodeError as e:
|
||||
log(f"Invalid JSON: {e}")
|
||||
# Only respond if there's an id we can extract
|
||||
try:
|
||||
partial = json.loads(line[:line.rfind("}")+1] if "}" in line else "{}")
|
||||
rid = partial.get("id") if isinstance(partial, dict) else None
|
||||
except Exception:
|
||||
rid = None
|
||||
if rid is not None:
|
||||
_send_response(rid, None, {
|
||||
"code": -32700,
|
||||
"message": f"Parse error: {e}",
|
||||
})
|
||||
except Exception as e:
|
||||
log(f"Fatal error: {traceback.format_exc()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user