From 825ab176bb10218575ae7d5a2845d60db7f74058 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 6 Aug 2026 15:52:09 +0800 Subject: [PATCH] refactor(events)!: simplify event identity and buffering Replace append-only polling with condition-backed deque consumption. Add graph, execution, event, and node execution identities with snapshot-stable sequencing. Remove unused AgentNode event DTOs and the custom event codec. BREAKING CHANGE: graph events now expose graph_id and event id separately, node correlation uses node_execution_id, and AgentNode event APIs are removed. --- src/graphon/engine/engine.py | 20 +- src/graphon/engine/event/node_failure.py | 6 +- src/graphon/engine/event/processor.py | 2 - src/graphon/engine/event/stream.py | 223 ++++++------------ .../engine/filter/builtin/response_stream.py | 38 ++- src/graphon/engine/layer/base.py | 2 +- src/graphon/engine/worker/worker.py | 7 +- src/graphon/engine_events/__init__.py | 4 - src/graphon/engine_events/agent.py | 16 -- src/graphon/engine_events/base.py | 20 +- src/graphon/enums.py | 3 - src/graphon/node_events/__init__.py | 2 - src/graphon/node_events/agent.py | 18 -- src/graphon/nodes/base/node.py | 64 ++--- src/graphon/runtime/execution.py | 13 + .../test_cooperative_container_execution.py | 48 +++- tests/engine/test_dispatch_patterns.py | 32 +-- tests/engine/test_event_stream.py | 156 ++++++++++++ tests/engine/test_raw_engine_events.py | 6 +- tests/engine/test_response_stream_filter.py | 38 ++- .../engine/test_serializable_graph_runtime.py | 20 +- tests/engine_events/test_traversal_events.py | 12 +- tests/node_events/test_node_event_aliases.py | 4 +- tests/nodes/tool/test_tool_node.py | 4 +- tests/runtime/test_graph_runtime_state.py | 31 +++ tests/workflows/test_full_engine_events.py | 30 ++- 26 files changed, 519 insertions(+), 300 deletions(-) delete mode 100644 src/graphon/engine_events/agent.py delete mode 100644 src/graphon/node_events/agent.py create mode 100644 tests/engine/test_event_stream.py diff --git a/src/graphon/engine/engine.py b/src/graphon/engine/engine.py index 31299a9..c6f8b91 100644 --- a/src/graphon/engine/engine.py +++ b/src/graphon/engine/engine.py @@ -61,6 +61,7 @@ def __init__( command_channel: CommandChannel | None = None, workers: int = 5, container_handler_factories: Sequence[ContainerHandlerFactory] = (), + execution_id: str | None = None, ) -> None: """Build an engine for one graph execution. @@ -76,6 +77,7 @@ def __init__( A process-local in-memory channel is created when omitted. workers: Fixed number of worker threads to create while running. container_handler_factories: Additional container handler factories. + execution_id: Existing execution identity to preserve when supplied. Raises: ValueError: If ``workers`` is not a positive integer. @@ -95,6 +97,17 @@ def __init__( # Graph execution tracks the overall execution state self._graph_execution = self._graph_runtime_state.graph_execution + if execution_id is None and not self._graph_execution.started: + for selector in ( + ("sys", "workflow_execution_id"), + ("sys", "workflow_run_id"), + ): + segment = self._graph_runtime_state.variable_pool.get(selector) + if segment is not None and segment.text: + execution_id = segment.text + break + if execution_id: + self._graph_execution.execution_id = execution_id # Queue for state-transition work generated by workers. dispatch_queue: queue.Queue[DispatchTask] = queue.Queue() @@ -110,7 +123,12 @@ def __init__( self._scheduler = root_frame.scheduler # === Event Streaming === - self._event_stream = EventStream(self._layers) + self._event_stream = EventStream( + self._layers, + graph_id=self._graph_execution.workflow_id, + execution_id=self._graph_execution.execution_id, + next_sequence=self._graph_execution.next_event_sequence, + ) # === Command Processing === # Processes external commands (e.g., abort requests) diff --git a/src/graphon/engine/event/node_failure.py b/src/graphon/engine/event/node_failure.py index c4fab6f..eee0418 100644 --- a/src/graphon/engine/event/node_failure.py +++ b/src/graphon/engine/event/node_failure.py @@ -118,7 +118,7 @@ def _handle_retry( time.sleep(node.retry_config.retry_interval_seconds) return NodeRunRetryEvent( - id=event.id, + node_execution_id=event.node_execution_id, node_title=node.title, node_id=event.node_id, node_type=event.node_type, @@ -147,7 +147,7 @@ def _handle_fail_branch(self, event: NodeRunFailedEvent) -> NodeRunExceptionEven } return NodeRunExceptionEvent( - id=event.id, + node_execution_id=event.node_execution_id, node_id=event.node_id, node_type=event.node_type, start_at=event.start_at, @@ -189,7 +189,7 @@ def _handle_default_value(self, event: NodeRunFailedEvent) -> NodeRunExceptionEv } return NodeRunExceptionEvent( - id=event.id, + node_execution_id=event.node_execution_id, node_id=event.node_id, node_type=event.node_type, start_at=event.start_at, diff --git a/src/graphon/engine/event/processor.py b/src/graphon/engine/event/processor.py index c5f42f9..e1dea23 100644 --- a/src/graphon/engine/event/processor.py +++ b/src/graphon/engine/event/processor.py @@ -5,7 +5,6 @@ from functools import singledispatchmethod from typing import final -from graphon.engine_events.agent import NodeRunAgentLogEvent from graphon.engine_events.base import NodeEvent from graphon.engine_events.iteration import ( NodeRunIterationFailedEvent, @@ -177,7 +176,6 @@ def _( | NodeRunLoopNextEvent | NodeRunLoopSucceededEvent | NodeRunLoopFailedEvent - | NodeRunAgentLogEvent | NodeRunModelPollingProgressEvent | NodeRunRetrieverResourceEvent | NodeRunReasoningChunkEvent diff --git a/src/graphon/engine/event/stream.py b/src/graphon/engine/event/stream.py index 678953a..aaa0879 100644 --- a/src/graphon/engine/event/stream.py +++ b/src/graphon/engine/event/stream.py @@ -2,9 +2,9 @@ import logging import threading -import time -from collections.abc import Generator -from contextlib import contextmanager +from collections import deque +from collections.abc import Callable, Generator +from datetime import UTC, datetime from typing import final from graphon.engine_events.base import EngineEvent @@ -14,172 +14,81 @@ _logger = logging.getLogger(__name__) -@final -class ReadWriteLock: - """A read-write lock implementation that allows multiple concurrent readers - but only one writer at a time. - """ - - def __init__(self) -> None: - self._read_ready = threading.Condition(threading.RLock()) - self._readers = 0 - - def acquire_read(self) -> None: - """Acquire a read lock.""" - _ = self._read_ready.acquire() - try: - self._readers += 1 - finally: - self._read_ready.release() - - def release_read(self) -> None: - """Release a read lock.""" - _ = self._read_ready.acquire() - try: - self._readers -= 1 - if self._readers == 0: - self._read_ready.notify_all() - finally: - self._read_ready.release() - - def acquire_write(self) -> None: - """Acquire a write lock.""" - _ = self._read_ready.acquire() - while self._readers > 0: - _ = self._read_ready.wait() - - def release_write(self) -> None: - """Release a write lock.""" - self._read_ready.release() - - @contextmanager - def read_lock(self) -> Generator: - """Return a context manager for read locking.""" - self.acquire_read() - try: - yield - finally: - self.release_read() - - @contextmanager - def write_lock(self) -> Generator: - """Return a context manager for write locking.""" - self.acquire_write() - try: - yield - finally: - self.release_write() - - @final class EventStream: - """Collect, buffer, and stream engine events. - - The stream is the single event boundary between the engine and external - consumers. It also notifies the engine's layers as events arrive. - """ - - def __init__(self, layers: list[Layer]) -> None: - """Initialize an event stream bound to the engine's live layer list. - - The list is retained by reference so layers registered after engine - construction are visible to the stream without a second configuration - phase. Collected events are buffered until :meth:`emit_events` yields - them, while lifecycle events can notify the same layers without being - added to that buffer. - - Args: - layers: Mutable list of layers owned by the engine. - - """ - self._events: list[EngineEvent] = [] - self._lock = ReadWriteLock() + """Collect, buffer, and stream engine events.""" + + def __init__( + self, + layers: list[Layer], + graph_id: str = "", + execution_id: str = "", + next_sequence: Callable[[], int] | None = None, + ) -> None: + self._graph_id = graph_id + self._execution_id = execution_id + self._events: deque[EngineEvent] = deque() + self._condition = threading.Condition() self._layers = layers - self._execution_complete = threading.Event() + self._execution_complete = False + self._next_sequence = next_sequence + self._local_sequence = 0 def notify_layers(self, event: EngineEvent) -> None: - """Notify all layers about an event without buffering it. - - Layer exceptions are caught and logged so one extension cannot disrupt - event delivery to the remaining layers or the engine itself. - - Args: - event: Event to send to every registered layer. - - """ - for layer in self._layers: - try: - layer.on_event(event) - except Exception: - _logger.exception("Error in layer on_event, layer_type=%s", type(layer)) + """Stamp an unbuffered lifecycle event and notify registered layers.""" + with self._condition: + self._stamp(event) + self._notify_layers(event) def collect(self, event: EngineEvent) -> None: - """Thread-safe method to collect an event. - - Args: - event: The event to collect - - """ - with self._lock.write_lock(): + """Buffer one event and wake its consumer.""" + with self._condition: + if self._execution_complete: + msg = "Cannot collect events after execution is complete" + raise RuntimeError(msg) + self._stamp(event) self._events.append(event) - - # NOTE: `notify_layers` is intentionally called outside the critical section - # to minimize lock contention and avoid blocking other readers or writers. - self.notify_layers(event) - - def _get_new_events(self, start_index: int) -> list[EngineEvent]: - """Get new events starting from a specific index. - - Args: - start_index: The index to start from - - Returns: - List of new events - - """ - with self._lock.read_lock(): - return list(self._events[start_index:]) - - def _event_count(self) -> int: - """Get the current count of collected events. - - Returns: - Number of collected events - - """ - with self._lock.read_lock(): - return len(self._events) + # Layers observe stream order before the consumer can wake. + self._notify_layers(event) + self._condition.notify() def mark_complete(self) -> None: - """Mark execution as complete to stop the event emission generator.""" - self._execution_complete.set() + """Mark execution complete and wake all waiting consumers.""" + with self._condition: + self._execution_complete = True + self._condition.notify_all() def reset(self) -> None: - """Discard events and completion state from the previous engine run.""" - with self._lock.write_lock(): + """Discard buffered events and completion state from the previous run.""" + with self._condition: self._events.clear() - self._execution_complete.clear() + self._execution_complete = False def emit_events(self) -> Generator[EngineEvent, None, None]: - """Generator that yields events as they're collected. - - Yields: - EngineEvent instances as they're processed - - """ - yielded_count = 0 - - while ( - not self._execution_complete.is_set() or yielded_count < self._event_count() - ): - # Get new events since last yield - new_events = self._get_new_events(yielded_count) - - # Yield any new events - for event in new_events: - yield event - yielded_count += 1 - - # Small sleep to avoid busy waiting - if not self._execution_complete.is_set() and not new_events: - time.sleep(0.001) + """Yield events in collection order, releasing each after consumption.""" + while True: + with self._condition: + self._condition.wait_for( + lambda: self._events or self._execution_complete + ) + if not self._events: + return + event = self._events.popleft() + yield event # ruff:ignore[unnecessary-assign-before-yield] + + def _stamp(self, event: EngineEvent) -> None: + event.graph_id = self._graph_id + event.execution_id = self._execution_id + if self._next_sequence is None: + self._local_sequence += 1 + event.sequence = self._local_sequence + else: + event.sequence = self._next_sequence() + event.emitted_at = datetime.now(UTC) + + def _notify_layers(self, event: EngineEvent) -> None: + for layer in self._layers: + try: + layer.on_event(event) + except Exception: + _logger.exception("Error in layer on_event, layer_type=%s", type(layer)) diff --git a/src/graphon/engine/filter/builtin/response_stream.py b/src/graphon/engine/filter/builtin/response_stream.py index 561e52d..a67560d 100644 --- a/src/graphon/engine/filter/builtin/response_stream.py +++ b/src/graphon/engine/filter/builtin/response_stream.py @@ -195,6 +195,8 @@ class _ResponseStreamFilterState(BaseModel): type: Literal["ResponseStreamFilter"] = Field(default="ResponseStreamFilter") version: str = Field(default="1.0") response_nodes: Sequence[str] = Field(default_factory=list) + graph_id: str = "" + execution_id: str = "" active_session: _ResponseSessionState | None = None waiting_sessions: Sequence[_ResponseSessionState] = Field(default_factory=list) pending_sessions: Sequence[_ResponseSessionState] = Field(default_factory=list) @@ -216,6 +218,8 @@ def __init__(self, *, pass_unmatched_chunks: bool = False) -> None: self._reset_run_state() def _reset_run_state(self) -> None: + self._graph_id = "" + self._execution_id = "" self._active_session: _ResponseSession | None = None self._waiting_sessions: deque[_ResponseSession] = deque() self._stream_buffers = _StreamBuffers() @@ -250,12 +254,14 @@ def on_event(self, event: EngineEvent) -> Iterable[EngineEvent]: self._ensure_initialized() match event: case GraphRunStartedEvent(): + self._graph_id = event.graph_id + self._execution_id = event.execution_id output: Iterable[EngineEvent] = [ event, *self._activate_initial_sessions(), ] case NodeRunStartedEvent(): - self._node_execution_ids[event.node_id] = event.id + self._node_execution_ids[event.node_id] = event.node_execution_id output = [event] case NodeRunStreamChunkEvent(): output = self._handle_stream_chunk(event) @@ -282,6 +288,8 @@ def dumps(self) -> str: state = _ResponseStreamFilterState( response_nodes=sorted(self._response_nodes), + graph_id=self._graph_id, + execution_id=self._execution_id, active_session=self._serialize_session(self._active_session), waiting_sessions=[ session_state @@ -360,6 +368,8 @@ def _apply_state(self, state: _ResponseStreamFilterState) -> None: ) self._active_session = active_session + self._graph_id = state.graph_id + self._execution_id = state.execution_id self._waiting_sessions = waiting_sessions self._stream_buffers = stream_buffers self._response_nodes = response_nodes @@ -608,7 +618,7 @@ def _session_references_reasoning_source( for segment in session.template.segments ) - def _get_or_create_execution_id(self, node_id: _NodeID) -> str: + def _get_or_create_node_execution_id(self, node_id: _NodeID) -> str: if node_id not in self._node_execution_ids: self._node_execution_ids[node_id] = str(uuid4()) return self._node_execution_ids[node_id] @@ -616,7 +626,7 @@ def _get_or_create_execution_id(self, node_id: _NodeID) -> str: def _create_stream_chunk_event( self, node_id: _NodeID, - execution_id: str, + node_execution_id: str, selector: Sequence[str], chunk: str, is_final: bool = False, @@ -625,7 +635,9 @@ def _create_stream_chunk_event( if selector and selector[0] not in graph.nodes and self._active_session: response_node = graph.nodes[self._active_session.node_id] return NodeRunStreamChunkEvent( - id=execution_id, + graph_id=self._graph_id, + execution_id=self._execution_id, + node_execution_id=node_execution_id, node_id=response_node.id, node_type=response_node.node_type, selector=list(selector), @@ -635,7 +647,9 @@ def _create_stream_chunk_event( node = graph.nodes[node_id] return NodeRunStreamChunkEvent( - id=execution_id, + graph_id=self._graph_id, + execution_id=self._execution_id, + node_execution_id=node_execution_id, node_id=node.id, node_type=node.node_type, selector=list(selector), @@ -656,7 +670,7 @@ def _process_variable_segment( output_node_id = self._active_session.node_id else: output_node_id = source_selector_prefix - execution_id = self._get_or_create_execution_id(output_node_id) + node_execution_id = self._get_or_create_node_execution_id(output_node_id) has_stream_events = self._stream_buffers.has_events(segment.selector) while self._stream_buffers.has_unread(segment.selector): @@ -668,7 +682,9 @@ def _process_variable_segment( response_node = self._bound_graph.nodes[self._active_session.node_id] events.append( NodeRunStreamChunkEvent( - id=execution_id, + graph_id=self._graph_id, + execution_id=self._execution_id, + node_execution_id=node_execution_id, node_id=response_node.id, node_type=response_node.node_type, container_id=event.container_id, @@ -693,7 +709,7 @@ def _process_variable_segment( events.append( self._create_stream_chunk_event( node_id=output_node_id, - execution_id=execution_id, + node_execution_id=node_execution_id, selector=segment.selector, chunk=value.markdown, is_final=is_last_segment, @@ -713,14 +729,16 @@ def _process_text_segment( raise RuntimeError(msg) current_response_node = self._bound_graph.nodes[active_session.node_id] - execution_id = self._get_or_create_execution_id(current_response_node.id) + node_execution_id = self._get_or_create_node_execution_id( + current_response_node.id + ) is_last_segment = ( active_session.index == len(active_session.template.segments) - 1 ) return [ self._create_stream_chunk_event( node_id=current_response_node.id, - execution_id=execution_id, + node_execution_id=node_execution_id, selector=self._get_text_segment_selector(current_response_node.id), chunk=segment.text, is_final=is_last_segment, diff --git a/src/graphon/engine/layer/base.py b/src/graphon/engine/layer/base.py index 003d654..a241c07 100644 --- a/src/graphon/engine/layer/base.py +++ b/src/graphon/engine/layer/base.py @@ -114,7 +114,7 @@ def on_node_run_end( """Called after a node finishes execution. The node's execution ID is available via `node._node_execution_id` and matches - the `id` field in all events emitted by this node execution. + the `node_execution_id` field in all events emitted by this node execution. Args: node: The node instance that just finished execution diff --git a/src/graphon/engine/worker/worker.py b/src/graphon/engine/worker/worker.py index cd4b474..096e612 100644 --- a/src/graphon/engine/worker/worker.py +++ b/src/graphon/engine/worker/worker.py @@ -272,7 +272,10 @@ def _consume_node_events( ) ) return None, True - if isinstance(event, NodeRunStartedEvent) and event.id == node.execution_id: + if ( + isinstance(event, NodeRunStartedEvent) + and event.node_execution_id == node.execution_id + ): self._current_node_started_at = event.start_at self._dispatch_queue.put( NodeEventTask(frame_id=self._current_frame_id, event=event) @@ -323,7 +326,7 @@ def _build_fallback_failure_event( failure_time = datetime.now(UTC).replace(tzinfo=None) error_message = str(error) return NodeRunFailedEvent( - id=node.execution_id, + node_execution_id=node.execution_id, node_id=node.id, node_type=node.node_type, error=error_message, diff --git a/src/graphon/engine_events/__init__.py b/src/graphon/engine_events/__init__.py index 17283f3..4755af4 100644 --- a/src/graphon/engine_events/__init__.py +++ b/src/graphon/engine_events/__init__.py @@ -1,6 +1,3 @@ -# Agent events -from .agent import NodeRunAgentLogEvent - # Base events from .base import EngineEvent, NodeEvent @@ -63,7 +60,6 @@ "GraphRunStartedEvent", "GraphRunSucceededEvent", "NodeEvent", - "NodeRunAgentLogEvent", "NodeRunExceptionEvent", "NodeRunFailedEvent", "NodeRunHumanInputFormFilledEvent", diff --git a/src/graphon/engine_events/agent.py b/src/graphon/engine_events/agent.py deleted file mode 100644 index 9e3dce5..0000000 --- a/src/graphon/engine_events/agent.py +++ /dev/null @@ -1,16 +0,0 @@ -from collections.abc import Mapping - -from pydantic import Field - -from .base import NodeEvent - - -class NodeRunAgentLogEvent(NodeEvent): - message_id: str = Field(..., description="message id") - label: str = Field(..., description="label") - node_execution_id: str = Field(..., description="node execution id") - parent_id: str | None = Field(..., description="parent id") - error: str | None = Field(..., description="error") - status: str = Field(..., description="status") - data: Mapping[str, object] = Field(..., description="data") - metadata: Mapping[str, object] = Field(default_factory=dict) diff --git a/src/graphon/engine_events/base.py b/src/graphon/engine_events/base.py index a8fdb0f..604009d 100644 --- a/src/graphon/engine_events/base.py +++ b/src/graphon/engine_events/base.py @@ -1,4 +1,8 @@ -from pydantic import BaseModel, Field +from datetime import UTC, datetime +from typing import Literal +from uuid import uuid4 + +from pydantic import BaseModel, Field, computed_field from graphon.enums import NodeType from graphon.node_events.base import NodeRunResult @@ -7,11 +11,23 @@ class EngineEvent(BaseModel): """Base model for events emitted by the engine.""" + id: str = Field(default_factory=lambda: str(uuid4())) + graph_id: str = "" + execution_id: str = "" + schema_version: Literal["1.0"] = "1.0" + sequence: int = Field(default=0, ge=0) + emitted_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) + + @computed_field + @property + def event_type(self) -> str: + return type(self).__name__ + class NodeEvent(EngineEvent): """Engine event associated with one node execution.""" - id: str = Field(..., description="node execution id") + node_execution_id: str = Field(..., description="node execution id") node_id: str node_type: NodeType container_id: str = "" diff --git a/src/graphon/enums.py b/src/graphon/enums.py index 1c73c22..300e41c 100644 --- a/src/graphon/enums.py +++ b/src/graphon/enums.py @@ -44,7 +44,6 @@ class BuiltinNodeTypes: VARIABLE_ASSIGNER: ClassVar[NodeType] = "assigner" DOCUMENT_EXTRACTOR: ClassVar[NodeType] = "document-extractor" LIST_OPERATOR: ClassVar[NodeType] = "list-operator" - AGENT: ClassVar[NodeType] = "agent" HUMAN_INPUT: ClassVar[NodeType] = "human-input" @@ -72,7 +71,6 @@ class BuiltinNodeTypes: BuiltinNodeTypes.VARIABLE_ASSIGNER, BuiltinNodeTypes.DOCUMENT_EXTRACTOR, BuiltinNodeTypes.LIST_OPERATOR, - BuiltinNodeTypes.AGENT, BuiltinNodeTypes.HUMAN_INPUT, ) @@ -211,7 +209,6 @@ class WorkflowNodeExecutionMetadataKey(StrEnum): TOTAL_PRICE = "total_price" CURRENCY = "currency" TOOL_INFO = "tool_info" - AGENT_LOG = "agent_log" ITERATION_ID = "iteration_id" ITERATION_INDEX = "iteration_index" LOOP_ID = "loop_id" diff --git a/src/graphon/node_events/__init__.py b/src/graphon/node_events/__init__.py index fe896a1..43ce0fd 100644 --- a/src/graphon/node_events/__init__.py +++ b/src/graphon/node_events/__init__.py @@ -1,4 +1,3 @@ -from .agent import AgentLogEvent from .base import NodeEventPayload, NodeRunResult from .iteration import ( IterationFailedEvent, @@ -27,7 +26,6 @@ ) __all__ = [ - "AgentLogEvent", "HumanInputFormFilledEvent", "HumanInputFormTimeoutEvent", "IterationFailedEvent", diff --git a/src/graphon/node_events/agent.py b/src/graphon/node_events/agent.py deleted file mode 100644 index 02d5e0a..0000000 --- a/src/graphon/node_events/agent.py +++ /dev/null @@ -1,18 +0,0 @@ -from collections.abc import Mapping -from typing import Any - -from pydantic import Field - -from .base import NodeEventPayload - - -class AgentLogEvent(NodeEventPayload): - message_id: str = Field(..., description="id") - label: str = Field(..., description="label") - node_execution_id: str = Field(..., description="node execution id") - parent_id: str | None = Field(..., description="parent id") - error: str | None = Field(..., description="error") - status: str = Field(..., description="status") - data: Mapping[str, Any] = Field(..., description="data") - metadata: Mapping[str, Any] = Field(default_factory=dict, description="metadata") - node_id: str = Field(..., description="node id") diff --git a/src/graphon/nodes/base/node.py b/src/graphon/nodes/base/node.py index dc70dea..9aa8f21 100644 --- a/src/graphon/nodes/base/node.py +++ b/src/graphon/nodes/base/node.py @@ -9,7 +9,6 @@ from types import MappingProxyType from typing import Any, ClassVar, assert_never, get_args, get_origin -from graphon.engine_events.agent import NodeRunAgentLogEvent from graphon.engine_events.base import NodeEvent from graphon.engine_events.iteration import ( NodeRunIterationFailedEvent, @@ -46,7 +45,6 @@ NodeType, WorkflowNodeExecutionStatus, ) -from graphon.node_events.agent import AgentLogEvent from graphon.node_events.base import ( NodeEventPayload, NodeRunResult, @@ -661,7 +659,7 @@ def run( # Create and push start event with required fields start_event = NodeRunStartedEvent( - id=execution_id, + node_execution_id=execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.title, @@ -734,7 +732,7 @@ def _normalize_run_event( if isinstance(event, NodeEventPayload): return self._dispatch(event) if not event.container_id: - event.id = self.execution_id + event.node_execution_id = self.execution_id return event def _build_run_failed_event(self, error: Exception) -> NodeRunFailedEvent: @@ -745,7 +743,7 @@ def _build_run_failed_event(self, error: Exception) -> NodeRunFailedEvent: ) finished_at = datetime.now(UTC).replace(tzinfo=None) return NodeRunFailedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, start_at=self._start_at, @@ -842,7 +840,7 @@ def _convert_node_run_result_to_node_event( match status: case WorkflowNodeExecutionStatus.FAILED: return NodeRunFailedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, start_at=self._start_at, @@ -852,7 +850,7 @@ def _convert_node_run_result_to_node_event( ) case WorkflowNodeExecutionStatus.SUCCEEDED: return NodeRunSucceededEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, start_at=self._start_at, @@ -880,7 +878,7 @@ def _dispatch(self, event: NodeEventPayload) -> NodeEvent: @_dispatch.register def _(self, event: StreamChunkEvent) -> NodeRunStreamChunkEvent: return NodeRunStreamChunkEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, selector=event.selector, @@ -891,7 +889,7 @@ def _(self, event: StreamChunkEvent) -> NodeRunStreamChunkEvent: @_dispatch.register def _(self, event: StreamReasoningEvent) -> NodeRunReasoningChunkEvent: return NodeRunReasoningChunkEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, selector=[self._node_id, "reasoning_content"], @@ -902,7 +900,7 @@ def _(self, event: StreamReasoningEvent) -> NodeRunReasoningChunkEvent: @_dispatch.register def _(self, event: ModelPollingProgressEvent) -> NodeRunModelPollingProgressEvent: return NodeRunModelPollingProgressEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, attempt=event.attempt, @@ -920,7 +918,7 @@ def _( match status: case WorkflowNodeExecutionStatus.SUCCEEDED: return NodeRunSucceededEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, start_at=self._start_at, @@ -929,7 +927,7 @@ def _( ) case WorkflowNodeExecutionStatus.FAILED: return NodeRunFailedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, start_at=self._start_at, @@ -953,7 +951,7 @@ def _( @_dispatch.register def _(self, event: VariableUpdatedEvent) -> NodeRunVariableUpdatedEvent: return NodeRunVariableUpdatedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, variable=event.variable, @@ -962,33 +960,17 @@ def _(self, event: VariableUpdatedEvent) -> NodeRunVariableUpdatedEvent: @_dispatch.register def _(self, event: PauseRequestedEvent) -> NodeRunPauseRequestedEvent: return NodeRunPauseRequestedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.PAUSED), reason=event.reason, ) - @_dispatch.register - def _(self, event: AgentLogEvent) -> NodeRunAgentLogEvent: - return NodeRunAgentLogEvent( - id=self.execution_id, - node_id=self._node_id, - node_type=self.node_type, - message_id=event.message_id, - label=event.label, - node_execution_id=event.node_execution_id, - parent_id=event.parent_id, - error=event.error, - status=event.status, - data=event.data, - metadata=event.metadata, - ) - @_dispatch.register def _(self, event: HumanInputFormFilledEvent) -> NodeRunHumanInputFormFilledEvent: return NodeRunHumanInputFormFilledEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=event.node_title, @@ -1001,7 +983,7 @@ def _(self, event: HumanInputFormFilledEvent) -> NodeRunHumanInputFormFilledEven @_dispatch.register def _(self, event: HumanInputFormTimeoutEvent) -> NodeRunHumanInputFormTimeoutEvent: return NodeRunHumanInputFormTimeoutEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=event.node_title, @@ -1011,7 +993,7 @@ def _(self, event: HumanInputFormTimeoutEvent) -> NodeRunHumanInputFormTimeoutEv @_dispatch.register def _(self, event: LoopStartedEvent) -> NodeRunLoopStartedEvent: return NodeRunLoopStartedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1024,7 +1006,7 @@ def _(self, event: LoopStartedEvent) -> NodeRunLoopStartedEvent: @_dispatch.register def _(self, event: LoopNextEvent) -> NodeRunLoopNextEvent: return NodeRunLoopNextEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1035,7 +1017,7 @@ def _(self, event: LoopNextEvent) -> NodeRunLoopNextEvent: @_dispatch.register def _(self, event: LoopSucceededEvent) -> NodeRunLoopSucceededEvent: return NodeRunLoopSucceededEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1049,7 +1031,7 @@ def _(self, event: LoopSucceededEvent) -> NodeRunLoopSucceededEvent: @_dispatch.register def _(self, event: LoopFailedEvent) -> NodeRunLoopFailedEvent: return NodeRunLoopFailedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1064,7 +1046,7 @@ def _(self, event: LoopFailedEvent) -> NodeRunLoopFailedEvent: @_dispatch.register def _(self, event: IterationStartedEvent) -> NodeRunIterationStartedEvent: return NodeRunIterationStartedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1077,7 +1059,7 @@ def _(self, event: IterationStartedEvent) -> NodeRunIterationStartedEvent: @_dispatch.register def _(self, event: IterationNextEvent) -> NodeRunIterationNextEvent: return NodeRunIterationNextEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1088,7 +1070,7 @@ def _(self, event: IterationNextEvent) -> NodeRunIterationNextEvent: @_dispatch.register def _(self, event: IterationSucceededEvent) -> NodeRunIterationSucceededEvent: return NodeRunIterationSucceededEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1102,7 +1084,7 @@ def _(self, event: IterationSucceededEvent) -> NodeRunIterationSucceededEvent: @_dispatch.register def _(self, event: IterationFailedEvent) -> NodeRunIterationFailedEvent: return NodeRunIterationFailedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, node_title=self.node_data.title, @@ -1117,7 +1099,7 @@ def _(self, event: IterationFailedEvent) -> NodeRunIterationFailedEvent: @_dispatch.register def _(self, event: RunRetrieverResourceEvent) -> NodeRunRetrieverResourceEvent: return NodeRunRetrieverResourceEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self._node_id, node_type=self.node_type, retriever_resources=event.retriever_resources, diff --git a/src/graphon/runtime/execution.py b/src/graphon/runtime/execution.py index cc822cc..419f681 100644 --- a/src/graphon/runtime/execution.py +++ b/src/graphon/runtime/execution.py @@ -70,6 +70,8 @@ class GraphExecutionState(BaseModel): version: Literal["2.0"] workflow_id: str + execution_id: str = Field(default_factory=lambda: str(uuid4())) + last_event_sequence: int = Field(default=0, ge=0) started: bool completed: bool aborted: bool @@ -97,6 +99,8 @@ class GraphExecution: """ workflow_id: str + execution_id: str = field(default_factory=lambda: str(uuid4())) + last_event_sequence: int = 0 started: bool = False completed: bool = False aborted: bool = False @@ -157,6 +161,11 @@ def get_or_create_node_execution( ) return self.node_executions[key] + def next_event_sequence(self) -> int: + """Return the next sequence number for this graph execution.""" + self.last_event_sequence += 1 + return self.last_event_sequence + def dumps(self) -> str: """Serialize the aggregate state into a JSON string.""" node_states = [ @@ -172,6 +181,8 @@ def dumps(self) -> str: state = GraphExecutionState( version="2.0", workflow_id=self.workflow_id, + execution_id=self.execution_id, + last_event_sequence=self.last_event_sequence, started=self.started, completed=self.completed, aborted=self.aborted, @@ -235,6 +246,8 @@ def from_snapshot(cls, data: str) -> GraphExecution: return cls( workflow_id=state.workflow_id, + execution_id=state.execution_id, + last_event_sequence=state.last_event_sequence, started=state.started, completed=state.completed, aborted=state.aborted, diff --git a/tests/engine/test_cooperative_container_execution.py b/tests/engine/test_cooperative_container_execution.py index e54bdf0..52f3976 100644 --- a/tests/engine/test_cooperative_container_execution.py +++ b/tests/engine/test_cooperative_container_execution.py @@ -4,6 +4,7 @@ from threading import Event, Lock, Thread from types import SimpleNamespace from typing import cast +from unittest.mock import MagicMock import pytest @@ -24,12 +25,15 @@ ) from graphon.engine_events.base import EngineEvent, NodeEvent from graphon.engine_events.node import ( + NodeRunExceptionEvent, NodeRunFailedEvent, + NodeRunRetryEvent, NodeRunStartedEvent, NodeRunSucceededEvent, ) from graphon.enums import ( BuiltinNodeTypes, + ErrorStrategy, NodeExecutionType, WorkflowNodeExecutionStatus, ) @@ -102,6 +106,44 @@ def on_node_run_end( self.end_events.append(result_event) +def test_error_handler_preserves_node_execution_but_not_event_id() -> None: + node = SimpleNamespace( + retry=True, + retry_config=SimpleNamespace(max_retries=1, retry_interval_seconds=0), + error_strategy=None, + title="Code", + ) + graph_execution = MagicMock() + graph_execution.get_or_create_node_execution.return_value.retry_count = 0 + handler = NodeFailureHandler( + cast(Graph, SimpleNamespace(nodes={"node": node})), + graph_execution, + ) + failed = NodeRunFailedEvent( + id="failed-event", + node_execution_id="node-run", + node_id="node", + node_type=BuiltinNodeTypes.CODE, + error="failed", + start_at=datetime.now(UTC).replace(tzinfo=None), + node_run_result=NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="failed", + ), + ) + + retry = handler.handle(frame_id="root", event=failed) + assert isinstance(retry, NodeRunRetryEvent) + + node.retry = False + node.error_strategy = ErrorStrategy.FAIL_BRANCH + exception = handler.handle(frame_id="root", event=failed) + assert isinstance(exception, NodeRunExceptionEvent) + + assert retry.node_execution_id == exception.node_execution_id == "node-run" + assert len({failed.id, retry.id, exception.id}) == 3 + + def test_ready_queue_round_trips_start_and_resume_tasks() -> None: queue_ = InMemoryReadyQueue() result = _container_result() @@ -212,7 +254,7 @@ def run( ) -> Generator[NodeEvent | LoopFrameRequest, object, None]: started_at = datetime.now(UTC).replace(tzinfo=None) yield NodeRunStartedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, node_title="Loop", @@ -249,7 +291,7 @@ def resume_container( }, ) yield NodeRunSucceededEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, start_at=started_at, @@ -309,7 +351,7 @@ def resume_container( frame_id=run_state.frame_id, node_id=run_state.node_id, ) - assert node_execution.execution_id == started.event.id + assert node_execution.execution_id == started.event.node_execution_id assert run_state.started_at == started.event.start_at assert layer.end_events == [] diff --git a/tests/engine/test_dispatch_patterns.py b/tests/engine/test_dispatch_patterns.py index 79c93fb..121f671 100644 --- a/tests/engine/test_dispatch_patterns.py +++ b/tests/engine/test_dispatch_patterns.py @@ -635,7 +635,7 @@ def run(self) -> Generator[NodeRunSucceededEvent, None, None]: raise TimeoutError(msg) now = datetime.now(UTC).replace(tzinfo=None) yield NodeRunSucceededEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, start_at=now, @@ -712,7 +712,7 @@ def run(self) -> Generator[NodeRunSucceededEvent, None, None]: barrier.wait(timeout=1) now = datetime.now(UTC).replace(tzinfo=None) yield NodeRunSucceededEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, start_at=now, @@ -889,7 +889,7 @@ def test_pause_requested_event_defers_current_task_for_resume() -> None: NodeEventTask( frame_id="child-frame", event=NodeRunPauseRequestedEvent( - id="human-run", + node_execution_id="human-run", node_id="human", node_type=BuiltinNodeTypes.HUMAN_INPUT, reason=HitlRequired( @@ -1223,7 +1223,7 @@ def bind_execution_id(self, execution_id: str) -> None: def run(self) -> Generator[NodeRunStartedEvent, None, None]: yield NodeRunStartedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, node_title="Start", @@ -1283,7 +1283,7 @@ def bind_execution_id(self, execution_id: str) -> None: def run(self) -> Generator[NodeRunStartedEvent, None, None]: yield NodeRunStartedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, node_title=self.title, @@ -1351,7 +1351,7 @@ def bind_execution_id(self, execution_id: str) -> None: def run(self) -> Generator[NodeRunStartedEvent, None, None]: yield NodeRunStartedEvent( - id=self.execution_id, + node_execution_id=self.execution_id, node_id=self.id, node_type=self.node_type, node_title="Answer", @@ -1401,13 +1401,13 @@ def run(self) -> Generator[NodeRunStartedEvent, None, None]: assert isinstance(event, NodeEventTask) assert isinstance(event.event, NodeRunStartedEvent) - assert event.event.id != root_execution.execution_id - assert event.event.id == child_execution.execution_id + assert event.event.node_execution_id != root_execution.execution_id + assert event.event.node_execution_id == child_execution.execution_id def test_dispatcher_preserves_task_event_for_dispatch() -> None: event = NodeRunStartedEvent( - id="run-1", + node_execution_id="run-1", node_id="start", node_type=BuiltinNodeTypes.CODE, node_title="Start", @@ -1454,7 +1454,7 @@ def dispatch(self, event: object) -> None: def test_event_processor_dispatches_task_event_payload() -> None: event = NodeRunStartedEvent( - id="run-1", + node_execution_id="run-1", node_id="start", node_type=BuiltinNodeTypes.CODE, node_title="Start", @@ -1521,7 +1521,7 @@ def test_event_processor_stamps_frame_owner_on_node_and_edge_events() -> None: frame_registry=frame_registry, ) event = NodeRunSucceededEvent( - id="run-child", + node_execution_id="run-child", node_id="child", node_type=BuiltinNodeTypes.CODE, start_at=datetime.now(UTC).replace(tzinfo=None), @@ -1629,7 +1629,7 @@ def test_parallel_iteration_preserves_aggregate_and_response_order() -> None: # NodeEventTask( frame_id="iteration-invocation:iteration:1", event=NodeRunSucceededEvent( - id="iteration-start-run-1", + node_execution_id="iteration-start-run-1", node_id="iteration-start", node_type=BuiltinNodeTypes.ITERATION_START, start_at=datetime.now(UTC).replace(tzinfo=None), @@ -1658,7 +1658,7 @@ def test_parallel_iteration_preserves_aggregate_and_response_order() -> None: # NodeEventTask( frame_id="iteration-invocation:iteration:2", event=NodeRunSucceededEvent( - id="iteration-start-run-2", + node_execution_id="iteration-start-run-2", node_id="iteration-start", node_type=BuiltinNodeTypes.ITERATION_START, start_at=datetime.now(UTC).replace(tzinfo=None), @@ -1675,7 +1675,7 @@ def test_parallel_iteration_preserves_aggregate_and_response_order() -> None: # NodeEventTask( frame_id="iteration-invocation:iteration:0", event=NodeRunSucceededEvent( - id="iteration-start-run-0", + node_execution_id="iteration-start-run-0", node_id="iteration-start", node_type=BuiltinNodeTypes.ITERATION_START, start_at=datetime.now(UTC).replace(tzinfo=None), @@ -1861,7 +1861,7 @@ def test_execution_limits_layer_sends_abort_when_limit_is_exceeded( layer.on_event( NodeRunSucceededEvent( - id="node-run-1", + node_execution_id="node-run-1", node_id="node-1", node_type=BuiltinNodeTypes.CODE, start_at=datetime.now(UTC).replace(tzinfo=None), @@ -1883,7 +1883,7 @@ class CustomNodeRunStartedEvent(NodeRunStartedEvent): layer.on_event( CustomNodeRunStartedEvent( - id="node-run-1", + node_execution_id="node-run-1", node_id="node-1", node_type=BuiltinNodeTypes.CODE, node_title="Code", diff --git a/tests/engine/test_event_stream.py b/tests/engine/test_event_stream.py new file mode 100644 index 0000000..0d30a13 --- /dev/null +++ b/tests/engine/test_event_stream.py @@ -0,0 +1,156 @@ +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC +from threading import Barrier, Event, Thread +from unittest.mock import MagicMock + +import pytest + +from graphon.engine.event.stream import EventStream +from graphon.engine.layer import Layer +from graphon.engine_events.base import EngineEvent +from graphon.engine_events.graph import GraphRunStartedEvent + + +def test_emit_events_waits_for_completion_and_drains_pending_events() -> None: + stream = EventStream([], graph_id="graph", execution_id="execution") + event = GraphRunStartedEvent() + emitted: list[EngineEvent] = [] + completed = Event() + + def consume() -> None: + emitted.extend(stream.emit_events()) + completed.set() + + consumer = Thread(target=consume) + consumer.start() + stream.collect(event) + stream.mark_complete() + + assert completed.wait(timeout=1) + consumer.join() + assert emitted == [event] + + +def test_concurrent_collection_assigns_sequences_in_emission_order() -> None: + layer = MagicMock(spec=Layer) + stream = EventStream([layer], graph_id="graph", execution_id="execution") + events = [GraphRunStartedEvent() for _ in range(8)] + barrier = Barrier(len(events)) + + def collect(event: EngineEvent) -> None: + barrier.wait() + stream.collect(event) + + with ThreadPoolExecutor(max_workers=len(events)) as executor: + list(executor.map(collect, events)) + + stream.mark_complete() + emitted = list(stream.emit_events()) + + assert [event.sequence for event in emitted] == list(range(1, len(events) + 1)) + assert len({event.id for event in emitted}) == len(events) + assert all(event.graph_id == "graph" for event in emitted) + assert all(event.execution_id == "execution" for event in emitted) + assert all(event.emitted_at.tzinfo is UTC for event in emitted) + notified = [call.args[0] for call in layer.on_event.call_args_list] + assert all( + layer_event is emitted_event + for layer_event, emitted_event in zip(notified, emitted, strict=True) + ) + + +def test_layer_observes_event_before_consumer() -> None: + consumer_received = Event() + emitted: list[EngineEvent] = [] + layer = MagicMock(spec=Layer) + stream = EventStream([layer], graph_id="graph", execution_id="execution") + + def consume() -> None: + emitted.append(next(stream.emit_events())) + consumer_received.set() + + consumer = Thread(target=consume, daemon=True) + observed_during_callback: list[bool] = [] + + def on_event(_: EngineEvent) -> None: + consumer.start() + observed_during_callback.append(consumer_received.wait(timeout=0.1)) + + layer.on_event.side_effect = on_event + event = GraphRunStartedEvent() + + stream.collect(event) + consumer.join(timeout=1) + + assert observed_during_callback == [False] + assert emitted == [event] + + +def test_consumed_event_is_removed_from_stream_buffer() -> None: + stream = EventStream([], graph_id="graph", execution_id="execution") + event = GraphRunStartedEvent() + stream.collect(event) + emitted = stream.emit_events() + + assert next(emitted) is event + assert not stream._events + + stream.mark_complete() + assert list(emitted) == [] + + +def test_collect_buffers_before_synchronously_notifying_layers() -> None: + failing_layer = MagicMock(spec=Layer) + failing_layer.on_event.side_effect = RuntimeError("layer failed") + recording_layer = MagicMock(spec=Layer) + stream = EventStream( + [failing_layer, recording_layer], + graph_id="graph", + execution_id="execution", + ) + buffered_when_notified: list[tuple[EngineEvent, ...]] = [] + recording_layer.on_event.side_effect = lambda _: buffered_when_notified.append( + tuple(stream._events) + ) + event = GraphRunStartedEvent() + + stream.collect(event) + + failing_layer.on_event.assert_called_once_with(event) + recording_layer.on_event.assert_called_once_with(event) + assert buffered_when_notified == [(event,)] + + +def test_reset_discards_pending_events_without_resetting_sequence() -> None: + stream = EventStream([], graph_id="graph", execution_id="execution") + lifecycle_event = GraphRunStartedEvent() + stale_event = GraphRunStartedEvent() + stream.notify_layers(lifecycle_event) + stream.collect(stale_event) + stream.mark_complete() + + stream.reset() + current_event = GraphRunStartedEvent() + stream.collect(current_event) + stream.mark_complete() + + assert list(stream.emit_events()) == [current_event] + assert [lifecycle_event.sequence, stale_event.sequence, current_event.sequence] == [ + 1, + 2, + 3, + ] + + +def test_completion_rejects_late_collection_but_allows_terminal_notification() -> None: + stream = EventStream([], graph_id="graph", execution_id="execution") + stream.mark_complete() + terminal_event = GraphRunStartedEvent() + late_event = GraphRunStartedEvent() + + stream.notify_layers(terminal_event) + + with pytest.raises(RuntimeError, match="after execution is complete"): + stream.collect(late_event) + assert terminal_event.sequence == 1 + assert late_event.sequence == 0 diff --git a/tests/engine/test_raw_engine_events.py b/tests/engine/test_raw_engine_events.py index 5484638..08a6d0f 100644 --- a/tests/engine/test_raw_engine_events.py +++ b/tests/engine/test_raw_engine_events.py @@ -72,7 +72,7 @@ def test_event_processor_collects_raw_stream_chunk_without_coordinator() -> None ), ) chunk = NodeRunStreamChunkEvent( - id="run-1", + node_execution_id="run-1", node_id="node-1", node_type=BuiltinNodeTypes.CODE, selector=["node-1", "answer"], @@ -102,7 +102,7 @@ def test_event_processor_collects_reasoning_chunk_without_warning( ), ) chunk = NodeRunReasoningChunkEvent( - id="run-1", + node_execution_id="run-1", node_id="node-1", node_type=BuiltinNodeTypes.CODE, selector=["node-1", "reasoning_content"], @@ -150,7 +150,7 @@ def test_event_processor_collects_traversal_events_before_node_success() -> None ), ) success = NodeRunSucceededEvent( - id="run-1", + node_execution_id="run-1", node_id="node-1", node_type=BuiltinNodeTypes.CODE, start_at=_now(), diff --git a/tests/engine/test_response_stream_filter.py b/tests/engine/test_response_stream_filter.py index 4860e42..fa02790 100644 --- a/tests/engine/test_response_stream_filter.py +++ b/tests/engine/test_response_stream_filter.py @@ -186,7 +186,7 @@ def _stream_chunk( is_final: bool = True, ) -> NodeRunStreamChunkEvent: return NodeRunStreamChunkEvent( - id=run_id, + node_execution_id=run_id, node_id=source_id, node_type=BuiltinNodeTypes.CODE, selector=list(selector or [source_id, "answer"]), @@ -204,7 +204,7 @@ def _reasoning_chunk( is_final: bool = False, ) -> NodeRunReasoningChunkEvent: return NodeRunReasoningChunkEvent( - id=run_id, + node_execution_id=run_id, node_id=source_id, node_type=BuiltinNodeTypes.CODE, selector=list(selector or [source_id, "reasoning_content"]), @@ -499,7 +499,7 @@ def test_response_stream_filter_reorders_buffered_stream_chunks_after_edge_taken event_filter.initialize(_context(graph)) started = NodeRunStartedEvent( - id="source-run", + node_execution_id="source-run", node_id="source", node_type=BuiltinNodeTypes.CODE, node_title="Source", @@ -555,7 +555,7 @@ def test_response_stream_filter_uses_retry_execution_id_for_scalar_value() -> No event_filter = ResponseStreamFilter() event_filter.initialize(_context(graph, variable_pool)) retry = NodeRunRetryEvent( - id="retry-run", + node_execution_id="retry-run", node_id="source", node_type=BuiltinNodeTypes.CODE, node_title="Source", @@ -568,7 +568,9 @@ def test_response_stream_filter_uses_retry_execution_id_for_scalar_value() -> No output = list(event_filter.on_event(_edge_taken())) chunks = [event for event in output if isinstance(event, NodeRunStreamChunkEvent)] - assert [(event.id, event.chunk) for event in chunks] == [("retry-run", "saved")] + assert [(event.node_execution_id, event.chunk) for event in chunks] == [ + ("retry-run", "saved") + ] def test_response_stream_filter_initialize_resets_run_state() -> None: @@ -612,6 +614,32 @@ def test_response_stream_filter_round_trips_resume_state() -> None: assert [event.chunk for event in chunks] == ["resumed"] +def test_response_stream_filter_restores_identity_for_synthesized_chunks() -> None: + graph = _variable_response_graph() + variable_pool = VariablePool() + variable_pool.add(["source", "answer"], StringSegment(value="resumed")) + context = _context(graph, variable_pool) + first_filter = ResponseStreamFilter() + first_filter.initialize(context) + started = GraphRunStartedEvent( + id="start-event", + graph_id="graph-1", + execution_id="execution-1", + ) + assert list(first_filter.on_event(started)) == [started] + + restored_filter = ResponseStreamFilter() + restored_filter.initialize(context) + restored_filter.loads(first_filter.dumps()) + output = list(restored_filter.on_event(_edge_taken())) + + chunks = [event for event in output if isinstance(event, NodeRunStreamChunkEvent)] + assert len(chunks) == 1 + assert chunks[0].graph_id == "graph-1" + assert chunks[0].execution_id == "execution-1" + assert chunks[0].id != started.id + + def test_response_stream_filter_can_load_before_filter_chain_initializes() -> None: graph = _variable_response_graph() context = _context(graph) diff --git a/tests/engine/test_serializable_graph_runtime.py b/tests/engine/test_serializable_graph_runtime.py index 32e897a..76da2fc 100644 --- a/tests/engine/test_serializable_graph_runtime.py +++ b/tests/engine/test_serializable_graph_runtime.py @@ -231,6 +231,7 @@ def _hitl_engine( *, runtime_state: RuntimeState, callback: HITLCallback, + execution_id: str | None = None, ) -> Engine: plan = inspect(dsl) graph_config = plan.document.graph_config @@ -262,6 +263,7 @@ def _hitl_engine( graph=graph, graph_runtime_state=runtime_state, workers=2, + execution_id=execution_id, ) @@ -599,6 +601,7 @@ def pause_after_one_round(context: HITLContext) -> Completed | PauseRequested: _loop_dsl(), runtime_state=_new_runtime_state({}), callback=pause_after_one_round, + execution_id="explicit-execution", ) ) paused_state = RuntimeState.from_snapshot(snapshot) @@ -664,7 +667,14 @@ def pause_after_one_round(context: HITLContext) -> Completed | PauseRequested: ] == [0, 1, 2] assert len(paused_successes) == 1 assert len(resumed_successes) == 2 - assert loop_started.id == loop_succeeded.id + assert loop_started.node_execution_id == loop_succeeded.node_execution_id + all_events = [*paused_events, *resumed_events] + assert {event.graph_id for event in all_events} == {"workflow"} + assert {event.execution_id for event in all_events} == {"explicit-execution"} + assert len({event.id for event in all_events}) == len(all_events) + assert [event.sequence for event in all_events] == list( + range(1, len(all_events) + 1) + ) assert not any( isinstance(event, NodeRunLoopStartedEvent) for event in resumed_events ) @@ -690,6 +700,10 @@ def pause_second_round(context: HITLContext) -> Completed | PauseRequested: paused_events = list(engine.run()) resumed_events = list(engine.run()) + all_events = [*paused_events, *resumed_events] + assert [event.sequence for event in all_events] == list( + range(1, len(all_events) + 1) + ) assert isinstance(paused_events[-1], GraphRunPausedEvent) assert isinstance(resumed_events[0], GraphRunStartedEvent) assert resumed_events[0].reason == WorkflowStartReason.RESUMPTION @@ -797,11 +811,11 @@ def pause_with_active_sibling( for event in resumed_successes ) == [1, 2] assert next( - event.id + event.node_execution_id for event in paused_events if isinstance(event, NodeRunIterationStartedEvent) ) == next( - event.id + event.node_execution_id for event in resumed_events if isinstance(event, NodeRunIterationSucceededEvent) ) diff --git a/tests/engine_events/test_traversal_events.py b/tests/engine_events/test_traversal_events.py index a0ac9e5..6ecd00d 100644 --- a/tests/engine_events/test_traversal_events.py +++ b/tests/engine_events/test_traversal_events.py @@ -1,4 +1,10 @@ -from graphon.engine_events import GraphEdgeSkippedEvent, GraphEdgeTakenEvent +from graphon.engine_events import ( + EngineEvent, + GraphEdgeSkippedEvent, + GraphEdgeTakenEvent, +) + +_ENVELOPE_FIELDS = {*EngineEvent.model_fields, "event_type"} def test_graph_edge_taken_event_exports_payload() -> None: @@ -9,7 +15,7 @@ def test_graph_edge_taken_event_exports_payload() -> None: source_handle="success", ) - assert event.model_dump() == { + assert event.model_dump(exclude=_ENVELOPE_FIELDS) == { "edge_id": "edge-1", "source_node_id": "source", "target_node_id": "target", @@ -25,7 +31,7 @@ def test_graph_edge_skipped_event_exports_payload() -> None: target_node_id="other", ) - assert event.model_dump() == { + assert event.model_dump(exclude=_ENVELOPE_FIELDS) == { "edge_id": "edge-2", "source_node_id": "source", "target_node_id": "other", diff --git a/tests/node_events/test_node_event_aliases.py b/tests/node_events/test_node_event_aliases.py index 8a0e30b..98f4c47 100644 --- a/tests/node_events/test_node_event_aliases.py +++ b/tests/node_events/test_node_event_aliases.py @@ -21,7 +21,7 @@ def test_variable_alias_still_validates_in_event_models() -> None: node_event = VariableUpdatedEvent.model_validate(payload) graph_event = NodeRunVariableUpdatedEvent.model_validate({ **payload, - "id": "evt-1", + "node_execution_id": "node-execution-1", "node_id": "start", "node_type": "start", "node_run_result": NodeRunResult( @@ -39,7 +39,7 @@ def test_pause_reason_alias_still_validates_in_event_models() -> None: node_event = PauseRequestedEvent.model_validate(payload) graph_event = NodeRunPauseRequestedEvent.model_validate({ **payload, - "id": "evt-2", + "node_execution_id": "node-execution-2", "node_id": "start", "node_type": "start", "node_run_result": NodeRunResult( diff --git a/tests/nodes/tool/test_tool_node.py b/tests/nodes/tool/test_tool_node.py index dd6ab42..e87dff2 100644 --- a/tests/nodes/tool/test_tool_node.py +++ b/tests/nodes/tool/test_tool_node.py @@ -438,8 +438,10 @@ def test_run_passes_variable_pool_and_bound_execution_id_to_runtime( } assert node.execution_id == "bound-execution" assert isinstance(events[0], NodeRunStartedEvent) - assert events[0].id == "bound-execution" + assert events[0].node_execution_id == "bound-execution" assert isinstance(events[-1], NodeRunSucceededEvent) + assert events[-1].node_execution_id == "bound-execution" + assert events[0].id != events[-1].id def test_run_requires_runtime_adapter_to_accept_execution_id() -> None: diff --git a/tests/runtime/test_graph_runtime_state.py b/tests/runtime/test_graph_runtime_state.py index f482eae..3f4f881 100644 --- a/tests/runtime/test_graph_runtime_state.py +++ b/tests/runtime/test_graph_runtime_state.py @@ -311,6 +311,7 @@ def test_workflow_id_creates_graph_execution(self) -> None: assert isinstance(execution, GraphExecution) assert execution.workflow_id == "workflow" + assert execution.execution_id assert state.graph_execution is execution def test_graph_configuration_rejects_different_graph(self) -> None: @@ -397,6 +398,9 @@ def test_dumps_and_loads_roundtrip(self) -> None: state.ready_queue.put(StartTask(frame_id="root", node_id="node-A")) graph_execution = state.graph_execution + graph_execution.execution_id = "execution-123" + assert graph_execution.next_event_sequence() == 1 + assert graph_execution.next_event_sequence() == 2 graph_execution.exceptions_count = 4 graph_execution.started = True graph_execution.error = ValueError("saved failure") @@ -421,6 +425,9 @@ def test_dumps_and_loads_roundtrip(self) -> None: restored_execution = restored.graph_execution assert restored_execution.workflow_id == "wf-123" + assert restored_execution.execution_id == "execution-123" + assert restored_execution.last_event_sequence == 2 + assert restored_execution.next_event_sequence() == 3 assert restored_execution.exceptions_count == 4 assert restored_execution.started is True assert isinstance(restored_execution.error, RuntimeError) @@ -479,6 +486,8 @@ def test_version_1_snapshot_migrates_to_frame_aware_version_2(self) -> None: assert json.loads(migrated["ready_queue"])["version"] == "2.0" assert json.loads(migrated["deferred_ready_tasks"])["version"] == "2.0" assert json.loads(migrated["graph_execution"])["version"] == "2.0" + assert restored.graph_execution.execution_id + assert restored.graph_execution.last_event_sequence == 0 assert migrated["graph_node_states"] == {"ready": NodeState.TAKEN} assert migrated["graph_edge_states"] == {"edge": NodeState.SKIPPED} assert restored.ready_queue.drain() == [ @@ -549,3 +558,25 @@ def test_snapshot_restore_preserves_file_segments(self) -> None: assert restored_file.value.filename == "resume.pdf" assert isinstance(restored_files, ArrayFileSegment) assert restored_files.value[0].filename == "resume.pdf" + + +def test_version_2_graph_execution_without_sequence_defaults_to_zero() -> None: + execution = GraphExecution.from_snapshot( + json.dumps({ + "version": "2.0", + "workflow_id": "workflow", + "execution_id": "execution", + "started": True, + "completed": False, + "aborted": False, + "paused": True, + "pause_reasons": [], + "error": None, + "exceptions_count": 0, + "node_executions": [], + }) + ) + + assert execution.execution_id == "execution" + assert execution.last_event_sequence == 0 + assert execution.next_event_sequence() == 1 diff --git a/tests/workflows/test_full_engine_events.py b/tests/workflows/test_full_engine_events.py index 9cc9ce4..102b459 100644 --- a/tests/workflows/test_full_engine_events.py +++ b/tests/workflows/test_full_engine_events.py @@ -1,8 +1,10 @@ from __future__ import annotations from collections.abc import Mapping, Sequence +from datetime import UTC from threading import Event, Thread -from typing import Any +from typing import Any, cast +from unittest.mock import MagicMock import pytest import yaml @@ -11,6 +13,7 @@ from graphon.engine import Engine from graphon.engine.container_handler import LoopContainerHandler from graphon.engine.frame import FrameRegistry +from graphon.engine.layer import Layer from graphon.engine.ready_queue.entities import StartTask from graphon.engine.ready_queue.in_memory import InMemoryReadyQueue from graphon.engine_events.base import EngineEvent @@ -224,7 +227,14 @@ def test_full_answer_graph_is_verified_from_events() -> None: edges=[_edge("start", "answer")], ) - events = run_workflow(dsl, start_inputs={"name": "Graphon"}) + engine = loads( + dsl, + workflow_id="workflow-envelope", + start_inputs={"name": "Graphon"}, + ) + layer = MagicMock(spec=Layer) + engine.add_layer(cast(Layer, layer)) + events = list(engine.run()) assert event_path(events) == [ _event("GraphRunStartedEvent"), @@ -235,6 +245,22 @@ def test_full_answer_graph_is_verified_from_events() -> None: _event("NodeRunSucceededEvent", "answer"), _event("GraphRunSucceededEvent"), ] + assert {event.graph_id for event in events} == {"workflow-envelope"} + assert len({event.execution_id for event in events}) == 1 + assert events[0].execution_id + assert len({event.id for event in events}) == len(events) + assert [event.sequence for event in events] == list(range(1, len(events) + 1)) + assert [event.event_type for event in events] == [ + type(event).__name__ for event in events + ] + assert {event.schema_version for event in events} == {"1.0"} + assert all(event.emitted_at.tzinfo is UTC for event in events) + layer_events = [call.args[0] for call in layer.on_event.call_args_list] + assert len(layer_events) == len(events) + assert all( + layer_event is emitted_event + for layer_event, emitted_event in zip(layer_events, events, strict=True) + ) outputs = final_outputs(events) assert set(outputs) == {"answer", "files"} assert outputs["answer"] == "Hello Graphon"