Normalize flow outputs and retain dynamic subflow jobs

src/jobflow/core/flow.py

@@ -931,6 +931,22 @@ def get_flow(
     return flow
 
 
+def _normalize_flow_output(value):
+    """Replace Jobs and Flows in nested containers with their outputs."""
+    if isinstance(value, (jobflow.Job, jobflow.Flow)):
+        return value.output
+    if isinstance(value, list):
+        return [_normalize_flow_output(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_normalize_flow_output(item) for item in value)
+    if isinstance(value, dict):
+        return {
+            _normalize_flow_output(key): _normalize_flow_output(item)
+            for key, item in value.items()
+        }
+    return value
+
+
 class DecoratedFlow(Flow):
     """A DecoratedFlow is a Flow that is returned on using the @flow decorator."""
 
@@ -973,7 +989,7 @@ def __init__(self, fn, *args, **kwargs):
                 f"of your @flow decorated function.",
                 stacklevel=2,
             )
-            output = output.output
+        output = _normalize_flow_output(output)
 
         super().__init__(name=name, jobs=children_list, output=output)
 

src/jobflow/core/job.py

@@ -604,7 +604,12 @@ def run(self, store: jobflow.JobStore, job_dir: Path = None) -> Response:
         from datetime import datetime
 
         from jobflow import CURRENT_JOB
-        from jobflow.core.flow import get_flow
+        from jobflow.core.flow import (
+            Flow,
+            _normalize_flow_output,
+            flow_build_context,
+            get_flow,
+        )
         from jobflow.core.schemas import JobStoreDocument
 
         index_str = f", {self.index}" if self.index != 1 else ""
@@ -626,7 +631,13 @@ def run(self, store: jobflow.JobStore, job_dir: Path = None) -> Response:
         if bound is not None and not isinstance(bound, types.ModuleType):
             function = types.MethodType(function, bound)
 
-        response = function(*self.function_args, **self.function_kwargs)
+        children = []
+        with flow_build_context(children):
+            response = function(*self.function_args, **self.function_kwargs)
+
+        children = [child for child in children if child.host is None]
+        if children and not isinstance(response, Response):
+            response = Flow(jobs=children, output=_normalize_flow_output(response))
         response = Response.from_job_returns(
             response, self.output_schema, job_dir=job_dir
         )

tests/core/test_flow_decorator.py

@@ -159,6 +159,19 @@ def my_flow(a, b):
     assert result[flow1.output.uuid][1].output == 7
 
 
+def test_flow_normalizes_nested_job_outputs():
+    """Test Jobs nested in decorated flow outputs resolve to references."""
+    from jobflow import flow
+
+    @flow
+    def my_flow():
+        return {"nested": (add(1, 2),)}
+
+    flow1 = my_flow()
+
+    assert flow1.output == {"nested": (flow1.jobs[0].output,)}
+
+
 def test_flow_returns_list():
     """Test that a flow that returns a list of OutputReferences
     can be created and run."""

tests/core/test_job.py

@@ -83,7 +83,6 @@ def test_job_run(capsys, memory_jobstore, memory_data_jobstore):
     response = test_job.run(memory_jobstore)
     assert capsys.readouterr().out == "I am a job\n"
     assert isinstance(response, Response)
-
     # test run with outputs
     test_job = Job(add, function_args=(1,), function_kwargs={"b": 2})
     response = test_job.run(memory_jobstore)
@@ -1447,3 +1446,23 @@ def add_configured(a, b):
         f"Expected job2.config.manager_config to be {{'key': 'original'}}, "
         f"but got {job2.config.manager_config!r} — shared instance bug confirmed."
     )
+
+
+def test_job_collects_dynamic_subflow(memory_jobstore):
+    """Test a job retains every job created by a returned dynamic subflow."""
+    from jobflow import Flow, job
+
+    @job
+    def add_job(a, b):
+        return a + b
+
+    @job
+    def dynamic_subflow():
+        first = add_job(1, 2)
+        return add_job(first.output, 3)
+
+    response = dynamic_subflow().run(memory_jobstore)
+
+    assert isinstance(response.replace, Flow)
+    assert len(response.replace) == 3
+    assert response.replace.output.uuid == response.replace[-2].uuid