even MORE code review

stateful to not deadhang, graph context -> output, more explicit state errors, other misc changes idfk any of those good thing theyre fixed tho
This commit is contained in:
lda
2026-05-26 02:33:41 +07:00 Verified
parent 9d11f78111
commit 95ec73d518
12 changed files with 159 additions and 14 deletions
+3 -2
View File
@@ -62,8 +62,9 @@ implementation state.
`WorkflowBuilder.prepare_subgraph()` and `WorkflowBuilder.resume()` make the
local runnable/resumable path available without core-runtime plumbing.
Saved interrupting artifacts can now pause and resume through
`run_deployment`/`resume_run` while the MCP server process stays alive.
Persisted resume remains future work.
`run_deployment`/`resume_run` for the duration of the MCP server process
(in-memory only). Persisted resume across process restarts remains future
work.
- **Concurrent foreach**: implemented in core with explicit scheduling,
reducer/merge semantics, item error policy, async handler batching, and
quiescent interrupt behavior. Remaining work is polish and future reuse of
+3
View File
@@ -37,6 +37,9 @@ Final workflow output has two projection modes:
- If `Workflow.output` is empty, legacy projection copies same-named top-level
state keys listed in `workflow.output_schema.properties`.
Prefer explicit `Workflow.output` bindings for new workflows. The same-named
state projection exists for compatibility with older plans.
The projected payload is then validated against `workflow.output_schema`.
## Why This Matters
+2 -3
View File
@@ -207,7 +207,7 @@ class WorkflowBuilder:
input_schema: SchemaLike
state_schema: StateSchemaLike
output_schema: SchemaLike
outcomes: Sequence[str] | None = None
outcomes: Sequence[str] = ("ok",)
start: str | None = None
reducers: ReducerCatalog | Mapping[str, ReducerDefinition] | None = None
node_specs: dict[str, NodeSpec[Any, Any]] = field(default_factory=dict)
@@ -222,7 +222,6 @@ class WorkflowBuilder:
self.input_schema = schema_ref_from(self.input_schema)
self.state_schema = state_schema_from(self.state_schema)
self.output_schema = schema_ref_from(self.output_schema)
self.outcomes = list(self.outcomes or ["ok"])
@overload
def use(
@@ -871,7 +870,7 @@ class WorkflowBuilder:
input_schema=cast(SchemaRef, self.input_schema),
state_schema=cast(StateSchema, self.state_schema),
output_schema=cast(SchemaRef, self.output_schema),
outcomes=list(self.outcomes or ["ok"]),
outcomes=list(self.outcomes),
node_defs=node_defs,
start=self.start,
nodes=self.nodes,
+4 -2
View File
@@ -363,8 +363,10 @@ def _state_write_from_metadata(raw: object) -> StateWrite:
visible_value=visible_value,
reducer=ReducerRef.model_validate(reducer),
)
except Exception as exc:
raise WorkflowExecutionError("malformed pending foreach write") from exc
except WorkflowExecutionError:
raise
except (TypeError, ValueError) as exc:
raise WorkflowExecutionError(f"malformed pending foreach write: {exc}") from exc
def _state_write_to_metadata(write: StateWrite) -> dict[str, Any]:
+6
View File
@@ -216,7 +216,13 @@ def _lineage_chain(
) -> Iterator[LineageState]:
lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
reverse_chain: list[LineageState] = []
seen: set[str] = set()
while True:
if lineage.id in seen:
raise WorkflowExecutionError(
f"cycle detected in lineage chain at {lineage.id!r}"
)
seen.add(lineage.id)
reverse_chain.append(lineage)
if lineage.parent_id is None:
break
+1
View File
@@ -253,6 +253,7 @@ def _finish_subgraph(
prepared.workflow,
child_scope.committed_state,
workflow_input=child_scope.workflow_input,
context=frame_context_values(child_frame),
)
validate_payload_against_schema(
prepared.workflow.output_schema,
+16 -4
View File
@@ -141,7 +141,7 @@ class _SessionOwner:
tool_name: str,
payload: dict[str, object],
) -> CallToolResult:
"""Submit a tool call for execution in the transport owner task."""
"""Submit a call and fail promptly if its transport owner exits."""
task = self._task
if task is None:
raise RuntimeError("persistent MCP session is not started")
@@ -152,7 +152,13 @@ class _SessionOwner:
await self._requests.put(
_ToolCallRequest(tool_name=tool_name, payload=payload, result=result)
)
return await result
done, _pending = await asyncio.wait(
{result, task}, return_when=asyncio.FIRST_COMPLETED
)
if result in done:
return result.result()
await task
raise RuntimeError("persistent MCP session stopped unexpectedly")
async def close(self) -> None:
"""Ask the owner task to close the MCP transport in its own scope."""
@@ -187,5 +193,11 @@ class _SessionOwner:
except BaseException as exc:
if not ready.done():
ready.set_exception(exc)
else:
raise
return
# Calls already queued behind the failing request cannot otherwise
# observe that their sole transport owner has exited.
while not self._requests.empty():
pending = self._requests.get_nowait()
if pending is not None and not pending.result.done():
pending.result.set_exception(exc)
raise
+3 -2
View File
@@ -604,8 +604,9 @@ def register_workflow_tools(server: FastMCP[Any], service: WfMcpService) -> None
title="Run Workflow Deployment",
description=(
"Run deployment_id with workflow_input and return status, terminal "
"outcome, output, diagnostics, and trace_count. Debug traces can include resolved "
"inputs and state changes; pass trace_range only when needed."
"outcome when completed, output, diagnostics, and trace_count. "
"Debug traces can include resolved inputs and state changes; pass "
"trace_range only when needed."
),
)
async def run_deployment(
+16
View File
@@ -10,6 +10,7 @@ from wf_core.runtime.foreach_state import (
ForeachBarrierState,
ItemErrorRecord,
PendingItemResult,
_state_write_from_metadata,
)
from wf_core.runtime.lineage import LineageStateView, lineage_writes_for_frame
from wf_core.runtime.ops.state import StatePatch
@@ -88,6 +89,21 @@ def test_foreach_barrier_state_round_trips_reducer_write_records() -> None:
assert write.reducer.name == "wf.std.add"
def test_pending_write_metadata_preserves_invalid_reducer_detail() -> None:
with pytest.raises(WorkflowExecutionError, match="mutually exclusive"):
_state_write_from_metadata(
{
"path": {"root": "state", "parts": ["count"]},
"incoming_value": 1,
"visible_value": 1,
"reducer": {
"name": "wf.std.add",
"ref": {"source": "wf.std", "capability_key": "add"},
},
}
)
def test_lineage_state_view_materializes_visible_values_without_mutating_base() -> None:
base_state = {"count": 2, "nested": {"value": "old"}}
view = LineageStateView(
+11
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
import pytest
from wf_core import (
END,
Edge,
@@ -9,6 +11,7 @@ from wf_core import (
StateField,
StateSchema,
Workflow,
WorkflowExecutionError,
)
from wf_core.models.reducers import ReducerRef
from wf_core.paths import StatePath
@@ -158,6 +161,14 @@ def test_state_view_for_frame_overlays_writes_onto_frame_scope_state() -> None:
assert run.state["value"] == "root"
def test_lineage_state_view_rejects_corrupt_parent_cycle() -> None:
run = create_run_state(_minimal_workflow(), {"value": "root"})
run.lineages["root"].parent_id = "root"
with pytest.raises(WorkflowExecutionError, match="cycle"):
lineage_state_view(run, scope_id="root", lineage_id="root")
def test_lineage_state_view_handles_deep_ancestry_without_recursion() -> None:
run = create_run_state(_minimal_workflow(), {"value": "root"})
parent_id = "root"
+30 -1
View File
@@ -24,8 +24,9 @@ from wf_core import (
resume_workflow_async,
resume_workflow,
)
from wf_core.models.steps import InputPathBinding, Step
from wf_core.paths import GraphSourcePath, LocalPath
from wf_core.validation.issues import ValidationIssueCode
from wf_core.models.steps import Step
def test_subgraph_step_validates_boundary_bindings_and_outcomes() -> None:
@@ -158,6 +159,34 @@ def test_subgraph_step_executes_caller_prepared_saved_child_ref() -> None:
assert run.trace[-1].step_type == "subgraph"
def test_subgraph_step_projects_child_context_output() -> None:
child = Workflow(
name="context_child",
input_schema=_schema({"text": {"type": "string"}}),
state_schema=StateSchema.from_field_map({}),
output_schema=_schema({"answer": {"type": "string"}}),
output=[
InputPathBinding(
target=LocalPath.of("answer"),
path=GraphSourcePath.parse("context.scope_id"),
)
],
outcomes=["ok"],
start="done",
nodes=[EndNode(id="done", type="end", outcome="ok")],
edges=[],
)
run = execute_workflow(
_workflow(output_schema=_schema({"answer": {"type": "string"}})),
{"text": "hello"},
{},
subgraphs={"child.workflow": PreparedSubgraph(workflow=child, registry={})},
)
assert run.output["answer"] == "root:subgraph:child"
def test_subgraph_step_executes_prepared_async_child() -> None:
async def answer(payload: dict[str, object], _ctx: object) -> dict[str, object]:
return {"answer": f"async:{payload['text']}"}
+64
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from typing import Any, cast
@@ -11,6 +12,7 @@ from wf_core import RuntimeContext
from wf_mcp.capabilities import DiscoveredTool
from wf_mcp.models import AuthRecord, ConnectionConfig
from wf_mcp.runtime import McpRuntimePool, PersistentMcpSession
from wf_mcp.runtime.factory import PersistentSessionFactory
from wf_mcp.sdk import ToolCallResult
from wf_mcp.workflow import wrap_discovered_tool
@@ -73,6 +75,36 @@ class FakeStatefulClient:
self.closed = True
class OwnerCrash(BaseException):
"""Simulate transport-owner death outside normal per-request exceptions."""
@dataclass(slots=True)
class CrashingClient:
started: asyncio.Event
crash: asyncio.Event
async def call_tool(
self, tool_name: str, payload: dict[str, object]
) -> CallToolResult:
self.started.set()
await self.crash.wait()
raise OwnerCrash("transport owner died")
class CrashingSessionFactory(PersistentSessionFactory):
def __init__(self, client: CrashingClient) -> None:
self.client = client
async def _create_with_stack(
self,
stack: AsyncExitStack,
connection: ConnectionConfig,
auth: AuthRecord | None,
) -> ClientSession:
return cast(ClientSession, self.client)
def _tool(name: str) -> DiscoveredTool:
return DiscoveredTool(
name=name,
@@ -199,3 +231,35 @@ def test_runtime_pool_replaces_session_when_fingerprint_changes() -> None:
assert len(created_clients) == 2
assert created_clients[0].closed is True
def test_persistent_session_fails_inflight_and_queued_calls_if_owner_dies() -> None:
connection = ConnectionConfig(
id="failing.default",
server="failing",
account="default",
metadata={},
)
async def exercise() -> tuple[
BaseException | ToolCallResult, BaseException | ToolCallResult
]:
started = asyncio.Event()
crash = asyncio.Event()
session = await CrashingSessionFactory(
CrashingClient(started=started, crash=crash)
).create(connection, None)
first = asyncio.create_task(session.call_tool("first", {}))
await started.wait()
second = asyncio.create_task(session.call_tool("second", {}))
await asyncio.sleep(0)
crash.set()
return await asyncio.wait_for(
asyncio.gather(first, second, return_exceptions=True),
timeout=0.2,
)
results = asyncio.run(exercise())
assert isinstance(results[0], OwnerCrash)
assert isinstance(results[1], OwnerCrash)