diff --git a/auth/oauth_callback_server.py b/auth/oauth_callback_server.py index 27e804cd..f6f55eca 100644 --- a/auth/oauth_callback_server.py +++ b/auth/oauth_callback_server.py @@ -13,6 +13,7 @@ import uvicorn from fastapi import FastAPI, Request +from fastapi.responses import FileResponse, JSONResponse from typing import Optional from urllib.parse import urlparse @@ -39,6 +40,8 @@ def __init__(self, port: int = 8000, base_uri: str = "http://localhost"): # Setup the callback route self._setup_callback_route() + # Setup attachment serving route + self._setup_attachment_route() def _setup_callback_route(self): """Setup the OAuth callback route.""" @@ -89,6 +92,35 @@ async def oauth_callback(request: Request): logger.error(error_message_detail, exc_info=True) return create_server_error_response(str(e)) + def _setup_attachment_route(self): + """Setup the attachment serving route.""" + from core.attachment_storage import get_attachment_storage + + @self.app.get("/attachments/{file_id}") + async def serve_attachment(file_id: str, request: Request): + """Serve a stored attachment file.""" + storage = get_attachment_storage() + metadata = storage.get_attachment_metadata(file_id) + + if not metadata: + return JSONResponse( + {"error": "Attachment not found or expired"}, + status_code=404 + ) + + file_path = storage.get_attachment_path(file_id) + if not file_path: + return JSONResponse( + {"error": "Attachment file not found"}, + status_code=404 + ) + + return FileResponse( + path=str(file_path), + filename=metadata["filename"], + media_type=metadata["mime_type"] + ) + def start(self) -> tuple[bool, str]: """ Start the minimal OAuth server. diff --git a/core/attachment_storage.py b/core/attachment_storage.py new file mode 100644 index 00000000..9fbffa87 --- /dev/null +++ b/core/attachment_storage.py @@ -0,0 +1,216 @@ +""" +Temporary attachment storage for Gmail attachments. + +Stores attachments in ./tmp directory and provides HTTP URLs for access. +Files are automatically cleaned up after expiration (default 1 hour). +""" + +import base64 +import logging +import uuid +from pathlib import Path +from typing import Optional, Dict +from datetime import datetime, timedelta + +logger = logging.getLogger(__name__) + +# Default expiration: 1 hour +DEFAULT_EXPIRATION_SECONDS = 3600 + +# Storage directory +STORAGE_DIR = Path("./tmp/attachments") +STORAGE_DIR.mkdir(parents=True, exist_ok=True) + + +class AttachmentStorage: + """Manages temporary storage of email attachments.""" + + def __init__(self, expiration_seconds: int = DEFAULT_EXPIRATION_SECONDS): + self.expiration_seconds = expiration_seconds + self._metadata: Dict[str, Dict] = {} + + def save_attachment( + self, + base64_data: str, + filename: Optional[str] = None, + mime_type: Optional[str] = None, + ) -> str: + """ + Save an attachment and return a unique file ID. + + Args: + base64_data: Base64-encoded attachment data + filename: Original filename (optional) + mime_type: MIME type (optional) + + Returns: + Unique file ID (UUID string) + """ + # Generate unique file ID + file_id = str(uuid.uuid4()) + + # Decode base64 data + try: + file_bytes = base64.urlsafe_b64decode(base64_data) + except Exception as e: + logger.error(f"Failed to decode base64 attachment data: {e}") + raise ValueError(f"Invalid base64 data: {e}") + + # Determine file extension from filename or mime type + extension = "" + if filename: + extension = Path(filename).suffix + elif mime_type: + # Basic mime type to extension mapping + mime_to_ext = { + "image/jpeg": ".jpg", + "image/png": ".png", + "image/gif": ".gif", + "application/pdf": ".pdf", + "application/zip": ".zip", + "text/plain": ".txt", + "text/html": ".html", + } + extension = mime_to_ext.get(mime_type, "") + + # Save file + file_path = STORAGE_DIR / f"{file_id}{extension}" + try: + file_path.write_bytes(file_bytes) + logger.info(f"Saved attachment {file_id} ({len(file_bytes)} bytes) to {file_path}") + except Exception as e: + logger.error(f"Failed to save attachment to {file_path}: {e}") + raise + + # Store metadata + expires_at = datetime.now() + timedelta(seconds=self.expiration_seconds) + self._metadata[file_id] = { + "file_path": str(file_path), + "filename": filename or f"attachment{extension}", + "mime_type": mime_type or "application/octet-stream", + "size": len(file_bytes), + "created_at": datetime.now(), + "expires_at": expires_at, + } + + return file_id + + def get_attachment_path(self, file_id: str) -> Optional[Path]: + """ + Get the file path for an attachment ID. + + Args: + file_id: Unique file ID + + Returns: + Path object if file exists and not expired, None otherwise + """ + if file_id not in self._metadata: + logger.warning(f"Attachment {file_id} not found in metadata") + return None + + metadata = self._metadata[file_id] + file_path = Path(metadata["file_path"]) + + # Check if expired + if datetime.now() > metadata["expires_at"]: + logger.info(f"Attachment {file_id} has expired, cleaning up") + self._cleanup_file(file_id) + return None + + # Check if file exists + if not file_path.exists(): + logger.warning(f"Attachment file {file_path} does not exist") + del self._metadata[file_id] + return None + + return file_path + + def get_attachment_metadata(self, file_id: str) -> Optional[Dict]: + """ + Get metadata for an attachment. + + Args: + file_id: Unique file ID + + Returns: + Metadata dict if exists and not expired, None otherwise + """ + if file_id not in self._metadata: + return None + + metadata = self._metadata[file_id].copy() + + # Check if expired + if datetime.now() > metadata["expires_at"]: + self._cleanup_file(file_id) + return None + + return metadata + + def _cleanup_file(self, file_id: str) -> None: + """Remove file and metadata.""" + if file_id in self._metadata: + file_path = Path(self._metadata[file_id]["file_path"]) + try: + if file_path.exists(): + file_path.unlink() + logger.debug(f"Deleted expired attachment file: {file_path}") + except Exception as e: + logger.warning(f"Failed to delete attachment file {file_path}: {e}") + del self._metadata[file_id] + + def cleanup_expired(self) -> int: + """ + Clean up expired attachments. + + Returns: + Number of files cleaned up + """ + now = datetime.now() + expired_ids = [ + file_id + for file_id, metadata in self._metadata.items() + if now > metadata["expires_at"] + ] + + for file_id in expired_ids: + self._cleanup_file(file_id) + + return len(expired_ids) + + +# Global instance +_attachment_storage: Optional[AttachmentStorage] = None + + +def get_attachment_storage() -> AttachmentStorage: + """Get the global attachment storage instance.""" + global _attachment_storage + if _attachment_storage is None: + _attachment_storage = AttachmentStorage() + return _attachment_storage + + +def get_attachment_url(file_id: str) -> str: + """ + Generate a URL for accessing an attachment. + + Args: + file_id: Unique file ID + + Returns: + Full URL to access the attachment + """ + import os + from core.config import WORKSPACE_MCP_PORT, WORKSPACE_MCP_BASE_URI + + # Use external URL if set (for reverse proxy scenarios) + external_url = os.getenv("WORKSPACE_EXTERNAL_URL") + if external_url: + base_url = external_url.rstrip("/") + else: + base_url = f"{WORKSPACE_MCP_BASE_URI}:{WORKSPACE_MCP_PORT}" + + return f"{base_url}/attachments/{file_id}" + diff --git a/core/server.py b/core/server.py index 86efb2a7..121a80ed 100644 --- a/core/server.py +++ b/core/server.py @@ -2,7 +2,7 @@ from typing import List, Optional from importlib import metadata -from fastapi.responses import HTMLResponse, JSONResponse +from fastapi.responses import HTMLResponse, JSONResponse, FileResponse from starlette.applications import Starlette from starlette.requests import Request from starlette.middleware import Middleware @@ -156,6 +156,33 @@ async def health_check(request: Request): "transport": get_transport_mode() }) +@server.custom_route("/attachments/{file_id}", methods=["GET"]) +async def serve_attachment(file_id: str, request: Request): + """Serve a stored attachment file.""" + from core.attachment_storage import get_attachment_storage + + storage = get_attachment_storage() + metadata = storage.get_attachment_metadata(file_id) + + if not metadata: + return JSONResponse( + {"error": "Attachment not found or expired"}, + status_code=404 + ) + + file_path = storage.get_attachment_path(file_id) + if not file_path: + return JSONResponse( + {"error": "Attachment file not found"}, + status_code=404 + ) + + return FileResponse( + path=str(file_path), + filename=metadata["filename"], + media_type=metadata["mime_type"] + ) + async def legacy_oauth2_callback(request: Request) -> HTMLResponse: state = request.query_params.get("state") code = request.query_params.get("code") diff --git a/gmail/gmail_tools.py b/gmail/gmail_tools.py index 5c69a3ef..ee69ceaf 100644 --- a/gmail/gmail_tools.py +++ b/gmail/gmail_tools.py @@ -688,20 +688,89 @@ async def get_gmail_attachment_content( # Format response with attachment data size_bytes = attachment.get('size', 0) size_kb = size_bytes / 1024 if size_bytes else 0 + base64_data = attachment.get('data', '') + + # Check if we're in stateless mode (can't save files) + from auth.oauth_config import is_stateless_mode + if is_stateless_mode(): + result_lines = [ + "Attachment downloaded successfully!", + f"Message ID: {message_id}", + f"Size: {size_kb:.1f} KB ({size_bytes} bytes)", + "\n⚠️ Stateless mode: File storage disabled.", + "\nBase64-encoded content (first 100 characters shown):", + f"{base64_data[:100]}...", + "\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch." + ] + logger.info(f"[get_gmail_attachment_content] Successfully downloaded {size_kb:.1f} KB attachment (stateless mode)") + return "\n".join(result_lines) - result_lines = [ - "Attachment downloaded successfully!", - f"Message ID: {message_id}", - f"Size: {size_kb:.1f} KB ({size_bytes} bytes)", - "\nBase64-encoded content (first 100 characters shown):", - f"{attachment['data'][:100]}...", - "\n\nThe full base64-encoded attachment data is available.", - "To save: decode the base64 data and write to a file with the appropriate extension.", - "\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch." - ] - - logger.info(f"[get_gmail_attachment_content] Successfully downloaded {size_kb:.1f} KB attachment") - return "\n".join(result_lines) + # Save attachment and generate URL + try: + from core.attachment_storage import get_attachment_storage, get_attachment_url + + storage = get_attachment_storage() + + # Try to get filename and mime type from message (optional - attachment IDs are ephemeral) + filename = None + mime_type = None + try: + # Quick metadata fetch to try to get attachment info + # Note: This might fail if attachment IDs changed, but worth trying + message_metadata = await asyncio.to_thread( + service.users() + .messages() + .get(userId="me", id=message_id, format="metadata") + .execute + ) + payload = message_metadata.get("payload", {}) + attachments = _extract_attachments(payload) + for att in attachments: + if att.get("attachmentId") == attachment_id: + filename = att.get("filename") + mime_type = att.get("mimeType") + break + except Exception: + # If we can't get metadata, use defaults + logger.debug(f"Could not fetch attachment metadata for {attachment_id}, using defaults") + + # Save attachment + file_id = storage.save_attachment( + base64_data=base64_data, + filename=filename, + mime_type=mime_type + ) + + # Generate URL + attachment_url = get_attachment_url(file_id) + + result_lines = [ + "Attachment downloaded successfully!", + f"Message ID: {message_id}", + f"Size: {size_kb:.1f} KB ({size_bytes} bytes)", + f"\n📎 Download URL: {attachment_url}", + "\nThe attachment has been saved and is available at the URL above.", + "The file will expire after 1 hour.", + "\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch." + ] + + logger.info(f"[get_gmail_attachment_content] Successfully saved {size_kb:.1f} KB attachment as {file_id}") + return "\n".join(result_lines) + + except Exception as e: + logger.error(f"[get_gmail_attachment_content] Failed to save attachment: {e}", exc_info=True) + # Fallback to showing base64 preview + result_lines = [ + "Attachment downloaded successfully!", + f"Message ID: {message_id}", + f"Size: {size_kb:.1f} KB ({size_bytes} bytes)", + "\n⚠️ Failed to save attachment file. Showing preview instead.", + "\nBase64-encoded content (first 100 characters shown):", + f"{base64_data[:100]}...", + f"\nError: {str(e)}", + "\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch." + ] + return "\n".join(result_lines) @server.tool()