diff --git a/agent_memory_server/config.py b/agent_memory_server/config.py index 95464d0..5af59ab 100644 --- a/agent_memory_server/config.py +++ b/agent_memory_server/config.py @@ -427,6 +427,7 @@ class Settings(BaseSettings): oauth2_audience: str | None = None oauth2_jwks_url: str | None = None oauth2_algorithms: list[str] = ["RS256"] + oauth2_resource_host: str | None = None # Token Authentication settings token_auth_enabled: bool = False diff --git a/agent_memory_server/main.py b/agent_memory_server/main.py index b1b3cb5..163595e 100644 --- a/agent_memory_server/main.py +++ b/agent_memory_server/main.py @@ -130,6 +130,17 @@ async def lifespan(app: FastAPI): app.include_router(memory_router) +@app.get("/.well-known/oauth-authorization-server") +async def oauth_authorization_server_metadata(): + """RFC 9728 Protected Resource Metadata for MCP OIDC discovery.""" + if not settings.oauth2_issuer_url: + return {"error": "OIDC not configured"} + return { + "resource": f"https://{settings.oauth2_resource_host or 'localhost'}", + "authorization_servers": [settings.oauth2_issuer_url], + } + + def on_start_logger(port: int): """Log startup information""" print("\n-----------------------------------") diff --git a/agent_memory_server/mcp.py b/agent_memory_server/mcp.py index 5bf55e0..e416a25 100644 --- a/agent_memory_server/mcp.py +++ b/agent_memory_server/mcp.py @@ -3,6 +3,8 @@ from typing import Any import ulid +from mcp.server.auth.provider import AccessToken, TokenVerifier +from mcp.server.auth.settings import AuthSettings from mcp.server.fastmcp import FastMCP as _FastMCPBase from agent_memory_server import working_memory as working_memory_core @@ -179,22 +181,16 @@ async def run_sse_async(self): ).serve() def streamable_http_app(self): - """Return a Starlette app for streamable-http with namespace routing.""" - from mcp.server.streamable_http_manager import StreamableHTTPSessionManager - from starlette.applications import Starlette + """Return a Starlette app for streamable-http with namespace routing. + + Extends the parent's streamable_http_app() by adding a namespace-prefixed + route while preserving auth middleware and protected resource routes. + """ from starlette.requests import Request from starlette.routing import Route - if self._session_manager is None: - self._session_manager = StreamableHTTPSessionManager( - app=self._mcp_server, - event_store=self._event_store, - retry_interval=self._retry_interval, - json_response=self.settings.json_response, - stateless=self.settings.stateless_http, - security_settings=self.settings.transport_security, - ) - + # Get the parent's app which includes auth middleware and well-known routes + app = super().streamable_http_app() session_manager = self._session_manager mcp_instance = self @@ -212,14 +208,9 @@ async def __call__(self, scope, receive, send): handler = _NamespaceAwareHandler() path = self.settings.streamable_http_path - return Starlette( - debug=self.settings.debug, - routes=[ - Route(path, endpoint=handler), - Route(f"/{{namespace}}{path}", endpoint=handler), - ], - lifespan=lambda app: session_manager.run(), - ) + # Add namespace-prefixed route to the existing app + app.routes.append(Route(f"/{{namespace}}{path}", endpoint=handler)) + return app async def run_streamable_http_async(self): """Start streamable HTTP MCP server.""" @@ -255,12 +246,51 @@ async def run_stdio_async(self): """ +class JWTTokenVerifier(TokenVerifier): + """Verify JWT tokens from Hydra using the existing auth module.""" + + async def verify_token(self, token: str) -> AccessToken | None: + from agent_memory_server.auth import verify_jwt + + try: + user_info = verify_jwt(token) + scopes = user_info.scope.split() if user_info.scope else [] + return AccessToken( + token=token, + client_id=user_info.sub, + scopes=scopes, + expires_at=user_info.exp, + ) + except Exception: + return None + + +def _build_mcp_auth_kwargs() -> dict[str, Any]: + """Build auth kwargs for FastMCP if OAuth2 is configured.""" + if ( + settings.auth_mode != "oauth2" + or not settings.oauth2_issuer_url + or not settings.oauth2_resource_host + ): + return {} + + resource_url = f"https://{settings.oauth2_resource_host}" + return { + "auth": AuthSettings( + issuer_url=settings.oauth2_issuer_url, + resource_server_url=resource_url, + ), + "token_verifier": JWTTokenVerifier(), + } + + mcp_app = FastMCP( "Redis Agent Memory Server", host=settings.mcp_host, port=settings.mcp_port, instructions=INSTRUCTIONS, default_namespace=settings.default_mcp_namespace, + **_build_mcp_auth_kwargs(), )