fix(s16): wait for parallel agents before closing the journal (#571)

Signed-off-by: Bryce <1557635468@qq.com>

s16_workflow_runtime/README.ja.md

@@ -110,14 +110,33 @@ script は少数の orchestration primitive だけを公開する `ExecutionStat
 
 各 item が同じ stage を独立して通る場合は `pipeline` を使えます。item A が stage 3 にいる間、item B はまだ stage 1 かもしれません。次の処理が前の group の全結果を必要とする場合は `parallel` を使います。
 
+ある分岐が失敗した場合、どちらの primitive も開始済みの他の分岐が終了するまで
+待ってからエラーを送出します。`asyncio.to_thread()` で実行中のモデル要求は、
+待機中の coroutine をキャンセルしても停止しません。各分岐が終了するまで
+journal を開いておくことで、成功した呼び出しの結果を resume 用に保存でき、
+最後の progress event の後に最終 task notification を出せます。
+
+```python
+async def parallel(self, thunks):
+    results = await asyncio.gather(
+        *[thunk() for thunk in thunks], return_exceptions=True
+    )
+    for result in results:
+        if isinstance(result, BaseException):
+            raise result
+    return results
+```
+
 ```python
 async def pipeline(self, items, *stages):
     async def run_item(item, idx):
         value = item
-        for stage in stages:                       # 各 item がすべての stage を独立して完走
+        for stage in stages:
             value = await stage(value, item, idx)
         return value
-    return await asyncio.gather(*[run_item(it, i) for i, it in enumerate(items)])
+    return await self.parallel([
+        lambda it=it, i=i: run_item(it, i) for i, it in enumerate(items)
+    ])
 ```
 
 ## 構造化出力: Subagent に散文を返させない

s16_workflow_runtime/README.md

@@ -110,14 +110,33 @@ A script receives an `ExecutionState` exposing a small set of orchestration prim
 
 Use `pipeline` when each item independently crosses the same stages. Item A may reach stage three while item B is still in stage one. Use `parallel` when the next step needs every result from the preceding group.
 
+If a branch fails, both primitives wait for the other started branches before
+raising the error. A model request running in `asyncio.to_thread()` does not stop
+just because its awaiter is cancelled. Keeping the journal open until the
+branches settle lets successful calls save their results for resume, and keeps
+the final task notification after the last progress event.
+
+```python
+async def parallel(self, thunks):
+    results = await asyncio.gather(
+        *[thunk() for thunk in thunks], return_exceptions=True
+    )
+    for result in results:
+        if isinstance(result, BaseException):
+            raise result
+    return results
+```
+
 ```python
 async def pipeline(self, items, *stages):
     async def run_item(item, idx):
         value = item
-        for stage in stages:                       # Each item independently completes every stage
+        for stage in stages:
             value = await stage(value, item, idx)
         return value
-    return await asyncio.gather(*[run_item(it, i) for i, it in enumerate(items)])
+    return await self.parallel([
+        lambda it=it, i=i: run_item(it, i) for i, it in enumerate(items)
+    ])
 ```
 
 ## Structured Output: Do Not Let Subagents Return Essays

s16_workflow_runtime/README.zh.md

@@ -110,14 +110,32 @@ def validate_meta(meta):
 
 每个 item 都要独立经过相同步骤时，可以使用 `pipeline`。item A 跑到第 3 阶段时，item B 可能还在第 1 阶段；下一步必须同时使用上一阶段全部结果时，再使用 `parallel` 等待所有调用完成。
 
+某个分支失败时，两种原语都会先等待其他已启动的分支结束，再抛出错误。
+通过 `asyncio.to_thread()` 发出的模型请求，不会因为等待它的协程被取消就停止。
+等各分支结束后再关闭 journal，成功调用的结果才能保存并用于 resume，
+最终任务通知也才能出现在最后一条进度事件之后。
+
+```python
+async def parallel(self, thunks):
+    results = await asyncio.gather(
+        *[thunk() for thunk in thunks], return_exceptions=True
+    )
+    for result in results:
+        if isinstance(result, BaseException):
+            raise result
+    return results
+```
+
 ```python
 async def pipeline(self, items, *stages):
     async def run_item(item, idx):
         value = item
-        for stage in stages:                       # 每个 item 独立跑完所有 stage
+        for stage in stages:
             value = await stage(value, item, idx)
         return value
-    return await asyncio.gather(*[run_item(it, i) for i, it in enumerate(items)])
+    return await self.parallel([
+        lambda it=it, i=i: run_item(it, i) for i, it in enumerate(items)
+    ])
 ```
 
 ## 结构化输出：别让子 agent 回来写散文

s16_workflow_runtime/code.py

@@ -513,7 +513,14 @@ class ExecutionState:
 
     async def parallel(self, thunks):
         """BARRIER: run all thunks concurrently and fail if any thunk fails."""
-        return await asyncio.gather(*[thunk() for thunk in thunks])
+        # Keep the journal open until every started branch has settled.
+        results = await asyncio.gather(
+            *[thunk() for thunk in thunks], return_exceptions=True
+        )
+        for result in results:
+            if isinstance(result, BaseException):
+                raise result
+        return results
 
     async def pipeline(self, items, *stages):
         """Per-item staged flow, NO barrier between stages: item A can be in
@@ -524,7 +531,9 @@ class ExecutionState:
             for stage in stages:
                 value = await stage(value, item, idx)
             return value
-        return await asyncio.gather(*[run_item(it, i) for i, it in enumerate(items)])
+        return await self.parallel([
+            lambda it=it, i=i: run_item(it, i) for i, it in enumerate(items)
+        ])
 
     async def workflow(self, name, args=None):
         """Run a saved workflow inline as a child (one level), sharing this run's

tests/test_workflow_goal_lessons.py

@@ -459,6 +459,79 @@ def test_parallel_agent_calls_do_not_block_the_event_loop(tmp_path: Path) -> Non
     assert task.usage == {"agents": 2, "tokens": 2}
 
 
+@pytest.mark.parametrize("primitive", ["parallel", "pipeline"])
+def test_workflow_failure_waits_for_agents_and_preserves_resume_results(
+    primitive: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    workflow = load_lesson(
+        "workflow_failure_drain_test",
+        ROOT / "s16_workflow_runtime" / "code.py",
+    )
+    monkeypatch.setattr(workflow, "STORE", tmp_path)
+    slow_started = threading.Event()
+    failure_raised = threading.Event()
+    release_slow = threading.Event()
+    calls = {"slow": 0, "fail": 0}
+
+    class ControlledRunner:
+        def run(self, prompt, schema=None, label=None):
+            calls[label] += 1
+            if label == "slow":
+                slow_started.set()
+                assert release_slow.wait(timeout=5)
+            elif calls[label] == 1:
+                assert slow_started.wait(timeout=5)
+                failure_raised.set()
+                raise RuntimeError("agent failed")
+            return workflow.RunnerOutput(label, 1)
+
+    monkeypatch.setattr(workflow, "RUNNER_FACTORY", ControlledRunner)
+
+    async def script(ctx, _args):
+        if primitive == "parallel":
+            return await ctx.parallel([
+                lambda: ctx.agent("slow", label="slow"),
+                lambda: ctx.agent("fail", label="fail"),
+            ])
+
+        async def stage(_value, item, _index):
+            return await ctx.agent(item, label=item)
+
+        return await ctx.pipeline(["slow", "fail"], stage)
+
+    async def run():
+        meta = {"name": "failure-drain", "description": "test"}
+        pending = asyncio.create_task(workflow.WorkflowTool().call(meta, script))
+        try:
+            assert await asyncio.to_thread(failure_raised.wait, 5)
+            done, _ = await asyncio.wait({pending}, timeout=0.05)
+            assert not done, "workflow finalized while an agent was still running"
+        finally:
+            release_slow.set()
+            first = await pending
+            # Drain any orphaned children as well when checking the unfixed code.
+            children = asyncio.all_tasks() - {asyncio.current_task()}
+            await asyncio.gather(*children, return_exceptions=True)
+
+        assert first["task"].status == "failed"
+        assert first["result"] == {"error": "agent failed"}
+        run_id = first["task"].run_id
+        snapshot = json.loads((tmp_path / f"{run_id}.json").read_text())
+        assert snapshot["task"]["usage"] == {"agents": 1, "tokens": 1}
+        records = [json.loads(line) for line in
+                   (tmp_path / f"{run_id}.journal.jsonl").read_text().splitlines()]
+        assert [record["value"] for record in records] == ["slow"]
+
+        resumed = await workflow.WorkflowTool().call(
+            meta, script, resume_from_run_id=run_id
+        )
+        assert resumed["task"].status == "completed"
+        assert resumed["result"] == ["slow", "fail"]
+        assert calls == {"slow": 1, "fail": 2}
+
+    asyncio.run(run())
+
+
 def test_workflow_default_entry_extends_the_real_s15_host(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None: