feat(graph-engine): make layer runtime state non-null and bound early (#30552)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

api/core/workflow/graph_engine/graph_engine.py

@@ -212,9 +212,16 @@ class GraphEngine:
             if id(node.graph_runtime_state) != expected_state_id:
                 raise ValueError(f"GraphRuntimeState consistency violation: Node '{node.id}' has a different instance")
 
+    def _bind_layer_context(
+        self,
+        layer: GraphEngineLayer,
+    ) -> None:
+        layer.initialize(ReadOnlyGraphRuntimeStateWrapper(self._graph_runtime_state), self._command_channel)
+
     def layer(self, layer: GraphEngineLayer) -> "GraphEngine":
         """Add a layer for extending functionality."""
         self._layers.append(layer)
+        self._bind_layer_context(layer)
         return self
 
     def run(self) -> Generator[GraphEngineEvent, None, None]:
@@ -301,14 +308,7 @@ class GraphEngine:
     def _initialize_layers(self) -> None:
         """Initialize layers with context."""
         self._event_manager.set_layers(self._layers)
-        # Create a read-only wrapper for the runtime state
-        read_only_state = ReadOnlyGraphRuntimeStateWrapper(self._graph_runtime_state)
         for layer in self._layers:
-            try:
-                layer.initialize(read_only_state, self._command_channel)
-            except Exception as e:
-                logger.warning("Failed to initialize layer %s: %s", layer.__class__.__name__, e)
-
             try:
                 layer.on_graph_start()
             except Exception as e:

api/core/workflow/graph_engine/layers/README.md

@@ -8,7 +8,7 @@ Pluggable middleware for engine extensions.
 
 Abstract base class for layers.
 
-- `initialize()` - Receive runtime context
+- `initialize()` - Receive runtime context (runtime state is bound here and always available to hooks)
 - `on_graph_start()` - Execution start hook
 - `on_event()` - Process all events
 - `on_graph_end()` - Execution end hook
@@ -34,6 +34,9 @@ engine.layer(debug_layer)
 engine.run()
 ```
 
+`engine.layer()` binds the read-only runtime state before execution, so
+`graph_runtime_state` is always available inside layer hooks.
+
 ## Custom Layers
 
 ```python

api/core/workflow/graph_engine/layers/base.py

@@ -13,6 +13,14 @@ from core.workflow.nodes.base.node import Node
 from core.workflow.runtime import ReadOnlyGraphRuntimeState
 
 
+class GraphEngineLayerNotInitializedError(Exception):
+    """Raised when a layer's runtime state is accessed before initialization."""
+
+    def __init__(self, layer_name: str | None = None) -> None:
+        name = layer_name or "GraphEngineLayer"
+        super().__init__(f"{name} runtime state is not initialized. Bind the layer to a GraphEngine before access.")
+
+
 class GraphEngineLayer(ABC):
     """
     Abstract base class for GraphEngine layers.
@@ -28,22 +36,27 @@ class GraphEngineLayer(ABC):
 
     def __init__(self) -> None:
         """Initialize the layer. Subclasses can override with custom parameters."""
-        self.graph_runtime_state: ReadOnlyGraphRuntimeState | None = None
+        self._graph_runtime_state: ReadOnlyGraphRuntimeState | None = None
         self.command_channel: CommandChannel | None = None
 
+    @property
+    def graph_runtime_state(self) -> ReadOnlyGraphRuntimeState:
+        if self._graph_runtime_state is None:
+            raise GraphEngineLayerNotInitializedError(type(self).__name__)
+        return self._graph_runtime_state
+
     def initialize(self, graph_runtime_state: ReadOnlyGraphRuntimeState, command_channel: CommandChannel) -> None:
         """
         Initialize the layer with engine dependencies.
 
-        Called by GraphEngine before execution starts to inject the read-only runtime state
-        and command channel. This allows layers to observe engine context and send
-        commands, but prevents direct state modification.
-
+        Called by GraphEngine to inject the read-only runtime state and command channel.
+        This is invoked when the layer is registered with a `GraphEngine` instance.
+        Implementations should be idempotent.
         Args:
             graph_runtime_state: Read-only view of the runtime state
             command_channel: Channel for sending commands to the engine
         """
-        self.graph_runtime_state = graph_runtime_state
+        self._graph_runtime_state = graph_runtime_state
         self.command_channel = command_channel
 
     @abstractmethod

api/core/workflow/graph_engine/layers/debug_logging.py

@@ -109,10 +109,8 @@ class DebugLoggingLayer(GraphEngineLayer):
         self.logger.info("=" * 80)
         self.logger.info("🚀 GRAPH EXECUTION STARTED")
         self.logger.info("=" * 80)
-
-        if self.graph_runtime_state:
-            # Log initial state
-            self.logger.info("Initial State:")
+        # Log initial state
+        self.logger.info("Initial State:")
 
     @override
     def on_event(self, event: GraphEngineEvent) -> None:
@@ -243,8 +241,7 @@ class DebugLoggingLayer(GraphEngineLayer):
         self.logger.info("  Node retries: %s", self.retry_count)
 
         # Log final state if available
-        if self.graph_runtime_state and self.include_outputs:
-            if self.graph_runtime_state.outputs:
-                self.logger.info("Final outputs: %s", self._format_dict(self.graph_runtime_state.outputs))
+        if self.include_outputs and self.graph_runtime_state.outputs:
+            self.logger.info("Final outputs: %s", self._format_dict(self.graph_runtime_state.outputs))
 
         self.logger.info("=" * 80)

api/core/workflow/graph_engine/layers/persistence.py

@@ -337,8 +337,6 @@ class WorkflowPersistenceLayer(GraphEngineLayer):
         if update_finished:
             execution.finished_at = naive_utc_now()
         runtime_state = self.graph_runtime_state
-        if runtime_state is None:
-            return
         execution.total_tokens = runtime_state.total_tokens
         execution.total_steps = runtime_state.node_run_steps
         execution.outputs = execution.outputs or runtime_state.outputs
@@ -404,6 +402,4 @@ class WorkflowPersistenceLayer(GraphEngineLayer):
 
     def _system_variables(self) -> Mapping[str, Any]:
         runtime_state = self.graph_runtime_state
-        if runtime_state is None:
-            return {}
         return runtime_state.variable_pool.get_by_prefix(SYSTEM_VARIABLE_NODE_ID)