open-edison 0.1.19__py3-none-any.whl → 0.1.29__py3-none-any.whl

src/middleware/session_tracking.py

@@ -12,7 +12,7 @@ from contextvars import ContextVar
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import mcp.types as mt
 from fastmcp.prompts.prompt import FunctionPrompt
@@ -27,8 +27,13 @@ from sqlalchemy import JSON, Column, Integer, String, create_engine, event
 from sqlalchemy.orm import Session, declarative_base
 from sqlalchemy.sql import select
 
+from src import events
 from src.config import get_config_dir  # type: ignore[reportMissingImports]
-from src.middleware.data_access_tracker import DataAccessTracker
+from src.middleware.data_access_tracker import (
+    DataAccessTracker,
+    SecurityError,
+)
+from src.permissions import Permissions
 from src.telemetry import (
     record_prompt_used,
     record_resource_used,
@@ -67,7 +72,10 @@ class MCPSessionModel(Base):  # type: ignore
     data_access_summary = Column(JSON)  # type: ignore
 
 
-current_session_id_ctxvar: ContextVar[str | None] = ContextVar("current_session_id", default=None)
+current_session_id_ctxvar: ContextVar[str | None] = ContextVar(
+    "current_session_id",
+    default=cast(str | None, None),  # noqa: B039
+)
 
 
 def get_current_session_data_tracker() -> DataAccessTracker | None:
@@ -101,7 +109,7 @@ def create_db_session() -> Generator[Session, None, None]:
     engine = create_engine(f"sqlite:///{db_path}")
 
     # Ensure changes are flushed to the main database file (avoid WAL for sql.js compatibility)
-    @event.listens_for(engine, "connect")
+    @event.listens_for(engine, "connect")  # noqa
     def _set_sqlite_pragmas(dbapi_connection, connection_record):  # type: ignore[no-untyped-def] # noqa
         cur = dbapi_connection.cursor()  # type: ignore[attr-defined]
         try:
@@ -171,7 +179,7 @@ def get_session_from_db(session_id: str) -> MCPSession:
                     "has_external_communication", False
                 )
             # Restore ACL highest level if present
-            if isinstance(summary_data, dict) and "acl" in summary_data:
+            if isinstance(summary_data, dict) and "acl" in summary_data:  # type: ignore
                 acl_summary: Any = summary_data.get("acl")  # type: ignore
                 if isinstance(acl_summary, dict):
                     highest = acl_summary.get("highest_acl_level")  # type: ignore
@@ -214,7 +222,7 @@ class SessionTrackingMiddleware(Middleware):
         return session, session_id
 
     # General hooks for on_request, on_message, etc.
-    async def on_request(
+    async def on_request(  # noqa
         self,
         context: MiddlewareContext[mt.Request[Any, Any]],  # type: ignore
         call_next: CallNext[mt.Request[Any, Any], Any],  # type: ignore
@@ -225,17 +233,25 @@ class SessionTrackingMiddleware(Middleware):
         # Get or create session stats
         _, _session_id = self._get_or_create_session_stats(context)
 
-        return await call_next(context)  # type: ignore
+        try:
+            return await call_next(context)  # type: ignore
+        except Exception:
+            log.exception("MCP request handling failed")
+            raise
 
     # Hooks for Tools
-    async def on_list_tools(
+    async def on_list_tools(  # noqa
         self,
         context: MiddlewareContext[Any],  # type: ignore
         call_next: CallNext[Any, Any],  # type: ignore
     ) -> Any:
         log.debug("🔍 on_list_tools")
         # Get the original response
-        response = await call_next(context)
+        try:
+            response = await call_next(context)
+        except Exception:
+            log.exception("MCP list_tools failed")
+            raise
         log.trace(f"🔍 on_list_tools response: {response}")
 
         session_id = current_session_id_ctxvar.get()
@@ -247,8 +263,11 @@ class SessionTrackingMiddleware(Middleware):
 
         # Filter out specific tools or return empty list
         allowed_tools: list[FunctionTool | ProxyTool | Any] = []
+        perms = Permissions()
         for tool in response:
-            log.trace(f"🔍 Processing tool listing {tool.name}")
+            # Due to proxy & server naming
+            tool_name = tool.key
+            log.trace(f"🔍 Processing tool listing {tool_name}")
             if isinstance(tool, FunctionTool):
                 log.trace("🔍 Tool is built-in")
                 log.trace(f"🔍 Tool is a FunctionTool: {tool}")
@@ -260,20 +279,18 @@ class SessionTrackingMiddleware(Middleware):
                 log.trace(f"🔍 Tool is a unknown type: {tool}")
                 continue
 
-            log.trace(f"🔍 Getting permissions for tool {tool.name}")
-            permissions = session.data_access_tracker.get_tool_permissions(tool.name)
-            log.trace(f"🔍 Tool permissions: {permissions}")
-            if permissions["enabled"]:
+            log.trace(f"🔍 Getting permissions for tool {tool_name}")
+            if perms.is_tool_enabled(tool_name):
                 allowed_tools.append(tool)
             else:
                 log.warning(
-                    f"🔍 Tool {tool.name} is disabled on not configured and will not be allowed"
+                    f"🔍 Tool {tool_name} is disabled or not configured and will not be allowed"
                 )
                 continue
 
         return allowed_tools  # type: ignore
 
-    async def on_call_tool(
+    async def on_call_tool(  # noqa
         self,
         context: MiddlewareContext[mt.CallToolRequestParams],  # type: ignore
         call_next: CallNext[mt.CallToolRequestParams, ToolResult],  # type: ignore
@@ -296,7 +313,28 @@ class SessionTrackingMiddleware(Middleware):
 
         assert session.data_access_tracker is not None
         log.debug(f"🔍 Analyzing tool {context.message.name} for security implications")
-        session.data_access_tracker.add_tool_call(context.message.name)
+        try:
+            session.data_access_tracker.add_tool_call(context.message.name)
+        except SecurityError as e:
+            # Publish pre-block event enriched with session_id then wait up to 30s for approval
+            events.fire_and_forget(
+                {
+                    "type": "mcp_pre_block",
+                    "kind": "tool",
+                    "name": context.message.name,
+                    "session_id": session_id,
+                    "error": str(e),
+                }
+            )
+            approved = await events.wait_for_approval(
+                session_id, "tool", context.message.name, timeout_s=30.0
+            )
+            if not approved:
+                raise
+            # Approved: apply effects and proceed
+            session.data_access_tracker.apply_effects_after_manual_approval(
+                "tool", context.message.name
+            )
         # Telemetry: record tool call
         record_tool_call(context.message.name)
 
@@ -330,7 +368,7 @@ class SessionTrackingMiddleware(Middleware):
         return await call_next(context)  # type: ignore
 
     # Hooks for Resources
-    async def on_list_resources(
+    async def on_list_resources(  # noqa
         self,
         context: MiddlewareContext[Any],  # type: ignore
         call_next: CallNext[Any, Any],  # type: ignore
@@ -338,7 +376,11 @@ class SessionTrackingMiddleware(Middleware):
         """Process resource access and track security implications."""
         log.trace("🔍 on_list_resources")
         # Get the original response
-        response = await call_next(context)
+        try:
+            response = await call_next(context)
+        except Exception:
+            log.exception("MCP list_resources failed")
+            raise
         log.trace(f"🔍 on_list_resources response: {response}")
 
         session_id = current_session_id_ctxvar.get()
@@ -350,6 +392,7 @@ class SessionTrackingMiddleware(Middleware):
 
         # Filter out specific tools or return empty list
         allowed_resources: list[FunctionResource | ProxyResource | Any] = []
+        perms = Permissions()
         for resource in response:
             resource_name = str(resource.uri)
             log.trace(f"🔍 Processing resource listing {resource_name}")
@@ -365,19 +408,17 @@ class SessionTrackingMiddleware(Middleware):
                 continue
 
             log.trace(f"🔍 Getting permissions for resource {resource_name}")
-            permissions = session.data_access_tracker.get_resource_permissions(resource_name)
-            log.trace(f"🔍 Resource permissions: {permissions}")
-            if permissions["enabled"]:
+            if perms.is_resource_enabled(resource_name):
                 allowed_resources.append(resource)
             else:
                 log.warning(
-                    f"🔍 Resource {resource_name} is disabled on not configured and will not be allowed"
+                    f"🔍 Resource {resource_name} is disabled or not configured and will not be allowed"
                 )
                 continue
 
         return allowed_resources  # type: ignore
 
-    async def on_read_resource(
+    async def on_read_resource(  # noqa
         self,
         context: MiddlewareContext[Any],  # type: ignore
         call_next: CallNext[Any, Any],  # type: ignore
@@ -386,7 +427,11 @@ class SessionTrackingMiddleware(Middleware):
         session_id = current_session_id_ctxvar.get()
         if session_id is None:
             log.warning("No session ID found for resource access tracking")
-            return await call_next(context)
+            try:
+                return await call_next(context)
+            except Exception:
+                log.exception("MCP read_resource failed")
+                raise
 
         session = get_session_from_db(session_id)
         log.trace(f"Adding resource access to session {session_id}")
@@ -396,7 +441,26 @@ class SessionTrackingMiddleware(Middleware):
         resource_name = str(context.message.uri)
 
         log.debug(f"🔍 Analyzing resource {resource_name} for security implications")
-        _ = session.data_access_tracker.add_resource_access(resource_name)
+        try:
+            _ = session.data_access_tracker.add_resource_access(resource_name)
+        except SecurityError as e:
+            events.fire_and_forget(
+                {
+                    "type": "mcp_pre_block",
+                    "kind": "resource",
+                    "name": resource_name,
+                    "session_id": session_id,
+                    "error": str(e),
+                }
+            )
+            approved = await events.wait_for_approval(
+                session_id, "resource", resource_name, timeout_s=30.0
+            )
+            if not approved:
+                raise
+            session.data_access_tracker.apply_effects_after_manual_approval(
+                "resource", resource_name
+            )
         record_resource_used(resource_name)
 
         # Update database session
@@ -409,10 +473,14 @@ class SessionTrackingMiddleware(Middleware):
             db_session.commit()
 
         log.trace(f"Resource access {resource_name} added to session {session_id}")
-        return await call_next(context)
+        try:
+            return await call_next(context)
+        except Exception:
+            log.exception("MCP read_resource failed")
+            raise
 
     # Hooks for Prompts
-    async def on_list_prompts(
+    async def on_list_prompts(  # noqa
         self,
         context: MiddlewareContext[Any],  # type: ignore
         call_next: CallNext[Any, Any],  # type: ignore
@@ -420,7 +488,11 @@ class SessionTrackingMiddleware(Middleware):
         """Process resource access and track security implications."""
         log.debug("🔍 on_list_prompts")
         # Get the original response
-        response = await call_next(context)
+        try:
+            response = await call_next(context)
+        except Exception:
+            log.exception("MCP list_prompts failed")
+            raise
         log.debug(f"🔍 on_list_prompts response: {response}")
 
         session_id = current_session_id_ctxvar.get()
@@ -432,6 +504,7 @@ class SessionTrackingMiddleware(Middleware):
 
         # Filter out specific tools or return empty list
         allowed_prompts: list[ProxyPrompt | Any] = []
+        perms = Permissions()
         for prompt in response:
             prompt_name = str(prompt.name)
             log.trace(f"🔍 Processing prompt listing {prompt_name}")
@@ -447,19 +520,17 @@ class SessionTrackingMiddleware(Middleware):
                 continue
 
             log.trace(f"🔍 Getting permissions for prompt {prompt_name}")
-            permissions = session.data_access_tracker.get_prompt_permissions(prompt_name)
-            log.trace(f"🔍 Prompt permissions: {permissions}")
-            if permissions["enabled"]:
+            if perms.is_prompt_enabled(prompt_name):
                 allowed_prompts.append(prompt)
             else:
                 log.warning(
-                    f"🔍 Prompt {prompt_name} is disabled on not configured and will not be allowed"
+                    f"🔍 Prompt {prompt_name} is disabled or not configured and will not be allowed"
                 )
                 continue
 
         return allowed_prompts  # type: ignore
 
-    async def on_get_prompt(
+    async def on_get_prompt(  # noqa
         self,
         context: MiddlewareContext[Any],  # type: ignore
         call_next: CallNext[Any, Any],  # type: ignore
@@ -468,7 +539,11 @@ class SessionTrackingMiddleware(Middleware):
         session_id = current_session_id_ctxvar.get()
         if session_id is None:
             log.warning("No session ID found for prompt access tracking")
-            return await call_next(context)
+            try:
+                return await call_next(context)
+            except Exception:
+                log.exception("MCP get_prompt failed")
+                raise
 
         session = get_session_from_db(session_id)
         log.trace(f"Adding prompt access to session {session_id}")
@@ -477,7 +552,24 @@ class SessionTrackingMiddleware(Middleware):
         prompt_name = context.message.name
 
         log.debug(f"🔍 Analyzing prompt {prompt_name} for security implications")
-        _ = session.data_access_tracker.add_prompt_access(prompt_name)
+        try:
+            _ = session.data_access_tracker.add_prompt_access(prompt_name)
+        except SecurityError as e:
+            events.fire_and_forget(
+                {
+                    "type": "mcp_pre_block",
+                    "kind": "prompt",
+                    "name": prompt_name,
+                    "session_id": session_id,
+                    "error": str(e),
+                }
+            )
+            approved = await events.wait_for_approval(
+                session_id, "prompt", prompt_name, timeout_s=30.0
+            )
+            if not approved:
+                raise
+            session.data_access_tracker.apply_effects_after_manual_approval("prompt", prompt_name)
         record_prompt_used(prompt_name)
 
         # Update database session
@@ -490,4 +582,8 @@ class SessionTrackingMiddleware(Middleware):
             db_session.commit()
 
         log.trace(f"Prompt access {prompt_name} added to session {session_id}")
-        return await call_next(context)
+        try:
+            return await call_next(context)
+        except Exception:
+            log.exception("MCP get_prompt failed")
+            raise

src/oauth_manager.py

@@ -0,0 +1,281 @@
+"""
+OAuth Manager for OpenEdison MCP Gateway
+
+Handles OAuth 2.1 authentication for MCP servers using FastMCP's built-in OAuth support.
+Provides detection, token management, and authentication flow coordination.
+"""
+
+import asyncio
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+from fastmcp.client.auth.oauth import (
+    FileTokenStorage,
+    OAuth,
+    check_if_auth_required,
+    default_cache_dir,
+)
+from loguru import logger as log
+
+
+class OAuthStatus(Enum):
+    """OAuth authentication status for MCP servers."""
+
+    UNKNOWN = "unknown"  # noqa
+    NOT_REQUIRED = "not_required"
+    NEEDS_AUTH = "needs_auth"
+    AUTHENTICATED = "authenticated"
+    ERROR = "error"
+    EXPIRED = "expired"  # noqa
+
+
+@dataclass
+class OAuthServerInfo:
+    """OAuth information for an MCP server."""
+
+    server_name: str
+    mcp_url: str
+    status: OAuthStatus
+    scopes: list[str] | None = None
+    client_name: str = "OpenEdison MCP Gateway"
+    error_message: str | None = None
+    token_expires_at: str | None = None
+    has_refresh_token: bool = False
+
+
+class OAuthManager:
+    """
+    Manages OAuth authentication for MCP servers.
+
+    This class provides a centralized interface for:
+    - Detecting which MCP servers require OAuth
+    - Managing OAuth tokens and credentials
+    - Providing OAuth authentication objects for FastMCP clients
+    - Handling token refresh and expiration
+    """
+
+    def __init__(self, cache_dir: Path | None = None):
+        """
+        Initialize OAuth manager.
+
+        Args:
+            cache_dir: Directory for token cache. Defaults to FastMCP's default.
+        """
+        self.cache_dir = cache_dir or default_cache_dir()
+        self._oauth_info: dict[str, OAuthServerInfo] = {}
+
+        log.info(f"🔐 OAuth Manager initialized with cache dir: {self.cache_dir}")
+
+    async def check_oauth_requirement(
+        self, server_name: str, mcp_url: str | None, timeout_seconds: float = 10.0
+    ) -> OAuthServerInfo:
+        """
+        Check if an MCP server requires OAuth authentication.
+
+        Args:
+            server_name: Name of the MCP server
+            mcp_url: URL of the MCP endpoint (None for local servers)
+            timeout_seconds: Timeout for the check request
+
+        Returns:
+            OAuthServerInfo with detection results
+        """
+        log.info(f"🔍 Checking OAuth requirement for {server_name}")
+
+        # If no mcp_url provided, this is a local server - no OAuth needed
+        if not mcp_url:
+            info = OAuthServerInfo(
+                server_name=server_name, mcp_url="", status=OAuthStatus.NOT_REQUIRED
+            )
+            log.info(f"✅ {server_name} is a local server - no OAuth required")
+            self._oauth_info[server_name] = info
+            return info
+
+        log.info(f"🔍 Checking OAuth requirement for remote server {server_name} at {mcp_url}")
+
+        try:
+            # Check if auth is required (with timeout)
+            auth_required = await asyncio.wait_for(
+                check_if_auth_required(mcp_url), timeout=timeout_seconds
+            )
+
+            if not auth_required:
+                info = OAuthServerInfo(
+                    server_name=server_name, mcp_url=mcp_url, status=OAuthStatus.NOT_REQUIRED
+                )
+                log.info(f"✅ {server_name} does not require OAuth")
+                self._oauth_info[server_name] = info
+                return info
+
+            # OAuth is required, proceed with token checking
+            log.info(f"🔐 {server_name} requires OAuth authentication")
+
+            # Check if we have existing valid tokens
+            token_storage = FileTokenStorage(server_url=mcp_url, cache_dir=self.cache_dir)
+            existing_tokens = await token_storage.get_tokens()
+
+            status = OAuthStatus.NEEDS_AUTH
+            token_expires_at = None
+            has_refresh_token = False
+
+            if existing_tokens:
+                # Check if tokens are still valid
+                # Note: FastMCP's FileTokenStorage doesn't expose expiration directly,
+                # so we'll attempt to use the tokens and see if they work
+                has_refresh_token = bool(existing_tokens.refresh_token)
+                if existing_tokens.access_token:
+                    # We have tokens, assume they're valid for now
+                    # The actual validation will happen when the client tries to use them
+                    status = OAuthStatus.AUTHENTICATED
+                    # Try to get expiration time if available
+                    try:
+                        expires_at = getattr(existing_tokens, "expires_at", None)
+                        if expires_at:
+                            token_expires_at = str(expires_at)
+                        else:
+                            expires_in = getattr(existing_tokens, "expires_in", None)
+                            if expires_in:
+                                # If expires_in is available, we can calculate expiration
+                                from datetime import datetime, timedelta
+
+                                expiry = datetime.now() + timedelta(seconds=expires_in)
+                                token_expires_at = expiry.isoformat()
+                    except Exception:
+                        # If we can't get expiration info, that's ok - token_expires_at will be None
+                        pass
+
+            info = OAuthServerInfo(
+                server_name=server_name,
+                mcp_url=mcp_url,
+                status=status,
+                scopes=None,  # We don't have metadata discovery, so no scopes info
+                token_expires_at=token_expires_at,
+                has_refresh_token=has_refresh_token,
+            )
+
+            log.info(f"🔐 {server_name} OAuth status: {status.value}")
+            self._oauth_info[server_name] = info
+            return info
+
+        except TimeoutError:
+            info = OAuthServerInfo(
+                server_name=server_name,
+                mcp_url=mcp_url,
+                status=OAuthStatus.ERROR,
+                error_message=f"OAuth check timed out after {timeout_seconds}s",
+            )
+            log.warning(f"⏰ OAuth check for {server_name} timed out")
+            self._oauth_info[server_name] = info
+            return info
+
+        except Exception as e:
+            info = OAuthServerInfo(
+                server_name=server_name,
+                mcp_url=mcp_url,
+                status=OAuthStatus.ERROR,
+                error_message=str(e),
+            )
+            log.error(f"❌ OAuth check for {server_name} failed: {e}")
+            self._oauth_info[server_name] = info
+            return info
+
+    def get_oauth_auth(
+        self,
+        server_name: str,
+        mcp_url: str,
+        scopes: list[str] | None = None,
+        client_name: str | None = None,
+    ) -> OAuth | None:
+        """
+        Get OAuth authentication object for FastMCP client.
+
+        Args:
+            server_name: Name of the MCP server
+            mcp_url: URL of the MCP endpoint
+            scopes: OAuth scopes to request
+            client_name: Client name for OAuth registration
+
+        Returns:
+            OAuth authentication object, or None if OAuth not required
+        """
+        info = self._oauth_info.get(server_name)
+
+        if not info or info.status == OAuthStatus.NOT_REQUIRED:
+            return None
+
+        if info.status == OAuthStatus.ERROR:
+            log.warning(f"⚠️ Cannot create OAuth auth for {server_name}: {info.error_message}")
+            return None
+
+        try:
+            oauth = OAuth(
+                mcp_url=mcp_url,
+                scopes=scopes or info.scopes,
+                client_name=client_name or info.client_name,
+                token_storage_cache_dir=self.cache_dir,
+                callback_port=50001,
+            )
+            log.info(f"🔐 Created OAuth auth for {server_name}")
+            return oauth
+
+        except Exception as e:
+            log.error(f"❌ Failed to create OAuth auth for {server_name}: {e}")
+            return None
+
+    def clear_tokens(self, server_name: str, mcp_url: str) -> bool:
+        """
+        Clear stored OAuth tokens for a server.
+
+        Args:
+            server_name: Name of the MCP server
+            mcp_url: URL of the MCP endpoint
+
+        Returns:
+            True if tokens were cleared successfully
+        """
+        try:
+            token_storage = FileTokenStorage(server_url=mcp_url, cache_dir=self.cache_dir)
+            token_storage.clear()
+
+            # Update our cached info
+            if server_name in self._oauth_info:
+                self._oauth_info[server_name].status = OAuthStatus.NEEDS_AUTH
+                self._oauth_info[server_name].token_expires_at = None
+                self._oauth_info[server_name].has_refresh_token = False
+
+            log.info(f"🗑️ Cleared OAuth tokens for {server_name}")
+            return True
+
+        except Exception as e:
+            log.error(f"❌ Failed to clear tokens for {server_name}: {e}")
+            return False
+
+    def get_server_info(self, server_name: str) -> OAuthServerInfo | None:
+        """Get OAuth info for a server."""
+        return self._oauth_info.get(server_name)
+
+    async def refresh_server_status(self, server_name: str, mcp_url: str) -> OAuthServerInfo:
+        """
+        Refresh OAuth status for a server.
+
+        Args:
+            server_name: Name of the MCP server
+            mcp_url: URL of the MCP endpoint
+
+        Returns:
+            Updated OAuthServerInfo
+        """
+        return await self.check_oauth_requirement(server_name, mcp_url)
+
+
+# Global OAuth manager instance
+_oauth_manager: OAuthManager | None = None
+
+
+def get_oauth_manager() -> OAuthManager:
+    """Get the global OAuth manager instance."""
+    global _oauth_manager
+    if _oauth_manager is None:
+        _oauth_manager = OAuthManager()
+    return _oauth_manager