agno 2.2.13__py3-none-any.whl

agno/os/mcp.py

@@ -0,0 +1,292 @@
+"""Router for MCP interface providing Model Context Protocol endpoints."""
+
+import logging
+from typing import TYPE_CHECKING, List, Optional, cast
+from uuid import uuid4
+
+from fastmcp import FastMCP
+from fastmcp.server.http import (
+    StarletteWithLifespan,
+)
+
+from agno.db.base import AsyncBaseDb, SessionType
+from agno.db.schemas import UserMemory
+from agno.os.routers.memory.schemas import (
+    UserMemorySchema,
+)
+from agno.os.schema import (
+    AgentSummaryResponse,
+    ConfigResponse,
+    InterfaceResponse,
+    SessionSchema,
+    TeamSummaryResponse,
+    WorkflowSummaryResponse,
+)
+from agno.os.utils import (
+    get_agent_by_id,
+    get_db,
+    get_team_by_id,
+    get_workflow_by_id,
+)
+from agno.run.agent import RunOutput
+from agno.run.team import TeamRunOutput
+from agno.run.workflow import WorkflowRunOutput
+
+if TYPE_CHECKING:
+    from agno.os.app import AgentOS
+
+logger = logging.getLogger(__name__)
+
+
+def get_mcp_server(
+    os: "AgentOS",
+) -> StarletteWithLifespan:
+    """Attach MCP routes to the provided router."""
+
+    # Create an MCP server
+    mcp = FastMCP(os.name or "AgentOS")
+
+    @mcp.tool(
+        name="get_agentos_config",
+        description="Get the configuration of the AgentOS",
+        tags={"core"},
+        output_schema=ConfigResponse.model_json_schema(),
+    )  # type: ignore
+    async def config() -> ConfigResponse:
+        return ConfigResponse(
+            os_id=os.id or "AgentOS",
+            description=os.description,
+            available_models=os.config.available_models if os.config else [],
+            databases=[db.id for db_list in os.dbs.values() for db in db_list],
+            chat=os.config.chat if os.config else None,
+            session=os._get_session_config(),
+            memory=os._get_memory_config(),
+            knowledge=os._get_knowledge_config(),
+            evals=os._get_evals_config(),
+            metrics=os._get_metrics_config(),
+            agents=[AgentSummaryResponse.from_agent(agent) for agent in os.agents] if os.agents else [],
+            teams=[TeamSummaryResponse.from_team(team) for team in os.teams] if os.teams else [],
+            workflows=[WorkflowSummaryResponse.from_workflow(w) for w in os.workflows] if os.workflows else [],
+            interfaces=[
+                InterfaceResponse(type=interface.type, version=interface.version, route=interface.prefix)
+                for interface in os.interfaces
+            ],
+        )
+
+    @mcp.tool(name="run_agent", description="Run an agent", tags={"core"})  # type: ignore
+    async def run_agent(agent_id: str, message: str) -> RunOutput:
+        agent = get_agent_by_id(agent_id, os.agents)
+        if agent is None:
+            raise Exception(f"Agent {agent_id} not found")
+        return await agent.arun(message)
+
+    @mcp.tool(name="run_team", description="Run a team", tags={"core"})  # type: ignore
+    async def run_team(team_id: str, message: str) -> TeamRunOutput:
+        team = get_team_by_id(team_id, os.teams)
+        if team is None:
+            raise Exception(f"Team {team_id} not found")
+        return await team.arun(message)
+
+    @mcp.tool(name="run_workflow", description="Run a workflow", tags={"core"})  # type: ignore
+    async def run_workflow(workflow_id: str, message: str) -> WorkflowRunOutput:
+        workflow = get_workflow_by_id(workflow_id, os.workflows)
+        if workflow is None:
+            raise Exception(f"Workflow {workflow_id} not found")
+        return await workflow.arun(message)
+
+    # Session Management Tools
+    @mcp.tool(name="get_sessions_for_agent", description="Get list of sessions for an agent", tags={"session"})  # type: ignore
+    async def get_sessions_for_agent(
+        agent_id: str,
+        db_id: str,
+        user_id: Optional[str] = None,
+        sort_by: str = "created_at",
+        sort_order: str = "desc",
+    ):
+        db = await get_db(os.dbs, db_id)
+        if isinstance(db, AsyncBaseDb):
+            db = cast(AsyncBaseDb, db)
+            sessions = await db.get_sessions(
+                session_type=SessionType.AGENT,
+                component_id=agent_id,
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+        else:
+            sessions = db.get_sessions(
+                session_type=SessionType.AGENT,
+                component_id=agent_id,
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+
+        return {
+            "data": [SessionSchema.from_dict(session) for session in sessions],  # type: ignore
+        }
+
+    @mcp.tool(name="get_sessions_for_team", description="Get list of sessions for a team", tags={"session"})  # type: ignore
+    async def get_sessions_for_team(
+        team_id: str,
+        db_id: str,
+        user_id: Optional[str] = None,
+        sort_by: str = "created_at",
+        sort_order: str = "desc",
+    ):
+        db = await get_db(os.dbs, db_id)
+        if isinstance(db, AsyncBaseDb):
+            db = cast(AsyncBaseDb, db)
+            sessions = await db.get_sessions(
+                session_type=SessionType.TEAM,
+                component_id=team_id,
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+        else:
+            sessions = db.get_sessions(
+                session_type=SessionType.TEAM,
+                component_id=team_id,
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+
+        return {
+            "data": [SessionSchema.from_dict(session) for session in sessions],  # type: ignore
+        }
+
+    @mcp.tool(name="get_sessions_for_workflow", description="Get list of sessions for a workflow", tags={"session"})  # type: ignore
+    async def get_sessions_for_workflow(
+        workflow_id: str,
+        db_id: str,
+        user_id: Optional[str] = None,
+        sort_by: str = "created_at",
+        sort_order: str = "desc",
+    ):
+        db = await get_db(os.dbs, db_id)
+        if isinstance(db, AsyncBaseDb):
+            db = cast(AsyncBaseDb, db)
+            sessions = await db.get_sessions(
+                session_type=SessionType.WORKFLOW,
+                component_id=workflow_id,
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+        else:
+            sessions = db.get_sessions(
+                session_type=SessionType.WORKFLOW,
+                component_id=workflow_id,
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+
+        return {
+            "data": [SessionSchema.from_dict(session) for session in sessions],  # type: ignore
+        }
+
+    # Memory Management Tools
+    @mcp.tool(name="create_memory", description="Create a new user memory", tags={"memory"})  # type: ignore
+    async def create_memory(
+        db_id: str,
+        memory: str,
+        user_id: str,
+        topics: Optional[List[str]] = None,
+    ) -> UserMemorySchema:
+        db = await get_db(os.dbs, db_id)
+        user_memory = db.upsert_user_memory(
+            memory=UserMemory(
+                memory_id=str(uuid4()),
+                memory=memory,
+                topics=topics or [],
+                user_id=user_id,
+            ),
+            deserialize=False,
+        )
+        if not user_memory:
+            raise Exception("Failed to create memory")
+
+        return UserMemorySchema.from_dict(user_memory)  # type: ignore
+
+    @mcp.tool(name="get_memories_for_user", description="Get a list of memories for a user", tags={"memory"})  # type: ignore
+    async def get_memories_for_user(
+        user_id: str,
+        sort_by: str = "updated_at",
+        sort_order: str = "desc",
+        db_id: Optional[str] = None,
+    ):
+        db = await get_db(os.dbs, db_id)
+        if isinstance(db, AsyncBaseDb):
+            db = cast(AsyncBaseDb, db)
+            user_memories = await db.get_user_memories(
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+        else:
+            user_memories = db.get_user_memories(
+                user_id=user_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                deserialize=False,
+            )
+        return {
+            "data": [UserMemorySchema.from_dict(user_memory) for user_memory in user_memories],  # type: ignore
+        }
+
+    @mcp.tool(name="update_memory", description="Update a memory", tags={"memory"})  # type: ignore
+    async def update_memory(
+        db_id: str,
+        memory_id: str,
+        memory: str,
+        user_id: str,
+    ) -> UserMemorySchema:
+        db = await get_db(os.dbs, db_id)
+        if isinstance(db, AsyncBaseDb):
+            db = cast(AsyncBaseDb, db)
+            user_memory = await db.upsert_user_memory(
+                memory=UserMemory(
+                    memory_id=memory_id,
+                    memory=memory,
+                    user_id=user_id,
+                ),
+                deserialize=False,
+            )
+        else:
+            user_memory = db.upsert_user_memory(
+                memory=UserMemory(
+                    memory_id=memory_id,
+                    memory=memory,
+                    user_id=user_id,
+                ),
+                deserialize=False,
+            )
+        if not user_memory:
+            raise Exception("Failed to update memory")
+
+        return UserMemorySchema.from_dict(user_memory)  # type: ignore
+
+    @mcp.tool(name="delete_memory", description="Delete a memory by ID", tags={"memory"})  # type: ignore
+    async def delete_memory(
+        db_id: str,
+        memory_id: str,
+    ) -> None:
+        db = await get_db(os.dbs, db_id)
+        if isinstance(db, AsyncBaseDb):
+            db = cast(AsyncBaseDb, db)
+            await db.delete_user_memory(memory_id=memory_id)
+        else:
+            db.delete_user_memory(memory_id=memory_id)
+
+    mcp_app = mcp.http_app(path="/mcp")
+    return mcp_app

agno/os/middleware/__init__.py

@@ -0,0 +1,7 @@
+from agno.os.middleware.jwt import (
+    JWTMiddleware,
+)
+
+__all__ = [
+    "JWTMiddleware",
+]

agno/os/middleware/jwt.py

@@ -0,0 +1,233 @@
+import fnmatch
+from enum import Enum
+from os import getenv
+from typing import List, Optional
+
+import jwt
+from fastapi import Request, Response
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from agno.utils.log import log_debug
+
+
+class TokenSource(str, Enum):
+    """Enum for JWT token source options."""
+
+    HEADER = "header"
+    COOKIE = "cookie"
+    BOTH = "both"  # Try header first, then cookie
+
+
+class JWTMiddleware(BaseHTTPMiddleware):
+    """
+    JWT Middleware for validating tokens and storing JWT claims in request state.
+
+    This middleware:
+    1. Extracts JWT token from Authorization header, cookies, or both
+    2. Decodes and validates the token
+    3. Stores JWT claims in request.state for easy access in endpoints
+
+    Token Sources:
+    - "header": Extract from Authorization header (default)
+    - "cookie": Extract from HTTP cookie
+    - "both": Try header first, then cookie as fallback
+
+    Claims are stored as:
+    - request.state.user_id: User ID from configured claim
+    - request.state.session_id: Session ID from configured claim
+    - request.state.dependencies: Dictionary of dependency claims
+    - request.state.session_state: Dictionary of session state claims
+    - request.state.authenticated: Boolean authentication status
+
+    """
+
+    def __init__(
+        self,
+        app,
+        secret_key: Optional[str] = None,
+        algorithm: str = "HS256",
+        token_source: TokenSource = TokenSource.HEADER,
+        token_header_key: str = "Authorization",
+        cookie_name: str = "access_token",
+        validate: bool = True,
+        excluded_route_paths: Optional[List[str]] = None,
+        scopes_claim: Optional[str] = None,
+        user_id_claim: str = "sub",
+        session_id_claim: str = "session_id",
+        dependencies_claims: Optional[List[str]] = None,
+        session_state_claims: Optional[List[str]] = None,
+    ):
+        """
+        Initialize the JWT middleware.
+
+        Args:
+            app: The FastAPI app instance
+            secret_key: The secret key to use for JWT validation (optional, will use JWT_SECRET_KEY environment variable if not provided)
+            algorithm: The algorithm to use for JWT validation
+            token_header_key: The key to use for the Authorization header (only used when token_source is header)
+            token_source: Where to extract the JWT token from (header, cookie, or both)
+            cookie_name: The name of the cookie containing the JWT token (only used when token_source is cookie/both)
+            validate: Whether to validate the JWT token
+            excluded_route_paths: A list of route paths to exclude from JWT validation
+            scopes_claim: The claim to use for scopes extraction
+            user_id_claim: The claim to use for user ID extraction
+            session_id_claim: The claim to use for session ID extraction
+            dependencies_claims: A list of claims to extract from the JWT token for dependencies
+            session_state_claims: A list of claims to extract from the JWT token for session state
+        """
+        super().__init__(app)
+        self.secret_key = secret_key or getenv("JWT_SECRET_KEY")
+        if not self.secret_key:
+            raise ValueError("Secret key is required")
+        self.algorithm = algorithm
+        self.token_header_key = token_header_key
+        self.token_source = token_source
+        self.cookie_name = cookie_name
+        self.validate = validate
+        self.excluded_route_paths = excluded_route_paths
+        self.scopes_claim = scopes_claim
+        self.user_id_claim = user_id_claim
+        self.session_id_claim = session_id_claim
+        self.dependencies_claims = dependencies_claims or []
+        self.session_state_claims = session_state_claims or []
+
+    def _extract_token_from_header(self, request: Request) -> Optional[str]:
+        """Extract JWT token from Authorization header."""
+        authorization = request.headers.get(self.token_header_key, "")
+        if not authorization:
+            return None
+
+        try:
+            # Remove the "Bearer " prefix (if present)
+            _, token = authorization.split(" ", 1)
+            return token
+        except ValueError:
+            return None
+
+    def _extract_token_from_cookie(self, request: Request) -> Optional[str]:
+        """Extract JWT token from cookie."""
+        return request.cookies.get(self.cookie_name)
+
+    def _extract_token(self, request: Request) -> Optional[str]:
+        """Extract JWT token based on configured token source."""
+        if self.token_source == TokenSource.HEADER:
+            return self._extract_token_from_header(request)
+        elif self.token_source == TokenSource.COOKIE:
+            return self._extract_token_from_cookie(request)
+        elif self.token_source == TokenSource.BOTH:
+            # Try header first, then cookie
+            token = self._extract_token_from_header(request)
+            if token is None:
+                token = self._extract_token_from_cookie(request)
+            return token
+        else:
+            log_debug(f"Unknown token source: {self.token_source}")
+            return None
+
+    def _get_missing_token_error_message(self) -> str:
+        """Get appropriate error message for missing token based on token source."""
+        if self.token_source == TokenSource.HEADER:
+            return "Authorization header missing"
+        elif self.token_source == TokenSource.COOKIE:
+            return f"JWT cookie '{self.cookie_name}' missing"
+        elif self.token_source == TokenSource.BOTH:
+            return f"JWT token missing from both Authorization header and '{self.cookie_name}' cookie"
+        else:
+            return "JWT token missing"
+
+    def _is_route_excluded(self, path: str) -> bool:
+        """Check if a route path matches any of the excluded patterns."""
+        if not self.excluded_route_paths:
+            return False
+
+        for excluded_path in self.excluded_route_paths:
+            # Support both exact matches and wildcard patterns
+            if fnmatch.fnmatch(path, excluded_path):
+                return True
+
+        return False
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        if self._is_route_excluded(request.url.path):
+            return await call_next(request)
+
+        # Extract JWT token from configured source (header, cookie, or both)
+        token = self._extract_token(request)
+
+        if not token:
+            if self.validate:
+                error_msg = self._get_missing_token_error_message()
+                return JSONResponse(status_code=401, content={"detail": error_msg})
+            return await call_next(request)
+
+        # Decode JWT token
+        try:
+            payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])  # type: ignore
+
+            # Extract scopes claims
+            scopes = []
+            if self.scopes_claim in payload:
+                extracted_scopes = payload[self.scopes_claim]
+                if isinstance(extracted_scopes, str):
+                    scopes = extracted_scopes.split(" ")
+                else:
+                    scopes = extracted_scopes
+            if scopes:
+                request.state.scopes = scopes
+
+            # Extract user information
+            if self.user_id_claim in payload:
+                user_id = payload[self.user_id_claim]
+                request.state.user_id = user_id
+            if self.session_id_claim in payload:
+                session_id = payload[self.session_id_claim]
+                request.state.session_id = session_id
+            else:
+                session_id = None
+
+            # Extract dependency claims
+            dependencies = {}
+            for claim in self.dependencies_claims:
+                if claim in payload:
+                    dependencies[claim] = payload[claim]
+
+            if dependencies:
+                request.state.dependencies = dependencies
+
+            # Extract session state claims
+            session_state = {}
+            for claim in self.session_state_claims:
+                if claim in payload:
+                    session_state[claim] = payload[claim]
+
+            if session_state:
+                request.state.session_state = session_state
+
+            request.state.token = token
+            request.state.authenticated = True
+
+            log_debug(f"JWT decoded successfully for user: {user_id}")
+            if dependencies:
+                log_debug(f"Extracted dependencies: {dependencies}")
+            if session_state:
+                log_debug(f"Extracted session state: {session_state}")
+
+        except jwt.ExpiredSignatureError:
+            if self.validate:
+                return JSONResponse(status_code=401, content={"detail": "Token has expired"})
+            request.state.authenticated = False
+            request.state.token = token
+
+        except jwt.InvalidTokenError as e:
+            if self.validate:
+                return JSONResponse(status_code=401, content={"detail": f"Invalid token: {str(e)}"})
+            request.state.authenticated = False
+            request.state.token = token
+        except Exception as e:
+            if self.validate:
+                return JSONResponse(status_code=401, content={"detail": f"Error decoding token: {str(e)}"})
+            request.state.authenticated = False
+            request.state.token = token
+
+        return await call_next(request)