From ff34d781b65c6e981ff2544000a8b552e4c804d7 Mon Sep 17 00:00:00 2001
From: Edward Linscott <edwardlinscott@gmail.com>
Date: Thu, 2 Jul 2026 13:56:53 +0200
Subject: [PATCH] Adopt node-graph's GraphTaskHandle for @task.graph decorators

node-graph #150 made ``build`` graph-only by handle type, removing it
from ``BaseHandle``. The graph decorator still returned a plain
``TaskHandle``, so every ``@task.graph`` handle lost
``build``/``run``/``run_get_graph``/``submit`` when paired with
node-graph > 0.6.5.

Mirror node-graph's own decorator: the graph decorator now returns a
``GraphTaskHandle`` carrying the four graph-materializing methods.
Plain task handles drop them - they only ever raised there, and plain
handles keep node-graph's direct-execute ``run`` semantics.

Requires node-graph with #150 (unreleased at time of writing).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 src/aiida_workgraph/decorator.py |  3 ++-
 src/aiida_workgraph/task.py      | 29 +++++++++++++++++++++++++----
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/aiida_workgraph/decorator.py b/src/aiida_workgraph/decorator.py
index 7627d553..a3eef220 100644
--- a/src/aiida_workgraph/decorator.py
+++ b/src/aiida_workgraph/decorator.py
@@ -178,8 +178,9 @@ def decorator_graph(
 
         def decorator(func) -> TaskHandle:
             from aiida_workgraph.tasks.graph_task import _build_graph_task_taskspec
+            from aiida_workgraph.task import GraphTaskHandle
 
-            handle = TaskHandle(
+            handle = GraphTaskHandle(
                 _build_graph_task_taskspec(
                     func,
                     identifier=identifier,
diff --git a/src/aiida_workgraph/task.py b/src/aiida_workgraph/task.py
index 71cdcfa8..227544e5 100644
--- a/src/aiida_workgraph/task.py
+++ b/src/aiida_workgraph/task.py
@@ -11,7 +11,7 @@
 from node_graph.task_spec import TaskSpec
 
 if TYPE_CHECKING:
-    pass
+    from aiida_workgraph.workgraph import WorkGraph
 
 
 class Task(GraphTask):
@@ -231,15 +231,36 @@ def __call__(self, *args, **kwargs):
 
         return outputs
 
-    def run(self, /, *args, **kwargs):
+
+class GraphTaskHandle(TaskHandle):
+    """Handle for ``@task.graph`` callables.
+
+    node-graph #150 made ``build`` graph-only by handle type (it moved from
+    ``BaseHandle`` to node-graph's ``GraphTaskHandle``). Mirror that here:
+    ``build``/``run``/``run_get_graph``/``submit`` all materialize a WorkGraph,
+    which only makes sense for graph specs — on plain task handles they only
+    ever raised.
+
+    ``build`` delegates to node-graph's implementation unbound rather than
+    inheriting it, because node-graph's ``GraphTaskHandle.__init__`` takes
+    ``spec`` only, while ours must also wire ``get_current_graph`` and
+    ``graph_class=WorkGraph``.
+    """
+
+    def build(self, /, *args: Any, **kwargs: Any) -> WorkGraph:
+        from node_graph.task_spec import GraphTaskHandle as _NGGraphTaskHandle
+
+        return _NGGraphTaskHandle.build(self, *args, **kwargs)
+
+    def run(self, /, *args: Any, **kwargs: Any) -> Any:
         graph = self.build(*args, **kwargs)
         return graph.run()
 
-    def run_get_graph(self, /, *args, **kwargs):
+    def run_get_graph(self, /, *args: Any, **kwargs: Any) -> tuple[Any, WorkGraph]:
         graph = self.build(*args, **kwargs)
         return graph.run(), graph
 
-    def submit(self, /, *args, **kwargs):
+    def submit(self, /, *args: Any, **kwargs: Any) -> WorkGraph:
         graph = self.build(*args, **kwargs)
         graph.submit()
         return graph