autobyteus 1.1.3__py3-none-any.whl → 1.1.4__py3-none-any.whl

autobyteus/workflow/shutdown_steps/agent_team_shutdown_step.py

@@ -0,0 +1,42 @@
+# file: autobyteus/autobyteus/workflow/shutdown_steps/agent_team_shutdown_step.py
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from autobyteus.workflow.shutdown_steps.base_workflow_shutdown_step import BaseWorkflowShutdownStep
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.context.workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+class AgentTeamShutdownStep(BaseWorkflowShutdownStep):
+    """Shutdown step to gracefully stop all running agents in the workflow."""
+    async def execute(self, context: 'WorkflowContext') -> bool:
+        workflow_id = context.workflow_id
+        logger.info(f"Workflow '{workflow_id}': Executing AgentTeamShutdownStep.")
+        
+        team_manager = context.team_manager
+        if not team_manager:
+            logger.warning(f"Workflow '{workflow_id}': No TeamManager found, cannot shut down agents.")
+            return True
+
+        # Get the list of all created agents from the single source of truth.
+        all_agents = team_manager.get_all_agents()
+        running_agents = [agent for agent in all_agents if agent.is_running]
+        
+        if not running_agents:
+            logger.info(f"Workflow '{workflow_id}': No running agents to shut down.")
+            return True
+        
+        logger.info(f"Workflow '{workflow_id}': Shutting down {len(running_agents)} running agents.")
+        stop_tasks = [agent.stop(timeout=10.0) for agent in running_agents]
+        results = await asyncio.gather(*stop_tasks, return_exceptions=True)
+        
+        all_successful = True
+        for agent, result in zip(running_agents, results):
+            if isinstance(result, Exception):
+                logger.error(f"Workflow '{workflow_id}': Error stopping agent '{agent.agent_id}': {result}", exc_info=result)
+                all_successful = False
+        
+        return all_successful

autobyteus/workflow/shutdown_steps/base_workflow_shutdown_step.py

@@ -0,0 +1,16 @@
+# file: autobyteus/autobyteus/workflow/shutdown_steps/base_workflow_shutdown_step.py
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.context.workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+class BaseWorkflowShutdownStep(ABC):
+    """Abstract base class for individual steps in the workflow shutdown process."""
+    @abstractmethod
+    async def execute(self, context: 'WorkflowContext') -> bool:
+        """Executes the shutdown step."""
+        raise NotImplementedError("Subclasses must implement the 'execute' method.")

autobyteus/workflow/shutdown_steps/bridge_cleanup_step.py

@@ -0,0 +1,28 @@
+# file: autobyteus/autobyteus/workflow/shutdown_steps/bridge_cleanup_step.py
+import logging
+from typing import TYPE_CHECKING
+
+from autobyteus.workflow.shutdown_steps.base_workflow_shutdown_step import BaseWorkflowShutdownStep
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.context.workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+class BridgeCleanupStep(BaseWorkflowShutdownStep):
+    """Shutdown step to gracefully stop all AgentEventBridge instances via the multiplexer."""
+    async def execute(self, context: 'WorkflowContext') -> bool:
+        workflow_id = context.workflow_id
+        logger.info(f"Workflow '{workflow_id}': Executing BridgeCleanupStep.")
+        
+        multiplexer = context.multiplexer
+        if not multiplexer:
+            logger.warning(f"Workflow '{workflow_id}': No AgentEventMultiplexer found, cannot shut down event bridges.")
+            return True
+
+        try:
+            await multiplexer.shutdown()
+            return True
+        except Exception as e:
+            logger.error(f"Workflow '{workflow_id}': Error shutting down agent event bridges via multiplexer: {e}", exc_info=True)
+            return False

autobyteus/workflow/shutdown_steps/sub_workflow_shutdown_step.py

@@ -0,0 +1,41 @@
+# file: autobyteus/autobyteus/workflow/shutdown_steps/sub_workflow_shutdown_step.py
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from autobyteus.workflow.shutdown_steps.base_workflow_shutdown_step import BaseWorkflowShutdownStep
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.context.workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+class SubWorkflowShutdownStep(BaseWorkflowShutdownStep):
+    """Shutdown step to gracefully stop all running sub-workflows."""
+    async def execute(self, context: 'WorkflowContext') -> bool:
+        workflow_id = context.workflow_id
+        logger.info(f"Workflow '{workflow_id}': Executing SubWorkflowShutdownStep.")
+        
+        team_manager = context.team_manager
+        if not team_manager:
+            logger.warning(f"Workflow '{workflow_id}': No TeamManager found, cannot shut down sub-workflows.")
+            return True
+
+        all_sub_workflows = team_manager.get_all_sub_workflows()
+        running_sub_workflows = [wf for wf in all_sub_workflows if wf.is_running]
+        
+        if not running_sub_workflows:
+            logger.info(f"Workflow '{workflow_id}': No running sub-workflows to shut down.")
+            return True
+        
+        logger.info(f"Workflow '{workflow_id}': Shutting down {len(running_sub_workflows)} running sub-workflows.")
+        stop_tasks = [wf.stop(timeout=20.0) for wf in running_sub_workflows]
+        results = await asyncio.gather(*stop_tasks, return_exceptions=True)
+        
+        all_successful = True
+        for wf, result in zip(running_sub_workflows, results):
+            if isinstance(result, Exception):
+                logger.error(f"Workflow '{workflow_id}': Error stopping sub-workflow '{wf.name}': {result}", exc_info=result)
+                all_successful = False
+        
+        return all_successful

autobyteus/workflow/shutdown_steps/workflow_shutdown_orchestrator.py

@@ -0,0 +1,35 @@
+# file: autobyteus/autobyteus/workflow/shutdown_steps/workflow_shutdown_orchestrator.py
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+from autobyteus.workflow.shutdown_steps.base_workflow_shutdown_step import BaseWorkflowShutdownStep
+from autobyteus.workflow.shutdown_steps.bridge_cleanup_step import BridgeCleanupStep
+from autobyteus.workflow.shutdown_steps.sub_workflow_shutdown_step import SubWorkflowShutdownStep
+from autobyteus.workflow.shutdown_steps.agent_team_shutdown_step import AgentTeamShutdownStep
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.context.workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+class WorkflowShutdownOrchestrator:
+    """Orchestrates the workflow's shutdown process."""
+    def __init__(self, steps: Optional[List[BaseWorkflowShutdownStep]] = None):
+        self.shutdown_steps = steps or [
+            BridgeCleanupStep(),
+            SubWorkflowShutdownStep(),
+            AgentTeamShutdownStep(),
+        ]
+
+    async def run(self, context: 'WorkflowContext') -> bool:
+        workflow_id = context.workflow_id
+        logger.info(f"Workflow '{workflow_id}': Shutdown orchestrator starting.")
+        
+        all_successful = True
+        for step in self.shutdown_steps:
+            if not await step.execute(context):
+                logger.error(f"Workflow '{workflow_id}': Shutdown step {step.__class__.__name__} failed.")
+                all_successful = False
+        
+        logger.info(f"Workflow '{workflow_id}': Shutdown orchestration completed.")
+        return all_successful

autobyteus/workflow/streaming/__init__.py

@@ -0,0 +1,26 @@
+# file: autobyteus/autobyteus/workflow/streaming/__init__.py
+"""
+Components related to workflow output streaming.
+"""
+from .workflow_event_notifier import WorkflowExternalEventNotifier
+from .workflow_event_stream import WorkflowEventStream
+from .workflow_stream_events import WorkflowStreamEvent, WorkflowStreamDataPayload
+from .workflow_stream_event_payloads import (
+    BaseWorkflowSpecificPayload,
+    WorkflowPhaseTransitionData,
+    AgentEventRebroadcastPayload,
+)
+from .agent_event_bridge import AgentEventBridge
+from .agent_event_multiplexer import AgentEventMultiplexer
+
+__all__ = [
+    "WorkflowExternalEventNotifier",
+    "WorkflowEventStream",
+    "WorkflowStreamEvent",
+    "WorkflowStreamDataPayload",
+    "BaseWorkflowSpecificPayload",
+    "WorkflowPhaseTransitionData",
+    "AgentEventRebroadcastPayload",
+    "AgentEventBridge",
+    "AgentEventMultiplexer",
+]

autobyteus/workflow/streaming/agent_event_bridge.py

@@ -0,0 +1,48 @@
+# file: autobyteus/autobyteus/workflow/streaming/agent_event_bridge.py
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from autobyteus.agent.streaming.agent_event_stream import AgentEventStream
+
+if TYPE_CHECKING:
+    from autobyteus.agent.agent import Agent
+    from autobyteus.workflow.streaming.workflow_event_notifier import WorkflowExternalEventNotifier
+
+logger = logging.getLogger(__name__)
+
+class AgentEventBridge:
+    """
+    A dedicated component that bridges events from a single agent's event stream
+    to the main workflow notifier. This is the core of the multiplexing pattern.
+    """
+    def __init__(self, agent: 'Agent', agent_name: str, notifier: 'WorkflowExternalEventNotifier', loop: asyncio.AbstractEventLoop):
+        self._agent_name = agent_name
+        self._notifier = notifier
+        self._stream = AgentEventStream(agent)
+        self._task: asyncio.Task = loop.create_task(self._run())
+        logger.info(f"AgentEventBridge created and task started for agent '{agent_name}'.")
+
+    async def _run(self):
+        """The background task that consumes from the stream and publishes to the notifier."""
+        try:
+            async for event in self._stream.all_events():
+                self._notifier.publish_agent_event(self._agent_name, event)
+        except asyncio.CancelledError:
+            logger.info(f"AgentEventBridge task for '{self._agent_name}' was cancelled.")
+        except Exception as e:
+            logger.error(f"Error in AgentEventBridge for '{self._agent_name}': {e}", exc_info=True)
+        finally:
+            logger.debug(f"AgentEventBridge task for '{self._agent_name}' is finishing.")
+
+    async def cancel(self):
+        """Gracefully stops the bridge."""
+        logger.info(f"Cancelling AgentEventBridge for '{self._agent_name}'.")
+        if not self._task.done():
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass # Expected
+        await self._stream.close()
+        logger.info(f"AgentEventBridge for '{self._agent_name}' cancelled successfully.")

autobyteus/workflow/streaming/agent_event_multiplexer.py

@@ -0,0 +1,70 @@
+# file: autobyteus/autobyteus/workflow/streaming/agent_event_multiplexer.py
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Dict, Optional
+
+from autobyteus.workflow.streaming.agent_event_bridge import AgentEventBridge
+from autobyteus.workflow.streaming.workflow_event_bridge import WorkflowEventBridge
+
+if TYPE_CHECKING:
+    from autobyteus.agent.agent import Agent
+    from autobyteus.workflow.agentic_workflow import AgenticWorkflow
+    from autobyteus.workflow.streaming.workflow_event_notifier import WorkflowExternalEventNotifier
+    from autobyteus.workflow.runtime.workflow_worker import WorkflowWorker
+
+logger = logging.getLogger(__name__)
+
+class AgentEventMultiplexer:
+    """
+    Manages the lifecycle of event bridges for all nodes (agents and sub-workflows).
+    It creates, tracks, and shuts down the bridges that forward node events
+    to the workflow's main event stream.
+    """
+    def __init__(self, workflow_id: str, notifier: 'WorkflowExternalEventNotifier', worker_ref: 'WorkflowWorker'):
+        self._workflow_id = workflow_id
+        self._notifier = notifier
+        self._worker = worker_ref
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._agent_bridges: Dict[str, AgentEventBridge] = {}
+        self._workflow_bridges: Dict[str, WorkflowEventBridge] = {}
+        logger.info(f"AgentEventMultiplexer initialized for workflow '{self._workflow_id}'.")
+
+    def _get_loop(self) -> asyncio.AbstractEventLoop:
+        """Retrieves the event loop from the worker on-demand."""
+        if self._loop is None or self._loop.is_closed():
+            self._loop = self._worker.get_worker_loop()
+            if self._loop is None:
+                raise RuntimeError(f"Workflow worker loop for workflow '{self._workflow_id}' is not available or not running.")
+        return self._loop
+
+    def start_bridging_agent_events(self, agent: 'Agent', agent_name: str):
+        """Creates and starts an AgentEventBridge for a direct agent node."""
+        if agent_name in self._agent_bridges:
+            logger.warning(f"Event bridge for agent '{agent_name}' already exists. Skipping creation.")
+            return
+
+        bridge = AgentEventBridge(agent=agent, agent_name=agent_name, notifier=self._notifier, loop=self._get_loop())
+        self._agent_bridges[agent_name] = bridge
+        logger.info(f"AgentEventMultiplexer started agent event bridge for '{agent_name}'.")
+
+    def start_bridging_workflow_events(self, sub_workflow: 'AgenticWorkflow', node_name: str):
+        """Creates and starts a WorkflowEventBridge for a sub-workflow node."""
+        if node_name in self._workflow_bridges:
+            logger.warning(f"Event bridge for sub-workflow '{node_name}' already exists. Skipping creation.")
+            return
+            
+        bridge = WorkflowEventBridge(sub_workflow=sub_workflow, sub_workflow_node_name=node_name, parent_notifier=self._notifier, loop=self._get_loop())
+        self._workflow_bridges[node_name] = bridge
+        logger.info(f"AgentEventMultiplexer started workflow event bridge for '{node_name}'.")
+
+    async def shutdown(self):
+        """Gracefully shuts down all active event bridges."""
+        logger.info(f"AgentEventMultiplexer for '{self._workflow_id}' shutting down all event bridges.")
+        agent_bridge_tasks = [b.cancel() for b in self._agent_bridges.values()]
+        workflow_bridge_tasks = [b.cancel() for b in self._workflow_bridges.values()]
+        
+        await asyncio.gather(*(agent_bridge_tasks + workflow_bridge_tasks), return_exceptions=True)
+        
+        self._agent_bridges.clear()
+        self._workflow_bridges.clear()
+        logger.info(f"All event bridges for workflow '{self._workflow_id}' have been shut down by multiplexer.")

autobyteus/workflow/streaming/workflow_event_bridge.py

@@ -0,0 +1,50 @@
+# file: autobyteus/autobyteus/workflow/streaming/workflow_event_bridge.py
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from autobyteus.workflow.streaming.workflow_event_stream import WorkflowEventStream
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.agentic_workflow import AgenticWorkflow
+    from autobyteus.workflow.streaming.workflow_event_notifier import WorkflowExternalEventNotifier
+
+logger = logging.getLogger(__name__)
+
+class WorkflowEventBridge:
+    """
+    A dedicated component that bridges events from a sub-workflow's event stream
+    to the parent workflow's notifier.
+    """
+    def __init__(self, sub_workflow: 'AgenticWorkflow', sub_workflow_node_name: str, parent_notifier: 'WorkflowExternalEventNotifier', loop: asyncio.AbstractEventLoop):
+        self._sub_workflow = sub_workflow
+        self._sub_workflow_node_name = sub_workflow_node_name
+        self._parent_notifier = parent_notifier
+        self._stream = WorkflowEventStream(sub_workflow)
+        self._task: asyncio.Task = loop.create_task(self._run())
+        logger.info(f"WorkflowEventBridge created and task started for sub-workflow '{sub_workflow_node_name}'.")
+
+    async def _run(self):
+        """The background task that consumes from the sub-workflow stream and re-publishes."""
+        try:
+            async for event in self._stream.all_events():
+                # Re-broadcast the event to the parent, adding the sub-workflow context.
+                self._parent_notifier.publish_sub_workflow_event(self._sub_workflow_node_name, event)
+        except asyncio.CancelledError:
+            logger.info(f"WorkflowEventBridge task for '{self._sub_workflow_node_name}' was cancelled.")
+        except Exception as e:
+            logger.error(f"Error in WorkflowEventBridge for '{self._sub_workflow_node_name}': {e}", exc_info=True)
+        finally:
+            logger.debug(f"WorkflowEventBridge task for '{self._sub_workflow_node_name}' is finishing.")
+
+    async def cancel(self):
+        """Gracefully stops the bridge."""
+        logger.info(f"Cancelling WorkflowEventBridge for '{self._sub_workflow_node_name}'.")
+        if not self._task.done():
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass # Expected
+        await self._stream.close()
+        logger.info(f"WorkflowEventBridge for '{self._sub_workflow_node_name}' cancelled successfully.")

autobyteus/workflow/streaming/workflow_event_notifier.py

@@ -0,0 +1,83 @@
+# file: autobyteus/autobyteus/workflow/streaming/workflow_event_notifier.py
+import logging
+from typing import Optional, Dict, Any, TYPE_CHECKING
+
+from autobyteus.events.event_emitter import EventEmitter
+from autobyteus.events.event_types import EventType
+from autobyteus.workflow.phases.workflow_operational_phase import WorkflowOperationalPhase
+from autobyteus.agent.streaming.stream_events import StreamEvent as AgentStreamEvent
+from .workflow_stream_events import WorkflowStreamEvent, AgentEventRebroadcastPayload, WorkflowPhaseTransitionData, SubWorkflowEventRebroadcastPayload
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.runtime.workflow_runtime import WorkflowRuntime
+    from autobyteus.workflow.streaming.workflow_stream_events import WorkflowStreamEvent as WorkflowStreamEventTypeHint
+
+logger = logging.getLogger(__name__)
+
+class WorkflowExternalEventNotifier(EventEmitter):
+    """
+    Responsible for emitting unified WorkflowStreamEvents for consumption by
+    external listeners (like a UI or the WorkflowEventStream).
+    """
+    def __init__(self, workflow_id: str, runtime_ref: 'WorkflowRuntime'):
+        super().__init__()
+        self.workflow_id = workflow_id
+        self.runtime_ref = runtime_ref
+        logger.debug(f"WorkflowExternalEventNotifier initialized for workflow '{self.workflow_id}'.")
+
+    def _emit_event(self, event: 'WorkflowStreamEventTypeHint'):
+        """
+        Emits a fully-formed WorkflowStreamEvent.
+        A new generic event type is used for the underlying pub/sub system to carry
+        the unified event object.
+        """
+        self.emit(EventType.WORKFLOW_STREAM_EVENT, payload=event)
+
+    def notify_phase_change(self, new_phase: WorkflowOperationalPhase, old_phase: Optional[WorkflowOperationalPhase], extra_data: Optional[Dict[str, Any]] = None):
+        """
+        Notifies of a workflow phase transition by creating and emitting a
+        'WORKFLOW' sourced event.
+        """
+        payload_dict = {
+            "new_phase": new_phase,
+            "old_phase": old_phase,
+            "error_message": extra_data.get("error_message") if extra_data else None,
+        }
+        filtered_payload_dict = {k: v for k, v in payload_dict.items() if v is not None}
+        
+        event = WorkflowStreamEvent(
+            workflow_id=self.workflow_id,
+            event_source_type="WORKFLOW",
+            data=WorkflowPhaseTransitionData(**filtered_payload_dict)
+        )
+        self._emit_event(event)
+    
+    def publish_agent_event(self, agent_name: str, agent_event: AgentStreamEvent):
+        """
+        Wraps an event from a direct member agent and publishes it on the main workflow stream
+        as an 'AGENT' sourced event.
+        """
+        event = WorkflowStreamEvent(
+            workflow_id=self.workflow_id,
+            event_source_type="AGENT",
+            data=AgentEventRebroadcastPayload(
+                agent_name=agent_name,
+                agent_event=agent_event
+            )
+        )
+        self._emit_event(event)
+        
+    def publish_sub_workflow_event(self, sub_workflow_node_name: str, sub_workflow_event: 'WorkflowStreamEventTypeHint'):
+        """
+        Wraps an event from a sub-workflow and publishes it on the parent workflow stream
+        as a 'SUB_WORKFLOW' sourced event.
+        """
+        event = WorkflowStreamEvent(
+            workflow_id=self.workflow_id,
+            event_source_type="SUB_WORKFLOW",
+            data=SubWorkflowEventRebroadcastPayload(
+                sub_workflow_node_name=sub_workflow_node_name,
+                sub_workflow_event=sub_workflow_event
+            )
+        )
+        self._emit_event(event)

autobyteus/workflow/streaming/workflow_event_stream.py

@@ -0,0 +1,33 @@
+# file: autobyteus/autobyteus/workflow/streaming/workflow_event_stream.py
+import asyncio
+import queue
+from typing import AsyncIterator, TYPE_CHECKING
+
+from autobyteus.events.event_types import EventType
+from autobyteus.workflow.streaming.workflow_stream_events import WorkflowStreamEvent
+from autobyteus.agent.streaming.queue_streamer import stream_queue_items
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.agentic_workflow import AgenticWorkflow
+
+_SENTINEL = object()
+
+class WorkflowEventStream:
+    """Consumes events from a WorkflowExternalEventNotifier for a specific workflow."""
+    def __init__(self, workflow: 'AgenticWorkflow'):
+        self.workflow_id = workflow.workflow_id
+        self._internal_q = queue.Queue()
+        self._notifier = workflow._runtime.notifier
+        self._notifier.subscribe(EventType.WORKFLOW_STREAM_EVENT, self._handle_event)
+
+    def _handle_event(self, payload: WorkflowStreamEvent, **kwargs):
+        if isinstance(payload, WorkflowStreamEvent) and payload.workflow_id == self.workflow_id:
+            self._internal_q.put(payload)
+
+    async def close(self):
+        self._notifier.unsubscribe(EventType.WORKFLOW_STREAM_EVENT, self._handle_event)
+        await asyncio.get_running_loop().run_in_executor(None, self._internal_q.put, _SENTINEL)
+
+    def all_events(self) -> AsyncIterator[WorkflowStreamEvent]:
+        """The primary method to consume all structured events from the workflow."""
+        return stream_queue_items(self._internal_q, _SENTINEL, f"workflow_{self.workflow_id}_stream")

autobyteus/workflow/streaming/workflow_stream_event_payloads.py

@@ -0,0 +1,28 @@
+# file: autobyteus/autobyteus/workflow/streaming/workflow_stream_event_payloads.py
+from typing import Optional, Any
+from pydantic import BaseModel, Field
+from autobyteus.workflow.phases.workflow_operational_phase import WorkflowOperationalPhase
+from autobyteus.agent.streaming.stream_events import StreamEvent as AgentStreamEvent
+# Need to use a forward reference string to avoid circular import at runtime
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from autobyteus.workflow.streaming.workflow_stream_events import WorkflowStreamEvent
+
+# --- Payloads for events originating from the "WORKFLOW" source ---
+class BaseWorkflowSpecificPayload(BaseModel):
+    pass
+
+class WorkflowPhaseTransitionData(BaseWorkflowSpecificPayload):
+    new_phase: WorkflowOperationalPhase
+    old_phase: Optional[WorkflowOperationalPhase] = None
+    error_message: Optional[str] = None
+
+# --- Payload for events originating from the "AGENT" source ---
+class AgentEventRebroadcastPayload(BaseModel):
+    agent_name: str # The friendly name, e.g., "Researcher_1"
+    agent_event: AgentStreamEvent # The original, unmodified event from the agent
+
+# --- Payload for events originating from the "SUB_WORKFLOW" source ---
+class SubWorkflowEventRebroadcastPayload(BaseModel):
+    sub_workflow_node_name: str # The friendly name of the sub-workflow node
+    sub_workflow_event: "WorkflowStreamEvent" = Field(..., description="The original, unmodified event from the sub-workflow's stream")

autobyteus/workflow/streaming/workflow_stream_events.py

@@ -0,0 +1,45 @@
+# file: autobyteus/autobyteus/workflow/streaming/workflow_stream_events.py
+import datetime
+import uuid
+from typing import Literal, Union
+from pydantic import BaseModel, Field, model_validator
+
+from .workflow_stream_event_payloads import WorkflowPhaseTransitionData, AgentEventRebroadcastPayload, SubWorkflowEventRebroadcastPayload
+
+# A union of all possible payloads for a "WORKFLOW" sourced event.
+WorkflowSpecificPayload = Union[WorkflowPhaseTransitionData]
+
+# The top-level discriminated union for the main event stream's payload.
+WorkflowStreamDataPayload = Union[WorkflowSpecificPayload, AgentEventRebroadcastPayload, SubWorkflowEventRebroadcastPayload]
+
+class WorkflowStreamEvent(BaseModel):
+    event_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    timestamp: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
+    workflow_id: str
+    event_source_type: Literal["WORKFLOW", "AGENT", "SUB_WORKFLOW"]
+    data: WorkflowStreamDataPayload
+
+    @model_validator(mode='after')
+    def check_data_matches_source_type(self) -> 'WorkflowStreamEvent':
+        is_agent_event = self.event_source_type == "AGENT"
+        is_agent_payload = isinstance(self.data, AgentEventRebroadcastPayload)
+
+        is_sub_workflow_event = self.event_source_type == "SUB_WORKFLOW"
+        is_sub_workflow_payload = isinstance(self.data, SubWorkflowEventRebroadcastPayload)
+
+        is_workflow_event = self.event_source_type == "WORKFLOW"
+        is_workflow_payload = isinstance(self.data, WorkflowPhaseTransitionData)
+
+        if is_agent_event and not is_agent_payload:
+            raise ValueError("event_source_type is 'AGENT' but data is not an AgentEventRebroadcastPayload")
+        
+        if is_sub_workflow_event and not is_sub_workflow_payload:
+            raise ValueError("event_source_type is 'SUB_WORKFLOW' but data is not a SubWorkflowEventRebroadcastPayload")
+        
+        if is_workflow_event and not is_workflow_payload:
+            raise ValueError("event_source_type is 'WORKFLOW' but data is not a valid workflow-specific payload")
+
+        return self
+
+# This is necessary for Pydantic v2 to correctly handle the recursive model
+SubWorkflowEventRebroadcastPayload.model_rebuild()

autobyteus/workflow/utils/__init__.py

@@ -0,0 +1,9 @@
+# file: autobyteus/autobyteus/workflow/utils/__init__.py
+"""
+Utility functions for interacting with workflows.
+"""
+from .wait_for_idle import wait_for_workflow_to_be_idle
+
+__all__ = [
+    "wait_for_workflow_to_be_idle",
+]

autobyteus/workflow/utils/wait_for_idle.py

@@ -0,0 +1,46 @@
+# file: autobyteus/autobyteus/workflow/utils/wait_for_idle.py
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from autobyteus.workflow.streaming.workflow_event_stream import WorkflowEventStream
+from autobyteus.workflow.phases.workflow_operational_phase import WorkflowOperationalPhase
+
+if TYPE_CHECKING:
+    from autobyteus.workflow.agentic_workflow import AgenticWorkflow
+
+logger = logging.getLogger(__name__)
+
+async def _wait_loop(streamer: WorkflowEventStream, workflow_id: str):
+    """Internal helper to listen for the IDLE or ERROR event."""
+    async for event in streamer.all_events():
+        if event.event_source_type == "WORKFLOW" and event.data.new_phase == WorkflowOperationalPhase.IDLE:
+            logger.info(f"Workflow '{workflow_id}' has become idle.")
+            return
+        if event.event_source_type == "WORKFLOW" and event.data.new_phase == WorkflowOperationalPhase.ERROR:
+             error_message = f"Workflow '{workflow_id}' entered an error state while waiting for idle: {event.data.error_message}"
+             logger.error(error_message)
+             raise RuntimeError(error_message)
+
+async def wait_for_workflow_to_be_idle(workflow: 'AgenticWorkflow', timeout: float = 60.0):
+    """
+    Waits for a workflow to complete its bootstrapping and enter the IDLE state.
+
+    Args:
+        workflow: The workflow instance to monitor.
+        timeout: The maximum time in seconds to wait.
+
+    Raises:
+        asyncio.TimeoutError: If the workflow does not become idle within the timeout period.
+        RuntimeError: If the workflow enters an error state.
+    """
+    if workflow.get_current_phase() == WorkflowOperationalPhase.IDLE:
+        return
+    
+    logger.info(f"Waiting for workflow '{workflow.workflow_id}' to become idle (timeout: {timeout}s)...")
+    
+    streamer = WorkflowEventStream(workflow)
+    try:
+        await asyncio.wait_for(_wait_loop(streamer, workflow.workflow_id), timeout=timeout)
+    finally:
+        await streamer.close()