budget_guardrail.py
4.64 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
"""c_budget_check (03-C §3.12) + p_budget_guardrail (04-P §4.4): the
human-configured hard budget ceiling. This is NOT the same guardrail as
Claude Code's own plan-quota exhaustion (p_usage_limit_wait, already built
in U-010/U-012/U-013) -- 04-P is explicit that the two must never share a
blocked_reason: budget_guardrail_config triggering is a human-configured
limit that needs human judgment to raise (blocked_reason='budget_exceeded'),
while plan-quota exhaustion self-recovers on a timer and only escalates to
a human after a timeout (blocked_reason='usage_limit_wait_timeout').
"""
import json
import uuid
from src import org_loop_route
class BudgetExceededError(Exception):
"""BUDGET_EXCEEDED -- 03-C §3.12."""
def _write_status_event(conn, task_id, project_id, from_status, to_status):
conn.execute(
"INSERT INTO execution_event (event_id, project_id, task_id, event_type, payload) "
"VALUES (?, ?, ?, 'task_status_changed', ?)",
(str(uuid.uuid4()), project_id, task_id, json.dumps({"from_status": from_status, "to_status": to_status})),
)
def list_budget_configs(conn):
"""Read-only snapshot for v_dashboard's budget bar (02-V §2.2: "预算用量
条：引用 budget_guardrail_config.current_usage/limit_value") -- 03-C only
documents a PUT for this table (§3.13, U-030's settings-page scope); a
read path is needed for the dashboard and didn't exist anywhere, so it's
added here rather than invented ad hoc in the HTTP layer."""
rows = conn.execute(
"SELECT config_id, scope, limit_type, limit_value, current_usage, updated_at "
"FROM budget_guardrail_config ORDER BY config_id"
).fetchall()
return [
{
"config_id": r[0], "scope": r[1], "limit_type": r[2],
"limit_value": r[3], "current_usage": r[4], "updated_at": r[5],
}
for r in rows
]
def check_budget(conn, task_id, config_id, estimated_usage):
"""Checked before/during c_work_graph_execute (03-C §3.12). On success,
accumulates current_usage (04-P monitoring->monitoring: "累加
current_usage") and returns {"allowed": True, "remaining_quota"}. On
reaching/exceeding the limit, transitions the task to
blocked_pending_human (blocked_reason='budget_exceeded') and raises
BudgetExceededError instead of returning allowed=False silently,
matching 03-C's own errors list ("触发后对应任务转 blocked_pending_human").
"""
row = conn.execute(
"SELECT limit_value, current_usage FROM budget_guardrail_config WHERE config_id = ?", (config_id,)
).fetchone()
if row is None:
raise ValueError("no such budget_guardrail_config: " + config_id)
limit_value, current_usage = row
projected = current_usage + estimated_usage
if projected >= limit_value:
task_row = conn.execute("SELECT project_id, status FROM task WHERE task_id = ?", (task_id,)).fetchone()
conn.execute(
"UPDATE task SET status = 'blocked_pending_human', blocked_reason = 'budget_exceeded', "
"updated_at = CURRENT_TIMESTAMP WHERE task_id = ?",
(task_id,),
)
_write_status_event(conn, task_id, task_row[0], task_row[1], "blocked_pending_human")
conn.commit()
raise BudgetExceededError(
"config {} projected usage {} would reach/exceed limit {}".format(config_id, projected, limit_value)
)
conn.execute("UPDATE budget_guardrail_config SET current_usage = ? WHERE config_id = ?", (projected, config_id))
conn.commit()
return {"allowed": True, "remaining_quota": limit_value - projected}
def update_budget_limit_and_resume(conn, config_id, new_limit_value, updated_by):
"""04-P §4.4: paused -> monitoring, guard "人工在 v_budget_settings 调高额度".
This is the minimal backend mechanism REQ-10's boundary case needs, not
the full settings-page controller (c_budget_config_update's complete
field set is U-030's job, PH3) -- reuses U-016's manual_resume (via
org_loop_route.route_org_loop) for the actual status restoration rather
than reimplementing it.
"""
conn.execute(
"UPDATE budget_guardrail_config SET limit_value = ?, updated_by = ?, updated_at = CURRENT_TIMESTAMP "
"WHERE config_id = ?",
(new_limit_value, updated_by, config_id),
)
conn.commit()
blocked_tasks = conn.execute(
"SELECT task_id FROM task WHERE status = 'blocked_pending_human' AND blocked_reason = 'budget_exceeded'"
).fetchall()
resumed = []
for (task_id,) in blocked_tasks:
org_loop_route.route_org_loop(conn, task_id, "manual_resume")
resumed.append(task_id)
return {"resumed_task_ids": resumed}