glaip-sdk 0.1.3__py3-none-any.whl → 0.6.19__py3-none-any.whl

glaip_sdk/utils/rendering/{steps.py → steps/event_processor.py}

@@ -1,4 +1,4 @@
-"""Rendering utilities.
+"""SSE event processing mixin for StepManager.
 
 Authors:
     Raymond Christopher (raymond.christopher@gdplabs.id)
@@ -7,373 +7,27 @@ Authors:
 from __future__ import annotations
 
 import logging
-from collections.abc import Iterator
+from collections.abc import Mapping
 from copy import deepcopy
 from time import monotonic
 from typing import Any
 
-from glaip_sdk.icons import ICON_AGENT_STEP, ICON_DELEGATE, ICON_TOOL_STEP
 from glaip_sdk.utils.rendering.models import Step
-from glaip_sdk.utils.rendering.step_tree_state import StepTreeState
+from glaip_sdk.utils.rendering.timing import coerce_server_time
 
 logger = logging.getLogger(__name__)
-UNKNOWN_STEP_DETAIL = "Unknown step detail"
 
 
-class StepManager:
-    """Manages the lifecycle and organization of execution steps.
+COERCION_FAILED_KEY = "_meta_coercion_failed_"
 
-    Tracks step creation, parent-child relationships, and execution state
-    with automatic pruning of old steps when limits are reached.
-    """
-
-    def __init__(self, max_steps: int = 200) -> None:
-        """Initialize the step manager.
-
-        Args:
-            max_steps: Maximum number of steps to retain before pruning
-        """
-        normalised_max = int(max_steps) if isinstance(max_steps, (int, float)) else 0
-        self.state = StepTreeState(max_steps=normalised_max)
-        self.by_id: dict[str, Step] = self.state.step_index
-        self.key_index: dict[tuple, str] = {}
-        self.slot_counter: dict[tuple, int] = {}
-        self.max_steps = normalised_max
-        self._last_running: dict[tuple, str] = {}
-        self._step_aliases: dict[str, str] = {}
-        self.root_agent_id: str | None = None
-        self._scope_anchors: dict[str, list[str]] = {}
-        self._step_scope_map: dict[str, str] = {}
-
-    def set_root_agent(self, agent_id: str | None) -> None:
-        """Record the root agent identifier for scope-aware parenting."""
-        if isinstance(agent_id, str) and agent_id.strip():
-            self.root_agent_id = agent_id.strip()
-
-    def _alloc_slot(
-        self,
-        task_id: str | None,
-        context_id: str | None,
-        kind: str,
-        name: str,
-    ) -> int:
-        k = (task_id, context_id, kind, name)
-        self.slot_counter[k] = self.slot_counter.get(k, 0) + 1
-        return self.slot_counter[k]
-
-    def _key(
-        self,
-        task_id: str | None,
-        context_id: str | None,
-        kind: str,
-        name: str,
-        slot: int,
-    ) -> tuple[str | None, str | None, str, str, int]:
-        return (task_id, context_id, kind, name, slot)
-
-    def _make_id(
-        self,
-        task_id: str | None,
-        context_id: str | None,
-        kind: str,
-        name: str,
-        slot: int,
-    ) -> str:
-        return f"{task_id or 't'}::{context_id or 'c'}::{kind}::{name}::{slot}"
-
-    def start_or_get(
-        self,
-        *,
-        task_id: str | None,
-        context_id: str | None,
-        kind: str,
-        name: str,
-        parent_id: str | None = None,
-        args: dict[str, object] | None = None,
-    ) -> Step:
-        """Start a new step or return existing running step with same parameters.
-
-        Args:
-            task_id: Task identifier
-            context_id: Context identifier
-            kind: Step kind (tool, delegate, agent)
-            name: Step name
-            parent_id: Parent step ID if this is a child step
-            args: Step arguments
-
-        Returns:
-            The Step instance (new or existing)
-        """
-        existing = self.find_running(task_id=task_id, context_id=context_id, kind=kind, name=name)
-        if existing:
-            if args and existing.args != args:
-                existing.args = args
-            return existing
-        slot = self._alloc_slot(task_id, context_id, kind, name)
-        key = self._key(task_id, context_id, kind, name, slot)
-        step_id = self._make_id(task_id, context_id, kind, name, slot)
-        st = Step(
-            step_id=step_id,
-            kind=kind,
-            name=name,
-            parent_id=parent_id,
-            task_id=task_id,
-            context_id=context_id,
-            args=args or {},
-        )
-        self.by_id[step_id] = st
-        if parent_id:
-            self.children.setdefault(parent_id, []).append(step_id)
-        else:
-            self.order.append(step_id)
-        self.key_index[key] = step_id
-        self.state.retained_ids.add(step_id)
-        self._prune_steps()
-        self._last_running[(task_id, context_id, kind, name)] = step_id
-        return st
-
-    def _calculate_total_steps(self) -> int:
-        """Calculate total number of steps."""
-        return len(self.order) + sum(len(v) for v in self.children.values())
-
-    def _get_subtree_size(self, root_id: str) -> int:
-        """Get the size of a subtree (including root)."""
-        subtree = [root_id]
-        stack = list(self.children.get(root_id, []))
-        while stack:
-            x = stack.pop()
-            subtree.append(x)
-            stack.extend(self.children.get(x, []))
-        return len(subtree)
-
-    def _remove_subtree(self, root_id: str) -> None:
-        """Remove a complete subtree from all data structures."""
-        for step_id in self._collect_subtree_ids(root_id):
-            self._purge_step_references(step_id)
-
-    def _collect_subtree_ids(self, root_id: str) -> list[str]:
-        """Return a flat list of step ids contained within a subtree."""
-        stack = [root_id]
-        collected: list[str] = []
-        while stack:
-            sid = stack.pop()
-            collected.append(sid)
-            stack.extend(self.children.pop(sid, []))
-        return collected
-
-    def _purge_step_references(self, step_id: str) -> None:
-        """Remove a single step id from all indexes and helper structures."""
-        st = self.by_id.pop(step_id, None)
-        if st:
-            key = (st.task_id, st.context_id, st.kind, st.name)
-            self._last_running.pop(key, None)
-            self.state.retained_ids.discard(step_id)
-            self.state.discard_running(step_id)
-        self._remove_parent_links(step_id)
-        if step_id in self.order:
-            self.order.remove(step_id)
-        self.state.buffered_children.pop(step_id, None)
-        self.state.pending_branch_failures.discard(step_id)
-
-    def _remove_parent_links(self, child_id: str) -> None:
-        """Detach a child id from any parent lists."""
-        for parent, kids in self.children.copy().items():
-            if child_id not in kids:
-                continue
-            kids.remove(child_id)
-            if not kids:
-                self.children.pop(parent, None)
-
-    def _should_prune_steps(self, total: int) -> bool:
-        """Check if steps should be pruned."""
-        if self.max_steps <= 0:
-            return False
-        return total > self.max_steps
 
-    def _get_oldest_step_id(self) -> str | None:
-        """Get the oldest step ID for pruning."""
-        return self.order[0] if self.order else None
+class StepEventMixin:
+    """Mixin providing SSE event processing capabilities for StepManager.
 
-    def _prune_steps(self) -> None:
-        """Prune steps when total exceeds maximum."""
-        total = self._calculate_total_steps()
-        if not self._should_prune_steps(total):
-            return
-
-        while self._should_prune_steps(total) and self.order:
-            sid = self._get_oldest_step_id()
-            if not sid:
-                break
-
-            subtree_size = self._get_subtree_size(sid)
-            self._remove_subtree(sid)
-            total -= subtree_size
-
-    def remove_step(self, step_id: str) -> None:
-        """Remove a single step from the tree and cached indexes."""
-        step = self.by_id.pop(step_id, None)
-        if not step:
-            return
-
-        if step.parent_id:
-            self.state.unlink_child(step.parent_id, step_id)
-        else:
-            self.state.unlink_root(step_id)
-
-        self.children.pop(step_id, None)
-        self.state.buffered_children.pop(step_id, None)
-        self.state.retained_ids.discard(step_id)
-        self.state.pending_branch_failures.discard(step_id)
-        self.state.discard_running(step_id)
-
-        self.key_index = {key: sid for key, sid in self.key_index.items() if sid != step_id}
-        for key, last_sid in self._last_running.copy().items():
-            if last_sid == step_id:
-                self._last_running.pop(key, None)
-
-        aliases = [alias for alias, target in self._step_aliases.items() if alias == step_id or target == step_id]
-        for alias in aliases:
-            self._step_aliases.pop(alias, None)
-
-    def get_child_count(self, step_id: str) -> int:
-        """Get the number of child steps for a given step.
-
-        Args:
-            step_id: The parent step ID
-
-        Returns:
-            Number of child steps
-        """
-        return len(self.children.get(step_id, []))
-
-    def find_running(
-        self,
-        *,
-        task_id: str | None,
-        context_id: str | None,
-        kind: str,
-        name: str,
-    ) -> Step | None:
-        """Find a currently running step with the given parameters.
-
-        Args:
-            task_id: Task identifier
-            context_id: Context identifier
-            kind: Step kind (tool, delegate, agent)
-            name: Step name
-
-        Returns:
-            The running Step if found, None otherwise
-        """
-        key = (task_id, context_id, kind, name)
-        step_id = self._last_running.get(key)
-        if step_id:
-            st = self.by_id.get(step_id)
-            if st and st.status != "finished":
-                return st
-        for sid in reversed(list(self._iter_all_steps())):
-            st = self.by_id.get(sid)
-            if (
-                st
-                and (st.task_id, st.context_id, st.kind, st.name)
-                == (
-                    task_id,
-                    context_id,
-                    kind,
-                    name,
-                )
-                and st.status != "finished"
-            ):
-                return st
-        return None
-
-    def finish(
-        self,
-        *,
-        task_id: str | None,
-        context_id: str | None,
-        kind: str,
-        name: str,
-        output: object | None = None,
-        duration_raw: float | None = None,
-    ) -> Step:
-        """Finish a step with the given parameters.
-
-        Args:
-            task_id: Task identifier
-            context_id: Context identifier
-            kind: Step kind (tool, delegate, agent)
-            name: Step name
-            output: Step output data
-            duration_raw: Raw duration in seconds
-
-        Returns:
-            The finished Step instance
-
-        Raises:
-            RuntimeError: If no matching step is found
-        """
-        st = self.find_running(task_id=task_id, context_id=context_id, kind=kind, name=name)
-        if not st:
-            # Try to find any existing step with matching parameters, even if not running
-            for sid in reversed(list(self._iter_all_steps())):
-                st_check = self.by_id.get(sid)
-                if (
-                    st_check
-                    and st_check.task_id == task_id
-                    and st_check.context_id == context_id
-                    and st_check.kind == kind
-                    and st_check.name == name
-                ):
-                    st = st_check
-                    break
-
-            # If still no step found, create a new one
-            if not st:
-                st = self.start_or_get(task_id=task_id, context_id=context_id, kind=kind, name=name)
-
-        if output:
-            st.output = output
-        st.finish(duration_raw)
-        key = (task_id, context_id, kind, name)
-        if self._last_running.get(key) == st.step_id:
-            self._last_running.pop(key, None)
-        return st
-
-    def _iter_all_steps(self) -> Iterator[str]:
-        for root in self.order:
-            yield root
-            stack = list(self.children.get(root, []))
-            while stack:
-                sid = stack.pop()
-                yield sid
-                stack.extend(self.children.get(sid, []))
-
-    def iter_tree(self) -> Iterator[tuple[str, tuple[bool, ...]]]:
-        """Expose depth-first traversal info for rendering."""
-        yield from self.state.iter_visible_tree()
-
-    @property
-    def order(self) -> list[str]:
-        """Root step ordering accessor backed by StepTreeState."""
-        return self.state.root_order
-
-    @order.setter
-    def order(self, value: list[str]) -> None:
-        self.state.root_order = list(value)
-
-    @property
-    def children(self) -> dict[str, list[str]]:
-        """Child mapping accessor backed by StepTreeState."""
-        return self.state.child_map
-
-    @children.setter
-    def children(self, value: dict[str, list[str]]) -> None:
-        self.state.child_map = value
-
-    # ------------------------------------------------------------------
-    # SSE-aware helpers
-    # ------------------------------------------------------------------
+    This mixin adds methods to process server-sent events (SSE) and update
+    step state accordingly. It handles event parsing, step creation/updates,
+    parent-child relationships, and duration tracking.
+    """
 
     def apply_event(self, event: dict[str, Any]) -> Step:
         """Apply an SSE step event and return the updated step."""
@@ -463,6 +117,7 @@ class StepManager:
 
     def _apply_single_event(self, event: dict[str, Any]) -> Step:
         metadata, step_id, tool_info, args = self._parse_event_payload(event)
+        metadata_failed = bool(metadata.pop(COERCION_FAILED_KEY, False))
         tool_name = self._resolve_tool_name(tool_info, metadata, step_id)
         kind = self._derive_step_kind(tool_name, metadata)
         parent_hint = self._coerce_parent_id(metadata.get("previous_step_ids"))
@@ -479,14 +134,17 @@ class StepManager:
         self._link_step(step, parent_id)
 
         self.state.retained_ids.add(step.step_id)
-        step.display_label = self._compose_display_label(step.kind, tool_name, args, metadata)
+        if metadata_failed:
+            step.metadata = {}
+        else:
+            step.metadata = dict(metadata)
         self._flush_buffered_children(step.step_id)
         self._apply_pending_branch_flags(step.step_id)
 
         status = self._normalise_status(metadata.get("status"), event.get("status"), event.get("task_state"))
         status = self._apply_failure_state(step, status, event)
 
-        server_time = self._coerce_server_time(metadata.get("time"))
+        server_time = coerce_server_time(metadata.get("time"))
         self._update_server_timestamps(step, server_time, status)
 
         self._apply_duration(
@@ -504,16 +162,17 @@ class StepManager:
             status=status,
         )
 
-        step.status_icon = self._status_icon_for_step(step)
         self._update_parallel_tracking(step)
         self._update_running_index(step)
         self._prune_steps()
         return step
 
     def _parse_event_payload(self, event: dict[str, Any]) -> tuple[dict[str, Any], str, dict[str, Any], dict[str, Any]]:
-        metadata = event.get("metadata") or {}
-        if not isinstance(metadata, dict):
+        metadata_raw = event.get("metadata") or {}
+        metadata, metadata_reliable = self._coerce_event_metadata(metadata_raw)
+        if not metadata:
             raise ValueError("Step event missing metadata payload")
+        metadata[COERCION_FAILED_KEY] = not metadata_reliable
 
         step_id = metadata.get("step_id")
         if not isinstance(step_id, str) or not step_id:
@@ -531,6 +190,39 @@ class StepManager:
 
         return metadata, step_id, tool_info, args
 
+    def _coerce_event_metadata(self, metadata: Any) -> tuple[dict[str, Any], bool]:
+        """Return a dict copy of event metadata with graceful fallbacks."""
+        if isinstance(metadata, dict):
+            return metadata, True
+        if isinstance(metadata, Mapping):
+            try:
+                return dict(metadata), True
+            except Exception:
+                logger.debug("Failed to coerce mapping metadata; falling back to key-by-key copy", exc_info=True)
+                return self._copy_mapping_fields(metadata), False
+        # All other payloads are treated as empty
+        return {}, False
+
+    @staticmethod
+    def _copy_mapping_fields(metadata: Mapping[str, Any]) -> dict[str, Any]:
+        """Copy known fields from a mapping without iteration."""
+        copied: dict[str, Any] = {}
+        for key in (
+            "step_id",
+            "tool_info",
+            "status",
+            "kind",
+            "previous_step_ids",
+            "time",
+            "agent_name",
+            "task_id",
+            "context_id",
+        ):
+            value = metadata.get(key)  # type: ignore[attr-defined]
+            if value is not None:
+                copied[key] = value
+        return copied
+
     def _resolve_tool_name(self, tool_info: dict[str, Any], metadata: dict[str, Any], step_id: str) -> str:
         name = tool_info.get("name")
         if not name:
@@ -738,65 +430,6 @@ class StepManager:
             return False
         return all(ch in "0123456789abcdefABCDEF" for ch in stripped)
 
-    def _step_icon_for_kind(self, step_kind: str) -> str:
-        if step_kind == "agent":
-            return ICON_AGENT_STEP
-        if step_kind == "delegate":
-            return ICON_DELEGATE
-        if step_kind == "thinking":
-            return "💭"
-        return ICON_TOOL_STEP
-
-    def _humanize_tool_name(self, raw_name: str | None) -> str:
-        if not raw_name:
-            return UNKNOWN_STEP_DETAIL
-        name = raw_name
-        if name.startswith("delegate_to_"):
-            name = name.removeprefix("delegate_to_")
-        elif name.startswith("delegate_"):
-            name = name.removeprefix("delegate_")
-        cleaned = name.replace("_", " ").replace("-", " ").strip()
-        if not cleaned:
-            return UNKNOWN_STEP_DETAIL
-        return cleaned[:1].upper() + cleaned[1:]
-
-    def _compose_display_label(
-        self,
-        step_kind: str,
-        tool_name: str | None,
-        args: dict[str, Any],
-        metadata: dict[str, Any],
-    ) -> str:
-        icon = self._step_icon_for_kind(step_kind)
-        body = self._resolve_label_body(step_kind, tool_name, metadata)
-        label = f"{icon} {body}".strip()
-        if isinstance(args, dict) and args:
-            label = f"{label} —"
-        return label or UNKNOWN_STEP_DETAIL
-
-    def _resolve_label_body(
-        self,
-        step_kind: str,
-        tool_name: str | None,
-        metadata: dict[str, Any],
-    ) -> str:
-        if step_kind == "thinking":
-            thinking_text = metadata.get("thinking_and_activity_info")
-            if isinstance(thinking_text, str) and thinking_text.strip():
-                return thinking_text.strip()
-            return "Thinking…"
-
-        if step_kind == "delegate":
-            return self._humanize_tool_name(tool_name)
-
-        if step_kind == "agent":
-            agent_name = metadata.get("agent_name")
-            if isinstance(agent_name, str) and agent_name.strip():
-                return agent_name.strip()
-
-        friendly = self._humanize_tool_name(tool_name)
-        return friendly
-
     def _normalise_status(
         self,
         metadata_status: Any,
@@ -994,15 +627,6 @@ class StepManager:
         slug = slug.replace("-", "_").strip()
         return slug or None
 
-    @staticmethod
-    def _coerce_server_time(value: Any) -> float | None:
-        if isinstance(value, (int, float)):
-            return float(value)
-        try:
-            return float(value)
-        except (TypeError, ValueError):
-            return None
-
     def _update_server_timestamps(self, step: Step, server_time: float | None, status: str) -> None:
         if server_time is None:
             return
@@ -1038,7 +662,6 @@ class StepManager:
         step = self.by_id.get(step_id)
         if step:
             step.branch_failed = True
-            step.status_icon = "warning"
         self.state.pending_branch_failures.discard(step_id)
 
     def _set_branch_warning(self, parent_id: str | None) -> None:
@@ -1047,7 +670,6 @@ class StepManager:
         parent = self.by_id.get(parent_id)
         if parent:
             parent.branch_failed = True
-            parent.status_icon = "warning"
         else:
             self.state.pending_branch_failures.add(parent_id)
 
@@ -1079,15 +701,6 @@ class StepManager:
             if current:
                 current.is_parallel = is_parallel
 
-    def _status_icon_for_step(self, step: Step) -> str:
-        if step.status == "finished":
-            return "warning" if step.branch_failed else "success"
-        if step.status == "failed":
-            return "failed"
-        if step.status == "stopped":
-            return "warning"
-        return "spinner"
-
     def _canonicalize_step_id(self, step_id: str, tool_info: dict[str, Any]) -> str:
         alias = self._lookup_alias(step_id)
         if alias: