agno 2.0.4__py3-none-any.whl → 2.0.6__py3-none-any.whl

agno/os/app.py

@@ -1,11 +1,12 @@
 from contextlib import asynccontextmanager
 from functools import partial
 from os import getenv
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Set, Union
 from uuid import uuid4
 
-from fastapi import FastAPI, HTTPException
+from fastapi import APIRouter, FastAPI, HTTPException
 from fastapi.responses import JSONResponse
+from fastapi.routing import APIRoute
 from rich import box
 from rich.panel import Panel
 from starlette.middleware.cors import CORSMiddleware
@@ -30,13 +31,15 @@ from agno.os.interfaces.base import BaseInterface
 from agno.os.router import get_base_router, get_websocket_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
 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.settings import AgnoAPISettings
-from agno.os.utils import generate_id
 from agno.team.team import Team
+from agno.utils.log import logger
+from agno.utils.string import generate_id, generate_id_from_name
 from agno.workflow.workflow import Workflow
 
 
@@ -70,8 +73,30 @@ class AgentOS:
         fastapi_app: Optional[FastAPI] = None,
         lifespan: Optional[Any] = None,
         enable_mcp: bool = False,
+        replace_routes: bool = True,
         telemetry: bool = True,
     ):
+        """Initialize AgentOS.
+
+        Args:
+            os_id: Unique identifier for this AgentOS instance
+            name: Name of the AgentOS instance
+            description: Description of the AgentOS instance
+            version: Version of the AgentOS instance
+            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
+            interfaces: List of interfaces to include in the OS
+            config: Configuration file path or AgentOSConfig instance
+            settings: API settings for the OS
+            fastapi_app: Optional custom FastAPI app to use instead of creating a new one
+            lifespan: Optional lifespan context manager for the FastAPI app
+            enable_mcp: Whether to enable MCP (Model Context Protocol)
+            replace_routes: If False and using a custom fastapi_app, skip AgentOS routes that
+                          conflict with existing routes, preferring the user's custom routes.
+                          If True (default), AgentOS routes will override conflicting custom routes.
+            telemetry: Whether to enable telemetry
+        """
         if not agents and not workflows and not teams:
             raise ValueError("Either agents, teams or workflows must be provided.")
 
@@ -92,11 +117,13 @@ class AgentOS:
 
         self.interfaces = interfaces or []
 
-        self.os_id: Optional[str] = os_id
+        self.os_id = os_id
         self.name = name
         self.version = version
         self.description = description
 
+        self.replace_routes = replace_routes
+
         self.telemetry = telemetry
 
         self.enable_mcp = enable_mcp
@@ -146,7 +173,10 @@ class AgentOS:
             for workflow in self.workflows:
                 # TODO: track MCP tools in workflow members
                 if not workflow.id:
-                    workflow.id = generate_id(workflow.name)
+                    workflow.id = generate_id_from_name(workflow.name)
+
+        if not self.os_id:
+            self.os_id = generate_id(self.name) if self.name else str(uuid4())
 
         if self.telemetry:
             from agno.api.os import OSLaunch, log_os_telemetry
@@ -207,18 +237,29 @@ class AgentOS:
             else:
                 self.fastapi_app = self._make_app(lifespan=self.lifespan)
 
-        # Add routes
-        self.fastapi_app.include_router(get_base_router(self, settings=self.settings))
-        self.fastapi_app.include_router(get_websocket_router(self, settings=self.settings))
-        self.fastapi_app.include_router(get_health_router())
+        # Add routes with conflict detection
+        self._add_router(get_base_router(self, settings=self.settings))
+        self._add_router(get_websocket_router(self, settings=self.settings))
+        self._add_router(get_health_router())
+        self._add_router(get_home_router(self))
 
         for interface in self.interfaces:
             interface_router = interface.get_router()
-            self.fastapi_app.include_router(interface_router)
+            self._add_router(interface_router)
 
         self._auto_discover_databases()
         self._auto_discover_knowledge_instances()
-        self._setup_routers()
+
+        routers = [
+            get_session_router(dbs=self.dbs),
+            get_memory_router(dbs=self.dbs),
+            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),
+        ]
+
+        for router in routers:
+            self._add_router(router)
 
         # Mount MCP if needed
         if self.enable_mcp and self.mcp_app:
@@ -266,6 +307,94 @@ class AgentOS:
 
         return app.routes
 
+    def _get_existing_route_paths(self) -> Dict[str, List[str]]:
+        """Get all existing route paths and methods from the FastAPI app.
+
+        Returns:
+            Dict[str, List[str]]: Dictionary mapping paths to list of HTTP methods
+        """
+        if not self.fastapi_app:
+            return {}
+
+        existing_paths: Dict[str, Any] = {}
+        for route in self.fastapi_app.routes:
+            if isinstance(route, APIRoute):
+                path = route.path
+                methods = list(route.methods) if route.methods else []
+                if path in existing_paths:
+                    existing_paths[path].extend(methods)
+                else:
+                    existing_paths[path] = methods
+        return existing_paths
+
+    def _add_router(self, router: APIRouter) -> None:
+        """Add a router to the FastAPI app, avoiding route conflicts.
+
+        Args:
+            router: The APIRouter to add
+        """
+        if not self.fastapi_app:
+            return
+
+        # Get existing routes
+        existing_paths = self._get_existing_route_paths()
+
+        # Check for conflicts
+        conflicts = []
+        conflicting_routes = []
+
+        for route in router.routes:
+            if isinstance(route, APIRoute):
+                full_path = route.path
+                route_methods = list(route.methods) if route.methods else []
+
+                if full_path in existing_paths:
+                    conflicting_methods: Set[str] = set(route_methods) & set(existing_paths[full_path])
+                    if conflicting_methods:
+                        conflicts.append({"path": full_path, "methods": list(conflicting_methods), "route": route})
+                        conflicting_routes.append(route)
+
+        if conflicts and self._app_set:
+            if self.replace_routes:
+                # Log warnings but still add all routes (AgentOS routes will override)
+                for conflict in conflicts:
+                    methods_str = ", ".join(conflict["methods"])  # type: ignore
+                    logger.warning(
+                        f"Route conflict detected: {methods_str} {conflict['path']} - "
+                        f"AgentOS route will override existing custom route"
+                    )
+
+                # Remove conflicting routes
+                for route in self.fastapi_app.routes:
+                    for conflict in conflicts:
+                        if isinstance(route, APIRoute):
+                            if route.path == conflict["path"] and list(route.methods) == list(conflict["methods"]):
+                                self.fastapi_app.routes.pop(self.fastapi_app.routes.index(route))
+
+                self.fastapi_app.include_router(router)
+
+            else:
+                # Skip conflicting AgentOS routes, prefer user's existing routes
+                for conflict in conflicts:
+                    methods_str = ", ".join(conflict["methods"])  # type: ignore
+                    logger.debug(
+                        f"Skipping conflicting AgentOS route: {methods_str} {conflict['path']} - "
+                        f"Using existing custom route instead"
+                    )
+
+                # Create a new router without the conflicting routes
+                filtered_router = APIRouter()
+                for route in router.routes:
+                    if route not in conflicting_routes:
+                        filtered_router.routes.append(route)
+
+                # Use the filtered router if it has any routes left
+                if filtered_router.routes:
+                    self.fastapi_app.include_router(filtered_router)
+        else:
+            # No conflicts, add router normally
+            self.fastapi_app.include_router(router)
+
     def _get_telemetry_data(self) -> Dict[str, Any]:
         """Get the telemetry data for the OS"""
         return {
@@ -293,18 +422,25 @@ class AgentOS:
     def _auto_discover_databases(self) -> None:
         """Auto-discover the databases used by all contextual agents, teams and workflows."""
         dbs = {}
+        knowledge_dbs = {}  # Track databases specifically used for knowledge
 
         for agent in self.agents or []:
             if agent.db:
                 dbs[agent.db.id] = agent.db
             if agent.knowledge and agent.knowledge.contents_db:
-                dbs[agent.knowledge.contents_db.id] = agent.knowledge.contents_db
+                knowledge_dbs[agent.knowledge.contents_db.id] = agent.knowledge.contents_db
+                # Also add to general dbs if it's used for both purposes
+                if agent.knowledge.contents_db.id not in dbs:
+                    dbs[agent.knowledge.contents_db.id] = agent.knowledge.contents_db
 
         for team in self.teams or []:
             if team.db:
                 dbs[team.db.id] = team.db
             if team.knowledge and team.knowledge.contents_db:
-                dbs[team.knowledge.contents_db.id] = team.knowledge.contents_db
+                knowledge_dbs[team.knowledge.contents_db.id] = team.knowledge.contents_db
+                # Also add to general dbs if it's used for both purposes
+                if team.knowledge.contents_db.id not in dbs:
+                    dbs[team.knowledge.contents_db.id] = team.knowledge.contents_db
 
         for workflow in self.workflows or []:
             if workflow.db:
@@ -317,6 +453,7 @@ class AgentOS:
                 dbs[interface.team.db.id] = interface.team.db
 
         self.dbs = dbs
+        self.knowledge_dbs = knowledge_dbs
 
     def _auto_discover_knowledge_instances(self) -> None:
         """Auto-discover the knowledge instances used by all contextual agents, teams and workflows."""
@@ -381,16 +518,17 @@ class AgentOS:
         if knowledge_config.dbs is None:
             knowledge_config.dbs = []
 
-        multiple_dbs: bool = len(self.dbs.keys()) > 1
+        multiple_knowledge_dbs: bool = len(self.knowledge_dbs.keys()) > 1
         dbs_with_specific_config = [db.db_id for db in knowledge_config.dbs]
 
-        for db_id in self.dbs.keys():
+        # Only add databases that are actually used for knowledge contents
+        for db_id in self.knowledge_dbs.keys():
             if db_id not in dbs_with_specific_config:
                 knowledge_config.dbs.append(
                     DatabaseConfig(
                         db_id=db_id,
                         domain_config=KnowledgeDomainConfig(
-                            display_name="Knowledge" if not multiple_dbs else "Knowledge in database " + db_id
+                            display_name="Knowledge" if not multiple_knowledge_dbs else "Knowledge in database " + db_id
                         ),
                     )
                 )
@@ -441,29 +579,6 @@ class AgentOS:
 
         return evals_config
 
-    def _setup_routers(self) -> None:
-        """Add all routers to the FastAPI app."""
-        if not self.dbs or not self.fastapi_app:
-            return
-
-        routers = [
-            get_session_router(dbs=self.dbs),
-            get_memory_router(dbs=self.dbs),
-            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),
-        ]
-
-        for router in routers:
-            self.fastapi_app.include_router(router)
-
-    def set_os_id(self) -> str:
-        # If os_id is already set, keep it instead of overriding with UUID
-        if self.os_id is None:
-            self.os_id = str(uuid4())
-
-        return self.os_id
-
     def serve(
         self,
         app: Union[str, FastAPI],
@@ -485,13 +600,18 @@ class AgentOS:
         from rich.align import Align
         from rich.console import Console, Group
 
-        aligned_endpoint = Align.center(f"[bold cyan]{public_endpoint}[/bold cyan]")
-        connection_endpoint = f"\n\n[bold dark_orange]Running on:[/bold dark_orange] http://{host}:{port}"
+        panel_group = []
+        panel_group.append(Align.center(f"[bold cyan]{public_endpoint}[/bold cyan]"))
+        panel_group.append(
+            Align.center(f"\n\n[bold dark_orange]OS running on:[/bold dark_orange] http://{host}:{port}")
+        )
+        if bool(self.settings.os_security_key):
+            panel_group.append(Align.center("\n\n[bold chartreuse3]:lock: Security Enabled[/bold chartreuse3]"))
 
         console = Console()
         console.print(
             Panel(
-                Group(aligned_endpoint, connection_endpoint),
+                Group(*panel_group),
                 title="AgentOS",
                 expand=False,
                 border_style="dark_orange",

agno/os/interfaces/agui/utils.py

@@ -2,10 +2,9 @@
 
 import json
 import uuid
-from collections import deque
 from collections.abc import Iterator
 from dataclasses import dataclass
-from typing import AsyncIterator, Deque, List, Optional, Set, Tuple, Union
+from typing import AsyncIterator, List, Set, Tuple, Union
 
 from ag_ui.core import (
     BaseEvent,
@@ -34,39 +33,22 @@ from agno.utils.message import get_text_from_message
 class EventBuffer:
     """Buffer to manage event ordering constraints, relevant when mapping Agno responses to AG-UI events."""
 
-    buffer: Deque[BaseEvent]
-    blocking_tool_call_id: Optional[str]  # The tool call that's currently blocking the buffer
     active_tool_call_ids: Set[str]  # All currently active tool calls
     ended_tool_call_ids: Set[str]  # All tool calls that have ended
 
     def __init__(self):
-        self.buffer = deque()
-        self.blocking_tool_call_id = None
         self.active_tool_call_ids = set()
         self.ended_tool_call_ids = set()
 
-    def is_blocked(self) -> bool:
-        """Check if the buffer is currently blocked by an active tool call."""
-        return self.blocking_tool_call_id is not None
-
     def start_tool_call(self, tool_call_id: str) -> None:
-        """Start a new tool call, marking it the current blocking tool call if needed."""
+        """Start a new tool call."""
         self.active_tool_call_ids.add(tool_call_id)
-        if self.blocking_tool_call_id is None:
-            self.blocking_tool_call_id = tool_call_id
 
-    def end_tool_call(self, tool_call_id: str) -> bool:
-        """End a tool call, marking it as ended and unblocking the buffer if needed."""
+    def end_tool_call(self, tool_call_id: str) -> None:
+        """End a tool call."""
         self.active_tool_call_ids.discard(tool_call_id)
         self.ended_tool_call_ids.add(tool_call_id)
 
-        # Unblock the buffer if the current blocking tool call is the one ending
-        if tool_call_id == self.blocking_tool_call_id:
-            self.blocking_tool_call_id = None
-            return True
-
-        return False
-
 
 def convert_agui_messages_to_agno_messages(messages: List[AGUIMessage]) -> List[Message]:
     """Convert AG-UI messages to Agno messages."""
@@ -169,6 +151,12 @@ def _create_events_from_chunk(
 
     # Handle starting a new tool call
     elif chunk.event == RunEvent.tool_call_started:
+        # End the current text message if one is active before starting tool calls
+        if message_started:
+            end_message_event = TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=message_id)
+            events_to_emit.append(end_message_event)
+            message_started = False  # Reset message_started state
+
         if chunk.tool is not None:  # type: ignore
             tool_call = chunk.tool  # type: ignore
             start_event = ToolCallStartEvent(
@@ -195,7 +183,7 @@ def _create_events_from_chunk(
                     type=EventType.TOOL_CALL_END,
                     tool_call_id=tool_call.tool_call_id,  # type: ignore
                 )
-                events_to_emit.append(end_event)  # type: ignore
+                events_to_emit.append(end_event)
 
                 if tool_call.result is not None:
                     result_event = ToolCallResultEvent(
@@ -205,27 +193,17 @@ def _create_events_from_chunk(
                         role="tool",
                         message_id=str(uuid.uuid4()),
                     )
-                    events_to_emit.append(result_event)  # type: ignore
-
-                if tool_call.result is not None:
-                    result_event = ToolCallResultEvent(
-                        type=EventType.TOOL_CALL_RESULT,
-                        tool_call_id=tool_call.tool_call_id,  # type: ignore
-                        content=str(tool_call.result),
-                        role="tool",
-                        message_id=str(uuid.uuid4()),
-                    )
-                    events_to_emit.append(result_event)  # type: ignore
+                    events_to_emit.append(result_event)
 
     # Handle reasoning
     elif chunk.event == RunEvent.reasoning_started:
-        step_started_event = StepStartedEvent(type=EventType.STEP_STARTED, step_name="reasoning")  # type: ignore
-        events_to_emit.append(step_started_event)  # type: ignore
+        step_started_event = StepStartedEvent(type=EventType.STEP_STARTED, step_name="reasoning")
+        events_to_emit.append(step_started_event)
     elif chunk.event == RunEvent.reasoning_completed:
-        step_started_event = StepFinishedEvent(type=EventType.STEP_FINISHED, step_name="reasoning")  # type: ignore
-        events_to_emit.append(step_started_event)  # type: ignore
+        step_finished_event = StepFinishedEvent(type=EventType.STEP_FINISHED, step_name="reasoning")
+        events_to_emit.append(step_finished_event)
 
-    return events_to_emit, message_started  # type: ignore
+    return events_to_emit, message_started
 
 
 def _create_completion_events(
@@ -251,7 +229,7 @@ def _create_completion_events(
     # End the message and run, denoting the end of the session
     if message_started:
         end_message_event = TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=message_id)
-        events_to_emit.append(end_message_event)  # type: ignore
+        events_to_emit.append(end_message_event)
 
     # emit frontend tool calls, i.e. external_execution=True
     if isinstance(chunk, RunPausedEvent) and chunk.tools is not None:
@@ -265,14 +243,14 @@ def _create_completion_events(
                 tool_call_name=tool.tool_name,
                 parent_message_id=message_id,
             )
-            events_to_emit.append(start_event)  # type: ignore
+            events_to_emit.append(start_event)
 
             args_event = ToolCallArgsEvent(
                 type=EventType.TOOL_CALL_ARGS,
                 tool_call_id=tool.tool_call_id,
                 delta=json.dumps(tool.tool_args),
             )
-            events_to_emit.append(args_event)  # type: ignore
+            events_to_emit.append(args_event)
 
             end_event = ToolCallEndEvent(
                 type=EventType.TOOL_CALL_END,
@@ -280,85 +258,25 @@ def _create_completion_events(
             )
             events_to_emit.append(end_event)
 
-    # emit frontend tool calls, i.e. external_execution=True
-    if isinstance(chunk, RunPausedEvent) and chunk.tools is not None:
-        for tool in chunk.tools:
-            if tool.tool_call_id is None or tool.tool_name is None:
-                continue
-
-            start_event = ToolCallStartEvent(
-                type=EventType.TOOL_CALL_START,
-                tool_call_id=tool.tool_call_id,
-                tool_call_name=tool.tool_name,
-                parent_message_id=message_id,
-            )
-            events_to_emit.append(start_event)  # type: ignore
-
-            args_event = ToolCallArgsEvent(
-                type=EventType.TOOL_CALL_ARGS,
-                tool_call_id=tool.tool_call_id,
-                delta=json.dumps(tool.tool_args),
-            )
-            events_to_emit.append(args_event)  # type: ignore
-
-            end_event = ToolCallEndEvent(
-                type=EventType.TOOL_CALL_END,
-                tool_call_id=tool.tool_call_id,
-            )
-            events_to_emit.append(end_event)  # type: ignore
-
     run_finished_event = RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=run_id)
-    events_to_emit.append(run_finished_event)  # type: ignore
+    events_to_emit.append(run_finished_event)
 
-    return events_to_emit  # type: ignore
+    return events_to_emit
 
 
 def _emit_event_logic(event: BaseEvent, event_buffer: EventBuffer) -> List[BaseEvent]:
-    """Process an event through the buffer and return events to actually emit."""
-    events_to_emit: List[BaseEvent] = []
-
-    if event_buffer.is_blocked():
-        # Handle events related to the current blocking tool call
-        if event.type == EventType.TOOL_CALL_ARGS:
-            if hasattr(event, "tool_call_id") and event.tool_call_id in event_buffer.active_tool_call_ids:  # type: ignore
-                events_to_emit.append(event)
-            else:
-                event_buffer.buffer.append(event)
-        elif event.type == EventType.TOOL_CALL_END:
-            tool_call_id = getattr(event, "tool_call_id", None)
-            if tool_call_id and tool_call_id == event_buffer.blocking_tool_call_id:
-                events_to_emit.append(event)
-                event_buffer.end_tool_call(tool_call_id)
-                # Flush buffered events after ending the blocking tool call
-                while event_buffer.buffer:
-                    buffered_event = event_buffer.buffer.popleft()
-                    # Recursively process buffered events
-                    nested_events = _emit_event_logic(buffered_event, event_buffer)
-                    events_to_emit.extend(nested_events)
-            elif tool_call_id and tool_call_id in event_buffer.active_tool_call_ids:
-                event_buffer.buffer.append(event)
-                event_buffer.end_tool_call(tool_call_id)
-            else:
-                event_buffer.buffer.append(event)
-        # Handle all other events
-        elif event.type == EventType.TOOL_CALL_START:
-            event_buffer.buffer.append(event)
-        else:
-            event_buffer.buffer.append(event)
-    # If the buffer is not blocked, emit the events normally
-    else:
-        if event.type == EventType.TOOL_CALL_START:
-            tool_call_id = getattr(event, "tool_call_id", None)
-            if tool_call_id:
-                event_buffer.start_tool_call(tool_call_id)
-            events_to_emit.append(event)
-        elif event.type == EventType.TOOL_CALL_END:
-            tool_call_id = getattr(event, "tool_call_id", None)
-            if tool_call_id:
-                event_buffer.end_tool_call(tool_call_id)
-            events_to_emit.append(event)
-        else:
-            events_to_emit.append(event)
+    """Process an event and return events to actually emit."""
+    events_to_emit: List[BaseEvent] = [event]
+
+    # Update the event buffer state for tracking purposes
+    if event.type == EventType.TOOL_CALL_START:
+        tool_call_id = getattr(event, "tool_call_id", None)
+        if tool_call_id:
+            event_buffer.start_tool_call(tool_call_id)
+    elif event.type == EventType.TOOL_CALL_END:
+        tool_call_id = getattr(event, "tool_call_id", None)
+        if tool_call_id:
+            event_buffer.end_tool_call(tool_call_id)
 
     return events_to_emit
 
@@ -370,23 +288,22 @@ def stream_agno_response_as_agui_events(
     message_id = str(uuid.uuid4())
     message_started = False
     event_buffer = EventBuffer()
+    stream_completed = False
+
+    completion_chunk = None
 
     for chunk in response_stream:
-        # Handle the lifecycle end event
+        # Check if this is a completion event
         if (
             chunk.event == RunEvent.run_completed
             or chunk.event == TeamRunEvent.run_completed
             or chunk.event == RunEvent.run_paused
         ):
-            completion_events = _create_completion_events(
-                chunk, event_buffer, message_started, message_id, thread_id, run_id
-            )
-            for event in completion_events:
-                events_to_emit = _emit_event_logic(event_buffer=event_buffer, event=event)
-                for emit_event in events_to_emit:
-                    yield emit_event
+            # Store completion chunk but don't process it yet
+            completion_chunk = chunk
+            stream_completed = True
         else:
-            # Process regular chunk
+            # Process regular chunk immediately
             events_from_chunk, message_started = _create_events_from_chunk(
                 chunk, message_id, message_started, event_buffer
             )
@@ -396,6 +313,30 @@ def stream_agno_response_as_agui_events(
                 for emit_event in events_to_emit:
                     yield emit_event
 
+    # Process ONLY completion cleanup events, not content from completion chunk
+    if completion_chunk:
+        completion_events = _create_completion_events(
+            completion_chunk, event_buffer, message_started, message_id, thread_id, run_id
+        )
+        for event in completion_events:
+            events_to_emit = _emit_event_logic(event_buffer=event_buffer, event=event)
+            for emit_event in events_to_emit:
+                yield emit_event
+
+    # Ensure completion events are always emitted even when stream ends naturally
+    if not stream_completed:
+        # Create a synthetic completion event to ensure proper cleanup
+        from agno.run.agent import RunCompletedEvent
+
+        synthetic_completion = RunCompletedEvent()
+        completion_events = _create_completion_events(
+            synthetic_completion, event_buffer, message_started, message_id, thread_id, run_id
+        )
+        for event in completion_events:
+            events_to_emit = _emit_event_logic(event_buffer=event_buffer, event=event)
+            for emit_event in events_to_emit:
+                yield emit_event
+
 
 # Async version - thin wrapper
 async def async_stream_agno_response_as_agui_events(
@@ -407,23 +348,22 @@ async def async_stream_agno_response_as_agui_events(
     message_id = str(uuid.uuid4())
     message_started = False
     event_buffer = EventBuffer()
+    stream_completed = False
+
+    completion_chunk = None
 
     async for chunk in response_stream:
-        # Handle the lifecycle end event
+        # Check if this is a completion event
         if (
             chunk.event == RunEvent.run_completed
             or chunk.event == TeamRunEvent.run_completed
             or chunk.event == RunEvent.run_paused
         ):
-            completion_events = _create_completion_events(
-                chunk, event_buffer, message_started, message_id, thread_id, run_id
-            )
-            for event in completion_events:
-                events_to_emit = _emit_event_logic(event_buffer=event_buffer, event=event)
-                for emit_event in events_to_emit:
-                    yield emit_event
+            # Store completion chunk but don't process it yet
+            completion_chunk = chunk
+            stream_completed = True
         else:
-            # Process regular chunk
+            # Process regular chunk immediately
             events_from_chunk, message_started = _create_events_from_chunk(
                 chunk, message_id, message_started, event_buffer
             )
@@ -432,3 +372,27 @@ async def async_stream_agno_response_as_agui_events(
                 events_to_emit = _emit_event_logic(event_buffer=event_buffer, event=event)
                 for emit_event in events_to_emit:
                     yield emit_event
+
+    # Process ONLY completion cleanup events, not content from completion chunk
+    if completion_chunk:
+        completion_events = _create_completion_events(
+            completion_chunk, event_buffer, message_started, message_id, thread_id, run_id
+        )
+        for event in completion_events:
+            events_to_emit = _emit_event_logic(event_buffer=event_buffer, event=event)
+            for emit_event in events_to_emit:
+                yield emit_event
+
+    # Ensure completion events are always emitted even when stream ends naturally
+    if not stream_completed:
+        # Create a synthetic completion event to ensure proper cleanup
+        from agno.run.agent import RunCompletedEvent
+
+        synthetic_completion = RunCompletedEvent()
+        completion_events = _create_completion_events(
+            synthetic_completion, event_buffer, message_started, message_id, thread_id, run_id
+        )
+        for event in completion_events:
+            events_to_emit = _emit_event_logic(event_buffer=event_buffer, event=event)
+            for emit_event in events_to_emit:
+                yield emit_event