Skip to content

Commit dc4caef

Browse files
committed
fix(workflows-mcp): extract only declared outputs for Workflow block events
_get_block_outputs now checks isinstance(block_output, Execution) and returns only block_output.outputs instead of model_dump() on the entire child execution tree. This prevents 100KB+ block_completed events for Workflow blocks while keeping Shell/LLMCall unchanged.
1 parent 50290d2 commit dc4caef

2 files changed

Lines changed: 147 additions & 1 deletion

File tree

src/workflows_mcp/engine/workflow_runner.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1426,9 +1426,20 @@ def _extract_execution_state(job_result: dict[str, Any]) -> "ExecutionState":
14261426
)
14271427

14281428
def _get_block_outputs(self, block_output: Any) -> dict[str, Any]:
1429-
"""Extract output dict from BlockOutput."""
1429+
"""Extract output dict from BlockOutput.
1430+
1431+
For regular blocks (Shell, LLMCall), block_output is a BlockOutput model
1432+
and model_dump() produces a small dict with stdout/response.
1433+
1434+
For Workflow blocks, block_output is an Execution object. Using model_dump()
1435+
would serialize the ENTIRE child execution tree (all child blocks, metadata,
1436+
nested executions). Instead, we extract only the declared workflow outputs.
1437+
"""
14301438
if not block_output:
14311439
return {}
1440+
# Workflow blocks return Execution — only include declared outputs, not the full tree
1441+
if isinstance(block_output, Execution):
1442+
return block_output.outputs or {}
14321443
return block_output.model_dump()
14331444

14341445

tests/test_get_block_outputs.py

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
"""
2+
Unit tests for WorkflowRunner._get_block_outputs.
3+
4+
Verifies that Workflow blocks (returning Execution objects) emit only
5+
declared outputs, not the full child execution tree via model_dump().
6+
"""
7+
8+
from pydantic import BaseModel, Field
9+
10+
from workflows_mcp.engine.execution import Execution
11+
from workflows_mcp.engine.metadata import Metadata
12+
from workflows_mcp.engine.workflow_runner import WorkflowRunner
13+
14+
15+
class FakeBlockOutput(BaseModel):
16+
"""Mimics ShellOutput/LLMCallOutput for testing."""
17+
18+
stdout: str = Field(default="")
19+
exit_code: int = Field(default=0)
20+
21+
22+
def _make_runner() -> WorkflowRunner:
23+
"""Create a bare WorkflowRunner (no callback needed for _get_block_outputs)."""
24+
return WorkflowRunner()
25+
26+
27+
def test_get_block_outputs_none_returns_empty_dict() -> None:
28+
"""None input returns empty dict."""
29+
runner = _make_runner()
30+
assert runner._get_block_outputs(None) == {}
31+
32+
33+
def test_get_block_outputs_regular_block_uses_model_dump() -> None:
34+
"""Regular BlockOutput (Shell/LLMCall) returns full model_dump()."""
35+
runner = _make_runner()
36+
output = FakeBlockOutput(stdout="hello world", exit_code=0)
37+
result = runner._get_block_outputs(output)
38+
39+
assert result == {"stdout": "hello world", "exit_code": 0}
40+
41+
42+
def test_get_block_outputs_execution_returns_only_declared_outputs() -> None:
43+
"""Execution (Workflow block) returns only declared outputs, not full tree.
44+
45+
This is the core fix: Workflow blocks return Execution objects.
46+
Before the fix, model_dump() serialized the ENTIRE child execution tree
47+
including all child blocks, metadata, and nested executions (~100KB+).
48+
After the fix, only the declared outputs (~few KB) are returned.
49+
"""
50+
runner = _make_runner()
51+
52+
# Simulate a child workflow execution with nested blocks and large content
53+
child_execution = Execution(
54+
inputs={"source_url": "https://example.com/large-doc.pdf"},
55+
outputs={
56+
"document_text": "extracted text (trimmed for test)",
57+
"metadata_json": '{"title": "Test Document"}',
58+
},
59+
blocks={
60+
"extract": Execution(
61+
inputs={"url": "https://example.com/large-doc.pdf"},
62+
outputs={"stdout": "A" * 84000}, # Simulates 84K document text
63+
metadata=Metadata.create_leaf_success(
64+
type="Shell",
65+
id="extract",
66+
duration_ms=1500,
67+
started_at="2026-01-01T00:00:00Z",
68+
wave=0,
69+
execution_order=0,
70+
index=0,
71+
depth=1,
72+
),
73+
),
74+
"metadata": Execution(
75+
inputs={"text": "some text"},
76+
outputs={"response": '{"title": "Test Document"}'},
77+
metadata=Metadata.create_leaf_success(
78+
type="LLMCall",
79+
id="metadata",
80+
duration_ms=2000,
81+
started_at="2026-01-01T00:00:01Z",
82+
wave=1,
83+
execution_order=1,
84+
index=1,
85+
depth=1,
86+
),
87+
),
88+
},
89+
)
90+
91+
result = runner._get_block_outputs(child_execution)
92+
93+
# Should return ONLY declared outputs
94+
assert result == {
95+
"document_text": "extracted text (trimmed for test)",
96+
"metadata_json": '{"title": "Test Document"}',
97+
}
98+
99+
# Crucially: should NOT contain the full child block tree
100+
assert "blocks" not in result
101+
assert "inputs" not in result
102+
assert "metadata" not in result
103+
# And should NOT contain the 84K content from child blocks
104+
assert "A" * 84000 not in str(result)
105+
106+
107+
def test_get_block_outputs_execution_empty_outputs() -> None:
108+
"""Execution with empty outputs returns empty dict."""
109+
runner = _make_runner()
110+
111+
child_execution = Execution(
112+
inputs={"key": "value"},
113+
outputs={},
114+
blocks={"some_block": Execution(inputs={}, outputs={"data": "big"})},
115+
)
116+
117+
result = runner._get_block_outputs(child_execution)
118+
assert result == {}
119+
120+
121+
def test_get_block_outputs_does_not_affect_regular_block_with_extra_fields() -> None:
122+
"""Regular BlockOutput with extra fields (custom outputs) still works."""
123+
124+
class CustomOutput(BaseModel):
125+
model_config = {"extra": "allow"}
126+
stdout: str = ""
127+
exit_code: int = 0
128+
129+
runner = _make_runner()
130+
output = CustomOutput(stdout="ok", exit_code=0, custom_field="custom_value")
131+
result = runner._get_block_outputs(output)
132+
133+
assert result["stdout"] == "ok"
134+
assert result["exit_code"] == 0
135+
assert result["custom_field"] == "custom_value"

0 commit comments

Comments
 (0)