From 3fa557425811960995629c4e4781e5f1ed4b9357 Mon Sep 17 00:00:00 2001 From: frankbria Date: Tue, 7 Jul 2026 22:01:43 -0700 Subject: [PATCH 1/2] fix(costs): eliminate N+1 in /tasks endpoint (#750) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The costs /tasks endpoint opened a new SQLite connection per row for title lookup and ran a dominant-agent subquery per task — ~1000 connection opens + ~2000 queries at limit=1000. - tasks.get_titles(): batch title resolution in one WHERE id IN (...) query, replacing per-row tasks.get() in _query_top_tasks. - get_top_tasks_by_cost(): fold the dominant-agent lookup into a single windowed CTE (ROW_NUMBER over per-task agent call counts), removing the per-task subquery. Ties now break on agent_id for determinism. Query count is O(1) regardless of limit. Behavior preserved; adds tests for dominant-agent selection and batch title lookup. --- codeframe/core/tasks.py | 35 ++++++++ .../repositories/token_repository.py | 79 +++++++++++-------- codeframe/ui/routers/costs_v2.py | 28 ++++--- tests/ui/test_costs_v2.py | 56 +++++++++++++ 4 files changed, 152 insertions(+), 46 deletions(-) diff --git a/codeframe/core/tasks.py b/codeframe/core/tasks.py index 459308cc..c7b46666 100644 --- a/codeframe/core/tasks.py +++ b/codeframe/core/tasks.py @@ -190,6 +190,41 @@ def get(workspace: Workspace, task_id: str) -> Optional[Task]: return _row_to_task(row) +def get_titles(workspace: Workspace, task_ids) -> dict[str, str]: + """Batch-resolve task titles by id in a single query (#750). + + Avoids the N+1 of calling ``get()`` per id (each opens its own SQLite + connection). Returns a dict mapping id -> title for the ids that exist; + missing ids are simply absent so callers supply their own placeholder. + + Args: + workspace: Workspace to query + task_ids: Iterable of task ids (falsy entries are ignored) + + Returns: + {task_id: title} for every id that resolves to a task + """ + # Preserve order-free dedupe; drop empties so a bad id can't skew the IN(). + ids = list(dict.fromkeys(tid for tid in task_ids if tid)) + if not ids: + return {} + # ponytail: limit is capped at 1000 upstream, well under SQLite's 32766 + # bound-variable ceiling (3.32+), so one IN() is safe without chunking. + placeholders = ",".join("?" * len(ids)) + conn = get_db_connection(workspace) + try: + cursor = conn.cursor() + cursor.execute( + f"SELECT id, title FROM tasks WHERE workspace_id = ? AND id IN ({placeholders})", + (workspace.id, *ids), + ) + rows = cursor.fetchall() + finally: + conn.close() + + return {row[0]: row[1] for row in rows} + + def find_by_prefix(workspace: Workspace, prefix: str) -> list[Task]: """Resolve tasks whose id starts with ``prefix`` (SQL ``LIKE``, no cap). diff --git a/codeframe/platform_store/repositories/token_repository.py b/codeframe/platform_store/repositories/token_repository.py index f02303a8..0aa88df1 100644 --- a/codeframe/platform_store/repositories/token_repository.py +++ b/codeframe/platform_store/repositories/token_repository.py @@ -396,52 +396,61 @@ def get_top_tasks_by_cost( raise ValueError("limit must be a positive integer") start_iso, end_iso = self._window_iso_bounds(days) + # Single windowed query — the dominant agent per task is picked with + # ROW_NUMBER() instead of an N+1 per-task subquery (#750). O(1) queries + # regardless of `limit`. Ties on call count break on agent_id for a + # deterministic result (previously "arbitrary"). cursor = self.conn.cursor() cursor.execute( """ + WITH per_task AS ( + SELECT + task_id, + COALESCE(SUM(input_tokens), 0) AS input_tokens, + COALESCE(SUM(output_tokens), 0) AS output_tokens, + COALESCE(SUM(estimated_cost_usd), 0.0) AS total_cost_usd + FROM token_usage + WHERE task_id IS NOT NULL + AND timestamp >= ? + AND timestamp < ? + GROUP BY task_id + ORDER BY total_cost_usd DESC + LIMIT ? + ), + agent_ranked AS ( + SELECT + task_id, + agent_id, + ROW_NUMBER() OVER ( + PARTITION BY task_id + ORDER BY COUNT(*) DESC, agent_id ASC + ) AS rn + FROM token_usage + WHERE task_id IS NOT NULL + AND timestamp >= ? + AND timestamp < ? + GROUP BY task_id, agent_id + ) SELECT - task_id, - COALESCE(SUM(input_tokens), 0) AS input_tokens, - COALESCE(SUM(output_tokens), 0) AS output_tokens, - COALESCE(SUM(estimated_cost_usd), 0.0) AS total_cost_usd - FROM token_usage - WHERE task_id IS NOT NULL - AND timestamp >= ? - AND timestamp < ? - GROUP BY task_id - ORDER BY total_cost_usd DESC - LIMIT ? + p.task_id AS task_id, + COALESCE(a.agent_id, '') AS agent_id, + p.input_tokens AS input_tokens, + p.output_tokens AS output_tokens, + p.total_cost_usd AS total_cost_usd + FROM per_task p + LEFT JOIN agent_ranked a + ON a.task_id = p.task_id AND a.rn = 1 + ORDER BY p.total_cost_usd DESC """, - (start_iso, end_iso, limit), + (start_iso, end_iso, limit, start_iso, end_iso), ) rows = cursor.fetchall() - # TODO(perf): the dominant-agent lookup is N+1 against the limit. - # Acceptable at limit=10 (analytics view) and even limit=1000 (badge - # map for a board). Fold into a single CTE if the cap grows further. result: List[Dict[str, Any]] = [] for row in rows: - task_id = row["task_id"] - # Find the most-used agent for this task in the same window. - cursor.execute( - """ - SELECT agent_id, COUNT(*) AS calls - FROM token_usage - WHERE task_id = ? - AND timestamp >= ? - AND timestamp < ? - GROUP BY agent_id - ORDER BY calls DESC - LIMIT 1 - """, - (task_id, start_iso, end_iso), - ) - agent_row = cursor.fetchone() - agent_id = agent_row["agent_id"] if agent_row else "" - result.append({ - "task_id": task_id, - "agent_id": agent_id, + "task_id": row["task_id"], + "agent_id": row["agent_id"], "input_tokens": int(row["input_tokens"] or 0), "output_tokens": int(row["output_tokens"] or 0), "total_cost_usd": float(row["total_cost_usd"] or 0.0), diff --git a/codeframe/ui/routers/costs_v2.py b/codeframe/ui/routers/costs_v2.py index cd0b0e75..1a055a92 100644 --- a/codeframe/ui/routers/costs_v2.py +++ b/codeframe/ui/routers/costs_v2.py @@ -223,18 +223,24 @@ def _query_top_tasks( finally: conn.close() + # Batch-resolve titles in one query instead of one connection per row (#750). + task_id_strs = [ + str(row["task_id"]) if row["task_id"] is not None else "" for row in rows + ] + try: + titles = tasks_module.get_titles(workspace, task_id_strs) + except Exception: + # Lookup failure is non-fatal — every row falls back to a placeholder. + logger.debug("costs/tasks: batch title lookup failed", exc_info=True) + titles = {} + entries: List[Dict[str, Any]] = [] - for row in rows: - raw_id = row["task_id"] - task_id_str = str(raw_id) if raw_id is not None else "" - title = _placeholder_task_title(task_id_str) - try: - task = tasks_module.get(workspace, task_id_str) - if task is not None: - title = task.title - except Exception: - # Lookup failures are non-fatal — keep the placeholder title. - logger.debug("costs/tasks: task lookup failed for %s", task_id_str, exc_info=True) + for row, task_id_str in zip(rows, task_id_strs): + title = ( + titles[task_id_str] + if task_id_str in titles + else _placeholder_task_title(task_id_str) + ) entries.append({ "task_id": task_id_str, diff --git a/tests/ui/test_costs_v2.py b/tests/ui/test_costs_v2.py index d841ac7a..0080c3ea 100644 --- a/tests/ui/test_costs_v2.py +++ b/tests/ui/test_costs_v2.py @@ -293,6 +293,62 @@ def test_aggregates_by_agent(self, test_client): assert agents[1]["agent_id"] == "codex" +class TestCostsTasksDominantAgent: + """Dominant-agent selection is now a single windowed query (#750).""" + + def test_reports_most_used_agent_per_task(self, test_client): + from codeframe.core import tasks as tasks_module + + workspace = test_client.workspace + task = tasks_module.create(workspace, title="Multi-agent", description="") + # codex made 3 calls, claude-code made 1 → codex is dominant. + for _ in range(3): + _record_usage_text_task(workspace, task_id=task.id, agent_id="codex", cost=0.01) + _record_usage_text_task(workspace, task_id=task.id, agent_id="claude-code", cost=0.01) + + response = test_client.get("/api/v2/costs/tasks") + entry = response.json()["tasks"][0] + assert entry["task_id"] == task.id + assert entry["agent_id"] == "codex" + + def test_multiple_tasks_each_get_own_dominant_agent(self, test_client): + from codeframe.core import tasks as tasks_module + + workspace = test_client.workspace + t1 = tasks_module.create(workspace, title="Task one", description="") + t2 = tasks_module.create(workspace, title="Task two", description="") + _record_usage_text_task(workspace, task_id=t1.id, agent_id="codex", cost=1.00) + _record_usage_text_task(workspace, task_id=t2.id, agent_id="claude-code", cost=0.50) + + response = test_client.get("/api/v2/costs/tasks") + by_id = {e["task_id"]: e for e in response.json()["tasks"]} + assert by_id[t1.id]["agent_id"] == "codex" + assert by_id[t2.id]["agent_id"] == "claude-code" + assert by_id[t1.id]["task_title"] == "Task one" + assert by_id[t2.id]["task_title"] == "Task two" + + +class TestGetTitlesBatch: + """tasks.get_titles resolves many ids in one query (#750).""" + + def test_batch_resolves_existing_and_omits_missing(self, test_workspace): + from codeframe.core import tasks as tasks_module + + a = tasks_module.create(test_workspace, title="Alpha", description="") + b = tasks_module.create(test_workspace, title="Beta", description="") + + titles = tasks_module.get_titles( + test_workspace, [a.id, b.id, "does-not-exist"] + ) + assert titles == {a.id: "Alpha", b.id: "Beta"} + + def test_empty_input_returns_empty_dict(self, test_workspace): + from codeframe.core import tasks as tasks_module + + assert tasks_module.get_titles(test_workspace, []) == {} + assert tasks_module.get_titles(test_workspace, [""]) == {} + + class TestCostsTasksDaysValidation: def test_below_minimum_rejected(self, test_client): response = test_client.get("/api/v2/costs/tasks?days=0") From c2b66a916351c5ac5d9987aeccfe05cb7a224bdc Mon Sep 17 00:00:00 2001 From: frankbria Date: Tue, 7 Jul 2026 22:11:04 -0700 Subject: [PATCH 2/2] chore(costs): log batch title-lookup failure at warning (CodeRabbit) A batch failure degrades every row's title to a placeholder, so surface it at warning level for parity with the sqlite3.Error path above. --- codeframe/ui/routers/costs_v2.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/codeframe/ui/routers/costs_v2.py b/codeframe/ui/routers/costs_v2.py index 1a055a92..79ef20cc 100644 --- a/codeframe/ui/routers/costs_v2.py +++ b/codeframe/ui/routers/costs_v2.py @@ -231,7 +231,9 @@ def _query_top_tasks( titles = tasks_module.get_titles(workspace, task_id_strs) except Exception: # Lookup failure is non-fatal — every row falls back to a placeholder. - logger.debug("costs/tasks: batch title lookup failed", exc_info=True) + # Log at warning (not debug): a batch failure degrades ALL titles, so an + # operator should see it — parity with the sqlite3.Error path above. + logger.warning("costs/tasks: batch title lookup failed", exc_info=True) titles = {} entries: List[Dict[str, Any]] = []