budget_guardrail.py
7.35 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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
"""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
_VALID_SCOPES = ("global", "per_task", "per_day")
_VALID_LIMIT_TYPES = ("token_count", "call_count")
class BudgetExceededError(Exception):
"""BUDGET_EXCEEDED -- 03-C §3.12."""
class InvalidParamError(Exception):
"""INVALID_PARAM -- 03-C §3.13: limit_value<=0, or scope/limit_type not
one of budget_guardrail_config's own enum values (checked here rather
than left to surface as a raw sqlite CHECK-constraint IntegrityError --
same reasoning as U-023's bug_record/task_id existence checks)."""
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 _resume_budget_blocked_tasks(conn):
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
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 mechanism REQ-10's boundary case needs -- updates
an *existing* row's limit_value only, silently no-ops if config_id
doesn't exist yet. upsert_budget_config() (U-030) is the fuller
v_budget_settings controller that also handles first-time creation.
"""
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()
return {"resumed_task_ids": _resume_budget_blocked_tasks(conn)}
def upsert_budget_config(conn, scope, limit_type, limit_value, updated_by, config_id="default"):
"""c_budget_config_update (03-C §3.13): v_budget_settings' submit_save.
Creates the row on first use (current_usage starts at 0) or updates an
existing one (current_usage is preserved -- this isn't a reset, just a
limit/scope/type change). Raising limit_value past current_usage is
exactly REQ-10's "人工在 v_budget_settings 调高额度" resume trigger
(04-P §4.4), so this also resumes any budget_exceeded blocked tasks --
same as update_budget_limit_and_resume, just reachable from a fresh
config too, not only an existing row."""
if scope not in _VALID_SCOPES:
raise InvalidParamError("scope must be one of {}, got {!r}".format(_VALID_SCOPES, scope))
if limit_type not in _VALID_LIMIT_TYPES:
raise InvalidParamError("limit_type must be one of {}, got {!r}".format(_VALID_LIMIT_TYPES, limit_type))
if limit_value <= 0:
raise InvalidParamError("limit_value must be > 0, got {!r}".format(limit_value))
exists = conn.execute(
"SELECT 1 FROM budget_guardrail_config WHERE config_id = ?", (config_id,)
).fetchone() is not None
if exists:
conn.execute(
"UPDATE budget_guardrail_config SET scope = ?, limit_type = ?, limit_value = ?, updated_by = ?, "
"updated_at = CURRENT_TIMESTAMP WHERE config_id = ?",
(scope, limit_type, limit_value, updated_by, config_id),
)
else:
conn.execute(
"INSERT INTO budget_guardrail_config (config_id, scope, limit_type, limit_value, current_usage, "
"updated_by) VALUES (?, ?, ?, ?, 0, ?)",
(config_id, scope, limit_type, limit_value, updated_by),
)
conn.commit()
resumed = _resume_budget_blocked_tasks(conn)
row = conn.execute(
"SELECT config_id, scope, limit_type, limit_value, current_usage, updated_at "
"FROM budget_guardrail_config WHERE config_id = ?", (config_id,)
).fetchone()
return {
"config": {
"config_id": row[0], "scope": row[1], "limit_type": row[2],
"limit_value": row[3], "current_usage": row[4], "updated_at": row[5],
},
"resumed_task_ids": resumed,
}