stemtrace 0.1.0__py3-none-any.whl

stemtrace/server/ui/frontend/src/hooks/WebSocketContext.tsx

@@ -0,0 +1,143 @@
+/**
+ * WebSocket context for sharing connection status globally.
+ */
+
+import { useQueryClient } from '@tanstack/react-query'
+import {
+  createContext,
+  type ReactNode,
+  useCallback,
+  useContext,
+  useEffect,
+  useRef,
+  useState,
+} from 'react'
+import type { TaskEvent } from '@/api/client'
+
+type ConnectionStatus = 'connecting' | 'connected' | 'disconnected' | 'error'
+
+interface WebSocketContextValue {
+  connectionStatus: ConnectionStatus
+  isConnected: boolean
+  lastEvent: TaskEvent | null
+  events: TaskEvent[]
+}
+
+const WebSocketContext = createContext<WebSocketContextValue | null>(null)
+
+// WebSocket URL
+const WS_URL = import.meta.env.DEV
+  ? `ws://${window.location.hostname}:8000/stemtrace/ws`
+  : `${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/stemtrace/ws`
+
+// Reconnect interval when disconnected (3 seconds for quick recovery)
+const RECONNECT_INTERVAL = 3000
+
+interface WebSocketProviderProps {
+  children: ReactNode
+}
+
+export function WebSocketProvider({ children }: WebSocketProviderProps) {
+  const [connectionStatus, setConnectionStatus] = useState<ConnectionStatus>('connecting')
+  const [lastEvent, setLastEvent] = useState<TaskEvent | null>(null)
+  const [events, setEvents] = useState<TaskEvent[]>([])
+  const queryClient = useQueryClient()
+  const wsRef = useRef<WebSocket | null>(null)
+  const reconnectTimeoutRef = useRef<number | null>(null)
+
+  const connect = useCallback(() => {
+    // Don't reconnect if already connected or connecting
+    if (
+      wsRef.current?.readyState === WebSocket.OPEN ||
+      wsRef.current?.readyState === WebSocket.CONNECTING
+    ) {
+      return
+    }
+
+    setConnectionStatus('connecting')
+
+    try {
+      const ws = new WebSocket(WS_URL)
+      wsRef.current = ws
+
+      ws.onopen = () => {
+        setConnectionStatus('connected')
+        console.log('[WebSocket] Connected')
+      }
+
+      ws.onmessage = (event) => {
+        try {
+          const data = JSON.parse(event.data) as TaskEvent
+          setLastEvent(data)
+          setEvents((prev) => [data, ...prev].slice(0, 100)) // Keep last 100
+
+          // Invalidate queries to refresh data
+          queryClient.invalidateQueries({ queryKey: ['tasks'] })
+          queryClient.invalidateQueries({ queryKey: ['graphs'] })
+        } catch (e) {
+          console.error('[WebSocket] Failed to parse message:', e)
+        }
+      }
+
+      ws.onclose = () => {
+        setConnectionStatus('disconnected')
+        wsRef.current = null
+        console.log('[WebSocket] Disconnected, will retry in 3s...')
+
+        // Schedule reconnect
+        if (reconnectTimeoutRef.current) {
+          clearTimeout(reconnectTimeoutRef.current)
+        }
+        reconnectTimeoutRef.current = window.setTimeout(() => {
+          connect()
+        }, RECONNECT_INTERVAL)
+      }
+
+      ws.onerror = () => {
+        setConnectionStatus('error')
+        // onclose will be called after onerror, which will trigger reconnect
+      }
+    } catch (e) {
+      setConnectionStatus('error')
+      console.error('[WebSocket] Failed to connect:', e)
+
+      // Schedule reconnect on error
+      if (reconnectTimeoutRef.current) {
+        clearTimeout(reconnectTimeoutRef.current)
+      }
+      reconnectTimeoutRef.current = window.setTimeout(() => {
+        connect()
+      }, RECONNECT_INTERVAL)
+    }
+  }, [queryClient])
+
+  useEffect(() => {
+    connect()
+
+    return () => {
+      if (reconnectTimeoutRef.current) {
+        clearTimeout(reconnectTimeoutRef.current)
+      }
+      if (wsRef.current) {
+        wsRef.current.close()
+      }
+    }
+  }, [connect])
+
+  const value: WebSocketContextValue = {
+    connectionStatus,
+    isConnected: connectionStatus === 'connected',
+    lastEvent,
+    events,
+  }
+
+  return <WebSocketContext.Provider value={value}>{children}</WebSocketContext.Provider>
+}
+
+export function useWebSocketContext(): WebSocketContextValue {
+  const context = useContext(WebSocketContext)
+  if (!context) {
+    throw new Error('useWebSocketContext must be used within a WebSocketProvider')
+  }
+  return context
+}

stemtrace/server/ui/frontend/src/hooks/useInfiniteScroll.ts

@@ -0,0 +1,56 @@
+/**
+ * Hook for infinite scroll behavior with intersection observer.
+ */
+
+import { useCallback, useEffect, useRef } from 'react'
+
+interface UseInfiniteScrollOptions {
+  hasNextPage: boolean | undefined
+  isFetchingNextPage: boolean
+  fetchNextPage: () => void
+  rootMargin?: string
+}
+
+export function useInfiniteScroll({
+  hasNextPage,
+  isFetchingNextPage,
+  fetchNextPage,
+  rootMargin = '200px',
+}: UseInfiniteScrollOptions) {
+  const observerRef = useRef<IntersectionObserver | null>(null)
+
+  const lastElementRef = useCallback(
+    (node: HTMLElement | null) => {
+      if (isFetchingNextPage) return
+
+      if (observerRef.current) {
+        observerRef.current.disconnect()
+      }
+
+      observerRef.current = new IntersectionObserver(
+        (entries) => {
+          if (entries[0].isIntersecting && hasNextPage) {
+            fetchNextPage()
+          }
+        },
+        { rootMargin },
+      )
+
+      if (node) {
+        observerRef.current.observe(node)
+      }
+    },
+    [isFetchingNextPage, hasNextPage, fetchNextPage, rootMargin],
+  )
+
+  // Cleanup on unmount
+  useEffect(() => {
+    return () => {
+      if (observerRef.current) {
+        observerRef.current.disconnect()
+      }
+    }
+  }, [])
+
+  return { lastElementRef }
+}

stemtrace/server/ui/frontend/src/index.css

@@ -0,0 +1,104 @@
+@import "tailwindcss";
+@import "react-day-picker/style.css";
+
+@theme {
+  --font-sans: "JetBrains Mono", ui-monospace, monospace;
+
+  /* Task state colors */
+  --color-pending: #6b7280;
+  --color-received: #6b7280;
+  --color-started: #3b82f6;
+  --color-success: #22c55e;
+  --color-failure: #ef4444;
+  --color-retry: #f59e0b;
+  --color-revoked: #8b5cf6;
+  --color-rejected: #dc2626;
+}
+
+body {
+  @apply antialiased;
+}
+
+.task-state-badge {
+  @apply inline-flex items-center px-2 py-0.5 rounded text-xs font-medium;
+}
+
+.task-state-pending {
+  @apply bg-gray-500/20 text-gray-400;
+}
+
+.task-state-received {
+  @apply bg-gray-500/20 text-gray-400;
+}
+
+.task-state-started {
+  @apply bg-blue-500/20 text-blue-400;
+}
+
+.task-state-success {
+  @apply bg-green-500/20 text-green-400;
+}
+
+.task-state-failure {
+  @apply bg-red-500/20 text-red-400;
+}
+
+.task-state-retry {
+  @apply bg-amber-500/20 text-amber-400;
+}
+
+.task-state-revoked {
+  @apply bg-purple-500/20 text-purple-400;
+}
+
+.task-state-rejected {
+  @apply bg-red-600/20 text-red-500;
+}
+
+/* React Flow customization */
+.react-flow__node {
+  @apply rounded-lg border border-slate-700 bg-slate-800 shadow-lg;
+}
+
+.react-flow__edge-path {
+  @apply stroke-slate-500;
+}
+
+.react-flow__controls {
+  @apply bg-slate-800 border-slate-700;
+}
+
+.react-flow__controls-button {
+  @apply bg-slate-800 border-slate-700 text-slate-400 hover:bg-slate-700;
+}
+
+/* Hide handles on header label nodes inside containers */
+.header-label-node .react-flow__handle {
+  display: none !important;
+}
+
+/* Hide handles on container nodes (GROUP/CHORD) - edges connect but handles hidden */
+.container-node .react-flow__handle {
+  opacity: 0 !important;
+  width: 8px;
+  height: 8px;
+}
+
+/* Hide left handles on container child nodes (only show right for edges) */
+.container-child .react-flow__handle-left,
+.container-child .react-flow__handle-top,
+.container-child .react-flow__handle-bottom {
+  display: none !important;
+}
+
+/* Make right handle on container children subtle */
+.container-child .react-flow__handle-right {
+  width: 6px;
+  height: 6px;
+  opacity: 0.5;
+}
+
+/* Allow task nodes to expand based on content, but not container nodes */
+.react-flow__node-default:not(.container-node) {
+  width: auto !important;
+}

stemtrace/server/ui/frontend/src/main.tsx

@@ -0,0 +1,44 @@
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
+import { createRouter, RouterProvider } from '@tanstack/react-router'
+import { StrictMode } from 'react'
+import { createRoot } from 'react-dom/client'
+import { WebSocketProvider } from './hooks/WebSocketContext'
+import { routeTree } from './routeTree.gen'
+import './index.css'
+
+// Create query client
+const queryClient = new QueryClient({
+  defaultOptions: {
+    queries: {
+      staleTime: 5000,
+      refetchOnWindowFocus: false,
+    },
+  },
+})
+
+// Create router
+const router = createRouter({
+  routeTree,
+  context: {
+    queryClient,
+  },
+  defaultPreload: 'intent',
+})
+
+// Register router for type safety
+declare module '@tanstack/react-router' {
+  interface Register {
+    router: typeof router
+  }
+}
+
+// Render
+createRoot(document.getElementById('root')!).render(
+  <StrictMode>
+    <QueryClientProvider client={queryClient}>
+      <WebSocketProvider>
+        <RouterProvider router={router} />
+      </WebSocketProvider>
+    </QueryClientProvider>
+  </StrictMode>,
+)

stemtrace/server/ui/frontend/src/routes/__root.tsx

@@ -0,0 +1,103 @@
+import type { QueryClient } from '@tanstack/react-query'
+import { createRootRouteWithContext, Link, Outlet } from '@tanstack/react-router'
+import { TanStackRouterDevtools } from '@tanstack/router-devtools'
+import { useHealth } from '@/api/queries'
+import { Tooltip } from '@/components/Tooltip'
+import { useWebSocketContext } from '@/hooks/WebSocketContext'
+
+interface RouterContext {
+  queryClient: QueryClient
+}
+
+export const Route = createRootRouteWithContext<RouterContext>()({
+  component: RootLayout,
+})
+
+function RootLayout() {
+  // Get WebSocket connection status from context
+  const { connectionStatus, isConnected } = useWebSocketContext()
+
+  // Check if server is reachable via health endpoint
+  const { data: health, isError: isServerDown } = useHealth()
+
+  // Determine display status
+  // Priority: Server down > WS connected > WS connecting > Polling
+  let statusLabel: string
+  let statusColor: string
+  let statusTooltip: string
+
+  if (isServerDown) {
+    statusLabel = 'Offline'
+    statusColor = 'bg-red-500'
+    statusTooltip = 'Server is unreachable. Check if stemtrace is running.'
+  } else if (isConnected) {
+    statusLabel = 'Live'
+    statusColor = 'bg-green-500'
+    statusTooltip = 'Real-time updates via WebSocket'
+  } else if (connectionStatus === 'connecting') {
+    statusLabel = 'Connecting'
+    statusColor = 'bg-amber-500 animate-pulse'
+    statusTooltip = 'Connecting to WebSocket...'
+  } else {
+    statusLabel = 'Polling'
+    statusColor = 'bg-amber-500'
+    statusTooltip = 'Updates every 5s. For real-time updates, configure WebSocket proxy.'
+  }
+
+  return (
+    <div className="min-h-screen bg-slate-950">
+      {/* Header */}
+      <header className="border-b border-slate-800 bg-slate-900/50 backdrop-blur sticky top-0 z-50">
+        <div className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8">
+          <div className="flex items-center justify-between h-14">
+            {/* Logo */}
+            <Link to="/" className="flex items-baseline gap-2">
+              <span className="text-lg font-bold bg-gradient-to-r from-green-400 to-emerald-500 bg-clip-text text-transparent">
+                stemtrace
+              </span>
+              {health?.version && <span className="text-xs text-slate-500">v{health.version}</span>}
+            </Link>
+
+            {/* Navigation */}
+            <nav className="flex items-center gap-6">
+              <Link
+                to="/"
+                className="text-sm text-slate-400 hover:text-slate-100 transition-colors [&.active]:text-slate-100"
+              >
+                Tasks
+              </Link>
+              <Link
+                to="/graphs"
+                className="text-sm text-slate-400 hover:text-slate-100 transition-colors [&.active]:text-slate-100"
+              >
+                Graphs
+              </Link>
+              <Link
+                to="/registry"
+                className="text-sm text-slate-400 hover:text-slate-100 transition-colors [&.active]:text-slate-100"
+              >
+                Registry
+              </Link>
+            </nav>
+
+            {/* Status indicator with tooltip */}
+            <Tooltip content={statusTooltip} placement="bottom">
+              <div className="flex items-center gap-2 cursor-help">
+                <span className={`w-2 h-2 rounded-full ${statusColor}`} />
+                <span className="text-xs text-slate-500">{statusLabel}</span>
+              </div>
+            </Tooltip>
+          </div>
+        </div>
+      </header>
+
+      {/* Main content */}
+      <main className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 py-8">
+        <Outlet />
+      </main>
+
+      {/* Dev tools */}
+      {import.meta.env.DEV && <TanStackRouterDevtools position="bottom-right" />}
+    </div>
+  )
+}

stemtrace/server/ui/frontend/src/routes/graph.$rootId.tsx

@@ -0,0 +1,51 @@
+import { createFileRoute, Link } from '@tanstack/react-router'
+import { useGraph } from '@/api/queries'
+import { TaskGraph } from '@/components/TaskGraph'
+
+export const Route = createFileRoute('/graph/$rootId')({
+  component: GraphDetailPage,
+})
+
+function GraphDetailPage() {
+  const { rootId } = Route.useParams()
+  const { data, isLoading, error } = useGraph(rootId)
+
+  if (isLoading) {
+    return (
+      <div className="flex items-center justify-center h-96">
+        <div className="animate-spin w-8 h-8 border-2 border-slate-600 border-t-green-500 rounded-full" />
+      </div>
+    )
+  }
+
+  if (error || !data) {
+    return (
+      <div className="text-center py-12">
+        <p className="text-red-400">Failed to load graph</p>
+        <Link to="/graphs" className="text-blue-400 hover:underline mt-2 inline-block">
+          Back to graphs
+        </Link>
+      </div>
+    )
+  }
+
+  return (
+    <div className="space-y-6">
+      {/* Breadcrumb */}
+      <nav className="flex items-center gap-2 text-sm">
+        <Link to="/graphs" className="text-slate-400 hover:text-slate-200">
+          Graphs
+        </Link>
+        <span className="text-slate-600">/</span>
+        <span className="text-slate-200 font-mono">{rootId.slice(0, 8)}</span>
+      </nav>
+
+      {/* Graph visualization */}
+      <div className="bg-slate-900 rounded-xl border border-slate-800 overflow-hidden">
+        <div className="h-[600px]">
+          <TaskGraph nodes={data.nodes} rootId={data.root_id} />
+        </div>
+      </div>
+    </div>
+  )
+}

stemtrace/server/ui/frontend/src/routes/graphs.tsx

@@ -0,0 +1,186 @@
+import { createFileRoute, Link, useNavigate } from '@tanstack/react-router'
+import type { DateRange } from 'react-day-picker'
+import { useGraphsInfinite } from '@/api/queries'
+import { DateRangePicker } from '@/components/DateRangePicker'
+import { TaskStateBadge } from '@/components/TaskStateBadge'
+import { useInfiniteScroll } from '@/hooks/useInfiniteScroll'
+import { formatDateParam, formatDuration, formatTime, parseDateParam } from '@/utils/format'
+
+interface SearchParams {
+  from_date?: string
+  to_date?: string
+}
+
+export const Route = createFileRoute('/graphs')({
+  component: GraphsPage,
+  validateSearch: (search: Record<string, unknown>): SearchParams => ({
+    from_date: typeof search.from_date === 'string' ? search.from_date : undefined,
+    to_date: typeof search.to_date === 'string' ? search.to_date : undefined,
+  }),
+})
+
+function GraphsPage() {
+  const navigate = useNavigate()
+  const { from_date, to_date } = Route.useSearch()
+
+  const { data, isLoading, error, hasNextPage, isFetchingNextPage, fetchNextPage } =
+    useGraphsInfinite({
+      from_date,
+      to_date,
+    })
+
+  const { lastElementRef } = useInfiniteScroll({
+    hasNextPage,
+    isFetchingNextPage,
+    fetchNextPage,
+  })
+
+  const dateRange: DateRange | undefined =
+    from_date || to_date
+      ? {
+          from: parseDateParam(from_date),
+          to: parseDateParam(to_date),
+        }
+      : undefined
+
+  const handleDateChange = (range: DateRange | undefined) => {
+    navigate({
+      to: '/graphs',
+      search: {
+        from_date: formatDateParam(range?.from),
+        to_date: formatDateParam(range?.to),
+      },
+      replace: true,
+    })
+  }
+
+  const allGraphs = data?.pages.flatMap((page) => page.graphs) ?? []
+  const total = data?.pages[0]?.total ?? 0
+
+  return (
+    <div className="space-y-6">
+      <div className="flex items-start justify-between">
+        <div>
+          <h1 className="text-2xl font-bold text-slate-100">Task Graphs</h1>
+          <p className="text-sm text-slate-500 mt-1">
+            Visualize task execution flows and dependencies
+          </p>
+        </div>
+        <DateRangePicker value={dateRange} onChange={handleDateChange} />
+      </div>
+
+      {isLoading ? (
+        <GraphListSkeleton />
+      ) : error ? (
+        <div className="text-center py-12 text-red-400">Failed to load graphs</div>
+      ) : !allGraphs.length ? (
+        <EmptyState />
+      ) : (
+        <div className="grid gap-4">
+          {allGraphs.map((graph, index) => {
+            const isLast = index === allGraphs.length - 1
+            const startTime = formatTime(graph.first_seen)
+            const duration = formatDuration(graph.duration_ms)
+
+            return (
+              <Link
+                key={graph.task_id}
+                ref={isLast ? lastElementRef : undefined}
+                to="/graph/$rootId"
+                params={{ rootId: graph.task_id }}
+                className="bg-slate-900 rounded-xl border border-slate-800 p-6 hover:border-slate-700 transition-colors"
+              >
+                <div className="flex items-start justify-between">
+                  <div>
+                    <h3 className="text-lg font-medium text-slate-100">{graph.name}</h3>
+                    <p className="text-sm text-slate-500 font-mono mt-1">{graph.task_id}</p>
+                  </div>
+                  <TaskStateBadge state={graph.state} />
+                </div>
+
+                <div className="flex items-center gap-4 mt-4 text-sm text-slate-400">
+                  <span>{graph.children.length} child tasks</span>
+                  {startTime && (
+                    <>
+                      <span className="text-slate-600">·</span>
+                      <span>{startTime}</span>
+                    </>
+                  )}
+                  {duration && (
+                    <>
+                      <span className="text-slate-600">·</span>
+                      <span>{duration}</span>
+                    </>
+                  )}
+                </div>
+              </Link>
+            )
+          })}
+
+          {/* Loading indicator */}
+          {isFetchingNextPage && <LoadingMore />}
+
+          {/* Pagination info */}
+          <div className="text-center text-sm text-slate-500 pt-4">
+            Showing {allGraphs.length} of {total} graphs
+          </div>
+        </div>
+      )}
+    </div>
+  )
+}
+
+function LoadingMore() {
+  return (
+    <div className="flex justify-center py-4">
+      <div className="flex items-center gap-2 text-sm text-slate-400">
+        <svg
+          className="animate-spin h-4 w-4"
+          xmlns="http://www.w3.org/2000/svg"
+          fill="none"
+          viewBox="0 0 24 24"
+          aria-hidden="true"
+        >
+          <circle
+            className="opacity-25"
+            cx="12"
+            cy="12"
+            r="10"
+            stroke="currentColor"
+            strokeWidth="4"
+          />
+          <path
+            className="opacity-75"
+            fill="currentColor"
+            d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
+          />
+        </svg>
+        Loading more...
+      </div>
+    </div>
+  )
+}
+
+function EmptyState() {
+  return (
+    <div className="text-center py-16 bg-slate-900/50 rounded-xl border border-slate-800">
+      <div className="text-4xl mb-4">🌱</div>
+      <h3 className="text-lg font-medium text-slate-200">No task graphs yet</h3>
+      <p className="text-sm text-slate-500 mt-1">Run some Celery tasks to see them appear here</p>
+    </div>
+  )
+}
+
+function GraphListSkeleton() {
+  return (
+    <div className="grid gap-4 animate-pulse">
+      {[...Array(3)].map((_, i) => (
+        <div key={i} className="bg-slate-900 rounded-xl border border-slate-800 p-6">
+          <div className="h-5 w-48 bg-slate-800 rounded" />
+          <div className="h-4 w-64 bg-slate-800 rounded mt-2" />
+          <div className="h-4 w-24 bg-slate-800 rounded mt-4" />
+        </div>
+      ))}
+    </div>
+  )
+}