agno 2.0.3__py3-none-any.whl → 2.0.5__py3-none-any.whl

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

agno/os/router.py

@@ -18,13 +18,12 @@ from pydantic import BaseModel
 from agno.agent.agent import Agent
 from agno.media import Audio, Image, Video
 from agno.media import File as FileMedia
-from agno.os.auth import get_authentication_dependency
+from agno.os.auth import get_authentication_dependency, validate_websocket_token
 from agno.os.schema import (
     AgentResponse,
     AgentSummaryResponse,
     BadRequestResponse,
     ConfigResponse,
-    HealthResponse,
     InterfaceResponse,
     InternalServerErrorResponse,
     Model,
@@ -49,7 +48,7 @@ from agno.os.utils import (
 from agno.run.agent import RunErrorEvent, RunOutput, RunOutputEvent
 from agno.run.team import RunErrorEvent as TeamRunErrorEvent
 from agno.run.team import TeamRunOutputEvent
-from agno.run.workflow import WorkflowErrorEvent, WorkflowRunOutputEvent
+from agno.run.workflow import WorkflowErrorEvent, WorkflowRunOutput, WorkflowRunOutputEvent
 from agno.team.team import Team
 from agno.utils.log import log_debug, log_error, log_warning, logger
 from agno.workflow.workflow import Workflow
@@ -110,6 +109,7 @@ class WebSocketManager:
     """Manages WebSocket connections for workflow runs"""
 
     active_connections: Dict[str, WebSocket]  # {run_id: websocket}
+    authenticated_connections: Dict[WebSocket, bool]  # {websocket: is_authenticated}
 
     def __init__(
         self,
@@ -117,22 +117,51 @@ class WebSocketManager:
     ):
         # Store active connections: {run_id: websocket}
         self.active_connections = active_connections or {}
+        # Track authentication state for each websocket
+        self.authenticated_connections = {}
 
-    async def connect(self, websocket: WebSocket):
+    async def connect(self, websocket: WebSocket, requires_auth: bool = True):
         """Accept WebSocket connection"""
         await websocket.accept()
         logger.debug("WebSocket connected")
 
-        # Send connection confirmation
+        # If auth is not required, mark as authenticated immediately
+        self.authenticated_connections[websocket] = not requires_auth
+
+        # Send connection confirmation with auth requirement info
         await websocket.send_text(
             json.dumps(
                 {
                     "event": "connected",
-                    "message": "Connected to workflow events",
+                    "message": (
+                        "Connected to workflow events. Please authenticate to continue."
+                        if requires_auth
+                        else "Connected to workflow events. Authentication not required."
+                    ),
+                    "requires_auth": requires_auth,
                 }
             )
         )
 
+    async def authenticate_websocket(self, websocket: WebSocket):
+        """Mark a WebSocket connection as authenticated"""
+        self.authenticated_connections[websocket] = True
+        logger.debug("WebSocket authenticated")
+
+        # Send authentication confirmation
+        await websocket.send_text(
+            json.dumps(
+                {
+                    "event": "authenticated",
+                    "message": "Authentication successful. You can now send commands.",
+                }
+            )
+        )
+
+    def is_authenticated(self, websocket: WebSocket) -> bool:
+        """Check if a WebSocket connection is authenticated"""
+        return self.authenticated_connections.get(websocket, False)
+
     async def register_workflow_websocket(self, run_id: str, websocket: WebSocket):
         """Register a workflow run with its WebSocket connection"""
         self.active_connections[run_id] = websocket
@@ -141,9 +170,26 @@ class WebSocketManager:
     async def disconnect_by_run_id(self, run_id: str):
         """Remove WebSocket connection by run_id"""
         if run_id in self.active_connections:
+            websocket = self.active_connections[run_id]
             del self.active_connections[run_id]
+            # Clean up authentication state
+            if websocket in self.authenticated_connections:
+                del self.authenticated_connections[websocket]
             logger.debug(f"WebSocket disconnected for run_id: {run_id}")
 
+    async def disconnect_websocket(self, websocket: WebSocket):
+        """Remove WebSocket connection and clean up all associated state"""
+        # Remove from authenticated connections
+        if websocket in self.authenticated_connections:
+            del self.authenticated_connections[websocket]
+
+        # Remove from active connections
+        runs_to_remove = [run_id for run_id, ws in self.active_connections.items() if ws == websocket]
+        for run_id in runs_to_remove:
+            del self.active_connections[run_id]
+
+        logger.debug("WebSocket disconnected and cleaned up")
+
     async def get_websocket_for_run(self, run_id: str) -> Optional[WebSocket]:
         """Get WebSocket connection for a workflow run"""
         return self.active_connections.get(run_id)
@@ -288,7 +334,7 @@ async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os:
                 session_id = str(uuid4())
 
         # Execute workflow in background with streaming
-        await workflow.arun(
+        workflow_result = await workflow.arun(
             input=user_message,
             session_id=session_id,
             user_id=user_id,
@@ -298,6 +344,10 @@ async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os:
             websocket=websocket,
         )
 
+        workflow_run_output = cast(WorkflowRunOutput, workflow_result)
+
+        await websocket_manager.register_workflow_websocket(workflow_run_output.run_id, websocket)  # type: ignore
+
     except Exception as e:
         logger.error(f"Error executing workflow via WebSocket: {e}")
         await websocket.send_text(json.dumps({"event": "error", "error": str(e)}))
@@ -334,6 +384,77 @@ async def workflow_response_streamer(
         return
 
 
+def get_websocket_router(
+    os: "AgentOS",
+    settings: AgnoAPISettings = AgnoAPISettings(),
+) -> APIRouter:
+    """
+    Create WebSocket router without HTTP authentication dependencies.
+    WebSocket endpoints handle authentication internally via message-based auth.
+    """
+    ws_router = APIRouter()
+
+    @ws_router.websocket(
+        "/workflows/ws",
+        name="workflow_websocket",
+    )
+    async def workflow_websocket_endpoint(websocket: WebSocket):
+        """WebSocket endpoint for receiving real-time workflow events"""
+        requires_auth = bool(settings.os_security_key)
+        await websocket_manager.connect(websocket, requires_auth=requires_auth)
+
+        try:
+            while True:
+                data = await websocket.receive_text()
+                message = json.loads(data)
+                action = message.get("action")
+
+                # Handle authentication first
+                if action == "authenticate":
+                    token = message.get("token")
+                    if not token:
+                        await websocket.send_text(json.dumps({"event": "auth_error", "error": "Token is required"}))
+                        continue
+
+                    if validate_websocket_token(token, settings):
+                        await websocket_manager.authenticate_websocket(websocket)
+                    else:
+                        await websocket.send_text(json.dumps({"event": "auth_error", "error": "Invalid token"}))
+                        continue
+
+                # Check authentication for all other actions (only when required)
+                elif requires_auth and not websocket_manager.is_authenticated(websocket):
+                    await websocket.send_text(
+                        json.dumps(
+                            {
+                                "event": "auth_required",
+                                "error": "Authentication required. Send authenticate action with valid token.",
+                            }
+                        )
+                    )
+                    continue
+
+                # Handle authenticated actions
+                elif action == "ping":
+                    await websocket.send_text(json.dumps({"event": "pong"}))
+
+                elif action == "start-workflow":
+                    # Handle workflow execution directly via WebSocket
+                    await handle_workflow_via_websocket(websocket, message, os)
+
+                else:
+                    await websocket.send_text(json.dumps({"event": "error", "error": f"Unknown action: {action}"}))
+
+        except Exception as e:
+            if "1012" not in str(e):
+                logger.error(f"WebSocket error: {e}")
+        finally:
+            # Clean up the websocket connection
+            await websocket_manager.disconnect_websocket(websocket)
+
+    return ws_router
+
+
 def get_base_router(
     os: "AgentOS",
     settings: AgnoAPISettings = AgnoAPISettings(),
@@ -346,7 +467,6 @@ def get_base_router(
     - Agent management and execution
     - Team collaboration and coordination
     - Workflow automation and orchestration
-    - Real-time WebSocket communications
 
     All endpoints include detailed documentation, examples, and proper error handling.
     """
@@ -362,24 +482,6 @@ def get_base_router(
     )
 
     # -- Main Routes ---
-
-    @router.get(
-        "/health",
-        tags=["Core"],
-        operation_id="health_check",
-        summary="Health Check",
-        description="Check the health status of the AgentOS API. Returns a simple status indicator.",
-        response_model=HealthResponse,
-        responses={
-            200: {
-                "description": "API is healthy and operational",
-                "content": {"application/json": {"example": {"status": "ok"}}},
-            }
-        },
-    )
-    async def health_check() -> HealthResponse:
-        return HealthResponse(status="ok")
-
     @router.get(
         "/config",
         response_model=ConfigResponse,
@@ -1228,35 +1330,6 @@ def get_base_router(
 
     # -- Workflow routes ---
 
-    @router.websocket(
-        "/workflows/ws",
-        name="workflow_websocket",
-    )
-    async def workflow_websocket_endpoint(websocket: WebSocket):
-        """WebSocket endpoint for receiving real-time workflow events"""
-        await websocket_manager.connect(websocket)
-
-        try:
-            while True:
-                data = await websocket.receive_text()
-                message = json.loads(data)
-                action = message.get("action")
-
-                if action == "ping":
-                    await websocket.send_text(json.dumps({"event": "pong"}))
-
-                elif action == "start-workflow":
-                    # Handle workflow execution directly via WebSocket
-                    await handle_workflow_via_websocket(websocket, message, os)
-        except Exception as e:
-            if "1012" not in str(e):
-                logger.error(f"WebSocket error: {e}")
-        finally:
-            # Clean up any run_ids associated with this websocket
-            runs_to_remove = [run_id for run_id, ws in websocket_manager.active_connections.items() if ws == websocket]
-            for run_id in runs_to_remove:
-                await websocket_manager.disconnect_by_run_id(run_id)
-
     @router.get(
         "/workflows",
         response_model=List[WorkflowSummaryResponse],