agno 2.3.16__py3-none-any.whl → 2.3.18__py3-none-any.whl

agno/os/routers/workflows/router.py

@@ -1,5 +1,5 @@
 import json
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union
 from uuid import uuid4
 
 from fastapi import (
@@ -15,7 +15,13 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 
 from agno.exceptions import InputCheckError, OutputCheckError
-from agno.os.auth import get_authentication_dependency, require_resource_access, validate_websocket_token
+from agno.os.auth import (
+    get_auth_token_from_request,
+    get_authentication_dependency,
+    require_resource_access,
+    validate_websocket_token,
+)
+from agno.os.managers import event_buffer, websocket_manager
 from agno.os.routers.workflows.schema import WorkflowResponse
 from agno.os.schema import (
     BadRequestResponse,
@@ -31,110 +37,17 @@ from agno.os.utils import (
     get_request_kwargs,
     get_workflow_by_id,
 )
-from agno.run.workflow import WorkflowErrorEvent, WorkflowRunOutput
-from agno.utils.log import log_warning, logger
+from agno.run.base import RunStatus
+from agno.run.workflow import WorkflowErrorEvent
+from agno.utils.log import log_debug, log_warning, logger
+from agno.utils.serialize import json_serializer
+from agno.workflow.remote import RemoteWorkflow
 from agno.workflow.workflow import Workflow
 
 if TYPE_CHECKING:
     from agno.os.app import AgentOS
 
 
-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,
-        active_connections: Optional[Dict[str, WebSocket]] = None,
-    ):
-        # 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, requires_auth: bool = True):
-        """Accept WebSocket connection"""
-        await websocket.accept()
-        logger.debug("WebSocket connected")
-
-        # 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. 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
-        logger.debug(f"Registered WebSocket for run_id: {run_id}")
-
-    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)
-
-
-# Global manager instance
-websocket_manager = WebSocketManager(
-    active_connections={},
-)
-
-
 async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os: "AgentOS"):
     """Handle workflow execution directly via WebSocket"""
     try:
@@ -153,6 +66,12 @@ async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os:
             await websocket.send_text(json.dumps({"event": "error", "error": f"Workflow {workflow_id} not found"}))
             return
 
+        if isinstance(workflow, RemoteWorkflow):
+            await websocket.send_text(
+                json.dumps({"event": "error", "error": "Remote workflows are not supported via WebSocket"})
+            )
+            return
+
         # Generate session_id if not provided
         # Use workflow's default session_id if not provided in message
         if not session_id:
@@ -162,7 +81,7 @@ async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os:
                 session_id = str(uuid4())
 
         # Execute workflow in background with streaming
-        workflow_result = await workflow.arun(  # type: ignore
+        await workflow.arun(  # type: ignore
             input=user_message,
             session_id=session_id,
             user_id=user_id,
@@ -172,9 +91,9 @@ 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
+        # NOTE: Don't register the original websocket in the manager
+        # It's already handled by the WebSocketHandler passed to the workflow
+        # The manager is ONLY for reconnected clients (see handle_workflow_subscription)
 
     except (InputCheckError, OutputCheckError) as e:
         await websocket.send_text(
@@ -200,12 +119,175 @@ async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os:
         await websocket.send_text(json.dumps(error_payload))
 
 
+async def handle_workflow_subscription(websocket: WebSocket, message: dict, os: "AgentOS"):
+    """
+    Handle subscription/reconnection to an existing workflow run.
+
+    Allows clients to reconnect after page refresh or disconnection and catch up on missed events.
+    """
+    try:
+        run_id = message.get("run_id")
+        workflow_id = message.get("workflow_id")
+        session_id = message.get("session_id")
+        last_event_index = message.get("last_event_index")  # 0-based index of last received event
+
+        if not run_id:
+            await websocket.send_text(json.dumps({"event": "error", "error": "run_id is required for subscription"}))
+            return
+
+        # Check if run exists in event buffer
+        buffer_status = event_buffer.get_run_status(run_id)
+
+        if buffer_status is None:
+            # Run not in buffer - check database
+            if workflow_id and session_id:
+                workflow = get_workflow_by_id(workflow_id, os.workflows)
+                if workflow and isinstance(workflow, Workflow):
+                    workflow_run = await workflow.aget_run_output(run_id, session_id)
+
+                    if workflow_run:
+                        # Run exists in DB - send all events from DB
+                        if workflow_run.events:
+                            await websocket.send_text(
+                                json.dumps(
+                                    {
+                                        "event": "replay",
+                                        "run_id": run_id,
+                                        "status": workflow_run.status.value if workflow_run.status else "unknown",
+                                        "total_events": len(workflow_run.events),
+                                        "message": "Run completed. Replaying all events from database.",
+                                    }
+                                )
+                            )
+
+                            # Send events one by one
+                            for idx, event in enumerate(workflow_run.events):
+                                # Convert event to dict and add event_index
+                                event_dict = event.model_dump() if hasattr(event, "model_dump") else event.to_dict()
+                                event_dict["event_index"] = idx
+                                if "run_id" not in event_dict:
+                                    event_dict["run_id"] = run_id
+
+                                await websocket.send_text(json.dumps(event_dict, default=json_serializer))
+                        else:
+                            await websocket.send_text(
+                                json.dumps(
+                                    {
+                                        "event": "replay",
+                                        "run_id": run_id,
+                                        "status": workflow_run.status.value if workflow_run.status else "unknown",
+                                        "total_events": 0,
+                                        "message": "Run completed but no events stored.",
+                                    }
+                                )
+                            )
+                        return
+
+            # Run not found anywhere
+            await websocket.send_text(
+                json.dumps({"event": "error", "error": f"Run {run_id} not found in buffer or database"})
+            )
+            return
+
+        # Run is in buffer (still active or recently completed)
+        if buffer_status in [RunStatus.completed, RunStatus.error, RunStatus.cancelled]:
+            # Run finished - send all events from buffer
+            all_events = event_buffer.get_events(run_id, last_event_index=None)
+
+            await websocket.send_text(
+                json.dumps(
+                    {
+                        "event": "replay",
+                        "run_id": run_id,
+                        "status": buffer_status.value,
+                        "total_events": len(all_events),
+                        "message": f"Run {buffer_status.value}. Replaying all events.",
+                    }
+                )
+            )
+
+            # Send all events
+            for idx, buffered_event in enumerate(all_events):
+                # Convert event to dict and add event_index
+                event_dict = (
+                    buffered_event.model_dump() if hasattr(buffered_event, "model_dump") else buffered_event.to_dict()
+                )
+                event_dict["event_index"] = idx
+                if "run_id" not in event_dict:
+                    event_dict["run_id"] = run_id
+
+                await websocket.send_text(json.dumps(event_dict))
+            return
+
+        # Run is still active - send missed events and subscribe to new ones
+        missed_events = event_buffer.get_events(run_id, last_event_index)
+        current_event_count = event_buffer.get_event_count(run_id)
+
+        if missed_events:
+            # Send catch-up notification
+            await websocket.send_text(
+                json.dumps(
+                    {
+                        "event": "catch_up",
+                        "run_id": run_id,
+                        "status": "running",
+                        "missed_events": len(missed_events),
+                        "current_event_count": current_event_count,
+                        "message": f"Catching up on {len(missed_events)} missed events.",
+                    }
+                )
+            )
+
+            # Send missed events
+            start_index = (last_event_index + 1) if last_event_index is not None else 0
+            for idx, buffered_event in enumerate(missed_events):
+                # Convert event to dict and add event_index
+                event_dict = (
+                    buffered_event.model_dump() if hasattr(buffered_event, "model_dump") else buffered_event.to_dict()
+                )
+                event_dict["event_index"] = start_index + idx
+                if "run_id" not in event_dict:
+                    event_dict["run_id"] = run_id
+
+                await websocket.send_text(json.dumps(event_dict))
+
+        # Register websocket for future events
+        await websocket_manager.register_websocket(run_id, websocket)
+
+        # Send subscription confirmation
+        await websocket.send_text(
+            json.dumps(
+                {
+                    "event": "subscribed",
+                    "run_id": run_id,
+                    "status": "running",
+                    "current_event_count": current_event_count,
+                    "message": "Subscribed to workflow run. You will receive new events as they occur.",
+                }
+            )
+        )
+
+        log_debug(f"Client subscribed to workflow run {run_id} (last_event_index: {last_event_index})")
+
+    except Exception as e:
+        logger.error(f"Error handling workflow subscription: {e}")
+        await websocket.send_text(
+            json.dumps(
+                {
+                    "event": "error",
+                    "error": f"Subscription failed: {str(e)}",
+                }
+            )
+        )
+
+
 async def workflow_response_streamer(
-    workflow: Workflow,
-    input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel]] = None,
+    workflow: Union[Workflow, RemoteWorkflow],
+    input: Union[str, Dict[str, Any], List[Any], BaseModel],
     session_id: Optional[str] = None,
     user_id: Optional[str] = None,
     background_tasks: Optional[BackgroundTasks] = None,
+    auth_token: Optional[str] = None,
     **kwargs: Any,
 ) -> AsyncGenerator:
     try:
@@ -213,12 +295,21 @@ async def workflow_response_streamer(
         if background_tasks is not None:
             kwargs["background_tasks"] = background_tasks
 
-        run_response = workflow.arun(
+        if "stream_events" in kwargs:
+            stream_events = kwargs.pop("stream_events")
+        else:
+            stream_events = True
+
+        # Pass auth_token for remote workflows
+        if auth_token and isinstance(workflow, RemoteWorkflow):
+            kwargs["auth_token"] = auth_token
+
+        run_response = workflow.arun(  # type: ignore
             input=input,
             session_id=session_id,
             user_id=user_id,
             stream=True,
-            stream_events=True,
+            stream_events=stream_events,
             **kwargs,
         )
 
@@ -368,6 +459,10 @@ def get_websocket_router(
                     # Handle workflow execution directly via WebSocket
                     await handle_workflow_via_websocket(websocket, message, os)
 
+                elif action == "reconnect":
+                    # Subscribe/reconnect to an existing workflow run
+                    await handle_workflow_subscription(websocket, message, os)
+
                 else:
                     await websocket.send_text(json.dumps({"event": "error", "error": f"Unknown action: {action}"}))
 
@@ -479,8 +574,10 @@ def get_workflow_router(
         workflow = get_workflow_by_id(workflow_id, os.workflows)
         if workflow is None:
             raise HTTPException(status_code=404, detail="Workflow not found")
-
-        return await WorkflowResponse.from_workflow(workflow)
+        if isinstance(workflow, RemoteWorkflow):
+            return await workflow.get_workflow_config()
+        else:
+            return await WorkflowResponse.from_workflow(workflow=workflow)
 
     @router.post(
         "/workflows/{workflow_id}/runs",
@@ -528,25 +625,25 @@ def get_workflow_router(
     ):
         kwargs = await get_request_kwargs(request, create_workflow_run)
 
-        if hasattr(request.state, "user_id"):
+        if hasattr(request.state, "user_id") and request.state.user_id is not None:
             if user_id:
                 log_warning("User ID parameter passed in both request state and kwargs, using request state")
             user_id = request.state.user_id
-        if hasattr(request.state, "session_id"):
+        if hasattr(request.state, "session_id") and request.state.session_id is not None:
             if session_id:
                 log_warning("Session ID parameter passed in both request state and kwargs, using request state")
             session_id = request.state.session_id
-        if hasattr(request.state, "session_state"):
+        if hasattr(request.state, "session_state") and request.state.session_state is not None:
             session_state = request.state.session_state
             if "session_state" in kwargs:
                 log_warning("Session state parameter passed in both request state and kwargs, using request state")
             kwargs["session_state"] = session_state
-        if hasattr(request.state, "dependencies"):
+        if hasattr(request.state, "dependencies") and request.state.dependencies is not None:
             dependencies = request.state.dependencies
             if "dependencies" in kwargs:
                 log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
             kwargs["dependencies"] = dependencies
-        if hasattr(request.state, "metadata"):
+        if hasattr(request.state, "metadata") and request.state.metadata is not None:
             metadata = request.state.metadata
             if "metadata" in kwargs:
                 log_warning("Metadata parameter passed in both request state and kwargs, using request state")
@@ -563,6 +660,9 @@ def get_workflow_router(
             logger.debug("Creating new session")
             session_id = str(uuid4())
 
+        # Extract auth token for remote workflows
+        auth_token = get_auth_token_from_request(request)
+
         # Return based on stream parameter
         try:
             if stream:
@@ -573,11 +673,16 @@ def get_workflow_router(
                         session_id=session_id,
                         user_id=user_id,
                         background_tasks=background_tasks,
+                        auth_token=auth_token,
                         **kwargs,
                     ),
                     media_type="text/event-stream",
                 )
             else:
+                # Pass auth_token for remote workflows
+                if auth_token and isinstance(workflow, RemoteWorkflow):
+                    kwargs["auth_token"] = auth_token
+
                 run_response = await workflow.arun(
                     input=message,
                     session_id=session_id,
@@ -616,8 +721,9 @@ def get_workflow_router(
         if workflow is None:
             raise HTTPException(status_code=404, detail="Workflow not found")
 
-        if not workflow.cancel_run(run_id=run_id):
-            raise HTTPException(status_code=500, detail="Failed to cancel run")
+        cancelled = workflow.cancel_run(run_id=run_id)
+        if not cancelled:
+            raise HTTPException(status_code=500, detail="Failed to cancel run - run not found or already completed")
 
         return JSONResponse(content={}, status_code=200)
 

agno/os/routers/workflows/schema.py

@@ -4,11 +4,75 @@ from pydantic import BaseModel, Field
 
 from agno.os.routers.agents.schema import AgentResponse
 from agno.os.routers.teams.schema import TeamResponse
-from agno.os.utils import get_workflow_input_schema_dict
 from agno.workflow.agent import WorkflowAgent
 from agno.workflow.workflow import Workflow
 
 
+def _generate_schema_from_params(params: Dict[str, Any]) -> Dict[str, Any]:
+    """Convert function parameters to JSON schema"""
+    properties: Dict[str, Any] = {}
+    required: List[str] = []
+
+    for param_name, param_info in params.items():
+        # Skip the default 'message' parameter for custom kwargs workflows
+        if param_name == "message":
+            continue
+
+        # Map Python types to JSON schema types
+        param_type = param_info.get("annotation", "str")
+        default_value = param_info.get("default")
+        is_required = param_info.get("required", False)
+
+        # Convert Python type annotations to JSON schema types
+        if param_type == "str":
+            properties[param_name] = {"type": "string"}
+        elif param_type == "bool":
+            properties[param_name] = {"type": "boolean"}
+        elif param_type == "int":
+            properties[param_name] = {"type": "integer"}
+        elif param_type == "float":
+            properties[param_name] = {"type": "number"}
+        elif "List" in str(param_type):
+            properties[param_name] = {"type": "array", "items": {"type": "string"}}
+        else:
+            properties[param_name] = {"type": "string"}  # fallback
+
+        # Add default value if present
+        if default_value is not None:
+            properties[param_name]["default"] = default_value
+
+        # Add to required if no default value
+        if is_required and default_value is None:
+            required.append(param_name)
+
+    schema = {"type": "object", "properties": properties}
+
+    if required:
+        schema["required"] = required
+
+    return schema
+
+
+def get_workflow_input_schema_dict(workflow: Workflow) -> Optional[Dict[str, Any]]:
+    """Get input schema as dictionary for API responses"""
+
+    # Priority 1: Explicit input_schema (Pydantic model)
+    if workflow.input_schema is not None:
+        try:
+            return workflow.input_schema.model_json_schema()
+        except Exception:
+            return None
+
+    # Priority 2: Auto-generate from custom kwargs
+    if workflow.steps and callable(workflow.steps):
+        custom_params = workflow.run_parameters
+        if custom_params and len(custom_params) > 1:  # More than just 'message'
+            return _generate_schema_from_params(custom_params)
+
+    # Priority 3: No schema (expects string message)
+    return None
+
+
 class WorkflowResponse(BaseModel):
     id: Optional[str] = Field(None, description="Unique identifier for the workflow")
     name: Optional[str] = Field(None, description="Name of the workflow")

agno/os/schema.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, Generic, List, Optional, TypeVar, Union
 from pydantic import BaseModel, ConfigDict, Field
 
 from agno.agent import Agent
+from agno.agent.remote import RemoteAgent
 from agno.db.base import SessionType
 from agno.os.config import (
     ChatConfig,
@@ -21,7 +22,9 @@ from agno.os.utils import (
     get_session_name,
 )
 from agno.session import AgentSession, TeamSession, WorkflowSession
+from agno.team.remote import RemoteTeam
 from agno.team.team import Team
+from agno.workflow.remote import RemoteWorkflow
 from agno.workflow.workflow import Workflow
 
 
@@ -102,7 +105,7 @@ class AgentSummaryResponse(BaseModel):
     db_id: Optional[str] = Field(None, description="Database identifier")
 
     @classmethod
-    def from_agent(cls, agent: Agent) -> "AgentSummaryResponse":
+    def from_agent(cls, agent: Union[Agent, RemoteAgent]) -> "AgentSummaryResponse":
         return cls(id=agent.id, name=agent.name, description=agent.description, db_id=agent.db.id if agent.db else None)
 
 
@@ -113,8 +116,9 @@ class TeamSummaryResponse(BaseModel):
     db_id: Optional[str] = Field(None, description="Database identifier")
 
     @classmethod
-    def from_team(cls, team: Team) -> "TeamSummaryResponse":
-        return cls(id=team.id, name=team.name, description=team.description, db_id=team.db.id if team.db else None)
+    def from_team(cls, team: Union[Team, RemoteTeam]) -> "TeamSummaryResponse":
+        db_id = team.db.id if team.db else None
+        return cls(id=team.id, name=team.name, description=team.description, db_id=db_id)
 
 
 class WorkflowSummaryResponse(BaseModel):
@@ -124,12 +128,13 @@ class WorkflowSummaryResponse(BaseModel):
     db_id: Optional[str] = Field(None, description="Database identifier")
 
     @classmethod
-    def from_workflow(cls, workflow: Workflow) -> "WorkflowSummaryResponse":
+    def from_workflow(cls, workflow: Union[Workflow, RemoteWorkflow]) -> "WorkflowSummaryResponse":
+        db_id = workflow.db.id if workflow.db else None
         return cls(
             id=workflow.id,
             name=workflow.name,
             description=workflow.description,
-            db_id=workflow.db.id if workflow.db else None,
+            db_id=db_id,
         )
 
 
@@ -182,18 +187,39 @@ class SessionSchema(BaseModel):
 
     @classmethod
     def from_dict(cls, session: Dict[str, Any]) -> "SessionSchema":
-        session_name = get_session_name(session)
+        session_name = session.get("session_name")
+        if not session_name:
+            session_name = get_session_name(session)
         session_data = session.get("session_data", {}) or {}
+
+        created_at = session.get("created_at", 0)
+        updated_at = session.get("updated_at", 0)
+
+        # Handle created_at and updated_at as either ISO 8601 string or timestamp
+        def parse_datetime(val):
+            if isinstance(val, str):
+                try:
+                    # Accept both with and without Z
+                    if val.endswith("Z"):
+                        val = val[:-1] + "+00:00"
+                    return datetime.fromisoformat(val)
+                except Exception:
+                    return None
+            elif isinstance(val, (int, float)):
+                try:
+                    return datetime.fromtimestamp(val, tz=timezone.utc)
+                except Exception:
+                    return None
+            return None
+
+        created_at = parse_datetime(session.get("created_at", 0))
+        updated_at = parse_datetime(session.get("updated_at", 0))
         return cls(
             session_id=session.get("session_id", ""),
             session_name=session_name,
             session_state=session_data.get("session_state", None),
-            created_at=datetime.fromtimestamp(session.get("created_at", 0), tz=timezone.utc)
-            if session.get("created_at")
-            else None,
-            updated_at=datetime.fromtimestamp(session.get("updated_at", 0), tz=timezone.utc)
-            if session.get("updated_at")
-            else None,
+            created_at=created_at,
+            updated_at=updated_at,
         )