External/roboco
1
0
Fork 0
mirror of https://github.com/rennf93/roboco.git synced 2026-08-03 07:23:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

Chore: reduce analytics complexity (#100)

Browse Source 
* refactor(analytics): reduce cyclomatic complexity in usage/pricing/rollup

Collapse the three near-identical get_by_* aggregation methods in
UsageService into a shared _aggregate_by helper parameterized by group
column and key name, and centralize token null-coalescing in a
_row_tokens helper. Extract the per-row upsert in _sweep_daily_rollup
into _upsert_rollup_row, and the pricing-table lookup into
_lookup_prices. All blocks now rank <= B and both modules rank A, so the
xenon gate passes; behavior is unchanged and existing tests stay green.

* feat(billing): make token pricing provider-aware

Distinguish three cases when a model has no per-token rate: a non-Anthropic
model (local Ollama, or an Ollama Cloud ":cloud" model billed by flat
subscription / GPU-time) legitimately has no per-token cost and returns 0.0
silently; an unpriced Anthropic ("claude"-named) model also returns 0.0 but
logs a warning, since that is real spend being undercounted and catches new or
renamed Claude models missing from the table. Folds the old ollama/ prefix
special-case into the general non-Anthropic path so there is one code path,
and replaces the blanket 'no pricing data' warning that fired even for
self-hosted models.

* fix(tasks): preserve ownership when force-unclaiming to pending

The stale-claim reaper and the dependency-blocked release both routed through
_force_unclaim_to_pending, which nulled assigned_to and left the task in a
pending state owned by nobody — no dispatcher re-spawns an ownerless pending
task, so it went dormant. The dispatcher-side claimed_by fallback only masked
half the cases.

Capture the owner before releasing the claim and keep both assigned_to and
claimed_by pointed at it (mirroring the unblock restore), releasing only the
live claim (active_claimant_id + heartbeat) and the WorkSession. The same agent
now resumes the task once it re-dispatches. Updates the reaper test that
asserted the old orphaning behavior and adds owner-preservation coverage for
both the reaper and dependency-release paths.

* fix(tasks): unblock restores the owner into both ownership fields

Audit follow-up to the force-unclaim ownership fix. unblock() only restored
assigned_to from blocker_raised_by, which block() stashes solely from
assigned_to. A task claimed via give_me_work (claimed_by set, assigned_to null)
therefore unblocked into a split-owner state — assigned_to null but claimed_by
set — that both the dev dispatcher and the PM pool-router race to pick up. It
also left claimed_by pointing at the resolver PM after an escalation.

Resolve the owner as blocker_raised_by or assigned_to or claimed_by and write
it to both fields, matching the force-unclaim and reassign convention so the
original worker resumes cleanly. Adds coverage for the give_me_work-claim case
and asserts owner restoration on the existing in_progress-resume test.

* test(orchestrator): cover dev owner resolution and the claimed_by fallback

_resolve_dev_owner_uuid had no coverage. Add the status-dependent precedence
(claimed/blocked prefer the live claimant; other statuses prefer the
PM-assigned owner) and the half-reap fallback where a pending task with
assigned_to nulled still resolves its owner from claimed_by instead of going
dormant.

* fix(tasks): wire the pre-block snapshot so unblock(restore=True) works

The restore=True path on a PM unblock was a no-op: pre_block_state /
pre_block_assignee (migration 006) were read by unblock_with_restore but never
written, so it always fell through to legacy unblock() and the restore flag did
nothing.

Snapshot the resting status + owner at every block entry (dependency block,
soft block, escalation) before mutating, capturing only the first block in a
chain so a re-block doesn't overwrite the original state. Escalation snapshots
the outgoing owner, not the escalation target, so restore returns the original
worker. The restore path applies the same branchless guard legacy unblock()
relies on — a snapshotted in_progress with no branch diverts to pending instead
of looping the dispatcher — and is extracted into _apply_pre_block_restore to
keep complexity under the gate. Adds coverage for snapshot capture, restore,
the branchless divert, and escalation owner restoration.

* test(tasks): update orphan-reconciler and dependency-release tests for owner preservation

Both the startup orphan reconciler and the dependency-blocked claim release
route through unclaim_for_reaper / _force_unclaim_to_pending, which now preserve
the owner instead of nulling assigned_to. Update the two tests that asserted the
old orphaning behavior to assert the owner is kept (so the same agent resumes)
while the live claim is released.

* chore(tests): scrub internal work-item labels from test names, docstrings, comments

Rename four test files that carried audit work-item IDs in their filenames
(test_p0_7_branch_atomicity, test_p2_8_orphan_reconciler,
test_p2_9_autogen_prompt_layer, test_p2_7_attempt_id) to describe what they
test, and strip the matching P-/D-/S- cluster labels from docstrings, comments,
and assertion messages across the test suite and two orchestrator comments.
These are internal references with no meaning in the codebase; behavior is
unchanged.

* style: reformat assertion line shortened by the internal-ref scrub

* build: waive unreachable torch CVE-2025-3000 in pip-audit gate

torch is a transitive CPU-pinned dep (piragi / sentence-transformers) never
loaded at runtime — the stack uses Ollama over HTTP for all embeddings/LLM, so
the vulnerable torch.jit.script path is unreachable. CVE-2025-3000 is MEDIUM,
local-only, with no published fix. Documented --ignore-vuln waiver; revisit when
a fixed torch ships.

* fix(orchestrator): route unplaceable pending tasks to main-pm instead of dropping them

_get_routing_target returned None when a 'dev'-classified task had no cell
agent (no team, or a non-cell team like fullstack/system) or when the routing
classification was unrecognized. _route_unassigned_pm_task logged 'no routing
target found' and returned, leaving the task ownerless and pending — and no
dispatcher re-spawns an unrouted pending task, so it went dormant for 10+ min
until the stuck-task detector caught it.

Fall back to main-pm (the same default cell_pm routing and escalation already
use) so the task is always owned and triaged, never stranded. Logs the fallback
so unplaceable tasks stay visible. Adds a test asserting no (routing, team)
combination ever resolves to None.

* fix(panel): make intake chat markdown inherit the bubble's text color

MarkdownBody is shared by the assistant (text-foreground) and user
(text-primary-foreground) bubbles. [&_*]:!text-inherit only colored the prose
div's descendants, so the prose div itself kept the prose typography body color
(gray) and children inherited that — unreadable on the muted assistant bubble.
Add !text-inherit on the prose div itself so it inherits the bubble's color
too; descendants then inherit the correct foreground. Fixes both bubbles without
hardcoding a color.

* fix(prompter): keep a board-reviewed product on the board team so Approve & Start shows

A product coordination root confirmed via 'Board review & Start' is assigned to
a board reviewer (product-owner) for review, but create_task_from_draft set
team=main_pm for every product unconditionally. The CEO's Approve & Start gate
keys on team=board, so the button never appeared — and because the owner stayed
a board agent while the team said main_pm, the dispatcher routed it to the board
path (nothing left to do after review) and the task stranded at pending, with
the board agent fruitlessly trying to escalate it up.

Route a product by its assignee: a board reviewer keeps it team=board (so the
gate appears and approve_and_start later hands it to Main PM), while a main-pm
assignee — the 'Approve & Start' straight-through path — is team=main_pm. Adds
_assignee_is_board mirroring TaskService's board-role check, and a test.

---------

Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
2026-06-11 04:36:17 +02:00
committed by GitHub
co-authored by Renn F
parent b3057628b0
commit ff35a646fa
26 changed files with 814 additions and 353 deletions
Download Patch File Download Diff File Expand all files Collapse all files
roboco/billing/pricing.py
+52 -20
View File
@@ -4,9 +4,20 @@ Token pricing for Claude API models.
Implements per-model USD cost calculation based on Anthropic's published
pricing. All prices are in USD per 1 million tokens.
Unknown model names return 0.0 without raising so callers don't need to
guard against missing pricing data. Self-hosted Ollama models always
return 0.0 (no API cost) — matched by the ``ollama/`` prefix convention.
Pricing is provider-aware. A model name resolves to one of three cases:
* **Anthropic** — priced from the table below by substring match.
* **Non-Anthropic** — local self-hosted Ollama models (``ollama/`` prefix or
bare model tags) and Ollama Cloud models (``:cloud`` tag). These have **no
per-token cost**: local inference runs on owned hardware, and Ollama Cloud
is billed by flat subscription / GPU-time rather than per token. Both return
an intentional ``0.0`` — not an error condition, so they are not warned on.
* **Unpriced Anthropic** — a ``claude``-named model with no table entry (a new
or renamed Claude model we have not priced yet). This is real spend we would
otherwise undercount, so it logs a warning and returns ``0.0``.
Every path returns ``0.0`` rather than raising, so callers don't need to guard
against missing pricing data.
"""
from __future__ import annotations
@@ -15,6 +26,11 @@ import structlog
logger = structlog.get_logger(__name__)
# Substrings that identify an Anthropic (Claude) model. Used only to decide
# whether an *unpriced* model is a Claude model we forgot to price (warn) vs a
# non-Anthropic model that legitimately has no per-token cost (don't warn).
_ANTHROPIC_FRAGMENTS = ("claude", "opus", "sonnet", "haiku")
# ---------------------------------------------------------------------------
# Per-model pricing table
# Format: model_name_fragment → (input_usd_per_1m, output_usd_per_1m,
@@ -48,6 +64,26 @@ _PRICING: list[tuple[str, float, float, float, float]] = [
_MILLION = 1_000_000.0
def _is_anthropic_model(lower: str) -> bool:
"""True if the lowercased model name looks like an Anthropic (Claude) model."""
return any(fragment in lower for fragment in _ANTHROPIC_FRAGMENTS)
def _lookup_prices(lower: str) -> tuple[float, float, float, float] | None:
"""Return the (input, output, cache_read, cache_write) rates for a model.
Matches ``lower`` (a lowercased model name) against the pricing table by
substring, longest fragment wins. Returns None when no fragment matches.
"""
best_fragment_len = 0
best_prices: tuple[float, float, float, float] | None = None
for fragment, inp_price, out_price, cr_price, cw_price in _PRICING:
if fragment in lower and len(fragment) > best_fragment_len:
best_fragment_len = len(fragment)
best_prices = (inp_price, out_price, cr_price, cw_price)
return best_prices
def calculate_cost(
model: str,
tokens_input: int,
@@ -58,8 +94,10 @@ def calculate_cost(
"""Calculate the estimated USD cost for a model invocation.
Matches the model name against the known pricing table using substring
search (longest match wins). Unknown models return 0.0 without raising.
Self-hosted Ollama models (``ollama/`` prefix) always return 0.0.
search (longest match wins). Provider-aware (see module docstring):
non-Anthropic models (local Ollama, Ollama Cloud) have no per-token cost
and return 0.0 silently; an unpriced Anthropic model returns 0.0 but logs
a warning since it represents real spend we are failing to count.
Args:
model: Model name or short alias (e.g. ``"claude-sonnet-4-6"``,
@@ -70,7 +108,7 @@ def calculate_cost(
tokens_cache_write: Prompt-cache write tokens (charged at reduced rate).
Returns:
Estimated cost in USD as a float. Returns 0.0 for unknown models
Estimated cost in USD as a float. Returns 0.0 for unpriced models
rather than raising.
"""
if not model:
@@ -78,21 +116,15 @@ def calculate_cost(
lower = model.lower()
# Self-hosted Ollama models have no API cost.
if lower.startswith("ollama/"):
return 0.0
# Find the best (longest fragment) match
best_fragment_len = 0
best_prices: tuple[float, float, float, float] | None = None
for fragment, inp_price, out_price, cr_price, cw_price in _PRICING:
if fragment in lower and len(fragment) > best_fragment_len:
best_fragment_len = len(fragment)
best_prices = (inp_price, out_price, cr_price, cw_price)
best_prices = _lookup_prices(lower)
if best_prices is None:
logger.warning("No pricing data found for model", model=model)
# No per-token rate. Warn only for Anthropic models (real spend we are
# undercounting); non-Anthropic models are local/subscription-billed
# and have no per-token cost, so an intentional 0.0 is correct.
if _is_anthropic_model(lower):
logger.warning("No pricing data found for Anthropic model", model=model)
else:
logger.debug("Non-Anthropic model has no per-token cost", model=model)
return 0.0
inp_price, out_price, cr_price, cw_price = best_prices
roboco/runtime/orchestrator.py
+73 -61
View File
@@ -588,8 +588,8 @@ class AgentOrchestrator:
# Self-heal: roll back orphan claims left over from a prior crash.
# Tasks that show CLAIMED/IN_PROGRESS but have NO
# branch_name set indicate _finalize_claim flushed the status before
# branch creation failed in a pre-P0-7 run. Without this, the next
# claim attempt fails non-idempotent on `git checkout -b`.
# branch creation failed (before claim-rollback was atomic). Without
# this, the next claim attempt fails non-idempotent on `git checkout -b`.
await self._reconcile_orphan_claims_on_startup()
# Note: Per-agent settings are now generated at spawn time
@@ -3330,7 +3330,7 @@ class AgentOrchestrator:
"""
try:
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentSpawnSessionTable, DailyUsageRollupTable
from roboco.db.tables import AgentSpawnSessionTable
except ImportError:
return
@@ -3382,59 +3382,7 @@ class AgentOrchestrator:
rows = result.fetchall()
for row in rows:
date_val = row.date
agent_slug = row.agent_slug
team = row.team
model = row.model
# Look for existing rollup row
existing_result = await db.execute(
select(DailyUsageRollupTable).where(
DailyUsageRollupTable.date == date_val,
DailyUsageRollupTable.agent_slug == agent_slug,
DailyUsageRollupTable.team == team,
DailyUsageRollupTable.model == model,
)
)
existing = existing_result.scalar_one_or_none()
tokens_input = int(row.tokens_input or 0)
tokens_output = int(row.tokens_output or 0)
tokens_cache_read = int(row.tokens_cache_read or 0)
tokens_cache_write = int(row.tokens_cache_write or 0)
total_cost = float(row.total_cost_usd or 0.0)
session_count = int(row.session_count or 0)
if existing is not None:
from sqlalchemy import update
await db.execute(
update(DailyUsageRollupTable)
.where(DailyUsageRollupTable.id == existing.id)
.values(
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
total_cost_usd=total_cost,
session_count=session_count,
)
)
else:
new_row = DailyUsageRollupTable(
id=_uuid4(),
date=date_val,
agent_slug=agent_slug,
team=team,
model=model,
tokens_input=tokens_input,
tokens_output=tokens_output,
tokens_cache_read=tokens_cache_read,
tokens_cache_write=tokens_cache_write,
total_cost_usd=total_cost,
session_count=session_count,
)
db.add(new_row)
await self._upsert_rollup_row(db, row, _uuid4)
await db.commit()
logger.debug("Daily usage rollup complete", rows_processed=len(rows))
@@ -3442,6 +3390,50 @@ class AgentOrchestrator:
except Exception as exc:
logger.warning("Daily usage rollup failed", error=str(exc))
async def _upsert_rollup_row(self, db: Any, row: Any, uuid4: Any) -> None:
"""Insert or update a single daily_usage_rollups row from an aggregate.
Looks up the existing rollup for (date, agent_slug, team, model) and
either updates its summed columns or inserts a fresh row.
"""
from sqlalchemy import select, update
from roboco.db.tables import DailyUsageRollupTable
key = {
"date": row.date,
"agent_slug": row.agent_slug,
"team": row.team,
"model": row.model,
}
values = {
"tokens_input": int(row.tokens_input or 0),
"tokens_output": int(row.tokens_output or 0),
"tokens_cache_read": int(row.tokens_cache_read or 0),
"tokens_cache_write": int(row.tokens_cache_write or 0),
"total_cost_usd": float(row.total_cost_usd or 0.0),
"session_count": int(row.session_count or 0),
}
existing_result = await db.execute(
select(DailyUsageRollupTable).where(
DailyUsageRollupTable.date == key["date"],
DailyUsageRollupTable.agent_slug == key["agent_slug"],
DailyUsageRollupTable.team == key["team"],
DailyUsageRollupTable.model == key["model"],
)
)
existing = existing_result.scalar_one_or_none()
if existing is not None:
await db.execute(
update(DailyUsageRollupTable)
.where(DailyUsageRollupTable.id == existing.id)
.values(**values)
)
else:
db.add(DailyUsageRollupTable(id=uuid4(), **key, **values))
async def restore_waiting_records(self) -> int:
"""Load persisted waiting records into memory on orchestrator start.
@@ -4502,11 +4494,31 @@ Start by:
if routing == "cell_pm":
return self._TEAM_PM_MAP.get(team, "main-pm") if team else "main-pm"
# Dev routing - requires agent selection
if routing == "dev" and team:
return self._select_agent_for_cell(team, "dev")
# Dev routing - select a cell agent.
if routing == "dev":
agent = self._select_agent_for_cell(team, "dev") if team else None
if agent:
return agent
# No cell agent — team is missing or a non-cell team (fullstack /
# system). Fall back to main-pm to triage rather than leaving the
# task ownerless-and-dormant: the dispatcher never re-spawns an
# unrouted pending task, so a None here strands it. Mirrors the
# cell_pm / escalation `... or "main-pm"` default.
logger.warning(
"dev routing found no cell agent; falling back to main-pm",
task_id=task.get("id"),
team=team,
)
return "main-pm"
return None
# Unrecognized routing classification — never strand the task; main-pm
# triages it instead of it going dormant.
logger.warning(
"unrecognized routing classification; falling back to main-pm",
routing=routing,
task_id=task.get("id"),
)
return "main-pm"
def _build_main_pm_triage_prompt(self, task: dict[str, Any]) -> str:
"""Build prompt for MAIN PM to triage and distribute to Cell PMs."""
@@ -4728,7 +4740,7 @@ Start now: evidence(task_id="{task_id}")
A task in CLAIMED/IN_PROGRESS with ``branch_name IS NULL`` is an
orphan: ``_finalize_claim`` flushed the status before branch creation
failed (or before the P0-7 rollback fix landed). The next claim then
failed (or before claim rollback became atomic). The next claim then
fails non-idempotent on ``git checkout -b`` because the on-disk
branch may exist while the DB state is stale.
roboco/services/prompter.py
+41 -10
View File
@@ -26,13 +26,21 @@ from sqlalchemy import select
from roboco.config import settings
from roboco.db.tables import (
AgentTable,
PrompterMessageTable,
PrompterSessionTable,
TaskDraftTable,
TaskTable,
)
from roboco.foundation.identity import CELL_TEAMS
from roboco.models.base import Complexity, TaskNature, TaskStatus, TaskType, Team
from roboco.models.base import (
AgentRole,
Complexity,
TaskNature,
TaskStatus,
TaskType,
Team,
)
from roboco.models.task import TaskCreateRequest
from roboco.services.base import NotFoundError, ServiceError, ValidationError
@@ -41,6 +49,13 @@ if TYPE_CHECKING:
logger = structlog.get_logger()
# A board/advisory assignee means a product coordination root is still in board
# review — it stays team=board until the CEO's Approve & Start hands it to Main
# PM. Mirrors `_BOARD_ADVISORY_ROLES` in TaskService.
_BOARD_REVIEW_ROLES: frozenset[AgentRole] = frozenset(
{AgentRole.PRODUCT_OWNER, AgentRole.HEAD_MARKETING, AgentRole.AUDITOR}
)
# ---------------------------------------------------------------------------
# Input types
@@ -389,6 +404,13 @@ class PrompterService:
)
return UUID(str(task.id))
async def _assignee_is_board(self, agent_id: UUID) -> bool:
"""True if ``agent_id`` is a board/advisory role (PO / marketing / auditor)."""
result = await self._session.execute(
select(AgentTable.role).where(AgentTable.id == agent_id)
)
return result.scalar_one_or_none() in _BOARD_REVIEW_ROLES
async def create_task_from_draft(
self,
draft_data: dict[str, Any],
@@ -434,21 +456,30 @@ class PrompterService:
_lead, task_type, nature, complexity = self._coerce_draft_enums(draft_data)
# Adaptive routing: a product target is a board-led coordination root
# owned by the Main PM (who fans out per cell); a project target is a
# single-cell executable task owned by the cell doing the work.
if resolved_product_id is not None:
team = Team.MAIN_PM
else:
team = self._lead_cell_team(draft_data, default=_lead)
# Explicit assignment (from the confirm button) wins; else fall back to
# any assignee carried on the draft.
# any assignee carried on the draft. Resolved before team routing — the
# owner decides the team for a product.
resolved_assigned_to: UUID | None = assigned_to
if resolved_assigned_to is None and draft_data.get("assigned_to"):
with contextlib.suppress(ValueError):
resolved_assigned_to = UUID(str(draft_data["assigned_to"]))
# Adaptive routing. A project target is a single-cell executable task
# owned by the lead cell. A product target is a board-led coordination
# root whose team follows the start mode (encoded in the assignee): the
# "Board review & Start" path assigns a board reviewer, so it must stay
# team=board until approved — otherwise the CEO's Approve & Start gate,
# which keys on team=board, never appears and the task strands. "Approve
# & Start" (assignee main-pm) and the post-approval state are team=main_pm.
if resolved_product_id is None:
team = self._lead_cell_team(draft_data, default=_lead)
elif resolved_assigned_to is not None and await self._assignee_is_board(
resolved_assigned_to
):
team = Team.BOARD
else:
team = Team.MAIN_PM
req = TaskCreateRequest(
title=draft_data["title"],
description=draft_data["description"],
roboco/services/task.py
+79 -12
View File
@@ -2283,17 +2283,32 @@ class TaskService(BaseService):
Shared core of ``unclaim_for_reaper`` and
``release_dependency_blocked_claim``. Routes through
``_validate_and_set_status`` so the state machine records the
transition, clears assignee/heartbeat/claimant, and abandons the active
WorkSession (best-effort, tagged with ``reason``) so a re-claim doesn't
trip the uniqueness constraint. Bypasses ownership/role checks — the
system itself is performing the transition. Returns True iff the task
was actually released (False when missing or not in a releasable state).
transition, releases the live claim (heartbeat + active claimant) and
abandons the active WorkSession (best-effort, tagged with ``reason``) so
a re-claim doesn't trip the uniqueness constraint. Bypasses
ownership/role checks — the system itself is performing the transition.
Returns True iff the task was actually released (False when missing or
not in a releasable state).
Ownership is **preserved**, not cleared: both callers release a task its
owner should resume — the reaper's holder is dead but will respawn, and
a dependency-blocked task continues with the same agent once the
upstream lands. Leaving ``assigned_to``/``claimed_by`` pointed at the
owner (mirroring the unblock restore) keeps the task from landing in an
ownerless ``pending`` limbo that no dispatcher re-spawns. The earlier
behaviour nulled ``assigned_to``, which is exactly that orphaning bug.
"""
task = await self.get(task_id)
if task is None:
return False
if task.status not in (TaskStatus.CLAIMED, TaskStatus.IN_PROGRESS):
return False
# Capture the owner before releasing the claim. A claimed/in_progress
# task is always owned via claimed_by/active_claimant_id (and usually
# assigned_to); fall back across them so the row never goes ownerless.
owner = cast(
"Any", task.assigned_to or task.claimed_by or task.active_claimant_id
)
try:
self._validate_and_set_status(task, TaskStatus.PENDING, None)
except TaskLifecycleError:
@@ -2303,7 +2318,8 @@ class TaskService(BaseService):
task.work_session_id, reason=reason
)
task.work_session_id = cast("Any", None)
task.assigned_to = cast("Any", None)
task.assigned_to = owner
task.claimed_by = owner
task.last_heartbeat_at = None
task.active_claimant_id = cast("Any", None)
await self.session.flush()
@@ -2494,6 +2510,26 @@ class TaskService(BaseService):
except TaskLifecycleError:
return None
def _snapshot_pre_block(self, task: TaskTable) -> None:
"""Record the pre-block status + owner so unblock(restore=True) works.
Captures the resting state a task is leaving so a PM ``unblock`` with
``restore=True`` can return it exactly there. Call this *before*
mutating status/ownership at every block entry (dependency block,
soft block, escalation). Only the first block in a chain snapshots —
a re-block (e.g. escalating an already-blocked task) must not overwrite
the original resting state with ``blocked``. Mirrors the ``not already
set`` guard used for ``blocker_raised_by``.
"""
if task.pre_block_state:
return
task.pre_block_state = (
task.status.value
if isinstance(task.status, TaskStatus)
else str(task.status)
)
task.pre_block_assignee = cast("Any", task.assigned_to or task.claimed_by)
async def block(
self,
task_id: UUID,
@@ -2522,6 +2558,7 @@ class TaskService(BaseService):
# the unblock path has a consistent source of truth.
if task.assigned_to and not task.blocker_raised_by:
task.blocker_raised_by = task.assigned_to
self._snapshot_pre_block(task)
self._validate_and_set_status(task, TaskStatus.BLOCKED, agent_role)
await self.session.flush()
@@ -2603,6 +2640,7 @@ class TaskService(BaseService):
# Remember the raiser so `unblock` can restore the task to them.
if task.assigned_to and not task.blocker_raised_by:
task.blocker_raised_by = task.assigned_to
self._snapshot_pre_block(task)
self._validate_and_set_status(task, TaskStatus.BLOCKED, agent_role)
await self.session.flush()
@@ -2652,12 +2690,22 @@ class TaskService(BaseService):
if task.status != TaskStatus.BLOCKED:
return None
# Restore the raiser so the orchestrator dispatcher (which
# includes `in_progress` tasks in its pickup list) respawns the
# original agent — not the PM who merely resolved the block.
if task.blocker_raised_by:
task.assigned_to = cast("Any", task.blocker_raised_by)
task.blocker_raised_by = None
# Restore ownership so the dispatcher respawns the original worker, not
# the PM who merely resolved the block. blocker_raised_by holds the
# pre-escalation dev; fall back to the surviving claim owner so a task
# claimed via give_me_work (which has no assigned_to to stash) is never
# left with a split owner — assigned_to null but claimed_by set, which
# the dev dispatcher and the PM pool-router would both try to grab.
# Keep both fields on the owner, mirroring _force_unclaim_to_pending and
# reassign; this also clears a stale claimed_by left pointing at the
# resolver PM after an escalation.
owner = cast(
"Any", task.blocker_raised_by or task.assigned_to or task.claimed_by
)
task.blocker_raised_by = None
if owner is not None:
task.assigned_to = owner
task.claimed_by = owner
# Clear resolver metadata — only meaningful while BLOCKED.
task.blocker_resolver_type = None
# A task with a branch was claimed before it blocked, so resume it
@@ -3618,6 +3666,7 @@ class TaskService(BaseService):
else str(task.status)
)
pre_block_owner = cast("Any", task.claimed_by)
self._snapshot_pre_block(task)
task.assigned_to = cast("Any", target_agent_id)
task.claimed_by = cast("Any", target_agent_id)
task.status = TaskStatus.BLOCKED
@@ -5891,6 +5940,24 @@ class TaskService(BaseService):
except ValueError:
return await self.unblock(task_id, agent_role="cell_pm")
return await self._apply_pre_block_restore(task, restored_status)
async def _apply_pre_block_restore(
self, task: TaskTable, restored_status: TaskStatus
) -> TaskTable:
"""Restore a blocked task to its snapshotted status + owner.
Sets status directly (bypassing the strict transition validator) and so
emits the audit explicitly, applies the branchless guard legacy
unblock() relies on, restores ownership from the snapshot, and clears
the pre-block snapshot fields.
"""
# A task with no branch cannot resume in_progress — the dispatcher
# refuses a branchless in_progress task and loops — so divert it to
# pending, exactly as legacy unblock() does.
if restored_status == TaskStatus.IN_PROGRESS and not task.branch_name:
restored_status = TaskStatus.PENDING
pre_status = (
task.status.value
if isinstance(task.status, TaskStatus)
roboco/services/usage.py
+87 -194
View File
@@ -15,11 +15,26 @@ from sqlalchemy import func, select
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import InstrumentedAttribute
from roboco.db.tables import AgentSpawnSessionTable, DailyUsageRollupTable
from roboco.services.base import BaseService
def _row_tokens(row: Any) -> tuple[int, int, int, int]:
"""Extract (input, output, cache_read, cache_write) token counts as ints.
Centralizes the null-coalescing that would otherwise be repeated across
every aggregation method.
"""
return (
int(row.tokens_input or 0),
int(row.tokens_output or 0),
int(row.tokens_cache_read or 0),
int(row.tokens_cache_write or 0),
)
def _parse_period(period: str) -> tuple[datetime, int]:
"""Parse period string into (start_dt, hours).
@@ -199,211 +214,89 @@ class UsageService(BaseService):
return points
# =========================================================================
# BY-AGENT
# BY-DIMENSION (agent / team / model)
# =========================================================================
async def _aggregate_by(
self,
group_column: InstrumentedAttribute[Any],
key_name: str,
period: str,
) -> list[dict[str, Any]]:
"""Aggregate token usage grouped by an arbitrary column.
Shared implementation behind get_by_agent/get_by_team/get_by_model.
``key_name`` is the dict key the grouping value is emitted under
(e.g. "agent_slug"). Rows are ordered by input+output desc and each
item carries pct_of_total computed against the grand total.
"""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
group_column.label(key_name),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(group_column)
.order_by(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
).desc()
)
)
rows = result.fetchall()
grand_total = sum(sum(_row_tokens(r)) for r in rows)
items = []
for r in rows:
ti, to_, tcr, tcw = _row_tokens(r)
total = ti + to_ + tcr + tcw
items.append(
{
key_name: getattr(r, key_name),
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": total,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
"pct_of_total": round(total / grand_total * 100, 2)
if grand_total > 0
else 0.0,
}
)
return items
async def get_by_agent(self, period: str = "24h") -> list[dict[str, Any]]:
"""Return per-agent token usage with pct_of_total."""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
AgentSpawnSessionTable.agent_slug,
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(AgentSpawnSessionTable.agent_slug)
.order_by(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
).desc()
)
return await self._aggregate_by(
AgentSpawnSessionTable.agent_slug, "agent_slug", period
)
rows = result.fetchall()
grand_total = sum(
int(r.tokens_input or 0)
+ int(r.tokens_output or 0)
+ int(r.tokens_cache_read or 0)
+ int(r.tokens_cache_write or 0)
for r in rows
)
items = []
for r in rows:
ti = int(r.tokens_input or 0)
to_ = int(r.tokens_output or 0)
tcr = int(r.tokens_cache_read or 0)
tcw = int(r.tokens_cache_write or 0)
total = ti + to_ + tcr + tcw
items.append(
{
"agent_slug": r.agent_slug,
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": total,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
"pct_of_total": round(total / grand_total * 100, 2)
if grand_total > 0
else 0.0,
}
)
return items
# =========================================================================
# BY-TEAM
# =========================================================================
async def get_by_team(self, period: str = "24h") -> list[dict[str, Any]]:
"""Return per-team token usage with pct_of_total."""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
AgentSpawnSessionTable.team,
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(AgentSpawnSessionTable.team)
.order_by(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
).desc()
)
)
rows = result.fetchall()
grand_total = sum(
int(r.tokens_input or 0)
+ int(r.tokens_output or 0)
+ int(r.tokens_cache_read or 0)
+ int(r.tokens_cache_write or 0)
for r in rows
)
items = []
for r in rows:
ti = int(r.tokens_input or 0)
to_ = int(r.tokens_output or 0)
tcr = int(r.tokens_cache_read or 0)
tcw = int(r.tokens_cache_write or 0)
total = ti + to_ + tcr + tcw
items.append(
{
"team": r.team,
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": total,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
"pct_of_total": round(total / grand_total * 100, 2)
if grand_total > 0
else 0.0,
}
)
return items
# =========================================================================
# BY-MODEL
# =========================================================================
return await self._aggregate_by(AgentSpawnSessionTable.team, "team", period)
async def get_by_model(self, period: str = "24h") -> list[dict[str, Any]]:
"""Return per-model token usage with pct_of_total."""
start_dt, _ = _parse_period(period)
result = await self.session.execute(
select(
AgentSpawnSessionTable.model,
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_input), 0).label(
"tokens_input"
),
func.coalesce(func.sum(AgentSpawnSessionTable.tokens_output), 0).label(
"tokens_output"
),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_read), 0
).label("tokens_cache_read"),
func.coalesce(
func.sum(AgentSpawnSessionTable.tokens_cache_write), 0
).label("tokens_cache_write"),
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost_usd"),
)
.where(
AgentSpawnSessionTable.started_at >= start_dt,
AgentSpawnSessionTable.ended_at.isnot(None),
)
.group_by(AgentSpawnSessionTable.model)
.order_by(
func.sum(
AgentSpawnSessionTable.tokens_input
+ AgentSpawnSessionTable.tokens_output
).desc()
)
)
rows = result.fetchall()
grand_total = sum(
int(r.tokens_input or 0)
+ int(r.tokens_output or 0)
+ int(r.tokens_cache_read or 0)
+ int(r.tokens_cache_write or 0)
for r in rows
)
items = []
for r in rows:
ti = int(r.tokens_input or 0)
to_ = int(r.tokens_output or 0)
tcr = int(r.tokens_cache_read or 0)
tcw = int(r.tokens_cache_write or 0)
total = ti + to_ + tcr + tcw
items.append(
{
"model": r.model,
"tokens_input": ti,
"tokens_output": to_,
"total_tokens": total,
"cost_usd": round(float(r.cost_usd or 0.0), 6),
"pct_of_total": round(total / grand_total * 100, 2)
if grand_total > 0
else 0.0,
}
)
return items
return await self._aggregate_by(AgentSpawnSessionTable.model, "model", period)
# =========================================================================
# PROJECTION