feat(trace): simplify state handling logic

Author: IvanAnishchuk

tests/infra/trace/traced_spec.py

@@ -111,18 +111,18 @@ def _spec_call_hook(
             args/kwargs: Arguments passed to the function.
         """
         method_name = wrapped.__name__
-        state: STATE | None = None
 
         # A. Prepare arguments: bind to signature and serialize
         bound_args = self._bind_args(wrapped, args, kwargs)
+        state: STATE | None = bound_args.arguments.get("state")
 
         # Process arguments and auto-register any new SSZ objects as artifacts
         serial_params: SERIALIZED_KWARGS = {
             k: self._process_arg(v) for k, v in bound_args.arguments.items()
         }
 
         # B. Identify State object and handle external state mutations
-        self._capture_pre_state(state := bound_args.arguments.get("state"))
+        self._capture_pre_state(state)
 
         # C. Execute the real function with original args/kwargs
         result: RAW_ARGS = wrapped(*args, **kwargs)
@@ -132,8 +132,7 @@ def _spec_call_hook(
         self._record_spec_call_step(method_name, serial_params, serial_result)
 
         # E. Update tracked state if mutated
-        if state is not None:
-            self._capture_post_state(state)
+        self._capture_post_state(state)
 
         return result