mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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:
@@ -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();
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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},
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user