stemtrace 0.1.0__py3-none-any.whl

stemtrace/server/consumer.py

@@ -0,0 +1,133 @@
+"""Broker-agnostic event consumer."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import TYPE_CHECKING
+
+from stemtrace.library.transports import get_transport
+
+if TYPE_CHECKING:
+    from stemtrace.core.ports import EventTransport
+    from stemtrace.server.store import GraphStore
+
+logger = logging.getLogger(__name__)
+
+
+class EventConsumer:
+    """Background consumer that reads events and updates the GraphStore."""
+
+    def __init__(
+        self,
+        broker_url: str,
+        store: GraphStore,
+        *,
+        prefix: str = "stemtrace",
+        ttl: int = 86400,
+    ) -> None:
+        """Initialize consumer with broker URL and target store."""
+        self._broker_url = broker_url
+        self._store = store
+        self._prefix = prefix
+        self._ttl = ttl
+        self._transport: EventTransport | None = None
+        self._thread: threading.Thread | None = None
+        self._stop_event = threading.Event()
+
+    @property
+    def is_running(self) -> bool:
+        """Whether the consumer thread is alive."""
+        return self._thread is not None and self._thread.is_alive()
+
+    def start(self) -> None:
+        """Start consuming in background thread. Idempotent."""
+        if self._thread is not None:
+            return
+
+        self._stop_event.clear()
+        self._transport = get_transport(
+            self._broker_url, prefix=self._prefix, ttl=self._ttl
+        )
+        self._thread = threading.Thread(
+            target=self._consume_loop,
+            name="stemtrace-consumer",
+            daemon=True,
+        )
+        self._thread.start()
+        logger.info("Event consumer started")
+
+    def stop(self, timeout: float = 5.0) -> None:
+        """Stop the consumer gracefully."""
+        if self._thread is None:
+            return
+
+        self._stop_event.set()
+        self._thread.join(timeout=timeout)
+        if self._thread.is_alive():
+            logger.warning("Consumer thread did not stop gracefully")
+        self._thread = None
+        self._transport = None
+        logger.info("Event consumer stopped")
+
+    def _consume_loop(self) -> None:
+        if self._transport is None:
+            return
+
+        logger.debug("Consumer loop starting, reading from %s", self._broker_url)
+
+        try:
+            for event in self._transport.consume():
+                if self._stop_event.is_set():
+                    break
+
+                try:
+                    self._store.add_event(event)
+                    logger.debug("Consumed event: %s (%s)", event.task_id, event.state)
+                except Exception:
+                    logger.exception("Error processing event %s", event.task_id)
+        except Exception:
+            if not self._stop_event.is_set():
+                logger.exception("Consumer loop error")
+
+
+class AsyncEventConsumer:
+    """Async context manager wrapper for EventConsumer."""
+
+    def __init__(
+        self,
+        broker_url: str,
+        store: GraphStore,
+        *,
+        prefix: str = "stemtrace",
+        ttl: int = 86400,
+    ) -> None:
+        """Initialize async consumer wrapper with broker URL and target store."""
+        self._consumer = EventConsumer(broker_url, store, prefix=prefix, ttl=ttl)
+
+    @property
+    def is_running(self) -> bool:
+        """Whether the consumer thread is alive."""
+        return self._consumer.is_running
+
+    async def __aenter__(self) -> AsyncEventConsumer:
+        """Start consumer on context enter."""
+        self._consumer.start()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
+        """Stop consumer on context exit."""
+        self._consumer.stop()
+
+    def start(self) -> None:
+        """Start the consumer."""
+        self._consumer.start()
+
+    def stop(self, timeout: float = 5.0) -> None:
+        """Stop the consumer."""
+        self._consumer.stop(timeout=timeout)

stemtrace/server/fastapi/__init__.py

@@ -0,0 +1,13 @@
+"""FastAPI integration module."""
+
+from stemtrace.server.fastapi.auth import no_auth, require_api_key, require_basic_auth
+from stemtrace.server.fastapi.extension import StemtraceExtension
+from stemtrace.server.fastapi.router import create_router
+
+__all__ = [
+    "StemtraceExtension",
+    "create_router",
+    "no_auth",
+    "require_api_key",
+    "require_basic_auth",
+]

stemtrace/server/fastapi/auth.py

@@ -0,0 +1,77 @@
+"""Authentication helpers for stemtrace API."""
+
+from __future__ import annotations
+
+import secrets
+from typing import Annotated, Any
+
+from fastapi import Depends, HTTPException, status
+from fastapi.security import APIKeyHeader, HTTPBasic, HTTPBasicCredentials
+
+_basic_auth = HTTPBasic(auto_error=False)
+
+
+def require_basic_auth(username: str, password: str) -> Any:
+    """Create dependency requiring HTTP Basic auth."""
+
+    def _verify(
+        credentials: Annotated[HTTPBasicCredentials | None, Depends(_basic_auth)],
+    ) -> str:
+        if credentials is None:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Authentication required",
+                headers={"WWW-Authenticate": "Basic"},
+            )
+
+        username_ok = secrets.compare_digest(
+            credentials.username.encode("utf-8"),
+            username.encode("utf-8"),
+        )
+        password_ok = secrets.compare_digest(
+            credentials.password.encode("utf-8"),
+            password.encode("utf-8"),
+        )
+
+        if not (username_ok and password_ok):
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid credentials",
+                headers={"WWW-Authenticate": "Basic"},
+            )
+
+        return credentials.username
+
+    return Depends(_verify)
+
+
+def require_api_key(api_key: str, *, header_name: str = "X-API-Key") -> Any:
+    """Create dependency requiring API key in header."""
+    key_header = APIKeyHeader(name=header_name, auto_error=False)
+
+    async def _verify(
+        provided_key: str | None = Depends(key_header),
+    ) -> str:
+        if provided_key is None:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="API key required",
+            )
+
+        if not secrets.compare_digest(
+            provided_key.encode("utf-8"),
+            api_key.encode("utf-8"),
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid API key",
+            )
+
+        return provided_key
+
+    return Depends(_verify)
+
+
+def no_auth() -> Any:
+    """Dependency that allows all requests."""
+    return Depends(lambda: None)

stemtrace/server/fastapi/extension.py

@@ -0,0 +1,121 @@
+"""Full-featured FastAPI extension with embedded consumer."""
+
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+from stemtrace.server.consumer import AsyncEventConsumer
+from stemtrace.server.fastapi.router import create_router
+from stemtrace.server.store import GraphStore
+from stemtrace.server.ui.static import get_static_router
+from stemtrace.server.websocket import WebSocketManager
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from fastapi import APIRouter, FastAPI
+
+
+class StemtraceExtension:
+    """Complete FastAPI integration: store, consumer, WebSocket, and router.
+
+    Example:
+        flow = StemtraceExtension(broker_url="redis://localhost:6379")
+        app = FastAPI(lifespan=flow.lifespan)
+        app.include_router(flow.router, prefix="/stemtrace")
+    """
+
+    def __init__(
+        self,
+        broker_url: str,
+        *,
+        embedded_consumer: bool = True,
+        serve_ui: bool = True,
+        prefix: str = "stemtrace",
+        ttl: int = 86400,
+        max_nodes: int = 10000,
+        auth_dependency: Any = None,
+    ) -> None:
+        """Initialize extension with broker URL and optional configuration."""
+        self._broker_url = broker_url
+        self._embedded_consumer = embedded_consumer
+        self._serve_ui = serve_ui
+        self._prefix = prefix
+        self._ttl = ttl
+        self._auth_dependency = auth_dependency
+
+        self._store = GraphStore(max_nodes=max_nodes)
+        self._ws_manager = WebSocketManager()
+        self._consumer: AsyncEventConsumer | None = None
+
+        if embedded_consumer:
+            self._consumer = AsyncEventConsumer(
+                broker_url,
+                self._store,
+                prefix=prefix,
+                ttl=ttl,
+            )
+
+        self._store.add_listener(self._ws_manager.queue_event)
+
+    @property
+    def store(self) -> GraphStore:
+        """The in-memory graph store."""
+        return self._store
+
+    @property
+    def consumer(self) -> AsyncEventConsumer | None:
+        """Event consumer, or None if embedded_consumer=False."""
+        return self._consumer
+
+    @property
+    def ws_manager(self) -> WebSocketManager:
+        """The WebSocket connection manager."""
+        return self._ws_manager
+
+    @property
+    def router(self) -> APIRouter:
+        """Pre-configured router. Mount with app.include_router()."""
+        router = create_router(
+            store=self._store,
+            consumer=self._consumer,
+            ws_manager=self._ws_manager,
+            auth_dependency=self._auth_dependency,
+        )
+
+        if self._serve_ui:
+            ui_router = get_static_router()
+            if ui_router is not None:
+                router.include_router(ui_router)
+
+        return router
+
+    @asynccontextmanager
+    async def lifespan(self, app: FastAPI) -> AsyncIterator[None]:  # noqa: ARG002
+        """FastAPI lifespan: starts consumer and WebSocket on startup."""
+        await self._ws_manager.start_broadcast_loop()
+        if self._consumer is not None:
+            self._consumer.start()
+
+        try:
+            yield
+        finally:
+            if self._consumer is not None:
+                self._consumer.stop()
+            await self._ws_manager.stop_broadcast_loop()
+
+    def compose_lifespan(self, other_lifespan: Any = None) -> Any:
+        """Compose with another lifespan context manager."""
+
+        @asynccontextmanager
+        async def combined(app: FastAPI) -> AsyncIterator[None]:
+            """Combined lifespan context manager."""
+            async with self.lifespan(app):
+                if other_lifespan is not None:
+                    async with other_lifespan(app):
+                        yield
+                else:
+                    yield
+
+        return combined

stemtrace/server/fastapi/router.py

@@ -0,0 +1,42 @@
+"""FastAPI router factory for stemtrace."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from fastapi import APIRouter
+
+from stemtrace.server.api.routes import create_api_router
+from stemtrace.server.api.websocket import create_websocket_router
+from stemtrace.server.store import GraphStore
+from stemtrace.server.websocket import WebSocketManager
+
+if TYPE_CHECKING:
+    from stemtrace.server.consumer import AsyncEventConsumer
+
+
+def create_router(
+    store: GraphStore | None = None,
+    consumer: AsyncEventConsumer | None = None,
+    ws_manager: WebSocketManager | None = None,
+    *,
+    auth_dependency: Any = None,
+) -> APIRouter:
+    """Create API router. For embedded consumer, use StemtraceExtension."""
+    if store is None:
+        store = GraphStore()
+    if ws_manager is None:
+        ws_manager = WebSocketManager()
+
+    router = APIRouter()
+    api_router = create_api_router(store, consumer, ws_manager)
+    ws_router = create_websocket_router(ws_manager)
+
+    dependencies: list[Any] = []
+    if auth_dependency is not None:
+        dependencies.append(auth_dependency)
+
+    router.include_router(api_router, dependencies=dependencies)
+    router.include_router(ws_router)
+
+    return router

stemtrace/server/store.py

@@ -0,0 +1,260 @@
+"""Thread-safe in-memory graph store."""
+
+from __future__ import annotations
+
+import contextlib
+import threading
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+from stemtrace.core.graph import NodeType, TaskGraph, TaskNode
+
+
+def _ensure_tz_aware(dt: datetime) -> datetime:
+    """Ensure datetime is timezone-aware (assume UTC if naive)."""
+    if dt.tzinfo is None:
+        return dt.replace(tzinfo=timezone.utc)
+    return dt
+
+
+def _ensure_end_of_day(dt: datetime) -> datetime:
+    """Ensure to_date is end of day and timezone-aware.
+
+    When a date-only value (YYYY-MM-DD) is parsed, it becomes midnight.
+    For to_date filtering, we want end of day to include the full day.
+    """
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    # If time is midnight (date-only input), set to end of day
+    if dt.hour == 0 and dt.minute == 0 and dt.second == 0 and dt.microsecond == 0:
+        dt = dt.replace(hour=23, minute=59, second=59, microsecond=999999)
+    return dt
+
+
+# Fallback for nodes with no events (synthetic nodes)
+_MIN_DATETIME = datetime.min.replace(tzinfo=timezone.utc)
+
+
+def _get_node_timestamp(node: TaskNode, graph: TaskGraph) -> datetime:
+    """Get the most recent timestamp for a node (including children for synthetic nodes)."""
+    if node.events:
+        return node.events[-1].timestamp
+
+    # For synthetic nodes (GROUP/CHORD), use the latest child timestamp
+    if node.node_type in (NodeType.GROUP, NodeType.CHORD) and node.children:
+        child_timestamps: list[datetime] = []
+        for child_id in node.children:
+            child = graph.get_node(child_id)
+            if child and child.events:
+                child_timestamps.append(child.events[-1].timestamp)
+        if child_timestamps:
+            return max(child_timestamps)
+
+    return _MIN_DATETIME
+
+
+def _get_first_timestamp(node: TaskNode, graph: TaskGraph) -> datetime:
+    """Get the first timestamp for a node (including children for synthetic nodes)."""
+    if node.events:
+        return node.events[0].timestamp
+
+    # For synthetic nodes (GROUP/CHORD), use the earliest child timestamp
+    if node.node_type in (NodeType.GROUP, NodeType.CHORD) and node.children:
+        child_timestamps: list[datetime] = []
+        for child_id in node.children:
+            child = graph.get_node(child_id)
+            if child and child.events:
+                child_timestamps.append(child.events[0].timestamp)
+        if child_timestamps:
+            return min(child_timestamps)
+
+    return _MIN_DATETIME
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from stemtrace.core.events import TaskEvent, TaskState
+
+
+class GraphStore:
+    """Thread-safe in-memory store for TaskGraph with LRU eviction."""
+
+    def __init__(self, max_nodes: int = 10000) -> None:
+        """Initialize store with optional maximum node limit for LRU eviction."""
+        self._graph = TaskGraph()
+        self._lock = threading.RLock()
+        self._max_nodes = max_nodes
+        self._listeners: list[Callable[[TaskEvent], None]] = []
+
+    def add_event(self, event: TaskEvent) -> None:
+        """Add event to graph and notify listeners."""
+        with self._lock:
+            self._graph.add_event(event)
+            self._maybe_evict()
+
+        for listener in self._listeners:
+            with contextlib.suppress(Exception):
+                listener(event)
+
+    def get_node(self, task_id: str) -> TaskNode | None:
+        """Get node by ID, or None if not found."""
+        with self._lock:
+            return self._graph.get_node(task_id)
+
+    def get_nodes(
+        self,
+        *,
+        limit: int = 100,
+        offset: int = 0,
+        state: TaskState | None = None,
+        name_contains: str | None = None,
+        from_date: datetime | None = None,
+        to_date: datetime | None = None,
+    ) -> tuple[list[TaskNode], int]:
+        """Get nodes with optional filtering, most recent first.
+
+        Returns:
+            Tuple of (filtered nodes, total count matching filters).
+        """
+        with self._lock:
+            nodes = list(self._graph.nodes.values())
+
+        if state is not None:
+            nodes = [n for n in nodes if n.state == state]
+        if name_contains is not None:
+            name_lower = name_contains.lower()
+            nodes = [n for n in nodes if name_lower in n.name.lower()]
+        if from_date is not None:
+            from_dt = _ensure_tz_aware(from_date)
+            nodes = [n for n in nodes if n.events and n.events[-1].timestamp >= from_dt]
+        if to_date is not None:
+            to_dt = _ensure_end_of_day(to_date)
+            nodes = [n for n in nodes if n.events and n.events[0].timestamp <= to_dt]
+
+        nodes.sort(
+            key=lambda n: n.events[-1].timestamp if n.events else _MIN_DATETIME,
+            reverse=True,
+        )
+        total = len(nodes)
+        return nodes[offset : offset + limit], total
+
+    def get_root_nodes(
+        self,
+        *,
+        limit: int = 50,
+        offset: int = 0,
+        from_date: datetime | None = None,
+        to_date: datetime | None = None,
+    ) -> tuple[list[TaskNode], int]:
+        """Get root nodes (no parent), most recent first.
+
+        Returns:
+            Tuple of (filtered root nodes, total count matching filters).
+        """
+        with self._lock:
+            root_nodes = [
+                self._graph.nodes[rid]
+                for rid in self._graph.root_ids
+                if rid in self._graph.nodes
+            ]
+
+            # Date filtering - use child timestamps for synthetic nodes
+            if from_date is not None:
+                from_dt = _ensure_tz_aware(from_date)
+                root_nodes = [
+                    n
+                    for n in root_nodes
+                    if _get_node_timestamp(n, self._graph) >= from_dt
+                ]
+            if to_date is not None:
+                to_dt = _ensure_end_of_day(to_date)
+                root_nodes = [
+                    n
+                    for n in root_nodes
+                    if _get_first_timestamp(n, self._graph) <= to_dt
+                ]
+
+            # Sort while holding lock since we need access to graph for children
+            root_nodes.sort(
+                key=lambda n: _get_node_timestamp(n, self._graph),
+                reverse=True,
+            )
+            total = len(root_nodes)
+        return root_nodes[offset : offset + limit], total
+
+    def get_children(self, task_id: str) -> list[TaskNode]:
+        """Get child nodes of a task."""
+        with self._lock:
+            node = self._graph.get_node(task_id)
+            if node is None:
+                return []
+            return [
+                self._graph.nodes[cid]
+                for cid in node.children
+                if cid in self._graph.nodes
+            ]
+
+    def get_graph_from_root(self, root_id: str) -> dict[str, TaskNode]:
+        """Get all nodes in subgraph starting from root."""
+        with self._lock:
+            root = self._graph.get_node(root_id)
+            if root is None:
+                return {}
+
+            result: dict[str, TaskNode] = {}
+            to_visit = [root_id]
+
+            while to_visit:
+                current_id = to_visit.pop()
+                if current_id in result:
+                    continue
+                node = self._graph.get_node(current_id)
+                if node is None:
+                    continue
+                result[current_id] = node
+                to_visit.extend(node.children)
+
+            return result
+
+    def add_listener(self, callback: Callable[[TaskEvent], None]) -> None:
+        """Register callback for new events (used by WebSocket manager)."""
+        self._listeners.append(callback)
+
+    def remove_listener(self, callback: Callable[[TaskEvent], None]) -> None:
+        """Unregister an event listener."""
+        with contextlib.suppress(ValueError):
+            self._listeners.remove(callback)
+
+    @property
+    def node_count(self) -> int:
+        """Current node count."""
+        with self._lock:
+            return len(self._graph.nodes)
+
+    def get_unique_task_names(self) -> set[str]:
+        """Get all unique task names seen in events."""
+        with self._lock:
+            return {node.name for node in self._graph.nodes.values()}
+
+    def _maybe_evict(self) -> None:
+        """Evict oldest 10% when over capacity. Call with lock held."""
+        if len(self._graph.nodes) <= self._max_nodes:
+            return
+
+        nodes_by_age = sorted(
+            self._graph.nodes.values(),
+            key=lambda n: n.events[0].timestamp if n.events else _MIN_DATETIME,
+        )
+
+        to_remove = len(nodes_by_age) - int(self._max_nodes * 0.9)
+        for node in nodes_by_age[:to_remove]:
+            if node.parent_id and node.parent_id in self._graph.nodes:
+                parent = self._graph.nodes[node.parent_id]
+                if node.task_id in parent.children:
+                    parent.children.remove(node.task_id)
+
+            if node.task_id in self._graph.root_ids:
+                self._graph.root_ids.remove(node.task_id)
+
+            del self._graph.nodes[node.task_id]

stemtrace/server/ui/__init__.py

@@ -0,0 +1,5 @@
+"""UI module - React SPA for task visualization."""
+
+from stemtrace.server.ui.static import get_static_router, is_ui_available
+
+__all__ = ["get_static_router", "is_ui_available"]