From 5a14ecfcb27948fc94bb3a04bbb0127548eb335b Mon Sep 17 00:00:00 2001
From: forthfate <forthfate@naver.com>
Date: Sat, 12 Sep 2026 00:35:28 +0900
Subject: [PATCH] Trace graph steps automatically

---
 backend/orbit_sdk.py            | 21 ++++++++++++++++++++-
 backend/tests/test_orbit_sdk.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/backend/orbit_sdk.py b/backend/orbit_sdk.py
index dc76383..f7958b2 100644
--- a/backend/orbit_sdk.py
+++ b/backend/orbit_sdk.py
@@ -17,6 +17,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
+from functools import wraps
 from pathlib import Path
 from typing import Any, Callable, Literal
 
@@ -107,7 +108,21 @@ def register(handler: Callable[..., Any]) -> Callable[..., Any]:
             )
             self._nodes[node_id] = node
             setattr(handler, "__orbit_graph_node__", node)
-            return handler
+
+            @wraps(handler)
+            def instrumented(*args: Any, **kwargs: Any) -> Any:
+                context = next(
+                    (value for value in (*args, *kwargs.values()) if isinstance(value, RunnerContext)),
+                    None,
+                )
+                if context is None:
+                    return handler(*args, **kwargs)
+                with context.function(node_id):
+                    return handler(*args, **kwargs)
+
+            setattr(instrumented, "__orbit_graph_node__", node)
+            setattr(handler, "__orbit_graph_wrapper__", instrumented)
+            return instrumented
 
         return register
 
@@ -241,6 +256,7 @@ def function(self, function_id: str):
         if not function_id.strip():
             raise ValueError("function_id must not be empty")
         started = datetime.now(UTC)
+        self.log(f"workflow function started: {function_id}")
         self.emit_result(
             {
                 "workflow_functions": [
@@ -255,6 +271,7 @@ def function(self, function_id: str):
         try:
             yield
         except BaseException:
+            self.log(f"workflow function failed: {function_id}")
             self.emit_result(
                 {
                     "workflow_functions": [
@@ -269,6 +286,7 @@ def function(self, function_id: str):
             )
             raise
         else:
+            self.log(f"workflow function succeeded: {function_id}")
             self.emit_result(
                 {
                     "workflow_functions": [
@@ -1536,6 +1554,7 @@ def main(self) -> None:
             os.environ.get("ORBIT_EXECUTION_MODE", "run"),
             int(os.environ.get("ORBIT_LOOP_INDEX", "1")),
         )
+        handler = getattr(handler, "__orbit_graph_wrapper__", handler)
         before = context.git_head()
         try:
             handler(context)
diff --git a/backend/tests/test_orbit_sdk.py b/backend/tests/test_orbit_sdk.py
index aa2b3ee..88c9016 100644
--- a/backend/tests/test_orbit_sdk.py
+++ b/backend/tests/test_orbit_sdk.py
@@ -6,6 +6,7 @@
 from base64 import b64encode
 
 import orbit_sdk as sdk
+import pytest
 
 
 def test_graph_declarations_export_nodes_and_typed_edges():
@@ -62,6 +63,37 @@ def test_function_trace_emits_successful_function_evidence(tmp_path, capsys):
     assert "collect-source-evidence" in capsys.readouterr().out
 
 
+def test_graph_step_automatically_traces_its_execution(tmp_path, capsys):
+    graph = sdk.Graph()
+
+    @graph.step("collect-source-evidence")
+    def collect(ctx) -> None:
+        ctx.log("Collected source evidence")
+
+    collect(context(tmp_path, iteration=1))
+
+    output = capsys.readouterr().out
+    assert "workflow function started: collect-source-evidence" in output
+    assert "workflow function succeeded: collect-source-evidence" in output
+    assert '"status": "running"' in output
+    assert '"status": "succeeded"' in output
+
+
+def test_graph_step_automatically_traces_failures(tmp_path, capsys):
+    graph = sdk.Graph()
+
+    @graph.step("collect-source-evidence")
+    def collect(ctx) -> None:
+        raise RuntimeError("evidence unavailable")
+
+    with pytest.raises(RuntimeError, match="evidence unavailable"):
+        collect(context(tmp_path, iteration=1))
+
+    output = capsys.readouterr().out
+    assert "workflow function failed: collect-source-evidence" in output
+    assert '"status": "failed"' in output
+
+
 def context(project, *, iteration: int, run_id: str = "run-123"):
     return sdk.RunnerContext(
         phase="execute",