fix(live): keep trace git metadata and late spans visible

backend/rest/routers/live.py

@@ -8,43 +8,61 @@
 import json
 import logging
 import time
+from datetime import UTC, datetime
 
 from fastapi import APIRouter, Request
 from starlette.responses import StreamingResponse
 
 from rest.routers.deps import ProjectAccess
+from shared.config import settings
 
 logger = logging.getLogger(__name__)
 
 router = APIRouter(prefix="/projects/{project_id}/traces/{trace_id}", tags=["Live"])
 
 HEARTBEAT_INTERVAL = 15  # seconds
 MAX_STREAM_SECONDS = 600  # 10 minutes — hard ceiling for idle connections
-
-
-def _is_trace_complete_in_clickhouse(project_id: str, trace_id: str) -> bool:
-    """Return True if ClickHouse has a root span (no parent) with an end_time for this trace."""
+# Read once from settings so tests can still patch the module constant. Kept
+# above the SDK's 5s default flush interval so a quiet window can't expire
+# between two batches of the same live trace.
+TRACE_COMPLETE_QUIET_SECONDS = settings.trace_complete_quiet_seconds
+
+
+def _completion_state_in_clickhouse(
+    project_id: str, trace_id: str
+) -> tuple[datetime | None, datetime | None]:
+    """Return (root end time, last ingest time) for this trace.
+
+    The root end time is None while the root span is still open. A root span
+    with an end time is a completion candidate, not proof that no more
+    descendant spans can arrive — some streaming handlers finish the root span
+    before background work emits its children, so the last ingest time
+    (max ch_update_time) tells us whether spans are still actively arriving.
+    """
     from db.clickhouse.client import get_clickhouse_client
 
     ch_client = get_clickhouse_client()
     # Dedup ReplacingMergeTree rows without FINAL: keep the latest version per
-    # span_id (filtered to this trace), then test the root-complete condition on
-    # the deduped rows.
+    # span_id (filtered to this trace), then aggregate on the deduped rows.
     result = ch_client.query(
         """
-        SELECT count() FROM (
-            SELECT parent_span_id, span_end_time FROM spans
+        SELECT
+            maxIf(span_end_time, isNull(parent_span_id)),
+            max(ch_update_time)
+        FROM (
+            SELECT parent_span_id, span_end_time, ch_update_time FROM spans
             WHERE project_id = {project_id:String}
               AND trace_id   = {trace_id:String}
             ORDER BY ch_update_time DESC
             LIMIT 1 BY span_id
         )
-        WHERE isNull(parent_span_id)
-          AND isNotNull(span_end_time)
         """,
         parameters={"project_id": project_id, "trace_id": trace_id},
     )
-    return result.result_rows[0][0] > 0
+    rows = result.result_rows
+    if not rows:
+        return None, None
+    return rows[0][0], rows[0][1]
 
 
 @router.get("/live")
@@ -58,12 +76,12 @@ async def live_trace_stream(
 
     The client receives:
     - `event: spans` with span data as each batch is ingested
-    - `event: trace_complete` when the root span finishes
+    - `event: trace_complete` after root completion and a short quiet window
     - Heartbeat comments every 15s to keep the connection alive
 
-    For traces that are already complete when the client connects, this
-    endpoint detects that via ClickHouse and emits trace_complete immediately,
-    after forwarding any span events that arrived on Redis concurrently.
+    A root span with an end time is treated as a completion candidate. The
+    stream stays open for a short quiet window so late descendant spans still
+    reach the client before trace_complete closes the frontend stream.
     """
 
     async def event_generator():
@@ -79,53 +97,78 @@ async def event_generator():
             await pubsub.subscribe(channel)
             logger.info(f"SSE client subscribed to {channel}")
 
-            # Check whether the trace is already done in ClickHouse.
-            already_complete = await asyncio.to_thread(
-                _is_trace_complete_in_clickhouse, project_id, trace_id
+            # Check whether ClickHouse already has a root span with an end time.
+            # That starts the quiet window, but it does not immediately close the
+            # stream: distributed traces can still receive descendant spans after
+            # the root wrapper has finished.
+            root_end_time, last_ingest_time = await asyncio.to_thread(
+                _completion_state_in_clickhouse, project_id, trace_id
             )
 
-            if already_complete:
-                # Forward any span events that were published to Redis while we
-                # were doing the ClickHouse check (concurrent Celery tasks).
-                # These spans are guaranteed to be in ClickHouse already because
-                # process_s3_traces always writes to ClickHouse before publishing
-                # to Redis.
-                while True:
-                    msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=0)
-                    if msg is None:
-                        break
-                    if msg["type"] == "message":
-                        data = json.loads(msg["data"])
-                        if data.get("type") == "spans":
-                            yield f"event: spans\ndata: {msg['data']}\n\n"
-
-                yield "event: trace_complete\ndata: {}\n\n"
-                return
-
-            # Live trace: stream normally until trace_complete or timeout.
+            completion_deadline = None
+            if root_end_time is not None:
+                # Anchor the quiet window to the LATEST of root end and last
+                # span ingest: an old idle trace closes immediately instead of
+                # holding the SSE connection and Redis subscription for the
+                # full window, while a trace whose descendants are still
+                # arriving (root ended early) keeps its late-span protection.
+                # ClickHouse stores naive UTC timestamps; clamp age at 0
+                # against clock skew.
+                anchor = max(root_end_time, last_ingest_time or root_end_time)
+                now_utc = datetime.now(UTC).replace(tzinfo=None)
+                age = max(0.0, (now_utc - anchor).total_seconds())
+                completion_deadline = time.monotonic() + max(
+                    0.0, TRACE_COMPLETE_QUIET_SECONDS - age
+                )
+
+            # Live trace: stream until a completion candidate remains quiet long
+            # enough, or until the hard timeout.
             deadline = time.monotonic() + MAX_STREAM_SECONDS
 
             while time.monotonic() < deadline:
                 if await request.is_disconnected():
                     break
 
+                now = time.monotonic()
+                if completion_deadline is not None:
+                    quiet_remaining = completion_deadline - now
+                    if quiet_remaining <= 0:
+                        yield "event: trace_complete\ndata: {}\n\n"
+                        return
+                    timeout = min(HEARTBEAT_INTERVAL, quiet_remaining)
+                else:
+                    timeout = HEARTBEAT_INTERVAL
+
                 message = await pubsub.get_message(
                     ignore_subscribe_messages=True,
-                    timeout=HEARTBEAT_INTERVAL,
+                    timeout=timeout,
                 )
 
                 if message is not None and message["type"] == "message":
                     data = json.loads(message["data"])
                     event_type = data.get("type", "spans")
 
+                    if event_type == "trace_complete":
+                        completion_deadline = time.monotonic() + TRACE_COMPLETE_QUIET_SECONDS
+                        continue
+
                     yield f"event: {event_type}\ndata: {message['data']}\n\n"
 
-                    if event_type == "trace_complete":
-                        break
+                    if event_type == "spans" and completion_deadline is not None:
+                        completion_deadline = time.monotonic() + TRACE_COMPLETE_QUIET_SECONDS
                 else:
+                    if completion_deadline is not None and time.monotonic() >= completion_deadline:
+                        yield "event: trace_complete\ndata: {}\n\n"
+                        return
                     yield ": heartbeat\n\n"
             else:
-                yield "event: stream_timeout\ndata: {}\n\n"
+                if completion_deadline is not None:
+                    # A completed root kept the quiet window resetting all the
+                    # way to the hard ceiling: the trace is done, so close as
+                    # complete — the frontend's completion path refetches.
+                    yield "event: trace_complete\ndata: {}\n\n"
+                else:
+                    yield "event: stream_timeout\ndata: {}\n\n"
 
         except asyncio.CancelledError:
             pass

backend/shared/config.py

@@ -163,6 +163,11 @@ class Settings(BaseSettings):
     )
     internal_api_secret: str = ""
 
+    # Live SSE: how long a completed root span must stay quiet before the
+    # stream emits trace_complete. Must exceed the SDK's flush interval
+    # (5s default) or the window can expire between two batches of a live trace.
+    trace_complete_quiet_seconds: float = 10.0
+
     # Service-specific settings
     clickhouse: ClickHouseSettings = ClickHouseSettings()
     s3: S3Settings = S3Settings()

backend/worker/otel_transform.py

@@ -325,6 +325,7 @@ def transform_otel_to_clickhouse(
     # Shortest ids_path = closest to root. Used to correct eager trace names
     # when the first span in a batch isn't the closest-to-root span for that trace.
     _trace_name_candidates: dict[str, tuple[int, str]] = {}
+    trace_git_attrs: dict[str, dict[str, str | None]] = {}
 
     # camelCase: resourceSpans
     resource_spans = otel_data.get("resourceSpans", [])
@@ -654,9 +655,13 @@ def transform_otel_to_clickhouse(
                 # Priority: root span values overwrite, child span values only set if empty
                 span_user_id = _extract_user_id(span_attrs)
                 span_session_id = _extract_session_id(span_attrs)
+                span_git_ref = span_attrs.get("traceroot.git.ref")
+                span_git_repo = span_attrs.get("traceroot.git.repo")
 
                 if trace_id not in trace_attrs:
                     trace_attrs[trace_id] = {"user_id": None, "session_id": None}
+                if trace_id not in trace_git_attrs:
+                    trace_git_attrs[trace_id] = {"git_ref": None, "git_repo": None}
 
                 if not parent_span_id:
                     # Root span: always use its values if present (overwrites child values)
@@ -666,6 +671,12 @@ def transform_otel_to_clickhouse(
                     trace_attrs[trace_id]["session_id"] = (
                         span_session_id or trace_attrs[trace_id]["session_id"]
                     )
+                    trace_git_attrs[trace_id]["git_ref"] = (
+                        span_git_ref or trace_git_attrs[trace_id]["git_ref"]
+                    )
+                    trace_git_attrs[trace_id]["git_repo"] = (
+                        span_git_repo or trace_git_attrs[trace_id]["git_repo"]
+                    )
                 else:
                     # Child span: only set if not already set (first child wins)
                     trace_attrs[trace_id]["user_id"] = (
@@ -674,6 +685,12 @@ def transform_otel_to_clickhouse(
                     trace_attrs[trace_id]["session_id"] = (
                         trace_attrs[trace_id]["session_id"] or span_session_id
                     )
+                    trace_git_attrs[trace_id]["git_ref"] = (
+                        trace_git_attrs[trace_id]["git_ref"] or span_git_ref
+                    )
+                    trace_git_attrs[trace_id]["git_repo"] = (
+                        trace_git_attrs[trace_id]["git_repo"] or span_git_repo
+                    )
 
                 # Eager trace creation:
                 # Create a "shallow" trace record on the FIRST span we see for
@@ -692,6 +709,10 @@ def transform_otel_to_clickhouse(
                         "user_id": trace_attrs[trace_id]["user_id"],
                         "session_id": trace_attrs[trace_id]["session_id"],
                     }
+                    if trace_git_attrs[trace_id]["git_ref"] is not None:
+                        traces[trace_id]["git_ref"] = trace_git_attrs[trace_id]["git_ref"]
+                    if trace_git_attrs[trace_id]["git_repo"] is not None:
+                        traces[trace_id]["git_repo"] = trace_git_attrs[trace_id]["git_repo"]
 
                 # Track the best-known root name for this trace using the span
                 # closest to the root (shortest ids_path). Batches may contain
@@ -715,9 +736,6 @@ def transform_otel_to_clickhouse(
 
                 if not parent_span_id:
                     # Root span arrived — upgrade to full trace with rich metadata
-                    git_ref = span_attrs.get("traceroot.git.ref")
-                    git_repo = span_attrs.get("traceroot.git.repo")
-
                     traces[trace_id].update(
                         {
                             "trace_start_time": start_time,
@@ -729,10 +747,10 @@ def transform_otel_to_clickhouse(
 
                     if environment is not None:
                         traces[trace_id]["environment"] = environment
-                    if git_ref is not None:
-                        traces[trace_id]["git_ref"] = git_ref
-                    if git_repo is not None:
-                        traces[trace_id]["git_repo"] = git_repo
+                    if trace_git_attrs[trace_id]["git_ref"] is not None:
+                        traces[trace_id]["git_ref"] = trace_git_attrs[trace_id]["git_ref"]
+                    if trace_git_attrs[trace_id]["git_repo"] is not None:
+                        traces[trace_id]["git_repo"] = trace_git_attrs[trace_id]["git_repo"]
 
                     # Extract trace-level metadata
                     trace_metadata = span_attrs.get("traceroot.trace.metadata")
@@ -772,4 +790,14 @@ def transform_otel_to_clickhouse(
             if attrs["session_id"] and not traces[trace_id].get("session_id"):
                 traces[trace_id]["session_id"] = attrs["session_id"]
 
+    # Git repo/ref are stamped on every SDK span, but the root span often arrives
+    # last in live streaming. Promote the first child values so repo/ref are visible
+    # while the trace is still running, then let root values overwrite when present.
+    for trace_id, attrs in trace_git_attrs.items():
+        if trace_id in traces:
+            if attrs["git_ref"] is not None:
+                traces[trace_id]["git_ref"] = attrs["git_ref"]
+            if attrs["git_repo"] is not None:
+                traces[trace_id]["git_repo"] = attrs["git_repo"]
+
     return list(traces.values()), spans