✨ Map: expose gathered zone outputs to the client

Zone/Map tasks have no AiiDA process node (`pk=None`), so
`Task.update_state` could never populate their outputs on the client
side. After `wg.run()`, `map_zone.outputs.<name>` returned an empty
dict even though the engine had the correct gathered results.

Three changes:

- `gather()` creates dynamic output namespaces so the client can assign
  per-prefix keys after the run.

- `update_map_task_state` persists the gathered result node PKs in
  `task_map_info[name]['result_pks']` on the process node.

- `WorkGraph.update()` reads those PKs back and populates the zone
  task's output sockets via `_populate_zone_outputs`.

The regression test checks the gathered outputs both in-session and
after a fresh `WorkGraph.load`, since only the latter goes purely
through the persisted state.

Author: GeigerJ2

src/aiida_workgraph/engine/task_state.py

@@ -3,7 +3,7 @@
 from typing_extensions import assert_never
 from aiida.orm.utils.serialize import serialize
 from aiida_workgraph.orm.utils import deserialize_safe
-from aiida.orm import ProcessNode, Data
+from aiida.orm import Data, Node, ProcessNode
 from aiida_workgraph.enums import TERMINAL_TASK_STATES, RuntimeInfoKey, TaskState
 from node_graph.socket import BaseSocket, TaskSocketNamespace
 
@@ -350,6 +350,24 @@ def update_map_task_state(self, name: str) -> None:
                     default=None,
                 )
             self.ctx._task_results[name][link.to_socket._name] = results
+        # Persist the gathered-result PKs on the process node so the client
+        # can reconstruct zone outputs after wg.update().  Zone tasks have no
+        # process node of their own, so without this the client cannot recover
+        # the per-prefix result nodes.
+        #
+        # Unstored nodes have no pk, so they are skipped rather than persisted.
+        result_pks: dict[str, dict[str, int]] = {}
+        for socket_name, val in self.ctx._task_results[name].items():
+            if socket_name.startswith('_'):
+                continue
+            if isinstance(val, dict):
+                result_pks[socket_name] = {
+                    prefix: node.pk for prefix, node in val.items() if isinstance(node, Node) and node.pk is not None
+                }
+        if result_pks:
+            map_info = self.process.node.get_task_map_info(name) or {}
+            map_info['result_pks'] = result_pks
+            self.set_task_runtime_info(name, 'map_info', map_info)
         self.set_task_runtime_info(name, 'state', TaskState.FINISHED)
         self.process.report(f'Task: {name} finished.')
         self.update_meta_tasks(name)

src/aiida_workgraph/tasks/builtins.py

@@ -138,7 +138,13 @@ def gather(self, sockets: Dict[str, BaseSocket]) -> BaseSocket:
         gather_item = self.gather_item_task
         for name in sockets:
             gather_item.add_input_spec('workgraph.any', name=name)
-            self.add_output_spec('workgraph.namespace', name=name)
+            # The gathered output namespace must be dynamic so the client
+            # can populate it with per-prefix keys after the run completes.
+            self.add_output_spec(
+                'workgraph.namespace',
+                name=name,
+                meta=SocketMeta(dynamic=True),
+            )
         gather_item.set_inputs(sockets)
         return gather_item.outputs
 

src/aiida_workgraph/workgraph.py

@@ -323,13 +323,45 @@ def update(self) -> None:
                 continue
             self.tasks[name].update_state(data)
 
+        # Zone/Map tasks have no process node (pk=None), so update_state
+        # cannot populate their outputs.  The engine persists the gathered
+        # result PKs in task_map_info; reconstruct outputs from those.
+        self._populate_zone_outputs(processes_data)
+
         if self.widget is not None:
             states = {name: data['state'] for name, data in processes_data.items()}
             self.widget.states = states
 
         if self.process.is_finished_ok:
             self.outputs._set_socket_value(resolve_node_link_managers(self.process.outputs))
 
+    def _populate_zone_outputs(self, processes_data: Dict[str, Any]) -> None:
+        """Populate outputs for zone/Map tasks from persisted result PKs.
+
+        Zone tasks have no AiiDA process node, so ``Task.update_state`` cannot
+        load their outputs.  The engine persists the gathered result node PKs in
+        ``task_map_info[name]['result_pks']``; this method loads those nodes and
+        sets them on the corresponding output sockets.
+        """
+        import aiida.orm
+
+        for name, data in processes_data.items():
+            if name not in self.tasks:
+                continue
+            # Only handle tasks with no process node that are finished.
+            if data['pk'] is not None or data['state'] != TaskState.FINISHED:
+                continue
+            map_info = self.process.get_task_map_info(name)
+            if not map_info:
+                continue
+            result_pks = map_info.get('result_pks', {})
+            task = self.tasks[name]
+            for socket_name, pk_map in result_pks.items():
+                if socket_name not in task.outputs._sockets:
+                    continue
+                values = {prefix: aiida.orm.load_node(pk) for prefix, pk in pk_map.items()}
+                task.outputs[socket_name]._value = values
+
     @property
     def pk(self) -> Optional[int]:
         return self.process.pk if self.process else None

tests/test_map.py

@@ -127,3 +127,31 @@ def test_map_zone_failed_iteration_fails_the_zone():
     assert map_zone.state == 'FAILED'
     assert wg.process.exit_status == 302
     assert 'key_1_maybe_fail' in wg.process.exit_message
+    # A failed zone must expose nothing rather than a short namespace.
+    assert map_zone.outputs.sum1._value == {}
+
+
+def test_map_zone_outputs_visible_to_client():
+    """Gathered zone outputs are readable from the client after the run.
+
+    A Map zone has no process node of its own, so ``Task.update_state`` cannot
+    populate its outputs and ``map_zone.outputs.<name>`` used to come back empty
+    even though the engine had gathered the results. The engine now persists the
+    gathered node PKs in ``task_map_info``, which the client reads back. Checked
+    both in-session and via a fresh ``WorkGraph.load``, since only the latter
+    goes purely through the persisted state.
+    """
+    n = 3
+    with WorkGraph('map_zone_outputs') as wg:
+        data = generate_data(n=n).data
+        with Map(data) as map_zone:
+            out1 = add(x=map_zone.value, y=1).result
+            map_zone.gather({'sum1': out1})
+        wg.run()
+
+    expected = {'key_0': 1, 'key_1': 2, 'key_2': 3}
+    gathered = map_zone.outputs.sum1._value
+    assert {prefix: node.value for prefix, node in gathered.items()} == expected
+
+    reloaded = WorkGraph.load(wg.process.pk).tasks[map_zone.name].outputs.sum1._value
+    assert {prefix: node.value for prefix, node in reloaded.items()} == expected