mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
feat(budgets): per-task and per-project cost budgets (flag-gated) (#654)
* fix(notifications): exponential backoff + CAS claim for expired-unacked re-escalation The sweep re-escalated every expired unacked ack-required notification on every ~60s tick, forever — the live incident: 3 fresh blocker escalations + Telegram DMs per minute from a static stale pile. Now each notification carries reescalation_count / last_reescalated_at / reescalation_delivered_count (migration 079): first fire at expiry, then doubling intervals from 1h capped at 24h, hard stop after ROBOCO_NOTIFICATION_MAX_REESCALATIONS (default 5) with one permanent log carrying attempts-vs-delivered so 'seen and ignored' is distinguishable from 'route never worked'. The due/wait/capped decision is a pure function in foundation/policy/communications.py. Per adversarial review, the attempt slot is claimed by compare-and-set (UPDATE ... WHERE reescalation_count = :n) BEFORE delivery — the previous draft leaned on the 60s dedup window, which never engages for BLOCKER_ESCALATION (_LOOP_PRONE_TYPES excludes it), so concurrent sweeps would have double-delivered. A lost claim skips delivery outright. Legacy rows read as count=0 and keep today's first-fire semantics. 61 tests incl. a two-session CAS race and a real alembic upgrade/downgrade round trip. * feat(budgets): per-task and per-project cost budgets (flag-gated) tasks.budget_usd + projects.monthly_budget_usd (migration 080, chained on 079; adds ix_agent_spawn_sessions_task_id since both enforcement seams filter on bare task_id). Behind ROBOCO_TASK_BUDGETS_ENABLED (default off, feature-flags card) — verifiably inert when off. Claim-time: a project-month-spend guard applies to WORK-STARTING claims only (i_will_work_on / i_will_plan) — per adversarial review, review/ doc/gate/inbound-PR claims are exempt so in-flight work can always finish reviewing and merging at cap. Spend counts closed sessions' estimated_cost_usd PLUS open sessions priced live from token snapshots (the original closed-only sum read parallel long sessions as $0). Sweep-side: the existing budget sweep also prices the active task's spend vs budget_usd (TaskType defaults when null); on breach the task is BLOCKED (HUMAN resolver, budget marker) BEFORE the graceful stop so the unclaim no-ops and the dispatcher never respawns onto it, and the CEO notification names both recovery steps. unblock on a budget-blocked task re-checks live spend and refuses while still over — no silent re-breach loop. Panel: budget inputs in both dialogs (0 rejected — a zero budget silently blocks everything), spend logic consolidated in TaskService.task_spend_usd. 42 new tests incl. a real-DB spend-query suite and a two-tick non-refire sweep test. --------- Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
@@ -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")
|
||||
@@ -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> = {}): 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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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<Set<string>>(
|
||||
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<string, string[]> = {};
|
||||
@@ -800,6 +815,26 @@ function EditProjectForm({
|
||||
|
||||
{showAutonomy && (
|
||||
<>
|
||||
<div className="grid gap-2">
|
||||
<HelpTip label="Calendar-month cap on this project's summed agent-spawn spend; a claim is refused once reached. Requires the task-budgets flag armed fleet-wide (ROBOCO_TASK_BUDGETS_ENABLED). Leave blank for no cap.">
|
||||
<Label htmlFor="monthly_budget_usd">Monthly Budget (USD)</Label>
|
||||
</HelpTip>
|
||||
<Input
|
||||
id="monthly_budget_usd"
|
||||
type="number"
|
||||
min="0.01"
|
||||
step="0.01"
|
||||
value={monthlyBudgetUsd}
|
||||
onChange={(e) => setMonthlyBudgetUsd(e.target.value)}
|
||||
placeholder="No cap"
|
||||
/>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
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.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center justify-between">
|
||||
<HelpTip label="Opens a fix task automatically when this repo's default-branch CI goes red. Also requires the CI-watch engine armed fleet-wide (ROBOCO_CI_WATCH_ENABLED) to actually run.">
|
||||
<Label htmlFor="ci_watch_enabled">
|
||||
|
||||
@@ -176,4 +176,26 @@ describe("FeatureFlagsCard — M42 off-transition confirm + pending-keys Set", (
|
||||
expect(screen.getByText(/wikilinked Obsidian vault/i)).toBeInTheDocument();
|
||||
expect(screen.getByText(/board-review drafts/i)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// task_budgets_enabled joined FLAG_DESCRIPTIONS + FLAG_TOOLTIPS together —
|
||||
// guard both stay in sync the same way the vault/docs-sync case above does.
|
||||
it("renders the description and tooltip for task_budgets_enabled", async () => {
|
||||
getFeatureFlags.mockResolvedValueOnce({
|
||||
flags: [
|
||||
{
|
||||
key: "task_budgets_enabled",
|
||||
label: "Task/project cost budgets",
|
||||
enabled: false,
|
||||
},
|
||||
],
|
||||
note: "Changes take effect on the next backend restart.",
|
||||
});
|
||||
render(withQueryClient(<FeatureFlagsCard />));
|
||||
|
||||
expect(
|
||||
await screen.findByText(/per-project monthly and per-task cost caps/i),
|
||||
).toBeInTheDocument();
|
||||
const label = screen.getByText("Task/project cost budgets");
|
||||
expect(label.getAttribute("data-state")).toBe("closed");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -55,6 +55,8 @@ const FLAG_DESCRIPTIONS: Record<string, string> = {
|
||||
"Enforce a per-project architectural standard (.roboco/conventions.yml): inject the map, attach baseline constraints, and block i_am_done / pr_pass on misplaced definitions or lint suppressions.",
|
||||
possibilities_matrix_enabled:
|
||||
"When a task's work is already done (commits + open PR + all acceptance criteria addressed + no open findings), submit it for QA in one i_am_done call instead of 3-6 turns — skips the retroactive plan, journal tracing, and local quality (CI-green proxy) gates. Off by default: the standard path is unchanged until you arm this.",
|
||||
task_budgets_enabled:
|
||||
"Enforce per-project monthly and per-task cost caps (USD). A claim is refused once a project's monthly budget is reached; an active task whose own budget (or its task-type default) is breached is stopped and blocked, and you're notified. Set the caps on the project edit dialog and a task's detail page — a project/task with no cap set is unaffected either way.",
|
||||
rag_auto_update_enabled:
|
||||
"Keep the knowledge base index refreshed automatically.",
|
||||
transcript_prune_enabled:
|
||||
@@ -126,6 +128,8 @@ const FLAG_TOOLTIPS: Record<string, string> = {
|
||||
conventions_enabled: "Enforces each project's architectural placement rules.",
|
||||
possibilities_matrix_enabled:
|
||||
"Fast-paths work that's already been done elsewhere.",
|
||||
task_budgets_enabled:
|
||||
"Caps agent spend per project (monthly) and per task.",
|
||||
rag_auto_update_enabled:
|
||||
"Keeps the RAG knowledge index automatically refreshed.",
|
||||
transcript_prune_enabled:
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
|
||||
const { mutateAsync } = vi.hoisted(() => ({
|
||||
mutateAsync: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
vi.mock("@/hooks/use-tasks", () => ({
|
||||
useUpdateTask: () => ({ mutateAsync, isPending: false }),
|
||||
}));
|
||||
|
||||
vi.mock("sonner", () => ({ toast: { success: vi.fn(), error: vi.fn() } }));
|
||||
|
||||
vi.mock("@/components/agents/agent-selector", () => ({
|
||||
AgentSelector: () => null,
|
||||
}));
|
||||
vi.mock("@/components/projects/project-selector", () => ({
|
||||
ProjectSelector: () => null,
|
||||
}));
|
||||
vi.mock("../markdown-editor", () => ({ MarkdownEditor: () => null }));
|
||||
vi.mock("../acceptance-criteria-editor", () => ({
|
||||
AcceptanceCriteriaEditor: () => null,
|
||||
}));
|
||||
|
||||
// Collapsible: always render children open, so the Budget field (inside
|
||||
// Advanced Options) is reachable without simulating the toggle click.
|
||||
vi.mock("@/components/ui/collapsible", () => ({
|
||||
Collapsible: ({ children }: { children: React.ReactNode }) => <>{children}</>,
|
||||
CollapsibleTrigger: ({ children }: { children: React.ReactNode }) => (
|
||||
<>{children}</>
|
||||
),
|
||||
CollapsibleContent: ({ children }: { children: React.ReactNode }) => (
|
||||
<>{children}</>
|
||||
),
|
||||
}));
|
||||
|
||||
// Select: a native <select> stub, matching create-task-dialog's convention.
|
||||
vi.mock("@/components/ui/select", () => ({
|
||||
Select: ({
|
||||
value,
|
||||
onValueChange,
|
||||
children,
|
||||
}: {
|
||||
value: string;
|
||||
onValueChange?: (v: string) => void;
|
||||
children: React.ReactNode;
|
||||
}) => (
|
||||
<select value={value} onChange={(e) => onValueChange?.(e.target.value)}>
|
||||
{children}
|
||||
</select>
|
||||
),
|
||||
SelectTrigger: ({ children }: { children: React.ReactNode }) => (
|
||||
<>{children}</>
|
||||
),
|
||||
SelectValue: () => null,
|
||||
SelectContent: ({ children }: { children: React.ReactNode }) => (
|
||||
<>{children}</>
|
||||
),
|
||||
SelectItem: ({
|
||||
value,
|
||||
children,
|
||||
}: {
|
||||
value: string;
|
||||
children: React.ReactNode;
|
||||
}) => <option value={value}>{children}</option>,
|
||||
}));
|
||||
|
||||
import { EditTaskDialog } from "../edit-task-dialog";
|
||||
import { mockTasks } from "@/lib/mock-data";
|
||||
import { toast } from "sonner";
|
||||
|
||||
const task = mockTasks[0];
|
||||
|
||||
function submit() {
|
||||
fireEvent.submit(document.querySelector("form")!);
|
||||
}
|
||||
|
||||
function budgetInput(): HTMLInputElement {
|
||||
return screen.getByPlaceholderText("Task-type default") as HTMLInputElement;
|
||||
}
|
||||
|
||||
describe("EditTaskDialog — Budget (USD) input", () => {
|
||||
beforeEach(() => {
|
||||
mutateAsync.mockClear();
|
||||
});
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("renders empty when the task has no budget_usd", () => {
|
||||
render(
|
||||
<EditTaskDialog
|
||||
task={{ ...task, budget_usd: null }}
|
||||
open={true}
|
||||
onOpenChange={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(budgetInput().value).toBe("");
|
||||
});
|
||||
|
||||
it("pre-fills the stored budget_usd", () => {
|
||||
render(
|
||||
<EditTaskDialog
|
||||
task={{ ...task, budget_usd: 3.5 }}
|
||||
open={true}
|
||||
onOpenChange={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(budgetInput().value).toBe("3.5");
|
||||
});
|
||||
|
||||
it("rejects 0 with an inline error and does not submit", async () => {
|
||||
render(
|
||||
<EditTaskDialog
|
||||
task={{ ...task, budget_usd: null }}
|
||||
open={true}
|
||||
onOpenChange={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
fireEvent.change(budgetInput(), { 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 () => {
|
||||
render(
|
||||
<EditTaskDialog
|
||||
task={{ ...task, budget_usd: null }}
|
||||
open={true}
|
||||
onOpenChange={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
fireEvent.change(budgetInput(), { target: { value: "-1" } });
|
||||
submit();
|
||||
|
||||
await waitFor(() => {
|
||||
expect(toast.error).toHaveBeenCalledWith(
|
||||
expect.stringMatching(/greater than 0/i),
|
||||
);
|
||||
});
|
||||
expect(mutateAsync).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("submits null when left empty (use the task-type default)", async () => {
|
||||
render(
|
||||
<EditTaskDialog
|
||||
task={{ ...task, budget_usd: 3.5 }}
|
||||
open={true}
|
||||
onOpenChange={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
fireEvent.change(budgetInput(), { target: { value: "" } });
|
||||
submit();
|
||||
|
||||
await waitFor(() => expect(mutateAsync).toHaveBeenCalledTimes(1));
|
||||
const { updates } = mutateAsync.mock.calls[0][0] as {
|
||||
updates: Record<string, unknown>;
|
||||
};
|
||||
expect(updates.budget_usd).toBeNull();
|
||||
});
|
||||
|
||||
it("submits a positive budget as a number", async () => {
|
||||
render(
|
||||
<EditTaskDialog
|
||||
task={{ ...task, budget_usd: null }}
|
||||
open={true}
|
||||
onOpenChange={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
fireEvent.change(budgetInput(), { target: { value: "2.5" } });
|
||||
submit();
|
||||
|
||||
await waitFor(() => expect(mutateAsync).toHaveBeenCalledTimes(1));
|
||||
const { updates } = mutateAsync.mock.calls[0][0] as {
|
||||
updates: Record<string, unknown>;
|
||||
};
|
||||
expect(updates.budget_usd).toBe(2.5);
|
||||
});
|
||||
});
|
||||
@@ -120,6 +120,9 @@ function EditTaskDialogInner({
|
||||
? new Date(task.target_date).toISOString().slice(0, 16)
|
||||
: "",
|
||||
);
|
||||
const [budgetUsd, setBudgetUsd] = useState<string>(
|
||||
task.budget_usd != null ? String(task.budget_usd) : "",
|
||||
);
|
||||
const [advancedOpen, setAdvancedOpen] = useState(false);
|
||||
|
||||
const updateTask = useUpdateTask();
|
||||
@@ -143,6 +146,15 @@ function EditTaskDialogInner({
|
||||
}
|
||||
setAcError(undefined);
|
||||
|
||||
const trimmedBudget = budgetUsd.trim();
|
||||
const parsedBudget = trimmedBudget ? Number(trimmedBudget) : null;
|
||||
if (trimmedBudget && (Number.isNaN(parsedBudget) || parsedBudget! <= 0)) {
|
||||
toast.error(
|
||||
"Budget must be greater than 0 — leave it empty for the task-type default",
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
const trimmedCriteria = acceptanceCriteria
|
||||
.map((c) => c.trim())
|
||||
.filter(Boolean);
|
||||
@@ -164,6 +176,7 @@ function EditTaskDialogInner({
|
||||
project_id: projectId,
|
||||
assigned_to: assignedTo,
|
||||
target_date: targetDate ? new Date(targetDate).toISOString() : null,
|
||||
budget_usd: parsedBudget,
|
||||
...(criteriaChanged && { acceptance_criteria: trimmedCriteria }),
|
||||
},
|
||||
});
|
||||
@@ -338,6 +351,26 @@ function EditTaskDialogInner({
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Budget (USD) */}
|
||||
<div className="space-y-2">
|
||||
<HelpTip label="Caps this task's own agent-spawn spend; only enforced when the task-budgets feature flag is on. Empty = use the task-type default.">
|
||||
<Label>Budget (USD)</Label>
|
||||
</HelpTip>
|
||||
<Input
|
||||
type="number"
|
||||
min="0.01"
|
||||
step="0.01"
|
||||
placeholder="Task-type default"
|
||||
value={budgetUsd}
|
||||
onChange={(e) => setBudgetUsd(e.target.value)}
|
||||
/>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Must be greater than 0 — a 0 budget would block the task
|
||||
before it spends a cent. Leave blank for the task-type
|
||||
default.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
{/* Git Configuration Section */}
|
||||
<div className="space-y-4 pt-4 border-t">
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<string, string[]> | 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<string, string[]>;
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
+194
-15
@@ -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,26 +7754,95 @@ Start by:
|
||||
AgentState.WAITING_SHORT,
|
||||
):
|
||||
continue
|
||||
await self._check_budget_for_agent(
|
||||
client, agent_id, instance, task_budgets_on
|
||||
)
|
||||
|
||||
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)
|
||||
if data is None or not data.get("halt"):
|
||||
continue
|
||||
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"),
|
||||
halt_threshold=data.get("halt_threshold"),
|
||||
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:
|
||||
# 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",
|
||||
stop_reason=stop_reason,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
|
||||
@@ -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)
|
||||
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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user