lda.chat/lda-wf
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: close run budget and MCP migration gaps

Browse Source
This commit is contained in:
lda
2026-09-06 19:21:10 +07:00 Verified
parent 6d5b6741fb
commit e36e462fd6
16 changed files with 267 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/superpowers/specs/2026-06-27-draft-semantic-authoring-boundary.md
+2 -2
View File
@@ -85,8 +85,8 @@ semantic operation produces one patch and consumes one revision.
clients do not need to know about the internal service split.
The service boundary is intentionally not capability-only. The current draft
model also represents `end`, `condition`, `interrupt`, `foreach`, `when`,
`choose`, `match`, and subgraph steps, and core may gain more step kinds. This
model also represents `end`, `interrupt`, `foreach`, `when`, `choose`, `match`,
and subgraph steps, and core may gain more step kinds. This
slice adds semantic operations only where required, but new step-kind helpers
belong in `WorkflowDraftAuthoringApi` rather than a parallel authoring system.
src/wf_api/runs.py
+4 -3
View File
@@ -83,6 +83,9 @@ class WorkflowRunApi:
max_steps: int | None = None,
) -> RunResult:
trace_values = _trace_range_values(trace_range)
limits = (
RunLimits(max_steps=max_steps) if max_steps is not None else RunLimits()
)
deployment, artifact, diagnostics, tree = (
self.deployments.deployment_validation(deployment_id)
)
@@ -92,12 +95,10 @@ class WorkflowRunApi:
artifact=artifact,
status="unrunnable",
diagnostics=diagnostics,
max_steps=limits.max_steps,
)
plan = raw_plan_from_artifact(artifact)
limits = (
RunLimits(max_steps=max_steps) if max_steps is not None else RunLimits()
)
run = await self.context.runtime.run_workflow_from_plan(
plan,
workflow_input,
src/wf_core/runtime/ops/flow.py
+11 -8
View File
@@ -22,11 +22,14 @@ from wf_core.runtime.scheduler import (
)
from wf_core.tokens import END
# Sentinel for ``append_trace()``: copy the named frame's admitted step number
# (failing closed when unassigned). Interrupt resume passes its stored
# activation number explicitly instead, so one activation keeps one number
# across its interrupt and resume-completion entries without a second admission.
_FROM_FRAME: Any = object()
class _FromFrame:
"""Sentinel type for resolving a trace step number from its frame."""
# Interrupt resume passes its stored activation number explicitly, so one
# activation keeps one number across interrupt and completion trace entries.
_FROM_FRAME = _FromFrame()
def append_trace(
@@ -40,7 +43,7 @@ def append_trace(
next_node_id: str,
output: dict[str, Any],
state_changes: dict[str, Any],
step_number: int | None | Any = _FROM_FRAME,
step_number: int | None | _FromFrame = _FROM_FRAME,
) -> None:
"""Append one trace entry carrying its admitted step number.
@@ -50,7 +53,7 @@ def append_trace(
``WorkflowExecutionError``. Pass ``step_number`` explicitly only to reuse a
persisted activation number (interrupt resume-completion).
"""
if step_number is _FROM_FRAME:
if isinstance(step_number, _FromFrame):
frame = run.frames.get(frame_id)
if frame is None:
raise WorkflowExecutionError(
@@ -85,7 +88,7 @@ def append_step_result_trace(
step_type: str,
next_node_id: str,
result: StepExecutionResult,
step_number: int | None | Any = _FROM_FRAME,
step_number: int | None | _FromFrame = _FROM_FRAME,
) -> None:
append_trace(
run,
src/wf_mcp/proxy/safe_names.py
+7 -1
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import hashlib
import re
from asyncio import Lock
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
@@ -39,6 +40,7 @@ class SafeToolNames(Transform):
self._original_to_safe: dict[str, str] = {}
self._server = server
self._primed = False
self._prime_lock = Lock()
async def list_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]:
return [
@@ -70,8 +72,12 @@ class SafeToolNames(Transform):
them; a failed prime falls through to the identity fallback so the
call still ends in a proper unknown-tool error.
"""
if self._server is None or self._primed:
if self._server is None:
return None
# A direct-call burst must share the first live listing. Publishing
# ``_primed`` without this barrier lets siblings observe empty maps.
async with self._prime_lock:
if not self._primed:
self._primed = True
try:
await self._server.list_tools()
src/wf_sources_mcp/client/source_client.py
+16 -19
View File
@@ -1,25 +1,25 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Protocol
from typing import TYPE_CHECKING, Any, Protocol, cast
from mcp.types import (
CallToolResult,
ClientNotification,
ClientRequest,
GetPromptResult,
ListPromptsResult,
ListResourcesResult,
ListToolsResult,
ReadResourceResult,
ServerResult,
client_notification_adapter,
client_request_adapter,
server_result_adapter,
)
from pydantic import TypeAdapter
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.raw_messages import (
RawRequest,
RawResult,
raw_notification,
raw_request,
)
if TYPE_CHECKING:
from wf_sources_mcp.catalog import (
@@ -58,9 +58,9 @@ class McpClientSession(Protocol):
# BaseSession stuff. not even complete signature, thats crazy
async def send_request(
self,
request: ClientRequest,
result_type: type[ServerResult] | TypeAdapter[ServerResult],
) -> ServerResult: ...
request: RawRequest,
result_type: type[RawResult],
) -> RawResult: ...
async def send_notification(self, notification: ClientNotification) -> None: ...
@@ -138,10 +138,8 @@ class McpSourceClient:
params: dict[str, Any] | None = None,
) -> dict[str, Any]:
result = await self.session.send_request(
client_request_adapter.validate_python(
{"method": method, "params": params}
),
server_result_adapter,
raw_request(method, params),
RawResult,
)
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
@@ -150,11 +148,10 @@ class McpSourceClient:
method: str,
params: dict[str, Any] | None = None,
) -> None:
await self.session.send_notification(
client_notification_adapter.validate_python(
{"method": method, "params": params}
)
)
notification = raw_notification(method, params)
# MCP 2's runtime accepts the generic Notification base class, while
# its public annotation still names only the standard-method union.
await self.session.send_notification(cast(ClientNotification, notification))
async def call_tool(
self,
src/wf_sources_mcp/raw_messages.py
+36
View File
@@ -0,0 +1,36 @@
"""Generic MCP messages for the deliberately untyped extension surface."""
from typing import Any
from mcp.types import Notification, Request
from pydantic import BaseModel, ConfigDict
type RawParams = dict[str, Any] | None
type RawRequest = Request[RawParams, str]
type RawNotification = Notification[RawParams, str]
class RawResult(BaseModel):
"""Preserve every field returned by an extension method."""
model_config = ConfigDict(extra="allow")
def raw_request(method: str, params: RawParams) -> RawRequest:
"""Build an extension request without narrowing it to standard methods."""
return Request[RawParams, str](method=method, params=params)
def raw_notification(method: str, params: RawParams) -> RawNotification:
"""Build an extension notification without narrowing its method name."""
return Notification[RawParams, str](method=method, params=params)
__all__ = [
"RawNotification",
"RawParams",
"RawRequest",
"RawResult",
"raw_notification",
"raw_request",
]
src/wf_sources_mcp/runtime/session.py
+12 -14
View File
@@ -2,13 +2,19 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any
from typing import Any, cast
from mcp.client.session import ClientSession
from mcp.types import ClientNotification
from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.raw_messages import (
RawResult,
raw_notification,
raw_request,
)
from wf_sources_mcp.sdk import ToolCallResult
from wf_sources_mcp.sdk.converters import tool_result_to_call_result
@@ -134,13 +140,9 @@ class PersistentMcpSession:
if self.invoke_method_callback is not None:
return await self.invoke_method_callback(method, params)
if self.client is not None:
from mcp.types import client_request_adapter, server_result_adapter
result = await self.client.send_request(
client_request_adapter.validate_python(
{"method": method, "params": params}
),
server_result_adapter,
raw_request(method, params),
RawResult,
)
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
raise RuntimeError("persistent MCP session has no method invoke transport")
@@ -155,13 +157,9 @@ class PersistentMcpSession:
await self.send_notification_callback(method, params)
return
if self.client is not None:
from mcp.types import client_notification_adapter
await self.client.send_notification(
client_notification_adapter.validate_python(
{"method": method, "params": params}
)
)
notification = raw_notification(method, params)
# MCP 2's annotation has not widened to its generic runtime shape.
await self.client.send_notification(cast(ClientNotification, notification))
return
raise RuntimeError("persistent MCP session has no notification send transport")
tests/core/test_structured_context_validation.py
+44 -43
View File
@@ -1,9 +1,17 @@
from __future__ import annotations
from collections.abc import Callable
import pytest
from wf_core import END, Edge, ForeachNode, NodeDef, NodeUse, SchemaRef, Workflow
from wf_core.models.schemas import StateField, StateSchema
from wf_core.models.steps import (
ConditionNode,
InputPathBinding,
InterruptNode,
SubgraphNode,
)
from wf_core.validation.issues import ValidationIssueCode
@@ -275,13 +283,7 @@ def test_node_input_surfaces_report_exact_model_paths(make_node, expected_path)
)
def test_all_model_surfaces_reject_missing_foreach_id() -> None:
from wf_core.models.steps import ConditionNode, InterruptNode, SubgraphNode
from wf_core.validation import validate_workflow
bad = "context.foreach.missing.item"
# Subgraph input
workflow = _base_workflow()
def _missing_foreach_subgraph(workflow: Workflow, bad: str) -> str:
workflow.nodes[2] = SubgraphNode.model_validate(
{
"id": "work",
@@ -290,16 +292,10 @@ def test_all_model_surfaces_reject_missing_foreach_id() -> None:
"input": [{"target": "order", "path": bad}],
}
)
report = validate_workflow(workflow)
assert (
_issue(
report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[2].input[0].path"
)
is not None
)
return "nodes[2].input[0].path"
# Condition check
workflow = _base_workflow()
def _missing_foreach_condition(workflow: Workflow, bad: str) -> str:
workflow.nodes[2] = ConditionNode.model_validate(
{"id": "work", "type": "condition", "check": {"op": "exists", "path": bad}}
)
@@ -314,25 +310,17 @@ def test_all_model_surfaces_reject_missing_foreach_id() -> None:
),
Edge.model_validate({"from": "customers", "outcome": "done", "to": END}),
]
report = validate_workflow(workflow)
assert (
_issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[2].check.path")
is not None
)
return "nodes[2].check.path"
# Foreach over
workflow = _base_workflow()
def _missing_foreach_over(workflow: Workflow, bad: str) -> str:
workflow.nodes[1] = ForeachNode.model_validate(
{"id": "orders", "type": "foreach", "over": bad, "as": "order"}
)
report = validate_workflow(workflow)
assert (
_issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[1].over")
is not None
)
return "nodes[1].over"
# Interrupt request
workflow = _base_workflow()
def _missing_foreach_interrupt(workflow: Workflow, bad: str) -> str:
workflow.nodes[2] = InterruptNode.model_validate(
{
"id": "work",
@@ -351,24 +339,37 @@ def test_all_model_surfaces_reject_missing_foreach_id() -> None:
),
Edge.model_validate({"from": "customers", "outcome": "done", "to": END}),
]
report = validate_workflow(workflow)
assert (
_issue(
report,
ValidationIssueCode.INVALID_CONTEXT_PATH,
"nodes[2].request[0].path",
)
is not None
)
return "nodes[2].request[0].path"
# Workflow output
from wf_core.models.steps import InputPathBinding as _IPB
def _missing_foreach_output(workflow: Workflow, bad: str) -> str:
workflow.output = [
InputPathBinding.model_validate({"target": "result", "path": bad})
]
return "output[0].path"
@pytest.mark.parametrize(
"configure_surface",
[
_missing_foreach_subgraph,
_missing_foreach_condition,
_missing_foreach_over,
_missing_foreach_interrupt,
_missing_foreach_output,
],
ids=["subgraph", "condition", "foreach", "interrupt", "workflow-output"],
)
def test_all_model_surfaces_reject_missing_foreach_id(
configure_surface: Callable[[Workflow, str], str],
) -> None:
from wf_core.validation import validate_workflow
workflow = _base_workflow()
workflow.output = [_IPB.model_validate({"target": "result", "path": bad})]
expected_path = configure_surface(workflow, "context.foreach.missing.item")
report = validate_workflow(workflow)
assert (
_issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "output[0].path")
_issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, expected_path)
is not None
)
tests/wf_api/test_runs.py
+48
View File
@@ -52,6 +52,24 @@ def _echo_service(root: Path) -> WfMcpService:
return service
def _unrunnable_service(root: Path) -> WfMcpService:
artifact_store = FileWorkflowArtifactStore(root)
artifact_store.save_artifact(echo_artifact())
artifact_store.save_deployment(
WorkflowDeployment(
id="echo.unbound",
artifact_id="echo",
artifact_version=1,
bindings=[],
)
)
return WfMcpService(
store=FileStore(root / "mcp"),
artifact_store=artifact_store,
run_store=FileRunStore(root / "mcp"),
)
def _interrupt_artifact() -> WorkflowArtifact:
return WorkflowArtifact(
id="approval",
@@ -162,6 +180,36 @@ async def test_run_deployment_rejects_non_positive_max_steps(tmp_path: Path) ->
)
async def test_invalid_budget_wins_over_unrunnable_deployment(tmp_path: Path) -> None:
api = WorkflowRunApi(
context_from_service(_unrunnable_service(tmp_path / "invalid_unrunnable"))
)
with pytest.raises(ValueError, match="positive"):
await api.run_deployment(
deployment_id="echo.unbound",
workflow_input={"text": "hello"},
max_steps=0,
)
async def test_unrunnable_deployment_reports_requested_budget(tmp_path: Path) -> None:
api = WorkflowRunApi(
context_from_service(_unrunnable_service(tmp_path / "requested_unrunnable"))
)
result = await api.run_deployment(
deployment_id="echo.unbound",
workflow_input={"text": "hello"},
max_steps=7,
)
assert result["status"] == "unrunnable"
assert result["max_steps"] == 7
assert result["steps_executed"] == 0
assert result["steps_remaining"] == 7
async def test_inspect_run_reports_effective_budget(tmp_path: Path) -> None:
api = WorkflowRunApi(context_from_service(_echo_service(tmp_path / "inspect")))
started = await api.run_deployment(
tests/wf_cli/test_app.py
+2
View File
@@ -1766,8 +1766,10 @@ def test_wf_draft_add_control_command_help_is_type_specific() -> None:
interrupt = runner.invoke(app, ["draft", "add", "interrupt", "--help"])
foreach = runner.invoke(app, ["draft", "add", "foreach", "--help"])
end = runner.invoke(app, ["draft", "add", "end", "--help"])
removed_join = runner.invoke(app, ["draft", "add", "join", "--help"])
assert interrupt.exit_code == foreach.exit_code == end.exit_code == 0
assert removed_join.exit_code != 0
assert "--request-schema-file" in interrupt.output
assert "--resume-schema-file" in interrupt.output
assert "--request" in interrupt.output
tests/wf_client/test_codec.py
+3
View File
@@ -152,6 +152,9 @@ def test_decode_run_result_returns_typed_domain_boundary() -> None:
assert result.run_id == "run-1"
assert result.output == {"result": "hello"}
assert result.max_steps == 10_000
assert result.steps_executed == 1
assert result.steps_remaining == 9_999
assert result.diagnostics == ()
tests/wf_mcp/test_safe_tool_names.py
+38
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import pytest
from fastmcp import FastMCP
from wf_mcp.proxy.safe_names import (
@@ -41,6 +42,43 @@ def test_safe_tool_names_hashes_overlength_names() -> None:
transform.assert_consistent()
async def test_concurrent_direct_calls_share_cold_start_priming(
monkeypatch: pytest.MonkeyPatch,
) -> None:
server: FastMCP[object] = FastMCP("safe-name-test")
for name in ("demo.one", "demo.two"):
def handler() -> None:
return None
server.tool(name=name)(handler)
transform = SafeToolNames(server)
server.add_transform(transform)
listing_started = asyncio.Event()
release_listing = asyncio.Event()
list_tools = server.list_tools
async def delayed_list_tools(*, run_middleware: bool = True):
listing_started.set()
await release_listing.wait()
return await list_tools(run_middleware=run_middleware)
monkeypatch.setattr(server, "list_tools", delayed_list_tools)
first = asyncio.create_task(server.get_tool("demo_one"))
await listing_started.wait()
second = asyncio.create_task(server.get_tool("demo_two"))
await asyncio.sleep(0)
release_listing.set()
first_tool, second_tool = await asyncio.gather(first, second)
assert first_tool is not None
assert second_tool is not None
assert first_tool.name == "demo_one"
assert second_tool.name == "demo_two"
def _server_with_tools(
*names: str,
transform: SafeToolNames | None = None,
tests/wf_sources_mcp/test_runtime.py
+11 -8
View File
@@ -20,10 +20,11 @@ from mcp.types import (
Tool,
server_result_adapter,
)
from pydantic import AnyUrl, TypeAdapter
from pydantic import AnyUrl
from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.raw_messages import RawRequest, RawResult
from wf_sources_mcp.runtime import (
McpRuntimePool,
PersistentMcpSession,
@@ -574,10 +575,11 @@ async def test_persistent_session_invoke_method_client_fallback() -> None:
class _MinimalClient:
async def send_request(
self,
request: ClientRequest,
result_type: type[ServerResult] | TypeAdapter[ServerResult],
) -> ServerResult:
return server_result_adapter.validate_python({"tools": []})
request: RawRequest,
result_type: type[RawResult],
) -> RawResult:
assert request.method == "test.method"
return result_type.model_validate({"extension": True})
session = PersistentMcpSession(
connection=_connection(),
@@ -585,8 +587,8 @@ async def test_persistent_session_invoke_method_client_fallback() -> None:
client=_MinimalClient(), # type: ignore[arg-type, ty:invalid-argument-type]
)
result = await session.invoke_method("tools/list")
assert result["tools"] == []
result = await session.invoke_method("test.method")
assert result["extension"] is True
@pytest.mark.asyncio
@@ -604,9 +606,10 @@ async def test_persistent_session_send_notification_client_fallback() -> None:
client=_MinimalClient(), # type: ignore[arg-type, ty:invalid-argument-type]
)
await session.send_notification("notifications/initialized")
await session.send_notification("test.event")
assert len(sent) == 1
assert sent[0].method == "test.event"
@pytest.mark.asyncio
tests/wf_sources_mcp/test_sdk_adapter.py
+7 -13
View File
@@ -5,11 +5,10 @@ from contextlib import asynccontextmanager
from typing import Any
import pytest
from mcp import GetPromptResult, ReadResourceResult, ServerResult
from mcp import GetPromptResult, ReadResourceResult
from mcp.types import (
CallToolResult,
ClientNotification,
ClientRequest,
ListPromptsResult,
ListResourcesResult,
ListToolsResult,
@@ -18,10 +17,10 @@ from mcp.types import (
TextContent,
Tool,
)
from pydantic import TypeAdapter
from wf_sources_mcp.client import McpSourceClient
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.raw_messages import RawRequest, RawResult
from wf_sources_mcp.sdk import BackendAdapter, McpSdkAdapter
from wf_sources_mcp.transports import StdioSourceTransport
@@ -38,7 +37,7 @@ def _connection() -> McpSourceConnection:
class _FakeSession:
def __init__(self) -> None:
self.notifications: list[ClientNotification] = []
self.requests: list[ClientRequest] = []
self.requests: list[RawRequest] = []
async def list_tools(self) -> ListToolsResult:
return ListToolsResult(
@@ -103,16 +102,11 @@ class _FakeSession:
async def send_request(
self,
request: ClientRequest,
result_type: type[ServerResult] | TypeAdapter[ServerResult],
) -> Any:
assert isinstance(result_type, TypeAdapter)
request: RawRequest,
result_type: type[RawResult],
) -> RawResult:
self.requests.append(request)
return type(
"ServerResultModel",
(),
{"model_dump": lambda _self, **_kwargs: {"ok": True}},
)()
return result_type.model_validate({"ok": True})
async def send_notification(self, notification: ClientNotification) -> None:
self.notifications.append(notification)
tests/wf_sources_mcp/test_source_client.py
+19 -18
View File
@@ -3,11 +3,10 @@ from __future__ import annotations
from typing import Any
import pytest
from mcp import GetPromptResult, ReadResourceResult, ServerResult
from mcp import GetPromptResult, ReadResourceResult
from mcp.types import (
CallToolResult,
ClientNotification,
ClientRequest,
ListPromptsResult,
ListResourcesResult,
ListToolsResult,
@@ -16,10 +15,10 @@ from mcp.types import (
TextContent,
Tool,
)
from pydantic import TypeAdapter
from wf_sources_mcp.client import McpSourceClient
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.raw_messages import RawRequest, RawResult
from wf_sources_mcp.transports import StdioSourceTransport
@@ -34,7 +33,7 @@ def _connection() -> McpSourceConnection:
class _FakeSession:
def __init__(self) -> None:
self.requests: list[ClientRequest] = []
self.requests: list[RawRequest] = []
self.notifications: list[ClientNotification] = []
async def list_tools(self) -> ListToolsResult:
@@ -101,16 +100,11 @@ class _FakeSession:
async def send_request(
self,
request: ClientRequest,
result_type: type[ServerResult] | TypeAdapter[ServerResult],
) -> Any:
assert isinstance(result_type, TypeAdapter)
request: RawRequest,
result_type: type[RawResult],
) -> RawResult:
self.requests.append(request)
return type(
"ClientResultModel",
(),
{"model_dump": lambda _self, **_kwargs: {"ok": True}},
)()
return result_type.model_validate({"ok": True})
async def send_notification(self, notification: ClientNotification) -> None:
self.notifications.append(notification)
@@ -158,16 +152,23 @@ async def test_source_client_reads_resources_and_prompts_as_payloads() -> None:
@pytest.mark.asyncio
async def test_source_client_invokes_methods_and_notifications() -> None:
async def test_source_client_invokes_extension_method() -> None:
session = _FakeSession()
source_client = McpSourceClient(session=session, connection=_connection())
result = await source_client.invoke_method("ping")
await source_client.send_notification("notifications/initialized")
result = await source_client.invoke_method("test.method", {"value": 1})
assert result == {"ok": True}
assert session.requests, "invoke_method should send a request"
assert session.notifications, "send_notification should send a notification"
assert session.requests[0].method == "test.method"
async def test_source_client_sends_extension_notification() -> None:
session = _FakeSession()
source_client = McpSourceClient(session=session, connection=_connection())
await source_client.send_notification("test.event", {"value": 1})
assert session.notifications[0].method == "test.event"
@pytest.mark.asyncio