datarobot-genai 0.1.59__tar.gz → 0.1.70__tar.gz

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.59
+Version: 0.1.70
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -32,7 +32,7 @@ Requires-Dist: aiohttp<4.0.0,>=3.9.0; extra == 'drmcp'
 Requires-Dist: aiosignal<2.0.0,>=1.3.1; extra == 'drmcp'
 Requires-Dist: boto3<2.0.0,>=1.34.0; extra == 'drmcp'
 Requires-Dist: datarobot-asgi-middleware<1.0.0,>=0.2.0; extra == 'drmcp'
-Requires-Dist: fastmcp==2.13.0.2; extra == 'drmcp'
+Requires-Dist: fastmcp<3.0.0,>=2.13.0.2; extra == 'drmcp'
 Requires-Dist: httpx<1.0.0,>=0.28.1; extra == 'drmcp'
 Requires-Dist: opentelemetry-api<2.0.0,>=1.22.0; extra == 'drmcp'
 Requires-Dist: opentelemetry-exporter-otlp-proto-http<2.0.0,>=1.22.0; extra == 'drmcp'

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.1.59"
+version = "0.1.70"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -84,7 +84,7 @@ drmcp = [
   "aiohttp>=3.9.0,<4.0.0",
   "aiohttp-retry>=2.8.3,<3.0.0",
   "aiosignal>=1.3.1,<2.0.0",
-  "fastmcp==2.13.0.2",
+  "fastmcp>=2.13.0.2,<3.0.0",
 ]
 
 [tool.hatch.build.targets.wheel]

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/src/datarobot_genai/core/agents/base.py

@@ -52,6 +52,7 @@ class BaseAgent(Generic[TTool], abc.ABC):
         verbose: bool | str | None = True,
         timeout: int | None = 90,
         authorization_context: dict[str, Any] | None = None,
+        forwarded_headers: dict[str, str] | None = None,
         **_: Any,
     ) -> None:
         self.api_key = api_key or os.environ.get("DATAROBOT_API_TOKEN")
@@ -68,6 +69,7 @@ class BaseAgent(Generic[TTool], abc.ABC):
             self.verbose = bool(verbose)
         self._mcp_tools: list[TTool] = []
         self._authorization_context = authorization_context or {}
+        self._forwarded_headers: dict[str, str] = forwarded_headers or {}
 
     def set_mcp_tools(self, tools: list[TTool]) -> None:
         self._mcp_tools = tools
@@ -86,6 +88,11 @@ class BaseAgent(Generic[TTool], abc.ABC):
         """Return the authorization context for this agent."""
         return self._authorization_context
 
+    @property
+    def forwarded_headers(self) -> dict[str, str]:
+        """Return the forwarded headers for this agent."""
+        return self._forwarded_headers
+
     def litellm_api_base(self, deployment_id: str | None) -> str:
         return get_api_base(self.api_base, deployment_id)
 

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/src/datarobot_genai/core/custom_model.py

@@ -139,6 +139,11 @@ def chat_entrypoint(
     completion_create_params["authorization_context"] = resolve_authorization_context(
         completion_create_params, **kwargs
     )
+    # Keep only allowed headers from the forwarded_headers.
+    incoming_headers = kwargs.get("headers", {}) or {}
+    allowed_headers = {"x-datarobot-api-token", "x-datarobot-api-key"}
+    forwarded_headers = {k: v for k, v in incoming_headers.items() if k.lower() in allowed_headers}
+    completion_create_params["forwarded_headers"] = forwarded_headers
 
     # Instantiate user agent with all supplied completion params including auth context
     agent = agent_cls(**completion_create_params)

datarobot_genai-0.1.70/src/datarobot_genai/core/mcp/common.py

@@ -0,0 +1,218 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import logging
+import re
+from http import HTTPStatus
+from typing import Any
+from typing import Literal
+
+import requests
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from pydantic import field_validator
+
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+
+logger = logging.getLogger(__name__)
+
+
+class MCPConfig(DataRobotAppFrameworkBaseSettings):
+    """Configuration for MCP server connection.
+
+    Derived values are exposed as properties rather than stored, avoiding
+    Pydantic field validation/serialization concerns for internal helpers.
+    """
+
+    external_mcp_url: str | None = None
+    external_mcp_headers: str | None = None
+    external_mcp_transport: Literal["sse", "streamable-http"] = "streamable-http"
+    mcp_deployment_id: str | None = None
+    datarobot_endpoint: str | None = None
+    datarobot_api_token: str | None = None
+    authorization_context: dict[str, Any] | None = None
+    forwarded_headers: dict[str, str] | None = None
+    mcp_server_port: int | None = None
+
+    _auth_context_handler: AuthContextHeaderHandler | None = None
+    _server_config: dict[str, Any] | None = None
+
+    @field_validator("external_mcp_headers", mode="before")
+    @classmethod
+    def validate_external_mcp_headers(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+
+        candidate = value.strip()
+
+        try:
+            json.loads(candidate)
+        except json.JSONDecodeError:
+            msg = "external_mcp_headers must be valid JSON"
+            logger.warning(msg)
+            return None
+
+        return candidate
+
+    @field_validator("mcp_deployment_id", mode="before")
+    @classmethod
+    def validate_mcp_deployment_id(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+
+        candidate = value.strip()
+
+        if not re.fullmatch(r"[0-9a-fA-F]{24}", candidate):
+            msg = "mcp_deployment_id must be a valid 24-character hex ID"
+            logger.warning(msg)
+            return None
+
+        return candidate
+
+    def _authorization_bearer_header(self) -> dict[str, str]:
+        """Return Authorization header with Bearer token or empty dict."""
+        if not self.datarobot_api_token:
+            return {}
+        auth = (
+            self.datarobot_api_token
+            if self.datarobot_api_token.startswith("Bearer ")
+            else f"Bearer {self.datarobot_api_token}"
+        )
+        return {"Authorization": auth}
+
+    @property
+    def auth_context_handler(self) -> AuthContextHeaderHandler:
+        if self._auth_context_handler is None:
+            self._auth_context_handler = AuthContextHeaderHandler()
+        return self._auth_context_handler
+
+    @property
+    def server_config(self) -> dict[str, Any] | None:
+        if self._server_config is None:
+            self._server_config = self._build_server_config()
+        return self._server_config
+
+    def _authorization_context_header(self) -> dict[str, str]:
+        """Return X-DataRobot-Authorization-Context header or empty dict."""
+        try:
+            return self.auth_context_handler.get_header(self.authorization_context)
+        except (LookupError, RuntimeError):
+            # Authorization context not available (e.g., in tests)
+            return {}
+
+    def _build_authenticated_headers(self) -> dict[str, str]:
+        """Build headers for authenticated requests.
+
+        Returns
+        -------
+            Dictionary containing forwarded headers (if available) and authentication headers.
+        """
+        headers: dict[str, str] = {}
+        if self.forwarded_headers:
+            headers.update(self.forwarded_headers)
+        headers.update(self._authorization_bearer_header())
+        headers.update(self._authorization_context_header())
+        return headers
+
+    def _check_localhost_server(self, url: str, timeout: float = 2.0) -> bool:
+        """Check if MCP server is running on localhost.
+
+        Parameters
+        ----------
+        url : str
+            The URL to check.
+        timeout : float, optional
+            Request timeout in seconds (default: 2.0).
+
+        Returns
+        -------
+        bool
+            True if server is running and responding with OK status, False otherwise.
+        """
+        try:
+            response = requests.get(url, timeout=timeout)
+            return (
+                response.status_code == HTTPStatus.OK
+                and response.json().get("message") == "DataRobot MCP Server is running"
+            )
+        except requests.RequestException as e:
+            logger.debug(f"Failed to connect to MCP server at {url}: {e}")
+            return False
+
+    def _build_server_config(self) -> dict[str, Any] | None:
+        """
+        Get MCP server configuration.
+
+        Returns
+        -------
+            Server configuration dict with url, transport, and optional headers,
+            or None if not configured.
+        """
+        if self.mcp_deployment_id:
+            # DataRobot deployment ID - requires authentication
+            if self.datarobot_endpoint is None:
+                raise ValueError(
+                    "When using a DataRobot hosted MCP deployment, datarobot_endpoint must be set."
+                )
+            if self.datarobot_api_token is None:
+                raise ValueError(
+                    "When using a DataRobot hosted MCP deployment, datarobot_api_token must be set."
+                )
+
+            base_url = self.datarobot_endpoint.rstrip("/")
+            if not base_url.endswith("/api/v2"):
+                base_url = f"{base_url}/api/v2"
+
+            url = f"{base_url}/deployments/{self.mcp_deployment_id}/directAccess/mcp"
+            headers = self._build_authenticated_headers()
+
+            logger.info(f"Using DataRobot hosted MCP deployment: {url}")
+
+            return {
+                "url": url,
+                "transport": "streamable-http",
+                "headers": headers,
+            }
+
+        if self.external_mcp_url:
+            # External MCP URL - no authentication needed
+            headers = {}
+
+            # Merge external headers if provided
+            if self.external_mcp_headers:
+                external_headers = json.loads(self.external_mcp_headers)
+                headers.update(external_headers)
+
+            logger.info(f"Using external MCP URL: {self.external_mcp_url}")
+
+            return {
+                "url": self.external_mcp_url.rstrip("/"),
+                "transport": self.external_mcp_transport,
+                "headers": headers,
+            }
+
+        # No MCP configuration found, setup localhost if running locally
+        if self.mcp_server_port:
+            url = f"http://localhost:{self.mcp_server_port}"
+            if self._check_localhost_server(url):
+                headers = self._build_authenticated_headers()
+                logger.info(f"Using localhost MCP server: {url}")
+                return {
+                    "url": f"{url}/mcp",
+                    "transport": "streamable-http",
+                    "headers": headers,
+                }
+            logger.warning(f"MCP server is not running or not responding at {url}")
+
+        return None

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/src/datarobot_genai/core/utils/auth.py

@@ -16,9 +16,14 @@ import warnings
 from typing import Any
 
 import jwt
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
+from datarobot.auth.identity import Identity
+from datarobot.auth.oauth import AsyncOAuthComponent
 from datarobot.auth.session import AuthCtx
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from datarobot.models.genai.agent.auth import ToolAuth
 from datarobot.models.genai.agent.auth import get_authorization_context
+from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +32,13 @@ class AuthContextConfig(DataRobotAppFrameworkBaseSettings):
     session_secret_key: str = ""
 
 
+class DRAppCtx(BaseModel):
+    """DataRobot application context from authorization metadata."""
+
+    email: str | None = None
+    api_key: str | None = None
+
+
 class AuthContextHeaderHandler:
     """Manages encoding and decoding of authorization context into JWT tokens.
 
@@ -146,6 +158,7 @@ class AuthContextHeaderHandler:
 
         auth_ctx_dict = self.decode(token)
         if not auth_ctx_dict:
+            logger.debug("Failed to decode auth context from token")
             return None
 
         try:
@@ -153,3 +166,54 @@ class AuthContextHeaderHandler:
         except Exception as e:
             logger.error(f"Failed to create AuthCtx from decoded token: {e}", exc_info=True)
             return None
+
+
+class AsyncOAuthTokenProvider:
+    """Manages OAuth access tokens using generic OAuth client."""
+
+    def __init__(self, auth_ctx: AuthCtx) -> None:
+        self.auth_ctx = auth_ctx
+        self.oauth_client = self._create_oauth_client()
+
+    def _get_identity(self, provider_type: str | None) -> Identity:
+        """Retrieve the appropriate identity from the authentication context."""
+        identities = [x for x in self.auth_ctx.identities if x.provider_identity_id is not None]
+
+        if not identities:
+            raise ValueError("No identities found in authorization context.")
+
+        if provider_type is None:
+            if len(identities) > 1:
+                raise ValueError(
+                    "Multiple identities found. Please specify 'provider_type' parameter."
+                )
+            return identities[0]
+
+        identity = next((id for id in identities if id.provider_type == provider_type), None)
+
+        if identity is None:
+            raise ValueError(f"No identity found for provider '{provider_type}'.")
+
+        return identity
+
+    async def get_token(self, auth_type: ToolAuth, provider_type: str | None = None) -> str:
+        """Get OAuth access token using the specified method."""
+        if auth_type != ToolAuth.OBO:
+            raise ValueError(
+                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
+            )
+
+        identity = self._get_identity(provider_type)
+        token_data = await self.oauth_client.refresh_access_token(
+            identity_id=identity.provider_identity_id
+        )
+        return token_data.access_token
+
+    def _create_oauth_client(self) -> AsyncOAuthComponent:
+        """Create either DataRobot or Authlib OAuth client based on
+        authorization context.
+
+        Note: at the moment, only DataRobot OAuth client is supported.
+        """
+        logger.debug("Using DataRobot OAuth client")
+        return DatarobotAsyncOAuthClient()

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/src/datarobot_genai/crewai/base.py

@@ -80,6 +80,37 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
         """
         raise NotImplementedError
 
+    def _extract_pipeline_interactions(self) -> MultiTurnSample | None:
+        """Extract pipeline interactions from event listener if available."""
+        if not hasattr(self, "event_listener"):
+            return None
+        try:
+            listener = getattr(self, "event_listener", None)
+            messages = getattr(listener, "messages", None) if listener is not None else None
+            return create_pipeline_interactions_from_messages(messages)
+        except Exception:
+            return None
+
+    def _extract_usage_metrics(self, crew_output: Any) -> UsageMetrics:
+        """Extract usage metrics from crew output."""
+        token_usage = getattr(crew_output, "token_usage", None)
+        if token_usage is not None:
+            return {
+                "completion_tokens": int(getattr(token_usage, "completion_tokens", 0)),
+                "prompt_tokens": int(getattr(token_usage, "prompt_tokens", 0)),
+                "total_tokens": int(getattr(token_usage, "total_tokens", 0)),
+            }
+        return default_usage_metrics()
+
+    def _process_crew_output(
+        self, crew_output: Any
+    ) -> tuple[str, MultiTurnSample | None, UsageMetrics]:
+        """Process crew output into response tuple."""
+        response_text = str(crew_output.raw)
+        pipeline_interactions = self._extract_pipeline_interactions()
+        usage_metrics = self._extract_usage_metrics(crew_output)
+        return response_text, pipeline_interactions, usage_metrics
+
     async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
         """Run the CrewAI workflow with the provided completion parameters."""
         user_prompt_content = extract_user_prompt_content(completion_create_params)
@@ -92,9 +123,8 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
 
         # Use MCP context manager to handle connection lifecycle
         with mcp_tools_context(
-            api_base=self.api_base,
-            api_key=self.api_key,
             authorization_context=self._authorization_context,
+            forwarded_headers=self.forwarded_headers,
         ) as mcp_tools:
             # Set MCP tools for all agents if MCP is not configured this is effectively a no-op
             self.set_mcp_tools(mcp_tools)
@@ -117,64 +147,13 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
                 async def _gen() -> AsyncGenerator[
                     tuple[str, MultiTurnSample | None, UsageMetrics]
                 ]:
-                    # Run kickoff in a worker thread.
                     crew_output = await asyncio.to_thread(
                         crew.kickoff,
                         inputs=self.make_kickoff_inputs(user_prompt_content),
                     )
-
-                    pipeline_interactions = None
-                    if hasattr(self, "event_listener"):
-                        try:
-                            listener = getattr(self, "event_listener", None)
-                            messages = (
-                                getattr(listener, "messages", None)
-                                if listener is not None
-                                else None
-                            )
-                            pipeline_interactions = create_pipeline_interactions_from_messages(
-                                messages
-                            )
-                        except Exception:
-                            pipeline_interactions = None
-
-                    token_usage = getattr(crew_output, "token_usage", None)
-                    if token_usage is not None:
-                        usage_metrics: UsageMetrics = {
-                            "completion_tokens": int(getattr(token_usage, "completion_tokens", 0)),
-                            "prompt_tokens": int(getattr(token_usage, "prompt_tokens", 0)),
-                            "total_tokens": int(getattr(token_usage, "total_tokens", 0)),
-                        }
-                    else:
-                        usage_metrics = default_usage_metrics()
-
-                    # Finalize stream with empty chunk carrying interactions and usage
-                    yield "", pipeline_interactions, usage_metrics
+                    yield self._process_crew_output(crew_output)
 
                 return _gen()
 
-            # Non-streaming: run to completion and return final result
             crew_output = crew.kickoff(inputs=self.make_kickoff_inputs(user_prompt_content))
-
-            response_text = str(crew_output.raw)
-
-            pipeline_interactions = None
-            if hasattr(self, "event_listener"):
-                try:
-                    listener = getattr(self, "event_listener", None)
-                    messages = getattr(listener, "messages", None) if listener is not None else None
-                    pipeline_interactions = create_pipeline_interactions_from_messages(messages)
-                except Exception:
-                    pipeline_interactions = None
-
-            token_usage = getattr(crew_output, "token_usage", None)
-            if token_usage is not None:
-                usage_metrics: UsageMetrics = {
-                    "completion_tokens": int(getattr(token_usage, "completion_tokens", 0)),
-                    "prompt_tokens": int(getattr(token_usage, "prompt_tokens", 0)),
-                    "total_tokens": int(getattr(token_usage, "total_tokens", 0)),
-                }
-            else:
-                usage_metrics = default_usage_metrics()
-
-            return response_text, pipeline_interactions, usage_metrics
+            return self._process_crew_output(crew_output)

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/src/datarobot_genai/crewai/mcp.py

@@ -29,15 +29,14 @@ from datarobot_genai.core.mcp.common import MCPConfig
 
 @contextmanager
 def mcp_tools_context(
-    api_base: str | None = None,
-    api_key: str | None = None,
     authorization_context: dict[str, Any] | None = None,
+    forwarded_headers: dict[str, str] | None = None,
 ) -> Generator[list[Any], None, None]:
     """Context manager for MCP tools that handles connection lifecycle."""
     config = MCPConfig(
-        api_base=api_base, api_key=api_key, authorization_context=authorization_context
+        authorization_context=authorization_context,
+        forwarded_headers=forwarded_headers,
     )
-
     # If no MCP server configured, return empty tools list
     if not config.server_config:
         print("No MCP server configured, using empty tools list", flush=True)
@@ -47,10 +46,8 @@ def mcp_tools_context(
     print(f"Connecting to MCP server: {config.server_config['url']}", flush=True)
 
     # Use MCPServerAdapter as context manager with the server config
-    adapter_setting = config.server_config.copy()
-    adapter_setting["transport"] = "streamable-http"
     try:
-        with MCPServerAdapter(adapter_setting) as tools:
+        with MCPServerAdapter(config.server_config) as tools:
             print(
                 f"Successfully connected to MCP server, got {len(tools)} tools",
                 flush=True,

{datarobot_genai-0.1.59 → datarobot_genai-0.1.70}/src/datarobot_genai/drmcp/core/auth.py

@@ -18,7 +18,6 @@ import logging
 from typing import Any
 
 from datarobot.auth.session import AuthCtx
-from datarobot.models.genai.agent.auth import OAuthAccessTokenProvider
 from datarobot.models.genai.agent.auth import ToolAuth
 from fastmcp.server.dependencies import get_context
 from fastmcp.server.dependencies import get_http_headers
@@ -27,12 +26,15 @@ from fastmcp.server.middleware import Middleware
 from fastmcp.server.middleware import MiddlewareContext
 from fastmcp.tools.tool import ToolResult
 
+from datarobot_genai.core.utils.auth import AsyncOAuthTokenProvider
 from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
-from datarobot_genai.drmcp import get_config
 
 logger = logging.getLogger(__name__)
 
 
+AUTH_CTX_KEY = "authorization_context"
+
+
 class OAuthMiddleWare(Middleware):
     """Middleware that parses `x-datarobot-authorization-context` for tool calls.
 
@@ -45,16 +47,8 @@ class OAuthMiddleWare(Middleware):
         Handler for encoding/decoding JWT tokens containing auth context.
     """
 
-    def __init__(self, secret_key: str | None = None) -> None:
-        """Initialize the middleware with authentication handler.
-
-        Parameters
-        ----------
-        secret_key : Optional[str]
-            Secret key for JWT validation. If None, uses the value from config.
-        """
-        secret_key = secret_key or get_config().session_secret_key
-        self.auth_handler = AuthContextHeaderHandler(secret_key)
+    def __init__(self, auth_handler: AuthContextHeaderHandler | None = None) -> None:
+        self.auth_handler = auth_handler or AuthContextHeaderHandler()
 
     async def on_call_tool(
         self, context: MiddlewareContext, call_next: CallNext[Any, ToolResult]
@@ -74,9 +68,12 @@ class OAuthMiddleWare(Middleware):
             The result from the tool execution.
         """
         auth_context = self._extract_auth_context()
+        if not auth_context:
+            logger.debug("No valid authorization context extracted from request headers.")
 
         if context.fastmcp_context is not None:
-            context.fastmcp_context.auth_context = auth_context
+            context.fastmcp_context.set_state(AUTH_CTX_KEY, auth_context)
+            logger.debug("Authorization context attached to state.")
 
         return await call_next(context)
 
@@ -99,8 +96,8 @@ class OAuthMiddleWare(Middleware):
             return None
 
 
-async def get_auth_context() -> AuthCtx:
-    """Retrieve the AuthCtx from the current request context, if available.
+async def must_get_auth_context() -> AuthCtx:
+    """Retrieve the AuthCtx from the current request context or raise error.
 
     Raises
     ------
@@ -113,14 +110,15 @@ async def get_auth_context() -> AuthCtx:
         The authorization context associated with the current request.
     """
     context = get_context()
-    auth_ctx = getattr(context, "auth_context", None)
+
+    auth_ctx = context.get_state(AUTH_CTX_KEY)
     if not auth_ctx:
-        raise RuntimeError("No authorization context found.")
+        raise RuntimeError("Could not retrieve authorization context from FastMCP context state.")
 
     return auth_ctx
 
 
-async def get_access_token(provider: str | None = None) -> str:
+async def get_access_token(provider_type: str | None = None) -> str:
     """Retrieve access token from the DataRobot OAuth Provider Service.
 
     OAuth access tokens can be retrieved only for providers where the user completed
@@ -132,7 +130,7 @@ async def get_access_token(provider: str | None = None) -> str:
 
     Parameters
     ----------
-    provider : str, optional
+    provider_type : str, optional
         The name of the OAuth provider. It should match the name of the provider configured
         during provider setup. If no value is provided and only one OAuth provider exists, that
         provider will be used. If multiple providers exist and none is specified, an error will be
@@ -142,12 +140,18 @@ async def get_access_token(provider: str | None = None) -> str:
     -------
     The oauth access token.
     """
-    token_provider = OAuthAccessTokenProvider(await get_auth_context())
-    access_token = token_provider.get_token(ToolAuth.OBO, provider)
-    return access_token
+    auth_ctx = await must_get_auth_context()
+    logger.debug("Retrieved authorization context")
+
+    oauth_token_provider = AsyncOAuthTokenProvider(auth_ctx)
+    oauth_access_token = await oauth_token_provider.get_token(
+        auth_type=ToolAuth.OBO,
+        provider_type=provider_type,
+    )
+    return oauth_access_token
 
 
-def initialize_oauth_middleware(mcp: Any, secret_key: str | None = None) -> None:
+def initialize_oauth_middleware(mcp: Any) -> None:
     """Initialize and register OAuth middleware with the MCP server.
 
     Parameters
@@ -157,6 +161,5 @@ def initialize_oauth_middleware(mcp: Any, secret_key: str | None = None) -> None
     secret_key : Optional[str]
         Secret key for JWT validation. If None, uses the value from config.
     """
-    middleware = OAuthMiddleWare(secret_key=secret_key)
-    mcp.add_middleware(middleware)
+    mcp.add_middleware(OAuthMiddleWare())
     logger.info("OAuth middleware registered successfully")