refactor(runner): unify repeat_prompt and multi-model capture into one path

Every prompt task now fans out into branches (the cross product of prompts x
models) that each capture into an isolated per-branch sink, regardless of which
axis fanned out. After the fan-out the runner:

- projects single-model branch results back into the shared store in branch
  order, preserving the implicit last-tool-result carry-over, run-level results,
  and session snapshot (now deterministic even for async repeat_prompt);
  multi-model tasks are excluded on purpose (no single "last" across models);
- captures the named output uniformly: a plain task publishes its single
  validated value, and any fan-out task (repeat_prompt, multiple models, or the
  cross product) publishes the per-branch fan-in list [{model, item, result}].

This removes the multi_model fork in the capture layer and fixes two
asymmetries: a single-model repeat_prompt can now fan in per-item results (not
just the last), and an async single-model repeat no longer races the shared
store. The capture model is documented in a design note above _aggregate_fanin
and in doc/GRAMMAR.md.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: anticomputer

doc/GRAMMAR.md

@@ -411,13 +411,23 @@ that later tasks consume by name.
 
 Three fields drive this:
 
-- `id` names a task. Its produced value (its final tool result) is exposed to
-  later tasks as `outputs.<id>`.
+- `id` names a task, exposing its output to later tasks as `outputs.<id>`. The
+  shape depends on whether the task fans out:
+  - A plain task (single model, no `repeat_prompt`) publishes its single
+    produced value (its final tool result).
+  - A task that fans out (`repeat_prompt`, multiple `models`, or their cross
+    product) publishes a per-branch fan-in list of
+    `{"model": <label>, "item": <index>, "result": <value>}` records, one per
+    branch. This is uniform across the item axis and the model axis, so a
+    single-model `repeat_prompt` and a multi-model task capture the same way.
 - `outputs` declares an inline JSON Schema (Draft 2020-12). When present, the
   task's value is validated against it before being stored. Validation is strict
   and does not coerce, so a value whose types do not match the contract is a
-  failure. A malformed schema is rejected when the taskflow is loaded, before
-  any model calls are made.
+  failure. On a fan-out task the schema is applied to each branch's `result`
+  (a violation is a failed branch under the `completion` policy); on a plain
+  task it is applied to the single value (a violation is a hard failure). A
+  malformed schema is rejected when the taskflow is loaded, before any model
+  calls are made.
 - `over` is an explicit iterable selector for `repeat_prompt`: a Jinja
   expression evaluated against the template context (so it yields a real list,
   not a re-parsed string).
@@ -473,11 +483,10 @@ Notes:
   `values`, ...) resolve to your data, not the method.
 - Without `id`/`outputs`/`over`, the implicit last-tool-result `repeat_prompt`
   behaviour is unchanged.
-- On multi-model tasks, `id` fans in per-branch results as a list of
-  `{model, item, result}` records (see "Multiple Models"). When an `outputs`
-  schema is present it is applied per branch: each `result` is the validated
-  value, and a branch that violates the schema is treated as a failed branch
-  under the task's `completion` policy.
+- Implicit carry-over (the next task's `repeat_prompt` reading the previous
+  task's last tool result) is fed by single-model tasks only. A multi-model
+  task does not feed it: there is no single "last" result across models, so a
+  downstream task must consume its output by name via `id`/`over`.
 
 ### Toolboxes / MCP Servers
 

src/seclab_taskflow_agent/runner.py

@@ -243,6 +243,42 @@ class _Branch:
     output_set: bool = False
 
 
+# ---------------------------------------------------------------------------
+# Task result/output capture model (unified across repeat_prompt + multi-model)
+#
+# Every prompt task fans out into one or more branches: the cross product of
+# the prompts to run (one for a plain task, N for ``repeat_prompt``) and the
+# resolved models (one for a single-model task, M for a multi-model task). Each
+# branch captures its own tool results into an isolated ``_Branch.sink`` (never
+# the shared store), so concurrent branches never interleave and capture is the
+# same shape regardless of which axis (items, models, or both) fanned out.
+#
+# After the branches run, the runner does two things:
+#
+#  1. Carry-over projection. For single-model tasks it replays the branch sinks
+#     into the shared run ``store`` in branch order, so the legacy implicit
+#     last-tool-result carry-over, the run-level result list, and the session
+#     snapshot behave exactly as before (and deterministically, even for async
+#     ``repeat_prompt``). Multi-model tasks are deliberately excluded: there is
+#     no well-defined "last" result across models, so their results are consumed
+#     by name via ``id``/``over``, not the implicit channel.
+#
+#  2. Named-output capture (``outputs.<id>``), by whether the task fanned out:
+#       - does NOT fan out (plain single-model task) -> the single decoded value
+#         (validated against the ``outputs`` schema here; a violation is a hard
+#         task failure). See :func:`_capture_task_output`.
+#       - fans out (``repeat_prompt`` or multiple models) -> the per-branch
+#         fan-in list ``[{model, item, result}]`` (each ``result`` validated
+#         per branch during the run; a violation is a failed branch reduced by
+#         the ``completion`` policy). See :func:`_aggregate_fanin` /
+#         :func:`_capture_branch_output`.
+#
+# "Fans out" is structural (``repeat_prompt or multi_model``), not a runtime
+# count, so a ``repeat_prompt`` over one item is still a one-element list rather
+# than surprising the author with a scalar.
+# ---------------------------------------------------------------------------
+
+
 def _aggregate_fanin(branches: list[_Branch]) -> list[dict[str, Any]]:
     """Aggregate each branch's final tool result into fan-in records.
 
@@ -980,6 +1016,13 @@ async def on_handoff_hook(context: RunContextWrapper[TContext], agent: Agent[TCo
             toolboxes_override = task.toolboxes or []
             env = task.env or {}
             repeat_prompt = task.repeat_prompt
+            # A task "fans out" when it runs more than one branch by intent:
+            # a repeat_prompt (over items) or a multi-model task (over models),
+            # or their cross product. This drives the unified output-capture
+            # model (see the capture design note above ``_aggregate_fanin``):
+            # a task that fans out yields a per-branch fan-in list, one that
+            # does not yields a single value.
+            fans_out = multi_model or repeat_prompt
             exclude_from_context = task.exclude_from_context
             async_task = task.async_task
             max_concurrent_tasks = task.async_limit
@@ -1053,6 +1096,11 @@ async def on_handoff_hook(context: RunContextWrapper[TContext], agent: Agent[TCo
                     outputs=store.outputs, over=over,
                 )
 
+                # run_prompts hands the branches it built back to the task loop
+                # (via this closure var) so projection + output capture happen
+                # once, after the retry loop, rather than per attempt.
+                captured_branches: list[_Branch] = []
+
                 async def run_prompts(async_task: bool = False, max_concurrent_tasks: int = 5) -> bool:
                     if run:
                         await render_model_output("** 🤖🐚 Executing Shell Task\n")
@@ -1112,31 +1160,26 @@ async def run_prompts(async_task: bool = False, max_concurrent_tasks: int = 5) -
                     concurrency = model_concurrency if multi_model else max_concurrent_tasks
 
                     async def _deploy(branch: _Branch, rm: ResolvedModel) -> bool:
-                        # Multi-model branches capture tool results into their
-                        # own private sink (never the shared store), keeping
-                        # concurrent branches isolated and enabling fan-in.
-                        # Single-model runs keep the shared sink for backwards
-                        # compatibility.
-                        if multi_model:
-                            async def _branch_tool_end(context, agent, tool, result) -> None:  # noqa: ANN001
-                                watchdog_ping()
-                                branch.sink.append(
-                                    normalize_openai_tool_output(result, tool_name=getattr(tool, "name", ""))
-                                )
-
-                            async def _branch_record(tr: ToolResult) -> None:
-                                watchdog_ping()
-                                branch.sink.append(tr)
-
-                            run_hooks = TaskRunHooks(
-                                on_tool_end=_branch_tool_end, on_tool_start=on_tool_start_hook
-                            )
-                            branch_record = _branch_record
-                        else:
-                            run_hooks = TaskRunHooks(
-                                on_tool_end=on_tool_end_hook, on_tool_start=on_tool_start_hook
+                        # Every branch captures its tool results into its own
+                        # isolated sink (never the shared store) so concurrent
+                        # branches never interleave and capture is uniform across
+                        # repeat_prompt and multi-model. Single-model results are
+                        # projected back into the shared store after the fan-out
+                        # (see the projection step in the task loop).
+                        async def _branch_tool_end(context, agent, tool, result) -> None:  # noqa: ANN001
+                            watchdog_ping()
+                            branch.sink.append(
+                                normalize_openai_tool_output(result, tool_name=getattr(tool, "name", ""))
                             )
-                            branch_record = record_tool_result
+
+                        async def _branch_record(tr: ToolResult) -> None:
+                            watchdog_ping()
+                            branch.sink.append(tr)
+
+                        run_hooks = TaskRunHooks(
+                            on_tool_end=_branch_tool_end, on_tool_start=on_tool_start_hook
+                        )
+                        branch_record = _branch_record
                         branch_ok = await deploy_task_agents(
                             available_tools,
                             branch.agents,
@@ -1159,12 +1202,15 @@ async def _branch_record(tr: ToolResult) -> None:
                             agent_hooks=TaskAgentHooks(on_handoff=on_handoff_hook),
                             stream_label=branch.label if multi_model else None,
                         )
-                        # Typed multi-model output: validate this branch's
-                        # result against the task's `outputs` schema. A missing
-                        # or schema-violating result makes it a failed branch,
+                        # Typed fan-out output: when the task fans out (repeat
+                        # and/or multi-model) and declares an `outputs` schema,
+                        # validate this branch's result against it. A missing or
+                        # schema-violating result makes it a failed branch,
                         # folded into the completion policy alongside run
-                        # failures; the validated value feeds fan-in.
-                        if multi_model and task_output_schema:
+                        # failures; the validated value feeds fan-in. (A plain,
+                        # non-fan-out task validates its single value later, in
+                        # the capture step.)
+                        if fans_out and task_output_schema:
                             return _capture_branch_output(
                                 branch,
                                 task_output_schema,
@@ -1181,11 +1227,10 @@ async def _branch_record(tr: ToolResult) -> None:
                         completion_policy=completion_policy,
                     )
 
-                    # Fan-in: expose each branch's final result as a named list
-                    # for downstream tasks (multi-model / cross-product only;
-                    # single-model typed-output capture happens in run_main).
-                    if multi_model and task_output_id:
-                        store.set_output(task_output_id, _aggregate_fanin(branches))
+                    # Hand the branches back to the task loop for projection +
+                    # output capture (done once, after the retry loop).
+                    nonlocal captured_branches
+                    captured_branches = branches
 
                     return complete
 
@@ -1252,27 +1297,43 @@ async def _branch_record(tr: ToolResult) -> None:
                     # --resume.
                     raise RuntimeError(f"Required task {task_name!r} did not complete")
 
-                # Capture this task's typed named output (M2). The task's
-                # output is its final tool result; when an ``outputs`` schema
-                # is declared it is validated before being stored under
-                # ``outputs.<id>`` for downstream tasks to consume by name. A
-                # declared output contract that the produced value violates is
-                # a hard failure, surfaced with the same session-saved / resume
-                # messaging as other task failures.
-                if task_output_id and not multi_model:
-                    try:
-                        _capture_task_output(
-                            store, task_output_id, task_output_schema, task_name
-                        )
-                    except Exception as exc:
-                        logging.error("Task %r output capture failed: %s", task_name, exc)
-                        session.mark_failed(f"Task {task_name!r} output capture failed: {exc}")
-                        await render_model_output(
-                            f"** 🤖❗ Output capture failed: {exc}\n"
-                            f"** 🤖💾 Session saved: {session.session_id}\n"
-                            f"** 🤖💡 Resume with: --resume {session.session_id}\n"
-                        )
-                        raise
+                # Projection: replay this task's per-branch tool results into the
+                # shared run store (single-model only) so the implicit
+                # last-tool-result carry-over, the run-level result list, and the
+                # session snapshot behave as before, and deterministically in
+                # branch order even for async repeat_prompt. Multi-model tasks
+                # are excluded on purpose (no well-defined "last" across models;
+                # consume them by name via id/over). Shell tasks have no branches
+                # and already wrote their result to the store directly.
+                if not multi_model:
+                    for _branch in captured_branches:
+                        for _tr in _branch.sink:
+                            store.record(_tr)
+
+                # Capture this task's typed named output, uniform across single-
+                # and multi-model (see the capture design note above
+                # ``_aggregate_fanin``). A task that fans out (repeat_prompt or
+                # multiple models) publishes the per-branch fan-in list; a plain
+                # task publishes its single value, validated against the
+                # ``outputs`` schema here (a violation is a hard failure, with
+                # the same session-saved / resume messaging as other failures).
+                if task_output_id:
+                    if fans_out:
+                        store.set_output(task_output_id, _aggregate_fanin(captured_branches))
+                    else:
+                        try:
+                            _capture_task_output(
+                                store, task_output_id, task_output_schema, task_name
+                            )
+                        except Exception as exc:
+                            logging.error("Task %r output capture failed: %s", task_name, exc)
+                            session.mark_failed(f"Task {task_name!r} output capture failed: {exc}")
+                            await render_model_output(
+                                f"** 🤖❗ Output capture failed: {exc}\n"
+                                f"** 🤖💾 Session saved: {session.session_id}\n"
+                                f"** 🤖💡 Resume with: --resume {session.session_id}\n"
+                            )
+                            raise
 
                 # Checkpoint after task (must_complete failures break above
                 # without advancing the resume cursor)

tests/test_runner_integration.py

@@ -444,3 +444,73 @@ async def failing_deploy(available_tools, agents, prompt, *, model=None, record_
         data = _session_data(tmp_path)
         assert data["finished"] is False
         assert "did not complete" in data["error"]
+
+
+class TestUnifiedCapture:
+    """The unified capture model: repeat_prompt and multi-model both fan in,
+    plain tasks stay scalar, and single-model results still feed carry-over."""
+
+    def test_single_model_repeat_fans_in_per_item(self, monkeypatch, tmp_path):
+        async def deploy(available_tools, agents, prompt, *, model=None, record_tool_result=None, **kw):
+            if record_tool_result is not None:
+                await record_tool_result(ToolResult(text=json.dumps({"echo": prompt})))
+            return True
+
+        task = TaskDefinition(
+            id="items", agents=["pkg.p"], user_prompt="do {{ result }}",
+            repeat_prompt=True, over="globals.xs",
+        )
+        _run(monkeypatch, tmp_path, _taskflow(task, globals_={"xs": ["a", "b", "c"]}), deploy_impl=deploy)
+        records = _session_outputs(tmp_path)["items"]
+        # A per-item fan-in list (previously only the last item was captured).
+        assert [r["item"] for r in records] == [0, 1, 2]
+        assert records[0]["result"] == {"echo": "do a"}
+        assert records[2]["result"] == {"echo": "do c"}
+
+    def test_plain_single_task_output_is_scalar(self, monkeypatch, tmp_path):
+        async def deploy(available_tools, agents, prompt, *, model=None, record_tool_result=None, **kw):
+            if record_tool_result is not None:
+                await record_tool_result(ToolResult(text=json.dumps({"answer": 42})))
+            return True
+
+        task = TaskDefinition(id="v", agents=["pkg.p"], user_prompt="hi")
+        _run(monkeypatch, tmp_path, _taskflow(task), deploy_impl=deploy)
+        # A task that does not fan out publishes its single value, not a list.
+        assert _session_outputs(tmp_path)["v"] == {"answer": 42}
+
+    def test_single_model_result_feeds_implicit_carryover(self, monkeypatch, tmp_path):
+        calls: list[str] = []
+
+        async def deploy(available_tools, agents, prompt, *, model=None, record_tool_result=None, **kw):
+            calls.append(prompt)
+            if record_tool_result is not None:
+                payload = ["x", "y"] if prompt.startswith("produce") else {"echo": prompt}
+                await record_tool_result(ToolResult(text=json.dumps(payload)))
+            return True
+
+        tasks = [
+            TaskDefinition(agents=["pkg.p"], user_prompt="produce the list"),
+            TaskDefinition(agents=["pkg.p"], user_prompt="handle {{ result }}", repeat_prompt=True),
+        ]
+        _run(monkeypatch, tmp_path, _taskflow_tasks(tasks), deploy_impl=deploy)
+        # Task B repeated once per element of Task A's carried-over result, which
+        # only works if Task A's branch result was projected into the store.
+        assert [c for c in calls if c.startswith("handle")] == ["handle x", "handle y"]
+
+    def test_async_repeat_fanin_is_ordered_despite_completion_order(self, monkeypatch, tmp_path):
+        async def deploy(available_tools, agents, prompt, *, model=None, record_tool_result=None, **kw):
+            # Earlier items sleep longer so they finish last; the fan-in must
+            # still be in item (submission) order, not completion order.
+            n = int(prompt.split()[-1])
+            await asyncio.sleep((3 - n) * 0.01)
+            if record_tool_result is not None:
+                await record_tool_result(ToolResult(text=json.dumps({"n": n})))
+            return True
+
+        task = TaskDefinition(
+            id="items", agents=["pkg.p"], user_prompt="item {{ result }}",
+            repeat_prompt=True, over="globals.xs", async_task=True, async_limit=3,
+        )
+        _run(monkeypatch, tmp_path, _taskflow(task, globals_={"xs": [0, 1, 2]}), deploy_impl=deploy)
+        records = _session_outputs(tmp_path)["items"]
+        assert [r["result"]["n"] for r in records] == [0, 1, 2]