Trace server improvements

trace_server/server.py

@@ -56,6 +56,7 @@
 TRACE_LOG_LEVEL     Log level: DEBUG, INFO, WARNING, ERROR (default: INFO).
 """
 
+import contextlib
 import gzip
 import io
 import json
@@ -144,6 +145,14 @@ def init_db() -> None:
         )
     """)
     db.execute("CREATE INDEX IF NOT EXISTS idx_span_issues_issue ON span_issues(jira_issue)")
+    with contextlib.suppress(sqlite3.OperationalError):
+        db.execute("ALTER TABLE spans ADD COLUMN workflow_name TEXT")
+    db.execute("CREATE INDEX IF NOT EXISTS idx_workflow_name ON spans(workflow_name)")
+    db.execute(
+        "UPDATE spans SET workflow_name = json_extract(attributes, '$.\"workflow.name\".stringValue') "
+        "WHERE workflow_name IS NULL "
+        "AND json_extract(attributes, '$.\"workflow.name\".stringValue') IS NOT NULL"
+    )
     # One-off backfill from spans.jira_issue into the new junction table
     if not db.execute("SELECT 1 FROM span_issues LIMIT 1").fetchone():
         cursor = db.execute("SELECT trace_id, span_id, jira_issue FROM spans WHERE jira_issue IS NOT NULL")
@@ -197,6 +206,7 @@ class SpanRow:
         "start_time",
         "status_code",
         "trace_id",
+        "workflow_name",
     )
 
     def __init__(
@@ -212,6 +222,7 @@ def __init__(
         jira_issues,
         agent_type,
         attributes,
+        workflow_name=None,
     ):
         self.trace_id = trace_id
         self.span_id = span_id
@@ -223,6 +234,7 @@ def __init__(
         self.jira_issues = jira_issues
         self.agent_type = agent_type
         self.attributes = attributes
+        self.workflow_name = workflow_name
 
     def as_tuple(self):
         return (
@@ -236,6 +248,7 @@ def as_tuple(self):
             ",".join(self.jira_issues) if self.jira_issues else None,
             self.agent_type,
             self.attributes,
+            self.workflow_name,
         )
 
 
@@ -341,6 +354,7 @@ def _extract_spans(otlp_data: dict) -> list[SpanRow]:
                         if name.endswith(("Agent", "Analyst"))
                         else None,
                         attributes=json.dumps(all_attrs),
+                        workflow_name=_get_val(all_attrs.get("workflow.name")),
                     )
                 )
 
@@ -358,8 +372,8 @@ def ingest_spans(otlp_data: dict) -> int:
     db.executemany(
         """INSERT OR REPLACE INTO spans
            (trace_id, span_id, parent_span_id, name, start_time, end_time,
-            status_code, jira_issue, agent_type, attributes)
-           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+            status_code, jira_issue, agent_type, attributes, workflow_name)
+           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
         [s.as_tuple() for s in spans],
     )
     span_keys = [(s.trace_id, s.span_id) for s in spans]
@@ -586,11 +600,17 @@ def query_recent_traces(since_ns: int, workflow: str | None, limit: int) -> list
     effective_limit = min(limit, MAX_LAST_TRACES)
 
     # Completed traces: root workflow span exists
-    root_conditions = ["parent_span_id = ''", "name LIKE '%Workflow'", "start_time >= ?"]
+    root_conditions = [
+        "parent_span_id = ''",
+        "name LIKE '%Workflow'",
+        "start_time >= ?",
+        "EXISTS (SELECT 1 FROM span_issues si WHERE si.trace_id = s.trace_id)",
+    ]
     root_bindings: list = [since_ns]
-    if workflow:
-        root_conditions.append("name = ?")
-        root_bindings.append(workflow)
+    wf_filter = workflow.removesuffix("Workflow").lower() if workflow else None
+    if wf_filter:
+        root_conditions.append("LOWER(name) LIKE ?")
+        root_bindings.append(f"%{wf_filter}%")
     root_where = " AND ".join(root_conditions)
 
     root_rows = db.execute(
@@ -608,13 +628,12 @@ def query_recent_traces(since_ns: int, workflow: str | None, limit: int) -> list
     # but no root Workflow span yet (NOT EXISTS uses idx_root_spans_v2).
     inprog_filter = ""
     inprog_bindings: list = [since_ns]
-    if workflow:
-        wf_base = workflow.removesuffix("Workflow").lower()
-        inprog_filter = " AND json_extract(s.attributes, '$.\"workflow.name\".stringValue') = ?"
-        inprog_bindings.append(wf_base)
+    if wf_filter:
+        inprog_filter = " AND LOWER(s.workflow_name) LIKE ?"
+        inprog_bindings.append(f"%{wf_filter}%")
     inprog_rows = db.execute(
         f"""SELECT s.trace_id, MIN(s.start_time) as first_start,
-                  MAX(json_extract(s.attributes, '$."workflow.name".stringValue')) as workflow_name
+                  MAX(s.workflow_name) as workflow_name
             FROM spans s INDEXED BY idx_start_time
             JOIN span_issues si ON s.trace_id = si.trace_id AND s.span_id = si.span_id
             WHERE s.start_time >= ?{inprog_filter}
@@ -623,6 +642,7 @@ def query_recent_traces(since_ns: int, workflow: str | None, limit: int) -> list
                 WHERE r.trace_id = s.trace_id AND r.parent_span_id = '' AND r.name LIKE '%Workflow'
               )
             GROUP BY s.trace_id
+            HAVING workflow_name IS NOT NULL
             ORDER BY first_start DESC
             LIMIT ?""",  # noqa: S608
         [*inprog_bindings, effective_limit],
@@ -674,7 +694,7 @@ def query_recent_traces(since_ns: int, workflow: str | None, limit: int) -> list
             continue
         counts = counts_by_trace.get(tid)
         wf_name = r["workflow_name"]
-        if wf_name:
+        if wf_name and not wf_name.endswith("Workflow"):
             wf_name = wf_name[0].upper() + wf_name[1:] + "Workflow"
         results.append(
             {

trace_server/static/app.js

@@ -676,8 +676,8 @@ function buildSpanTree(spans) {
     if (isRoot) {
       const wfAttr = (missingParents.get(p.span_id) || [])
         .map(c => getVal((c.attributes || {})['workflow.name'])).find(Boolean);
-      if (wfAttr) {
-        node.name = wfAttr[0].toUpperCase() + wfAttr.slice(1) + 'Workflow';
+      if (typeof wfAttr === 'string') {
+        node.name = wfAttr.endsWith('Workflow') ? wfAttr : wfAttr[0].toUpperCase() + wfAttr.slice(1) + 'Workflow';
       } else if (!node.name) {
         node.name = '(in progress)';
       }
@@ -1297,7 +1297,7 @@ function traceWorkflowName(spans) {
   if (root) return root.name;
   for (const s of spans) {
     const wf = getVal((s.attributes || {})['workflow.name']);
-    if (wf) return wf[0].toUpperCase() + wf.slice(1) + 'Workflow';
+    if (typeof wf === 'string') return wf.endsWith('Workflow') ? wf : wf[0].toUpperCase() + wf.slice(1) + 'Workflow';
   }
   return spans[0]?.name || 'trace';
 }

ymir/agents/backport_agent.py

@@ -827,7 +827,7 @@ async def main() -> None:
     ):
         upstream_patches = upstream_patches_raw.split(",")
         logger.info("Running in direct mode with environment variables")
-        with span_processor.start_transaction(jira_issue, workflow="backport"):
+        with span_processor.start_transaction(jira_issue, workflow="BackportWorkflow"):
             state = await run_workflow(
                 package=package,
                 dist_git_branch=branch,
@@ -930,7 +930,7 @@ async def retry(
 
             try:
                 logger.info(f"Starting backport processing for {backport_data.jira_issue}")
-                with span_processor.start_transaction(backport_data.jira_issue, workflow="backport"):
+                with span_processor.start_transaction(backport_data.jira_issue, workflow="BackportWorkflow"):
                     state = await run_workflow(
                         package=backport_data.package,
                         dist_git_branch=dist_git_branch,

ymir/agents/mr_consolidation_agent.py

@@ -1777,7 +1777,7 @@ async def main() -> None:
     if (package := os.getenv("PACKAGE")) and (branch := os.getenv("BRANCH")):
         release_strategy = os.getenv("RELEASE_STRATEGY", "per_commit")
         logger.info("Running in direct mode for %s/%s", package, branch)
-        with span_processor.start_transaction(None, workflow="mr_consolidation"):
+        with span_processor.start_transaction(None, workflow="MRConsolidationWorkflow"):
             state = await run_workflow(
                 package=package,
                 dist_git_branch=branch,
@@ -1821,7 +1821,7 @@ async def process_task(payload: bytes) -> None:
             try:
                 with span_processor.start_transaction(
                     jira_key,
-                    workflow="mr_consolidation",
+                    workflow="MRConsolidationWorkflow",
                 ):
                     job_strategy = job.release_strategy or os.getenv(
                         "RELEASE_STRATEGY",

ymir/agents/preliminary_testing_agent.py

@@ -404,7 +404,7 @@ async def main() -> None:
         sys.exit(1)
 
     logger.info("Running preliminary testing analysis for %s (dry_run=%s)", jira_issue, dry_run)
-    with span_processor.start_transaction(jira_issue, workflow="preliminary_testing"):
+    with span_processor.start_transaction(jira_issue, workflow="PreliminaryTestingWorkflow"):
         result = await run_preliminary_testing(
             jira_issue,
             dry_run=dry_run,

ymir/agents/rebase_agent.py

@@ -423,7 +423,7 @@ async def comment_in_jira(state):
         and (branch := os.getenv("BRANCH", None))
     ):
         logger.info("Running in direct mode with environment variables")
-        with span_processor.start_transaction(jira_issue, workflow="rebase"):
+        with span_processor.start_transaction(jira_issue, workflow="RebaseWorkflow"):
             state = await run_workflow(
                 package=package,
                 dist_git_branch=branch,
@@ -522,7 +522,7 @@ async def retry(
 
             try:
                 logger.info(f"Starting rebase processing for {rebase_data.jira_issue}")
-                with span_processor.start_transaction(rebase_data.jira_issue, workflow="rebase"):
+                with span_processor.start_transaction(rebase_data.jira_issue, workflow="RebaseWorkflow"):
                     state = await run_workflow(
                         package=rebase_data.package,
                         dist_git_branch=dist_git_branch,

ymir/agents/rebuild_agent.py

@@ -360,7 +360,7 @@ async def comment_in_jira(state):
         consolidated_raw = os.getenv("CONSOLIDATED_ISSUES", None)
         consolidated_issues = json.loads(consolidated_raw) if consolidated_raw else None
         logger.info("Running in direct mode with environment variables")
-        with span_processor.start_transaction(jira_issue, workflow="rebuild"):
+        with span_processor.start_transaction(jira_issue, workflow="RebuildWorkflow"):
             state = await run_workflow(
                 package=package,
                 dist_git_branch=branch,
@@ -483,7 +483,7 @@ async def retry(
                     await fix_await(redis.lpush(RedisQueues.ERROR_LIST.value, error))
 
             try:
-                with span_processor.start_transaction(rebuild_data.jira_issue, workflow="rebuild"):
+                with span_processor.start_transaction(rebuild_data.jira_issue, workflow="RebuildWorkflow"):
                     state = await run_workflow(
                         package=rebuild_data.package,
                         dist_git_branch=dist_git_branch,

ymir/agents/triage_agent.py

@@ -997,7 +997,7 @@ async def main() -> None:
 
     if jira_issue := os.getenv("JIRA_ISSUE", None):
         logger.info("Running in direct mode with environment variable")
-        with span_processor.start_transaction(jira_issue, workflow="triage"):
+        with span_processor.start_transaction(jira_issue, workflow="TriageWorkflow"):
             agent_factory = build_agent_factory_with_mock_repos(create_triage_agent, jira_issue)
             state = await run_workflow(
                 jira_issue,
@@ -1174,7 +1174,7 @@ async def retry(task, error, input=input, user_triggered=user_triggered):
 
             try:
                 logger.info(f"Starting triage processing for {input.issue}")
-                with span_processor.start_transaction(input.issue, workflow="triage"):
+                with span_processor.start_transaction(input.issue, workflow="TriageWorkflow"):
                     state = await run_workflow(
                         input.issue,
                         dry_run,