diff --git a/trace_server/server.py b/trace_server/server.py index 3a5a828a2..e04465884 100644 --- a/trace_server/server.py +++ b/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( { diff --git a/trace_server/static/app.js b/trace_server/static/app.js index 94d521db4..eefc7e128 100644 --- a/trace_server/static/app.js +++ b/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'; } diff --git a/ymir/agents/backport_agent.py b/ymir/agents/backport_agent.py index 2d3a6cbb8..22e78dae6 100644 --- a/ymir/agents/backport_agent.py +++ b/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, diff --git a/ymir/agents/mr_consolidation_agent.py b/ymir/agents/mr_consolidation_agent.py index 928883947..3aeef7cbe 100644 --- a/ymir/agents/mr_consolidation_agent.py +++ b/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", diff --git a/ymir/agents/preliminary_testing_agent.py b/ymir/agents/preliminary_testing_agent.py index 6811c404a..f881290af 100644 --- a/ymir/agents/preliminary_testing_agent.py +++ b/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, diff --git a/ymir/agents/rebase_agent.py b/ymir/agents/rebase_agent.py index da7ad370b..3f01aaf11 100644 --- a/ymir/agents/rebase_agent.py +++ b/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, diff --git a/ymir/agents/rebuild_agent.py b/ymir/agents/rebuild_agent.py index 01fd96cdc..d5486dc27 100644 --- a/ymir/agents/rebuild_agent.py +++ b/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, diff --git a/ymir/agents/triage_agent.py b/ymir/agents/triage_agent.py index cf058acf9..6a66306dc 100644 --- a/ymir/agents/triage_agent.py +++ b/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,