feat: trace status

brian14708 · brian14708 · commit 7b4cb9cff86f · 2025-10-16T13:44:35.000+08:00
diff --git a/src/duron/_core/invoke.py b/src/duron/_core/invoke.py
@@ -35,6 +35,7 @@
 from duron.log import derive_id, is_entry, random_id, set_annotations
 from duron.tracing import Tracer, current_tracer
 from duron.tracing._span import NULL_SPAN
+from duron.tracing._tracer import span
 from duron.typing import Unspecified, inspect_function
 
 if TYPE_CHECKING:
@@ -292,7 +293,8 @@ async def _invoke_prelude(
             elif type_ is StreamWriter:
                 _, extra_kwargs[name] = await ctx.create_stream(dtype, name=name)
         try:
-            result = await job_fn.fn(ctx, *args, **extra_kwargs, **kwargs)
+            with span("InvokeRun"):
+                result = await job_fn.fn(ctx, *args, **extra_kwargs, **kwargs)
             if promise_:
                 await ctx.complete_promise(promise_[0], result=None)
                 await promise_[1]
@@ -347,6 +349,8 @@ def __init__(
         self._tracer: Tracer | None = Tracer.current()
 
     async def close(self) -> None:
+        if self._tracer:
+            self._tracer.close()
         await self._send_traces(flush=True)
         if self._lease:
             await self._log.release_lease(self._lease)
@@ -387,6 +391,8 @@ async def run(self) -> object:
             return self._task.result()
 
         self._running = True
+        if self._tracer:
+            self._tracer.start()
         for msg in self._pending_msg:
             await self.enqueue_log(msg)
         self._pending_msg.clear()
diff --git a/src/duron/tracing/_events.py b/src/duron/tracing/_events.py
@@ -26,6 +26,8 @@ class SpanEnd(TypedDict):
     span_id: str
     ts: int
     attributes: NotRequired[Mapping[str, JSONValue]]
+    status: NotRequired[Literal["OK", "ERROR"]]
+    status_message: NotRequired[str]
 
 
 class Event(TypedDict):
diff --git a/src/duron/tracing/_span.py b/src/duron/tracing/_span.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final, Protocol
+from typing import TYPE_CHECKING, Final, Literal, Protocol
 from typing_extensions import Self, final
 
 if TYPE_CHECKING:
@@ -17,6 +17,13 @@ def record(
         /,
     ) -> None: ...
 
+    def set_status(
+        self,
+        status: Literal["OK", "ERROR"],
+        message: str | None = None,
+        /,
+    ) -> None: ...
+
 
 @final
 class _NullSpan:
@@ -29,6 +36,12 @@ def __enter__(self) -> Self:
     def record(_key: str, _value: JSONValue) -> None:
         return
 
+    @staticmethod
+    def set_status(
+        _status: Literal["OK", "ERROR"], _message: str | None = None
+    ) -> None:
+        return
+
     def __exit__(
         self,
         exc_type: type[BaseException] | None,
diff --git a/src/duron/tracing/_tracer.py b/src/duron/tracing/_tracer.py
@@ -1,13 +1,14 @@
 from __future__ import annotations
 
+import enum
 import logging
 import os
 import threading
 import time
 from contextvars import ContextVar
 from dataclasses import dataclass
 from hashlib import blake2b
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Literal, cast
 from typing_extensions import Self, override
 
 from duron.log import set_annotations
@@ -34,8 +35,22 @@
 )
 
 
+class TracerState(enum.Enum):
+    INIT = "init"
+    STARTED = "started"
+    CLOSED = "closed"
+
+
 class Tracer:
-    __slots__: tuple[str, ...] = ("_events", "_lock", "instance_id", "trace_id")
+    __slots__: tuple[str, ...] = (
+        "_events",
+        "_init_buffer",
+        "_lock",
+        "_open_spans",
+        "_state",
+        "instance_id",
+        "trace_id",
+    )
 
     def __init__(
         self,
@@ -48,10 +63,90 @@ def __init__(
         self.instance_id: str = instance_id or _trace_id()
         self._events: list[TraceEvent] = []
         self._lock = threading.Lock()
+        self._state: TracerState = TracerState.INIT
+        self._init_buffer: list[TraceEvent] = []
+        self._open_spans: dict[str, SpanStart] = {}
 
     def emit_event(self, event: TraceEvent) -> None:
         with self._lock:
-            self._events.append(event)
+            if self._state == TracerState.CLOSED:
+                return
+
+            if self._state == TracerState.INIT:
+                # Buffer events in INIT state
+                self._init_buffer.append(event)
+                # Track open/closed spans
+                if event["type"] == "span.start":
+                    self._open_spans[event["span_id"]] = event
+                elif event["type"] == "span.end":
+                    self._open_spans.pop(event["span_id"], None)
+            else:  # STARTED state
+                self._events.append(event)
+                # Track open spans even in STARTED state for close()
+                if event["type"] == "span.start":
+                    self._open_spans[event["span_id"]] = event
+                elif event["type"] == "span.end":
+                    self._open_spans.pop(event["span_id"], None)
+
+    def start(self) -> None:
+        """Transition to STARTED state, clear completed spans, emit remaining."""
+        with self._lock:
+            if self._state != TracerState.INIT:
+                return
+
+            # Find all span IDs that have been completed (have both start and end)
+            completed_span_ids: set[str] = set()
+            span_ends: set[str] = set()
+
+            for event in self._init_buffer:
+                if event["type"] == "span.end":
+                    span_ends.add(event["span_id"])
+
+            for event in self._init_buffer:
+                if event["type"] == "span.start" and event["span_id"] in span_ends:
+                    completed_span_ids.add(event["span_id"])
+
+            # Filter out completed spans (both start and end) and their related events
+            # Keep only span.start events for incomplete spans and other events
+            for event in self._init_buffer:
+                event_span_id = event.get("span_id")
+
+                # Skip completed span events (both start and end)
+                if event_span_id in completed_span_ids:
+                    continue
+
+                # Skip span.end events for incomplete spans (keep only starts)
+                if event["type"] == "span.end":
+                    continue
+
+                # Emit everything else (span.start for incomplete spans, other events)
+                self._events.append(event)
+
+            # Clear the init buffer
+            self._init_buffer.clear()
+
+            # Transition to STARTED state
+            self._state = TracerState.STARTED
+
+    def close(self) -> None:
+        """Transition to CLOSED state and mark all open spans as failed."""
+        with self._lock:
+            if self._state == TracerState.CLOSED or self._state == TracerState.INIT:
+                return
+
+            for span_id in list(self._open_spans.keys()):
+                end_event: SpanEnd = {
+                    "type": "span.end",
+                    "span_id": span_id,
+                    "ts": time.time_ns() // 1000,
+                    "status": "ERROR",
+                    "status_message": "tracer closed",
+                }
+                self._events.append(end_event)
+
+            self._open_spans.clear()
+
+            self._state = TracerState.CLOSED
 
     def pop_events(self, *, flush: bool) -> list[dict[str, JSONValue]]:
         with self._lock:
@@ -156,6 +251,8 @@ class _TracerSpan:
     attributes: dict[str, JSONValue] | None = None
     links: tuple[LinkRef, ...] | None = None
     _token: Token[_TracerSpan | None] | None = None
+    _status: Literal["OK", "ERROR"] | None = None
+    _status_message: str | None = None
 
     def __enter__(self) -> Self:
         start_ns = time.time_ns()
@@ -184,6 +281,12 @@ def record(self, key: str, value: JSONValue) -> None:
         else:
             self.attributes = {key: value}
 
+    def set_status(
+        self, status: Literal["OK", "ERROR"], message: str | None = None
+    ) -> None:
+        self._status = status
+        self._status_message = message
+
     def __exit__(
         self,
         exc_type: type[BaseException] | None,
@@ -199,6 +302,10 @@ def __exit__(
         if self.attributes:
             evnt["attributes"] = self.attributes
             self.attributes = None
+        if self._status:
+            evnt["status"] = self._status
+        if self._status_message:
+            evnt["status_message"] = self._status_message
 
         self.tracer.emit_event(evnt)
         if self._token:
diff --git a/tools/trace-ui/src/components/trace-view/detail-panel.tsx b/tools/trace-ui/src/components/trace-view/detail-panel.tsx
@@ -1,10 +1,12 @@
 import {
+  CheckCircle2,
   Clock,
   Hash,
   Layers,
   Link as LinkIcon,
   Tag,
   X,
+  XCircle,
   Zap,
 } from "lucide-react";
 
@@ -137,7 +139,23 @@ export function DetailPanel({
               <h3 className="text-lg font-semibold break-words text-slate-900 dark:text-slate-50">
                 {selectedSpan.name}
               </h3>
-              <div className="flex gap-2">
+              <div className="flex flex-wrap gap-2">
+                {selectedSpan.status && (
+                  <span
+                    className={`flex items-center gap-1 rounded border px-2 py-1 text-xs font-medium ${
+                      selectedSpan.status === "ERROR"
+                        ? "border-red-300 bg-red-50 text-red-800 dark:border-red-700 dark:bg-red-950 dark:text-red-200"
+                        : "border-green-300 bg-green-50 text-green-800 dark:border-green-700 dark:bg-green-950 dark:text-green-200"
+                    }`}
+                  >
+                    {selectedSpan.status === "ERROR" ? (
+                      <XCircle className="h-3 w-3" />
+                    ) : (
+                      <CheckCircle2 className="h-3 w-3" />
+                    )}
+                    {selectedSpan.status}
+                  </span>
+                )}
                 {selectedSpan.incomplete && (
                   <span className="rounded border border-red-300 bg-red-50 px-2 py-1 text-xs font-medium text-red-800 dark:border-red-700 dark:bg-red-950 dark:text-red-200">
                     Incomplete
@@ -161,6 +179,11 @@ export function DetailPanel({
                 from the trace end.
               </p>
             )}
+            {selectedSpan.status === "ERROR" && selectedSpan.statusMessage && (
+              <p className="mt-2 text-xs text-red-600 dark:text-red-400">
+                ❌ Error: {selectedSpan.statusMessage}
+              </p>
+            )}
           </div>
 
           {/* Timing Information */}
diff --git a/tools/trace-ui/src/components/trace-view/span-utils.ts b/tools/trace-ui/src/components/trace-view/span-utils.ts
@@ -3,6 +3,17 @@ import type { Span } from "@/lib/trace";
 // Color mapping for different span types
 // Soft pastel palette for visual comfort
 export const getSpanColor = (span: Span): string => {
+  // Error status takes priority over everything else
+  if (span.status === "ERROR") {
+    const errorBaseColor = "bg-red-400 dark:bg-red-500";
+    const errorGradientFrom = "from-red-400 dark:from-red-500";
+
+    if (span.incomplete) {
+      return `bg-gradient-to-r ${errorGradientFrom} to-transparent`;
+    }
+    return errorBaseColor;
+  }
+
   const type = span.attributes?.type as string | undefined;
 
   let baseColor: string;
diff --git a/tools/trace-ui/src/lib/trace.ts b/tools/trace-ui/src/lib/trace.ts
@@ -27,6 +27,8 @@ export interface Span {
   links?: Array<{ span_id: string; trace_id: string }>;
   events?: InstantEvent[]; // Instant events that occurred within this span
   incomplete?: boolean; // True if span only has start event, no end event
+  status?: "OK" | "ERROR";
+  statusMessage?: string;
 }
 
 // internal representation of a trace event from the log
@@ -40,6 +42,8 @@ interface TraceEvent {
   attributes?: Record<string, JSONValue>;
   links?: Array<{ span_id: string; trace_id: string }>;
   kind?: string;
+  status?: "OK" | "ERROR";
+  status_message?: string;
 }
 
 export interface TraceFile {
@@ -67,6 +71,8 @@ interface RawTraceEvent {
   attributes?: Record<string, JSONValue>;
   links?: Array<{ span_id: string; trace_id?: string }>;
   kind?: string;
+  status?: "OK" | "ERROR";
+  status_message?: string;
 }
 
 interface RawLog {
@@ -106,6 +112,8 @@ function rawEventToTraceEvent(
       trace_id: link.trace_id || traceId,
     })),
     kind: event.kind,
+    status: event.status,
+    status_message: event.status_message,
   };
 }
 
@@ -193,6 +201,8 @@ export function extractSpansFromEntries(entries: TraceEvent[]): Span[] {
         name: string;
         attributes?: Record<string, JSONValue>;
       }>;
+      status?: "OK" | "ERROR";
+      statusMessage?: string;
     }
   >();
 
@@ -233,6 +243,9 @@ export function extractSpansFromEntries(entries: TraceEvent[]): Span[] {
       if (event.attributes) {
         span.attributes = { ...span.attributes, ...event.attributes };
       }
+      // Capture status and status_message from span.end
+      span.status = event.status;
+      span.statusMessage = event.status_message;
     } else if (event.type === "event") {
       // Collect instant events
       const eventName = event.name ?? "event";
@@ -287,6 +300,8 @@ export function extractSpansFromEntries(entries: TraceEvent[]): Span[] {
       attributes: spanData.attributes,
       links: spanData.links,
       events: convertedEvents.length > 0 ? convertedEvents : undefined,
+      status: spanData.status,
+      statusMessage: spanData.statusMessage,
     };
 
     // Add incomplete flag if needed
