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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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']}"}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user