feat(graph_engine): add ready_queue state persistence to GraphRuntimeState

- Add ReadyQueueState TypedDict for type-safe queue serialization
- Add ready_queue attribute to GraphRuntimeState for initializing with pre-existing queue state
- Update GraphEngine to load ready_queue from GraphRuntimeState on initialization
- Implement proper type hints using ReadyQueueState for better type safety
- Add comprehensive tests for ready_queue loading functionality

The ready_queue is read-only after initialization and allows resuming workflow
execution with a pre-populated queue of nodes ready to execute.
This commit is contained in:
-LAN-
2025-09-15 03:05:10 +08:00
parent 0f15a2baca
commit b4ef1de30f
7 changed files with 159 additions and 16 deletions

View File

@ -4,6 +4,7 @@ import pytest
from core.workflow.entities.graph_runtime_state import GraphRuntimeState
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.graph_engine.ready_queue import ReadyQueueState
class TestGraphRuntimeState:
@ -109,3 +110,30 @@ class TestGraphRuntimeState:
# Original should remain unchanged
assert state.get_output("nested")["level1"]["level2"]["value"] == "test"
def test_ready_queue_property(self):
variable_pool = VariablePool()
# Test default empty ready_queue
state = GraphRuntimeState(variable_pool=variable_pool, start_at=time())
assert state.ready_queue == {}
# Test initialization with ready_queue data as ReadyQueueState
queue_data = ReadyQueueState(type="InMemoryReadyQueue", version="1.0", items=["node1", "node2"], maxsize=0)
state = GraphRuntimeState(variable_pool=variable_pool, start_at=time(), ready_queue=queue_data)
assert state.ready_queue == queue_data
# Test with different ready_queue data at initialization
another_queue_data = ReadyQueueState(
type="InMemoryReadyQueue",
version="1.0",
items=["node3", "node4", "node5"],
maxsize=0,
)
another_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time(), ready_queue=another_queue_data)
assert another_state.ready_queue == another_queue_data
# Test immutability - modifying retrieved queue doesn't affect internal state
retrieved_queue = state.ready_queue
retrieved_queue["items"].append("node6")
assert len(state.ready_queue["items"]) == 2 # Should still be 2, not 3

View File

@ -744,3 +744,78 @@ def test_event_sequence_validation_with_table_tests():
else:
assert result.event_sequence_match is True
assert result.success, f"Test {i + 1} failed: {result.event_mismatch_details or result.error}"
def test_ready_queue_state_loading():
"""
Test that the ready_queue state is properly loaded from GraphRuntimeState
during GraphEngine initialization.
"""
# Use the TableTestRunner to create a proper workflow instance
runner = TableTestRunner()
# Create a simple workflow
test_case = WorkflowTestCase(
fixture_path="simple_passthrough_workflow",
inputs={"query": "test"},
expected_outputs={"query": "test"},
description="Test ready_queue loading",
)
# Load the workflow fixture
workflow_runner = runner.workflow_runner
fixture_data = workflow_runner.load_fixture("simple_passthrough_workflow")
# Create graph and runtime state with pre-populated ready_queue
ready_queue_data = {
"type": "InMemoryReadyQueue",
"version": "1.0",
"items": ["node1", "node2", "node3"],
"maxsize": 0,
}
# We need to create the graph first, then create a new GraphRuntimeState with ready_queue
graph, original_runtime_state = workflow_runner.create_graph_from_fixture(fixture_data, query="test")
# Create a new GraphRuntimeState with the ready_queue data
from core.workflow.entities import GraphRuntimeState
from core.workflow.graph_engine.ready_queue import ReadyQueueState
# Convert ready_queue_data to ReadyQueueState
ready_queue_state = ReadyQueueState(**ready_queue_data)
graph_runtime_state = GraphRuntimeState(
variable_pool=original_runtime_state.variable_pool,
start_at=original_runtime_state.start_at,
ready_queue=ready_queue_state,
)
# Update all nodes to use the new GraphRuntimeState
for node in graph.nodes.values():
node.graph_runtime_state = graph_runtime_state
# Create GraphEngine
command_channel = InMemoryChannel()
engine = GraphEngine(
tenant_id="test-tenant",
app_id="test-app",
workflow_id="test-workflow",
user_id="test-user",
user_from=UserFrom.ACCOUNT,
invoke_from=InvokeFrom.DEBUGGER,
call_depth=0,
graph=graph,
graph_config={},
graph_runtime_state=graph_runtime_state,
command_channel=command_channel,
)
# Verify that the ready_queue was loaded from GraphRuntimeState
assert engine._ready_queue.qsize() == 3
# Verify the initial state matches what was provided
initial_queue_state = engine.graph_runtime_state.ready_queue
assert initial_queue_state["type"] == "InMemoryReadyQueue"
assert initial_queue_state["version"] == "1.0"
assert len(initial_queue_state["items"]) == 3
assert initial_queue_state["items"] == ["node1", "node2", "node3"]