agno 2.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/os/app.py

@@ -12,9 +12,11 @@ from rich.panel import Panel
 from starlette.requests import Request
 
 from agno.agent.agent import Agent
-from agno.db.base import BaseDb
+from agno.db.base import AsyncBaseDb, BaseDb
+from agno.knowledge.knowledge import Knowledge
 from agno.os.config import (
     AgentOSConfig,
+    AuthorizationConfig,
     DatabaseConfig,
     EvalsConfig,
     EvalsDomainConfig,
@@ -26,9 +28,12 @@ from agno.os.config import (
     MetricsDomainConfig,
     SessionConfig,
     SessionDomainConfig,
+    TracesConfig,
+    TracesDomainConfig,
 )
 from agno.os.interfaces.base import BaseInterface
-from agno.os.router import get_base_router, get_websocket_router
+from agno.os.router import get_base_router
+from agno.os.routers.agents import get_agent_router
 from agno.os.routers.evals import get_eval_router
 from agno.os.routers.health import get_health_router
 from agno.os.routers.home import get_home_router
@@ -36,16 +41,21 @@ from agno.os.routers.knowledge import get_knowledge_router
 from agno.os.routers.memory import get_memory_router
 from agno.os.routers.metrics import get_metrics_router
 from agno.os.routers.session import get_session_router
+from agno.os.routers.teams import get_team_router
+from agno.os.routers.traces import get_traces_router
+from agno.os.routers.workflows import get_websocket_router, get_workflow_router
 from agno.os.settings import AgnoAPISettings
 from agno.os.utils import (
     collect_mcp_tools_from_team,
     collect_mcp_tools_from_workflow,
     find_conflicting_routes,
     load_yaml_config,
+    resolve_origins,
+    setup_tracing_for_os,
     update_cors_middleware,
 )
 from agno.team.team import Team
-from agno.utils.log import logger
+from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.utils.string import generate_id, generate_id_from_name
 from agno.workflow.workflow import Workflow
 
@@ -64,6 +74,39 @@ async def mcp_lifespan(_, mcp_tools):
         await tool.close()
 
 
+@asynccontextmanager
+async def db_lifespan(app: FastAPI, agent_os: "AgentOS"):
+    """Initializes databases in the event loop"""
+    if agent_os.auto_provision_dbs:
+        agent_os._initialize_sync_databases()
+        await agent_os._initialize_async_databases()
+    yield
+
+
+def _combine_app_lifespans(lifespans: list) -> Any:
+    """Combine multiple FastAPI app lifespan context managers into one."""
+    if len(lifespans) == 1:
+        return lifespans[0]
+
+    from contextlib import asynccontextmanager
+
+    @asynccontextmanager
+    async def combined_lifespan(app):
+        async def _run_nested(index: int):
+            if index >= len(lifespans):
+                yield
+                return
+
+            async with lifespans[index](app):
+                async for _ in _run_nested(index + 1):
+                    yield
+
+        async for _ in _run_nested(0):
+            yield
+
+    return combined_lifespan
+
+
 class AgentOS:
     def __init__(
         self,
@@ -74,19 +117,23 @@ class AgentOS:
         agents: Optional[List[Agent]] = None,
         teams: Optional[List[Team]] = None,
         workflows: Optional[List[Workflow]] = None,
+        knowledge: Optional[List[Knowledge]] = None,
         interfaces: Optional[List[BaseInterface]] = None,
         a2a_interface: bool = False,
+        authorization: bool = False,
+        authorization_config: Optional[AuthorizationConfig] = None,
+        cors_allowed_origins: Optional[List[str]] = None,
         config: Optional[Union[str, AgentOSConfig]] = None,
         settings: Optional[AgnoAPISettings] = None,
         lifespan: Optional[Any] = None,
         enable_mcp_server: bool = False,
         base_app: Optional[FastAPI] = None,
         on_route_conflict: Literal["preserve_agentos", "preserve_base_app", "error"] = "preserve_agentos",
+        tracing: bool = False,
+        tracing_db: Optional[Union[BaseDb, AsyncBaseDb]] = None,
+        auto_provision_dbs: bool = True,
+        run_hooks_in_background: bool = False,
         telemetry: bool = True,
-        os_id: Optional[str] = None,  # Deprecated
-        enable_mcp: bool = False,  # Deprecated
-        fastapi_app: Optional[FastAPI] = None,  # Deprecated
-        replace_routes: Optional[bool] = None,  # Deprecated
     ):
         """Initialize AgentOS.
 
@@ -98,6 +145,7 @@ class AgentOS:
             agents: List of agents to include in the OS
             teams: List of teams to include in the OS
             workflows: List of workflows to include in the OS
+            knowledge: List of knowledge bases to include in the OS
             interfaces: List of interfaces to include in the OS
             a2a_interface: Whether to expose the OS agents and teams in an A2A server
             config: Configuration file path or AgentOSConfig instance
@@ -106,11 +154,19 @@ class AgentOS:
             enable_mcp_server: Whether to enable MCP (Model Context Protocol)
             base_app: Optional base FastAPI app to use for the AgentOS. All routes and middleware will be added to this app.
             on_route_conflict: What to do when a route conflict is detected in case a custom base_app is provided.
+            auto_provision_dbs: Whether to automatically provision databases
+            authorization: Whether to enable authorization
+            authorization_config: Configuration for the authorization middleware
+            cors_allowed_origins: List of allowed CORS origins (will be merged with default Agno domains)
+            tracing: If True, enables OpenTelemetry tracing for all agents and teams in the OS
+            tracing_db: Dedicated database for storing and reading traces. Recommended for multi-db setups.
+                       If not provided and tracing=True, the first available db from agents/teams/workflows is used.
+            run_hooks_in_background: If True, run agent/team pre/post hooks as FastAPI background tasks (non-blocking)
             telemetry: Whether to enable telemetry
 
         """
-        if not agents and not workflows and not teams:
-            raise ValueError("Either agents, teams or workflows must be provided.")
+        if not agents and not workflows and not teams and not knowledge:
+            raise ValueError("Either agents, teams, workflows or knowledge bases must be provided.")
 
         self.config = load_yaml_config(config) if isinstance(config, str) else config
 
@@ -119,22 +175,15 @@ class AgentOS:
         self.teams: Optional[List[Team]] = teams
         self.interfaces = interfaces or []
         self.a2a_interface = a2a_interface
-
+        self.knowledge = knowledge
         self.settings: AgnoAPISettings = settings or AgnoAPISettings()
-
+        self.auto_provision_dbs = auto_provision_dbs
         self._app_set = False
 
         if base_app:
             self.base_app: Optional[FastAPI] = base_app
             self._app_set = True
             self.on_route_conflict = on_route_conflict
-        elif fastapi_app:
-            self.base_app = fastapi_app
-            self._app_set = True
-            if replace_routes is not None:
-                self.on_route_conflict = "preserve_agentos" if replace_routes else "preserve_base_app"
-            else:
-                self.on_route_conflict = on_route_conflict
         else:
             self.base_app = None
             self._app_set = False
@@ -144,7 +193,7 @@ class AgentOS:
 
         self.name = name
 
-        self.id = id or os_id
+        self.id = id
         if not self.id:
             self.id = generate_id(self.name) if self.name else str(uuid4())
 
@@ -152,58 +201,138 @@ class AgentOS:
         self.description = description
 
         self.telemetry = telemetry
+        self.tracing = tracing
+        self.tracing_db = tracing_db
 
-        self.enable_mcp_server = enable_mcp or enable_mcp_server
+        self.enable_mcp_server = enable_mcp_server
         self.lifespan = lifespan
 
+        # RBAC
+        self.authorization = authorization
+        self.authorization_config = authorization_config
+
+        # CORS configuration - merge user-provided origins with defaults from settings
+        self.cors_allowed_origins = resolve_origins(cors_allowed_origins, self.settings.cors_origin_list)
+
+        # If True, run agent/team hooks as FastAPI background tasks
+        self.run_hooks_in_background = run_hooks_in_background
+
         # List of all MCP tools used inside the AgentOS
         self.mcp_tools: List[Any] = []
         self._mcp_app: Optional[Any] = None
 
-        if self.agents:
-            for agent in self.agents:
-                # Track all MCP tools to later handle their connection
-                if agent.tools:
-                    for tool in agent.tools:
-                        # Checking if the tool is a MCPTools or MultiMCPTools instance
-                        type_name = type(tool).__name__
-                        if type_name in ("MCPTools", "MultiMCPTools"):
-                            if tool not in self.mcp_tools:
-                                self.mcp_tools.append(tool)
+        self._initialize_agents()
+        self._initialize_teams()
+        self._initialize_workflows()
 
-                agent.initialize_agent()
+        if self.tracing:
+            self._setup_tracing()
 
-                # Required for the built-in routes to work
-                agent.store_events = True
+        if self.telemetry:
+            from agno.api.os import OSLaunch, log_os_telemetry
 
-        if self.teams:
-            for team in self.teams:
-                # Track all MCP tools recursively
-                collect_mcp_tools_from_team(team, self.mcp_tools)
+            log_os_telemetry(launch=OSLaunch(os_id=self.id, data=self._get_telemetry_data()))
 
-                team.initialize_team()
+    def _add_agent_os_to_lifespan_function(self, lifespan):
+        """
+        Inspect a lifespan function and wrap it to pass agent_os if it accepts it.
 
-                # Required for the built-in routes to work
-                team.store_events = True
+        Returns:
+            A wrapped lifespan that passes agent_os if the lifespan function expects it.
+        """
+        # Getting the actual function inside the lifespan
+        lifespan_function = lifespan
+        if hasattr(lifespan, "__wrapped__"):
+            lifespan_function = lifespan.__wrapped__
 
-                for member in team.members:
-                    if isinstance(member, Agent):
-                        member.team_id = None
-                        member.initialize_agent()
-                    elif isinstance(member, Team):
-                        member.initialize_team()
+        try:
+            from inspect import signature
 
-        if self.workflows:
-            for workflow in self.workflows:
-                # Track MCP tools recursively in workflow members
-                collect_mcp_tools_from_workflow(workflow, self.mcp_tools)
-                if not workflow.id:
-                    workflow.id = generate_id_from_name(workflow.name)
+            # Inspecting the lifespan function signature to find its parameters
+            sig = signature(lifespan_function)
+            params = list(sig.parameters.keys())
 
-        if self.telemetry:
-            from agno.api.os import OSLaunch, log_os_telemetry
+            # If the lifespan function expects the 'agent_os' parameter, add it
+            if "agent_os" in params:
+                return partial(lifespan, agent_os=self)
+            else:
+                return lifespan
 
-            log_os_telemetry(launch=OSLaunch(os_id=self.id, data=self._get_telemetry_data()))
+        except (ValueError, TypeError):
+            return lifespan
+
+    def resync(self, app: FastAPI) -> None:
+        """Resync the AgentOS to discover, initialize and configure: agents, teams, workflows, databases and knowledge bases."""
+        self._initialize_agents()
+        self._initialize_teams()
+        self._initialize_workflows()
+        self._auto_discover_databases()
+        self._auto_discover_knowledge_instances()
+
+        if self.enable_mcp_server:
+            from agno.os.mcp import get_mcp_server
+
+            self._mcp_app = get_mcp_server(self)
+
+        self._reprovision_routers(app=app)
+
+    def _reprovision_routers(self, app: FastAPI) -> None:
+        """Re-provision all routes for the AgentOS."""
+        updated_routers = [
+            get_session_router(dbs=self.dbs),
+            get_metrics_router(dbs=self.dbs),
+            get_knowledge_router(knowledge_instances=self.knowledge_instances),
+            get_traces_router(dbs=self.dbs),
+            get_memory_router(dbs=self.dbs),
+            get_eval_router(dbs=self.dbs, agents=self.agents, teams=self.teams),
+        ]
+
+        # Clear all previously existing routes
+        app.router.routes = [
+            route
+            for route in app.router.routes
+            if hasattr(route, "path")
+            and route.path in ["/docs", "/redoc", "/openapi.json", "/docs/oauth2-redirect"]
+            or route.path.startswith("/mcp")  # type: ignore
+        ]
+
+        # Add the built-in routes
+        self._add_built_in_routes(app=app)
+
+        # Add the updated routes
+        for router in updated_routers:
+            self._add_router(app, router)
+
+        # Mount MCP if needed
+        if self.enable_mcp_server and self._mcp_app:
+            app.mount("/", self._mcp_app)
+
+    def _add_built_in_routes(self, app: FastAPI) -> None:
+        """Add all AgentOSbuilt-in routes to the given app."""
+        # Add the home router if MCP server is not enabled
+        if not self.enable_mcp_server:
+            self._add_router(app, get_home_router(self))
+
+        self._add_router(app, get_health_router(health_endpoint="/health"))
+        self._add_router(app, get_base_router(self, settings=self.settings))
+        self._add_router(app, get_agent_router(self, settings=self.settings))
+        self._add_router(app, get_team_router(self, settings=self.settings))
+        self._add_router(app, get_workflow_router(self, settings=self.settings))
+        self._add_router(app, get_websocket_router(self, settings=self.settings))
+
+        # Add A2A interface if relevant
+        has_a2a_interface = False
+        for interface in self.interfaces:
+            if not has_a2a_interface and interface.__class__.__name__ == "A2A":
+                has_a2a_interface = True
+            interface_router = interface.get_router()
+            self._add_router(app, interface_router)
+        if self.a2a_interface and not has_a2a_interface:
+            from agno.os.interfaces.a2a import A2A
+
+            a2a_interface = A2A(agents=self.agents, teams=self.teams, workflows=self.workflows)
+            self.interfaces.append(a2a_interface)
+            self._add_router(app, a2a_interface.get_router())
 
     def _make_app(self, lifespan: Optional[Any] = None) -> FastAPI:
         # Adjust the FastAPI app lifespan to handle MCP connections if relevant
@@ -220,7 +349,7 @@ class AgentOS:
                         async with mcp_tools_lifespan(app):  # type: ignore
                             yield
 
-                app_lifespan = combined_lifespan  # type: ignore
+                app_lifespan = combined_lifespan
             else:
                 app_lifespan = mcp_tools_lifespan
 
@@ -234,53 +363,175 @@ class AgentOS:
             lifespan=app_lifespan,
         )
 
+    def _initialize_agents(self) -> None:
+        """Initialize and configure all agents for AgentOS usage."""
+        if not self.agents:
+            return
+        for agent in self.agents:
+            # Track all MCP tools to later handle their connection
+            if agent.tools:
+                for tool in agent.tools:
+                    # Checking if the tool is an instance of MCPTools, MultiMCPTools, or a subclass of those
+                    if hasattr(type(tool), "__mro__"):
+                        mro_names = {cls.__name__ for cls in type(tool).__mro__}
+                        if mro_names & {"MCPTools", "MultiMCPTools"}:
+                            if tool not in self.mcp_tools:
+                                self.mcp_tools.append(tool)
+
+            agent.initialize_agent()
+
+            # Required for the built-in routes to work
+            agent.store_events = True
+
+            # Propagate run_hooks_in_background setting from AgentOS to agents
+            agent._run_hooks_in_background = self.run_hooks_in_background
+
+    def _initialize_teams(self) -> None:
+        """Initialize and configure all teams for AgentOS usage."""
+        if not self.teams:
+            return
+
+        for team in self.teams:
+            # Track all MCP tools recursively
+            collect_mcp_tools_from_team(team, self.mcp_tools)
+
+            team.initialize_team()
+
+            for member in team.members:
+                if isinstance(member, Agent):
+                    member.team_id = None
+                    member.initialize_agent()
+                elif isinstance(member, Team):
+                    member.initialize_team()
+
+            # Required for the built-in routes to work
+            team.store_events = True
+
+            # Propagate run_hooks_in_background setting to team and all nested members
+            team.propagate_run_hooks_in_background(self.run_hooks_in_background)
+
+    def _initialize_workflows(self) -> None:
+        """Initialize and configure all workflows for AgentOS usage."""
+        if not self.workflows:
+            return
+
+        if self.workflows:
+            for workflow in self.workflows:
+                # Track MCP tools recursively in workflow members
+                collect_mcp_tools_from_workflow(workflow, self.mcp_tools)
+
+                if not workflow.id:
+                    workflow.id = generate_id_from_name(workflow.name)
+
+                # Required for the built-in routes to work
+                workflow.store_events = True
+
+                # Propagate run_hooks_in_background setting to workflow and all its step agents/teams
+                workflow.propagate_run_hooks_in_background(self.run_hooks_in_background)
+
+    def _setup_tracing(self) -> None:
+        """Set up OpenTelemetry tracing for this AgentOS.
+
+        Uses tracing_db if provided, otherwise falls back to the first available
+        database from agents/teams/workflows.
+        """
+        # Use tracing_db if explicitly provided
+        if self.tracing_db is not None:
+            setup_tracing_for_os(db=self.tracing_db)
+            return
+
+        # Fall back to finding the first available database
+        db: Optional[Union[BaseDb, AsyncBaseDb]] = None
+
+        for agent in self.agents or []:
+            if agent.db:
+                db = agent.db
+                break
+
+        if db is None:
+            for team in self.teams or []:
+                if team.db:
+                    db = team.db
+                    break
+
+        if db is None:
+            for workflow in self.workflows or []:
+                if workflow.db:
+                    db = workflow.db
+                    break
+
+        if db is None:
+            log_warning(
+                "tracing=True but no database found. "
+                "Provide 'tracing_db' parameter or 'db' parameter to at least one agent/team/workflow."
+            )
+            return
+
+        setup_tracing_for_os(db=db)
+
     def get_app(self) -> FastAPI:
         if self.base_app:
             fastapi_app = self.base_app
-        else:
-            if self.enable_mcp_server:
-                from contextlib import asynccontextmanager
 
+            # Initialize MCP server if enabled
+            if self.enable_mcp_server:
                 from agno.os.mcp import get_mcp_server
 
                 self._mcp_app = get_mcp_server(self)
 
-                final_lifespan = self._mcp_app.lifespan  # type: ignore
-                if self.lifespan is not None:
-                    # Combine both lifespans
-                    @asynccontextmanager
-                    async def combined_lifespan(app: FastAPI):
-                        # Run both lifespans
-                        async with self.lifespan(app):  # type: ignore
-                            async with self._mcp_app.lifespan(app):  # type: ignore
-                                yield
+            # Collect all lifespans that need to be combined
+            lifespans = []
 
-                    final_lifespan = combined_lifespan  # type: ignore
+            # The user provided lifespan
+            if self.lifespan:
+                # Wrap the user lifespan with agent_os parameter
+                wrapped_lifespan = self._add_agent_os_to_lifespan_function(self.lifespan)
+                lifespans.append(wrapped_lifespan)
 
-                fastapi_app = self._make_app(lifespan=final_lifespan)
-            else:
-                fastapi_app = self._make_app(lifespan=self.lifespan)
+            # The provided app's existing lifespan
+            if fastapi_app.router.lifespan_context:
+                lifespans.append(fastapi_app.router.lifespan_context)
 
-        # Add routes
-        self._add_router(fastapi_app, get_base_router(self, settings=self.settings))
-        self._add_router(fastapi_app, get_websocket_router(self, settings=self.settings))
-        self._add_router(fastapi_app, get_health_router())
-        self._add_router(fastapi_app, get_home_router(self))
+            # The MCP tools lifespan
+            if self.mcp_tools:
+                lifespans.append(partial(mcp_lifespan, mcp_tools=self.mcp_tools))
 
-        has_a2a_interface = False
-        for interface in self.interfaces:
-            if not has_a2a_interface and interface.__class__.__name__ == "A2A":
-                has_a2a_interface = True
-            interface_router = interface.get_router()
-            self._add_router(fastapi_app, interface_router)
+            # The /mcp server lifespan
+            if self.enable_mcp_server and self._mcp_app:
+                lifespans.append(self._mcp_app.lifespan)
 
-        # Add A2A interface if requested and not provided in self.interfaces
-        if self.a2a_interface and not has_a2a_interface:
-            from agno.os.interfaces.a2a import A2A
+            # The async database lifespan
+            lifespans.append(partial(db_lifespan, agent_os=self))
 
-            a2a_interface = A2A(agents=self.agents, teams=self.teams, workflows=self.workflows)
-            self.interfaces.append(a2a_interface)
-            self._add_router(fastapi_app, a2a_interface.get_router())
+            # Combine lifespans and set them in the app
+            if lifespans:
+                fastapi_app.router.lifespan_context = _combine_app_lifespans(lifespans)
+
+        else:
+            lifespans = []
+
+            # User provided lifespan
+            if self.lifespan:
+                lifespans.append(self._add_agent_os_to_lifespan_function(self.lifespan))
+
+            # MCP tools lifespan
+            if self.mcp_tools:
+                lifespans.append(partial(mcp_lifespan, mcp_tools=self.mcp_tools))
+
+            # MCP server lifespan
+            if self.enable_mcp_server:
+                from agno.os.mcp import get_mcp_server
+
+                self._mcp_app = get_mcp_server(self)
+                lifespans.append(self._mcp_app.lifespan)
+
+            # Async database initialization lifespan
+            lifespans.append(partial(db_lifespan, agent_os=self))  # type: ignore
+
+            final_lifespan = _combine_app_lifespans(lifespans) if lifespans else None
+            fastapi_app = self._make_app(lifespan=final_lifespan)
+
+        self._add_built_in_routes(app=fastapi_app)
 
         self._auto_discover_databases()
         self._auto_discover_knowledge_instances()
@@ -291,6 +542,7 @@ class AgentOS:
             get_eval_router(dbs=self.dbs, agents=self.agents, teams=self.teams),
             get_metrics_router(dbs=self.dbs),
             get_knowledge_router(knowledge_instances=self.knowledge_instances),
+            get_traces_router(dbs=self.dbs),
         ]
 
         for router in routers:
@@ -299,35 +551,76 @@ class AgentOS:
         # Mount MCP if needed
         if self.enable_mcp_server and self._mcp_app:
             fastapi_app.mount("/", self._mcp_app)
-        else:
-            # Add the home router
-            self._add_router(fastapi_app, get_home_router(self))
 
         if not self._app_set:
 
             @fastapi_app.exception_handler(HTTPException)
             async def http_exception_handler(_, exc: HTTPException) -> JSONResponse:
+                log_error(f"HTTP exception: {exc.status_code} {exc.detail}")
                 return JSONResponse(
                     status_code=exc.status_code,
                     content={"detail": str(exc.detail)},
                 )
 
-            async def general_exception_handler(request: Request, call_next):
-                try:
-                    return await call_next(request)
-                except Exception as e:
-                    return JSONResponse(
-                        status_code=e.status_code if hasattr(e, "status_code") else 500,  # type: ignore
-                        content={"detail": str(e)},
-                    )
+            @fastapi_app.exception_handler(Exception)
+            async def general_exception_handler(_: Request, exc: Exception) -> JSONResponse:
+                import traceback
 
-            fastapi_app.middleware("http")(general_exception_handler)
+                log_error(f"Unhandled exception:\n{traceback.format_exc(limit=5)}")
+
+                return JSONResponse(
+                    status_code=getattr(exc, "status_code", 500),
+                    content={"detail": str(exc)},
+                )
 
         # Update CORS middleware
-        update_cors_middleware(fastapi_app, self.settings.cors_origin_list)  # type: ignore
+        update_cors_middleware(fastapi_app, self.cors_allowed_origins)  # type: ignore
+
+        # Set agent_os_id and cors_allowed_origins on app state
+        # This allows middleware (like JWT) to access these values
+        fastapi_app.state.agent_os_id = self.id
+        fastapi_app.state.cors_allowed_origins = self.cors_allowed_origins
+
+        # Add JWT middleware if authorization is enabled
+        if self.authorization:
+            self._add_jwt_middleware(fastapi_app)
 
         return fastapi_app
 
+    def _add_jwt_middleware(self, fastapi_app: FastAPI) -> None:
+        from agno.os.middleware.jwt import JWTMiddleware, JWTValidator
+
+        verify_audience = False
+        jwks_file = None
+        verification_keys = None
+        algorithm = "RS256"
+
+        if self.authorization_config:
+            algorithm = self.authorization_config.algorithm or "RS256"
+            verification_keys = self.authorization_config.verification_keys
+            jwks_file = self.authorization_config.jwks_file
+            verify_audience = self.authorization_config.verify_audience or False
+
+        log_info(f"Adding JWT middleware for authorization (algorithm: {algorithm})")
+
+        # Create validator and store on app.state for WebSocket access
+        jwt_validator = JWTValidator(
+            verification_keys=verification_keys,
+            jwks_file=jwks_file,
+            algorithm=algorithm,
+        )
+        fastapi_app.state.jwt_validator = jwt_validator
+
+        # Add middleware to stack
+        fastapi_app.add_middleware(
+            JWTMiddleware,
+            verification_keys=verification_keys,
+            jwks_file=jwks_file,
+            algorithm=algorithm,
+            authorization=self.authorization,
+            verify_audience=verify_audience,
+        )
+
     def get_routes(self) -> List[Any]:
         """Retrieve all routes from the FastAPI app.
 
@@ -353,7 +646,7 @@ class AgentOS:
                 # Skip conflicting AgentOS routes, prefer user's existing routes
                 for conflict in conflicts:
                     methods_str = ", ".join(conflict["methods"])  # type: ignore
-                    logger.debug(
+                    log_debug(
                         f"Skipping conflicting AgentOS route: {methods_str} {conflict['path']} - "
                         f"Using existing custom route instead"
                     )
@@ -372,7 +665,7 @@ class AgentOS:
                 # Log warnings but still add all routes (AgentOS routes will override)
                 for conflict in conflicts:
                     methods_str = ", ".join(conflict["methods"])  # type: ignore
-                    logger.warning(
+                    log_warning(
                         f"Route conflict detected: {methods_str} {conflict['path']} - "
                         f"AgentOS route will override existing custom route"
                     )
@@ -404,11 +697,12 @@ class AgentOS:
         }
 
     def _auto_discover_databases(self) -> None:
-        """Auto-discover the databases used by all contextual agents, teams and workflows."""
-        from agno.db.base import BaseDb
+        """Auto-discover and initialize the databases used by all contextual agents, teams and workflows."""
 
-        dbs: Dict[str, BaseDb] = {}
-        knowledge_dbs: Dict[str, BaseDb] = {}  # Track databases specifically used for knowledge
+        dbs: Dict[str, List[Union[BaseDb, AsyncBaseDb]]] = {}
+        knowledge_dbs: Dict[
+            str, List[Union[BaseDb, AsyncBaseDb]]
+        ] = {}  # Track databases specifically used for knowledge
 
         for agent in self.agents or []:
             if agent.db:
@@ -426,68 +720,120 @@ class AgentOS:
             if workflow.db:
                 self._register_db_with_validation(dbs, workflow.db)
 
+        for knowledge_base in self.knowledge or []:
+            if knowledge_base.contents_db:
+                self._register_db_with_validation(knowledge_dbs, knowledge_base.contents_db)
+
         for interface in self.interfaces or []:
             if interface.agent and interface.agent.db:
                 self._register_db_with_validation(dbs, interface.agent.db)
             elif interface.team and interface.team.db:
                 self._register_db_with_validation(dbs, interface.team.db)
 
+        # Register tracing_db if provided (for traces reading)
+        if self.tracing_db is not None:
+            self._register_db_with_validation(dbs, self.tracing_db)
+
         self.dbs = dbs
         self.knowledge_dbs = knowledge_dbs
 
-    def _register_db_with_validation(self, registered_dbs: Dict[str, Any], db: BaseDb) -> None:
-        """Register a database in the contextual OS after validating it is not conflicting with registered databases"""
-        if db.id in registered_dbs:
-            existing_db = registered_dbs[db.id]
-            if not self._are_db_instances_compatible(existing_db, db):
-                raise ValueError(
-                    f"Database ID conflict detected: Two different database instances have the same ID '{db.id}'. "
-                    f"Database instances with the same ID must point to the same database with identical configuration."
+        # Initialize all discovered databases
+        if self.auto_provision_dbs:
+            self._pending_async_db_init = True
+
+    def _initialize_sync_databases(self) -> None:
+        """Initialize sync databases."""
+        from itertools import chain
+
+        unique_dbs = list(
+            {
+                id(db): db
+                for db in chain(
+                    chain.from_iterable(self.dbs.values()), chain.from_iterable(self.knowledge_dbs.values())
                 )
-        registered_dbs[db.id] = db
+            }.values()
+        )
 
-    def _are_db_instances_compatible(self, db1: BaseDb, db2: BaseDb) -> bool:
-        """
-        Return True if the two given database objects are compatible
-        Two database objects are compatible if they point to the same database with identical configuration.
-        """
-        # If they're the same object reference, they're compatible
-        if db1 is db2:
-            return True
+        for db in unique_dbs:
+            if isinstance(db, AsyncBaseDb):
+                continue  # Skip async dbs
 
-        if type(db1) is not type(db2):
-            return False
+            try:
+                if hasattr(db, "_create_all_tables") and callable(db._create_all_tables):
+                    db._create_all_tables()
+            except Exception as e:
+                log_warning(f"Failed to initialize {db.__class__.__name__} (id: {db.id}): {e}")
 
-        if hasattr(db1, "db_url") and hasattr(db2, "db_url"):
-            if db1.db_url != db2.db_url:  # type: ignore
-                return False
+    async def _initialize_async_databases(self) -> None:
+        """Initialize async databases."""
 
-        if hasattr(db1, "db_file") and hasattr(db2, "db_file"):
-            if db1.db_file != db2.db_file:  # type: ignore
-                return False
+        from itertools import chain
 
-        # If table names are different, they're not compatible
-        if (
-            db1.session_table_name != db2.session_table_name
-            or db1.memory_table_name != db2.memory_table_name
-            or db1.metrics_table_name != db2.metrics_table_name
-            or db1.eval_table_name != db2.eval_table_name
-            or db1.knowledge_table_name != db2.knowledge_table_name
-        ):
-            return False
+        unique_dbs = list(
+            {
+                id(db): db
+                for db in chain(
+                    chain.from_iterable(self.dbs.values()), chain.from_iterable(self.knowledge_dbs.values())
+                )
+            }.values()
+        )
 
-        return True
+        for db in unique_dbs:
+            if not isinstance(db, AsyncBaseDb):
+                continue  # Skip sync dbs
+
+            try:
+                if hasattr(db, "_create_all_tables") and callable(db._create_all_tables):
+                    await db._create_all_tables()
+            except Exception as e:
+                log_warning(f"Failed to initialize async {db.__class__.__name__} (id: {db.id}): {e}")
+
+    def _get_db_table_names(self, db: BaseDb) -> Dict[str, str]:
+        """Get the table names for a database"""
+        table_names = {
+            "session_table_name": db.session_table_name,
+            "culture_table_name": db.culture_table_name,
+            "memory_table_name": db.memory_table_name,
+            "metrics_table_name": db.metrics_table_name,
+            "evals_table_name": db.eval_table_name,
+            "knowledge_table_name": db.knowledge_table_name,
+        }
+        return {k: v for k, v in table_names.items() if v is not None}
+
+    def _register_db_with_validation(
+        self, registered_dbs: Dict[str, List[Union[BaseDb, AsyncBaseDb]]], db: Union[BaseDb, AsyncBaseDb]
+    ) -> None:
+        """Register a database in the contextual OS after validating it is not conflicting with registered databases"""
+        if db.id in registered_dbs:
+            registered_dbs[db.id].append(db)
+        else:
+            registered_dbs[db.id] = [db]
 
     def _auto_discover_knowledge_instances(self) -> None:
         """Auto-discover the knowledge instances used by all contextual agents, teams and workflows."""
-        knowledge_instances = []
+        seen_ids = set()
+        knowledge_instances: List[Knowledge] = []
+
+        def _add_knowledge_if_not_duplicate(knowledge: "Knowledge") -> None:
+            """Add knowledge instance if it's not already in the list (by object identity or db_id)."""
+            # Use database ID if available, otherwise use object ID as fallback
+            if not knowledge.contents_db:
+                return
+            if knowledge.contents_db.id in seen_ids:
+                return
+            seen_ids.add(knowledge.contents_db.id)
+            knowledge_instances.append(knowledge)
+
         for agent in self.agents or []:
             if agent.knowledge:
-                knowledge_instances.append(agent.knowledge)
+                _add_knowledge_if_not_duplicate(agent.knowledge)
 
         for team in self.teams or []:
             if team.knowledge:
-                knowledge_instances.append(team.knowledge)
+                _add_knowledge_if_not_duplicate(team.knowledge)
+
+        for knowledge_base in self.knowledge or []:
+            _add_knowledge_if_not_duplicate(knowledge_base)
 
         self.knowledge_instances = knowledge_instances
 
@@ -498,13 +844,15 @@ class AgentOS:
             session_config.dbs = []
 
         dbs_with_specific_config = [db.db_id for db in session_config.dbs]
-
-        for db_id in self.dbs.keys():
+        for db_id, dbs in self.dbs.items():
             if db_id not in dbs_with_specific_config:
+                # Collect unique table names from all databases with the same id
+                unique_tables = list(set(db.session_table_name for db in dbs))
                 session_config.dbs.append(
                     DatabaseConfig(
                         db_id=db_id,
                         domain_config=SessionDomainConfig(display_name=db_id),
+                        tables=unique_tables,
                     )
                 )
 
@@ -518,12 +866,15 @@ class AgentOS:
 
         dbs_with_specific_config = [db.db_id for db in memory_config.dbs]
 
-        for db_id in self.dbs.keys():
+        for db_id, dbs in self.dbs.items():
             if db_id not in dbs_with_specific_config:
+                # Collect unique table names from all databases with the same id
+                unique_tables = list(set(db.memory_table_name for db in dbs))
                 memory_config.dbs.append(
                     DatabaseConfig(
                         db_id=db_id,
                         domain_config=MemoryDomainConfig(display_name=db_id),
+                        tables=unique_tables,
                     )
                 )
 
@@ -557,12 +908,15 @@ class AgentOS:
 
         dbs_with_specific_config = [db.db_id for db in metrics_config.dbs]
 
-        for db_id in self.dbs.keys():
+        for db_id, dbs in self.dbs.items():
             if db_id not in dbs_with_specific_config:
+                # Collect unique table names from all databases with the same id
+                unique_tables = list(set(db.metrics_table_name for db in dbs))
                 metrics_config.dbs.append(
                     DatabaseConfig(
                         db_id=db_id,
                         domain_config=MetricsDomainConfig(display_name=db_id),
+                        tables=unique_tables,
                     )
                 )
 
@@ -576,17 +930,50 @@ class AgentOS:
 
         dbs_with_specific_config = [db.db_id for db in evals_config.dbs]
 
-        for db_id in self.dbs.keys():
+        for db_id, dbs in self.dbs.items():
             if db_id not in dbs_with_specific_config:
+                # Collect unique table names from all databases with the same id
+                unique_tables = list(set(db.eval_table_name for db in dbs))
                 evals_config.dbs.append(
                     DatabaseConfig(
                         db_id=db_id,
                         domain_config=EvalsDomainConfig(display_name=db_id),
+                        tables=unique_tables,
                     )
                 )
 
         return evals_config
 
+    def _get_traces_config(self) -> TracesConfig:
+        traces_config = self.config.traces if self.config and self.config.traces else TracesConfig()
+
+        if traces_config.dbs is None:
+            traces_config.dbs = []
+
+        dbs_with_specific_config = [db.db_id for db in traces_config.dbs]
+
+        # If tracing_db is explicitly set, only use that database for traces
+        if self.tracing_db is not None:
+            if self.tracing_db.id not in dbs_with_specific_config:
+                traces_config.dbs.append(
+                    DatabaseConfig(
+                        db_id=self.tracing_db.id,
+                        domain_config=TracesDomainConfig(display_name=self.tracing_db.id),
+                    )
+                )
+        else:
+            # Fall back to all discovered databases
+            for db_id in self.dbs.keys():
+                if db_id not in dbs_with_specific_config:
+                    traces_config.dbs.append(
+                        DatabaseConfig(
+                            db_id=db_id,
+                            domain_config=TracesDomainConfig(display_name=db_id),
+                        )
+                    )
+
+        return traces_config
+
     def serve(
         self,
         app: Union[str, FastAPI],
@@ -595,6 +982,7 @@ class AgentOS:
         port: int = 7777,
         reload: bool = False,
         workers: Optional[int] = None,
+        access_log: bool = False,
         **kwargs,
     ):
         import uvicorn
@@ -627,4 +1015,13 @@ class AgentOS:
             )
         )
 
-        uvicorn.run(app=app, host=host, port=port, reload=reload, workers=workers, **kwargs)
+        uvicorn.run(
+            app=app,
+            host=host,
+            port=port,
+            reload=reload,
+            workers=workers,
+            access_log=access_log,
+            lifespan="on",
+            **kwargs,
+        )