agno 2.0.10__py3-none-any.whl → 2.1.0__py3-none-any.whl

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)

agno/os/router.py

@@ -16,6 +16,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 
 from agno.agent.agent import Agent
+from agno.exceptions import InputCheckError, OutputCheckError
 from agno.media import Audio, Image, Video
 from agno.media import File as FileMedia
 from agno.os.auth import get_authentication_dependency, validate_websocket_token
@@ -74,7 +75,6 @@ async def _get_request_kwargs(request: Request, endpoint_func: Callable) -> Dict
     kwargs = {key: value for key, value in form_data.items() if key not in known_fields}
 
     # Handle JSON parameters. They are passed as strings and need to be deserialized.
-
     if session_state := kwargs.get("session_state"):
         try:
             session_state_dict = json.loads(session_state)  # type: ignore
@@ -82,6 +82,7 @@ async def _get_request_kwargs(request: Request, endpoint_func: Callable) -> Dict
         except json.JSONDecodeError:
             kwargs.pop("session_state")
             log_warning(f"Invalid session_state parameter couldn't be loaded: {session_state}")
+
     if dependencies := kwargs.get("dependencies"):
         try:
             dependencies_dict = json.loads(dependencies)  # type: ignore
@@ -245,7 +246,14 @@ async def agent_response_streamer(
         )
         async for run_response_chunk in run_response:
             yield format_sse_event(run_response_chunk)  # type: ignore
-
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = RunErrorEvent(
+            content=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
     except Exception as e:
         import traceback
 
@@ -274,6 +282,14 @@ async def agent_continue_response_streamer(
         )
         async for run_response_chunk in continue_response:
             yield format_sse_event(run_response_chunk)  # type: ignore
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = RunErrorEvent(
+            content=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
 
     except Exception as e:
         import traceback
@@ -281,6 +297,8 @@ async def agent_continue_response_streamer(
         traceback.print_exc(limit=3)
         error_response = RunErrorEvent(
             content=str(e),
+            error_type=e.type if hasattr(e, "type") else None,
+            error_id=e.error_id if hasattr(e, "error_id") else None,
         )
         yield format_sse_event(error_response)
         return
@@ -313,6 +331,14 @@ async def team_response_streamer(
         )
         async for run_response_chunk in run_response:
             yield format_sse_event(run_response_chunk)  # type: ignore
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = TeamRunErrorEvent(
+            content=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
 
     except Exception as e:
         import traceback
@@ -320,6 +346,8 @@ async def team_response_streamer(
         traceback.print_exc()
         error_response = TeamRunErrorEvent(
             content=str(e),
+            error_type=e.type if hasattr(e, "type") else None,
+            error_id=e.error_id if hasattr(e, "error_id") else None,
         )
         yield format_sse_event(error_response)
         return
@@ -366,9 +394,28 @@ async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os:
 
         await websocket_manager.register_workflow_websocket(workflow_run_output.run_id, websocket)  # type: ignore
 
+    except (InputCheckError, OutputCheckError) as e:
+        await websocket.send_text(
+            json.dumps(
+                {
+                    "event": "error",
+                    "error": str(e),
+                    "error_type": e.type,
+                    "error_id": e.error_id,
+                    "additional_data": e.additional_data,
+                }
+            )
+        )
     except Exception as e:
         logger.error(f"Error executing workflow via WebSocket: {e}")
-        await websocket.send_text(json.dumps({"event": "error", "error": str(e)}))
+        error_payload = {
+            "event": "error",
+            "error": str(e),
+            "error_type": e.type if hasattr(e, "type") else None,
+            "error_id": e.error_id if hasattr(e, "error_id") else None,
+        }
+        error_payload = {k: v for k, v in error_payload.items() if v is not None}
+        await websocket.send_text(json.dumps(error_payload))
 
 
 async def workflow_response_streamer(
@@ -391,12 +438,23 @@ async def workflow_response_streamer(
         async for run_response_chunk in run_response:
             yield format_sse_event(run_response_chunk)  # type: ignore
 
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = WorkflowErrorEvent(
+            error=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
+
     except Exception as e:
         import traceback
 
         traceback.print_exc()
         error_response = WorkflowErrorEvent(
             error=str(e),
+            error_type=e.type if hasattr(e, "type") else None,
+            error_id=e.error_id if hasattr(e, "error_id") else None,
         )
         yield format_sse_event(error_response)
         return
@@ -464,7 +522,7 @@ def get_websocket_router(
                     await websocket.send_text(json.dumps({"event": "error", "error": f"Unknown action: {action}"}))
 
         except Exception as e:
-            if "1012" not in str(e):
+            if "1012" not in str(e) and "1001" not in str(e):
                 logger.error(f"WebSocket error: {e}")
         finally:
             # Clean up the websocket connection
@@ -520,7 +578,7 @@ def get_base_router(
                 "content": {
                     "application/json": {
                         "example": {
-                            "os_id": "demo",
+                            "id": "demo",
                             "description": "Example AgentOS configuration",
                             "available_models": [],
                             "databases": ["9c884dc4-9066-448c-9074-ef49ec7eb73c"],
@@ -582,7 +640,7 @@ def get_base_router(
     )
     async def config() -> ConfigResponse:
         return ConfigResponse(
-            os_id=os.os_id or "Unnamed OS",
+            os_id=os.id or "Unnamed OS",
             description=os.description,
             available_models=os.config.available_models if os.config else [],
             databases=[db.id for db in os.dbs.values()],
@@ -669,7 +727,7 @@ def get_base_router(
                 "content": {
                     "text/event-stream": {
                         "examples": {
-                            "event_strea": {
+                            "event_stream": {
                                 "summary": "Example event stream response",
                                 "value": 'event: RunStarted\ndata: {"content": "Hello!", "run_id": "123..."}\n\n',
                             }
@@ -692,6 +750,25 @@ def get_base_router(
     ):
         kwargs = await _get_request_kwargs(request, create_agent_run)
 
+        if hasattr(request.state, "user_id"):
+            if user_id:
+                log_warning("User ID parameter passed in both request state and kwargs, using request state")
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            if session_id:
+                log_warning("Session ID parameter passed in both request state and kwargs, using request state")
+            session_id = request.state.session_id
+        if hasattr(request.state, "session_state"):
+            session_state = request.state.session_state
+            if "session_state" in kwargs:
+                log_warning("Session state parameter passed in both request state and kwargs, using request state")
+            kwargs["session_state"] = session_state
+        if hasattr(request.state, "dependencies"):
+            dependencies = request.state.dependencies
+            if "dependencies" in kwargs:
+                log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
+            kwargs["dependencies"] = dependencies
+
         agent = get_agent_by_id(agent_id, os.agents)
         if agent is None:
             raise HTTPException(status_code=404, detail="Agent not found")
@@ -774,21 +851,25 @@ def get_base_router(
                 media_type="text/event-stream",
             )
         else:
-            run_response = cast(
-                RunOutput,
-                await agent.arun(
-                    input=message,
-                    session_id=session_id,
-                    user_id=user_id,
-                    images=base64_images if base64_images else None,
-                    audio=base64_audios if base64_audios else None,
-                    videos=base64_videos if base64_videos else None,
-                    files=input_files if input_files else None,
-                    stream=False,
-                    **kwargs,
-                ),
-            )
-            return run_response.to_dict()
+            try:
+                run_response = cast(
+                    RunOutput,
+                    await agent.arun(
+                        input=message,
+                        session_id=session_id,
+                        user_id=user_id,
+                        images=base64_images if base64_images else None,
+                        audio=base64_audios if base64_audios else None,
+                        videos=base64_videos if base64_videos else None,
+                        files=input_files if input_files else None,
+                        stream=False,
+                        **kwargs,
+                    ),
+                )
+                return run_response.to_dict()
+
+            except InputCheckError as e:
+                raise HTTPException(status_code=400, detail=str(e))
 
     @router.post(
         "/agents/{agent_id}/runs/{run_id}/cancel",
@@ -849,11 +930,17 @@ def get_base_router(
     async def continue_agent_run(
         agent_id: str,
         run_id: str,
+        request: Request,
         tools: str = Form(...),  # JSON string of tools
         session_id: Optional[str] = Form(None),
         user_id: Optional[str] = Form(None),
         stream: bool = Form(True),
     ):
+        if hasattr(request.state, "user_id"):
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            session_id = request.state.session_id
+
         # Parse the JSON string manually
         try:
             tools_data = json.loads(tools) if tools else None
@@ -891,17 +978,21 @@ def get_base_router(
                 media_type="text/event-stream",
             )
         else:
-            run_response_obj = cast(
-                RunOutput,
-                await agent.acontinue_run(
-                    run_id=run_id,  # run_id from path
-                    updated_tools=updated_tools,
-                    session_id=session_id,
-                    user_id=user_id,
-                    stream=False,
-                ),
-            )
-            return run_response_obj.to_dict()
+            try:
+                run_response_obj = cast(
+                    RunOutput,
+                    await agent.acontinue_run(
+                        run_id=run_id,  # run_id from path
+                        updated_tools=updated_tools,
+                        session_id=session_id,
+                        user_id=user_id,
+                        stream=False,
+                    ),
+                )
+                return run_response_obj.to_dict()
+
+            except InputCheckError as e:
+                raise HTTPException(status_code=400, detail=str(e))
 
     @router.get(
         "/agents",
@@ -1041,6 +1132,25 @@ def get_base_router(
     ):
         kwargs = await _get_request_kwargs(request, create_team_run)
 
+        if hasattr(request.state, "user_id"):
+            if user_id:
+                log_warning("User ID parameter passed in both request state and kwargs, using request state")
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            if session_id:
+                log_warning("Session ID parameter passed in both request state and kwargs, using request state")
+            session_id = request.state.session_id
+        if hasattr(request.state, "session_state"):
+            session_state = request.state.session_state
+            if "session_state" in kwargs:
+                log_warning("Session state parameter passed in both request state and kwargs, using request state")
+            kwargs["session_state"] = session_state
+        if hasattr(request.state, "dependencies"):
+            dependencies = request.state.dependencies
+            if "dependencies" in kwargs:
+                log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
+            kwargs["dependencies"] = dependencies
+
         logger.debug(f"Creating team run: {message=} {session_id=} {monitor=} {user_id=} {team_id=} {files=} {kwargs=}")
 
         team = get_team_by_id(team_id, os.teams)
@@ -1122,18 +1232,22 @@ def get_base_router(
                 media_type="text/event-stream",
             )
         else:
-            run_response = await team.arun(
-                input=message,
-                session_id=session_id,
-                user_id=user_id,
-                images=base64_images if base64_images else None,
-                audio=base64_audios if base64_audios else None,
-                videos=base64_videos if base64_videos else None,
-                files=document_files if document_files else None,
-                stream=False,
-                **kwargs,
-            )
-            return run_response.to_dict()
+            try:
+                run_response = await team.arun(
+                    input=message,
+                    session_id=session_id,
+                    user_id=user_id,
+                    images=base64_images if base64_images else None,
+                    audio=base64_audios if base64_audios else None,
+                    videos=base64_videos if base64_videos else None,
+                    files=document_files if document_files else None,
+                    stream=False,
+                    **kwargs,
+                )
+                return run_response.to_dict()
+
+            except InputCheckError as e:
+                raise HTTPException(status_code=400, detail=str(e))
 
     @router.post(
         "/teams/{team_id}/runs/{run_id}/cancel",
@@ -1464,6 +1578,25 @@ def get_base_router(
     ):
         kwargs = await _get_request_kwargs(request, create_workflow_run)
 
+        if hasattr(request.state, "user_id"):
+            if user_id:
+                log_warning("User ID parameter passed in both request state and kwargs, using request state")
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            if session_id:
+                log_warning("Session ID parameter passed in both request state and kwargs, using request state")
+            session_id = request.state.session_id
+        if hasattr(request.state, "session_state"):
+            session_state = request.state.session_state
+            if "session_state" in kwargs:
+                log_warning("Session state parameter passed in both request state and kwargs, using request state")
+            kwargs["session_state"] = session_state
+        if hasattr(request.state, "dependencies"):
+            dependencies = request.state.dependencies
+            if "dependencies" in kwargs:
+                log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
+            kwargs["dependencies"] = dependencies
+
         # Retrieve the workflow by ID
         workflow = get_workflow_by_id(workflow_id, os.workflows)
         if workflow is None:
@@ -1497,6 +1630,9 @@ def get_base_router(
                     **kwargs,
                 )
                 return run_response.to_dict()
+
+        except InputCheckError as e:
+            raise HTTPException(status_code=400, detail=str(e))
         except Exception as e:
             # Handle unexpected runtime errors
             raise HTTPException(status_code=500, detail=f"Error running workflow: {str(e)}")

agno/os/routers/home.py

@@ -30,7 +30,7 @@ def get_home_router(os: "AgentOS") -> APIRouter:
                                 "value": {
                                     "name": "AgentOS API",
                                     "description": "AI Agent Operating System API",
-                                    "os_id": "demo-os",
+                                    "id": "demo-os",
                                     "version": "1.0.0",
                                 },
                             }
@@ -45,7 +45,7 @@ def get_home_router(os: "AgentOS") -> APIRouter:
         return {
             "name": "AgentOS API",
             "description": os.description or "AI Agent Operating System API",
-            "os_id": os.os_id or "agno-agentos",
+            "id": os.id or "agno-agentos",
             "version": os.version or "1.0.0",
         }