stemtrace 0.1.0__py3-none-any.whl

stemtrace/library/signals.py

@@ -0,0 +1,434 @@
+"""Celery signal handlers for task lifecycle events."""
+
+import logging
+import threading
+import traceback as tb_module
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from celery.signals import (
+    task_failure,
+    task_postrun,
+    task_prerun,
+    task_retry,
+    task_revoked,
+    task_sent,
+)
+
+from stemtrace.core.events import TaskEvent, TaskState
+from stemtrace.core.ports import EventTransport
+from stemtrace.library.config import get_config
+from stemtrace.library.scrubbing import (
+    DEFAULT_SENSITIVE_KEYS,
+    safe_serialize,
+    scrub_args,
+    scrub_dict,
+)
+
+if TYPE_CHECKING:
+    from celery import Task
+
+logger = logging.getLogger(__name__)
+
+_transport: EventTransport | None = None
+# Track task IDs that have received PENDING to avoid duplicates from retries
+_pending_emitted: set[str] = set()
+_pending_emitted_lock = threading.RLock()
+
+
+def _extract_chord_info(chord_attr: Any) -> tuple[str | None, str | None]:
+    """Extract chord group_id and callback task_id from chord attribute.
+
+    Celery's task.request.chord on HEADER tasks contains:
+    - The chord callback signature with options.group_id and options.task_id
+
+    Returns:
+        Tuple of (chord_group_id, callback_task_id)
+    """
+    if chord_attr is None:
+        return None, None
+
+    group_id: str | None = None
+    callback_id: str | None = None
+
+    # Celery Signature object
+    if hasattr(chord_attr, "options"):
+        opts = chord_attr.options
+        if isinstance(opts, dict):
+            group_id = opts.get("group_id") or opts.get("group")
+            callback_id = opts.get("task_id")
+
+    # Dict-like (Celery chord dict: {'task': ..., 'options': {'group_id': ..., 'task_id': ...}})
+    elif isinstance(chord_attr, dict):
+        opts = chord_attr.get("options", {})
+        if isinstance(opts, dict):
+            group_id = opts.get("group_id") or opts.get("group")
+            callback_id = opts.get("task_id")
+
+    return group_id, callback_id
+
+
+def _publish_event(event: TaskEvent) -> None:
+    """Publish event via transport. Fire-and-forget: logs errors, never raises."""
+    if _transport is None:
+        logger.warning("stemtrace not initialized, event dropped: %s", event.task_id)
+        return
+
+    try:
+        _transport.publish(event)
+    except Exception:
+        logger.warning(
+            "Failed to publish event for task %s", event.task_id, exc_info=True
+        )
+
+
+def _get_scrub_config() -> tuple[
+    frozenset[str], frozenset[str] | None, int, bool, bool, bool
+]:
+    """Get scrubbing configuration.
+
+    Returns:
+        Tuple of (sensitive_keys, safe_keys, max_size, scrub_enabled,
+                  capture_args, capture_result)
+    """
+    config = get_config()
+    if config is None:
+        return (DEFAULT_SENSITIVE_KEYS, None, 10240, True, True, True)
+
+    if config.scrub_sensitive_data:
+        sensitive = DEFAULT_SENSITIVE_KEYS | config.additional_sensitive_keys
+        safe = config.safe_keys if config.safe_keys else None
+    else:
+        sensitive = frozenset()
+        safe = None
+
+    return (
+        sensitive,
+        safe,
+        config.max_data_size,
+        config.scrub_sensitive_data,
+        config.capture_args,
+        config.capture_result,
+    )
+
+
+def _scrub_and_serialize_args(
+    args: tuple[Any, ...],
+) -> list[Any] | None:
+    """Scrub and serialize positional arguments."""
+    sensitive, safe, max_size, scrub_enabled, capture_args, _ = _get_scrub_config()
+    if not capture_args:
+        return None
+
+    if scrub_enabled:
+        scrubbed = scrub_args(args, sensitive, safe_keys=safe)
+    else:
+        scrubbed = list(args)
+
+    result: Any = safe_serialize(scrubbed, max_size, sensitive, safe_keys=safe)
+    # safe_serialize may return truncation message string or the list
+    if isinstance(result, list):
+        return result
+    return [result] if result is not None else scrubbed
+
+
+def _scrub_and_serialize_kwargs(
+    kwargs: dict[str, Any],
+) -> dict[str, Any] | None:
+    """Scrub and serialize keyword arguments."""
+    sensitive, safe, max_size, scrub_enabled, capture_args, _ = _get_scrub_config()
+    if not capture_args:
+        return None
+
+    if scrub_enabled:
+        scrubbed = scrub_dict(kwargs, sensitive, safe_keys=safe)
+    else:
+        scrubbed = kwargs
+
+    result: Any = safe_serialize(scrubbed, max_size, sensitive, safe_keys=safe)
+    # safe_serialize may return truncation message string or the dict
+    if isinstance(result, dict):
+        return result
+    # If truncated, wrap message in dict
+    return {"_truncated": result} if result is not None else scrubbed
+
+
+def _scrub_and_serialize_result(result: Any) -> Any | None:
+    """Scrub and serialize task result."""
+    sensitive, safe, max_size, _, _, capture_result = _get_scrub_config()
+    if not capture_result:
+        return None
+
+    return safe_serialize(result, max_size, sensitive, safe_keys=safe)
+
+
+def _format_exception(exc: BaseException | None, einfo: Any = None) -> str | None:
+    """Format exception to a string message."""
+    if exc is not None:
+        return f"{type(exc).__name__}: {exc}"
+    if einfo is not None:
+        return str(einfo.exception) if hasattr(einfo, "exception") else str(einfo)
+    return None
+
+
+def _format_traceback(einfo: Any = None) -> str | None:
+    """Format traceback from exception info."""
+    if einfo is None:
+        return None
+
+    # billiard ExceptionInfo has traceback attribute
+    if hasattr(einfo, "traceback"):
+        tb_str: str = einfo.traceback
+        return tb_str
+    # Standard exception info tuple
+    if hasattr(einfo, "tb"):
+        return "".join(tb_module.format_tb(einfo.tb))
+    return None
+
+
+def _on_task_prerun(
+    task_id: str,
+    task: "Task",
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    **_: Any,
+) -> None:
+    chord_id, chord_callback_id = _extract_chord_info(
+        getattr(task.request, "chord", None)
+    )
+    _publish_event(
+        TaskEvent(
+            task_id=task_id,
+            name=task.name,
+            state=TaskState.STARTED,
+            timestamp=datetime.now(timezone.utc),
+            parent_id=getattr(task.request, "parent_id", None),
+            root_id=getattr(task.request, "root_id", None),
+            group_id=getattr(task.request, "group", None),
+            chord_id=chord_id,
+            chord_callback_id=chord_callback_id,
+            retries=task.request.retries or 0,
+            args=_scrub_and_serialize_args(args),
+            kwargs=_scrub_and_serialize_kwargs(kwargs),
+        )
+    )
+
+
+def _on_task_postrun(
+    task_id: str,
+    task: "Task",
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    retval: Any,
+    state: str,
+    **_: Any,
+) -> None:
+    del args, kwargs
+    if state != "SUCCESS":
+        return
+
+    # Clean up PENDING tracking
+    _pending_emitted.discard(task_id)
+
+    chord_id, chord_callback_id = _extract_chord_info(
+        getattr(task.request, "chord", None)
+    )
+    _publish_event(
+        TaskEvent(
+            task_id=task_id,
+            name=task.name,
+            state=TaskState.SUCCESS,
+            timestamp=datetime.now(timezone.utc),
+            parent_id=getattr(task.request, "parent_id", None),
+            root_id=getattr(task.request, "root_id", None),
+            group_id=getattr(task.request, "group", None),
+            chord_id=chord_id,
+            chord_callback_id=chord_callback_id,
+            retries=task.request.retries or 0,
+            result=_scrub_and_serialize_result(retval),
+        )
+    )
+
+
+def _on_task_failure(
+    task_id: str,
+    exception: BaseException,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    traceback: Any,
+    einfo: Any,
+    sender: "Task",
+    **_: Any,
+) -> None:
+    del args, kwargs, traceback
+
+    # Clean up PENDING tracking
+    _pending_emitted.discard(task_id)
+
+    chord_id, chord_callback_id = _extract_chord_info(
+        getattr(sender.request, "chord", None)
+    )
+    _publish_event(
+        TaskEvent(
+            task_id=task_id,
+            name=sender.name,
+            state=TaskState.FAILURE,
+            timestamp=datetime.now(timezone.utc),
+            parent_id=getattr(sender.request, "parent_id", None),
+            root_id=getattr(sender.request, "root_id", None),
+            group_id=getattr(sender.request, "group", None),
+            chord_id=chord_id,
+            chord_callback_id=chord_callback_id,
+            retries=sender.request.retries or 0,
+            exception=_format_exception(exception),
+            traceback=_format_traceback(einfo),
+        )
+    )
+
+
+def _on_task_retry(
+    sender: "Task",
+    request: Any,
+    reason: Any,
+    einfo: Any,
+    **_: Any,
+) -> None:
+    # reason can be an exception or string
+    exc_message: str | None = None
+    if isinstance(reason, BaseException):
+        exc_message = _format_exception(reason)
+    elif reason is not None:
+        exc_message = str(reason)
+
+    chord_id, chord_callback_id = _extract_chord_info(getattr(request, "chord", None))
+    # Use current retry count (not +1) so RETRY groups with the STARTED that failed
+    # Timeline: STARTED(0) → RETRY(0) → STARTED(1) → RETRY(1) → ...
+    _publish_event(
+        TaskEvent(
+            task_id=request.id,
+            name=sender.name,
+            state=TaskState.RETRY,
+            timestamp=datetime.now(timezone.utc),
+            parent_id=getattr(request, "parent_id", None),
+            root_id=getattr(request, "root_id", None),
+            group_id=getattr(request, "group", None),
+            chord_id=chord_id,
+            chord_callback_id=chord_callback_id,
+            retries=request.retries or 0,
+            exception=exc_message,
+            traceback=_format_traceback(einfo),
+        )
+    )
+
+
+def _on_task_revoked(
+    request: Any,
+    terminated: bool,
+    signum: int | None,
+    expired: bool,
+    sender: "Task",
+    **_: Any,
+) -> None:
+    del terminated, signum, expired
+    chord_id, chord_callback_id = _extract_chord_info(getattr(request, "chord", None))
+    _publish_event(
+        TaskEvent(
+            task_id=request.id,
+            name=sender.name,
+            state=TaskState.REVOKED,
+            timestamp=datetime.now(timezone.utc),
+            parent_id=getattr(request, "parent_id", None),
+            root_id=getattr(request, "root_id", None),
+            group_id=getattr(request, "group", None),
+            chord_id=chord_id,
+            chord_callback_id=chord_callback_id,
+            retries=getattr(request, "retries", 0) or 0,
+        )
+    )
+
+
+def _on_task_sent(
+    sender: str | None = None,
+    task_id: str | None = None,
+    task: str | None = None,
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    headers: dict[str, Any] | None = None,
+    **_: Any,
+) -> None:
+    """Handle task_sent signal - fires when .delay() or .apply_async() is called.
+
+    Note: This fires on the SENDER side (where the task is queued), not on the worker.
+    It captures the PENDING state before a worker picks up the task.
+
+    For retries, we skip emitting PENDING since the task is already tracked
+    and will get RETRY + STARTED events.
+    """
+    if task_id is None:
+        return
+
+    # Skip PENDING for retries - check headers first, then our tracking set
+    if headers and headers.get("retries", 0) > 0:
+        return
+
+    # Skip if we've already emitted PENDING for this task (handles retry re-queues)
+    # Use lock to make check-then-add atomic and prevent duplicates from concurrent threads
+    with _pending_emitted_lock:
+        if task_id in _pending_emitted:
+            return
+        _pending_emitted.add(task_id)
+
+    task_name = task or sender or "unknown"
+
+    # Extract group_id from headers if available
+    group_id = headers.get("group") if headers else None
+
+    _publish_event(
+        TaskEvent(
+            task_id=task_id,
+            name=task_name,
+            state=TaskState.PENDING,
+            timestamp=datetime.now(timezone.utc),
+            group_id=group_id,
+            args=_scrub_and_serialize_args(args) if args else None,
+            kwargs=_scrub_and_serialize_kwargs(kwargs) if kwargs else None,
+        )
+    )
+
+
+def connect_signals(transport: EventTransport) -> None:
+    """Register signal handlers with the given transport."""
+    global _transport
+    _transport = transport
+
+    task_sent.connect(_on_task_sent)
+    task_prerun.connect(_on_task_prerun)
+    task_postrun.connect(_on_task_postrun)
+    task_failure.connect(_on_task_failure)
+    task_retry.connect(_on_task_retry)
+    task_revoked.connect(_on_task_revoked)
+
+    # Set publisher for bootsteps (RECEIVED events)
+    from stemtrace.library.bootsteps import _set_publisher
+
+    _set_publisher(_publish_event)
+
+    logger.info("stemtrace signal handlers connected")
+
+
+def disconnect_signals() -> None:
+    """Disconnect all signal handlers."""
+    global _transport
+    _transport = None
+
+    task_sent.disconnect(_on_task_sent)
+    task_prerun.disconnect(_on_task_prerun)
+    task_postrun.disconnect(_on_task_postrun)
+    task_failure.disconnect(_on_task_failure)
+    task_retry.disconnect(_on_task_retry)
+    task_revoked.disconnect(_on_task_revoked)
+
+    # Clear tracking state
+    _pending_emitted.clear()
+
+    logger.info("stemtrace signal handlers disconnected")

stemtrace/library/transports/__init__.py

@@ -0,0 +1,36 @@
+"""Broker-agnostic transport factory."""
+
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse
+
+from stemtrace.core.exceptions import UnsupportedBrokerError
+
+if TYPE_CHECKING:
+    from stemtrace.core.ports import EventTransport
+
+_SCHEME_ALIASES: dict[str, str] = {
+    "rediss": "redis",
+    "amqps": "amqp",
+}
+
+
+def get_transport(
+    url: str, prefix: str = "stemtrace", ttl: int = 86400
+) -> "EventTransport":
+    """Create a transport from a broker URL."""
+    scheme = urlparse(url).scheme.lower()
+    scheme = _SCHEME_ALIASES.get(scheme, scheme)
+
+    if scheme == "redis":
+        from stemtrace.library.transports.redis import RedisTransport
+
+        return RedisTransport.from_url(url, prefix=prefix, ttl=ttl)
+    elif scheme == "memory":
+        from stemtrace.library.transports.memory import MemoryTransport
+
+        return MemoryTransport()
+    else:
+        raise UnsupportedBrokerError(scheme)
+
+
+__all__ = ["get_transport"]

stemtrace/library/transports/memory.py

@@ -0,0 +1,33 @@
+"""In-memory transport for testing."""
+
+from collections.abc import Iterator
+from typing import ClassVar
+
+from typing_extensions import Self
+
+from stemtrace.core.events import TaskEvent
+
+
+class MemoryTransport:
+    """In-memory event transport. Events stored in class-level list for test inspection."""
+
+    events: ClassVar[list[TaskEvent]] = []
+
+    def publish(self, event: TaskEvent) -> None:
+        """Store event in memory."""
+        MemoryTransport.events.append(event)
+
+    def consume(self) -> Iterator[TaskEvent]:
+        """Yield all stored events."""
+        yield from MemoryTransport.events
+
+    @classmethod
+    def from_url(cls, url: str) -> Self:
+        """Create transport (ignores URL)."""
+        del url
+        return cls()
+
+    @classmethod
+    def clear(cls) -> None:
+        """Clear all stored events."""
+        cls.events.clear()

stemtrace/library/transports/redis.py

@@ -0,0 +1,92 @@
+"""Redis transport using Redis Streams."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import Self
+
+from stemtrace.core.events import TaskEvent
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from redis import Redis
+
+logger = logging.getLogger(__name__)
+
+
+class RedisTransport:
+    """Redis Streams-based event transport (XADD/XREAD)."""
+
+    def __init__(self, client: Redis[Any], prefix: str, ttl: int) -> None:
+        """Initialize Redis transport with client and stream configuration."""
+        self._client = client
+        self._ttl = ttl
+        self._stream_key = f"{prefix}:events"
+        self._maxlen = max(ttl, 10000)
+
+    @property
+    def client(self) -> Redis[Any]:
+        """The Redis client."""
+        return self._client
+
+    @property
+    def stream_key(self) -> str:
+        """The stream key."""
+        return self._stream_key
+
+    @property
+    def ttl(self) -> int:
+        """TTL in seconds."""
+        return self._ttl
+
+    def publish(self, event: TaskEvent) -> None:
+        """Publish event. Fire-and-forget: logs errors, never raises."""
+        try:
+            self._client.xadd(
+                self._stream_key,
+                {"data": event.model_dump_json()},
+                maxlen=self._maxlen,
+                approximate=True,
+            )
+        except Exception:
+            logger.warning(
+                "Failed to publish event %s to Redis", event.task_id, exc_info=True
+            )
+
+    def consume(self, last_id: str = "0") -> Iterator[TaskEvent]:
+        """Blocking iterator that yields events as they arrive."""
+        current_id = last_id
+        while True:
+            results = self._client.xread(
+                {self._stream_key: current_id},
+                block=5000,
+                count=100,
+            )
+            if not results:
+                continue
+
+            for _stream_name, messages in results:
+                for message_id, fields in messages:
+                    current_id = (
+                        message_id.decode()
+                        if isinstance(message_id, bytes)
+                        else message_id
+                    )
+                    data = fields.get(b"data") or fields.get("data")
+                    if data:
+                        data_str = data.decode() if isinstance(data, bytes) else data
+                        yield TaskEvent.model_validate_json(data_str)
+
+    @classmethod
+    def from_url(cls, url: str, prefix: str = "stemtrace", ttl: int = 86400) -> Self:
+        """Create transport from Redis URL."""
+        from redis import Redis as RedisClient
+
+        return cls(
+            client=RedisClient.from_url(url, decode_responses=False),
+            prefix=prefix,
+            ttl=ttl,
+        )

stemtrace/server/__init__.py

@@ -0,0 +1,24 @@
+"""Server component - Web server for visualization."""
+
+from stemtrace.server.consumer import AsyncEventConsumer, EventConsumer
+from stemtrace.server.fastapi import (
+    StemtraceExtension,
+    create_router,
+    no_auth,
+    require_api_key,
+    require_basic_auth,
+)
+from stemtrace.server.store import GraphStore
+from stemtrace.server.websocket import WebSocketManager
+
+__all__ = [
+    "AsyncEventConsumer",
+    "EventConsumer",
+    "GraphStore",
+    "StemtraceExtension",
+    "WebSocketManager",
+    "create_router",
+    "no_auth",
+    "require_api_key",
+    "require_basic_auth",
+]