Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
Toggle table of contents sidebar

# Source code for mindroot.coreplugins.mcp_.oauth_storage

    """OAuth token storage implementation for MCP servers."""

    from typing import Optional
    from datetime import datetime
    import json

    try:
        from mcp.client.auth import TokenStorage
        from mcp.shared.auth import OAuthToken, OAuthClientInformationFull
    except ImportError:
        # Fallback if MCP not available
        class TokenStorage:
            pass
        OAuthToken = None
        OAuthClientInformationFull = None

    class MCPTokenStorage(TokenStorage):
        """Token storage implementation for MCP OAuth flows.

        Stores tokens and client information in the MCP server configuration.
        """

        def __init__(self, server_name: str, mcp_manager):
            self.server_name = server_name
            self.mcp_manager = mcp_manager

        async def get_tokens(self) -> Optional[OAuthToken]:
            """Get stored OAuth tokens for the server."""
            if OAuthToken is None:
                return None

            if self.server_name not in self.mcp_manager.servers:
                return None

            server = self.mcp_manager.servers[self.server_name]

            if not server.access_token:
                return None

            # Parse token expiration
            expires_at = None
            if server.token_expires_at:
                try:
                    expires_at = datetime.fromisoformat(server.token_expires_at.replace('Z', '+00:00'))
                except ValueError:
                    pass

            return OAuthToken(
                access_token=server.access_token,
                refresh_token=server.refresh_token,
                expires_at=expires_at,
                scope=" ".join(server.scopes) if server.scopes else None
            )

        async def set_tokens(self, tokens: OAuthToken) -> None:
            """Store OAuth tokens for the server."""
            if self.server_name not in self.mcp_manager.servers:
                return

            server = self.mcp_manager.servers[self.server_name]

            # Update server with token information
            server.access_token = tokens.access_token
            server.refresh_token = tokens.refresh_token

            if hasattr(tokens, 'expires_at') and tokens.expires_at:
                server.token_expires_at = tokens.expires_at.isoformat()
            elif hasattr(tokens, 'expires_in') and tokens.expires_in:
                # Calculate expires_at from expires_in
                from datetime import datetime, timedelta
                expires_at = datetime.now() + timedelta(seconds=tokens.expires_in)
                server.token_expires_at = expires_at.isoformat()
            else:
                server.token_expires_at = None

            if tokens.scope:
                server.scopes = tokens.scope.split(" ")

            # Save configuration
            self.mcp_manager.save_config()

        async def get_client_info(self) -> Optional[OAuthClientInformationFull]:
            """Get stored OAuth client information."""
            if OAuthClientInformationFull is None:
                return None

            if self.server_name not in self.mcp_manager.servers:
                return None

            server = self.mcp_manager.servers[self.server_name]

            if not server.client_id:
                return None

            # Create client info from server configuration
            client_info_data = {
                "client_id": server.client_id,
                "client_name": f"MindRoot MCP Client - {server.name}",
                "redirect_uris": [server.redirect_uri] if server.redirect_uri else [],
                "grant_types": ["authorization_code", "refresh_token"],
                "response_types": ["code"],
                "scope": " ".join(server.scopes) if server.scopes else "user"
            }

            if server.client_secret:
                client_info_data["client_secret"] = server.client_secret

            return OAuthClientInformationFull(**client_info_data)

        async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
            """Store OAuth client information."""
            if self.server_name not in self.mcp_manager.servers:
                return

            server = self.mcp_manager.servers[self.server_name]

            # Update server with client information
            server.client_id = client_info.client_id

            if hasattr(client_info, 'client_secret') and client_info.client_secret:
                server.client_secret = client_info.client_secret

            if client_info.redirect_uris:
                server.redirect_uri = client_info.redirect_uris[0]

            if client_info.scope:
                server.scopes = client_info.scope.split(" ")

            # Save configuration
            self.mcp_manager.save_config()

        def clear_tokens(self):
            """Clear stored tokens for the server."""
            if self.server_name not in self.mcp_manager.servers:
                return

            server = self.mcp_manager.servers[self.server_name]
            server.access_token = None
            server.refresh_token = None
            server.token_expires_at = None

            # Save configuration
            self.mcp_manager.save_config()