Add tests and agent docs for kwargs_type input/output

.ai/modular.md

@@ -102,6 +102,37 @@ class HeliosChunkDenoiseStep(HeliosChunkLoopWrapper):
 
 Note: sub-blocks inside `LoopSequentialPipelineBlocks` receive `(components, block_state, i, t)` for denoise loops or `(components, block_state, k)` for chunk loops.
 
+## Key pattern: `kwargs_type` bags (`denoiser_input_fields`)
+
+The conditioning inputs a denoiser needs often vary by workflow — especially for omni models like Cosmos3, where the action workflow requires additional action conditioning, and a workflow that generates sound along with video requires additional sound inputs. Tag these outputs with `kwargs_type="denoiser_input_fields"` when they are written, and have the denoiser declare the bag once and receive everything tagged — this avoids creating a new denoiser block for each workflow just to list its specific inputs:
+
+```python
+# producer side: standard conditioning outputs already carry the tag via their templates
+OutputParam.template("prompt_embeds")  # kwargs_type="denoiser_input_fields"
+# workflow-specific fields declare it explicitly
+OutputParam(
+    "action_embeds",
+    kwargs_type="denoiser_input_fields",
+    type_hint=torch.Tensor,
+    description="Action conditioning fed into the transformer.",
+)
+
+# consumer side (the loop denoiser): declare the bag once
+InputParam.template("denoiser_input_fields")
+
+# inside the denoiser __call__: every tagged value arrives in one dict —
+# and also individually (block_state.prompt_embeds, block_state.action_embeds, ...)
+block_state.denoiser_input_fields  # {"prompt_embeds": ..., "action_embeds": ...}
+```
+
+The denoiser typically filters the bag against the transformer's forward signature and forwards the matches — so a new block can add conditioning just by tagging its output (no change to the denoiser), and tagged fields the transformer doesn't accept are silently ignored (see `qwenimage/denoise.py` or `helios/denoise.py`; `z_image/denoise.py` is a minimal bag consumer).
+
+How the tagging works (behavior is pinned down in `tests/modular_pipelines/test_modular_pipelines_custom_blocks.py::TestBlockKwargsTypeInputs`):
+
+- A value gets its tag when it is **written** to pipeline state: a block output is tagged if declared with `OutputParam(..., kwargs_type=...)`; a user-passed input is tagged if the pipeline-level `InputParam` it matches declares a kwargs_type.
+- Users can always pass the whole bag as a dict under the kwargs_type name — `pipe(denoiser_input_fields={"prompt_embeds": ...})` — and every entry gets tagged.
+- **Gotcha — standalone runs:** a named input declared *without* the kwargs_type lands in state by name but never reaches the bag. So when a denoise block runs standalone (without the upstream blocks whose tagged outputs normally fill the bag), passing those values as plain named inputs silently does nothing — they must go through the `denoiser_input_fields={...}` dict, or the block must declare them as named `InputParam(..., kwargs_type="denoiser_input_fields")` inputs.
+
 ## Key pattern: Workflow selection
 
 ```python

tests/modular_pipelines/test_modular_pipelines_custom_blocks.py

@@ -578,6 +578,110 @@ def test_loop_block_requirements_save_load(self, tmp_path):
         assert expected_requirements == config["requirements"]
 
 
+class DummyKwargsProducerStep(ModularPipelineBlocks):
+    """Takes `a` and `b` as regular named inputs and passes them through (with a `-producer`
+    suffix so tests can see the values went through this block), writing `a` back as an output
+    tagged with kwargs_type `typea` and `b` as an untagged output."""
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [InputParam(name="a", default=None), InputParam(name="b", default=None)]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [OutputParam("a", kwargs_type="typea"), OutputParam("b")]
+
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        block_state.a = f"{block_state.a}-producer"
+        block_state.b = f"{block_state.b}-producer"
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class DummyKwargsConsumerStep(ModularPipelineBlocks):
+    """Consumes the `typea` bag and the named input `b`, and verifies what it received against
+    the expected values passed as inputs (`expected_a`, `expected_b`, `expected_typea`)."""
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(kwargs_type="typea"),
+            InputParam(name="b", default=None),
+            InputParam(name="expected_a", default=None),
+            InputParam(name="expected_b", default=None),
+            InputParam(name="expected_typea", default=None),
+        ]
+
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        assert block_state.typea == block_state.expected_typea
+        assert block_state.b == block_state.expected_b
+        # values delivered through the bag are also set individually on block_state;
+        # `a` only exists here if it was delivered through the bag
+        if block_state.expected_a is None:
+            assert not hasattr(block_state, "a")
+        else:
+            assert block_state.a == block_state.expected_a
+        return components, state
+
+
+class TestBlockKwargsTypeInputs:
+    """Test how `kwargs_type` fields flow from user inputs and block outputs to consumer blocks.
+
+    This is the mechanism behind `denoiser_input_fields`: a block declaring
+    `InputParam(kwargs_type=...)` receives a dict of every state value *tagged* with that
+    kwargs_type. A value gets its tag when it is written to the pipeline state: an output
+    written by a block is tagged if the block declared it with
+    `OutputParam(..., kwargs_type=...)`, and a user-passed input is tagged if the pipeline's
+    `InputParam` for it declares a kwargs_type. A named input declared *without* a kwargs_type
+    therefore never reaches the bag, even though it is available in state by name.
+    """
+
+    def test_tagged_block_outputs_are_delivered_to_consumer(self):
+        blocks = SequentialPipelineBlocks.from_blocks_dict(
+            {"producer": DummyKwargsProducerStep(), "consumer": DummyKwargsConsumerStep()}
+        )
+        pipe = blocks.init_pipeline()
+
+        # `a` goes through the producer and is written back as a tagged output, so it reaches
+        # the consumer through the bag; `b` also goes through the producer, but its output is
+        # untagged: it reaches the consumer only as the named input, never through the bag
+        pipe(
+            a="testa",
+            b="testb",
+            expected_a="testa-producer",
+            expected_b="testb-producer",
+            expected_typea={"a": "testa-producer"},
+        )
+
+    def test_user_inputs_passed_by_name_do_not_reach_the_bag(self):
+        pipe = DummyKwargsConsumerStep().init_pipeline()
+
+        # the consumer only knows `a` through the `typea` bag: passing it by name does not
+        # reach the block at all. `b` is declared by name without a kwargs_type: it reaches
+        # the block as the named input, but never through the bag.
+        pipe(
+            a="testa",
+            b="testb",
+            expected_a=None,
+            expected_b="testb",
+            expected_typea={},
+        )
+
+    def test_kwargs_type_dict_input_is_delivered(self):
+        pipe = DummyKwargsConsumerStep().init_pipeline()
+
+        # the whole bag can be passed as a dict under the kwargs_type name: every entry is
+        # tagged individually, so `a` now reaches the consumer through the bag
+        pipe(
+            typea={"a": "testa"},
+            expected_a="testa",
+            expected_b=None,
+            expected_typea={"a": "testa"},
+        )
+
+
 @slow
 @nightly
 @require_torch