fix(tg): cockpit data correctness — real GLM pricing, display timezone, agent activity tracking (#666)

* fix(tg): cockpit data correctness — real GLM pricing, display timezone, agent activity tracking

Three root causes behind the Mini App/bot showing wrong numbers:

Pricing: glm-5.2 gets a grounded per-token rate (z.ai published pricing,
$1.40/$4.40/$0.26 per 1M, source+date in the table comment) so a GLM
fleet day stops reporting $0.00 for half a million tokens; ungrounded
Ollama-Cloud models render "subscription (untracked)" instead of a bare
zero (is_ollama_cloud_model, consumed directly by the cockpit). Side
effect, intended and documented: honestly-priced GLM now trips the
downgrade-only comparator for new qa/documenter complexity pins.

Display timezone: the cockpit bucketed days in UTC for a GMT+2 operator.
New pure foundation module display_time (resolve_zone/local_date/
trailing_dates/day_bounds_utc, DST-correct with tests for the 23h/25h
days) + ROBOCO_DISPLAY_TIMEZONE (IANA-validated, default UTC); the
cockpit's spend/velocity series bucket raw session/completion rows by
the display zone. The UTC-keyed rollup table and the main dashboard are
deliberately untouched.

Agent activity: AgentTable.status was never set to ACTIVE and
current_task_id was never written anywhere — "active: 0, working: []"
was structurally permanent. Every claim path now marks the claimant
ACTIVE with rollback symmetry (_finalize_claim for dev/PM claims,
_qa_or_doc_claim for QA/doc/PR-gate claims, pr_review_claim for external
review) and every release path clears it (pass/fail QA, pr_pass/pr_fail,
complete_review, advance-to-PM-review, reaper unclaim, voluntary
unclaim, reassign retarget, pool divert, admin transitions, unblock
restore-to-in-progress). The bot's /status shares the cockpit's fleet
derivation so the two surfaces can't disagree. Known ceiling, commented:
one current_task_id column shows a multi-root coordinator PM's most
recent claim only.

Drill: sonnet develop -> sonnet adversarial (refuted the original
chokepoint coverage claim; QA/doc/reviewer paths were unwired) ->
correction round (wired them all + restored a dropped assertion, deleted
a dead helper and the dead subscription_billed field) -> review.

* fix(db): post_update on AgentTable.current_task breaks the flush cycle

agents.current_task_id and tasks.assigned_to reference each other, so a
flush touching both rows — every claim now marks its agent ACTIVE — is
an instance-level circular dependency SQLAlchemy cannot topologically
sort. The e2e smoke's full verb paths (12 tests) hit it; the unit and
integration suites never flush both dirty rows with relationships
loaded. post_update emits the FK as a second UPDATE, the canonical fix
for mutually-referencing rows.

* fix(budgets): enforce only explicitly-set budgets — no per-TaskType defaults

The per-TaskType default cap table blocked an unbudgeted coordination
root one opus planning turn in ($1.50 PLANNING default vs. real
coordination spend) — a false positive by design the moment the fleet
runs a priced model. Budgets are now explicit-input only:
effective_task_budget_usd returns None for an unset budget_usd, the
budget sweep skips enforcement (and never prices spend) on None, and
the unblock re-check passes on None so clearing the budget field is
itself a valid resolution. The project monthly cap stays as the
explicit-input fleet-wide backstop. Panel copy tells the truth
("No cap" placeholder; empty = uncapped), and the TaskType default
table plus its resolver are deleted.

---------

Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
2026-07-23 21:09:50 +02:00
committed by GitHub
co-authored by Renn F
parent 0c1d450a05
commit a036c97985
35 changed files with 1364 additions and 191 deletions
+1 -1
View File
@@ -440,7 +440,7 @@ Agent backends are pluggable. `roboco/llm/providers/` defines an `AgentProvider`
**Delegation detail-fidelity (always-on, 2026-07-16).** Details no longer thin out at hand-off in either direction. DOWN: `delegate` refuses any child that doesn't declare `covers_parent_criteria` mapping onto the parent's real acceptance criteria (matched by id or exact text; an unresolvable ref is rejected naming the valid criteria, never silently dropped — previously the mapping was optional and coverage surfaced only at `submit_up`'s roll-up gate, after the whole wave had already run); the success envelope carries `parent_ac_coverage` `{covered, uncovered}` so a wave-planning PM sees remaining gaps in the same turn, while multi-wave planning stays legal. UP: `pass_review` requires `criteria_verified` — one `{criterion, evidence}` entry per task acceptance criterion (the findings ledger's id-or-exact-text matcher, soup-checked and length-capped evidence), rejecting with the unverified criteria named; entries render deterministically into `qa_notes` as `[AC] <criterion> — verified: <evidence>` lines, so a gestalt "looks good" pass is structurally impossible. Video briefs stopped being prose-only: an enumerable feature list (release `highlights`, or `input_props.highlights` carried onto a reject re-author) becomes its own acceptance criterion ("Every brief-named feature appears as its own fully readable scene: …", bounded to the AC caps; a re-author without highlights carries "every point in the CEO rejection feedback is visibly addressed"), so the dropped-scene class — a four-feature brief shipping three scenes past every gate — is caught by the QA per-AC stamp instead of the CEO's eyeball.
**Task/project cost budgets (default-off `ROBOCO_TASK_BUDGETS_ENABLED`).** `tasks.budget_usd` + `projects.monthly_budget_usd` (migration 080). Claim-time: a project-month-spend guard (`project_budget_exceeded_guard`, `roboco/services/gateway/claim_guards.py`) applies only to WORK-STARTING claims (`i_will_work_on` / `i_will_plan`) — review/doc/gate/inbound-PR claims are exempt so in-flight work can always finish reviewing and merging even at cap; spend counts closed sessions' `estimated_cost_usd` plus open sessions priced live from token snapshots. Sweep-side: the orchestrator's existing budget sweep also prices the active task's own spend against `budget_usd` (falling back to a `TaskType` default, `effective_task_budget_usd` in `roboco/foundation/policy/agent_loop.py`, when null); on breach the task is BLOCKED (HUMAN resolver, budget marker) BEFORE the graceful stop so the unclaim no-ops and the dispatcher never respawns onto it, and the CEO notification names both recovery steps — `unblock` on a budget-blocked task re-checks live spend and refuses while still over, so there's no silent re-breach loop. Off => neither cap is ever consulted regardless of field values. Panel: budget inputs on both the project and task dialogs (a `0` is rejected — it would silently block everything); spend math is consolidated in `TaskService.task_spend_usd`.
**Task/project cost budgets (default-off `ROBOCO_TASK_BUDGETS_ENABLED`).** `tasks.budget_usd` + `projects.monthly_budget_usd` (migration 080). Claim-time: a project-month-spend guard (`project_budget_exceeded_guard`, `roboco/services/gateway/claim_guards.py`) applies only to WORK-STARTING claims (`i_will_work_on` / `i_will_plan`) — review/doc/gate/inbound-PR claims are exempt so in-flight work can always finish reviewing and merging even at cap; spend counts closed sessions' `estimated_cost_usd` plus open sessions priced live from token snapshots. Sweep-side: the orchestrator's existing budget sweep also prices the active task's own spend against `budget_usd` (explicit-input only — a null budget means no cap, resolved by `effective_task_budget_usd` in `roboco/foundation/policy/agent_loop.py`; the earlier per-TaskType default table blocked an unbudgeted coordination root one opus turn in and was removed); on breach the task is BLOCKED (HUMAN resolver, budget marker) BEFORE the graceful stop so the unclaim no-ops and the dispatcher never respawns onto it, and the CEO notification names both recovery steps — `unblock` on a budget-blocked task re-checks live spend and refuses while still over, so there's no silent re-breach loop. Off => neither cap is ever consulted regardless of field values. Panel: budget inputs on both the project and task dialogs (a `0` is rejected — it would silently block everything); spend math is consolidated in `TaskService.task_spend_usd`.
**PR labeler (always-on).** `derive_pr_labels` (`roboco/foundation/policy/pr_labels.py`, pure) derives the org-structure label vocabulary every fleet PR now carries: `to {base_branch}` — the PR's REAL resolved target branch, verbatim (never assumed from `is_root_pr`, so a project with a renamed/non-standard trunk or an env-ladder rung gets an accurate label instead of a hardcoded `master`/`slave`), `root` for an assembled root PR, `MegaTask` for a batch-carrying task, and a layer label (`main-pm` for a Main-PM coordination root, `cell/{team}` for a cell-assembled PR, else `subtask/{team}` for a leaf dev PR). Applied best-effort at all three PR-opening sites in `GitService` so a human triaging the PR queue sees which tree and which org layer a PR belongs to at a glance.
@@ -56,7 +56,7 @@ const FLAG_DESCRIPTIONS: Record<string, string> = {
possibilities_matrix_enabled:
"When a task's work is already done (commits + open PR + all acceptance criteria addressed + no open findings), submit it for QA in one i_am_done call instead of 3-6 turns — skips the retroactive plan, journal tracing, and local quality (CI-green proxy) gates. Off by default: the standard path is unchanged until you arm this.",
task_budgets_enabled:
"Enforce per-project monthly and per-task cost caps (USD). A claim is refused once a project's monthly budget is reached; an active task whose own budget (or its task-type default) is breached is stopped and blocked, and you're notified. Set the caps on the project edit dialog and a task's detail page — a project/task with no cap set is unaffected either way.",
"Enforce per-project monthly and per-task cost caps (USD). A claim is refused once a project's monthly budget is reached; an active task whose own explicitly-set budget is breached is stopped and blocked, and you're notified. Set the caps on the project edit dialog and a task's detail page — a project/task with no cap set is never capped.",
rag_auto_update_enabled:
"Keep the knowledge base index refreshed automatically.",
transcript_prune_enabled:
@@ -81,7 +81,7 @@ function submit() {
}
function budgetInput(): HTMLInputElement {
return screen.getByPlaceholderText("Task-type default") as HTMLInputElement;
return screen.getByPlaceholderText("No cap") as HTMLInputElement;
}
describe("EditTaskDialog — Budget (USD) input", () => {
@@ -153,7 +153,7 @@ describe("EditTaskDialog — Budget (USD) input", () => {
expect(mutateAsync).not.toHaveBeenCalled();
});
it("submits null when left empty (use the task-type default)", async () => {
it("submits null when left empty (no cap)", async () => {
render(
<EditTaskDialog
task={{ ...task, budget_usd: 3.5 }}
@@ -156,7 +156,7 @@ function EditTaskDialogInner({
const parsedBudget = trimmedBudget ? Number(trimmedBudget) : null;
if (trimmedBudget && (Number.isNaN(parsedBudget) || parsedBudget! <= 0)) {
toast.error(
"Budget must be greater than 0 — leave it empty for the task-type default",
"Budget must be greater than 0 — leave it empty for no cap",
);
return;
}
@@ -359,21 +359,21 @@ function EditTaskDialogInner({
{/* Budget (USD) */}
<div className="space-y-2">
<HelpTip label="Caps this task's own agent-spawn spend; only enforced when the task-budgets feature flag is on. Empty = use the task-type default.">
<HelpTip label="Caps this task's own agent-spawn spend; only enforced when the task-budgets feature flag is on. Empty = no cap.">
<Label>Budget (USD)</Label>
</HelpTip>
<Input
type="number"
min="0.01"
step="0.01"
placeholder="Task-type default"
placeholder="No cap"
value={budgetUsd}
onChange={(e) => setBudgetUsd(e.target.value)}
/>
<p className="text-xs text-muted-foreground">
Must be greater than 0 a 0 budget would block the task
before it spends a cent. Leave blank for the task-type
default.
before it spends a cent. Leave blank for no cap: budgets
enforce only when explicitly set.
</p>
{spendUsd != null && (
<p className="text-xs text-muted-foreground" data-testid="task-spend">
@@ -228,6 +228,25 @@ describe("TgTodayTab", () => {
expect(screen.getByText(/idle · 21/)).toBeInTheDocument();
});
it("labels an untracked subscription-billed spend day instead of a bare $0", async () => {
get.mockResolvedValue({
data: brief({
spend: {
tokens_today: 456_221,
cost_today_usd: 0,
subscription_billed: true,
series: [1, 2, 3, 4, 5, 6, 0],
delta_pct: null,
},
}),
});
renderTab();
expect(await screen.findByText("≈$0")).toBeInTheDocument();
expect(screen.getByText(/subscription \(untracked\)/i)).toBeInTheDocument();
expect(screen.queryByText("$0.00")).not.toBeInTheDocument();
});
it("shows an error state when the brief fails to load", async () => {
get.mockRejectedValue(new Error("boom"));
renderTab();
+13 -2
View File
@@ -69,6 +69,11 @@ export interface TodayBrief {
spend: {
tokens_today: number;
cost_today_usd: number;
/** True when cost_today_usd is $0 because today's tokens ran on an
* Ollama Cloud model with no grounded per-token rate (subscription-
* billed, not literally free) the hero renders "untracked", never a
* bare misleading $0.00. */
subscription_billed: boolean;
series: number[];
delta_pct: number | null;
};
@@ -134,10 +139,16 @@ function SpendHero({
/>
</div>
<span className="tg-display block text-[44px] leading-none tabular-nums">
${cost.toFixed(2)}
{spend.subscription_billed ? "≈$0" : `$${cost.toFixed(2)}`}
</span>
<div className="mt-1.5 flex items-center gap-2">
<TgDeltaChip pct={spend.delta_pct} />
{spend.subscription_billed ? (
<span className="text-xs font-medium text-amber-500">
subscription (untracked)
</span>
) : (
<TgDeltaChip pct={spend.delta_pct} />
)}
<span className="text-xs tabular-nums text-muted-foreground">
{fmtTokens(spend.tokens_today)} tokens
</span>
+1
View File
@@ -180,6 +180,7 @@ export const DEMO_TODAY: TodayBrief = {
spend: {
tokens_today: 2_400_000,
cost_today_usd: 18.72,
subscription_billed: false,
series: [12.4, 9.1, 15.8, 11.2, 21.6, 14.9, 18.72],
delta_pct: 25.6,
},
+2 -2
View File
@@ -211,7 +211,7 @@ class TaskUpdate(BaseModel):
priority: int | None = Field(default=None, ge=0, le=3)
sequence: int | None = Field(default=None, ge=0) # Order within siblings
# Cost cap (ROBOCO_TASK_BUDGETS_ENABLED). An explicit null clears it back
# to "use the TaskType default" — handled at the route layer like the
# to uncapped (budgets enforce only when set) — handled at the route layer like the
# other _NULLABLE_TASK_FIELDS (TaskService.update() itself skips None).
# gt=0 — a 0/negative cap would block every claim immediately (#654).
budget_usd: float | None = Field(default=None, gt=0)
@@ -311,7 +311,7 @@ class TaskResponse(BaseModel):
status: TaskStatus
priority: int
sequence: int # Order number within siblings
# Cost cap (ROBOCO_TASK_BUDGETS_ENABLED). Null = use the TaskType default.
# Cost cap (ROBOCO_TASK_BUDGETS_ENABLED). Null = no cap (explicit-input only).
budget_usd: float | None = None
# This task's own accumulated agent-spawn spend (TaskService.task_spend_usd).
# Only populated by GET /tasks/{id} (an extra DB read) when
+65 -21
View File
@@ -4,7 +4,7 @@ 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.
Pricing is provider-aware. A model name resolves to one of four cases:
Pricing is provider-aware. A model name resolves to one of five cases:
* **Anthropic** priced from the table below by substring match.
* **Priced non-Anthropic** xAI Grok (``grok-build-*``, billed per token via
@@ -13,10 +13,20 @@ Pricing is provider-aware. A model name resolves to one of four cases:
metered), and Google Gemini (``gemini-2.5-*``, billed per token via the
Gemini API) are priced from the table too. Match by substring like the rest.
* **Free 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.
or bare model tags with no ``:cloud`` tag). These have **no per-token
cost**: local inference runs on owned hardware. Returns an intentional
``0.0`` not an error condition, so it is not warned on.
* **Subscription-billed Ollama Cloud** an Ollama Cloud model (``:cloud``
tag) is, like Grok/Codex/Gemini, really billed via flat subscription /
GPU-time rather than per token but that does not make it free the way
local inference is, so it is priced from the table at its API-equivalent
rate the same way grok-build is (see the GLM-5.2 entry below) whenever a
citable published rate exists. An Ollama Cloud model with NO table entry
(a rate we could not ground) still returns ``0.0``; ``is_ollama_cloud_model``
identifies this case for a caller that wants to render "subscription
(untracked)" instead of implying the run was free (see the TG cockpit's
own spend label, which calls it directly rather than through
:class:`CostResult`).
* **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``.
@@ -82,6 +92,22 @@ _PRICING: list[tuple[str, float, float, float, float]] = [
("gemini-2.5-pro", 1.25, 10.00, 1.25, 1.25),
("gemini-2.5-flash-lite", 0.10, 0.40, 0.10, 0.10),
("gemini-2.5-flash", 0.30, 2.50, 0.30, 0.30),
# Z.ai GLM-5.2 — priced non-Anthropic. Ollama Cloud's `glm-5.2:cloud` tag
# (roboco's own ROBOCO_LOCAL_LLM_MODEL default) is billed via flat
# subscription/GPU-time, not per token, but the same "attribute at the
# underlying API-equivalent rate" convention as grok-build/gpt-5.3-codex
# applies once a real published rate exists — a real fleet running on
# this model was undercounting spend as literal $0 otherwise. Official
# first-party pricing: https://docs.z.ai/guides/overview/pricing (fetched
# 2026-07-23): $1.40/1M input, $4.40/1M output, $0.26/1M cached input. No
# cache-write discount is published, so cache_write falls back to the
# input rate (same convention as grok-build/gpt-5.3-codex above).
# Side effect: input_price_per_million("glm-5.2:cloud") is now $1.40 (was
# $0.0) — pricier than haiku's $1.00 — so the cost-tiered
# complexity-override's downgrade-only comparator now REJECTS a new
# qa/documenter pin to glm-5.2:cloud that was previously allowed when it
# priced as free-tier. Intended: it really is costlier per input token.
("glm-5.2", 1.40, 4.40, 0.26, 1.40),
# Short aliases used in ROLE_MODEL_MAP / MODEL_MAP
("opus", 5.00, 25.00, 0.50, 6.25),
("sonnet", 3.00, 15.00, 0.30, 0.75),
@@ -109,6 +135,16 @@ def _is_anthropic_model(lower: str) -> bool:
return any(fragment in lower for fragment in _ANTHROPIC_FRAGMENTS)
def is_ollama_cloud_model(model: str) -> bool:
"""True for an Ollama Cloud model (the ``:cloud`` tag convention, e.g.
``glm-5.2:cloud``) distinguishes a subscription-billed cloud model from
a genuinely-free self-hosted Ollama model (``ollama/`` prefix or a bare
local tag with no ``:cloud`` suffix). Consumed by the TG cockpit
(``tg_cockpit.py``) to label an ungrounded cloud model's ``$0`` spend as
"subscription (untracked)" rather than implying the run was free."""
return ":cloud" in model.lower()
def _lookup_prices(lower: str) -> tuple[float, float, float, float] | None:
"""Return the (input, output, cache_read, cache_write) rates for a model.
@@ -135,8 +171,8 @@ def input_price_per_million(model: str) -> float:
policy) to rank two models against each other without needing a separate
explicit tier ordering: the input rate already orders the Claude tiers
(haiku < sonnet < opus) and prices Grok below Sonnet, so "costlier" reduces
to "higher input price". A model with no pricing-table match (self-hosted,
Ollama Cloud genuinely free per-token) returns ``0.0``, the cheapest
to "higher input price". A model with no pricing-table match (self-hosted
Ollama, or an ungrounded Ollama Cloud tag) returns ``0.0``, the cheapest
possible rank, so it can never be rejected as "costlier".
"""
if not model:
@@ -149,12 +185,14 @@ def input_price_per_million(model: str) -> float:
class CostResult:
"""Estimated cost plus pricing attribution (#65).
``cost_usd`` is ``0.0`` for both a genuinely-free non-Anthropic model (local
inference no per-token cost) and an unpriced Anthropic model (a Claude
model we forgot to price real spend we are failing to count). ``unpriced``
distinguishes them so a caller can surface the miss instead of silently
reporting ``$0``. ``is_anthropic`` records which family the model resolved
to. ``calculate_cost`` returns just the ``cost_usd`` float for existing
``cost_usd`` is ``0.0`` for both a genuinely-free non-Anthropic model
(local inference, or an ungrounded Ollama Cloud tag see
``is_ollama_cloud_model`` for distinguishing the two) and an unpriced
Anthropic model (a Claude model we forgot to price real spend we are
failing to count). ``unpriced`` flags only the latter case so a caller
can surface the miss instead of silently reporting ``$0``.
``is_anthropic`` records which family the model resolved to.
``calculate_cost`` returns just the ``cost_usd`` float for existing
callers; ``calculate_cost_result`` returns the full attribution.
"""
@@ -195,11 +233,13 @@ def calculate_cost_result(
"""Calculate the estimated USD cost for a model invocation, with attribution.
Matches the model name against the known pricing table using substring
search (longest match wins). Provider-aware (see module docstring):
non-Anthropic models (local Ollama, Ollama Cloud) have no per-token cost
and return ``cost_usd=0.0, unpriced=False``; an unpriced Anthropic model
returns ``cost_usd=0.0, unpriced=True`` and logs a warning since it
represents real spend we are failing to count.
search (longest match wins). Provider-aware (see module docstring): a
priced non-Anthropic model (Grok, Codex, Gemini, or a grounded Ollama
Cloud model like GLM-5.2) gets a real per-token cost; an ungrounded
non-Anthropic model (local Ollama, or an unpriced Ollama Cloud tag) has
no per-token cost and returns ``cost_usd=0.0, unpriced=False``; an
unpriced Anthropic model returns ``cost_usd=0.0, unpriced=True`` and logs
a warning since it represents real spend we are failing to count.
Args:
model: Model name or short alias (e.g. ``"claude-sonnet-5"``,
@@ -227,14 +267,18 @@ def calculate_cost_result(
best_prices = _lookup_prices(lower)
if best_prices is None:
# 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.
# undercounting); a non-Anthropic model with no rate is either
# genuinely free (self-hosted) or subscription-billed-but-untracked
# (Ollama Cloud) — both log at debug (see ``is_ollama_cloud_model``
# for a caller that needs to distinguish the two).
if is_anthropic:
logger.warning("No pricing data found for Anthropic model", model=model)
else:
logger.debug("Non-Anthropic model has no per-token cost", model=model)
# Unpriced only when it is an Anthropic model we forgot to price; a
# non-Anthropic model with no rate is intentionally free.
# non-Anthropic model with no rate is either genuinely free or
# subscription-billed-but-untracked (never "unpriced" — that label is
# reserved for real spend we are failing to count).
return CostResult(
cost_usd=0.0, unpriced=is_anthropic, is_anthropic=is_anthropic
)
+25 -2
View File
@@ -14,6 +14,7 @@ from functools import lru_cache
from pathlib import Path
from typing import Literal
from urllib.parse import urlparse
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from pydantic import Field, computed_field, field_validator, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -42,6 +43,28 @@ class Settings(BaseSettings):
environment: str = Field(
default="development", pattern="^(development|staging|production)$"
)
display_timezone: str = Field(
default="UTC",
description=(
"IANA timezone name (e.g. 'Europe/Berlin') used ONLY for "
"display-side 'today'/day-bucket derivations — currently the "
"Telegram cockpit's Today brief and bot commands. Default 'UTC' "
"is a no-op for every deployment that doesn't set this. DB "
"storage stays UTC canonical regardless; this never touches how "
"timestamps are written or how daily_usage_rollups are keyed."
),
)
@field_validator("display_timezone")
@classmethod
def _validate_display_timezone(cls, v: str) -> str:
try:
ZoneInfo(v)
except (ZoneInfoNotFoundError, ValueError) as exc:
raise ValueError(
f"display_timezone {v!r} is not a valid IANA timezone name"
) from exc
return v
# ==========================================================================
# API Server
@@ -391,8 +414,8 @@ class Settings(BaseSettings):
"Per-task and per-project cost budgets. When on: a claim is "
"refused once a project's monthly_budget_usd is reached (summed "
"agent-spawn spend across its tasks this calendar month), and the "
"budget sweep blocks an active task whose own budget_usd (or the "
"TaskType default) is breached, notifying the CEO. Off => neither "
"budget sweep blocks an active task whose own explicitly-set "
"budget_usd is breached, notifying the CEO. Off => neither "
"cap is ever consulted, regardless of project/task field values."
),
)
+7 -3
View File
@@ -138,9 +138,13 @@ class AgentTable(Base):
DateTime(timezone=True), onupdate=lambda: datetime.now(UTC), nullable=True
)
# Relationships - use lazy="joined" for single optional relationship
# Relationships - use lazy="joined" for single optional relationship.
# post_update: agents.current_task_id and tasks.assigned_to reference each
# other, so a flush touching both rows (a claim marking its agent ACTIVE)
# is an instance-level cycle SQLAlchemy cannot topologically sort —
# post_update breaks it by emitting this FK as a second UPDATE.
current_task: Mapped["TaskTable | None"] = relationship(
"TaskTable", foreign_keys=[current_task_id], lazy="joined"
"TaskTable", foreign_keys=[current_task_id], lazy="joined", post_update=True
)
@@ -210,7 +214,7 @@ class TaskTable(Base):
)
priority: Mapped[int] = mapped_column(Integer, nullable=False, default=2)
# Cost budget (migration 079, feature-flagged: ROBOCO_TASK_BUDGETS_ENABLED).
# Null falls back to the TaskType default when the flag is on; a pure
# Null = no cap: budgets enforce only when explicitly set; a pure
# no-op off. Enforced periodically by the orchestrator's budget sweep.
budget_usd: Mapped[float | None] = mapped_column(Float, nullable=True)
+10 -37
View File
@@ -26,8 +26,6 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Literal
from roboco.models.base import TaskType
@dataclass(frozen=True)
class BudgetPolicy:
@@ -79,44 +77,19 @@ class BudgetPolicy:
DEFAULT_BUDGET: BudgetPolicy = BudgetPolicy()
# Per-TaskType default $ budget (USD), consulted ONLY when a task's own
# `budget_usd` is null AND ROBOCO_TASK_BUDGETS_ENABLED is on (see
# roboco/config.py). Relative sizing reflects typical turn/tool-call weight:
# CODE is the most token-heavy (multi-file edits, gate runs, revisions);
# RESEARCH/DESIGN sit mid (web research + note/asset writing); PLANNING is
# lighter prose; DOCUMENTATION and ADMINISTRATIVE are the cheapest, mostly
# read-and-write-notes work.
TASK_TYPE_DEFAULT_BUDGET_USD: dict[TaskType, float] = {
TaskType.CODE: 5.0,
TaskType.RESEARCH: 2.0,
TaskType.DESIGN: 2.0,
TaskType.PLANNING: 1.5,
TaskType.DOCUMENTATION: 1.0,
TaskType.ADMINISTRATIVE: 0.5,
}
def effective_task_budget_usd(task: Any) -> float | None:
"""A task's $ cap: its own ``budget_usd``, or ``None`` for no cap at all.
def default_budget_usd_for(task_type: TaskType) -> float:
"""The TASK_TYPE_DEFAULT_BUDGET_USD entry for ``task_type``.
Falls back to the CODE tier (the most generous) for any TaskType this
dict has not been kept in sync with, so a future TaskType addition fails
open (a spend cap that's too generous) rather than crashing the sweep.
Budgets are enforcement-by-explicit-input only a task nobody set a
budget on is uncapped (the earlier per-TaskType default table blocked a
default-budget coordination root one opus planning turn in; 2026-07-23).
The project's ``monthly_budget_usd`` claim guard is the explicit-input
backstop for fleet-wide spend. The one place this resolution happens
the orchestrator's budget sweep and the ``unblock`` re-check both call
this and skip enforcement on ``None``.
"""
return TASK_TYPE_DEFAULT_BUDGET_USD.get(
task_type, TASK_TYPE_DEFAULT_BUDGET_USD[TaskType.CODE]
)
def effective_task_budget_usd(task: Any) -> float:
"""A task's effective $ cap: its own ``budget_usd``, or the TaskType
default when null. The one place this resolution happens the
orchestrator's budget sweep and the ``unblock`` re-check both call this
instead of re-deriving the null-fallback themselves."""
budget_usd = getattr(task, "budget_usd", None)
if budget_usd is not None:
return float(budget_usd)
return default_budget_usd_for(task.task_type)
return float(budget_usd) if budget_usd is not None else None
# Per-verb retry caps. Verbs that hit a tracing_gap or invalid_state and
+53
View File
@@ -0,0 +1,53 @@
"""Display-timezone day bucketing — pure, DB-free.
DB storage stays UTC canonical everywhere; these functions only decide which
calendar day (in the operator's configured ``display_timezone``) a UTC
instant belongs to, for read-side "today" / trailing-window aggregation
(currently the Telegram cockpit's Today brief + bot commands). No writes, no
ORM, no settings import callers pass the configured tz name in.
"""
from __future__ import annotations
from datetime import UTC, date, datetime, timedelta
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
def resolve_zone(tz_name: str) -> ZoneInfo:
"""``ZoneInfo`` for ``tz_name``, falling back to UTC on an unknown name.
Settings validation already rejects a bad name at load time; this is a
defensive fallback for a value that reached here some other way (e.g. a
stale env read before validation ran) rather than a 500.
"""
try:
return ZoneInfo(tz_name)
except (ZoneInfoNotFoundError, ValueError):
return ZoneInfo("UTC")
def local_date(instant: datetime, tz_name: str) -> date:
"""The calendar date ``instant`` (any tz-aware datetime) falls on in
``tz_name``."""
return instant.astimezone(resolve_zone(tz_name)).date()
def trailing_dates(
tz_name: str, days: int, *, now: datetime | None = None
) -> list[date]:
"""The last ``days`` calendar dates in ``tz_name``, oldest -> today
(inclusive of today)."""
today = local_date(now or datetime.now(UTC), tz_name)
return [today - timedelta(days=n) for n in reversed(range(days))]
def day_bounds_utc(tz_name: str, day: date) -> tuple[datetime, datetime]:
"""UTC ``[start, end)`` instants spanning ``day``'s midnight-to-midnight
window in ``tz_name`` correct across a DST transition day (a spring-
forward day is a real 23h UTC span, fall-back a real 25h one) because
``ZoneInfo`` re-resolves the offset from the wall-clock fields at
``.astimezone()`` time rather than freezing it at construction."""
zone = resolve_zone(tz_name)
start_local = datetime(day.year, day.month, day.day, tzinfo=zone)
end_local = start_local + timedelta(days=1)
return start_local.astimezone(UTC), end_local.astimezone(UTC)
+4 -4
View File
@@ -166,15 +166,15 @@ class Task(TimestampMixin):
default=2, ge=0, le=3, description="0=P0(highest), 3=P3(lowest)"
)
# Cost budget (feature-flagged: ROBOCO_TASK_BUDGETS_ENABLED). Null = fall
# back to the TaskType default (see foundation/policy/agent_loop.py
# TASK_TYPE_DEFAULT_BUDGET_USD) when the flag is on; a pure no-op off.
# Cost budget (feature-flagged: ROBOCO_TASK_BUDGETS_ENABLED). Null = no
# cap — budgets enforce only when explicitly set (see
# foundation/policy/agent_loop.py effective_task_budget_usd); a no-op off.
budget_usd: float | None = Field(
default=None,
gt=0,
description=(
"Cap on this task's own accumulated agent-spawn spend "
"(estimated_cost_usd). Null = use the TaskType default. Must be "
"(estimated_cost_usd). Null = no cap (explicit-input only). Must be "
"> 0 — a 0/negative cap would block every claim immediately."
),
)
+7 -4
View File
@@ -8053,11 +8053,12 @@ Start by:
async def _task_budget_breach(self, task_id_str: str) -> tuple[float, float] | None:
"""``(cap_usd, spend_usd)`` if this task's own $ budget is breached.
``None`` when unbreached OR the task has already left claimed/
``None`` when unbreached, when the task carries no explicit
``budget_usd`` (budgets are explicit-input only an unset budget
means no cap; ``effective_task_budget_usd`` is the shared resolver
``unblock``'s re-check uses), OR the task has already left claimed/
in_progress (a stale re-check racing the task's own progress — not a
breach). The cap is ``task.budget_usd``, falling back to the
TaskType default (``effective_task_budget_usd`` the same resolver
``unblock``'s re-check uses) when null. Spend is
breach). Spend is
``TaskService.task_spend_usd`` (closed-session cost + open-session
live-token pricing via ``calculate_cost`` a DB-only read off
``_sweep_token_snapshots``'s periodically-refreshed token columns, no
@@ -8082,6 +8083,8 @@ Start by:
):
return None
cap_usd = effective_task_budget_usd(task)
if cap_usd is None:
return None
spend_usd = await svc.task_spend_usd(task_id)
if spend_usd < cap_usd:
return None
@@ -1278,7 +1278,9 @@ class Choreographer:
return None
spend_usd = await self.task.task_spend_usd(t.id)
cap_usd = effective_task_budget_usd(t)
if spend_usd >= cap_usd:
# No explicit budget = no cap: the breach that stamped the marker was
# since resolved by clearing the budget field — unblock proceeds.
if cap_usd is not None and spend_usd >= cap_usd:
return Envelope.invalid_state(
message=(
f"task {t.id} is still over its cost budget: "
+219 -3
View File
@@ -1965,6 +1965,12 @@ class TaskService(BaseService):
# the reviewer deadlocks at post_pr_review (tracing_gap) and burns tokens
# into the do-server breaker. Cleared by complete_review.
task.active_claimant_id = cast("Any", reviewer_agent_id)
# Flip agent.status/current_task_id to ACTIVE/this-task (mirrors
# _finalize_claim / _qa_or_doc_claim) — the external-PR-review
# claim chokepoint (claim_pr_review), otherwise the reviewer never
# shows as active in the fleet either.
reviewer_agent = await self.session.get(AgentTable, reviewer_agent_id)
self._mark_agent_active_for_claim(reviewer_agent, task.id)
self._validate_and_set_status(
task, TaskStatus.CLAIMED, "pr_reviewer", audit_agent_id=reviewer_agent_id
)
@@ -2006,6 +2012,9 @@ class TaskService(BaseService):
task.assigned_to = None
task.claimed_by = None
task.active_claimant_id = cast("Any", None)
# The external-PR reviewer's claim ends here — release the fleet
# marker (mirrors pr_review_claim's own ACTIVE-marking).
await self._clear_agent_current_task(reviewer_id, task_id)
self._validate_and_set_status(
task,
TaskStatus.COMPLETED,
@@ -2963,7 +2972,11 @@ class TaskService(BaseService):
new_status in _REVIEW_QUEUE_STATES
and from_status != TaskStatus.BLOCKED.value
):
prior_claimant = to_python_uuid(task.active_claimant_id)
task.active_claimant_id = cast("Any", None)
# The stale claimant is no longer active on THIS task — release
# its fleet marker (mirrors every other claim-clearing site).
await self._clear_agent_current_task(prior_claimant, task_id)
task.status = new_status
await self.session.flush()
self._emit_status_transition_audit(
@@ -3291,6 +3304,18 @@ class TaskService(BaseService):
# throws — without restoring it, a retried claim sees the field set and
# the trust-but-verify short-circuits, never re-pushing the branch.
original_branch_name = task.branch_name
# agent.status/current_task_id too: this IS the one production
# chokepoint every claim verb (give_me_work -> i_will_work_on /
# i_will_plan, claim_review, claim_doc_task, claim_pr_review,
# claim_gate_review) routes through, so it's the single place to fix
# the "fleet always reports active: 0" bug — nothing else in the
# codebase ever set agent.status to ACTIVE or wrote
# current_task_id, so the Today brief's fleet/agents breakdown and
# /status's "Active agents" count were structurally incapable of
# reflecting real activity. Snapshot for the same rollback-on-
# branch-failure path as the task fields above.
original_agent_status = agent.status if agent else None
original_agent_task = agent.current_task_id if agent else None
now = datetime.now(UTC)
task.assigned_to = cast("Any", agent_id)
@@ -3308,6 +3333,7 @@ class TaskService(BaseService):
# declared but never written; now wired so the
# invariant is functional.
task.active_claimant_id = cast("Any", agent_id)
self._mark_agent_active_for_claim(agent, task.id)
agent_role = agent.role.value if agent and agent.role else None
if task.status in self._CLAIMABLE_STATUSES:
@@ -3332,6 +3358,9 @@ class TaskService(BaseService):
task.last_heartbeat_at = original_heartbeat
task.active_claimant_id = original_claimant_id
task.branch_name = original_branch_name
self._restore_agent_claim_snapshot(
agent, original_agent_status, original_agent_task
)
await self.session.flush()
# emit the reversal audit row so the journey doesn't diverge
# from real state. The forward ``task.claimed`` audit row was
@@ -3374,6 +3403,64 @@ class TaskService(BaseService):
self._background_tasks.add(bg_task)
bg_task.add_done_callback(self._background_tasks.discard)
def _mark_agent_active_for_claim(
self, agent: AgentTable | None, task_id: Any
) -> None:
"""Flip agent.status/current_task_id to ACTIVE/this-task on a
successful claim (or claim-equivalent resume) step. Nothing else in
the codebase ever set agent.status to ACTIVE or wrote
current_task_id, so the Today brief's fleet/agents breakdown and
/status's fleet counts were structurally incapable of reflecting
real activity (always active: 0).
Every production call site (every place an agent starts genuinely
working a task without a further claim call in between):
- ``_finalize_claim`` the dev/PM claim chokepoint (``claim()``),
with branch-failure rollback via ``_restore_agent_claim_snapshot``.
- ``_qa_or_doc_claim`` QA / Documenter / in-path PR-reviewer claim
(backs ``qa_claim``, ``doc_claim``, ``pr_gate_claim``). No
rollback needed no failure path follows the field writes.
- ``pr_review_claim`` the external-PR-review claim
(``claim_pr_review``).
- ``_apply_pre_block_restore`` ONLY when the pre-block snapshot
restores straight to IN_PROGRESS (a real resume with no fresh
claim call). A PENDING or review-queue restored status
deliberately does NOT mark active there see that method.
- ``unblock`` ONLY when the branch check resumes IN_PROGRESS
directly (same reasoning as the restore case above).
The release half (clearing current_task_id when a claim ends) is
``_clear_agent_current_task`` / ``_retarget_agent_claim`` see
their own docstrings for their call sites (``mark_agent_idle``,
every unclaim path, ``pass_qa``/``fail_qa``, ``pr_pass``/``pr_fail``,
``complete_review``, ``_maybe_advance_to_pm_review``,
``admin_set_status``, ``_admin_out_of_blocked``,
``_divert_owned_task_to_pool``, ``reassign_active_claim``).
ponytail: a single column can't represent a coordinator PM holding
several concurrent roots (PM claim-guard concurrency is intentional
see claim_guards.py) it just points at the MOST RECENT claim,
same ceiling the column already had before this fix (it was simply
never written at all). Upgrade path: a per-agent set of active task
ids, if multi-task fleet display is ever needed."""
if agent is None:
return
agent.status = AgentStatus.ACTIVE
agent.current_task_id = cast("Any", task_id)
def _restore_agent_claim_snapshot(
self,
agent: AgentTable | None,
status: AgentStatus | None,
current_task_id: Any,
) -> None:
"""Undo ``_mark_agent_active_for_claim`` when a branch-creation
failure rolls the whole claim back."""
if agent is None:
return
agent.status = cast("AgentStatus", status)
agent.current_task_id = cast("Any", current_task_id)
async def acquire_claim_lock(self, agent_id: UUID) -> None:
"""Take a per-agent transaction-scoped advisory lock.
@@ -4604,6 +4691,13 @@ class TaskService(BaseService):
task.claimed_by = owner
task.last_heartbeat_at = None
task.active_claimant_id = cast("Any", None)
# Ownership is preserved but the claim is released — the owner isn't
# actively working on this task right now (the reaper's holder is
# provably dead; the dependency-blocked owner is waiting on upstream
# work), so IDLE it rather than leaving it falsely reporting ACTIVE
# with a stale current_task_id until the eventual re-claim.
if owner is not None:
await self._retarget_agent_claim(cast("UUID", owner), None, task_id)
await self.session.flush()
return True
@@ -4657,9 +4751,13 @@ class TaskService(BaseService):
if task is None or task.assigned_to != agent_id:
return None
if task.status == TaskStatus.PENDING:
return await self._unclaim_pending_assignment(task)
result = await self._unclaim_pending_assignment(task)
await self._clear_agent_current_task(agent_id, task_id)
return result
if task.status == TaskStatus.BLOCKED:
return await self._unclaim_from_blocked(task)
result = await self._unclaim_from_blocked(task)
await self._clear_agent_current_task(agent_id, task_id)
return result
# verifying (self-verification before awaiting_qa) and needs_revision
# (QA/CEO sent it back) are both claims a dev can be sitting on when it
# decides to bail — a wedged dev retrying unclaim from either got a
@@ -4702,9 +4800,29 @@ class TaskService(BaseService):
task.work_session_id = cast("Any", None)
task.assigned_to = cast("Any", None)
task.active_claimant_id = cast("Any", None)
await self._clear_agent_current_task(agent_id, task_id)
await self.session.flush()
return task
async def _clear_agent_current_task(
self, agent_id: UUID | None, task_id: UUID
) -> None:
"""Clear the agent's ``current_task_id`` iff it still points at
``task_id`` every release path's specific side effect on the claim
marker (``_finalize_claim`` / ``_retarget_agent_claim`` /
``mark_agent_idle`` own writing it; this is the release half).
``agent_id=None`` (no prior claimant to release) is a no-op lets
every release call site skip its own None-guard. Leaves
``agent.status`` untouched: releasing THIS task doesn't mean the
agent is idle overall, only that it isn't the one working on it
anymore the next claim or ``i_am_idle`` call is authoritative for
status."""
if agent_id is None:
return
agent = await self.session.get(AgentTable, agent_id)
if agent is not None and agent.current_task_id == task_id:
agent.current_task_id = None
async def _unclaim_pending_assignment(self, task: TaskTable) -> TaskTable:
"""Release a never-claimed ``pending`` assignment (no status change).
@@ -5017,6 +5135,16 @@ class TaskService(BaseService):
# claim gate then holds it cleanly if its dependency is still unmet.
target = TaskStatus.IN_PROGRESS if task.branch_name else TaskStatus.PENDING
self._validate_and_set_status(task, target, agent_role)
owner_id = to_python_uuid(owner)
if target == TaskStatus.IN_PROGRESS and owner_id is not None:
# A real resume of active work with no fresh claim() call in
# between — mirrors _apply_pre_block_restore's own IN_PROGRESS
# case. The PENDING branch deliberately does NOT mark the owner
# active: it needs a fresh claim() to actually resume, and
# marking ahead of that would show it active before any
# container spawns.
owner_agent = await self.session.get(AgentTable, owner_id)
self._mark_agent_active_for_claim(owner_agent, task.id)
await self.session.flush()
self.log.info(
@@ -5260,6 +5388,9 @@ class TaskService(BaseService):
# route POST /pass-qa calls pass_qa directly — fix once at the
# shared transition so every caller is covered.
task.active_claimant_id = cast("Any", None)
# QA's claim on THIS task ends here — release the fleet marker
# (mirrors _qa_or_doc_claim's own ACTIVE-marking on the claim side).
await self._clear_agent_current_task(captured_qa_id, task_id)
task.qa_verified = True
# Reset docs so the documenter writes fresh docs for this cycle.
# DO NOT reset pr_created — the PR exists pre-QA under the current
@@ -5336,6 +5467,9 @@ class TaskService(BaseService):
# Store QA agent before reassigning
qa_agent_id = task.assigned_to
# QA's claim on THIS task ends here — release the fleet marker
# (mirrors _qa_or_doc_claim's own ACTIVE-marking on the claim side).
await self._clear_agent_current_task(to_python_uuid(qa_agent_id), task_id)
# Reassign to original developer so they can work on revisions
original_dev = extract_original_developer(task)
@@ -5634,6 +5768,14 @@ class TaskService(BaseService):
pr_created=task.pr_created,
)
if ready_for_pm:
# Capture the prior claimant (the documenter, or whoever held
# this review-queue claim) BEFORE it's overwritten below, so
# their fleet marker can be released once the reassign lands —
# they're no longer the active claimant, and the PM hasn't
# actually claimed (spawned) yet, so it must NOT be marked
# active in their place (that's `claim()`'s job, not this
# pre-assignment).
prior_claimant = to_python_uuid(task.active_claimant_id)
self._validate_and_set_status(
task, TaskStatus.AWAITING_PM_REVIEW, "documenter"
)
@@ -5650,6 +5792,7 @@ class TaskService(BaseService):
# check, which still sees the documenter as claimant.
task.claimed_by = cast("Any", owning_pm) if owning_pm else None
task.active_claimant_id = cast("Any", owning_pm) if owning_pm else None
await self._clear_agent_current_task(prior_claimant, task_id)
self.log.info(
"Documentation complete, awaiting PM review",
task_id=str(task_id),
@@ -10017,12 +10160,14 @@ class TaskService(BaseService):
if redirect.dev_notes_line is not None:
task.dev_notes = (task.dev_notes or "") + redirect.dev_notes_line
old_assignee = cast("UUID | None", task.claimed_by)
now = datetime.now(UTC)
task.assigned_to = cast("Any", effective_assignee)
task.claimed_by = cast("Any", effective_assignee)
task.claimed_at = now
task.last_heartbeat_at = now
task.active_claimant_id = cast("Any", effective_assignee)
await self._retarget_agent_claim(old_assignee, effective_assignee, task_id)
await self.session.flush()
self.log.info(
"Active task reassigned to a fresh claimant",
@@ -10031,8 +10176,39 @@ class TaskService(BaseService):
)
return task
async def _retarget_agent_claim(
self,
old_agent_id: UUID | None,
new_agent_id: UUID | None,
task_id: UUID,
) -> None:
"""Move the ACTIVE marker + ``current_task_id`` from the old
claimant to the new one on a handoff. The OTHER production
chokepoint (besides ``_finalize_claim``) that hands a task to an
agent ``reassign_active_claim`` needs the same write, or a
reassigned task keeps reporting its stale prior claimant as "the one
working on it" while the real new claimant never shows as active."""
if old_agent_id is not None and old_agent_id != new_agent_id:
old = await self.session.get(AgentTable, old_agent_id)
if old is not None and old.current_task_id == task_id:
old.status = AgentStatus.IDLE
old.current_task_id = None
if new_agent_id is None:
return
new_agent = await self.session.get(AgentTable, new_agent_id)
if new_agent is not None:
new_agent.status = AgentStatus.ACTIVE
new_agent.current_task_id = cast("Any", task_id)
async def mark_agent_idle(self, agent_id: UUID) -> None:
"""Set agent.status = IDLE + emit an ``agent.idle`` audit row.
"""Set agent.status = IDLE, clear current_task_id, + emit an
``agent.idle`` audit row.
Clearing ``current_task_id`` matters now that ``_finalize_claim`` /
``_retarget_agent_claim`` actually populate it on claim: without this,
an agent that goes idle (container about to shut down) would keep
reporting its last task as "currently working" forever, since
nothing else ever clears the column.
The audit row (target = the agent, details.agent_slug) is the idle
signal the member scorecard needs it lets the sweeper compute idle
@@ -10046,6 +10222,7 @@ class TaskService(BaseService):
if agent is None:
return
agent.status = AgentStatus.IDLE
agent.current_task_id = None
from roboco.db.tables import AuditLogTable
self.session.add(
@@ -10085,6 +10262,14 @@ class TaskService(BaseService):
role-scoped guard that only blocks a competing PR_REVIEWER claim;
a PM owning the root (non-reviewer claim) is intentionally
overridable by the first reviewer.
Also flips agent.status/current_task_id to ACTIVE/this-task
(``_mark_agent_active_for_claim``, mirroring ``_finalize_claim``)
this is the QA/Documenter/in-path-PR-reviewer claim chokepoint
(backs ``qa_claim``, ``doc_claim``, ``pr_gate_claim``), so without it
those three roles could never show as active in the fleet. No
branch-creation step follows (unlike ``_finalize_claim``), so there
is no failure path afterward to roll the marker back from.
"""
lock_result = await self.session.execute(
select(TaskTable)
@@ -10121,6 +10306,8 @@ class TaskService(BaseService):
# used by claimant_lock. Cleared by QA pass/fail
# and doc-complete when the review hand-off finishes.
task.active_claimant_id = cast("Any", agent_id)
agent = await self.session.get(AgentTable, agent_id)
self._mark_agent_active_for_claim(agent, task.id)
await self.session.flush()
return task
@@ -10320,6 +10507,9 @@ class TaskService(BaseService):
task.assigned_to = None
task.claimed_by = None
task.active_claimant_id = cast("Any", None)
# The gate reviewer's claim on THIS task ends here — release the
# fleet marker (mirrors _qa_or_doc_claim's own ACTIVE-marking).
await self._clear_agent_current_task(captured, task_id)
self._validate_and_set_status(
task,
TaskStatus.AWAITING_PM_REVIEW,
@@ -10366,6 +10556,9 @@ class TaskService(BaseService):
task.assigned_to = cast("Any", pm.id) if pm is not None else None
task.claimed_by = cast("Any", pm.id) if pm is not None else None
task.active_claimant_id = cast("Any", None)
# The gate reviewer's claim on THIS task ends here — release the
# fleet marker (mirrors _qa_or_doc_claim's own ACTIVE-marking).
await self._clear_agent_current_task(captured, task_id)
self._validate_and_set_status(
task,
TaskStatus.NEEDS_REVISION,
@@ -10511,6 +10704,19 @@ class TaskService(BaseService):
pre_status, restored_status, restored_owner = self._restore_block_ownership(
task, restored_status
)
restored_owner_id = to_python_uuid(restored_owner)
if restored_status == TaskStatus.IN_PROGRESS and restored_owner_id is not None:
# A real resume of active dev work with no fresh claim() call in
# between — _finalize_claim's own ACTIVE-marking never runs
# here, so this is the one restore outcome that needs the same
# write directly. PENDING (the branchless divert above) and any
# review-queue restored_status deliberately do NOT mark the
# owner active here: a fresh claim() / qa_claim / doc_claim /
# pr_gate_claim call is what actually resumes those, and marking
# ahead of that would show an agent active before its container
# ever spawns.
restored_agent = await self.session.get(AgentTable, restored_owner_id)
self._mark_agent_active_for_claim(restored_agent, task.id)
await self.session.flush()
# This restore path sets the status directly (bypassing the strict
# transition validator), so emit the audit explicitly — no status
@@ -10602,7 +10808,11 @@ class TaskService(BaseService):
# Review/queue targets are re-claimed via the claim verbs — a stale
# escalation claim surviving the override strands the next claimant
# (give_me_work hands out the task while its note() writes bounce).
prior_claimant = to_python_uuid(task.active_claimant_id)
self._clear_claim_and_block_snapshot(task)
# The stale claimant is no longer active on THIS task — release
# its fleet marker too (mirrors admin_set_status's own release).
await self._clear_agent_current_task(prior_claimant, cast("UUID", task.id))
return None
def _clear_claim_and_block_snapshot(self, task: TaskTable) -> None:
@@ -10846,6 +11056,12 @@ class TaskService(BaseService):
task.active_claimant_id = cast("Any", None)
task.status = TaskStatus.PENDING
task.dev_notes = (task.dev_notes or "") + note
# The refused owner is genuinely no longer engaged with this task at
# all (ownership cleared entirely, not just re-routed) — idle its
# fleet marker, mirroring _force_unclaim_to_pending's reaper release.
await self._retarget_agent_claim(
to_python_uuid(prior_owner), None, cast("UUID", task.id)
)
await self.session.flush()
self._emit_status_transition_audit(
task,
+25 -17
View File
@@ -43,14 +43,13 @@ from typing import TYPE_CHECKING, Any, ClassVar, cast
from uuid import UUID
import structlog
from sqlalchemy import func, select
from roboco.config import settings
from roboco.db.base import get_session_factory
from roboco.db.tables import AgentTable, AuditLogTable
from roboco.db.tables import AuditLogTable
from roboco.foundation.policy.content import markers
from roboco.foundation.policy.content.validators import reject_trivial
from roboco.models.base import AgentStatus, TaskStatus
from roboco.models.base import TaskStatus
from roboco.seeds.initial_data import AGENT_UUIDS
from roboco.services import telegram_bridge as bridge
from roboco.services.base import BaseService, ValidationError
@@ -68,7 +67,6 @@ from roboco.services.telegram_credentials import get_telegram_credentials_servic
from roboco.services.tg_cockpit import get_tg_cockpit_service
from roboco.services.tiktok_client import build_tiktok_poster
from roboco.services.tiktok_credentials import get_tiktok_credentials_service
from roboco.services.usage import get_usage_service
from roboco.services.video_post_service import TaskAlreadyCompletedError as _VideoDone
from roboco.services.video_post_service import (
VideoCaptionTooLongError,
@@ -564,21 +562,25 @@ class TelegramInboundEngine(BaseService):
)
async def _render_status(self) -> str:
"""Cheap snapshot: active-agent count + task counts by status — no
spend/strategy/pitch queries (that's the heavier cockpit summary).
Statuses render in lifecycle order (only the nonzero ones), not
alphabetically a CEO scanning on a phone reads top-to-bottom as the
pipeline, not as an a-z dump."""
"""Cheap snapshot: fleet status breakdown + task counts by status —
no spend/strategy/pitch queries (that's the heavier cockpit summary).
The fleet breakdown shares the SAME by_status derivation as
``/agents`` and the Today brief (``TgCockpitService.fleet`` ->
``DashboardService.get_all_agent_status``) so the three surfaces can
never disagree on what "active" means. Task statuses render in
lifecycle order (only the nonzero ones), not alphabetically a CEO
scanning on a phone reads top-to-bottom as the pipeline, not as an
a-z dump."""
counts = await get_task_service(self.session).count_by_status()
active_result = await self.session.execute(
select(func.count(AgentTable.id)).where(
AgentTable.status == AgentStatus.ACTIVE
)
fleet = await get_tg_cockpit_service(self.session).fleet()
by_status = fleet.get("by_status", {})
fleet_line = " · ".join(
f"<b>{by_status.get(status, 0)}</b> {status}"
for status in ("active", "idle", "offline")
)
active = active_result.scalar_one()
lines = [
"<b>🤖 Fleet</b>",
f"Active agents: <b>{active}</b>",
fleet_line,
"",
"<b>📋 Tasks</b>",
]
@@ -751,9 +753,15 @@ class TelegramInboundEngine(BaseService):
return _truncate("\n".join(lines))
async def _render_usage(self) -> str:
"""Today's spend from the day rollup — the Today brief's number."""
summary = await get_usage_service(self.session).get_today_summary()
"""Today's spend — the Today brief's own number (display-timezone
bucketed), so the bot and the Mini App never disagree on "today"."""
summary = await get_tg_cockpit_service(self.session).today_spend()
tokens = int(summary.get("tokens_today", 0))
if summary.get("subscription_billed"):
return (
f"<b>💸 Spend today</b>\n≈$0 — subscription (untracked) · "
f"{tokens:,} tokens"
)
cost = float(summary.get("cost_today_usd", 0.0))
return f"<b>💸 Spend today</b>\n${cost:,.2f} · {tokens:,} tokens"
+109 -40
View File
@@ -1,33 +1,43 @@
"""TgCockpitService — the Mini App's one-round-trip "Today" brief.
Composes existing read paths (TaskService listers, DashboardService's agent
snapshot, UsageService's day rollup) into a single phone-sized payload
snapshot, raw agent_spawn_sessions rows) into a single phone-sized payload
answering "does anything need me?". Deliberately DB-only and cheap: no live
GitHub calls, no release-readiness snapshot (that path clones + shells out to
git), no orchestrator singleton the ship-state red/green proxy is the set
of open ci_watch fix tasks, which exists precisely when a watched repo went
red.
"Today" and the trailing 7-day series bucket by calendar day in
``settings.display_timezone`` (default UTC, a no-op for anyone who hasn't set
it), not the server's UTC day — a GMT+2 CEO's evening activity used to land
on the wrong display day. This intentionally bypasses ``UsageService.
get_today_summary`` (which reads the UTC-keyed ``daily_usage_rollups`` table,
unchanged and still used by the main CEO dashboard) and instead reads raw
``agent_spawn_sessions`` rows directly, bucketing them in Python by the
display timezone the rollup table's key stays UTC-only, but the underlying
timestamped rows let the COCKPIT'S read side be timezone-correct without
touching the rollup pipeline.
"""
from __future__ import annotations
from datetime import UTC, date, datetime, timedelta
from typing import TYPE_CHECKING, Any
from sqlalchemy import cast as sql_cast
from sqlalchemy import func, select
from sqlalchemy.types import Date
from sqlalchemy import select
from roboco.billing.pricing import is_ollama_cloud_model
from roboco.config import settings
from roboco.db.tables import AgentSpawnSessionTable, TaskTable
from roboco.foundation.policy import display_time
from roboco.foundation.policy.content import markers
from roboco.models.base import TaskStatus
from roboco.services.base import BaseService
from roboco.services.dashboard import get_dashboard_service
from roboco.services.task import get_task_service
from roboco.services.usage import get_usage_service
if TYPE_CHECKING:
from datetime import date
from uuid import UUID
from sqlalchemy.ext.asyncio import AsyncSession
@@ -73,9 +83,9 @@ class TgCockpitService(BaseService):
}
def _window_dates(self) -> list[date]:
"""The last ``_SERIES_DAYS`` calendar dates (UTC), oldest → today."""
today = datetime.now(UTC).date()
return [today - timedelta(days=n) for n in reversed(range(_SERIES_DAYS))]
"""The last ``_SERIES_DAYS`` calendar dates in the display timezone,
oldest -> today."""
return display_time.trailing_dates(settings.display_timezone, _SERIES_DAYS)
async def _needs_you(self, tasks: TaskService) -> dict[str, Any]:
awaiting = await tasks.list_awaiting_ceo_approval()
@@ -141,62 +151,121 @@ class TgCockpitService(BaseService):
# Mirrors the CEO dashboard overview: a usage hiccup degrades the
# brief to zeros instead of failing the whole endpoint.
try:
summary = await get_usage_service(self.session).get_today_summary()
series = await self._spend_series()
today = series[-1] if series else 0.0
prior = series[-2] if len(series) >= 2 else 0.0 # noqa: PLR2004
days = self._window_dates()
(
cost_by_day,
tokens_by_day,
models_by_day,
) = await self._session_metrics_by_day(days)
series = [round(cost_by_day.get(d, 0.0), 4) for d in days]
today_cost = series[-1] if series else 0.0
prior_cost = series[-2] if len(series) >= 2 else 0.0 # noqa: PLR2004
today_tokens = tokens_by_day.get(days[-1], 0) if days else 0
today_models = models_by_day.get(days[-1], set()) if days else set()
# $0 with real tokens spent, on a subscription-billed-but-
# untracked model (an Ollama Cloud tag with no grounded rate) —
# "subscription (untracked)", never a bare misleading "$0".
subscription_billed = (
today_cost == 0.0
and today_tokens > 0
and any(is_ollama_cloud_model(m) for m in today_models)
)
return {
"tokens_today": int(summary.get("tokens_today", 0)),
"cost_today_usd": float(summary.get("cost_today_usd", 0.0)),
"tokens_today": today_tokens,
"cost_today_usd": today_cost,
"subscription_billed": subscription_billed,
"series": series,
"delta_pct": _pct_change(today, prior),
"delta_pct": _pct_change(today_cost, prior_cost),
}
except Exception: # pragma: no cover - defensive degradation
self.log.warning("today-brief usage summary failed", exc_info=True)
return {
"tokens_today": 0,
"cost_today_usd": 0.0,
"subscription_billed": False,
"series": [0.0] * _SERIES_DAYS,
"delta_pct": None,
}
async def _spend_series(self) -> list[float]:
"""Per-day cost (USD) over the trailing window, zero-filled — the
hero sparkline. Grouped by the spawn session's start date, matching
what today's spend counts."""
day = sql_cast(AgentSpawnSessionTable.started_at, Date).label("day")
async def today_spend(self) -> dict[str, Any]:
"""Just today's tokens/cost/subscription_billed, no series — the
bot's ``/usage`` command doesn't need the 7-day sparkline. Shares the
Today brief's own ``_spend`` computation so the two can never
silently disagree on what "today" means."""
spend = await self._spend()
return {
"tokens_today": spend["tokens_today"],
"cost_today_usd": spend["cost_today_usd"],
"subscription_billed": spend["subscription_billed"],
}
async def _session_metrics_by_day(
self, days: list[date]
) -> tuple[dict[date, float], dict[date, int], dict[date, set[str]]]:
"""One raw ``agent_spawn_sessions`` query spanning the whole window,
bucketed into display-timezone calendar days: cost, total tokens, and
the distinct models used per day. Backs both the spend hero (today +
the 7-day series) and the subscription-billed detection off the SAME
read, rather than two queries that could silently disagree on what
"today" means."""
tz = settings.display_timezone
start_utc, _ = display_time.day_bounds_utc(tz, days[0])
_, end_utc = display_time.day_bounds_utc(tz, days[-1])
result = await self.session.execute(
select(
day,
func.coalesce(
func.sum(AgentSpawnSessionTable.estimated_cost_usd), 0.0
).label("cost"),
AgentSpawnSessionTable.started_at,
AgentSpawnSessionTable.estimated_cost_usd,
AgentSpawnSessionTable.tokens_input,
AgentSpawnSessionTable.tokens_output,
AgentSpawnSessionTable.tokens_cache_read,
AgentSpawnSessionTable.tokens_cache_write,
AgentSpawnSessionTable.model,
).where(
AgentSpawnSessionTable.started_at >= start_utc,
AgentSpawnSessionTable.started_at < end_utc,
)
.where(
sql_cast(AgentSpawnSessionTable.started_at, Date)
>= self._window_dates()[0]
)
.group_by(day)
)
by_day = {row.day: float(row.cost) for row in result}
return [round(by_day.get(d, 0.0), 4) for d in self._window_dates()]
cost_by_day: dict[date, float] = {}
tokens_by_day: dict[date, int] = {}
models_by_day: dict[date, set[str]] = {}
for row in result:
d = display_time.local_date(row.started_at, tz)
cost_by_day[d] = cost_by_day.get(d, 0.0) + float(
row.estimated_cost_usd or 0.0
)
tokens_by_day[d] = tokens_by_day.get(d, 0) + (
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)
)
models_by_day.setdefault(d, set()).add(row.model or "")
return cost_by_day, tokens_by_day, models_by_day
async def _velocity(self) -> dict[str, Any]:
"""Per-day completed-task counts over the trailing window (the
'shipped this week' bars) plus the window total."""
'shipped this week' bars) plus the window total, bucketed by the
display timezone."""
try:
day = sql_cast(TaskTable.completed_at, Date).label("day")
days = self._window_dates()
tz = settings.display_timezone
start_utc, _ = display_time.day_bounds_utc(tz, days[0])
_, end_utc = display_time.day_bounds_utc(tz, days[-1])
result = await self.session.execute(
select(day, func.count(TaskTable.id).label("n"))
.where(
select(TaskTable.completed_at).where(
TaskTable.status == TaskStatus.COMPLETED,
TaskTable.completed_at.isnot(None),
sql_cast(TaskTable.completed_at, Date) >= self._window_dates()[0],
TaskTable.completed_at >= start_utc,
TaskTable.completed_at < end_utc,
)
.group_by(day)
)
by_day = {row.day: int(row.n) for row in result}
series = [by_day.get(d, 0) for d in self._window_dates()]
by_day: dict[date, int] = {}
for completed_at in result.scalars():
if completed_at is None: # excluded by the WHERE clause already
continue
d = display_time.local_date(completed_at, tz)
by_day[d] = by_day.get(d, 0) + 1
series = [by_day.get(d, 0) for d in days]
return {"series": series, "week_total": sum(series)}
except Exception: # pragma: no cover - defensive degradation
self.log.warning("today-brief velocity failed", exc_info=True)
@@ -756,11 +756,20 @@ async def test_pr_review_gate_pass_path(
assert claimed is not None
assert str(claimed.status) == Status.AWAITING_PR_REVIEW.value
assert claimed.assigned_to == reviewer.id
# pr_gate_claim (via _qa_or_doc_claim) flips the reviewer's fleet marker.
reviewer_row = await db_session.get(AgentTable, reviewer_id)
assert reviewer_row is not None
assert reviewer_row.status == AgentStatus.ACTIVE
assert reviewer_row.current_task_id == task.id
passed = await svc.pr_pass(reviewer_id, task.id, notes="integration verified")
assert passed is not None
assert str(passed.status) == Status.AWAITING_PM_REVIEW.value
assert passed.assigned_to is None # cleared so the PM-closure dispatch routes
# pr_pass releases the reviewer's fleet marker too.
reviewer_row = await db_session.get(AgentTable, reviewer_id)
assert reviewer_row is not None
assert reviewer_row.current_task_id is None
final = await svc.get(task.id)
assert final is not None
+16 -7
View File
@@ -13,6 +13,7 @@ from fastapi import FastAPI
from httpx import ASGITransport, AsyncClient
from roboco.api.deps import get_agent_context, get_db
from roboco.api.routes.provider import router as provider_router
from roboco.billing.pricing import input_price_per_million
from roboco.db.tables import (
ModelAssignmentTable,
ProviderConfigTable,
@@ -30,11 +31,19 @@ if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
def _first_model_for_type(provider_type: ModelProvider) -> str:
def _unpriced_model_for_type(provider_type: ModelProvider) -> str:
"""The first `provider_type` catalog entry pricing.py has NOT grounded a
real per-token rate for tests below want "an unpriced, free-tier
downgrade-safe model" specifically to exercise provider-readiness gating,
not pricing itself, so grounding a real rate for one catalog entry (e.g.
GLM-5.2) must not silently break them by picking that one."""
for entry in MODEL_CATALOG:
if entry.provider_type == provider_type:
if (
entry.provider_type == provider_type
and input_price_per_million(entry.model_name) == 0.0
):
return entry.model_name
raise RuntimeError(f"no catalog entry for {provider_type}")
raise RuntimeError(f"no unpriced catalog entry for {provider_type}")
def _make_app(
@@ -846,12 +855,12 @@ async def test_put_complexity_override_allows_same_tier_as_baseline(
async def test_put_complexity_override_rejects_disabled_provider(
app_client_with_ollama: AsyncClient,
) -> None:
"""qa's baseline (haiku) prices no cheaper than Ollama Cloud (unpriced,
treated as free-tier) so the downgrade-only check passes but the
"""qa's baseline (haiku) prices no cheaper than an unpriced Ollama Cloud
model (free-tier) so the downgrade-only check passes but the
OLLAMA_CLOUD provider is disabled (no key set) in this fixture's seeded
state, so the write-time readiness guard rejects it before it can
silently no-op to the legacy Anthropic path at spawn."""
ollama_model = _first_model_for_type(ModelProvider.OLLAMA_CLOUD)
ollama_model = _unpriced_model_for_type(ModelProvider.OLLAMA_CLOUD)
response = await app_client_with_ollama.put(
"/api/providers/complexity-overrides",
json={"role": "qa", "complexity": "low", "model_name": ollama_model},
@@ -875,7 +884,7 @@ async def test_put_complexity_override_warns_on_cross_family_once_provider_ready
json={"api_key": "test-key"},
headers=_HDR_PM,
)
ollama_model = _first_model_for_type(ModelProvider.OLLAMA_CLOUD)
ollama_model = _unpriced_model_for_type(ModelProvider.OLLAMA_CLOUD)
response = await app_client_with_ollama.put(
"/api/providers/complexity-overrides",
json={"role": "qa", "complexity": "low", "model_name": ollama_model},
+183 -6
View File
@@ -989,6 +989,12 @@ async def test_unblock_restores_to_in_progress(
# Owner restored into both fields so the dev dispatcher respawns it.
assert unblocked.assigned_to == task_setup["agent_id"]
assert unblocked.claimed_by == task_setup["agent_id"]
# A real resume with no fresh claim() call — unblock must flip the
# owner's fleet marker itself (mirrors _finalize_claim/_qa_or_doc_claim).
owner_row = await db_session.get(AgentTable, task_setup["agent_id"])
assert owner_row is not None
assert owner_row.status == AgentStatus.ACTIVE
assert owner_row.current_task_id == task.id
@pytest.mark.asyncio
@@ -1021,6 +1027,67 @@ async def test_unblock_keeps_owner_for_give_me_work_claim(
assert unblocked.claimed_by == task_setup["agent_id"]
@pytest.mark.asyncio
async def test_admin_set_status_into_review_queue_releases_agent_marker(
task_setup: dict, db_session: AsyncSession
) -> None:
"""A non-blocked admin override into a review/queue state clears the
stale claimant's active_claimant_id (M19) — it must release that
claimant's fleet marker too, or a dead escalation claim keeps reporting
an agent as active forever."""
svc = task_setup["svc"]
dev_id = task_setup["agent_id"]
task = await svc.create(_req(task_setup))
task.status = TaskStatus.IN_PROGRESS
task.assigned_to = dev_id
task.claimed_by = dev_id
task.active_claimant_id = dev_id
await db_session.flush()
dev_agent = await db_session.get(AgentTable, dev_id)
assert dev_agent is not None
dev_agent.status = AgentStatus.ACTIVE
dev_agent.current_task_id = task.id
await db_session.flush()
out = await svc.admin_set_status(task.id, TaskStatus.AWAITING_QA)
assert out is not None
assert out.active_claimant_id is None
dev_agent = await db_session.get(AgentTable, dev_id)
assert dev_agent is not None
assert dev_agent.current_task_id is None
@pytest.mark.asyncio
async def test_divert_owned_task_to_pool_idles_prior_owner(
task_setup: dict, db_session: AsyncSession
) -> None:
"""_divert_owned_task_to_pool clears ownership entirely — the refused
owner isn't engaged with this task at all anymore, so its fleet marker
must be released (mirrors _force_unclaim_to_pending's reaper release)."""
svc = task_setup["svc"]
owner_id = task_setup["agent_id"]
task = await svc.create(_req(task_setup))
task.status = TaskStatus.IN_PROGRESS
task.assigned_to = owner_id
task.claimed_by = owner_id
await db_session.flush()
owner_agent = await db_session.get(AgentTable, owner_id)
assert owner_agent is not None
owner_agent.status = AgentStatus.ACTIVE
owner_agent.current_task_id = task.id
await db_session.flush()
await svc._divert_owned_task_to_pool(task, note="test diversion")
assert task.status == TaskStatus.PENDING
assert task.assigned_to is None
owner_agent = await db_session.get(AgentTable, owner_id)
assert owner_agent is not None
assert owner_agent.status == AgentStatus.IDLE
assert owner_agent.current_task_id is None
# ---------------------------------------------------------------------------
# QA + completion happy paths
# ---------------------------------------------------------------------------
@@ -1068,13 +1135,21 @@ async def test_pass_qa_clears_active_claimant_for_doc_claim(
task.status = TaskStatus.AWAITING_QA
task.pr_number = 42
task.pr_url = "https://github.com/x/y/pull/42"
task.assigned_to = qa_id
task.claimed_by = qa_id
task.active_claimant_id = qa_id
await db_session.flush()
# qa_claim (via _qa_or_doc_claim) flips the QA agent's fleet marker —
# verify it, then verify pass_qa releases it symmetrically.
qa_claimed = await svc.qa_claim(qa_id, task.id)
assert qa_claimed is not None
qa_agent = await db_session.get(AgentTable, qa_id)
assert qa_agent is not None
assert qa_agent.status == AgentStatus.ACTIVE
assert qa_agent.current_task_id == task.id
passed = await svc.pass_qa(task.id, notes="LGTM", agent_role="qa")
assert passed is not None
assert passed.active_claimant_id is None
qa_agent = await db_session.get(AgentTable, qa_id)
assert qa_agent is not None
assert qa_agent.current_task_id is None
doc = AgentTable(
id=uuid4(),
name="Doc",
@@ -1093,6 +1168,10 @@ async def test_pass_qa_clears_active_claimant_for_doc_claim(
claimed = await svc.doc_claim(doc.id, task.id)
assert claimed is not None
assert to_uuid(claimed.active_claimant_id) == doc.id
doc_agent = await db_session.get(AgentTable, doc.id)
assert doc_agent is not None
assert doc_agent.status == AgentStatus.ACTIVE
assert doc_agent.current_task_id == task.id
@pytest.mark.asyncio
@@ -1106,13 +1185,15 @@ async def test_fail_qa_clears_active_claimant(
qa_id = task_setup["agent_id"]
task = await svc.create(_req(task_setup))
task.status = TaskStatus.AWAITING_QA
task.assigned_to = qa_id
task.claimed_by = qa_id
task.active_claimant_id = qa_id
await db_session.flush()
assert await svc.qa_claim(qa_id, task.id) is not None
failed = await svc.fail_qa(task.id, notes="please fix X")
assert failed is not None
assert failed.active_claimant_id is None
# fail_qa releases the QA agent's fleet marker too.
qa_agent = await db_session.get(AgentTable, qa_id)
assert qa_agent is not None
assert qa_agent.current_task_id is None
# ---------------------------------------------------------------------------
@@ -1368,6 +1449,70 @@ async def test_claim_pending_with_unmet_dependency_returns_none(
assert claimed.status == TaskStatus.CLAIMED
@pytest.mark.asyncio
async def test_claim_sets_agent_active_and_current_task(
task_setup: dict, db_session: AsyncSession
) -> None:
"""_finalize_claim is the one production chokepoint every claim verb
routes through before this fix, nothing ever wrote agent.status=ACTIVE
or current_task_id, so the fleet/Today-brief breakdown could never show
a real "active" agent or populate "working[]"."""
svc = task_setup["svc"]
task = await svc.create(_req(task_setup))
task.branch_name = "feature/backend/abcd1234"
await db_session.flush()
claimed = await svc.claim(task.id, task_setup["agent_id"])
assert claimed is not None
agent = await db_session.get(AgentTable, task_setup["agent_id"])
assert agent is not None
assert agent.status == AgentStatus.ACTIVE
assert agent.current_task_id == task.id
@pytest.mark.asyncio
async def test_mark_agent_idle_clears_current_task(
task_setup: dict, db_session: AsyncSession
) -> None:
"""Idling an agent must clear current_task_id — otherwise it keeps
reporting the last-claimed task as "currently working" forever, since
nothing else ever clears the column."""
svc = task_setup["svc"]
task = await svc.create(_req(task_setup))
task.branch_name = "feature/backend/abcd1234"
await db_session.flush()
await svc.claim(task.id, task_setup["agent_id"])
await svc.mark_agent_idle(task_setup["agent_id"])
agent = await db_session.get(AgentTable, task_setup["agent_id"])
assert agent is not None
assert agent.status == AgentStatus.IDLE
assert agent.current_task_id is None
@pytest.mark.asyncio
async def test_unclaim_for_agent_clears_current_task(
task_setup: dict, db_session: AsyncSession
) -> None:
"""A voluntary unclaim releases the claim marker on the AGENT too, not
just the task otherwise the fleet keeps showing the agent as working
on a task it just gave up."""
svc = task_setup["svc"]
task = await svc.create(_req(task_setup))
task.branch_name = "feature/backend/abcd1234"
await db_session.flush()
await svc.claim(task.id, task_setup["agent_id"])
result = await svc.unclaim_for_agent(task.id, task_setup["agent_id"])
assert result is not None
agent = await db_session.get(AgentTable, task_setup["agent_id"])
assert agent is not None
assert agent.current_task_id is None
# ---------------------------------------------------------------------------
# Sequence claim guardrail (CEO directive: sequence is the bar, independent
# of dependency_ids — see _claim_blocked_by_sequence).
@@ -1900,6 +2045,33 @@ async def test_unclaim_for_reaper_resets(
assert refreshed.claimed_by == task_setup["agent_id"]
@pytest.mark.asyncio
async def test_unclaim_for_reaper_idles_the_provably_dead_holder(
task_setup: dict, db_session: AsyncSession
) -> None:
"""The reaper's holder is provably dead (heartbeat past TTL) — it must
stop reporting ACTIVE with a stale current_task_id, or the fleet keeps
showing a dead agent as "working" until the eventual re-claim."""
svc = task_setup["svc"]
agent = await db_session.get(AgentTable, task_setup["agent_id"])
assert agent is not None
task = await svc.create(_req(task_setup))
task.status = TaskStatus.CLAIMED
task.assigned_to = task_setup["agent_id"]
task.claimed_by = task_setup["agent_id"]
task.active_claimant_id = task_setup["agent_id"]
agent.status = AgentStatus.ACTIVE
agent.current_task_id = task.id
await db_session.flush()
await svc.unclaim_for_reaper(task.id)
refreshed_agent = await db_session.get(AgentTable, task_setup["agent_id"])
assert refreshed_agent is not None
assert refreshed_agent.status == AgentStatus.IDLE
assert refreshed_agent.current_task_id is None
# ---------------------------------------------------------------------------
# resume_for_agent
# ---------------------------------------------------------------------------
@@ -2100,6 +2272,11 @@ async def test_pr_gate_claim_allows_first_reviewer_when_pm_owns_root(
assert task.active_claimant_id == reviewer.id
assert task.claimed_by == reviewer.id
assert task.assigned_to == reviewer.id
# pr_gate_claim (via _qa_or_doc_claim) flips the reviewer's fleet marker.
reviewer_row = await db_session.get(AgentTable, reviewer.id)
assert reviewer_row is not None
assert reviewer_row.status == AgentStatus.ACTIVE
assert reviewer_row.current_task_id == task.id
@pytest.mark.asyncio
@@ -2151,6 +2151,11 @@ async def test_docs_complete_advance_clears_stale_documenter_claim(
task.pr_number = 1
task.pr_url = "u"
task.pr_created = True
# Simulate the documenter having genuinely claimed it (what
# _qa_or_doc_claim's own ACTIVE-marking would have set).
documenter_agent = await db_session.get(AgentTable, documenter_id)
assert documenter_agent is not None
documenter_agent.current_task_id = task.id
await db_session.flush()
out = await svc.docs_complete(task.id, doc_notes="documented all flows")
@@ -2162,6 +2167,15 @@ async def test_docs_complete_advance_clears_stale_documenter_claim(
assert out.claimed_by == pm_agent.id
assert out.active_claimant_id == pm_agent.id
assert out.claimed_by != documenter_id
# The outgoing documenter's fleet marker is released...
documenter_agent = await db_session.get(AgentTable, documenter_id)
assert documenter_agent is not None
assert documenter_agent.current_task_id is None
# ...but the PM is NOT marked active by this pre-assignment — the PM
# hasn't actually claimed (spawned) yet, only claim() does that.
pm_row = await db_session.get(AgentTable, pm_agent.id)
assert pm_row is not None
assert pm_row.current_task_id is None
assert out.active_claimant_id != documenter_id
+99 -8
View File
@@ -22,6 +22,7 @@ from roboco.billing.pricing import (
calculate_cost,
calculate_cost_result,
input_price_per_million,
is_ollama_cloud_model,
)
# ---------------------------------------------------------------------------
@@ -68,6 +69,14 @@ _CODEX_OUTPUT = 14.00
_CODEX_CACHE_READ = 0.175
_CODEX_CACHE_WRITE = 1.75
# Z.ai GLM-5.2 — priced non-Anthropic (Ollama Cloud's `glm-5.2:cloud` tag,
# subscription-billed but attributed at the API-equivalent rate). Source:
# https://docs.z.ai/guides/overview/pricing (fetched 2026-07-23).
_GLM_INPUT = 1.40
_GLM_OUTPUT = 4.40
_GLM_CACHE_READ = 0.26
_GLM_CACHE_WRITE = 1.40
# Tolerance for floating-point comparisons
_TOL = 1e-4
@@ -386,6 +395,57 @@ class TestCodexTier:
assert _CODEX_OUTPUT > _CODEX_INPUT
# ---------------------------------------------------------------------------
# GLM-5.2 tier (Ollama Cloud — priced non-Anthropic, grounded in a citable
# published rate; see the module's pricing-table comment for the source).
# ---------------------------------------------------------------------------
class TestGlmTier:
"""glm-5.2:cloud pricing — Ollama Cloud, priced like grok-build/codex."""
def test_input_only(self) -> None:
cost = calculate_cost("glm-5.2:cloud", tokens_input=_M, tokens_output=0)
assert abs(cost - _GLM_INPUT) < _TOL
def test_output_only(self) -> None:
cost = calculate_cost("glm-5.2:cloud", tokens_input=0, tokens_output=_M)
assert abs(cost - _GLM_OUTPUT) < _TOL
def test_cached_input(self) -> None:
cost = calculate_cost(
"glm-5.2:cloud", tokens_input=0, tokens_output=0, tokens_cache_read=_M
)
assert abs(cost - _GLM_CACHE_READ) < _TOL
def test_cache_write(self) -> None:
cost = calculate_cost(
"glm-5.2:cloud", tokens_input=0, tokens_output=0, tokens_cache_write=_M
)
assert abs(cost - _GLM_CACHE_WRITE) < _TOL
def test_all_token_types(self) -> None:
cost = calculate_cost(
"glm-5.2:cloud",
tokens_input=_M,
tokens_output=_M,
tokens_cache_read=_M,
tokens_cache_write=_M,
)
expected = _GLM_INPUT + _GLM_OUTPUT + _GLM_CACHE_READ + _GLM_CACHE_WRITE
assert abs(cost - expected) < _TOL
def test_glm_is_not_treated_as_anthropic(self) -> None:
assert _is_anthropic_model("glm-5.2:cloud") is False
assert calculate_cost("glm-5.2:cloud", tokens_input=_M, tokens_output=0) > 0.0
def test_bare_glm_tag_without_cloud_suffix_still_prices(self) -> None:
"""The fragment match is on 'glm-5.2', independent of the ':cloud'
tag suffix a differently-tagged variant still resolves."""
cost = calculate_cost("glm-5.2", tokens_input=_M, tokens_output=0)
assert abs(cost - _GLM_INPUT) < _TOL
# ---------------------------------------------------------------------------
# Unknown / edge cases — must return 0.0 without raising
# ---------------------------------------------------------------------------
@@ -477,16 +537,23 @@ class TestSubstringMatchPriority:
class TestProviderAwareness:
"""Non-Anthropic models (local Ollama / Ollama Cloud) cost 0.0 per token."""
"""Genuinely-free local Ollama costs 0.0 per token; an ungrounded Ollama
Cloud model also costs 0.0 (we have no rate for it see
``is_ollama_cloud_model`` for the caller-side distinction from "free"). A
GROUNDED Ollama Cloud model (glm-5.2) is priced for real see
``TestGlmTier``."""
def test_ollama_prefixed_model_returns_zero(self) -> None:
"""Self-hosted Ollama models (``ollama/`` prefix) have no API cost."""
cost = calculate_cost("ollama/llama3", tokens_input=_M, tokens_output=_M)
assert cost == _ZERO_COST
def test_ollama_cloud_model_returns_zero(self) -> None:
"""Ollama Cloud (``:cloud`` tag) is subscription-billed, not per token."""
cost = calculate_cost("glm-5.2:cloud", tokens_input=_M, tokens_output=_M)
def test_ungrounded_ollama_cloud_model_returns_zero(self) -> None:
"""An Ollama Cloud (``:cloud`` tag) model with no table entry has no
rate to price from still 0.0 (not "unpriced"; see TestCostResult)."""
cost = calculate_cost(
"some-future-model:cloud", tokens_input=_M, tokens_output=_M
)
assert cost == _ZERO_COST
def test_bare_local_model_returns_zero(self) -> None:
@@ -589,6 +656,25 @@ def test_sonnet5_promo_active_on_or_before_2026_08_31(
)
class TestIsOllamaCloudModel:
"""The ':cloud' tag convention, shared by pricing.py's own attribution
logic and external callers (the TG cockpit's spend label)."""
def test_cloud_tagged_model_is_cloud(self) -> None:
assert is_ollama_cloud_model("glm-5.2:cloud") is True
assert is_ollama_cloud_model("SOME-MODEL:CLOUD") is True
def test_local_model_is_not_cloud(self) -> None:
assert is_ollama_cloud_model("ollama/llama3") is False
assert is_ollama_cloud_model("qwen3-embedding:0.6b") is False
def test_anthropic_model_is_not_cloud(self) -> None:
assert is_ollama_cloud_model("claude-sonnet-5") is False
def test_empty_string_is_not_cloud(self) -> None:
assert is_ollama_cloud_model("") is False
def test_sonnet5_reverts_to_list_rate_after_2026_08_31(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -635,12 +721,17 @@ class TestInputPricePerMillion:
)
def test_unpriced_non_anthropic_model_is_free_tier(self) -> None:
"""A self-hosted / Ollama Cloud model has no per-token rate — treated
as the cheapest possible tier, so it can never be rejected as
"costlier" by the downgrade-only policy."""
assert input_price_per_million("glm-5.2:cloud") == 0.0
"""A self-hosted / ungrounded-Ollama-Cloud model has no per-token
rate treated as the cheapest possible tier, so it can never be
rejected as "costlier" by the downgrade-only policy."""
assert input_price_per_million("some-future-model:cloud") == 0.0
assert input_price_per_million("my-custom-self-hosted-model:7b") == 0.0
def test_grounded_ollama_cloud_model_has_real_rate(self) -> None:
"""GLM-5.2 is now grounded in a real published rate, unlike an
unpriced Ollama Cloud model."""
assert input_price_per_million("glm-5.2:cloud") == _GLM_INPUT
def test_empty_model_returns_zero(self) -> None:
assert input_price_per_million("") == 0.0
@@ -0,0 +1,94 @@
"""display_time — pure display-timezone day bucketing, incl. DST boundaries."""
from __future__ import annotations
from datetime import UTC, date, datetime, timedelta
from roboco.foundation.policy.display_time import (
day_bounds_utc,
local_date,
resolve_zone,
trailing_dates,
)
class TestResolveZone:
def test_known_zone(self) -> None:
assert resolve_zone("Europe/Berlin").key == "Europe/Berlin"
def test_utc_default(self) -> None:
assert resolve_zone("UTC").key == "UTC"
def test_unknown_zone_falls_back_to_utc(self) -> None:
assert resolve_zone("Not/AZone").key == "UTC"
class TestLocalDate:
def test_utc_noop(self) -> None:
instant = datetime(2026, 7, 23, 10, 0, tzinfo=UTC)
assert local_date(instant, "UTC") == date(2026, 7, 23)
def test_gmt_plus_2_evening_utc_is_next_day_local(self) -> None:
"""22:30 UTC on the 22nd is 00:30 the NEXT day in GMT+2 — the exact
'CEO's evening activity lands on the wrong display day' bug."""
instant = datetime(2026, 7, 22, 22, 30, tzinfo=UTC)
assert local_date(instant, "Europe/Berlin") == date(2026, 7, 23)
def test_gmt_plus_2_early_morning_utc_is_prior_day_local(self) -> None:
# No — GMT+2 is AHEAD of UTC, so early UTC morning is still the same
# local day; use a clearly-behind zone for the "prior day" case.
instant = datetime(2026, 7, 23, 2, 0, tzinfo=UTC)
assert local_date(instant, "America/Los_Angeles") == date(2026, 7, 22)
class TestTrailingDates:
def test_seven_days_oldest_to_today(self) -> None:
now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
dates = trailing_dates("UTC", 7, now=now)
assert len(dates) == 7 # noqa: PLR2004
assert dates[-1] == date(2026, 7, 23)
assert dates[0] == date(2026, 7, 17)
assert dates == sorted(dates)
def test_timezone_shifts_which_day_is_today(self) -> None:
"""23:00 UTC on the 22nd is already the 23rd in Europe/Berlin."""
now = datetime(2026, 7, 22, 23, 0, tzinfo=UTC)
assert trailing_dates("UTC", 1, now=now) == [date(2026, 7, 22)]
assert trailing_dates("Europe/Berlin", 1, now=now) == [date(2026, 7, 23)]
class TestDayBoundsUtc:
def test_utc_day_is_exactly_24h(self) -> None:
start, end = day_bounds_utc("UTC", date(2026, 7, 23))
assert start == datetime(2026, 7, 23, 0, 0, tzinfo=UTC)
assert end == datetime(2026, 7, 24, 0, 0, tzinfo=UTC)
assert (end - start).total_seconds() == 24 * 3600
def test_gmt_plus_2_day_bounds(self) -> None:
start, end = day_bounds_utc("Europe/Berlin", date(2026, 7, 23))
# Summer time (CEST, UTC+2): local midnight is 22:00 UTC the day before.
assert start == datetime(2026, 7, 22, 22, 0, tzinfo=UTC)
assert end == datetime(2026, 7, 23, 22, 0, tzinfo=UTC)
def test_dst_spring_forward_day_is_23_hours(self) -> None:
"""Europe/Berlin springs forward on the last Sunday of March —
2026-03-29 02:00 CET -> 03:00 CEST so that local day is only 23h
of real UTC time, not 24."""
start, end = day_bounds_utc("Europe/Berlin", date(2026, 3, 29))
assert (end - start).total_seconds() == 23 * 3600
def test_dst_fall_back_day_is_25_hours(self) -> None:
"""Europe/Berlin falls back on the last Sunday of October —
2026-10-25 03:00 CEST -> 02:00 CET so that local day is 25h."""
start, end = day_bounds_utc("Europe/Berlin", date(2026, 10, 25))
assert (end - start).total_seconds() == 25 * 3600
def test_an_instant_at_start_falls_in_this_day(self) -> None:
start, _ = day_bounds_utc("Europe/Berlin", date(2026, 7, 23))
assert local_date(start, "Europe/Berlin") == date(2026, 7, 23)
def test_an_instant_just_before_end_falls_in_this_day(self) -> None:
_, end = day_bounds_utc("Europe/Berlin", date(2026, 7, 23))
assert local_date(end - timedelta(seconds=1), "Europe/Berlin") == date(
2026, 7, 23
)
+38 -1
View File
@@ -8,11 +8,12 @@ agents_config data) and the reaper-safe service write
from __future__ import annotations
from typing import cast
from unittest.mock import AsyncMock, MagicMock
from uuid import UUID, uuid4
import pytest
from roboco.models.base import TaskStatus
from roboco.models.base import AgentStatus, TaskStatus
from roboco.seeds.initial_data import AGENT_UUIDS
from roboco.services.gateway.choreographer._impl import Choreographer
from roboco.services.task import TaskService
@@ -98,6 +99,12 @@ def _build_task(**over: object) -> MagicMock:
def _service() -> TaskService:
session = MagicMock()
session.flush = AsyncMock()
# reassign_active_claim now retargets the agent-side claim marker
# (_retarget_agent_claim), which reads old/new agent rows via
# session.get — default to "no matching row" so tests that don't care
# about the agent side effect stay a no-op there, same as before this
# write existed.
session.get = AsyncMock(return_value=None)
return TaskService(session)
@@ -123,3 +130,33 @@ async def test_reassign_active_claim_refuses_non_active_status() -> None:
svc = _service()
object.__setattr__(svc, "get", AsyncMock(return_value=task))
assert await svc.reassign_active_claim(task.id, uuid4()) is None
@pytest.mark.asyncio
async def test_reassign_active_claim_retargets_agent_active_marker() -> None:
"""The old claimant's ACTIVE/current_task_id marker must move to the new
claimant otherwise the fleet keeps showing the SUPERSEDED agent as
working on this task, and never shows the real new claimant as active."""
old_id, new_id = uuid4(), uuid4()
task = _build_task(status=TaskStatus.IN_PROGRESS, claimed_by=old_id)
old_agent = MagicMock(status=AgentStatus.ACTIVE, current_task_id=task.id)
new_agent = MagicMock(status=AgentStatus.IDLE, current_task_id=None)
svc = _service()
object.__setattr__(svc, "get", AsyncMock(return_value=task))
async def _fake_get(_model: object, agent_id: object) -> object:
if agent_id == old_id:
return old_agent
if agent_id == new_id:
return new_agent
return None
cast("MagicMock", svc.session).get = AsyncMock(side_effect=_fake_get)
result = await svc.reassign_active_claim(task.id, new_id)
assert result is task
assert old_agent.status == AgentStatus.IDLE
assert old_agent.current_task_id is None
assert new_agent.status == AgentStatus.ACTIVE
assert new_agent.current_task_id == task.id
+8 -10
View File
@@ -236,16 +236,16 @@ async def test_handle_breach_skips_a_task_that_already_moved_on() -> None:
# ---------------------------------------------------------------------------
# _task_budget_breach: cap resolution (null -> TaskType default) + spend sum
# _task_budget_breach: explicit-input-only cap (null budget = never a breach)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_breach_falls_back_to_tasktype_default_when_budget_null() -> None:
"""Cap resolution (task.budget_usd null -> TaskType default) and spend
both delegate to TaskService now (task_spend_usd's own open-session
pricing is covered directly by its shared implementation see
test_project_month_spend_usd_db.py's real-DB open-session case)."""
async def test_null_budget_is_never_a_breach() -> None:
"""Budgets enforce only when explicitly set: a task with no budget_usd is
uncapped, regardless of spend the spend query is never even issued (the
old per-TaskType default table blocked a default-budget coordination root
one opus planning turn in)."""
orch = _make_orchestrator()
task_id = "44444444-4444-4444-4444-444444444444"
task = MagicMock(
@@ -262,10 +262,8 @@ async def test_breach_falls_back_to_tasktype_default_when_budget_null() -> None:
):
breach = await orch._task_budget_breach(task_id)
assert breach is not None
cap_usd, spend_usd = breach
assert cap_usd == 1.0 # TASK_TYPE_DEFAULT_BUDGET_USD[DOCUMENTATION]
assert spend_usd == _MOCK_TASK_SPEND_USD
assert breach is None
task_svc.task_spend_usd.assert_not_awaited()
@pytest.mark.asyncio
@@ -201,6 +201,9 @@ async def test_fail_qa_emits_auditor_alert(
"""``fail_qa`` calls the auditor rework producer with QA attribution."""
session = MagicMock()
session.flush = AsyncMock()
# fail_qa releases the QA agent's fleet marker via session.get — default
# to "no matching row" for a test that doesn't care about that side effect.
session.get = AsyncMock(return_value=None)
task = _mock_task(status=TaskStatus.AWAITING_QA)
task.orchestration_markers = {"original_developer": str(uuid4())}
@@ -235,6 +238,9 @@ async def test_pr_fail_emits_auditor_alert(
"""``pr_fail`` calls the auditor rework producer with reviewer attribution."""
session = MagicMock()
session.flush = AsyncMock()
# pr_fail releases the reviewer's fleet marker via session.get — default
# to "no matching row" for a test that doesn't care about that side effect.
session.get = AsyncMock(return_value=None)
reviewer_id = uuid4()
pm_id = uuid4()
task = _mock_task(status=TaskStatus.AWAITING_PR_REVIEW)
@@ -34,6 +34,11 @@ def _bind(svc: TaskService, name: str, value: object) -> None:
def _service() -> TaskService:
session = MagicMock()
session.flush = AsyncMock()
# reassign_active_claim now retargets the agent-side claim marker
# (_retarget_agent_claim), which reads agent rows via session.get —
# default to "no matching row" so tests that don't care about the
# agent side effect stay a no-op there.
session.get = AsyncMock(return_value=None)
return TaskService(session)
@@ -611,6 +616,9 @@ async def test_unblock_with_restore_emits_audit_event() -> None:
``AuditLogTable`` added to the session."""
session = MagicMock()
session.flush = AsyncMock()
# An IN_PROGRESS restore now looks up the restored owner via session.get
# to flip its ACTIVE marker — default to "no matching row".
session.get = AsyncMock(return_value=None)
added: list[object] = []
session.add.side_effect = added.append
svc = TaskService(session)
+10
View File
@@ -253,6 +253,12 @@ async def test_pr_review_claim_and_complete(db_session: AsyncSession) -> None:
assert claimed.active_claimant_id is not None
assert UUID(str(claimed.active_claimant_id)) == reviewer_id
assert claimed.claimed_at is not None
# The external-PR-review claim chokepoint flips the reviewer's fleet
# marker too — otherwise pr_reviewer agents never show as active.
reviewer_row = await db_session.get(AgentTable, reviewer_id)
assert reviewer_row is not None
assert reviewer_row.status == AgentStatus.ACTIVE
assert reviewer_row.current_task_id == task_id
# Re-claiming a non-pending task is a no-op.
assert await svc.pr_review_claim(reviewer_id, task_id) is None
@@ -269,6 +275,10 @@ async def test_pr_review_claim_and_complete(db_session: AsyncSession) -> None:
assert done.claimed_by is None
# Single-claimant lock cleared on completion (the review hand-off is done).
assert done.active_claimant_id is None
# complete_review releases the reviewer's fleet marker too.
reviewer_row = await db_session.get(AgentTable, reviewer_id)
assert reviewer_row is not None
assert reviewer_row.current_task_id is None
# Re-completing a completed task is a no-op.
assert await svc.complete_review(reviewer_id, task_id) is None
+92 -9
View File
@@ -567,12 +567,15 @@ async def test_wire_sibling_collision_dag_notifies_only_for_new_edges() -> None:
@pytest.mark.asyncio
async def test_mark_agent_idle_sets_status_idle() -> None:
agent = MagicMock(id=uuid4(), status=AgentStatus.ACTIVE)
agent = MagicMock(id=uuid4(), status=AgentStatus.ACTIVE, current_task_id=uuid4())
result = MagicMock()
result.scalar_one_or_none.return_value = agent
svc = _service_with(result)
await svc.mark_agent_idle(agent.id)
assert agent.status == AgentStatus.IDLE
# Otherwise the agent keeps reporting its last task as "currently
# working" forever — nothing else ever clears this column.
assert agent.current_task_id is None
# ---------------------------------------------------------------------------
@@ -587,6 +590,10 @@ async def test_qa_claim_sets_assignment_on_awaiting_qa() -> None:
result.scalar_one_or_none.return_value = task
session = MagicMock(flush=AsyncMock())
session.execute = AsyncMock(return_value=result)
# _qa_or_doc_claim now looks up the claiming agent via session.get to
# flip its ACTIVE marker — default to "no matching row" for a test that
# doesn't care about that side effect.
session.get = AsyncMock(return_value=None)
svc = TaskService(session)
qa_id = uuid4()
out = await svc.qa_claim(qa_id, task.id)
@@ -617,6 +624,7 @@ async def test_doc_claim_sets_assignment_on_awaiting_documentation() -> None:
result.scalar_one_or_none.return_value = task
session = MagicMock(flush=AsyncMock())
session.execute = AsyncMock(return_value=result)
session.get = AsyncMock(return_value=None)
svc = TaskService(session)
doc_id = uuid4()
out = await svc.doc_claim(doc_id, task.id)
@@ -677,7 +685,10 @@ async def test_unblock_with_restore_returns_to_pre_block_state() -> None:
blocker_resolver_type=BlockerResolverType.AGENT,
blocker_raised_by=pre_assignee,
)
svc = TaskService(MagicMock(flush=AsyncMock()))
# An IN_PROGRESS restore now looks up the restored owner via session.get
# to flip its ACTIVE marker — default to "no matching row".
session = MagicMock(flush=AsyncMock(), get=AsyncMock(return_value=None))
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
out = await svc.unblock_with_restore(uuid4(), task.id, restore=True)
assert out is task
@@ -721,13 +732,17 @@ async def test_unblock_no_branch_returns_to_pending() -> None:
task = _build_task(
status=TaskStatus.BLOCKED, branch_name=None, blocker_raised_by=raiser
)
svc = TaskService(MagicMock(flush=AsyncMock()))
session = MagicMock(flush=AsyncMock())
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
_bind(svc, "_index_lifecycle_event_background", AsyncMock())
out = await svc.unblock(task.id)
assert out is task
assert task.status == TaskStatus.PENDING
assert task.assigned_to == raiser
# A PENDING restore is NOT a resume — the owner isn't marked active here;
# a fresh claim() is what actually resumes it, so no agent lookup runs.
session.get.assert_not_called()
@pytest.mark.asyncio
@@ -747,7 +762,10 @@ async def test_admin_set_status_out_of_blocked_restores_pre_block_owner() -> Non
pre_block_state="in_progress",
pre_block_assignee=dev,
)
svc = TaskService(MagicMock(flush=AsyncMock()))
# An IN_PROGRESS restore now looks up the restored owner (dev) via
# session.get to flip its ACTIVE marker.
session = MagicMock(flush=AsyncMock(), get=AsyncMock(return_value=None))
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
out = await svc.admin_set_status(task.id, TaskStatus.IN_PROGRESS)
assert out is task
@@ -791,7 +809,10 @@ async def test_admin_set_status_into_review_queue_clears_active_claimant() -> No
claimed_by=dev,
active_claimant_id=dev,
)
svc = TaskService(MagicMock(flush=AsyncMock()))
# Clearing the stale claimant now looks it up via session.get to release
# its ACTIVE marker too.
session = MagicMock(flush=AsyncMock(), get=AsyncMock(return_value=None))
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
out = await svc.admin_set_status(task.id, TaskStatus.AWAITING_QA)
assert out is task
@@ -943,7 +964,10 @@ async def test_admin_set_status_blocked_to_review_state_clears_claim() -> None:
pre_block_state="awaiting_pm_review",
pre_block_assignee=pm,
)
svc = TaskService(MagicMock(flush=AsyncMock()))
# Clearing the stale claim now looks it up via session.get to release
# its ACTIVE marker too.
session = MagicMock(flush=AsyncMock(), get=AsyncMock(return_value=None))
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
out = await svc.admin_set_status(task.id, TaskStatus.AWAITING_PM_REVIEW)
assert out is task
@@ -970,7 +994,10 @@ async def test_admin_set_status_blocked_to_needs_revision_clears_claim() -> None
pre_block_state="awaiting_pm_review",
pre_block_assignee=pm,
)
svc = TaskService(MagicMock(flush=AsyncMock()))
# Clearing the stale claim now looks it up via session.get to release
# its ACTIVE marker too.
session = MagicMock(flush=AsyncMock(), get=AsyncMock(return_value=None))
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
out = await svc.admin_set_status(task.id, TaskStatus.NEEDS_REVISION)
assert out is task
@@ -1085,7 +1112,10 @@ async def test_admin_set_status_pre_block_restore_syncs_active_claimant() -> Non
pre_block_state="in_progress",
pre_block_assignee=dev,
)
svc = TaskService(MagicMock(flush=AsyncMock()))
# An IN_PROGRESS restore now looks up the restored owner (dev) via
# session.get to flip its ACTIVE marker.
session = MagicMock(flush=AsyncMock(), get=AsyncMock(return_value=None))
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
out = await svc.admin_set_status(task.id, TaskStatus.IN_PROGRESS)
assert out is task
@@ -1566,7 +1596,10 @@ async def test_unblock_with_branch_resumes_in_progress() -> None:
branch_name="feature/backend/abc12345",
blocker_raised_by=uuid4(),
)
svc = TaskService(MagicMock(flush=AsyncMock()))
# Resuming IN_PROGRESS now looks up the restored owner via session.get
# to flip its ACTIVE marker.
session = MagicMock(flush=AsyncMock(), get=AsyncMock(return_value=None))
svc = TaskService(session)
_bind(svc, "get", AsyncMock(return_value=task))
_bind(svc, "_index_lifecycle_event_background", AsyncMock())
out = await svc.unblock(task.id)
@@ -1809,6 +1842,56 @@ async def test_finalize_claim_rollback_emits_reversal_audit() -> None:
assert {"from": "claimed", "to": "pending"} in audit_calls
@pytest.mark.asyncio
async def test_finalize_claim_sets_agent_active_then_rolls_back_on_failure() -> None:
"""_finalize_claim must flip agent.status/current_task_id to ACTIVE/this
task BEFORE the branch step runs (previously nothing ever wrote these
fields at all the fleet-status bug), and roll them back to their
pre-claim values on a branch-creation failure, same as the task fields.
"""
session = MagicMock()
session.flush = AsyncMock()
svc = TaskService(session)
task = _build_task(
status=TaskStatus.PENDING,
branch_name=None,
project_id=uuid4(),
product_id=None,
batch_id=None,
parent_task_id=None,
cell_projects=[],
pr_created=False,
pr_number=None,
)
agent = MagicMock(
id=uuid4(),
role=AgentRole.DEVELOPER,
status=AgentStatus.IDLE,
current_task_id=None,
)
_bind(svc, "_emit_status_transition_audit", MagicMock())
captured: dict[str, object] = {}
async def _boom(_task: object, _agent_id: object) -> str:
captured["status"] = agent.status
captured["current_task_id"] = agent.current_task_id
raise RuntimeError("branch boom")
_bind(svc, "_ensure_branch_for_task", _boom)
with pytest.raises(RuntimeError, match="branch boom"):
await svc._finalize_claim(task, agent, agent.id)
# Set to ACTIVE/this-task before the branch step ran...
assert captured["status"] == AgentStatus.ACTIVE
assert captured["current_task_id"] == task.id
# ...and rolled back to the pre-claim values once branch creation failed.
assert agent.status == AgentStatus.IDLE
assert agent.current_task_id is None
@pytest.mark.asyncio
async def test_emit_status_transition_audit_writes_in_session_atomically() -> None:
"""The status-transition audit row is written into the CALLER's session (same
@@ -30,6 +30,11 @@ def _bind(svc: TaskService, name: str, value: object) -> None:
def _service() -> TaskService:
session = MagicMock()
session.flush = AsyncMock()
# reassign_active_claim now retargets the agent-side claim marker
# (_retarget_agent_claim), which reads agent rows via session.get —
# default to "no matching row" so tests that don't care about the
# agent side effect stay a no-op there.
session.get = AsyncMock(return_value=None)
return TaskService(session)
@@ -96,6 +96,33 @@ async def test_run_cycle_syncs_commands_exactly_once(
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_render_status_shares_fleet_derivation_with_render_agents(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The /status fleet line must come from the SAME by_status breakdown
/agents and the Today brief use (TgCockpitService.fleet) a second,
independent agent-status query previously could (and did) disagree."""
fleet: dict[str, Any] = {
"total": 27,
"by_status": {"active": 3, "idle": 20, "offline": 4},
"working": [],
}
cockpit = MagicMock(fleet=AsyncMock(return_value=fleet))
monkeypatch.setattr(ti, "get_tg_cockpit_service", lambda _s: cockpit)
tasks = MagicMock(
count_by_status=AsyncMock(return_value={"in_progress": 5, "pending": 2})
)
monkeypatch.setattr(ti, "get_task_service", lambda _s: tasks)
text = await _engine()._render_status()
assert "3</b> active" in text
assert "20</b> idle" in text
assert "4</b> offline" in text
assert "in_progress" in text
@pytest.mark.asyncio
async def test_render_agents_lists_working_agents(
monkeypatch: pytest.MonkeyPatch,
@@ -124,12 +151,16 @@ async def test_render_agents_lists_working_agents(
async def test_render_usage_formats_today_summary(
monkeypatch: pytest.MonkeyPatch,
) -> None:
usage = MagicMock(
get_today_summary=AsyncMock(
return_value={"tokens_today": 2_400_000, "cost_today_usd": 18.7}
cockpit = MagicMock(
today_spend=AsyncMock(
return_value={
"tokens_today": 2_400_000,
"cost_today_usd": 18.7,
"subscription_billed": False,
}
)
)
monkeypatch.setattr(ti, "get_usage_service", lambda _s: usage)
monkeypatch.setattr(ti, "get_tg_cockpit_service", lambda _s: cockpit)
text = await _engine()._render_usage()
@@ -137,6 +168,30 @@ async def test_render_usage_formats_today_summary(
assert "2,400,000 tokens" in text
@pytest.mark.asyncio
async def test_render_usage_labels_subscription_billed_spend(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""An untracked-subscription spend day (Ollama Cloud, no grounded rate)
must never render as a bare, misleading '$0.00'."""
cockpit = MagicMock(
today_spend=AsyncMock(
return_value={
"tokens_today": 456_221,
"cost_today_usd": 0.0,
"subscription_billed": True,
}
)
)
monkeypatch.setattr(ti, "get_tg_cockpit_service", lambda _s: cockpit)
text = await _engine()._render_usage()
assert "subscription (untracked)" in text
assert "456,221 tokens" in text
assert "$0.00" not in text
@pytest.mark.asyncio
async def test_render_blocked_sections_and_links(
monkeypatch: pytest.MonkeyPatch,
+131 -1
View File
@@ -9,12 +9,13 @@ assumed.
from __future__ import annotations
from datetime import UTC, date, datetime, timedelta
from typing import TYPE_CHECKING, cast
from uuid import uuid4
import pytest
from roboco.config import settings
from roboco.db.tables import AgentTable, TaskTable
from roboco.db.tables import AgentSpawnSessionTable, AgentTable, TaskTable
from roboco.foundation import identity as _foundation
from roboco.foundation.policy.content import markers
from roboco.models.base import (
@@ -43,6 +44,28 @@ SYSTEM_UUID = _foundation.AGENTS["system"].uuid
CI_WATCH_SOURCE = "ci_watch"
_COST_TOL = 0.005
_TOKENS_FLOOR = 1
def _spawn_session(
*, started_at: datetime, model: str, cost: float, tokens_input: int = 1000
) -> AgentSpawnSessionTable:
return AgentSpawnSessionTable(
id=uuid4(),
agent_slug=f"be-dev-{uuid4().hex[:6]}",
team="backend",
role="developer",
model=model,
task_id=None,
started_at=started_at,
ended_at=started_at + timedelta(minutes=5),
tokens_input=tokens_input,
tokens_output=0,
estimated_cost_usd=cost,
)
# 1 awaiting + 1 blocked + 4 held drafts (release/x/video/roadmap-item).
EXPECTED_NEEDS_YOU_TOTAL = 6
@@ -195,3 +218,110 @@ async def test_today_composes_needs_you_fleet_and_ship(
assert brief["ship"]["open_release_proposal"] is True
assert brief["ship"]["ci_fix_tasks"] == baseline["ship"]["ci_fix_tasks"] + 1
# ---------------------------------------------------------------------------
# Display-timezone bucketing (Issue 2) — "today" is display_timezone-aware,
# not always the server's UTC day. `_session_metrics_by_day` is called
# directly with an explicit historical `days` window so the test is fully
# deterministic (disconnected from the real "now").
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_session_metrics_by_day_buckets_by_display_timezone(
db_session: AsyncSession, monkeypatch: pytest.MonkeyPatch
) -> None:
"""23:30 UTC on the 15th is already 00:30 on the 16th in Europe/Berlin
(winter, CET = UTC+1) the exact 'evening activity lands on the wrong
display day' bug this fix targets."""
started = datetime(2026, 1, 15, 23, 30, tzinfo=UTC)
session_row = _spawn_session(started_at=started, model="claude-sonnet-5", cost=1.23)
db_session.add(session_row)
await db_session.flush()
svc = get_tg_cockpit_service(db_session)
monkeypatch.setattr(settings, "display_timezone", "UTC")
cost_utc, tokens_utc, _ = await svc._session_metrics_by_day([date(2026, 1, 15)])
assert cost_utc.get(date(2026, 1, 15), 0.0) >= _COST_TOL
assert tokens_utc.get(date(2026, 1, 15), 0) >= _TOKENS_FLOOR
monkeypatch.setattr(settings, "display_timezone", "Europe/Berlin")
cost_berlin, tokens_berlin, _ = await svc._session_metrics_by_day(
[date(2026, 1, 16)]
)
assert cost_berlin.get(date(2026, 1, 16), 0.0) >= _COST_TOL
assert tokens_berlin.get(date(2026, 1, 16), 0) >= _TOKENS_FLOOR
# And the SAME row must NOT double-count into the UTC calendar day under
# the Berlin bucketing — the 15th should now come up empty for this row.
cost_berlin_15, _, _ = await svc._session_metrics_by_day([date(2026, 1, 15)])
assert cost_berlin_15.get(date(2026, 1, 15), 0.0) < _COST_TOL
@pytest.mark.asyncio
async def test_window_dates_shifts_with_display_timezone(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""`_window_dates` (real 'now') must reflect the configured display
timezone, not always UTC."""
svc = get_tg_cockpit_service(cast("AsyncSession", None))
monkeypatch.setattr(settings, "display_timezone", "UTC")
utc_dates = svc._window_dates()
monkeypatch.setattr(settings, "display_timezone", "Pacific/Kiritimati")
# UTC+14 — the furthest-ahead real timezone; "today" there is never
# earlier, and is later whenever UTC hasn't crossed its own midnight yet.
kiritimati_dates = svc._window_dates()
assert kiritimati_dates[-1] >= utc_dates[-1]
# ---------------------------------------------------------------------------
# Ollama Cloud honesty-labeling (Issue 1) — an ungrounded ':cloud' model's $0
# is flagged subscription_billed, never rendered as a bare misleading "$0".
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_today_spend_flags_ungrounded_ollama_cloud_as_subscription_billed(
db_session: AsyncSession,
) -> None:
now = datetime.now(UTC)
db_session.add(
_spawn_session(started_at=now, model="some-future-model:cloud", cost=0.0)
)
await db_session.flush()
summary = await get_tg_cockpit_service(db_session).today_spend()
assert summary["cost_today_usd"] == pytest.approx(0.0)
assert summary["subscription_billed"] is True
@pytest.mark.asyncio
async def test_today_spend_not_subscription_billed_when_priced(
db_session: AsyncSession,
) -> None:
"""A real per-token cost (even from a priced Ollama Cloud model like
GLM-5.2) is never mislabeled as an untracked subscription figure."""
now = datetime.now(UTC)
db_session.add(_spawn_session(started_at=now, model="glm-5.2:cloud", cost=2.5))
await db_session.flush()
summary = await get_tg_cockpit_service(db_session).today_spend()
assert summary["subscription_billed"] is False
@pytest.mark.asyncio
async def test_today_spend_not_subscription_billed_for_local_ollama(
db_session: AsyncSession,
) -> None:
"""A genuinely-free self-hosted model (no ':cloud' tag) at $0 is just
free, not an untracked subscription."""
now = datetime.now(UTC)
db_session.add(_spawn_session(started_at=now, model="ollama/llama3", cost=0.0))
await db_session.flush()
summary = await get_tg_cockpit_service(db_session).today_spend()
assert summary["subscription_billed"] is False
+21
View File
@@ -200,3 +200,24 @@ def test_resolve_uvicorn_loop_factory_uvloop_returns_new_event_loop() -> None:
assert isinstance(loop, uvloop.Loop)
finally:
loop.close()
# ---------------------------------------------------------------------------
# display_timezone — the TG cockpit's day-bucketing timezone (Issue 2)
# ---------------------------------------------------------------------------
def test_display_timezone_defaults_to_utc() -> None:
"""Default is a no-op for every deployment that doesn't set it."""
assert Settings().display_timezone == "UTC"
def test_display_timezone_accepts_valid_iana_name() -> None:
assert Settings(display_timezone="Europe/Berlin").display_timezone == (
"Europe/Berlin"
)
def test_display_timezone_rejects_unknown_name() -> None:
with pytest.raises(ValidationError):
Settings(display_timezone="Not/AZone")