stemtrace 0.1.0__py3-none-any.whl

stemtrace/server/__main__.py

@@ -0,0 +1,136 @@
+"""CLI entry point for stemtrace server."""
+
+from typing import Annotated
+
+import typer
+
+app = typer.Typer(
+    name="stemtrace",
+    help="Celery task flow visualizer",
+    no_args_is_help=True,
+)
+
+
+@app.command()
+def server(
+    broker_url: Annotated[
+        str,
+        typer.Option(
+            "--broker-url",
+            "-b",
+            envvar="CELERY_FLOW_BROKER_URL",
+            help="Broker URL (default: redis://localhost:6379/0)",
+        ),
+    ] = "redis://localhost:6379/0",
+    host: Annotated[
+        str,
+        typer.Option("--host", "-h", help="Host to bind to"),
+    ] = "127.0.0.1",
+    port: Annotated[
+        int,
+        typer.Option("--port", "-p", help="Port to bind to"),
+    ] = 8000,
+    reload: Annotated[
+        bool,
+        typer.Option("--reload", help="Enable auto-reload (development)"),
+    ] = False,
+) -> None:
+    """Start the stemtrace web server with embedded consumer."""
+    import uvicorn
+    from fastapi import FastAPI
+
+    from stemtrace.server.fastapi.extension import StemtraceExtension
+
+    typer.echo(f"Starting stemtrace server on {host}:{port}")
+    typer.echo(f"Broker: {broker_url}")
+
+    from stemtrace.server.ui.static import get_static_router
+
+    # Create extension without UI (we'll serve UI at root separately)
+    extension = StemtraceExtension(broker_url=broker_url, serve_ui=False)
+    fastapi_app = FastAPI(
+        title="stemtrace",
+        lifespan=extension.lifespan,
+    )
+
+    # API and WebSocket at /stemtrace (frontend expects this path)
+    fastapi_app.include_router(extension.router, prefix="/stemtrace")
+
+    # UI at root for standalone CLI
+    ui_router = get_static_router()
+    if ui_router is not None:
+        fastapi_app.include_router(ui_router)
+
+    uvicorn.run(fastapi_app, host=host, port=port, reload=reload)
+
+
+@app.command()
+def consume(
+    broker_url: Annotated[
+        str,
+        typer.Option(
+            "--broker-url",
+            "-b",
+            envvar="CELERY_FLOW_BROKER_URL",
+            help="Broker URL (default: redis://localhost:6379/0)",
+        ),
+    ] = "redis://localhost:6379/0",
+    prefix: Annotated[
+        str,
+        typer.Option("--prefix", help="Stream key prefix"),
+    ] = "stemtrace",
+    ttl: Annotated[
+        int,
+        typer.Option("--ttl", help="Event TTL in seconds"),
+    ] = 86400,
+) -> None:
+    """Run the event consumer standalone (for external processing)."""
+    import signal
+    import sys
+
+    from stemtrace.server.consumer import EventConsumer
+    from stemtrace.server.store import GraphStore
+
+    typer.echo("Starting stemtrace consumer (standalone mode)")
+    typer.echo(f"Broker: {broker_url}")
+
+    store = GraphStore()
+    consumer = EventConsumer(broker_url, store, prefix=prefix, ttl=ttl)
+
+    def handle_signal(_signum: int, _frame: object) -> None:
+        """Handle shutdown signals gracefully."""
+        typer.echo("\nShutting down consumer...")
+        consumer.stop()
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, handle_signal)
+    signal.signal(signal.SIGTERM, handle_signal)
+
+    consumer.start()
+    typer.echo("Consumer running. Press Ctrl+C to stop.")
+
+    try:
+        signal.pause()
+    except AttributeError:
+        # Windows doesn't have signal.pause
+        import time
+
+        while consumer.is_running:
+            time.sleep(1)
+
+
+@app.command()
+def version() -> None:
+    """Show version information."""
+    from stemtrace import __version__
+
+    typer.echo(f"stemtrace {__version__}")
+
+
+def main() -> None:
+    """CLI entry point."""
+    app()
+
+
+if __name__ == "__main__":
+    main()

stemtrace/server/api/__init__.py

@@ -0,0 +1,29 @@
+"""API module - REST and WebSocket endpoints."""
+
+from stemtrace.server.api.routes import create_api_router
+from stemtrace.server.api.schemas import (
+    ErrorResponse,
+    GraphListResponse,
+    GraphNodeResponse,
+    GraphResponse,
+    HealthResponse,
+    TaskDetailResponse,
+    TaskEventResponse,
+    TaskListResponse,
+    TaskNodeResponse,
+)
+from stemtrace.server.api.websocket import create_websocket_router
+
+__all__ = [
+    "ErrorResponse",
+    "GraphListResponse",
+    "GraphNodeResponse",
+    "GraphResponse",
+    "HealthResponse",
+    "TaskDetailResponse",
+    "TaskEventResponse",
+    "TaskListResponse",
+    "TaskNodeResponse",
+    "create_api_router",
+    "create_websocket_router",
+]

stemtrace/server/api/routes.py

@@ -0,0 +1,270 @@
+"""REST API endpoints for stemtrace."""
+
+from __future__ import annotations
+
+import contextlib
+from datetime import datetime  # noqa: TC003 (FastAPI needs this at runtime)
+from typing import TYPE_CHECKING, Annotated
+
+from fastapi import APIRouter, HTTPException, Query
+
+from stemtrace.server.api.schemas import (
+    ErrorResponse,
+    GraphListResponse,
+    GraphNodeResponse,
+    GraphResponse,
+    HealthResponse,
+    RegisteredTaskResponse,
+    TaskDetailResponse,
+    TaskEventResponse,
+    TaskListResponse,
+    TaskNodeResponse,
+    TaskRegistryResponse,
+)
+
+if TYPE_CHECKING:
+    from stemtrace.core.events import TaskState
+    from stemtrace.core.graph import TaskNode
+    from stemtrace.server.consumer import AsyncEventConsumer
+    from stemtrace.server.store import GraphStore
+    from stemtrace.server.websocket import WebSocketManager
+
+
+def _node_to_response(node: TaskNode) -> TaskNodeResponse:
+    first_seen = node.events[0].timestamp if node.events else None
+    last_updated = node.events[-1].timestamp if node.events else None
+
+    duration_ms = None
+    if first_seen and last_updated and first_seen != last_updated:
+        duration_ms = int((last_updated - first_seen).total_seconds() * 1000)
+
+    return TaskNodeResponse(
+        task_id=node.task_id,
+        name=node.name,
+        state=node.state,
+        node_type=node.node_type,
+        group_id=node.group_id,
+        chord_id=node.chord_id,
+        parent_id=node.parent_id,
+        children=node.children,
+        events=[TaskEventResponse.model_validate(e) for e in node.events],
+        first_seen=first_seen,
+        last_updated=last_updated,
+        duration_ms=duration_ms,
+    )
+
+
+def _node_to_graph_response(
+    node: TaskNode,
+    all_nodes: dict[str, TaskNode] | None = None,
+) -> GraphNodeResponse:
+    first_seen = node.events[0].timestamp if node.events else None
+    last_updated = node.events[-1].timestamp if node.events else None
+
+    # For synthetic nodes (GROUP/CHORD), compute timing from children
+    if not node.events and node.children and all_nodes:
+        child_first_seen: list[datetime] = []
+        child_last_updated: list[datetime] = []
+        for child_id in node.children:
+            child = all_nodes.get(child_id)
+            if child and child.events:
+                child_first_seen.append(child.events[0].timestamp)
+                child_last_updated.append(child.events[-1].timestamp)
+        if child_first_seen:
+            first_seen = min(child_first_seen)
+        if child_last_updated:
+            last_updated = max(child_last_updated)
+
+    duration_ms = None
+    if first_seen and last_updated and first_seen != last_updated:
+        duration_ms = int((last_updated - first_seen).total_seconds() * 1000)
+
+    return GraphNodeResponse(
+        task_id=node.task_id,
+        name=node.name,
+        state=node.state,
+        node_type=node.node_type,
+        group_id=node.group_id,
+        chord_id=node.chord_id,
+        parent_id=node.parent_id,
+        children=node.children,
+        duration_ms=duration_ms,
+        first_seen=first_seen,
+        last_updated=last_updated,
+    )
+
+
+def create_api_router(
+    store: GraphStore,
+    consumer: AsyncEventConsumer | None = None,
+    ws_manager: WebSocketManager | None = None,
+) -> APIRouter:
+    """Create REST API router with task and graph endpoints."""
+    router = APIRouter(prefix="/api", tags=["stemtrace"])
+
+    @router.get("/health", response_model=HealthResponse)
+    async def health() -> HealthResponse:
+        """Return server health status and connection counts."""
+        from stemtrace import __version__
+
+        return HealthResponse(
+            status="ok",
+            version=__version__,
+            consumer_running=consumer.is_running if consumer else False,
+            websocket_connections=ws_manager.connection_count if ws_manager else 0,
+            node_count=store.node_count,
+        )
+
+    @router.get(
+        "/tasks",
+        response_model=TaskListResponse,
+        responses={400: {"model": ErrorResponse}},
+    )
+    async def list_tasks(
+        limit: Annotated[int, Query(ge=1, le=500)] = 100,
+        offset: Annotated[int, Query(ge=0)] = 0,
+        state: Annotated[str | None, Query(description="Filter by task state")] = None,
+        name: Annotated[
+            str | None, Query(description="Filter by name substring")
+        ] = None,
+        from_date: Annotated[
+            datetime | None, Query(description="Filter by start date (ISO format)")
+        ] = None,
+        to_date: Annotated[
+            datetime | None, Query(description="Filter by end date (ISO format)")
+        ] = None,
+    ) -> TaskListResponse:
+        """List tasks with optional filtering by state, name, and date range."""
+        from stemtrace.core.events import TaskState as TS
+
+        task_state: TaskState | None = None
+        if state is not None:
+            with contextlib.suppress(ValueError):
+                task_state = TS(state)
+
+        nodes, total = store.get_nodes(
+            limit=limit,
+            offset=offset,
+            state=task_state,
+            name_contains=name,
+            from_date=from_date,
+            to_date=to_date,
+        )
+        return TaskListResponse(
+            tasks=[_node_to_response(n) for n in nodes],
+            total=total,
+            limit=limit,
+            offset=offset,
+        )
+
+    @router.get(
+        "/tasks/registry",
+        response_model=TaskRegistryResponse,
+    )
+    async def get_task_registry(
+        query: Annotated[
+            str | None, Query(description="Filter by task name substring")
+        ] = None,
+    ) -> TaskRegistryResponse:
+        """List all discovered task definitions."""
+        unique_names = store.get_unique_task_names()
+
+        tasks: list[RegisteredTaskResponse] = []
+        for name in sorted(unique_names):
+            if query and query.lower() not in name.lower():
+                continue
+
+            parts = name.rsplit(".", 1)
+            module = parts[0] if len(parts) > 1 else None
+
+            tasks.append(
+                RegisteredTaskResponse(
+                    name=name,
+                    module=module,
+                    signature=None,
+                    docstring=None,
+                    bound=False,
+                )
+            )
+
+        return TaskRegistryResponse(tasks=tasks, total=len(tasks))
+
+    @router.get(
+        "/tasks/{task_id}",
+        response_model=TaskDetailResponse,
+        responses={404: {"model": ErrorResponse}},
+    )
+    async def get_task(task_id: str) -> TaskDetailResponse:
+        """Get detailed information for a specific task including children."""
+        node = store.get_node(task_id)
+        if node is None:
+            raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
+
+        children = store.get_children(task_id)
+        return TaskDetailResponse(
+            task=_node_to_response(node),
+            children=[_node_to_response(c) for c in children],
+        )
+
+    @router.get(
+        "/tasks/{task_id}/children",
+        response_model=list[TaskNodeResponse],
+        responses={404: {"model": ErrorResponse}},
+    )
+    async def get_task_children(task_id: str) -> list[TaskNodeResponse]:
+        """Get child tasks for a specific task."""
+        node = store.get_node(task_id)
+        if node is None:
+            raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
+
+        children = store.get_children(task_id)
+        return [_node_to_response(c) for c in children]
+
+    @router.get("/graphs", response_model=GraphListResponse)
+    async def list_graphs(
+        limit: Annotated[int, Query(ge=1, le=100)] = 50,
+        offset: Annotated[int, Query(ge=0)] = 0,
+        from_date: Annotated[
+            datetime | None, Query(description="Filter by start date (ISO format)")
+        ] = None,
+        to_date: Annotated[
+            datetime | None, Query(description="Filter by end date (ISO format)")
+        ] = None,
+    ) -> GraphListResponse:
+        """List task execution graphs (root tasks) with pagination and date filtering."""
+        roots, total = store.get_root_nodes(
+            limit=limit,
+            offset=offset,
+            from_date=from_date,
+            to_date=to_date,
+        )
+        # Build nodes dict for synthetic node timing computation
+        all_nodes: dict[str, TaskNode] = {}
+        for root in roots:
+            all_nodes[root.task_id] = root
+            for child in store.get_children(root.task_id):
+                all_nodes[child.task_id] = child
+        return GraphListResponse(
+            graphs=[_node_to_graph_response(r, all_nodes) for r in roots],
+            total=total,
+            limit=limit,
+            offset=offset,
+        )
+
+    @router.get(
+        "/graphs/{root_id}",
+        response_model=GraphResponse,
+        responses={404: {"model": ErrorResponse}},
+    )
+    async def get_graph(root_id: str) -> GraphResponse:
+        """Get complete task graph starting from a root task."""
+        graph = store.get_graph_from_root(root_id)
+        if not graph:
+            raise HTTPException(status_code=404, detail=f"Graph {root_id} not found")
+
+        return GraphResponse(
+            root_id=root_id,
+            nodes={tid: _node_to_graph_response(n, graph) for tid, n in graph.items()},
+        )
+
+    return router

stemtrace/server/api/schemas.py

@@ -0,0 +1,136 @@
+"""API response schemas for REST endpoints."""
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from stemtrace.core.events import TaskState
+from stemtrace.core.graph import NodeType
+
+
+class TaskEventResponse(BaseModel):
+    """Single task event with full details."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    task_id: str
+    name: str
+    state: TaskState
+    timestamp: datetime
+    parent_id: str | None = None
+    root_id: str | None = None
+    group_id: str | None = None
+    trace_id: str | None = None
+    retries: int = 0
+
+    # Enhanced event data
+    args: list[Any] | None = None
+    kwargs: dict[str, Any] | None = None
+    result: Any | None = None
+    exception: str | None = None
+    traceback: str | None = None
+
+
+class TaskNodeResponse(BaseModel):
+    """Task node with events and timing info."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    task_id: str
+    name: str
+    state: TaskState
+    node_type: NodeType = NodeType.TASK
+    group_id: str | None = None
+    chord_id: str | None = None
+    parent_id: str | None = None
+    children: list[str] = Field(default_factory=list)
+    events: list[TaskEventResponse] = Field(default_factory=list)
+    first_seen: datetime | None = None
+    last_updated: datetime | None = None
+    duration_ms: int | None = None
+
+
+class TaskListResponse(BaseModel):
+    """Paginated list of tasks."""
+
+    tasks: list[TaskNodeResponse]
+    total: int
+    limit: int
+    offset: int
+
+
+class TaskDetailResponse(BaseModel):
+    """Task with its children."""
+
+    task: TaskNodeResponse
+    children: list[TaskNodeResponse] = Field(default_factory=list)
+
+
+class GraphNodeResponse(BaseModel):
+    """Minimal node for graph visualization."""
+
+    task_id: str
+    name: str
+    state: TaskState
+    node_type: NodeType = NodeType.TASK
+    group_id: str | None = None
+    chord_id: str | None = None
+    parent_id: str | None = None
+    children: list[str] = Field(default_factory=list)
+    duration_ms: int | None = None
+    first_seen: datetime | None = None
+    last_updated: datetime | None = None
+
+
+class GraphResponse(BaseModel):
+    """Full task graph from root."""
+
+    root_id: str
+    nodes: dict[str, GraphNodeResponse]
+
+
+class GraphListResponse(BaseModel):
+    """Paginated list of root graphs."""
+
+    graphs: list[GraphNodeResponse]
+    total: int
+    limit: int
+    offset: int
+
+
+class HealthResponse(BaseModel):
+    """Health check status."""
+
+    status: str = "ok"
+    version: str
+    consumer_running: bool = False
+    websocket_connections: int = 0
+    node_count: int = 0
+
+
+class ErrorResponse(BaseModel):
+    """API error."""
+
+    detail: str
+    error_code: str | None = None
+
+
+# Task Registry schemas
+
+
+class RegisteredTaskResponse(BaseModel):
+    """A registered Celery task definition."""
+
+    name: str
+    signature: str | None = None
+    docstring: str | None = None
+    module: str | None = None
+    bound: bool = False
+
+
+class TaskRegistryResponse(BaseModel):
+    """List of all registered tasks."""
+
+    tasks: list[RegisteredTaskResponse]
+    total: int

stemtrace/server/api/websocket.py

@@ -0,0 +1,28 @@
+"""WebSocket endpoint for real-time event streaming."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+
+if TYPE_CHECKING:
+    from stemtrace.server.websocket import WebSocketManager
+
+
+def create_websocket_router(ws_manager: WebSocketManager) -> APIRouter:
+    """Create WebSocket router for real-time task events."""
+    router = APIRouter(tags=["stemtrace-websocket"])
+
+    @router.websocket("/ws")
+    async def websocket_endpoint(websocket: WebSocket) -> None:
+        """Handle WebSocket connection for real-time task events."""
+        await ws_manager.connect(websocket)
+        try:
+            await ws_manager.listen(websocket)
+        except WebSocketDisconnect:
+            pass
+        finally:
+            ws_manager.disconnect(websocket)
+
+    return router