diff --git a/alembic/versions/080_task_project_budgets.py b/alembic/versions/080_task_project_budgets.py new file mode 100644 index 00000000..f9bcd388 --- /dev/null +++ b/alembic/versions/080_task_project_budgets.py @@ -0,0 +1,59 @@ +"""Per-task and per-project cost budgets (feature-flagged, default off). + +``tasks.budget_usd`` caps one task's own accumulated agent-spawn spend; +``projects.monthly_budget_usd`` caps a project's calendar-month spend across +all its tasks. Both nullable and additive — null means "no cap" and is a pure +no-op regardless of ``ROBOCO_TASK_BUDGETS_ENABLED`` (the flag itself gates +whether the caps are ever consulted at all). Mirrors the ``ci_watch_enabled`` +per-project opt-in shape (migration 048): the column exists unconditionally, +the feature flag decides whether anything reads it. + +``ix_agent_spawn_sessions_task_id`` rides along: both the claim-time monthly- +spend query (``TaskService.project_month_spend_usd``) and the orchestrator's +per-tick task-spend sweep (``_task_spend_usd``) filter this table by +``task_id`` with no existing index — an unbounded per-row scan every minute +once the flag is armed. Added here rather than a separate migration since it +exists to serve these same two new read paths. + +Revision ID: 080_task_project_budgets +Revises: 079_notification_backoff +Create Date: 2026-07-22 + +NOTE: down_revision was re-chained from 078_project_codegen_command to +079_notification_backoff (PR #652 landed first at integration) — the +079_task_project_budgets -> 080_task_project_budgets rename + re-chain is +exactly the "may be re-chained at integration" case flagged in the original +docstring, not a fork of the tree. +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "080_task_project_budgets" +down_revision = "079_notification_backoff" +branch_labels: dict[str, str] | None = None +depends_on: dict[str, str] | None = None + + +def upgrade() -> None: + op.add_column( + "tasks", + sa.Column("budget_usd", sa.Float(), nullable=True), + ) + op.add_column( + "projects", + sa.Column("monthly_budget_usd", sa.Float(), nullable=True), + ) + op.create_index( + "ix_agent_spawn_sessions_task_id", + "agent_spawn_sessions", + ["task_id"], + ) + + +def downgrade() -> None: + op.drop_index("ix_agent_spawn_sessions_task_id", table_name="agent_spawn_sessions") + op.drop_column("projects", "monthly_budget_usd") + op.drop_column("tasks", "budget_usd") diff --git a/panel/src/components/projects/__tests__/edit-project-dialog.test.tsx b/panel/src/components/projects/__tests__/edit-project-dialog.test.tsx index 2bc66c0e..3183868d 100644 --- a/panel/src/components/projects/__tests__/edit-project-dialog.test.tsx +++ b/panel/src/components/projects/__tests__/edit-project-dialog.test.tsx @@ -5,6 +5,9 @@ import type { ReactNode } from "react"; import React from "react"; import { Team } from "@/types"; import type { Project } from "@/types"; +import { toast } from "sonner"; + +vi.mock("sonner", () => ({ toast: { success: vi.fn(), error: vi.fn() } })); // jsdom has no ResizeObserver; Radix Switch (the always-rendered "Active" // toggle) measures its thumb via one on mount — mirrors @@ -117,6 +120,7 @@ function makeProject(overrides: Partial = {}): Project { video_engine_enabled: false, dep_update_command: null, dep_update_paths: null, + monthly_budget_usd: null, sandbox_services: null, sandbox_extensions: null, workspace_path: null, @@ -385,3 +389,104 @@ describe("EditProjectDialog — Protected Branches", () => { expect(call.updates.protected_branches).toEqual(["master", "slave"]); }); }); + +describe("EditProjectDialog — Monthly Budget (USD)", () => { + beforeEach(() => { + vi.clearAllMocks(); + getCredentialsStatus.mockResolvedValue({ has_credentials: true }); + mutateAsync.mockResolvedValue(makeProject()); + useUpdateProject.mockReturnValue({ mutateAsync, isPending: false }); + }); + + function openAutonomySection() { + fireEvent.click( + screen.getByRole("button", { name: /Show Autonomous Maintenance/i }), + ); + } + + // fireEvent.submit(form) rather than clicking the Save button — this + // dialog's Tabs-wrapped form doesn't reliably translate a button click + // into a submit event under jsdom; submitting the form directly is the + // same idiom create-task-dialog.test.tsx already uses. + function submit() { + fireEvent.submit(document.querySelector("form")!); + } + + it("pre-fills the stored monthly_budget_usd", async () => { + renderDialog(makeProject({ monthly_budget_usd: 42 })); + await screen.findByRole("button", { name: /Save Changes/i }); + openAutonomySection(); + + expect(screen.getByLabelText(/Monthly Budget/i)).toHaveValue(42); + }); + + it("rejects 0 with an inline error and does not submit", async () => { + renderDialog(makeProject({ monthly_budget_usd: null })); + await screen.findByRole("button", { name: /Save Changes/i }); + openAutonomySection(); + + fireEvent.change(screen.getByLabelText(/Monthly Budget/i), { + target: { value: "0" }, + }); + submit(); + + await waitFor(() => { + expect(toast.error).toHaveBeenCalledWith( + expect.stringMatching(/greater than 0/i), + ); + }); + expect(mutateAsync).not.toHaveBeenCalled(); + }); + + it("rejects a negative budget the same way", async () => { + renderDialog(makeProject({ monthly_budget_usd: null })); + await screen.findByRole("button", { name: /Save Changes/i }); + openAutonomySection(); + + fireEvent.change(screen.getByLabelText(/Monthly Budget/i), { + target: { value: "-5" }, + }); + submit(); + + await waitFor(() => { + expect(toast.error).toHaveBeenCalledWith( + expect.stringMatching(/greater than 0/i), + ); + }); + expect(mutateAsync).not.toHaveBeenCalled(); + }); + + it("submits null when cleared (no cap)", async () => { + renderDialog(makeProject({ monthly_budget_usd: 42 })); + await screen.findByRole("button", { name: /Save Changes/i }); + openAutonomySection(); + + fireEvent.change(screen.getByLabelText(/Monthly Budget/i), { + target: { value: "" }, + }); + submit(); + + await waitFor(() => expect(mutateAsync).toHaveBeenCalled()); + const call = mutateAsync.mock.calls[0][0] as { + updates: { monthly_budget_usd?: number | null }; + }; + expect(call.updates.monthly_budget_usd).toBeNull(); + }); + + it("submits a positive cap as a number", async () => { + renderDialog(makeProject({ monthly_budget_usd: null })); + await screen.findByRole("button", { name: /Save Changes/i }); + openAutonomySection(); + + fireEvent.change(screen.getByLabelText(/Monthly Budget/i), { + target: { value: "100" }, + }); + submit(); + + await waitFor(() => expect(mutateAsync).toHaveBeenCalled()); + const call = mutateAsync.mock.calls[0][0] as { + updates: { monthly_budget_usd?: number | null }; + }; + expect(call.updates.monthly_budget_usd).toBe(100); + }); +}); diff --git a/panel/src/components/projects/edit-project-dialog.tsx b/panel/src/components/projects/edit-project-dialog.tsx index 0e34e52b..2d9c0750 100644 --- a/panel/src/components/projects/edit-project-dialog.tsx +++ b/panel/src/components/projects/edit-project-dialog.tsx @@ -237,6 +237,9 @@ function EditProjectForm({ const [depUpdatePaths, setDepUpdatePaths] = useState( (project.dep_update_paths || []).join(", "), ); + const [monthlyBudgetUsd, setMonthlyBudgetUsd] = useState( + project.monthly_budget_usd != null ? String(project.monthly_budget_usd) : "", + ); const sandboxServices = project.sandbox_services || []; const [sandboxSet, setSandboxSet] = useState>( new Set(sandboxServices), @@ -312,6 +315,15 @@ function EditProjectForm({ return; } + const trimmedBudget = monthlyBudgetUsd.trim(); + const parsedBudget = trimmedBudget ? Number(trimmedBudget) : null; + if (trimmedBudget && (Number.isNaN(parsedBudget) || parsedBudget! <= 0)) { + toast.error( + "Monthly budget must be greater than 0 — leave it empty for no cap", + ); + return; + } + // Build update payload const updates: ProjectUpdate = { name, @@ -342,6 +354,9 @@ function EditProjectForm({ .map((p) => p.trim()) .filter(Boolean) : undefined, + // Sent explicitly (never coerced to undefined) so clearing the input + // actually clears the stored cap instead of being dropped. + monthly_budget_usd: parsedBudget, sandbox_services: [...sandboxSet], sandbox_extensions: (() => { const extObj: Record = {}; @@ -800,6 +815,26 @@ function EditProjectForm({ {showAutonomy && ( <> +
+ + + + setMonthlyBudgetUsd(e.target.value)} + placeholder="No cap" + /> +

+ Claims are refused once this month's spend reaches the cap. + Must be greater than 0 — a 0 budget would block every claim + immediately. Leave blank for no cap. +

+
+
+ {/* Budget (USD) */} +
+ + + + setBudgetUsd(e.target.value)} + /> +

+ Must be greater than 0 — a 0 budget would block the task + before it spends a cent. Leave blank for the task-type + default. +

+
+ {/* Git Configuration Section */}
diff --git a/panel/src/lib/api/projects.ts b/panel/src/lib/api/projects.ts index 291ada52..580cca1b 100644 --- a/panel/src/lib/api/projects.ts +++ b/panel/src/lib/api/projects.ts @@ -174,6 +174,7 @@ export const projectsApi = { video_engine_enabled: false, dep_update_command: null, dep_update_paths: null, + monthly_budget_usd: null, sandbox_services: null, sandbox_extensions: null, workspace_path: null, diff --git a/panel/src/types/index.ts b/panel/src/types/index.ts index 2bb0eb57..0235f4d2 100644 --- a/panel/src/types/index.ts +++ b/panel/src/types/index.ts @@ -224,6 +224,8 @@ export interface Task { acceptance_criteria: string[]; status: TaskStatus; priority: number; // 0=P0(highest), 1=P1, 2=P2, 3=P3(lowest) + // Cost cap (ROBOCO_TASK_BUDGETS_ENABLED). null = use the task-type default. + budget_usd?: number | null; sequence: number; // Order number within siblings team: Team; created_by: string; @@ -1059,6 +1061,9 @@ export interface Project { video_engine_enabled: boolean; dep_update_command: string | null; dep_update_paths: string[] | null; + // Calendar-month cap on summed agent-spawn spend across this project's + // tasks; null = no cap. Only enforced when ROBOCO_TASK_BUDGETS_ENABLED is on. + monthly_budget_usd: number | null; sandbox_services: string[] | null; sandbox_extensions: Record | null; // Runtime state @@ -1123,6 +1128,8 @@ export interface ProjectUpdate { video_engine_enabled?: boolean; dep_update_command?: string; dep_update_paths?: string[]; + // null clears the cap (no cap); omitted leaves unchanged. + monthly_budget_usd?: number | null; sandbox_services?: string[]; sandbox_extensions?: Record; } diff --git a/roboco/api/routes/tasks.py b/roboco/api/routes/tasks.py index f56cd229..b3ad2aa9 100644 --- a/roboco/api/routes/tasks.py +++ b/roboco/api/routes/tasks.py @@ -213,7 +213,7 @@ _MIN_NOTES_CHARS = 20 # After TaskService.update() gains its not-None guard, null-clears for these # fields are handled at the route layer by direct setattr on the ORM object. _NULLABLE_TASK_FIELDS: frozenset[str] = frozenset( - {"assigned_to", "parent_task_id", "project_id"} + {"assigned_to", "parent_task_id", "project_id", "budget_usd"} ) # Structural / ownership fields a bare task owner (UPDATE_OWN) must NOT @@ -222,7 +222,8 @@ _NULLABLE_TASK_FIELDS: frozenset[str] = frozenset( # delegation plan. These are PM/ASSIGN-gated operations; the verb layer gates # them to PM roles (reassign/delegate/triage), so the REST PATCH surface must # not let an owner bypass that by setattr-ing them directly. Only a caller with -# the higher ASSIGN permission may set them. +# the higher ASSIGN permission may set them. budget_usd joins this set too — a +# self-serve budget raise on your own task would defeat the whole cap. _PRIVILEGED_UPDATE_FIELDS: frozenset[str] = frozenset( { "assigned_to", @@ -232,6 +233,7 @@ _PRIVILEGED_UPDATE_FIELDS: frozenset[str] = frozenset( "blocker_ids", "plan", "project_id", + "budget_usd", } ) diff --git a/roboco/api/schemas/project.py b/roboco/api/schemas/project.py index d99f4ad7..5ff3afea 100644 --- a/roboco/api/schemas/project.py +++ b/roboco/api/schemas/project.py @@ -57,6 +57,7 @@ class ProjectResponse(BaseModel): video_engine_enabled: bool = False dep_update_command: str | None = None dep_update_paths: list[str] | None = None + monthly_budget_usd: float | None = None sandbox_services: list[str] | None = None sandbox_extensions: dict[str, list[str]] | None = None @@ -211,6 +212,7 @@ class ProjectUpdateRequest(BaseModel): video_engine_enabled: bool | None = None dep_update_command: str | None = None dep_update_paths: list[str] | None = None + monthly_budget_usd: float | None = None sandbox_services: list[str] | None = None sandbox_extensions: dict[str, list[str]] | None = None @@ -303,6 +305,7 @@ def project_to_response(project: "ProjectTable") -> ProjectResponse: video_engine_enabled=bool(project.video_engine_enabled), dep_update_command=project.dep_update_command, dep_update_paths=project.dep_update_paths, + monthly_budget_usd=getattr(project, "monthly_budget_usd", None), sandbox_services=project.sandbox_services, sandbox_extensions=project.sandbox_extensions, workspace_path=project.workspace_path, diff --git a/roboco/api/schemas/tasks.py b/roboco/api/schemas/tasks.py index 52700883..1460281d 100644 --- a/roboco/api/schemas/tasks.py +++ b/roboco/api/schemas/tasks.py @@ -210,6 +210,10 @@ 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 + # other _NULLABLE_TASK_FIELDS (TaskService.update() itself skips None). + budget_usd: float | None = Field(default=None, ge=0) target_date: datetime | None = None estimated_complexity: Complexity | None = None @@ -306,6 +310,8 @@ class TaskResponse(BaseModel): status: TaskStatus priority: int sequence: int # Order number within siblings + # Cost cap (ROBOCO_TASK_BUDGETS_ENABLED). Null = use the TaskType default. + budget_usd: float | None = None nature: TaskNature # Technical or non-technical work # Task Type & Git Configuration (all tasks follow git workflow) @@ -881,6 +887,7 @@ def task_to_response(task: "TaskTable") -> TaskResponse: status=task.status, priority=task.priority, sequence=task.sequence, + budget_usd=getattr(task, "budget_usd", None), nature=task.nature, task_type=task.task_type, project_id=to_python_uuid(task.project_id), diff --git a/roboco/config.py b/roboco/config.py index c0ec518a..0bf18034 100644 --- a/roboco/config.py +++ b/roboco/config.py @@ -384,6 +384,17 @@ class Settings(BaseSettings): "gates. Off => the standard i_am_done path is unchanged." ), ) + task_budgets_enabled: bool = Field( + default=False, + description=( + "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 " + "cap is ever consulted, regardless of project/task field values." + ), + ) # ========================================================================== # Web Research (pluggable external search/fetch for Board + PM roles) diff --git a/roboco/db/tables.py b/roboco/db/tables.py index ce95ff37..fced701d 100644 --- a/roboco/db/tables.py +++ b/roboco/db/tables.py @@ -209,6 +209,10 @@ class TaskTable(Base): nullable=True, ) 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 + # no-op off. Enforced periodically by the orchestrator's budget sweep. + budget_usd: Mapped[float | None] = mapped_column(Float, nullable=True) # Task Type & Git Configuration (all tasks follow git workflow) task_type: Mapped[TaskType] = mapped_column( @@ -561,6 +565,11 @@ class ProjectTable(Base): ARRAY(String), nullable=True ) + # Cost budgets (migration 079, feature-flagged: ROBOCO_TASK_BUDGETS_ENABLED). + # Null = no cap, regardless of the flag. Enforced at claim time against + # this calendar month's summed agent-spawn spend across the project's tasks. + monthly_budget_usd: Mapped[float | None] = mapped_column(Float, nullable=True) + # Sandboxed per-agent-spawn DB/Redis opt-in. A project participates only # when sandbox_services is set (e.g. ["postgres", "redis"]); values are # validated by the Project pydantic model before reaching here. @@ -1908,6 +1917,9 @@ class AgentSpawnSessionTable(Base): Index("ix_agent_spawn_sessions_started_at", "started_at"), Index("ix_agent_spawn_sessions_ended_at", "ended_at"), Index("ix_agent_spawn_sessions_team", "team"), + # Migration 080: serves the per-project monthly-spend claim guard and + # the per-task budget sweep, both of which filter on bare task_id. + Index("ix_agent_spawn_sessions_task_id", "task_id"), ) diff --git a/roboco/foundation/policy/agent_loop.py b/roboco/foundation/policy/agent_loop.py index cd5c7c97..6b94b11a 100644 --- a/roboco/foundation/policy/agent_loop.py +++ b/roboco/foundation/policy/agent_loop.py @@ -24,7 +24,9 @@ cumulative cap alongside the window so pacing can't defeat the breaker. from __future__ import annotations from dataclasses import dataclass -from typing import Literal +from typing import Any, Literal + +from roboco.models.base import TaskType @dataclass(frozen=True) @@ -77,6 +79,46 @@ 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 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. + """ + 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) + + # Per-verb retry caps. Verbs that hit a tracing_gap or invalid_state and # get retried more than this many times in a 60s window will receive a # circuit_open envelope on the next attempt. diff --git a/roboco/foundation/policy/content/markers.py b/roboco/foundation/policy/content/markers.py index 1d40d01e..320f27b0 100644 --- a/roboco/foundation/policy/content/markers.py +++ b/roboco/foundation/policy/content/markers.py @@ -519,3 +519,26 @@ def mark_block_flip_notified(task: HasMarkers) -> None: set_marker( task, BLOCK_FLIP_COUNT, {"count": get_block_flip_count(task), "notified": True} ) + + +# --- budget-breach block ---------------------------------------------------- +# Stamped by the orchestrator's task-budget sweep the moment it BLOCKs a task +# for exceeding its $ budget (ROBOCO_TASK_BUDGETS_ENABLED). `unblock` consults +# it to re-check spend-vs-cap: still over refuses (naming the budget +# remediation) so a PM can't silently re-breach the same cap the next tick; +# under (the CEO raised it) clears the marker and lets the unblock through. +# No historical cap/spend stored here — the re-check always reads live values. + +BUDGET_BLOCKED = "budget_blocked" + + +def mark_budget_blocked(task: HasMarkers) -> None: + set_marker(task, BUDGET_BLOCKED, True) + + +def is_budget_blocked(task: HasMarkers) -> bool: + return bool(get_marker(task, BUDGET_BLOCKED, False)) + + +def clear_budget_blocked(task: HasMarkers) -> None: + clear_marker(task, BUDGET_BLOCKED) diff --git a/roboco/models/project.py b/roboco/models/project.py index bd8208af..808d6a7d 100644 --- a/roboco/models/project.py +++ b/roboco/models/project.py @@ -245,6 +245,18 @@ class Project(TimestampMixin): description="Lockfile globs to inspect (null → infer uv.lock/pnpm-lock.yaml)", ) + # Cost budgets (feature-flagged: ROBOCO_TASK_BUDGETS_ENABLED). Null = no + # cap, regardless of the flag — this is purely additive. + monthly_budget_usd: float | None = Field( + default=None, + ge=0, + description=( + "Calendar-month cap on this project's summed agent-spawn spend " + "(estimated_cost_usd). Null = no cap. Only enforced at claim time " + "when ROBOCO_TASK_BUDGETS_ENABLED is on." + ), + ) + # Sandboxed per-agent-spawn DB/Redis opt-in sandbox_services: list[str] | None = Field( default=None, @@ -319,6 +331,7 @@ class ProjectCreate(RobocoBase): build_command: str | None = None quality_command: str | None = None codegen_command: str | None = None + monthly_budget_usd: float | None = None class ProjectUpdate(RobocoBase): @@ -358,6 +371,7 @@ class ProjectUpdate(RobocoBase): video_engine_enabled: bool | None = None dep_update_command: str | None = None dep_update_paths: list[str] | None = None + monthly_budget_usd: float | None = None sandbox_services: list[str] | None = None sandbox_extensions: dict[str, list[str]] | None = None github_installation_id: int | None = Field( diff --git a/roboco/models/task.py b/roboco/models/task.py index daf846aa..2722d1ed 100644 --- a/roboco/models/task.py +++ b/roboco/models/task.py @@ -166,6 +166,17 @@ 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. + budget_usd: float | None = Field( + default=None, + description=( + "Cap on this task's own accumulated agent-spawn spend " + "(estimated_cost_usd). Null = use the TaskType default." + ), + ) + # Task Type & Git Configuration (all tasks follow git workflow) task_type: TaskType = Field( default=TaskType.CODE, description="Type of task (code, research, etc.)" @@ -432,6 +443,7 @@ class TaskUpdate(RobocoBase): description: str | None = None acceptance_criteria: list[str] | None = None priority: int | None = Field(default=None, ge=0, le=3) + budget_usd: float | None = Field(default=None, ge=0) status: TaskStatus | None = None assigned_to: UUID | None = None target_date: datetime | None = None diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index 752e141e..97a96597 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -7624,17 +7624,127 @@ Start by: return None return data if isinstance(data, dict) else None - async def _sweep_budget_exceeded(self) -> None: - """Stop agents whose per-session SDK budget reports halt=true. + 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. - Each agent's SDK server is reachable at + ``None`` when unbreached 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 + ``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 + fresh SDK round-trip needed for "cheaply available"). + """ + from roboco.db.base import get_db_context + from roboco.foundation.policy.agent_loop import effective_task_budget_usd + from roboco.models.base import TaskStatus + from roboco.services.task import TaskService + from roboco.utils.converters import InvalidIdentifierError, require_uuid + + try: + task_id = require_uuid(task_id_str) + except InvalidIdentifierError: + return None + async with get_db_context() as db: + svc = TaskService(db) + task = await svc.get(task_id) + if task is None or task.status not in ( + TaskStatus.CLAIMED, + TaskStatus.IN_PROGRESS, + ): + return None + cap_usd = effective_task_budget_usd(task) + spend_usd = await svc.task_spend_usd(task_id) + if spend_usd < cap_usd: + return None + return cap_usd, spend_usd + + async def _handle_task_budget_breach( + self, task_id_str: str, *, cap_usd: float, spend_usd: float + ) -> None: + """Block a task whose own $ budget is breached and notify the CEO. + + Runs BEFORE ``stop_agent`` so its ``release_claim`` unclaim (which + only fires from claimed/in_progress — see ``_force_unclaim_to_pending``) + finds the task already ``BLOCKED`` and no-ops: the task never bounces + through ``pending`` for an instant re-claim to re-burn the same + budget. ``blocker_resolver_type=HUMAN`` keeps the dispatcher from ever + respawning onto it (``_is_hitl_blocked``) — only the CEO raising the + cap (or cancelling) moves it forward. Best-effort: a failure here logs + and lets the caller's ``stop_agent`` proceed regardless — one more + tick of a live over-budget agent is the safer failure mode than + crashing the sweep. + """ + from roboco.db.base import get_db_context + from roboco.foundation.policy.content import markers + from roboco.models.base import BlockerResolverType, TaskStatus + from roboco.services.notification_delivery import ( + get_notification_delivery_service, + ) + from roboco.services.task import TaskService + from roboco.utils.converters import InvalidIdentifierError, require_uuid + + try: + task_id = require_uuid(task_id_str) + except InvalidIdentifierError as exc: + logger.warning( + "task budget breach had malformed task id", + task_id_str=task_id_str, + error=str(exc), + ) + return + try: + async with get_db_context() as db: + svc = TaskService(db) + task = await svc.get(task_id) + if task is None or task.status not in ( + TaskStatus.CLAIMED, + TaskStatus.IN_PROGRESS, + ): + return + task.blocker_resolver_type = BlockerResolverType.HUMAN + markers.mark_budget_blocked(task) + await svc.admin_set_status( + task_id, TaskStatus.BLOCKED, actor_role="system" + ) + delivery = get_notification_delivery_service(db) + await delivery.notify_ceo_of_budget_breach( + task=task, + task_id=task_id, + cap_usd=cap_usd, + spend_usd=spend_usd, + ) + except Exception as exc: + logger.warning( + "task budget-breach block/notify failed", + task_id=task_id_str, + error=str(exc), + ) + + async def _sweep_budget_exceeded(self) -> None: + """Stop agents whose per-session SDK budget reports halt=true, OR + (when ``ROBOCO_TASK_BUDGETS_ENABLED`` is on) whose active task's own + $ budget is breached. + + Tool-call halt: each agent's SDK server is reachable at `http://roboco-agent-{agent_id}:9000/budget/status` on the shared - agent network. A budget-exceeded agent gets a forced stop with a - `budget_exceeded` reason; the task is already being auto-substituted - by the post-tool hook on the agent side. + agent network; the task is already being auto-substituted by the + post-tool hook on the agent side, so a release_claim stop suffices. + + Task $ budget: `_task_budget_breach` compares accumulated spend + (agent_spawn_sessions) against task.budget_usd / the TaskType + default. Unlike the tool-call path this explicitly BLOCKs the task + (see `_handle_task_budget_breach`) before the agent is stopped, and + notifies the CEO — a breached $ cap needs a human decision (raise + the cap or leave it blocked), not just a silent re-queue. """ if not self._instances: return + from roboco.config import settings as _settings + + task_budgets_on = _settings.task_budgets_enabled async with httpx.AsyncClient( timeout=3.0, headers=_system_api_headers() ) as client: @@ -7644,33 +7754,102 @@ Start by: AgentState.WAITING_SHORT, ): continue - url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/budget/status" - data = await self._fetch_budget_status(client, url, agent_id) - if data is None or not data.get("halt"): - continue - logger.warning( - "Agent budget exceeded; terminating container", - agent_id=agent_id, - total_calls=data.get("total"), - halt_threshold=data.get("halt_threshold"), + await self._check_budget_for_agent( + client, agent_id, instance, task_budgets_on ) - try: - # release_claim=True: a budget-exceeded agent is terminated - # for cost overruns and will not continue its task, so hand - # the claim back to the pool now instead of waiting for the - # reaper's TTL. - await self.stop_agent( - agent_id, - graceful=True, - release_claim=True, - stop_reason="budget_sweep", - ) - except Exception as e: - logger.warning( - "Failed to stop budget-exceeded agent", - agent_id=agent_id, - error=str(e), - ) + + async def _check_budget_for_agent( + self, + client: httpx.AsyncClient, + agent_id: str, + instance: Any, + task_budgets_on: bool, + ) -> None: + """Stop one agent if its tool-call budget halted OR its task's $ + budget breached. Extracted from ``_sweep_budget_exceeded`` (xenon + complexity budget) — see that method's docstring for the two triggers. + """ + url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/budget/status" + data = await self._fetch_budget_status(client, url, agent_id) + task_breach = await self._maybe_task_budget_breach(instance, task_budgets_on) + tool_call_halt = bool(data and data.get("halt")) + if not tool_call_halt and task_breach is None: + return + await self._stop_budget_exceeded_agent(agent_id, instance, data, task_breach) + + async def _maybe_task_budget_breach( + self, instance: Any, task_budgets_on: bool + ) -> tuple[float, float] | None: + """``_task_budget_breach`` for the instance's task, or None inert.""" + if not task_budgets_on or not instance.current_task_id: + return None + return await self._task_budget_breach(instance.current_task_id) + + async def _resolve_budget_stop_reason( + self, + agent_id: str, + instance: Any, + data: dict[str, Any] | None, + task_breach: tuple[float, float] | None, + ) -> str: + """Block + notify on a task breach, else log the tool-call halt. + + Either way returns the ``stop_reason`` the caller passes to + ``stop_agent``. + """ + if task_breach is None or not instance.current_task_id: + logger.warning( + "Agent budget exceeded; terminating container", + agent_id=agent_id, + total_calls=data.get("total") if data else None, + halt_threshold=data.get("halt_threshold") if data else None, + ) + return "budget_sweep" + cap_usd, spend_usd = task_breach + logger.warning( + "Task budget exceeded; blocking task and terminating container", + agent_id=agent_id, + task_id=instance.current_task_id, + cap_usd=cap_usd, + spend_usd=round(spend_usd, 4), + ) + await self._handle_task_budget_breach( + instance.current_task_id, cap_usd=cap_usd, spend_usd=spend_usd + ) + return "budget_exceeded_task" + + async def _stop_budget_exceeded_agent( + self, + agent_id: str, + instance: Any, + data: dict[str, Any] | None, + task_breach: tuple[float, float] | None, + ) -> None: + """Resolve the stop reason (blocking + notifying on a task breach + first) then gracefully stop the agent, releasing its claim. + + release_claim=True mirrors the tool-call-halt path. On a task breach + the task is already BLOCKED (`_resolve_budget_stop_reason` ran the + block+notify first), so stop_agent's own release-to-pending unclaim + finds it out of claimed/in_progress and no-ops — it never bounces + through pending for an instant re-claim to re-burn the same budget. + """ + stop_reason = await self._resolve_budget_stop_reason( + agent_id, instance, data, task_breach + ) + try: + await self.stop_agent( + agent_id, + graceful=True, + release_claim=True, + stop_reason=stop_reason, + ) + except Exception as e: + logger.warning( + "Failed to stop budget-exceeded agent", + agent_id=agent_id, + error=str(e), + ) @staticmethod async def _inspect_container_state( diff --git a/roboco/services/gateway/choreographer/_impl.py b/roboco/services/gateway/choreographer/_impl.py index 49924e97..0f00d988 100644 --- a/roboco/services/gateway/choreographer/_impl.py +++ b/roboco/services/gateway/choreographer/_impl.py @@ -36,6 +36,7 @@ from roboco.services.gateway.choreographer.collision import build_collision_cont from roboco.services.gateway.claim_guards import ( already_active_guard, paused_tasks_guard, + project_budget_exceeded_guard, unmet_dependency_guard, ) from roboco.services.gateway.envelope import Envelope @@ -1148,6 +1149,7 @@ class Choreographer: task: Any, role_str: str | None = None, skip_dev_guards: bool = False, + check_project_budget: bool = False, ) -> Envelope | None: """Run concurrency-invariant claim guards. Returns rejection or None. @@ -1170,6 +1172,17 @@ class Choreographer: developer starting a code task do not apply (the dependency guard still runs). See ``claim_gate_review`` (#192). + ``check_project_budget`` (default False — explicit opt-in, not + inferred from role/verb) scopes the project monthly-budget guard to + genuinely work-STARTING claims: ``i_will_work_on`` / ``i_will_plan`` + pass True. A project at cap still has non-negotiable in-flight work + to finish — QA's ``claim_review``, the PR gate's + ``claim_gate_review``, ``claim_doc_task``, and inbound + ``claim_pr_review`` all pass False (the default) so reviewing/ + documenting/merging what's already been paid for never wedges behind + an exhausted cap whose incremental cost is negligible next to the + sunk spend. + Pre-gateway location: _helpers.py:124-204. """ if not skip_dev_guards and role_str not in self._COORDINATOR_ROLES: @@ -1179,39 +1192,107 @@ class Choreographer: paused = await self.task.list_paused_for_agent(agent_id) if guard := paused_tasks_guard(paused, task.id): return guard - dep_ids = list(task.dependency_ids) - if dep_ids: - unmet = await self.task.unmet_dependency_ids(dep_ids) - if guard := unmet_dependency_guard(task, unmet): - # Re-check before mutating: the read above is an unlocked SELECT, - # and an upstream dependency may have reached a terminal state - # (completed/cancelled) in the microseconds between that read and - # now. Dependencies are monotonic — unmet -> met only, terminal - # states never reopen — so a fresh read that now finds them met - # stays met, and the task can proceed. Releasing it anyway would - # needlessly clear its branch + abandon its WorkSession and - # bounce the assignee, only for the dependency-completion - # re-dispatch to re-dispatch + re-claim it a moment later. Skip - # the release and let the caller proceed (return None). The - # cross-task residual window (upstream completes between this - # re-check and the release below) is not closable by a row lock - # on the dependent — but the re-check narrows the window from - # [first read -> release] to [re-check -> release] and, in the - # common case, the first read already sees met (no guard). - fresh_unmet = await self.task.unmet_dependency_ids(dep_ids) - if not fresh_unmet: - return None - # Still unmet — park the dependency-gated task back to pending so - # the orchestrator stops respawning its assignee (the respawn - # loop targets only claimed/in_progress) and the dispatch - # dependency filter holds it until the upstream completes. No-op - # unless the task is currently claimed/in_progress. - await self.task.release_dependency_blocked_claim(task.id) - return guard + if guard := await self._dependency_claim_guard(task): + return guard + if check_project_budget and ( + guard := await self._project_budget_claim_guard(task) + ): + return guard if skip_dev_guards: return None return await self._lane_claim_guard(task) + async def _dependency_claim_guard(self, task: Any) -> Envelope | None: + """Refuse claim while the task has non-terminal dependencies. + + Extracted from ``_run_claim_guards`` (xenon return-count budget). + """ + dep_ids = list(task.dependency_ids) + if not dep_ids: + return None + unmet = await self.task.unmet_dependency_ids(dep_ids) + guard = unmet_dependency_guard(task, unmet) + if guard is None: + return None + # Re-check before mutating: the read above is an unlocked SELECT, and + # an upstream dependency may have reached a terminal state + # (completed/cancelled) in the microseconds between that read and now. + # Dependencies are monotonic — unmet -> met only, terminal states never + # reopen — so a fresh read that now finds them met stays met, and the + # task can proceed. Releasing it anyway would needlessly clear its + # branch + abandon its WorkSession and bounce the assignee, only for + # the dependency-completion re-dispatch to re-dispatch + re-claim it a + # moment later. Skip the release and let the caller proceed (return + # None). The cross-task residual window (upstream completes between + # this re-check and the release below) is not closable by a row lock + # on the dependent — but the re-check narrows the window from [first + # read -> release] to [re-check -> release] and, in the common case, + # the first read already sees met (no guard). + fresh_unmet = await self.task.unmet_dependency_ids(dep_ids) + if not fresh_unmet: + return None + # Still unmet — park the dependency-gated task back to pending so the + # orchestrator stops respawning its assignee (the respawn loop targets + # only claimed/in_progress) and the dispatch dependency filter holds it + # until the upstream completes. No-op unless the task is currently + # claimed/in_progress. + await self.task.release_dependency_blocked_claim(task.id) + return guard + + async def _project_budget_claim_guard(self, task: Any) -> Envelope | None: + """Refuse claim once the task's project has spent its monthly cap. + + Inert unless ``ROBOCO_TASK_BUDGETS_ENABLED`` is on AND the task's + project has ``monthly_budget_usd`` set — a task with no project (a + branchless coordination root) or a project with no cap never even + reaches the spend query. + """ + from roboco.config import settings as _settings + + if not _settings.task_budgets_enabled: + return None + project = getattr(task, "project", None) + if project is None: + return None + monthly_budget_usd = getattr(project, "monthly_budget_usd", None) + if monthly_budget_usd is None: + return None + month_spend_usd = await self.task.project_month_spend_usd(project.id) + return project_budget_exceeded_guard(task, monthly_budget_usd, month_spend_usd) + + async def _budget_unblock_guard(self, t: Any) -> Envelope | None: + """Refuse ``unblock`` on a budget-blocked task while still over cap. + + Inert unless ``ROBOCO_TASK_BUDGETS_ENABLED`` is on AND the task + carries the ``BUDGET_BLOCKED`` marker (stamped by the orchestrator's + task-budget sweep at breach time — see + ``AgentOrchestrator._handle_task_budget_breach``); a task blocked for + any other reason (dependency, manual escalation) never reaches the + spend query. Clears the marker on success so a later, unrelated block + on the same task is never mistaken for a stale budget breach. + """ + from roboco.config import settings as _settings + from roboco.foundation.policy.agent_loop import effective_task_budget_usd + + if not _settings.task_budgets_enabled or not markers.is_budget_blocked(t): + return None + spend_usd = await self.task.task_spend_usd(t.id) + cap_usd = effective_task_budget_usd(t) + if spend_usd >= cap_usd: + return Envelope.invalid_state( + message=( + f"task {t.id} is still over its cost budget: " + f"${cap_usd:,.2f} cap, ${spend_usd:,.2f} spent." + ), + remediate=( + "raise the task's Budget (USD) field (or the project's " + "Monthly Budget if that's the actual cap) before calling " + "unblock again" + ), + ) + markers.clear_budget_blocked(t) + return None + async def _lane_claim_guard(self, task: Any) -> Envelope | None: """Refuse a code leaf behind an earlier open same-assignee sibling. @@ -1361,11 +1442,13 @@ class Choreographer: verb=verb_name, ) # Concurrency guards still apply on resumption (paused / already-active - # in another task). + # in another task). check_project_budget=True: resuming i_will_work_on + # / i_will_plan is still a work-STARTING claim. if guard := await self._run_claim_guards( agent_id=agent_id, task=t, role_str=role_str, + check_project_budget=True, ): return await self._emit_rejection( self._with_briefing(guard, briefing).with_introspection( @@ -1477,6 +1560,7 @@ class Choreographer: agent_id=ctx.agent_id, task=t, role_str=role_str, + check_project_budget=True, ): return await self._emit_rejection( self._with_briefing(guard, briefing).with_introspection( @@ -6871,6 +6955,19 @@ class Choreographer: verb="unblock", ) + # Budget-breach block: re-check spend-vs-cap before letting the PM + # clear it. Without this, a PM unblock on a task the orchestrator's + # sweep blocked for a $ overrun would silently re-breach the same cap + # the very next tick if the CEO's raise didn't actually clear it (or + # never raised it at all). + if guard := await self._budget_unblock_guard(t): + return await self._emit_rejection( + guard.with_introspection(task=t, role=role), + agent_id=pm_agent_id, + task_id=task_id, + verb="unblock", + ) + # Write-then-gate: the PM's unblock reason is recorded as the # journal:decision the gate below requires, so a PM that didn't # pre-call note(scope='decision') doesn't stall in a tracing_gap diff --git a/roboco/services/gateway/choreographer/_protocol.py b/roboco/services/gateway/choreographer/_protocol.py index 38812e87..8364c097 100644 --- a/roboco/services/gateway/choreographer/_protocol.py +++ b/roboco/services/gateway/choreographer/_protocol.py @@ -122,6 +122,7 @@ class ChoreographerHelpers: task: Any, role_str: str | None = None, skip_dev_guards: bool = False, + check_project_budget: bool = False, ) -> Envelope | None: raise NotImplementedError diff --git a/roboco/services/gateway/claim_guards.py b/roboco/services/gateway/claim_guards.py index 7712fd79..5bf972fd 100644 --- a/roboco/services/gateway/claim_guards.py +++ b/roboco/services/gateway/claim_guards.py @@ -83,6 +83,32 @@ def paused_tasks_guard( ) +def project_budget_exceeded_guard( + target_task: Any, monthly_budget_usd: float | None, month_spend_usd: float +) -> Envelope | None: + """Refuse claim once the task's project has spent its monthly cap. + + Only fires when ``monthly_budget_usd`` is set (``None`` = no cap, the + guard is inert) — the caller resolves both the project's cap and this + calendar month's summed agent-spawn spend (a DB read) so this predicate + stays pure, mirroring ``unmet_dependency_guard``. At-or-over the cap + refuses; strictly under it passes. + """ + if monthly_budget_usd is None or month_spend_usd < monthly_budget_usd: + return None + return Envelope.invalid_state( + message=( + f"task {target_task.id}'s project has reached its monthly budget " + f"(${monthly_budget_usd:,.2f} cap, ${month_spend_usd:,.2f} spent " + "this calendar month)." + ), + remediate=( + "wait for next calendar month, or raise the project's Monthly " + "Budget (USD) field in project settings" + ), + ) + + def unmet_dependency_guard( target_task: Any, unmet_dependency_ids: list[UUID] ) -> Envelope | None: diff --git a/roboco/services/notification_delivery.py b/roboco/services/notification_delivery.py index d019bf19..77a30bf6 100644 --- a/roboco/services/notification_delivery.py +++ b/roboco/services/notification_delivery.py @@ -1164,6 +1164,58 @@ class NotificationDeliveryService(BaseService): task_id=task_id, subject=notification.subject, actionable=True ) + async def notify_ceo_of_budget_breach( + self, + *, + task: TaskTable, + task_id: UUID, + cap_usd: float, + spend_usd: float, + ) -> None: + """Create + deliver the CEO budget-breach notification. + + Shaped exactly like ``notify_ceo_of_escalation`` (APPROVAL/HIGH, + actionable Telegram keyboard) — the CEO's options are the same shape + (raise the cap, or leave it blocked). There is no human escalator + here (the orchestrator's budget sweep triggers this), so ``from_agent`` + falls back to the task's own assignee, mirroring + ``notify_ceo_of_completion``'s fallback. + + Names BOTH remediation steps: raising ``budget_usd`` alone does not + resume the task — a PM must still call ``unblock`` (or the panel + equivalent). ``unblock`` itself re-checks spend-vs-cap + (``markers.BUDGET_BLOCKED``) and refuses while still over, so a raise + that didn't actually clear the cap is caught there, not silently + re-breached. + """ + ceo = await self._get_ceo_agent() + if not ceo: + return + from_agent = cast("UUID", task.assigned_to) if task.assigned_to else ceo.id + notification = NotificationTable( + type=NotificationType.APPROVAL, + priority=NotificationPriority.HIGH, + from_agent=from_agent, + to_agents=[ceo.id], + subject=f"Budget exceeded: {task.title or 'Unknown task'}", + body=( + f"Task {task_display(task, task_id)} was stopped and blocked: " + f"its cost budget (${cap_usd:,.2f}) is exceeded — " + f"${spend_usd:,.2f} spent so far.\n\n" + "Two steps to resume it: 1) raise the task's Budget (USD) " + "field (task detail, or the project's Monthly Budget if " + "that's the cap), 2) have its PM call unblock — it stays " + "blocked until unblock runs, and refuses again if the cap " + "still isn't cleared. Or leave it blocked / cancel it." + ), + related_task_id=task_id, + requires_ack=ACK_REQUIRED_BY_TYPE[NotificationType.APPROVAL], + ) + await self._persist_and_deliver(notification) + await self._notify_telegram( + task_id=task_id, subject=notification.subject, actionable=True + ) + async def notify_ceo_of_completion(self, *, task: TaskTable, task_id: UUID) -> None: """CEO-facing completion notification with the granular effort breakdown. diff --git a/roboco/services/settings.py b/roboco/services/settings.py index 04962eec..f0af8c4f 100644 --- a/roboco/services/settings.py +++ b/roboco/services/settings.py @@ -70,6 +70,7 @@ FEATURE_FLAGS: tuple[tuple[str, str], ...] = ( "possibilities_matrix_enabled", "Possibilities matrix (work-already-done fast path)", ), + ("task_budgets_enabled", "Task/project cost budgets"), ("rag_auto_update_enabled", "RAG auto-update"), ("transcript_prune_enabled", "Transcript pruning"), ("gateway_health_enabled", "Gateway-health recovery"), diff --git a/roboco/services/task.py b/roboco/services/task.py index 7b32cd32..7fddd3fa 100644 --- a/roboco/services/task.py +++ b/roboco/services/task.py @@ -19,6 +19,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import InstanceState from roboco.db.tables import ( + AgentSpawnSessionTable, AgentTable, JournalEntryTable, JournalTable, @@ -8196,6 +8197,104 @@ class TaskService(BaseService): if dep_status not in terminal ] + async def project_month_spend_usd(self, project_id: UUID) -> float: + """This calendar month's summed agent-spawn cost for a project's tasks. + + Backs the claim-time monthly-budget guard (``ROBOCO_TASK_BUDGETS_ENABLED``). + ``agent_spawn_sessions.task_id`` is a plain ``String(36)``, not a real + FK (see ``AgentSpawnSessionTable``), so the join casts ``tasks.id`` to + text. A still-open session's ``estimated_cost_usd`` stays null until + close (``AgentOrchestrator._finalize_spawn_session``) — summing that + column alone undercounts N parallel long-running sessions as $0 while + they're open, which is exactly the case a claim-time budget guard must + not miss. So this fetches per-row token/cost columns (one query, no + server-side SUM) and prices any still-open row from its periodically- + refreshed token counts via the same ``calculate_cost`` the + orchestrator's task-budget sweep uses (``_task_spend_usd``). + + Month attribution is deliberately ``started_at``-bucketed: a session + that starts in one calendar month and closes in the next counts + entirely against its START month, never split pro-rata across the + boundary. Accepted simplification — a claim-time guard only needs to + be roughly current, not a settled ledger; the tiny sliver of a + month-spanning session's tail is a rounding error against a monthly + cap, not a correctness bug. + """ + from roboco.billing.pricing import calculate_cost + + month_start = datetime.now(UTC).replace( + day=1, hour=0, minute=0, second=0, microsecond=0 + ) + task_id_str = cast("Any", TaskTable.id).cast(String) + result = await self.session.execute( + select( + AgentSpawnSessionTable.estimated_cost_usd, + AgentSpawnSessionTable.ended_at, + AgentSpawnSessionTable.model, + AgentSpawnSessionTable.tokens_input, + AgentSpawnSessionTable.tokens_output, + AgentSpawnSessionTable.tokens_cache_read, + AgentSpawnSessionTable.tokens_cache_write, + ) + .select_from(AgentSpawnSessionTable) + .join(TaskTable, task_id_str == AgentSpawnSessionTable.task_id) + .where( + TaskTable.project_id == project_id, + AgentSpawnSessionTable.started_at >= month_start, + ) + ) + total = 0.0 + for row in result.all(): + if row.estimated_cost_usd is not None: + total += row.estimated_cost_usd + elif row.ended_at is None: + total += calculate_cost( + model=row.model, + tokens_input=row.tokens_input, + tokens_output=row.tokens_output, + tokens_cache_read=row.tokens_cache_read, + tokens_cache_write=row.tokens_cache_write, + ) + return total + + async def task_spend_usd(self, task_id: UUID) -> float: + """Total agent-spawn spend for one task: closed sessions' + ``estimated_cost_usd`` + the live cost of any still-open session. + + Scoped to one task instead of a project+month — otherwise identical + to ``project_month_spend_usd``'s open-session handling: a still-open + session's cost stays null until close, so it's priced from its + periodically-refreshed token counts via ``calculate_cost``. Backs the + orchestrator's per-task budget sweep and ``unblock``'s budget-breach + re-check (both need "is this task still over its own cap"). + """ + from roboco.billing.pricing import calculate_cost + + rows = ( + ( + await self.session.execute( + select(AgentSpawnSessionTable).where( + AgentSpawnSessionTable.task_id == str(task_id) + ) + ) + ) + .scalars() + .all() + ) + total = 0.0 + for row in rows: + if row.estimated_cost_usd is not None: + total += row.estimated_cost_usd + elif row.ended_at is None: + total += calculate_cost( + model=row.model, + tokens_input=row.tokens_input, + tokens_output=row.tokens_output, + tokens_cache_read=row.tokens_cache_read, + tokens_cache_write=row.tokens_cache_write, + ) + return total + async def inherit_unmet_dependencies( self, subtask_id: UUID, parent_id: UUID ) -> None: diff --git a/tests/unit/config/test_task_budgets_flag.py b/tests/unit/config/test_task_budgets_flag.py new file mode 100644 index 00000000..637fe1f6 --- /dev/null +++ b/tests/unit/config/test_task_budgets_flag.py @@ -0,0 +1,27 @@ +"""Per-task and per-project cost budgets are gated by a default-off config +flag, registered as a panel-tunable feature flag (mirrors possibilities_matrix).""" + +from __future__ import annotations + +import os +from unittest import mock + +from roboco.config import Settings +from roboco.services.settings import FEATURE_FLAGS, validate_setting + + +def test_task_budgets_disabled_by_default() -> None: + assert Settings().task_budgets_enabled is False + + +def test_task_budgets_reads_env_var() -> None: + with mock.patch.dict(os.environ, {"ROBOCO_TASK_BUDGETS_ENABLED": "true"}): + assert Settings().task_budgets_enabled is True + + +def test_task_budgets_flag_registered_in_feature_flags() -> None: + assert "task_budgets_enabled" in [key for key, _ in FEATURE_FLAGS] + + +def test_task_budgets_flag_validates_as_bool() -> None: + validate_setting("task_budgets_enabled", "true") diff --git a/tests/unit/gateway/test_budget_claim_guard.py b/tests/unit/gateway/test_budget_claim_guard.py new file mode 100644 index 00000000..2c3901f5 --- /dev/null +++ b/tests/unit/gateway/test_budget_claim_guard.py @@ -0,0 +1,112 @@ +"""Per-project monthly-budget claim guard (ROBOCO_TASK_BUDGETS_ENABLED). + +Two layers: the pure predicate ``project_budget_exceeded_guard`` (over/under/ +null cap) and the Choreographer's ``_project_budget_claim_guard`` wiring +(flag-off inert, no-cap inert, spend query only reached when both are set). +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from roboco.config import settings +from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps +from roboco.services.gateway.claim_guards import project_budget_exceeded_guard + +# --------------------------------------------------------------------------- +# Pure predicate: project_budget_exceeded_guard +# --------------------------------------------------------------------------- + + +def test_predicate_null_cap_is_inert() -> None: + task = MagicMock(id=uuid4()) + assert project_budget_exceeded_guard(task, None, 999.0) is None + + +def test_predicate_under_cap_passes() -> None: + task = MagicMock(id=uuid4()) + assert project_budget_exceeded_guard(task, 10.0, 5.0) is None + + +def test_predicate_over_cap_refuses() -> None: + task = MagicMock(id=uuid4()) + env = project_budget_exceeded_guard(task, 10.0, 15.0) + assert env is not None + body = env.as_dict() + assert body["error"] == "invalid_state" + assert "10.00" in body["message"] + assert "15.00" in body["message"] + assert "project settings" in body["remediate"] + + +def test_predicate_at_cap_refuses() -> None: + """At-or-over refuses — the cap itself is not headroom.""" + task = MagicMock(id=uuid4()) + assert project_budget_exceeded_guard(task, 10.0, 10.0) is not None + + +# --------------------------------------------------------------------------- +# Choreographer wiring: _project_budget_claim_guard +# --------------------------------------------------------------------------- + + +def _make_choreographer(*, project_month_spend_usd: float = 0.0) -> Choreographer: + task_svc = AsyncMock() + task_svc.project_month_spend_usd.return_value = project_month_spend_usd + base = { + "task": task_svc, + "work_session": AsyncMock(), + "git": AsyncMock(), + "a2a": AsyncMock(), + "journal": AsyncMock(), + "audit": AsyncMock(), + "evidence_repo": AsyncMock(), + } + return Choreographer(ChoreographerDeps(**base)) + + +def _task_with_project(monthly_budget_usd: float | None) -> MagicMock: + project = MagicMock(id=uuid4(), monthly_budget_usd=monthly_budget_usd) + return MagicMock(id=uuid4(), project=project) + + +@pytest.mark.asyncio +async def test_flag_off_is_inert_even_over_cap(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "task_budgets_enabled", False) + c = _make_choreographer(project_month_spend_usd=999.0) + task = _task_with_project(monthly_budget_usd=10.0) + assert await c._project_budget_claim_guard(task) is None + c.task.project_month_spend_usd.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_no_project_cap_is_inert(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "task_budgets_enabled", True) + c = _make_choreographer(project_month_spend_usd=999.0) + task = _task_with_project(monthly_budget_usd=None) + assert await c._project_budget_claim_guard(task) is None + c.task.project_month_spend_usd.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_under_cap_passes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "task_budgets_enabled", True) + c = _make_choreographer(project_month_spend_usd=4.0) + task = _task_with_project(monthly_budget_usd=10.0) + assert await c._project_budget_claim_guard(task) is None + + +@pytest.mark.asyncio +async def test_over_cap_refuses_naming_the_field( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(settings, "task_budgets_enabled", True) + c = _make_choreographer(project_month_spend_usd=12.0) + task = _task_with_project(monthly_budget_usd=10.0) + env = await c._project_budget_claim_guard(task) + assert env is not None + body = env.as_dict() + assert body["error"] == "invalid_state" + assert "Monthly Budget" in body["remediate"] diff --git a/tests/unit/gateway/test_budget_guard_scoped_to_work_claims.py b/tests/unit/gateway/test_budget_guard_scoped_to_work_claims.py new file mode 100644 index 00000000..d74b111b --- /dev/null +++ b/tests/unit/gateway/test_budget_guard_scoped_to_work_claims.py @@ -0,0 +1,335 @@ +"""The project monthly-budget claim guard is scoped to work-STARTING claims. + +A project at (or over) its monthly cap still has non-negotiable in-flight +work to finish: QA's ``claim_review``, the PR gate's ``claim_gate_review``, +``claim_doc_task``, and inbound ``claim_pr_review`` never even reach the +spend query (``check_project_budget`` defaults False at those four call +sites) — reviewing/documenting/merging what's already been paid for must +never wedge behind an exhausted cap. Only ``i_will_work_on`` / ``i_will_plan`` +(``check_project_budget=True``) — genuinely starting new spend — refuse. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from roboco.config import settings +from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps + +_STEPS = [ + { + "title": "Implement the change", + "description": ( + "edit the target file, add tests, run them, and stage the " + "change for commit on the task branch" + ), + } +] + + +@pytest.fixture(autouse=True) +def _budgets_on(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "task_budgets_enabled", True) + + +def _over_cap_project() -> MagicMock: + """A project with a $10 cap — the spend stub below always reports $999.""" + return MagicMock(id=uuid4(), monthly_budget_usd=10.0) + + +def _make_deps(task_svc: AsyncMock, **overrides: Any) -> ChoreographerDeps: + base: dict[str, Any] = { + "task": task_svc, + "work_session": AsyncMock(), + "git": AsyncMock(), + "a2a": AsyncMock(), + "journal": AsyncMock(), + "audit": AsyncMock(), + "evidence_repo": AsyncMock(), + } + base.update(overrides) + return ChoreographerDeps(**base) + + +# --------------------------------------------------------------------------- +# Review-ish claims: guard skipped — each succeeds despite an over-cap project. +# project_month_spend_usd is asserted NOT awaited: the guard is skipped +# entirely, not merely lucky (e.g. a cache hit). +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_claim_review_succeeds_at_cap() -> None: + task_svc = AsyncMock() + t = MagicMock( + id=uuid4(), + status="awaiting_qa", + assigned_to=uuid4(), + parent_task_id=uuid4(), + task_type="code", + dependency_ids=[], + team="backend", + pr_number=10, + pr_url="https://example/pr/10", + branch_name="feature/backend/abc", + batch_id=None, + project=_over_cap_project(), + ) + task_svc.get.return_value = t + task_svc.agent_for.return_value = MagicMock(role="qa", slug="be-qa") + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.unmet_dependency_ids = AsyncMock(return_value=[]) + task_svc.has_earlier_incomplete_code_sibling.return_value = False + task_svc.qa_claim = AsyncMock(return_value=t) + task_svc.project_month_spend_usd = AsyncMock(return_value=999.0) + c = Choreographer(_make_deps(task_svc)) + cc: Any = c + cc._build_qa_review_evidence = AsyncMock(return_value={}) + + env = await c.claim_review(uuid4(), t.id) + body = env.as_dict() + assert body.get("error") is None, body + task_svc.project_month_spend_usd.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_claim_gate_review_succeeds_at_cap() -> None: + task_svc = AsyncMock() + t = MagicMock( + id=uuid4(), + status="awaiting_pr_review", + assigned_to=uuid4(), + parent_task_id=uuid4(), + task_type="planning", + dependency_ids=[], + team="main_pm", + pr_number=139, + pr_url="https://example/pr/139", + branch_name="feature/main_pm/root", + batch_id=None, + project=_over_cap_project(), + ) + task_svc.get.return_value = t + task_svc.agent_for.return_value = MagicMock( + role="pr_reviewer", slug="be-pr-reviewer" + ) + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.unmet_dependency_ids = AsyncMock(return_value=[]) + task_svc.pr_gate_claim = AsyncMock(return_value=t) + task_svc.project_month_spend_usd = AsyncMock(return_value=999.0) + c = Choreographer(_make_deps(task_svc)) + cc: Any = c + cc._build_gate_review_evidence = AsyncMock(return_value={"pr_number": 139}) + + env = await c.claim_gate_review(uuid4(), t.id) + body = env.as_dict() + assert body.get("error") is None, body + task_svc.project_month_spend_usd.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_claim_doc_task_succeeds_at_cap() -> None: + task_svc = AsyncMock() + t = MagicMock( + id=uuid4(), + status="awaiting_documentation", + assigned_to=None, + parent_task_id=None, + task_type="documentation", + team="backend", + branch_name="feature/backend/abc", + quick_context=None, + documents=[], + commits=[{"sha": "abc123", "message": "[x] work"}], + pr_number=7, + pr_url="https://github.com/x/y/pull/7", + dev_notes="done", + acceptance_criteria_status=[], + work_session_id=uuid4(), + dependency_ids=[], + batch_id=None, + project=_over_cap_project(), + ) + task_svc.get.return_value = t + task_svc.agent_for.return_value = MagicMock(role="documenter", team="backend") + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.unmet_dependency_ids = AsyncMock(return_value=[]) + task_svc.doc_claim.return_value = t + task_svc.project_month_spend_usd = AsyncMock(return_value=999.0) + git_svc = AsyncMock() + git_svc.diff.return_value = "diff" + git_svc.list_changed_files.return_value = ["README.md"] + c = Choreographer(_make_deps(task_svc, git=git_svc)) + + env = await c.claim_doc_task(uuid4(), t.id) + body = env.as_dict() + assert body["error"] is None, body + task_svc.project_month_spend_usd.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_claim_pr_review_succeeds_at_cap() -> None: + task_svc = AsyncMock() + t = MagicMock( + id=uuid4(), + status="pending", + assigned_to=None, + parent_task_id=None, + task_type="code", + dependency_ids=[], + team="system", + batch_id=None, + project=_over_cap_project(), + ) + task_svc.get.return_value = t + task_svc.agent_for.return_value = MagicMock( + role="pr_reviewer", slug="be-pr-reviewer" + ) + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.unmet_dependency_ids = AsyncMock(return_value=[]) + task_svc.has_earlier_incomplete_code_sibling.return_value = False + task_svc.pr_review_claim = AsyncMock(return_value=t) + task_svc.project_month_spend_usd = AsyncMock(return_value=999.0) + c = Choreographer(_make_deps(task_svc)) + cc: Any = c + cc._build_pr_review_evidence = AsyncMock(return_value={}) + + env = await c.claim_pr_review(uuid4(), t.id) + body = env.as_dict() + assert body.get("error") is None, body + task_svc.project_month_spend_usd.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Work-starting claims: guard ON — both refuse at cap. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_i_will_work_on_refuses_at_cap() -> None: + agent_id = uuid4() + task_id = uuid4() + target = MagicMock( + id=task_id, + status="pending", + plan=None, + assigned_to=None, + parent_task_id=None, + sequence=0, + task_type="code", + team="backend", + dependency_ids=[], + project=_over_cap_project(), + ) + task_svc = AsyncMock() + task_svc.get.return_value = target + task_svc.agent_for.return_value = MagicMock( + id=agent_id, role="developer", team="backend", slug=None + ) + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.unmet_dependency_ids = AsyncMock(return_value=[]) + task_svc.project_month_spend_usd = AsyncMock(return_value=999.0) + c = Choreographer(_make_deps(task_svc)) + + env = await c.i_will_work_on(agent_id, task_id, plan="x", steps=_STEPS) + body = env.as_dict() + assert body["error"] == "invalid_state", body + assert "10.00" in body["message"] and "999.00" in body["message"] + task_svc.claim.assert_not_awaited() + task_svc.project_month_spend_usd.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_i_will_plan_refuses_at_cap() -> None: + pm_id = uuid4() + task_id = uuid4() + target = MagicMock( + id=task_id, + status="pending", + plan=None, + assigned_to=None, + parent_task_id=None, + sequence=0, + task_type="planning", + team="backend", + dependency_ids=[], + project=_over_cap_project(), + ) + task_svc = AsyncMock() + task_svc.get.return_value = target + task_svc.agent_for.return_value = MagicMock( + id=pm_id, role="cell_pm", team="backend", slug=None + ) + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.unmet_dependency_ids = AsyncMock(return_value=[]) + task_svc.project_month_spend_usd = AsyncMock(return_value=999.0) + c = Choreographer(_make_deps(task_svc)) + + env = await c.i_will_plan( + pm_id, + task_id, + plan="break it down", + rich_plan={ + "approach": ( + "Single-cell decomposition: backend handles the full scope; " + "no frontend or ux work required for this planning task. " + "be-dev-1 owns the change end to end; QA reviews after the " + "PR opens, documentation follows, then be-pm completes and " + "submits up. Strict sequencing, no cross-cell dependencies." + ), + "sub_tasks": [ + { + "title": "Backend planning slice", + "description": ( + "scope the change, assign be-dev-1, who implements " + "with tests and opens the leaf PR for QA review." + ), + } + ], + }, + ) + body = env.as_dict() + assert body["error"] == "invalid_state", body + task_svc.claim.assert_not_awaited() + task_svc.project_month_spend_usd.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Plumbing contract on _run_claim_guards itself. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_claim_guards_only_checks_budget_when_asked() -> None: + task_svc = AsyncMock() + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.unmet_dependency_ids = AsyncMock(return_value=[]) + c = Choreographer(_make_deps(task_svc)) + task = MagicMock(id=uuid4(), dependency_ids=[]) + calls = {"n": 0} + + async def _fake_guard(_t: Any) -> None: + calls["n"] += 1 + + cc: Any = c + cc._project_budget_claim_guard = _fake_guard + + await c._run_claim_guards(agent_id=uuid4(), task=task, skip_dev_guards=True) + assert calls["n"] == 0, ( + "budget guard must not run without check_project_budget=True" + ) + + await c._run_claim_guards( + agent_id=uuid4(), task=task, skip_dev_guards=True, check_project_budget=True + ) + assert calls["n"] == 1 diff --git a/tests/unit/gateway/test_budget_unblock_guard.py b/tests/unit/gateway/test_budget_unblock_guard.py new file mode 100644 index 00000000..84d6e234 --- /dev/null +++ b/tests/unit/gateway/test_budget_unblock_guard.py @@ -0,0 +1,165 @@ +"""unblock's budget-breach re-check (ROBOCO_TASK_BUDGETS_ENABLED). + +A task the orchestrator's budget sweep BLOCKed carries the BUDGET_BLOCKED +marker (`_handle_task_budget_breach`). `unblock` re-checks spend-vs-cap +before letting it through: still over refuses (naming the budget +remediation), so a PM can't silently re-breach the same cap the next tick; +under (the CEO raised the cap) clears the marker and the unblock proceeds +exactly as before. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from roboco.config import settings +from roboco.foundation.policy.content import markers +from roboco.models.base import TaskType +from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps + + +def _make_deps(task_svc: AsyncMock) -> ChoreographerDeps: + base: dict[str, Any] = { + "task": task_svc, + "work_session": AsyncMock(), + "git": AsyncMock(), + "a2a": AsyncMock(), + "journal": AsyncMock(), + "audit": AsyncMock(), + "evidence_repo": AsyncMock(), + } + base["journal"].has_decision_for_task.return_value = True + base["journal"].latest_decision_at.return_value = datetime.now(UTC) + return ChoreographerDeps(**base) + + +def _budget_blocked_task(*, budget_usd: float | None = 5.0) -> MagicMock: + t = MagicMock( + id=uuid4(), + status="blocked", + pre_block_state="in_progress", + pre_block_assignee=uuid4(), + pre_block_metadata={}, + dependency_ids=[], + task_type=TaskType.CODE, + budget_usd=budget_usd, + orchestration_markers=None, + ) + markers.mark_budget_blocked(t) + return t + + +@pytest.fixture(autouse=True) +def _budgets_on(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings, "task_budgets_enabled", True) + + +@pytest.mark.asyncio +async def test_unblock_refuses_while_still_over_cap() -> None: + t = _budget_blocked_task(budget_usd=5.0) + task_svc = AsyncMock() + task_svc.get.return_value = t + task_svc.unmet_dependency_ids.return_value = [] + task_svc.task_spend_usd.return_value = 7.0 + c = Choreographer(_make_deps(task_svc)) + + env = await c.unblock(uuid4(), t.id, "attempting to resume") + body = env.as_dict() + assert body["error"] == "invalid_state", body + assert "5.00" in body["message"] and "7.00" in body["message"] + assert "budget" in body["remediate"].lower() + task_svc.unblock_with_restore.assert_not_awaited() + # Marker survives — a retry that hasn't actually cleared the cap must + # still be caught. + assert markers.is_budget_blocked(t) is True + + +@pytest.mark.asyncio +async def test_unblock_succeeds_when_already_under_cap() -> None: + t = _budget_blocked_task(budget_usd=5.0) + task_svc = AsyncMock() + task_svc.get.return_value = t + task_svc.unmet_dependency_ids.return_value = [] + task_svc.task_spend_usd.return_value = 3.0 + task_svc.unblock_with_restore.return_value = t + c = Choreographer(_make_deps(task_svc)) + + env = await c.unblock(uuid4(), t.id, "spend never actually breached") + body = env.as_dict() + assert body.get("error") is None, body + task_svc.unblock_with_restore.assert_awaited_once() + assert markers.is_budget_blocked(t) is False + + +@pytest.mark.asyncio +async def test_raise_then_unblock_succeeds_after_a_prior_refusal() -> None: + """First attempt: still over -> refused. CEO raises budget_usd. Second + attempt: now under -> succeeds, marker cleared.""" + t = _budget_blocked_task(budget_usd=5.0) + task_svc = AsyncMock() + task_svc.get.return_value = t + task_svc.unmet_dependency_ids.return_value = [] + task_svc.task_spend_usd.return_value = 7.0 + task_svc.unblock_with_restore.return_value = t + c = Choreographer(_make_deps(task_svc)) + pm_id = uuid4() + + refused = await c.unblock(pm_id, t.id, "attempting resume") + assert refused.as_dict()["error"] == "invalid_state" + task_svc.unblock_with_restore.assert_not_awaited() + + # CEO raises the cap; re-block for a fresh attempt (the same task row, as + # it would be across two real requests). + t.budget_usd = 20.0 + t.status = "blocked" + succeeded = await c.unblock(pm_id, t.id, "raised the budget, resuming") + body = succeeded.as_dict() + assert body.get("error") is None, body + task_svc.unblock_with_restore.assert_awaited_once() + assert markers.is_budget_blocked(t) is False + + +@pytest.mark.asyncio +async def test_flag_off_ignores_the_marker(monkeypatch: pytest.MonkeyPatch) -> None: + """The marker alone must never gate anything with the flag off.""" + monkeypatch.setattr(settings, "task_budgets_enabled", False) + t = _budget_blocked_task(budget_usd=5.0) + task_svc = AsyncMock() + task_svc.get.return_value = t + task_svc.unmet_dependency_ids.return_value = [] + task_svc.unblock_with_restore.return_value = t + c = Choreographer(_make_deps(task_svc)) + + env = await c.unblock(uuid4(), t.id, "flag is off") + assert env.as_dict().get("error") is None, env.as_dict() + task_svc.task_spend_usd.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_non_budget_block_never_reaches_the_spend_query() -> None: + """A task blocked for an ordinary reason (no BUDGET_BLOCKED marker) skips + the guard entirely — it's not a budget block at all.""" + t = MagicMock( + id=uuid4(), + status="blocked", + pre_block_state="in_progress", + pre_block_assignee=uuid4(), + pre_block_metadata={}, + dependency_ids=[], + task_type=TaskType.CODE, + budget_usd=5.0, + orchestration_markers=None, + ) + task_svc = AsyncMock() + task_svc.get.return_value = t + task_svc.unmet_dependency_ids.return_value = [] + task_svc.unblock_with_restore.return_value = t + c = Choreographer(_make_deps(task_svc)) + + env = await c.unblock(uuid4(), t.id, "manual escalation resolved") + assert env.as_dict().get("error") is None, env.as_dict() + task_svc.task_spend_usd.assert_not_awaited() diff --git a/tests/unit/runtime/test_task_budget_sweep.py b/tests/unit/runtime/test_task_budget_sweep.py new file mode 100644 index 00000000..99a8ad30 --- /dev/null +++ b/tests/unit/runtime/test_task_budget_sweep.py @@ -0,0 +1,340 @@ +"""Task-level $ budget sweep (ROBOCO_TASK_BUDGETS_ENABLED). + +`_sweep_budget_exceeded` gains a second trigger alongside the existing +tool-call halt: when the flag is on and an active task's own $ budget is +breached, the task is BLOCKED + the CEO notified (`_handle_task_budget_breach`) +BEFORE the agent is gracefully stopped — never a mid-verb kill, and never left +to bounce through `pending` for an instant re-claim. +""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from roboco.config import settings +from roboco.models.base import BlockerResolverType, TaskStatus, TaskType +from roboco.runtime.orchestrator import AgentOrchestrator, AgentState + +_MOCK_TASK_SPEND_USD = 3.0 + + +def _make_orchestrator() -> AgentOrchestrator: + with patch.object(AgentOrchestrator, "__init__", return_value=None): + orch = AgentOrchestrator.__new__(AgentOrchestrator) + orch._instances = {} + orch._lock = MagicMock() + return orch + + +def _instance(task_id: str | None) -> MagicMock: + inst = MagicMock() + inst.state = AgentState.ACTIVE + inst.container_id = "deadbeef1234" + inst.current_task_id = task_id + inst.error_count = 0 + inst.config = MagicMock(git_context=None) + return inst + + +def _db_ctx(db: Any) -> Any: + @asynccontextmanager + async def _ctx() -> Any: + yield db + + return _ctx + + +@pytest.mark.asyncio +async def test_task_budget_breach_blocks_before_graceful_stop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + orch = _make_orchestrator() + task_id = "11111111-1111-1111-1111-111111111111" + orch._instances = {"be-dev-1": _instance(task_id)} + monkeypatch.setattr(settings, "task_budgets_enabled", True) + + with ( + patch.object( + AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None) + ), + patch.object(orch, "_task_budget_breach", AsyncMock(return_value=(5.0, 7.5))), + patch.object(orch, "_handle_task_budget_breach", AsyncMock()) as handle_mock, + patch.object(orch, "stop_agent", AsyncMock()) as stop_mock, + ): + await orch._sweep_budget_exceeded() + + # Block + notify runs, and runs BEFORE stop_agent (never a mid-verb kill — + # graceful=True, and the task is already blocked by the time the agent dies). + handle_mock.assert_awaited_once_with(task_id, cap_usd=5.0, spend_usd=7.5) + stop_mock.assert_awaited_once_with( + "be-dev-1", + graceful=True, + release_claim=True, + stop_reason="budget_exceeded_task", + ) + + +@pytest.mark.asyncio +async def test_flag_off_never_checks_task_budget( + monkeypatch: pytest.MonkeyPatch, +) -> None: + orch = _make_orchestrator() + task_id = "11111111-1111-1111-1111-111111111111" + orch._instances = {"be-dev-1": _instance(task_id)} + monkeypatch.setattr(settings, "task_budgets_enabled", False) + + with ( + patch.object( + AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None) + ), + patch.object(orch, "_task_budget_breach", AsyncMock()) as breach_mock, + patch.object(orch, "stop_agent", AsyncMock()) as stop_mock, + ): + await orch._sweep_budget_exceeded() + + breach_mock.assert_not_awaited() + stop_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_under_budget_is_a_no_op(monkeypatch: pytest.MonkeyPatch) -> None: + orch = _make_orchestrator() + task_id = "11111111-1111-1111-1111-111111111111" + orch._instances = {"be-dev-1": _instance(task_id)} + monkeypatch.setattr(settings, "task_budgets_enabled", True) + + with ( + patch.object( + AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None) + ), + patch.object(orch, "_task_budget_breach", AsyncMock(return_value=None)), + patch.object(orch, "_handle_task_budget_breach", AsyncMock()) as handle_mock, + patch.object(orch, "stop_agent", AsyncMock()) as stop_mock, + ): + await orch._sweep_budget_exceeded() + + handle_mock.assert_not_awaited() + stop_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_tool_call_halt_path_is_unchanged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The pre-existing tool-call halt trigger still fires with its own + stop_reason, independent of the new $ budget path.""" + orch = _make_orchestrator() + orch._instances = {"be-dev-1": _instance(None)} + monkeypatch.setattr(settings, "task_budgets_enabled", False) + + halt_status = {"halt": True, "total": 301, "halt_threshold": 300} + with ( + patch.object( + AgentOrchestrator, + "_fetch_budget_status", + AsyncMock(return_value=halt_status), + ), + patch.object(orch, "stop_agent", AsyncMock()) as stop_mock, + ): + await orch._sweep_budget_exceeded() + + stop_mock.assert_awaited_once_with( + "be-dev-1", graceful=True, release_claim=True, stop_reason="budget_sweep" + ) + + +@pytest.mark.asyncio +async def test_no_task_id_skips_task_budget_check( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A taskless spawn (current_task_id=None) never reaches the $ budget + check even with the flag on.""" + orch = _make_orchestrator() + orch._instances = {"be-dev-1": _instance(None)} + monkeypatch.setattr(settings, "task_budgets_enabled", True) + + with ( + patch.object( + AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None) + ), + patch.object(orch, "_task_budget_breach", AsyncMock()) as breach_mock, + patch.object(orch, "stop_agent", AsyncMock()) as stop_mock, + ): + await orch._sweep_budget_exceeded() + + breach_mock.assert_not_awaited() + stop_mock.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# _handle_task_budget_breach: the block + notify write itself +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_handle_breach_blocks_task_and_notifies_ceo() -> None: + orch = _make_orchestrator() + task_id = "22222222-2222-2222-2222-222222222222" + task = MagicMock(status=TaskStatus.IN_PROGRESS) + db = MagicMock() + + task_svc = MagicMock() + task_svc.get = AsyncMock(return_value=task) + task_svc.admin_set_status = AsyncMock() + delivery = MagicMock() + delivery.notify_ceo_of_budget_breach = AsyncMock() + + with ( + patch("roboco.db.base.get_db_context", _db_ctx(db)), + patch("roboco.services.task.TaskService", return_value=task_svc), + patch( + "roboco.services.notification_delivery.get_notification_delivery_service", + return_value=delivery, + ), + ): + await orch._handle_task_budget_breach(task_id, cap_usd=5.0, spend_usd=8.0) + + assert task.blocker_resolver_type == BlockerResolverType.HUMAN + task_svc.admin_set_status.assert_awaited_once() + args, _kwargs = task_svc.admin_set_status.call_args + assert args[1] == TaskStatus.BLOCKED + delivery.notify_ceo_of_budget_breach.assert_awaited_once_with( + task=task, task_id=args[0], cap_usd=5.0, spend_usd=8.0 + ) + + +@pytest.mark.asyncio +async def test_handle_breach_skips_a_task_that_already_moved_on() -> None: + """A stale re-check racing the task's own progress (e.g. it completed + between the read and the write) must not block/notify.""" + orch = _make_orchestrator() + task_id = "33333333-3333-3333-3333-333333333333" + task = MagicMock(status=TaskStatus.COMPLETED) + db = MagicMock() + + task_svc = MagicMock() + task_svc.get = AsyncMock(return_value=task) + task_svc.admin_set_status = AsyncMock() + delivery = MagicMock() + delivery.notify_ceo_of_budget_breach = AsyncMock() + + with ( + patch("roboco.db.base.get_db_context", _db_ctx(db)), + patch("roboco.services.task.TaskService", return_value=task_svc), + patch( + "roboco.services.notification_delivery.get_notification_delivery_service", + return_value=delivery, + ), + ): + await orch._handle_task_budget_breach(task_id, cap_usd=5.0, spend_usd=8.0) + + task_svc.admin_set_status.assert_not_awaited() + delivery.notify_ceo_of_budget_breach.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# _task_budget_breach: cap resolution (null -> TaskType default) + spend sum +# --------------------------------------------------------------------------- + + +@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).""" + orch = _make_orchestrator() + task_id = "44444444-4444-4444-4444-444444444444" + task = MagicMock( + status=TaskStatus.IN_PROGRESS, task_type=TaskType.DOCUMENTATION, budget_usd=None + ) + task_svc = MagicMock() + task_svc.get = AsyncMock(return_value=task) + task_svc.task_spend_usd = AsyncMock(return_value=_MOCK_TASK_SPEND_USD) + db = MagicMock() + + with ( + patch("roboco.db.base.get_db_context", _db_ctx(db)), + patch("roboco.services.task.TaskService", return_value=task_svc), + ): + 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 + + +@pytest.mark.asyncio +async def test_breach_none_when_task_left_claimed_or_in_progress() -> None: + """A stale re-check (the task already reached e.g. awaiting_qa) is not a + breach — the spend query is never even issued.""" + orch = _make_orchestrator() + task_id = "55555555-5555-5555-5555-555555555555" + task = MagicMock( + status=TaskStatus.AWAITING_QA, task_type=TaskType.CODE, budget_usd=1.0 + ) + task_svc = MagicMock() + task_svc.get = AsyncMock(return_value=task) + task_svc.task_spend_usd = AsyncMock() + db = MagicMock() + + with ( + patch("roboco.db.base.get_db_context", _db_ctx(db)), + patch("roboco.services.task.TaskService", return_value=task_svc), + ): + breach = await orch._task_budget_breach(task_id) + + assert breach is None + task_svc.task_spend_usd.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Repeated ticks: a blocked task must not re-fire the breach handling. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_repeated_ticks_do_not_refire_once_blocked( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Two consecutive _sweep_budget_exceeded ticks against the same still- + registered instance: tick 1 detects the breach and blocks the task; by + tick 2 the task is BLOCKED (no longer CLAIMED/IN_PROGRESS), so + _task_budget_breach's own status guard returns None — the sweep never + re-blocks / re-notifies / re-stops a task that's already been handled.""" + orch = _make_orchestrator() + task_id = "77777777-7777-7777-7777-777777777777" + orch._instances = {"be-dev-1": _instance(task_id)} + monkeypatch.setattr(settings, "task_budgets_enabled", True) + + in_progress_task = MagicMock( + status=TaskStatus.IN_PROGRESS, task_type=TaskType.CODE, budget_usd=5.0 + ) + # Simulates the task having been transitioned to BLOCKED by tick 1's + # (mocked-out) _handle_task_budget_breach before tick 2 re-checks it. + blocked_task = MagicMock( + status=TaskStatus.BLOCKED, task_type=TaskType.CODE, budget_usd=5.0 + ) + task_svc = MagicMock() + task_svc.get = AsyncMock(side_effect=[in_progress_task, blocked_task]) + task_svc.task_spend_usd = AsyncMock(return_value=7.0) + db = MagicMock() + + with ( + patch.object( + AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None) + ), + patch("roboco.db.base.get_db_context", _db_ctx(db)), + patch("roboco.services.task.TaskService", return_value=task_svc), + patch.object(orch, "_handle_task_budget_breach", AsyncMock()) as handle_mock, + patch.object(orch, "stop_agent", AsyncMock()) as stop_mock, + ): + await orch._sweep_budget_exceeded() # tick 1: breach detected + await orch._sweep_budget_exceeded() # tick 2: already blocked + + handle_mock.assert_awaited_once() + stop_mock.assert_awaited_once() diff --git a/tests/unit/services/test_project_month_spend_usd_db.py b/tests/unit/services/test_project_month_spend_usd_db.py new file mode 100644 index 00000000..cd013731 --- /dev/null +++ b/tests/unit/services/test_project_month_spend_usd_db.py @@ -0,0 +1,212 @@ +"""TaskService.project_month_spend_usd against a real Postgres DB. + +The join (agent_spawn_sessions.task_id, a plain String(36), to tasks.id via a +cast) and the month-boundary filter are exactly the kind of thing a mocked +unit test can't prove actually executes as real SQL. This also pins the +open-session live-token pricing fix: a still-open session's +estimated_cost_usd is null until close, so it must be priced from its token +columns via calculate_cost, not silently summed as $0. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +import pytest +from roboco.billing.pricing import calculate_cost +from roboco.db.tables import AgentSpawnSessionTable, AgentTable, ProjectTable, TaskTable +from roboco.models.base import ( + AgentRole, + AgentStatus, + Complexity, + TaskNature, + TaskStatus, + TaskType, + Team, +) +from roboco.services.task import TaskService + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncSession + +_MODEL = "claude-sonnet-5" + + +async def _seed_project(session: AsyncSession) -> tuple[UUID, UUID]: + """Seed a system agent + project. Returns ``(project_id, system_agent_id)`` + — the latter doubles as the task-row FK target below.""" + system_agent = AgentTable( + id=uuid4(), + name="System", + slug=f"system-{uuid4().hex[:8]}", + role=AgentRole.SYSTEM, + team=None, + status=AgentStatus.ACTIVE, + model_config={}, + system_prompt="system", + capabilities=[], + permissions={}, + metrics={}, + ) + session.add(system_agent) + await session.flush() + + project = ProjectTable( + id=uuid4(), + name="Budget Spend Test Project", + slug=f"budget-spend-{uuid4().hex[:8]}", + git_url="https://github.com/example/budget-spend.git", + default_branch="main", + protected_branches=["main"], + assigned_cell=Team.BACKEND, + created_by=system_agent.id, + is_active=True, + ) + session.add(project) + await session.flush() + return UUID(str(project.id)), UUID(str(system_agent.id)) + + +async def _seed_task(session: AsyncSession, project_id: UUID, created_by: UUID) -> UUID: + task = TaskTable( + id=uuid4(), + title="Budget spend fixture task", + description="A description long enough to satisfy any length floor.", + acceptance_criteria=["it exists"], + status=TaskStatus.IN_PROGRESS, + priority=2, + task_type=TaskType.CODE, + nature=TaskNature.TECHNICAL, + estimated_complexity=Complexity.MEDIUM, + created_by=created_by, + team=Team.BACKEND, + project_id=project_id, + ) + session.add(task) + await session.flush() + return UUID(str(task.id)) + + +def _spawn_session( + task_id: UUID, + *, + started_at: datetime, + estimated_cost_usd: float | None, + ended_at: datetime | None, + tokens: tuple[int, int] = (0, 0), +) -> AgentSpawnSessionTable: + tokens_input, tokens_output = tokens + return AgentSpawnSessionTable( + id=uuid4(), + agent_slug="be-dev-1", + team="backend", + role="developer", + model=_MODEL, + task_id=str(task_id), + started_at=started_at, + ended_at=ended_at, + estimated_cost_usd=estimated_cost_usd, + tokens_input=tokens_input, + tokens_output=tokens_output, + ) + + +@pytest.mark.asyncio +async def test_sums_closed_and_prices_open_session(db_session: AsyncSession) -> None: + """A closed session's estimated_cost_usd + an open session's live-token + price (calculate_cost) — not the open session silently counted as $0.""" + project_id, agent_id = await _seed_project(db_session) + task_id = await _seed_task(db_session, project_id, agent_id) + now = datetime.now(UTC) + + db_session.add( + _spawn_session( + task_id, + started_at=now - timedelta(hours=2), + estimated_cost_usd=2.5, + ended_at=now - timedelta(hours=1), + ) + ) + db_session.add( + _spawn_session( + task_id, + started_at=now - timedelta(minutes=30), + estimated_cost_usd=None, + ended_at=None, + tokens=(100_000, 50_000), + ) + ) + await db_session.flush() + + expected_open_cost = calculate_cost( + model=_MODEL, tokens_input=100_000, tokens_output=50_000 + ) + assert expected_open_cost > 0, ( + "fixture model must be priced for this to be a real test" + ) + + svc = TaskService(db_session) + total = await svc.project_month_spend_usd(project_id) + assert total == pytest.approx(2.5 + expected_open_cost) + + +@pytest.mark.asyncio +async def test_excludes_last_months_session(db_session: AsyncSession) -> None: + """A session that started before this calendar month's boundary must not + count, even though its cost is closed and non-zero.""" + project_id, agent_id = await _seed_project(db_session) + task_id = await _seed_task(db_session, project_id, agent_id) + month_start = datetime.now(UTC).replace( + day=1, hour=0, minute=0, second=0, microsecond=0 + ) + last_month = month_start - timedelta(days=1) + + db_session.add( + _spawn_session( + task_id, + started_at=last_month, + estimated_cost_usd=50.0, + ended_at=last_month + timedelta(hours=1), + ) + ) + await db_session.flush() + + svc = TaskService(db_session) + total = await svc.project_month_spend_usd(project_id) + assert total == 0.0 + + +@pytest.mark.asyncio +async def test_join_excludes_other_projects_tasks(db_session: AsyncSession) -> None: + """A session on ANOTHER project's task must never bleed into this + project's sum — proves the join filters by project_id, not just presence + in agent_spawn_sessions.""" + project_id, agent_id = await _seed_project(db_session) + other_project_id, other_agent_id = await _seed_project(db_session) + my_task_id = await _seed_task(db_session, project_id, agent_id) + other_task_id = await _seed_task(db_session, other_project_id, other_agent_id) + now = datetime.now(UTC) + + db_session.add( + _spawn_session( + my_task_id, + started_at=now - timedelta(hours=1), + estimated_cost_usd=1.0, + ended_at=now, + ) + ) + db_session.add( + _spawn_session( + other_task_id, + started_at=now - timedelta(hours=1), + estimated_cost_usd=999.0, + ended_at=now, + ) + ) + await db_session.flush() + + svc = TaskService(db_session) + total = await svc.project_month_spend_usd(project_id) + assert total == pytest.approx(1.0)