agno 2.0.0rc2__py3-none-any.whl → 2.3.0__py3-none-any.whl

agno/os/interfaces/agui/agui.py

@@ -1,6 +1,6 @@
 """Main class for the AG-UI app, used to expose an Agno Agent or Team in an AG-UI compatible format."""
 
-from typing import Optional
+from typing import List, Optional
 
 from fastapi.routing import APIRouter
 
@@ -15,16 +15,32 @@ class AGUI(BaseInterface):
 
     router: APIRouter
 
-    def __init__(self, agent: Optional[Agent] = None, team: Optional[Team] = None):
+    def __init__(
+        self,
+        agent: Optional[Agent] = None,
+        team: Optional[Team] = None,
+        prefix: str = "",
+        tags: Optional[List[str]] = None,
+    ):
+        """
+        Initialize the AGUI interface.
+
+        Args:
+            agent: The agent to expose via AG-UI
+            team: The team to expose via AG-UI
+            prefix: Custom prefix for the router (e.g., "/agui/v1", "/chat/public")
+            tags: Custom tags for the router (e.g., ["AGUI", "Chat"], defaults to ["AGUI"])
+        """
         self.agent = agent
         self.team = team
+        self.prefix = prefix
+        self.tags = tags or ["AGUI"]
 
-        if not self.agent and not self.team:
-            raise ValueError("AGUI requires an agent and a team")
+        if not (self.agent or self.team):
+            raise ValueError("AGUI requires an agent or a team")
 
-    def get_router(self, **kwargs) -> APIRouter:
-        # Cannot be overridden
-        self.router = APIRouter(tags=["AGUI"])
+    def get_router(self) -> APIRouter:
+        self.router = APIRouter(prefix=self.prefix, tags=self.tags)  # type: ignore
 
         self.router = attach_routes(router=self.router, agent=self.agent, team=self.team)
 

agno/os/interfaces/agui/router.py

@@ -19,6 +19,7 @@ from agno.agent.agent import Agent
 from agno.os.interfaces.agui.utils import (
     async_stream_agno_response_as_agui_events,
     convert_agui_messages_to_agno_messages,
+    validate_agui_state,
 )
 from agno.team.team import Team
 
@@ -34,12 +35,22 @@ async def run_agent(agent: Agent, run_input: RunAgentInput) -> AsyncIterator[Bas
         messages = convert_agui_messages_to_agno_messages(run_input.messages or [])
         yield RunStartedEvent(type=EventType.RUN_STARTED, thread_id=run_input.thread_id, run_id=run_id)
 
+        # Look for user_id in run_input.forwarded_props
+        user_id = None
+        if run_input.forwarded_props and isinstance(run_input.forwarded_props, dict):
+            user_id = run_input.forwarded_props.get("user_id")
+
+        # Validating the session state is of the expected type (dict)
+        session_state = validate_agui_state(run_input.state, run_input.thread_id)
+
         # Request streaming response from agent
         response_stream = agent.arun(
             input=messages,
             session_id=run_input.thread_id,
             stream=True,
-            stream_intermediate_steps=True,
+            stream_events=True,
+            user_id=user_id,
+            session_state=session_state,
         )
 
         # Stream the response content in AG-UI format
@@ -64,12 +75,22 @@ async def run_team(team: Team, input: RunAgentInput) -> AsyncIterator[BaseEvent]
         messages = convert_agui_messages_to_agno_messages(input.messages or [])
         yield RunStartedEvent(type=EventType.RUN_STARTED, thread_id=input.thread_id, run_id=run_id)
 
+        # Look for user_id in input.forwarded_props
+        user_id = None
+        if input.forwarded_props and isinstance(input.forwarded_props, dict):
+            user_id = input.forwarded_props.get("user_id")
+
+        # Validating the session state is of the expected type (dict)
+        session_state = validate_agui_state(input.state, input.thread_id)
+
         # Request streaming response from team
         response_stream = team.arun(
             input=messages,
             session_id=input.thread_id,
             stream=True,
-            stream_intermediate_steps=True,
+            stream_steps=True,
+            user_id=user_id,
+            session_state=session_state,
         )
 
         # Stream the response content in AG-UI format
@@ -89,7 +110,10 @@ def attach_routes(router: APIRouter, agent: Optional[Agent] = None, team: Option
 
     encoder = EventEncoder()
 
-    @router.post("/agui")
+    @router.post(
+        "/agui",
+        name="run_agent",
+    )
     async def run_agent_agui(run_input: RunAgentInput):
         async def event_generator():
             if agent:

agno/os/interfaces/agui/utils.py

@@ -2,13 +2,13 @@
 
 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 dataclasses import asdict, dataclass, is_dataclass
+from typing import Any, AsyncIterator, Dict, List, Optional, Set, Tuple, Union
 
 from ag_ui.core import (
     BaseEvent,
+    CustomEvent,
     EventType,
     RunFinishedEvent,
     StepFinishedEvent,
@@ -22,50 +22,96 @@ from ag_ui.core import (
     ToolCallStartEvent,
 )
 from ag_ui.core.types import Message as AGUIMessage
+from pydantic import BaseModel
 
 from agno.models.message import Message
 from agno.run.agent import RunContentEvent, RunEvent, RunOutputEvent, RunPausedEvent
 from agno.run.team import RunContentEvent as TeamRunContentEvent
 from agno.run.team import TeamRunEvent, TeamRunOutputEvent
+from agno.utils.log import log_warning
 from agno.utils.message import get_text_from_message
 
 
+def validate_agui_state(state: Any, thread_id: str) -> Optional[Dict[str, Any]]:
+    """Validate the given AGUI state is of the expected type (dict)."""
+    if state is None:
+        return None
+
+    if isinstance(state, dict):
+        return state
+
+    if isinstance(state, BaseModel):
+        try:
+            return state.model_dump()
+        except Exception:
+            pass
+
+    if is_dataclass(state):
+        try:
+            return asdict(state)  # type: ignore
+        except Exception:
+            pass
+
+    if hasattr(state, "to_dict") and callable(getattr(state, "to_dict")):
+        try:
+            result = state.to_dict()  # type: ignore
+            if isinstance(result, dict):
+                return result
+        except Exception:
+            pass
+
+    log_warning(f"AGUI state must be a dict, got {type(state).__name__}. State will be ignored. Thread: {thread_id}")
+    return None
+
+
 @dataclass
 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
+    current_text_message_id: str = ""  # ID of the current text message context (for tool call parenting)
+    next_text_message_id: str = ""  # Pre-generated ID for the next text message
+    pending_tool_calls_parent_id: str = ""  # Parent message ID for pending tool calls
 
     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
+        self.current_text_message_id = ""
+        self.next_text_message_id = str(uuid.uuid4())
+        self.pending_tool_calls_parent_id = ""
 
     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
+    def start_text_message(self) -> str:
+        """Start a new text message and return its ID."""
+        # Use the pre-generated next ID as current, and generate a new next ID
+        self.current_text_message_id = self.next_text_message_id
+        self.next_text_message_id = str(uuid.uuid4())
+        return self.current_text_message_id
+
+    def get_parent_message_id_for_tool_call(self) -> str:
+        """Get the message ID to use as parent for tool calls."""
+        # If we have a pending parent ID set (from text message end), use that
+        if self.pending_tool_calls_parent_id:
+            return self.pending_tool_calls_parent_id
+        # Otherwise use current text message ID
+        return self.current_text_message_id
+
+    def set_pending_tool_calls_parent_id(self, parent_id: str) -> None:
+        """Set the parent message ID for upcoming tool calls."""
+        self.pending_tool_calls_parent_id = parent_id
 
-        return False
+    def clear_pending_tool_calls_parent_id(self) -> None:
+        """Clear the pending parent ID when a new text message starts."""
+        self.pending_tool_calls_parent_id = ""
 
 
 def convert_agui_messages_to_agno_messages(messages: List[AGUIMessage]) -> List[Message]:
@@ -131,10 +177,18 @@ def _create_events_from_chunk(
     message_id: str,
     message_started: bool,
     event_buffer: EventBuffer,
-) -> Tuple[List[BaseEvent], bool]:
+) -> Tuple[List[BaseEvent], bool, str]:
     """
     Process a single chunk and return events to emit + updated message_started state.
-    Returns: (events_to_emit, new_message_started_state)
+
+    Args:
+        chunk: The event chunk to process
+        message_id: Current message identifier
+        message_started: Whether a message is currently active
+        event_buffer: Event buffer for tracking tool call state
+
+    Returns:
+        Tuple of (events_to_emit, new_message_started_state, message_id)
     """
     events_to_emit: List[BaseEvent] = []
 
@@ -151,6 +205,11 @@ def _create_events_from_chunk(
         # Handle the message start event, emitted once per message
         if not message_started:
             message_started = True
+            message_id = event_buffer.start_text_message()
+
+            # Clear pending tool calls parent ID when starting new text message
+            event_buffer.clear_pending_tool_calls_parent_id()
+
             start_event = TextMessageStartEvent(
                 type=EventType.TEXT_MESSAGE_START,
                 message_id=message_id,
@@ -167,15 +226,37 @@ def _create_events_from_chunk(
             )
             events_to_emit.append(content_event)  # type: ignore
 
-    # Handle starting a new tool call
-    elif chunk.event == RunEvent.tool_call_started:
+    # Handle starting a new tool
+    elif chunk.event == RunEvent.tool_call_started or chunk.event == TeamRunEvent.tool_call_started:
         if chunk.tool is not None:  # type: ignore
             tool_call = chunk.tool  # type: ignore
+
+            # End current text message and handle for tool calls
+            current_message_id = message_id
+            if message_started:
+                # End the current text message
+                end_message_event = TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=current_message_id)
+                events_to_emit.append(end_message_event)
+
+                # Set this message as the parent for any upcoming tool calls
+                # This ensures multiple sequential tool calls all use the same parent
+                event_buffer.set_pending_tool_calls_parent_id(current_message_id)
+
+                # Reset message started state and generate new message_id for future messages
+                message_started = False
+                message_id = str(uuid.uuid4())
+
+            # Get the parent message ID - this will use pending parent if set, ensuring multiple tool calls in sequence have the same parent
+            parent_message_id = event_buffer.get_parent_message_id_for_tool_call()
+
+            if not parent_message_id:
+                parent_message_id = current_message_id
+
             start_event = ToolCallStartEvent(
                 type=EventType.TOOL_CALL_START,
                 tool_call_id=tool_call.tool_call_id,  # type: ignore
                 tool_call_name=tool_call.tool_name,  # type: ignore
-                parent_message_id=message_id,
+                parent_message_id=parent_message_id,
             )
             events_to_emit.append(start_event)
 
@@ -187,7 +268,7 @@ def _create_events_from_chunk(
             events_to_emit.append(args_event)  # type: ignore
 
     # Handle tool call completion
-    elif chunk.event == RunEvent.tool_call_completed:
+    elif chunk.event == RunEvent.tool_call_completed or chunk.event == TeamRunEvent.tool_call_completed:
         if chunk.tool is not None:  # type: ignore
             tool_call = chunk.tool  # type: ignore
             if tool_call.tool_call_id not in event_buffer.ended_tool_call_ids:
@@ -195,7 +276,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 +286,34 @@ 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)
+
+    # Handle custom events
+    elif chunk.event == RunEvent.custom_event:
+        # Use the name of the event class if available, otherwise default to the CustomEvent
+        try:
+            custom_event_name = chunk.__class__.__name__
+        except Exception:
+            custom_event_name = chunk.event
 
-    return events_to_emit, message_started  # type: ignore
+        # Use the complete Agno event as value if parsing it works, else the event content field
+        try:
+            custom_event_value = chunk.to_dict()
+        except Exception:
+            custom_event_value = chunk.content  # type: ignore
+
+        custom_event = CustomEvent(name=custom_event_name, value=custom_event_value)
+        events_to_emit.append(custom_event)
+
+    return events_to_emit, message_started, message_id
 
 
 def _create_completion_events(
@@ -251,37 +339,36 @@ 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:
-        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,
+        # First, emit an assistant message for external tool calls
+        assistant_message_id = str(uuid.uuid4())
+        assistant_start_event = TextMessageStartEvent(
+            type=EventType.TEXT_MESSAGE_START,
+            message_id=assistant_message_id,
+            role="assistant",
+        )
+        events_to_emit.append(assistant_start_event)
+
+        # Add any text content if present for the assistant message
+        if chunk.content:
+            content_event = TextMessageContentEvent(
+                type=EventType.TEXT_MESSAGE_CONTENT,
+                message_id=assistant_message_id,
+                delta=str(chunk.content),
             )
-            events_to_emit.append(start_event)  # type: ignore
+            events_to_emit.append(content_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
+        # End the assistant message
+        assistant_end_event = TextMessageEndEvent(
+            type=EventType.TEXT_MESSAGE_END,
+            message_id=assistant_message_id,
+        )
+        events_to_emit.append(assistant_end_event)
 
-            end_event = ToolCallEndEvent(
-                type=EventType.TOOL_CALL_END,
-                tool_call_id=tool.tool_call_id,
-            )
-            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:
+        # Now emit the tool call events with the assistant message as parent
         for tool in chunk.tools:
             if tool.tool_call_id is None or tool.tool_name is None:
                 continue
@@ -290,75 +377,42 @@ def _create_completion_events(
                 type=EventType.TOOL_CALL_START,
                 tool_call_id=tool.tool_call_id,
                 tool_call_name=tool.tool_name,
-                parent_message_id=message_id,
+                parent_message_id=assistant_message_id,  # Use the assistant message as parent
             )
-            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,
                 tool_call_id=tool.tool_call_id,
             )
-            events_to_emit.append(end_event)  # type: ignore
+            events_to_emit.append(end_event)
 
     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
 
@@ -367,27 +421,26 @@ def stream_agno_response_as_agui_events(
     response_stream: Iterator[Union[RunOutputEvent, TeamRunOutputEvent]], thread_id: str, run_id: str
 ) -> Iterator[BaseEvent]:
     """Map the Agno response stream to AG-UI format, handling event ordering constraints."""
-    message_id = str(uuid.uuid4())
+    message_id = ""  # Will be set by EventBuffer when text message starts
     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
-            events_from_chunk, message_started = _create_events_from_chunk(
+            # Process regular chunk immediately
+            events_from_chunk, message_started, message_id = _create_events_from_chunk(
                 chunk, message_id, message_started, event_buffer
             )
 
@@ -396,6 +449,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(
@@ -404,27 +481,26 @@ async def async_stream_agno_response_as_agui_events(
     run_id: str,
 ) -> AsyncIterator[BaseEvent]:
     """Map the Agno response stream to AG-UI format, handling event ordering constraints."""
-    message_id = str(uuid.uuid4())
+    message_id = ""  # Will be set by EventBuffer when text message starts
     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
-            events_from_chunk, message_started = _create_events_from_chunk(
+            # Process regular chunk immediately
+            events_from_chunk, message_started, message_id = _create_events_from_chunk(
                 chunk, message_id, message_started, event_buffer
             )
 
@@ -432,3 +508,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