jason/google-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #189 from taylorwilsdon/WORKSPACE_MCP_STATELESS_MODE

Browse Source 
feat: Stateless Mode for Container Deployments & user_google_email removal
This commit is contained in:
Taylor Wilsdon
2025-08-23 13:05:19 -04:00
committed by GitHub
parent 0ab155e10e 38f445f35b
commit aeb0302d71
10 changed files with 388 additions and 443 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
README.md
View File

@@ -170,6 +170,7 @@ uv run main.py --tools gmail drive
| `GOOGLE_PSE_API_KEY` | API key for Custom Search |
| `GOOGLE_PSE_ENGINE_ID` | Search Engine ID for Custom Search |
| `MCP_ENABLE_OAUTH21` | Set to `true` for OAuth 2.1 support |
| `WORKSPACE_MCP_STATELESS_MODE` | Set to `true` for stateless operation (requires OAuth 2.1) |
</td></tr>
</table>
@@ -940,6 +941,36 @@ This architecture enables any OAuth 2.1 compliant client to authenticate users t
</details>
### Stateless Mode (Container-Friendly)
The server supports a stateless mode designed for containerized environments where file system writes should be avoided:
**Enabling Stateless Mode:**
```bash
# Stateless mode requires OAuth 2.1 to be enabled
export MCP_ENABLE_OAUTH21=true
export WORKSPACE_MCP_STATELESS_MODE=true
uv run main.py --transport streamable-http
```
**Key Features:**
- **No file system writes**: Credentials are never written to disk
- **No debug logs**: File-based logging is completely disabled
- **Memory-only sessions**: All tokens stored in memory via OAuth 2.1 session store
- **Container-ready**: Perfect for Docker, Kubernetes, and serverless deployments
- **Token per request**: Each request must include a valid Bearer token
**Requirements:**
- Must be used with `MCP_ENABLE_OAUTH21=true`
- Incompatible with single-user mode
- Clients must handle OAuth flow and send valid tokens with each request
This mode is ideal for:
- Cloud deployments where persistent storage is unavailable
- Multi-tenant environments requiring strict isolation
- Containerized applications with read-only filesystems
- Serverless functions and ephemeral compute environments
**MCP Inspector**: No additional configuration needed with desktop OAuth client.
**Claude Code Inspector**: No additional configuration needed with desktop OAuth client.

12
auth/google_auth.py
View File

@@ -17,7 +17,7 @@ from googleapiclient.errors import HttpError
from auth.scopes import SCOPES
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
from auth.oauth_config import get_oauth_config, is_stateless_mode
from core.config import (
get_transport_mode,
get_oauth_redirect_uri,
@@ -603,11 +603,16 @@ def get_credentials(
)
if not credentials and user_google_email:
if not is_stateless_mode():
logger.debug(
f"[get_credentials] No session credentials, trying credential store for user_google_email '{user_google_email}'."
)
store = get_credential_store()
credentials = store.get_credential(user_google_email)
else:
logger.debug(
f"[get_credentials] No session credentials, skipping file store in stateless mode for user_google_email '{user_google_email}'."
)
if credentials and session_id:
logger.debug(
@@ -661,10 +666,13 @@ def get_credentials(
f"[get_credentials] Credentials refreshed successfully. User: '{user_google_email}', Session: '{session_id}'"
)
# Save refreshed credentials
# Save refreshed credentials (skip file save in stateless mode)
if user_google_email: # Always save to credential store if email is known
if not is_stateless_mode():
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}")
# Also update OAuth21SessionStore
store = get_oauth21_session_store()

5
auth/oauth21_session_store.py
View File

@@ -17,11 +17,6 @@ from google.oauth2.credentials import Credentials
logger = logging.getLogger(__name__)
# =============================================================================
# Session Context Management (absorbed from session_context.py)
# =============================================================================
# Context variable to store the current session information
_current_session_context: contextvars.ContextVar[Optional['SessionContext']] = contextvars.ContextVar(
'current_session_context',

7
auth/oauth_common_handlers.py
View File

@@ -16,7 +16,7 @@ from google.oauth2.credentials import Credentials
from auth.oauth21_session_store import store_token_session
from auth.google_auth import get_credential_store
from auth.scopes import get_current_scopes
from auth.oauth_config import get_oauth_config
from auth.oauth_config import get_oauth_config, is_stateless_mode
from auth.oauth_error_handling import (
OAuthError, OAuthValidationError, OAuthConfigurationError,
create_oauth_error_response, validate_token_request,
@@ -180,12 +180,15 @@ async def handle_proxy_token_exchange(request: Request):
expiry=expiry
)
# Save credentials to file for legacy auth
# Save credentials to file for legacy auth (skip in stateless mode)
if not is_stateless_mode():
store = get_credential_store()
if not store.store_credential(user_email, credentials):
logger.error(f"Failed to save Google credentials for {user_email}")
else:
logger.info(f"Saved Google credentials for {user_email}")
else:
logger.info(f"Skipping credential file save in stateless mode for {user_email}")
except jwt.ExpiredSignatureError:
logger.error("ID token has expired - cannot extract user email")
except jwt.InvalidTokenError as e:

10
auth/oauth_config.py
View File

@@ -39,6 +39,11 @@ class OAuthConfig:
self.pkce_required = self.oauth21_enabled # PKCE is mandatory in OAuth 2.1
self.supported_code_challenge_methods = ["S256", "plain"] if not self.oauth21_enabled else ["S256"]
# Stateless mode configuration
self.stateless_mode = os.getenv("WORKSPACE_MCP_STATELESS_MODE", "false").lower() == "true"
if self.stateless_mode and not self.oauth21_enabled:
raise ValueError("WORKSPACE_MCP_STATELESS_MODE requires MCP_ENABLE_OAUTH21=true")
# Transport mode (will be set at runtime)
self._transport_mode = "stdio" # Default
@@ -340,3 +345,8 @@ def is_oauth21_enabled() -> bool:
def get_oauth_redirect_uri() -> str:
"""Get the primary OAuth redirect URI."""
return get_oauth_config().redirect_uri
def is_stateless_mode() -> bool:
"""Check if stateless mode is enabled."""
return get_oauth_config().stateless_mode

149
auth/service_decorator.py
View File

@@ -1,5 +1,6 @@
import inspect
import logging
import re
from functools import wraps
from typing import Dict, List, Optional, Any, Callable, Union, Tuple
@@ -38,9 +39,7 @@ from auth.scopes import (
CUSTOM_SEARCH_SCOPE,
)
# OAuth 2.1 integration is now handled by FastMCP auth
OAUTH21_INTEGRATION_AVAILABLE = True
logger = logging.getLogger(__name__)
# Authentication helper functions
def _get_auth_context(
@@ -236,7 +235,56 @@ async def get_authenticated_google_service_oauth21(
return service, user_google_email
logger = logging.getLogger(__name__)
def _extract_oauth21_user_email(authenticated_user: Optional[str], func_name: str) -> str:
"""
Extract user email for OAuth 2.1 mode.
Args:
authenticated_user: The authenticated user from context
func_name: Name of the function being decorated (for error messages)
Returns:
User email string
Raises:
Exception: If no authenticated user found in OAuth 2.1 mode
"""
if not authenticated_user:
raise Exception(
f"OAuth 2.1 mode requires an authenticated user for {func_name}, but none was found."
)
return authenticated_user
def _extract_oauth20_user_email(
args: tuple,
kwargs: dict,
wrapper_sig: inspect.Signature
) -> str:
"""
Extract user email for OAuth 2.0 mode from function arguments.
Args:
args: Positional arguments passed to wrapper
kwargs: Keyword arguments passed to wrapper
wrapper_sig: Function signature for parameter binding
Returns:
User email string
Raises:
Exception: If user_google_email parameter not found
"""
bound_args = wrapper_sig.bind(*args, **kwargs)
bound_args.apply_defaults()
user_google_email = bound_args.arguments.get("user_google_email")
if not user_google_email:
raise Exception(
"'user_google_email' parameter is required but was not found."
)
return user_google_email
def _remove_user_email_arg_from_docstring(docstring: str) -> str:
@@ -424,7 +472,16 @@ def require_google_service(
)
# Create a new signature for the wrapper that excludes the 'service' parameter.
# This is the signature that FastMCP will see.
# In OAuth 2.1 mode, also exclude 'user_google_email' since it's automatically determined.
if is_oauth21_enabled():
# Remove both 'service' and 'user_google_email' parameters
filtered_params = [
p for p in params[1:]
if p.name != 'user_google_email'
]
wrapper_sig = original_sig.replace(parameters=filtered_params)
else:
# Only remove 'service' parameter for OAuth 2.0 mode
wrapper_sig = original_sig.replace(parameters=params[1:])
@wraps(func)
@@ -432,17 +489,17 @@ def require_google_service(
# Note: `args` and `kwargs` are now the arguments for the *wrapper*,
# which does not include 'service'.
# Extract user_google_email from the arguments passed to the wrapper
bound_args = wrapper_sig.bind(*args, **kwargs)
bound_args.apply_defaults()
user_google_email = bound_args.arguments.get("user_google_email")
if not user_google_email:
# This should ideally not be reached if 'user_google_email' is a required parameter
raise Exception(
"'user_google_email' parameter is required but was not found."
# Get authentication context early to determine OAuth mode
authenticated_user, auth_method, mcp_session_id = _get_auth_context(
func.__name__
)
# Extract user_google_email based on OAuth mode
if is_oauth21_enabled():
user_google_email = _extract_oauth21_user_email(authenticated_user, func.__name__)
else:
user_google_email = _extract_oauth20_user_email(args, kwargs, wrapper_sig)
# Get service configuration from the decorator's arguments
if service_type not in SERVICE_CONFIGS:
raise Exception(f"Unknown service type: {service_type}")
@@ -457,11 +514,6 @@ def require_google_service(
try:
tool_name = func.__name__
# Get authentication context
authenticated_user, auth_method, mcp_session_id = _get_auth_context(
tool_name
)
# Log authentication status
logger.debug(
f"[{tool_name}] Auth: {authenticated_user or 'none'} via {auth_method or 'none'} (session: {mcp_session_id[:8] if mcp_session_id else 'none'})"
@@ -472,7 +524,9 @@ def require_google_service(
authenticated_user, mcp_session_id, tool_name
)
# Override user_google_email with authenticated user when using OAuth 2.1
# In OAuth 2.1 mode, user_google_email is already set to authenticated_user
# In OAuth 2.0 mode, we may need to override it
if not is_oauth21_enabled():
wrapper_params = list(wrapper_sig.parameters.keys())
user_google_email, args = _override_oauth21_user_email(
use_oauth21,
@@ -484,10 +538,6 @@ def require_google_service(
tool_name,
)
# Update bound_args for consistency
if use_oauth21 and authenticated_user and user_google_email == authenticated_user:
bound_args.arguments["user_google_email"] = authenticated_user
# Authenticate service
service, actual_user_email = await _authenticate_service(
use_oauth21,
@@ -509,6 +559,10 @@ def require_google_service(
raise
try:
# In OAuth 2.1 mode, we need to add user_google_email to kwargs since it was removed from signature
if is_oauth21_enabled():
kwargs["user_google_email"] = user_google_email
# Prepend the fetched service object to the original arguments
return await func(service, *args, **kwargs)
except RefreshError as e:
@@ -552,12 +606,31 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
"""
def decorator(func: Callable) -> Callable:
original_sig = inspect.signature(func)
# In OAuth 2.1 mode, remove user_google_email from the signature
if is_oauth21_enabled():
params = list(original_sig.parameters.values())
filtered_params = [
p for p in params
if p.name != 'user_google_email'
]
wrapper_sig = original_sig.replace(parameters=filtered_params)
else:
wrapper_sig = original_sig
@wraps(func)
async def wrapper(*args, **kwargs):
# Extract user_google_email
sig = inspect.signature(func)
param_names = list(sig.parameters.keys())
# Get authentication context early
tool_name = func.__name__
authenticated_user, _, mcp_session_id = _get_auth_context(tool_name)
# Extract user_google_email based on OAuth mode
if is_oauth21_enabled():
user_google_email = _extract_oauth21_user_email(authenticated_user, tool_name)
else:
# OAuth 2.0 mode: extract from arguments (original logic)
param_names = list(original_sig.parameters.keys())
user_google_email = None
if "user_google_email" in kwargs:
user_google_email = kwargs["user_google_email"]
@@ -588,17 +661,14 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
resolved_scopes = _resolve_scopes(scopes)
try:
tool_name = func.__name__
# Get authentication context
authenticated_user, _, mcp_session_id = _get_auth_context(tool_name)
# Detect OAuth version (simplified for multiple services)
use_oauth21 = (
is_oauth21_enabled() and authenticated_user is not None
)
# Override user_google_email with authenticated user when using OAuth 2.1
# In OAuth 2.0 mode, we may need to override user_google_email
if not is_oauth21_enabled():
param_names = list(original_sig.parameters.keys())
user_google_email, args = _override_oauth21_user_email(
use_oauth21,
authenticated_user,
@@ -634,6 +704,10 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
# Call the original function with refresh error handling
try:
# In OAuth 2.1 mode, we need to add user_google_email to kwargs since it was removed from signature
if is_oauth21_enabled():
kwargs["user_google_email"] = user_google_email
return await func(*args, **kwargs)
except RefreshError as e:
# Handle token refresh errors gracefully
@@ -642,6 +716,15 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
)
raise Exception(error_message)
# Set the wrapper's signature
wrapper.__signature__ = wrapper_sig
# Conditionally modify docstring to remove user_google_email parameter documentation
if is_oauth21_enabled():
logger.debug('OAuth 2.1 mode enabled, removing user_google_email from docstring')
if func.__doc__:
wrapper.__doc__ = _remove_user_email_arg_from_docstring(func.__doc__)
return wrapper
return decorator

216
auth/service_decorator_oauth21.py.bak
View File

@@ -1,216 +0,0 @@
"""
Enhanced Service Decorator with OAuth 2.1 Support
This module provides an enhanced version of the service decorator that can
extract and use OAuth 2.1 session context from FastMCP.
"""
import inspect
import logging
from functools import wraps
from typing import Dict, List, Optional, Any, Callable, Union
from datetime import datetime, timedelta
from google.auth.exceptions import RefreshError
from auth.service_decorator import (
SERVICE_CONFIGS,
SCOPE_GROUPS,
_resolve_scopes,
_get_cache_key,
_is_cache_valid,
_handle_token_refresh_error,
_get_cached_service,
_cache_service,
GoogleAuthenticationError,
)
from auth.oauth21_integration import get_authenticated_google_service_oauth21
logger = logging.getLogger(__name__)
def _extract_context_from_args(args: tuple, kwargs: dict, sig: inspect.Signature) -> Optional[Dict[str, Any]]:
"""
Extract FastMCP Context from function arguments.
Args:
args: Positional arguments
kwargs: Keyword arguments
sig: Function signature
Returns:
Context information if found
"""
param_names = list(sig.parameters.keys())
# Check for Context type annotation
for param_name, param in sig.parameters.items():
if param.annotation and "Context" in str(param.annotation):
# Found Context parameter
if param_name in kwargs:
ctx = kwargs[param_name]
else:
try:
param_index = param_names.index(param_name)
if param_index < len(args):
ctx = args[param_index]
else:
continue
except ValueError:
continue
# Extract relevant information from Context
if ctx:
context_info = {}
# Try to get session_id
if hasattr(ctx, "session_id"):
context_info["session_id"] = ctx.session_id
# Try to get request object
if hasattr(ctx, "request"):
context_info["request"] = ctx.request
# Try to get auth context from request state
if hasattr(ctx, "request") and hasattr(ctx.request, "state"):
if hasattr(ctx.request.state, "auth"):
context_info["auth_context"] = ctx.request.state.auth
return context_info if context_info else None
return None
def require_google_service_oauth21(
service_type: str,
scopes: Union[str, List[str]],
version: Optional[str] = None,
cache_enabled: bool = True,
fallback_to_legacy: bool = True
):
"""
Enhanced decorator that injects authenticated Google service with OAuth 2.1 support.
This decorator checks for FastMCP Context in the function parameters and uses
OAuth 2.1 session information if available, otherwise falls back to legacy auth.
Args:
service_type: Type of Google service (e.g., 'gmail', 'drive')
scopes: Required scopes or scope aliases
version: API version (optional, uses default if not specified)
cache_enabled: Whether to cache service instances
fallback_to_legacy: Whether to fall back to legacy auth if OAuth 2.1 fails
Usage:
@require_google_service_oauth21("gmail", "gmail_read")
async def search_emails(service, user_google_email: str, ctx: Context):
# service is automatically injected
# ctx provides OAuth 2.1 session context
"""
def decorator(func: Callable) -> Callable:
# Get service configuration
if service_type not in SERVICE_CONFIGS:
raise ValueError(f"Unknown service type: {service_type}")
service_config = SERVICE_CONFIGS[service_type]
service_name = service_config["service"]
service_version = version or service_config["version"]
# Resolve scopes
resolved_scopes = _resolve_scopes(scopes)
# Create wrapper with modified signature
sig = inspect.signature(func)
params = list(sig.parameters.values())
# Remove 'service' parameter from signature
wrapper_params = [p for p in params if p.name != 'service']
wrapper_sig = sig.replace(parameters=wrapper_params)
@wraps(func)
async def wrapper(*args, **kwargs):
# Extract user_google_email
user_google_email = None
if 'user_google_email' in kwargs:
user_google_email = kwargs['user_google_email']
else:
param_names = list(sig.parameters.keys())
try:
user_email_index = param_names.index('user_google_email')
if user_email_index < len(args):
user_google_email = args[user_email_index]
except ValueError:
pass
if not user_google_email:
raise ValueError("user_google_email parameter is required")
# Extract context information
context = _extract_context_from_args(args, kwargs, sig)
service = None
actual_user_email = user_google_email
# Check cache if enabled
if cache_enabled:
cache_key = _get_cache_key(user_google_email, service_name, service_version, resolved_scopes)
cached_result = _get_cached_service(cache_key)
if cached_result:
service, actual_user_email = cached_result
logger.debug(f"Using cached service for {user_google_email}")
if service is None:
try:
tool_name = func.__name__
# Try OAuth 2.1 authentication with context
if context:
logger.debug(f"Attempting OAuth 2.1 authentication for {tool_name}")
service, actual_user_email = await get_authenticated_google_service_oauth21(
service_name=service_name,
version=service_version,
tool_name=tool_name,
user_google_email=user_google_email,
required_scopes=resolved_scopes,
context=context,
)
elif fallback_to_legacy:
# Fall back to legacy authentication
logger.debug(f"Using legacy authentication for {tool_name}")
from auth.google_auth import get_authenticated_google_service
service, actual_user_email = await get_authenticated_google_service(
service_name=service_name,
version=service_version,
tool_name=tool_name,
user_google_email=user_google_email,
required_scopes=resolved_scopes,
)
else:
raise GoogleAuthenticationError(
"OAuth 2.1 context required but not found"
)
# Cache the service if enabled
if cache_enabled and service:
cache_key = _get_cache_key(user_google_email, service_name, service_version, resolved_scopes)
_cache_service(cache_key, service, actual_user_email)
except GoogleAuthenticationError as e:
raise Exception(str(e))
# Call the original function with the service object injected
try:
return await func(service, *args, **kwargs)
except RefreshError as e:
error_message = _handle_token_refresh_error(e, actual_user_email, service_name)
raise Exception(error_message)
# Set the wrapper's signature to the one without 'service'
wrapper.__signature__ = wrapper_sig
return wrapper
return decorator
# Alias for backward compatibility
require_google_service = require_google_service_oauth21

50
core/log_formatter.py
View File

@@ -5,7 +5,9 @@ Provides visually appealing log formatting with emojis and consistent styling
to match the safe_print output format.
"""
import logging
import os
import re
import sys
class EnhancedLogFormatter(logging.Formatter):
@@ -140,3 +142,51 @@ def setup_enhanced_logging(log_level: int = logging.INFO, use_colors: bool = Tru
console_handler.setFormatter(formatter)
console_handler.setLevel(log_level)
root_logger.addHandler(console_handler)
def configure_file_logging(logger_name: str = None) -> bool:
"""
Configure file logging based on stateless mode setting.
In stateless mode, file logging is completely disabled to avoid filesystem writes.
In normal mode, sets up detailed file logging to 'mcp_server_debug.log'.
Args:
logger_name: Optional name for the logger (defaults to root logger)
Returns:
bool: True if file logging was configured, False if skipped (stateless mode)
"""
# Check if stateless mode is enabled
stateless_mode = os.getenv("WORKSPACE_MCP_STATELESS_MODE", "false").lower() == "true"
if stateless_mode:
logger = logging.getLogger(logger_name)
logger.debug("File logging disabled in stateless mode")
return False
# Configure file logging for normal mode
try:
target_logger = logging.getLogger(logger_name)
log_file_dir = os.path.dirname(os.path.abspath(__file__))
# Go up one level since we're in core/ subdirectory
log_file_dir = os.path.dirname(log_file_dir)
log_file_path = os.path.join(log_file_dir, 'mcp_server_debug.log')
file_handler = logging.FileHandler(log_file_path, mode='a')
file_handler.setLevel(logging.DEBUG)
file_formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(process)d - %(threadName)s '
'[%(module)s.%(funcName)s:%(lineno)d] - %(message)s'
)
file_handler.setFormatter(file_formatter)
target_logger.addHandler(file_handler)
logger = logging.getLogger(logger_name)
logger.debug(f"Detailed file logging configured to: {log_file_path}")
return True
except Exception as e:
sys.stderr.write(f"CRITICAL: Failed to set up file logging to '{log_file_path}': {e}\n")
return False

34
fastmcp_server.py
View File

@@ -9,8 +9,8 @@ import os
import sys
from dotenv import load_dotenv
from auth.oauth_config import reload_oauth_config
from core.log_formatter import EnhancedLogFormatter
from auth.oauth_config import reload_oauth_config, is_stateless_mode
from core.log_formatter import EnhancedLogFormatter, configure_file_logging
from core.utils import check_credentials_directory_permissions
from core.server import server, set_transport_mode, configure_server_for_http
from core.tool_registry import set_enabled_tools as set_enabled_tool_names, wrap_server_tool_method, filter_server_tools
@@ -33,25 +33,8 @@ logging.basicConfig(
)
logger = logging.getLogger(__name__)
# Configure file logging
try:
root_logger = logging.getLogger()
log_file_dir = os.path.dirname(os.path.abspath(__file__))
log_file_path = os.path.join(log_file_dir, 'mcp_server_debug.log')
file_handler = logging.FileHandler(log_file_path, mode='a')
file_handler.setLevel(logging.DEBUG)
file_formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(process)d - %(threadName)s '
'[%(module)s.%(funcName)s:%(lineno)d] - %(message)s'
)
file_handler.setFormatter(file_formatter)
root_logger.addHandler(file_handler)
logger.debug(f"Detailed file logging configured to: {log_file_path}")
except Exception as e:
sys.stderr.write(f"CRITICAL: Failed to set up file logging to '{log_file_path}': {e}\n")
# Configure file logging based on stateless mode
configure_file_logging()
def configure_safe_logging():
"""Configure safe Unicode handling for logging."""
@@ -76,15 +59,18 @@ def configure_safe_logging():
# Configure safe logging
configure_safe_logging()
# Check credentials directory permissions
try:
# Check credentials directory permissions (skip in stateless mode)
if not is_stateless_mode():
try:
logger.info("🔍 Checking credentials directory permissions...")
check_credentials_directory_permissions()
logger.info("✅ Credentials directory permissions verified")
except (PermissionError, OSError) as e:
except (PermissionError, OSError) as e:
logger.error(f"❌ Credentials directory permission check failed: {e}")
logger.error(" Please ensure the service has write permissions to create/access the credentials directory")
sys.exit(1)
else:
logger.info("🔍 Skipping credentials directory check (stateless mode)")
# Set transport mode for HTTP (FastMCP CLI defaults to streamable-http)
set_transport_mode('streamable-http')

33
main.py
View File

@@ -7,7 +7,7 @@ from importlib import metadata
from dotenv import load_dotenv
from auth.oauth_config import reload_oauth_config
from core.log_formatter import EnhancedLogFormatter
from core.log_formatter import EnhancedLogFormatter, configure_file_logging
from core.utils import check_credentials_directory_permissions
from core.server import server, set_transport_mode, configure_server_for_http
from core.tool_tier_loader import resolve_tools_from_tier
@@ -27,24 +27,10 @@ logging.basicConfig(
)
logger = logging.getLogger(__name__)
try:
root_logger = logging.getLogger()
log_file_dir = os.path.dirname(os.path.abspath(__file__))
log_file_path = os.path.join(log_file_dir, 'mcp_server_debug.log')
configure_file_logging()
file_handler = logging.FileHandler(log_file_path, mode='a')
file_handler.setLevel(logging.DEBUG)
file_formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(process)d - %(threadName)s '
'[%(module)s.%(funcName)s:%(lineno)d] - %(message)s'
)
file_handler.setFormatter(file_formatter)
root_logger.addHandler(file_handler)
logger.debug(f"Detailed file logging configured to: {log_file_path}")
except Exception as e:
sys.stderr.write(f"CRITICAL: Failed to set up file logging to '{log_file_path}': {e}\n")
# Define stateless_mode for use in main() function
stateless_mode = os.getenv("WORKSPACE_MCP_STATELESS_MODE", "false").lower() == "true"
def safe_print(text):
# Don't print to stderr when running as MCP server via uvx to avoid JSON parsing errors
@@ -135,6 +121,7 @@ def main():
"USER_GOOGLE_EMAIL": os.getenv('USER_GOOGLE_EMAIL', 'Not Set'),
"MCP_SINGLE_USER_MODE": os.getenv('MCP_SINGLE_USER_MODE', 'false'),
"MCP_ENABLE_OAUTH21": os.getenv('MCP_ENABLE_OAUTH21', 'false'),
"WORKSPACE_MCP_STATELESS_MODE": os.getenv('WORKSPACE_MCP_STATELESS_MODE', 'false'),
"OAUTHLIB_INSECURE_TRANSPORT": os.getenv('OAUTHLIB_INSECURE_TRANSPORT', 'false'),
"GOOGLE_CLIENT_SECRET_PATH": os.getenv('GOOGLE_CLIENT_SECRET_PATH', 'Not Set'),
}
@@ -225,11 +212,16 @@ def main():
# Set global single-user mode flag
if args.single_user:
if stateless_mode:
safe_print("❌ Single-user mode is incompatible with stateless mode")
safe_print(" Stateless mode requires OAuth 2.1 which is multi-user")
sys.exit(1)
os.environ['MCP_SINGLE_USER_MODE'] = '1'
safe_print("🔐 Single-user mode enabled")
safe_print("")
# Check credentials directory permissions before starting
# Check credentials directory permissions before starting (skip in stateless mode)
if not stateless_mode:
try:
safe_print("🔍 Checking credentials directory permissions...")
check_credentials_directory_permissions()
@@ -240,6 +232,9 @@ def main():
safe_print(" Please ensure the service has write permissions to create/access the credentials directory")
logger.error(f"Failed credentials directory permission check: {e}")
sys.exit(1)
else:
safe_print("🔍 Skipping credentials directory check (stateless mode)")
safe_print("")
try:
# Set transport mode for OAuth callback handling