Skip to content

Commit b8b2acb

Browse files
DorianZhengDorianZheng
authored
Merge pull request #67 from boxlite-ai/feature/token-budget-enforcement
feat(budget): add per-agent token budget enforcement
2 parents ac6bb38 + 157521e commit b8b2acb

5 files changed

Lines changed: 464 additions & 0 deletions

File tree

‎src/agent/agent-impl.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ import { ChannelManager } from './channel-manager.js';
6161
import { GroupManager } from './group-manager.js';
6262
import { TaskManager } from './task-manager.js';
6363
import { MessageProcessor } from './message-processor.js';
64+
import { registerBudgetActions } from './budget-actions.js';
6465

6566
export { type Agent };
6667

@@ -442,6 +443,7 @@ export class AgentImpl
442443
dataDir: this.config.dataDir,
443444
assistantName: this.config.assistantName,
444445
});
446+
registerBudgetActions(this, this.db);
445447
logger.info({ agent: this.name }, 'Database initialized');
446448
this.groupMgr.loadState();
447449

‎src/agent/budget-actions.ts‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
import { z } from 'zod';
2+
3+
import type { Agent } from '../api/agent.js';
4+
import type { AgentDb } from '../db.js';
5+
6+
/** Compute the start of the current daily budget period (epoch ms). */
7+
function getDailyPeriodStart(resetHour: number): number {
8+
const now = new Date();
9+
const start = new Date(
10+
Date.UTC(
11+
now.getUTCFullYear(),
12+
now.getUTCMonth(),
13+
now.getUTCDate(),
14+
resetHour,
15+
0,
16+
0,
17+
0,
18+
),
19+
);
20+
if (start.getTime() > now.getTime()) {
21+
start.setUTCDate(start.getUTCDate() - 1);
22+
}
23+
return start.getTime();
24+
}
25+
26+
export function registerBudgetActions(agent: Agent, db: AgentDb): void {
27+
// ── budget_set ──────────────────────────────────────────────────
28+
29+
agent.action(
30+
'budget_set',
31+
'Configure the token spending budget for an agent group. Pass null to remove a limit.',
32+
{
33+
group_jid: z.string().describe('The group/chat JID to configure'),
34+
daily_limit_usd: z
35+
.number()
36+
.positive()
37+
.nullable()
38+
.optional()
39+
.describe('Daily spend limit in USD. null removes the daily limit.'),
40+
total_limit_usd: z
41+
.number()
42+
.positive()
43+
.nullable()
44+
.optional()
45+
.describe('All-time spend limit in USD. null removes the total limit.'),
46+
reset_hour: z
47+
.number()
48+
.int()
49+
.min(0)
50+
.max(23)
51+
.optional()
52+
.describe(
53+
'UTC hour at which the daily counter resets (0–23, default 0).',
54+
),
55+
},
56+
async (args) => {
57+
db.setBudgetConfig(args.group_jid, {
58+
daily_limit_usd: args.daily_limit_usd,
59+
total_limit_usd: args.total_limit_usd,
60+
reset_hour: args.reset_hour,
61+
});
62+
return { ok: true };
63+
},
64+
);
65+
66+
// ── budget_get ──────────────────────────────────────────────────
67+
68+
agent.action(
69+
'budget_get',
70+
'Query the current budget configuration, pause state, and spending for an agent group.',
71+
{
72+
group_jid: z.string().describe('The group/chat JID to query'),
73+
},
74+
async (args) => {
75+
const config = db.getBudgetConfig(args.group_jid);
76+
const state = db.getBudgetState(args.group_jid);
77+
const resetHour = config?.reset_hour ?? 0;
78+
const periodStart = getDailyPeriodStart(resetHour);
79+
const dailyCostUsd = db.getDailyUsageUsd(args.group_jid, periodStart);
80+
const totalCostUsd = db.getTotalUsageUsd(args.group_jid);
81+
82+
const dailyPct =
83+
config?.daily_limit_usd != null && config.daily_limit_usd > 0
84+
? dailyCostUsd / config.daily_limit_usd
85+
: null;
86+
const totalPct =
87+
config?.total_limit_usd != null && config.total_limit_usd > 0
88+
? totalCostUsd / config.total_limit_usd
89+
: null;
90+
91+
return {
92+
config: {
93+
daily_limit_usd: config?.daily_limit_usd ?? null,
94+
total_limit_usd: config?.total_limit_usd ?? null,
95+
reset_hour: resetHour,
96+
},
97+
state,
98+
usage: {
99+
daily_cost_usd: dailyCostUsd,
100+
total_cost_usd: totalCostUsd,
101+
daily_pct: dailyPct,
102+
total_pct: totalPct,
103+
},
104+
};
105+
},
106+
);
107+
108+
// ── budget_resume ───────────────────────────────────────────────
109+
110+
agent.action(
111+
'budget_resume',
112+
'Clear the budget-exceeded pause for an agent group, allowing it to run again.',
113+
{
114+
group_jid: z.string().describe('The group/chat JID to resume'),
115+
},
116+
async (args) => {
117+
db.clearBudgetPaused(args.group_jid);
118+
return { ok: true };
119+
},
120+
);
121+
}

‎src/agent/message-processor.ts‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,26 @@ import { buildMcpRuntimeConfig } from './mcp-runtime.js';
2020

2121
export { buildMcpRuntimeConfig };
2222

23+
/** Compute the start of the current daily budget period (epoch ms). */
24+
function getDailyPeriodStart(resetHour: number): number {
25+
const now = new Date();
26+
const start = new Date(
27+
Date.UTC(
28+
now.getUTCFullYear(),
29+
now.getUTCMonth(),
30+
now.getUTCDate(),
31+
resetHour,
32+
0,
33+
0,
34+
0,
35+
),
36+
);
37+
if (start.getTime() > now.getTime()) {
38+
start.setUTCDate(start.getUTCDate() - 1);
39+
}
40+
return start.getTime();
41+
}
42+
2343
function hasWakeTrigger(
2444
messages: Array<{ content: string; sender: string; is_from_me?: boolean }>,
2545
chatJid: string,
@@ -137,6 +157,11 @@ export class MessageProcessor {
137157
'Processing messages',
138158
);
139159

160+
// ── Budget enforcement ───────────────────────���────────────────
161+
const budgetBlocked = await this.checkAndEnforceBudget(chatJid);
162+
if (budgetBlocked) return false;
163+
// ─────────────────────────────────────────────────────────────
164+
140165
let idleTimer: ReturnType<typeof setTimeout> | null = null;
141166
const resetIdleTimer = () => {
142167
if (idleTimer) clearTimeout(idleTimer);
@@ -322,6 +347,129 @@ export class MessageProcessor {
322347
return true;
323348
}
324349

350+
/**
351+
* Check whether this group's token budget is exceeded or warning-level.
352+
* Returns true if the agent should be blocked (budget exceeded), false otherwise.
353+
* Emits 'budget.exceeded' or 'budget.warning' events as appropriate.
354+
* Sends a user-facing message when pausing.
355+
*/
356+
private async checkAndEnforceBudget(chatJid: string): Promise<boolean> {
357+
const config = this.ctx.db.getBudgetConfig(chatJid);
358+
if (
359+
!config ||
360+
(config.daily_limit_usd == null && config.total_limit_usd == null)
361+
) {
362+
return false; // no budget configured
363+
}
364+
365+
const state = this.ctx.db.getBudgetState(chatJid);
366+
const now = new Date().toISOString();
367+
368+
// If already paused, block immediately.
369+
if (state.paused) {
370+
this.ctx.emit('budget.exceeded', {
371+
agentId: this.ctx.id,
372+
jid: chatJid,
373+
limitType: state.paused_reason === 'total_limit' ? 'total' : 'daily',
374+
limitUsd:
375+
state.paused_reason === 'total_limit'
376+
? (config.total_limit_usd ?? 0)
377+
: (config.daily_limit_usd ?? 0),
378+
usedUsd: 0,
379+
timestamp: now,
380+
});
381+
const channel = findChannel(this.channelMgr.channelArray, chatJid);
382+
await channel?.sendMessage?.(
383+
chatJid,
384+
'Agent paused: token budget exceeded. Resume from Dune settings.',
385+
);
386+
return true;
387+
}
388+
389+
const periodStart = getDailyPeriodStart(config.reset_hour);
390+
const dailyCost = this.ctx.db.getDailyUsageUsd(chatJid, periodStart);
391+
const totalCost = this.ctx.db.getTotalUsageUsd(chatJid);
392+
393+
// Check hard limits.
394+
if (
395+
(config.daily_limit_usd != null && dailyCost >= config.daily_limit_usd) ||
396+
(config.total_limit_usd != null && totalCost >= config.total_limit_usd)
397+
) {
398+
const isDaily =
399+
config.daily_limit_usd != null && dailyCost >= config.daily_limit_usd;
400+
const limitType: 'daily' | 'total' = isDaily ? 'daily' : 'total';
401+
const limitUsd = isDaily
402+
? config.daily_limit_usd!
403+
: config.total_limit_usd!;
404+
const usedUsd = isDaily ? dailyCost : totalCost;
405+
406+
this.ctx.db.setBudgetPaused(
407+
chatJid,
408+
`${limitType}_limit` as 'daily_limit' | 'total_limit',
409+
);
410+
this.ctx.emit('budget.exceeded', {
411+
agentId: this.ctx.id,
412+
jid: chatJid,
413+
limitType,
414+
limitUsd,
415+
usedUsd,
416+
timestamp: now,
417+
});
418+
419+
const channel = findChannel(this.channelMgr.channelArray, chatJid);
420+
await channel?.sendMessage?.(
421+
chatJid,
422+
`Agent paused: ${limitType} token budget exceeded ($${usedUsd.toFixed(4)} of $${limitUsd.toFixed(2)} used). Resume from Dune settings.`,
423+
);
424+
425+
logger.warn(
426+
{ chatJid, limitType, limitUsd, usedUsd, agent: this.ctx.name },
427+
'Budget exceeded — agent paused',
428+
);
429+
return true;
430+
}
431+
432+
// Check warning threshold (80%).
433+
const warnDaily =
434+
config.daily_limit_usd != null &&
435+
dailyCost >= 0.8 * config.daily_limit_usd;
436+
const warnTotal =
437+
config.total_limit_usd != null &&
438+
totalCost >= 0.8 * config.total_limit_usd;
439+
440+
if (warnDaily || warnTotal) {
441+
const isDaily = warnDaily;
442+
const limitType: 'daily' | 'total' = isDaily ? 'daily' : 'total';
443+
const limitUsd = isDaily
444+
? config.daily_limit_usd!
445+
: config.total_limit_usd!;
446+
const usedUsd = isDaily ? dailyCost : totalCost;
447+
const pctUsed = usedUsd / limitUsd;
448+
449+
this.ctx.emit('budget.warning', {
450+
agentId: this.ctx.id,
451+
jid: chatJid,
452+
pctUsed,
453+
limitType,
454+
limitUsd,
455+
usedUsd,
456+
timestamp: now,
457+
});
458+
459+
logger.info(
460+
{
461+
chatJid,
462+
limitType,
463+
pctUsed: Math.round(pctUsed * 100),
464+
agent: this.ctx.name,
465+
},
466+
'Budget warning — approaching limit',
467+
);
468+
}
469+
470+
return false; // proceed normally
471+
}
472+
325473
/** Execute agent in a container for the given group. */
326474
async runAgent(
327475
group: InternalRegisteredGroup,

‎src/api/events.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ export interface AgentEvents extends Record<string, any[]> {
2929
'task.run.succeeded': [payload: TaskRunSucceededEvent];
3030
'task.run.failed': [payload: TaskRunFailedEvent];
3131
'task.run.skipped': [payload: TaskRunSkippedEvent];
32+
'budget.exceeded': [payload: BudgetExceededEvent];
33+
'budget.warning': [payload: BudgetWarningEvent];
3234
started: [];
3335
stopped: [];
3436
}
@@ -302,3 +304,35 @@ export interface TaskRunSkippedEvent {
302304
detail?: string;
303305
timestamp: string;
304306
}
307+
308+
// ── Budget events ─────────────────────────────────────────────────
309+
310+
/** Agent was paused because it exceeded its token budget. */
311+
export interface BudgetExceededEvent {
312+
agentId: string;
313+
/** Group/chat JID that exceeded its budget. */
314+
jid: string;
315+
/** Which limit was hit. */
316+
limitType: 'daily' | 'total';
317+
/** The configured limit in USD. */
318+
limitUsd: number;
319+
/** Amount spent at time of enforcement. */
320+
usedUsd: number;
321+
timestamp: string;
322+
}
323+
324+
/** Agent is approaching its token budget (≥80% used). */
325+
export interface BudgetWarningEvent {
326+
agentId: string;
327+
/** Group/chat JID. */
328+
jid: string;
329+
/** Fraction of the limit used (e.g. 0.83 = 83%). */
330+
pctUsed: number;
331+
/** Which limit triggered the warning. */
332+
limitType: 'daily' | 'total';
333+
/** The configured limit in USD. */
334+
limitUsd: number;
335+
/** Amount spent so far. */
336+
usedUsd: number;
337+
timestamp: string;
338+
}

0 commit comments

Comments
 (0)