Add support for tracing script execution (#48276)

* Add support for tracing script execution

* Tweak

homeassistant/components/automation/trace.py

@@ -2,148 +2,25 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
-import datetime as dt
-from itertools import count
-from typing import Any, Deque
 
-from homeassistant.components.trace.const import DATA_TRACE, STORED_TRACES
-from homeassistant.components.trace.utils import LimitedSizeDict
-from homeassistant.core import Context
-from homeassistant.helpers.trace import TraceElement, trace_id_set
-from homeassistant.util import dt as dt_util
+from homeassistant.components.trace import AutomationTrace, async_store_trace
 
 # mypy: allow-untyped-calls, allow-untyped-defs
 # mypy: no-check-untyped-defs, no-warn-return-any
 
 
-class AutomationTrace:
-    """Container for automation trace."""
-
-    _run_ids = count(0)
-
-    def __init__(
-        self,
-        key: tuple[str, str],
-        config: dict[str, Any],
-        context: Context,
-    ):
-        """Container for automation trace."""
-        self._action_trace: dict[str, Deque[TraceElement]] | None = None
-        self._condition_trace: dict[str, Deque[TraceElement]] | None = None
-        self._config: dict[str, Any] = config
-        self.context: Context = context
-        self._error: Exception | None = None
-        self._state: str = "running"
-        self.run_id: str = str(next(self._run_ids))
-        self._timestamp_finish: dt.datetime | None = None
-        self._timestamp_start: dt.datetime = dt_util.utcnow()
-        self._key: tuple[str, str] = key
-        self._variables: dict[str, Any] | None = None
-
-    def set_action_trace(self, trace: dict[str, Deque[TraceElement]]) -> None:
-        """Set action trace."""
-        self._action_trace = trace
-
-    def set_condition_trace(self, trace: dict[str, Deque[TraceElement]]) -> None:
-        """Set condition trace."""
-        self._condition_trace = trace
-
-    def set_error(self, ex: Exception) -> None:
-        """Set error."""
-        self._error = ex
-
-    def set_variables(self, variables: dict[str, Any]) -> None:
-        """Set variables."""
-        self._variables = variables
-
-    def finished(self) -> None:
-        """Set finish time."""
-        self._timestamp_finish = dt_util.utcnow()
-        self._state = "stopped"
-
-    def as_dict(self) -> dict[str, Any]:
-        """Return dictionary version of this AutomationTrace."""
-
-        result = self.as_short_dict()
-
-        action_traces = {}
-        condition_traces = {}
-        if self._action_trace:
-            for key, trace_list in self._action_trace.items():
-                action_traces[key] = [item.as_dict() for item in trace_list]
-
-        if self._condition_trace:
-            for key, trace_list in self._condition_trace.items():
-                condition_traces[key] = [item.as_dict() for item in trace_list]
-
-        result.update(
-            {
-                "action_trace": action_traces,
-                "condition_trace": condition_traces,
-                "config": self._config,
-                "context": self.context,
-                "variables": self._variables,
-            }
-        )
-        if self._error is not None:
-            result["error"] = str(self._error)
-        return result
-
-    def as_short_dict(self) -> dict[str, Any]:
-        """Return a brief dictionary version of this AutomationTrace."""
-
-        last_action = None
-        last_condition = None
-        trigger = None
-
-        if self._action_trace:
-            last_action = list(self._action_trace)[-1]
-        if self._condition_trace:
-            last_condition = list(self._condition_trace)[-1]
-        if self._variables:
-            trigger = self._variables.get("trigger", {}).get("description")
-
-        result = {
-            "last_action": last_action,
-            "last_condition": last_condition,
-            "run_id": self.run_id,
-            "state": self._state,
-            "timestamp": {
-                "start": self._timestamp_start,
-                "finish": self._timestamp_finish,
-            },
-            "trigger": trigger,
-            "domain": self._key[0],
-            "item_id": self._key[1],
-        }
-        if self._error is not None:
-            result["error"] = str(self._error)
-        if last_action is not None:
-            result["last_action"] = last_action
-            result["last_condition"] = last_condition
-
-        return result
-
-
 @contextmanager
 def trace_automation(hass, item_id, config, context):
-    """Trace action execution of automation with automation_id."""
-    key = ("automation", item_id)
-    trace = AutomationTrace(key, config, context)
-    trace_id_set((key, trace.run_id))
-
-    if key:
-        traces = hass.data[DATA_TRACE]
-        if key not in traces:
-            traces[key] = LimitedSizeDict(size_limit=STORED_TRACES)
-        traces[key][trace.run_id] = trace
+    """Trace action execution of automation with item_id."""
+    trace = AutomationTrace(item_id, config, context)
+    async_store_trace(hass, trace)
 
     try:
         yield trace
     except Exception as ex:  # pylint: disable=broad-except
-        if key:
+        if item_id:
             trace.set_error(ex)
         raise ex
     finally:
-        if key:
+        if item_id:
             trace.finished()

homeassistant/components/script/__init__.py

@@ -36,8 +36,11 @@ from homeassistant.helpers.script import (
     make_script_schema,
 )
 from homeassistant.helpers.service import async_set_service_schema
+from homeassistant.helpers.trace import trace_get, trace_path
 from homeassistant.loader import bind_hass
 
+from .trace import trace_script
+
 _LOGGER = logging.getLogger(__name__)
 
 DOMAIN = "script"
@@ -221,7 +224,7 @@ async def _async_process_config(hass, config, component):
         )
 
     script_entities = [
-        ScriptEntity(hass, object_id, cfg)
+        ScriptEntity(hass, object_id, cfg, cfg.raw_config)
         for object_id, cfg in config.get(DOMAIN, {}).items()
     ]
 
@@ -253,7 +256,7 @@ class ScriptEntity(ToggleEntity):
 
     icon = None
 
-    def __init__(self, hass, object_id, cfg):
+    def __init__(self, hass, object_id, cfg, raw_config):
         """Initialize the script."""
         self.object_id = object_id
         self.icon = cfg.get(CONF_ICON)
@@ -272,6 +275,7 @@ class ScriptEntity(ToggleEntity):
             variables=cfg.get(CONF_VARIABLES),
         )
         self._changed = asyncio.Event()
+        self._raw_config = raw_config
 
     @property
     def should_poll(self):
@@ -323,7 +327,7 @@ class ScriptEntity(ToggleEntity):
             {ATTR_NAME: self.script.name, ATTR_ENTITY_ID: self.entity_id},
             context=context,
         )
-        coro = self.script.async_run(variables, context)
+        coro = self._async_run(variables, context)
         if wait:
             await coro
             return
@@ -335,6 +339,16 @@ class ScriptEntity(ToggleEntity):
         self.hass.async_create_task(coro)
         await self._changed.wait()
 
+    async def _async_run(self, variables, context):
+        with trace_script(
+            self.hass, self.object_id, self._raw_config, context
+        ) as script_trace:
+            script_trace.set_variables(variables)
+            # Prepare tracing the execution of the script's sequence
+            script_trace.set_action_trace(trace_get())
+            with trace_path("sequence"):
+                return await self.script.async_run(variables, context)
+
     async def async_turn_off(self, **kwargs):
         """Stop running the script.
 

homeassistant/components/script/config.py

@@ -25,8 +25,21 @@ async def async_validate_config_item(hass, config, full_config=None):
     return config
 
 
+class ScriptConfig(dict):
+    """Dummy class to allow adding attributes."""
+
+    raw_config = None
+
+
 async def _try_async_validate_config_item(hass, object_id, config, full_config=None):
     """Validate config item."""
+    raw_config = None
+    try:
+        raw_config = dict(config)
+    except ValueError:
+        # Invalid config
+        pass
+
     try:
         cv.slug(object_id)
         config = await async_validate_config_item(hass, config, full_config)
@@ -34,6 +47,8 @@ async def _try_async_validate_config_item(hass, object_id, config, full_config=N
         async_log_exception(ex, DOMAIN, full_config or config, hass)
         return None
 
+    config = ScriptConfig(config)
+    config.raw_config = raw_config
     return config
 
 

homeassistant/components/script/manifest.json

@@ -2,6 +2,7 @@
   "domain": "script",
   "name": "Scripts",
   "documentation": "https://www.home-assistant.io/integrations/script",
+  "dependencies": ["trace"],
   "codeowners": [
     "@home-assistant/core"
   ],

homeassistant/components/script/trace.py

@@ -0,0 +1,23 @@
+"""Trace support for script."""
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+from homeassistant.components.trace import ScriptTrace, async_store_trace
+
+
+@contextmanager
+def trace_script(hass, item_id, config, context):
+    """Trace execution of a script."""
+    trace = ScriptTrace(item_id, config, context)
+    async_store_trace(hass, trace)
+
+    try:
+        yield trace
+    except Exception as ex:  # pylint: disable=broad-except
+        if item_id:
+            trace.set_error(ex)
+        raise ex
+    finally:
+        if item_id:
+            trace.finished()

homeassistant/components/trace/__init__.py

@@ -1,12 +1,188 @@
-"""Support for automation and script tracing and debugging."""
+"""Support for script and automation tracing and debugging."""
+from __future__ import annotations
+
+import datetime as dt
+from itertools import count
+from typing import Any, Deque
+
+from homeassistant.core import Context
+from homeassistant.helpers.trace import TraceElement, trace_id_set
+import homeassistant.util.dt as dt_util
+
 from . import websocket_api
-from .const import DATA_TRACE
+from .const import DATA_TRACE, STORED_TRACES
+from .utils import LimitedSizeDict
 
 DOMAIN = "trace"
 
 
 async def async_setup(hass, config):
     """Initialize the trace integration."""
-    hass.data.setdefault(DATA_TRACE, {})
+    hass.data[DATA_TRACE] = {}
     websocket_api.async_setup(hass)
     return True
+
+
+def async_store_trace(hass, trace):
+    """Store a trace if its item_id is valid."""
+    key = trace.key
+    if key[1]:
+        traces = hass.data[DATA_TRACE]
+        if key not in traces:
+            traces[key] = LimitedSizeDict(size_limit=STORED_TRACES)
+        traces[key][trace.run_id] = trace
+
+
+class ActionTrace:
+    """Base container for an script or automation trace."""
+
+    _run_ids = count(0)
+
+    def __init__(
+        self,
+        key: tuple[str, str],
+        config: dict[str, Any],
+        context: Context,
+    ):
+        """Container for script trace."""
+        self._action_trace: dict[str, Deque[TraceElement]] | None = None
+        self._config: dict[str, Any] = config
+        self.context: Context = context
+        self._error: Exception | None = None
+        self._state: str = "running"
+        self.run_id: str = str(next(self._run_ids))
+        self._timestamp_finish: dt.datetime | None = None
+        self._timestamp_start: dt.datetime = dt_util.utcnow()
+        self.key: tuple[str, str] = key
+        self._variables: dict[str, Any] | None = None
+        trace_id_set((key, self.run_id))
+
+    def set_action_trace(self, trace: dict[str, Deque[TraceElement]]) -> None:
+        """Set action trace."""
+        self._action_trace = trace
+
+    def set_error(self, ex: Exception) -> None:
+        """Set error."""
+        self._error = ex
+
+    def set_variables(self, variables: dict[str, Any]) -> None:
+        """Set variables."""
+        self._variables = variables
+
+    def finished(self) -> None:
+        """Set finish time."""
+        self._timestamp_finish = dt_util.utcnow()
+        self._state = "stopped"
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return dictionary version of this ActionTrace."""
+
+        result = self.as_short_dict()
+
+        action_traces = {}
+        if self._action_trace:
+            for key, trace_list in self._action_trace.items():
+                action_traces[key] = [item.as_dict() for item in trace_list]
+
+        result.update(
+            {
+                "action_trace": action_traces,
+                "config": self._config,
+                "context": self.context,
+                "variables": self._variables,
+            }
+        )
+        if self._error is not None:
+            result["error"] = str(self._error)
+        return result
+
+    def as_short_dict(self) -> dict[str, Any]:
+        """Return a brief dictionary version of this ActionTrace."""
+
+        last_action = None
+
+        if self._action_trace:
+            last_action = list(self._action_trace)[-1]
+
+        result = {
+            "last_action": last_action,
+            "run_id": self.run_id,
+            "state": self._state,
+            "timestamp": {
+                "start": self._timestamp_start,
+                "finish": self._timestamp_finish,
+            },
+            "domain": self.key[0],
+            "item_id": self.key[1],
+        }
+        if self._error is not None:
+            result["error"] = str(self._error)
+        if last_action is not None:
+            result["last_action"] = last_action
+
+        return result
+
+
+class AutomationTrace(ActionTrace):
+    """Container for automation trace."""
+
+    def __init__(
+        self,
+        item_id: str,
+        config: dict[str, Any],
+        context: Context,
+    ):
+        """Container for automation trace."""
+        key = ("automation", item_id)
+        super().__init__(key, config, context)
+        self._condition_trace: dict[str, Deque[TraceElement]] | None = None
+
+    def set_condition_trace(self, trace: dict[str, Deque[TraceElement]]) -> None:
+        """Set condition trace."""
+        self._condition_trace = trace
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return dictionary version of this AutomationTrace."""
+
+        result = super().as_dict()
+
+        condition_traces = {}
+
+        if self._condition_trace:
+            for key, trace_list in self._condition_trace.items():
+                condition_traces[key] = [item.as_dict() for item in trace_list]
+        result["condition_trace"] = condition_traces
+
+        return result
+
+    def as_short_dict(self) -> dict[str, Any]:
+        """Return a brief dictionary version of this AutomationTrace."""
+
+        result = super().as_short_dict()
+
+        last_condition = None
+        trigger = None
+
+        if self._condition_trace:
+            last_condition = list(self._condition_trace)[-1]
+        if self._variables:
+            trigger = self._variables.get("trigger", {}).get("description")
+
+        result["trigger"] = trigger
+        result["last_condition"] = last_condition
+
+        return result
+
+
+class ScriptTrace(ActionTrace):
+    """Container for automation trace."""
+
+    def __init__(
+        self,
+        item_id: str,
+        config: dict[str, Any],
+        context: Context,
+    ):
+        """Container for automation trace."""
+        key = ("script", item_id)
+        super().__init__(key, config, context)

homeassistant/components/trace/const.py

@@ -1,4 +1,4 @@
-"""Shared constants for automation and script tracing and debugging."""
+"""Shared constants for script and automation tracing and debugging."""
 
 DATA_TRACE = "trace"
-STORED_TRACES = 5  # Stored traces per automation
+STORED_TRACES = 5  # Stored traces per script or automation

homeassistant/components/trace/websocket_api.py

@@ -28,7 +28,7 @@ from .utils import TraceJSONEncoder
 
 # mypy: allow-untyped-calls, allow-untyped-defs
 
-TRACE_DOMAINS = ["automation"]
+TRACE_DOMAINS = ["automation", "script"]
 
 
 @callback

homeassistant/helpers/trace.py

@@ -71,7 +71,7 @@ trace_path_stack_cv: ContextVar[list[str] | None] = ContextVar(
 )
 # Copy of last variables
 variables_cv: ContextVar[Any | None] = ContextVar("variables_cv", default=None)
-# Automation ID + Run ID
+# (domain, item_id) + Run ID
 trace_id_cv: ContextVar[tuple[str, str] | None] = ContextVar(
     "trace_id_cv", default=None
 )