feat(core): Persist agent budget spend in the database (#40138)

Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
This commit is contained in:
Anne Aguirre 2026-10-05 15:11:29 +00:00 committed by GitHub
parent 5f73c7cdcb
commit a2fd4189f1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
23 changed files with 958 additions and 87 deletions

View File

@ -10,6 +10,8 @@ Auto-generated from the PostgreSQL migrations in @n8n/db. Do not edit by hand.
| ---- | ------- | ------- | ---- |
| [public.activity_event](public.activity_event.md) | 11 | | BASE TABLE |
| [public.agent_background_job](public.agent_background_job.md) | 20 | | BASE TABLE |
| [public.agent_budget_applied_call](public.agent_budget_applied_call.md) | 2 | | BASE TABLE |
| [public.agent_budget_spend](public.agent_budget_spend.md) | 4 | | BASE TABLE |
| [public.agent_channel_status](public.agent_channel_status.md) | 11 | | BASE TABLE |
| [public.agent_chat_attachments](public.agent_chat_attachments.md) | 12 | | BASE TABLE |
| [public.agent_chat_subscriptions](public.agent_chat_subscriptions.md) | 6 | | BASE TABLE |
@ -443,6 +445,16 @@ erDiagram
timestamp_3__with_time_zone updatedAt
varchar_36_ workflowId
}
"public.agent_budget_applied_call" {
uuid callId
timestamp_3__with_time_zone createdAt
}
"public.agent_budget_spend" {
timestamp_3__with_time_zone createdAt
varchar_128_ key
double_precision totalUsd
timestamp_3__with_time_zone updatedAt
}
"public.agent_channel_status" {
varchar_36_ agentId FK
integer attempts

View File

@ -0,0 +1,38 @@
# public.agent_budget_applied_call
## Columns
| Name | Type | Default | Nullable | Children | Parents | Comment |
| ---- | ---- | ------- | -------- | -------- | ------- | ------- |
| callId | uuid | | false | | | Model call id. A second insert of the same id does not add spend. |
| createdAt | timestamp(3) with time zone | CURRENT_TIMESTAMP(3) | false | | | Time this call id was stored. A later prune can delete old rows. |
## Constraints
| Name | Type | Definition |
| ---- | ---- | ---------- |
| PK_0d91039b74f93ae1bdfc485e2ba | PRIMARY KEY | PRIMARY KEY ("callId") |
| agent_budget_applied_call_callId_not_null | n | NOT NULL "callId" |
| agent_budget_applied_call_createdAt_not_null | n | NOT NULL "createdAt" |
## Indexes
| Name | Definition |
| ---- | ---------- |
| PK_0d91039b74f93ae1bdfc485e2ba | CREATE UNIQUE INDEX "PK_0d91039b74f93ae1bdfc485e2ba" ON public.agent_budget_applied_call USING btree ("callId") |
## Relations
```mermaid
erDiagram
"public.agent_budget_applied_call" {
uuid callId
timestamp_3__with_time_zone createdAt
}
```
---
> Generated by [tbls](https://github.com/k1LoW/tbls)

View File

@ -0,0 +1,45 @@
# public.agent_budget_spend
## Columns
| Name | Type | Default | Nullable | Children | Parents | Comment |
| ---- | ---- | ------- | -------- | -------- | ------- | ------- |
| createdAt | timestamp(3) with time zone | CURRENT_TIMESTAMP(3) | false | | | |
| key | varchar(128) | | false | | | Session key is the root thread id. Month key is {agentId}:{YYYY-MM}. |
| totalUsd | double precision | 0 | false | | | Running USD estimate for this ledger key. A check rejects a negative total. |
| updatedAt | timestamp(3) with time zone | CURRENT_TIMESTAMP(3) | false | | | |
## Constraints
| Name | Type | Definition |
| ---- | ---- | ---------- |
| CHK_agent_budget_spend_total_usd_non_negative | CHECK | CHECK (("totalUsd" >= (0)::double precision)) |
| PK_ca88e3ef81cc367fc55807e582c | PRIMARY KEY | PRIMARY KEY (key) |
| agent_budget_spend_createdAt_not_null | n | NOT NULL "createdAt" |
| agent_budget_spend_key_not_null | n | NOT NULL key |
| agent_budget_spend_totalUsd_not_null | n | NOT NULL "totalUsd" |
| agent_budget_spend_updatedAt_not_null | n | NOT NULL "updatedAt" |
## Indexes
| Name | Definition |
| ---- | ---------- |
| PK_ca88e3ef81cc367fc55807e582c | CREATE UNIQUE INDEX "PK_ca88e3ef81cc367fc55807e582c" ON public.agent_budget_spend USING btree (key) |
## Relations
```mermaid
erDiagram
"public.agent_budget_spend" {
timestamp_3__with_time_zone createdAt
varchar_128_ key
double_precision totalUsd
timestamp_3__with_time_zone updatedAt
}
```
---
> Generated by [tbls](https://github.com/k1LoW/tbls)

View File

@ -10,6 +10,8 @@ Auto-generated from the SQLite migrations in @n8n/db. Do not edit by hand.
| ---- | ------- | ------- | ---- |
| [activity_event](activity_event.md) | 11 | | table |
| [agent_background_job](agent_background_job.md) | 20 | | table |
| [agent_budget_applied_call](agent_budget_applied_call.md) | 2 | | table |
| [agent_budget_spend](agent_budget_spend.md) | 4 | | table |
| [agent_channel_status](agent_channel_status.md) | 11 | | table |
| [agent_chat_attachments](agent_chat_attachments.md) | 12 | | table |
| [agent_chat_subscriptions](agent_chat_subscriptions.md) | 6 | | table |
@ -429,6 +431,16 @@ erDiagram
datetime_3_ updatedAt
varchar_36_ workflowId
}
"agent_budget_applied_call" {
varchar callId PK
datetime_3_ createdAt
}
"agent_budget_spend" {
datetime_3_ createdAt
varchar_128_ key PK
REAL totalUsd
datetime_3_ updatedAt
}
"agent_channel_status" {
varchar_36_ agentId PK
INTEGER attempts

View File

@ -0,0 +1,48 @@
# agent_budget_applied_call
## Description
<details>
<summary><strong>Table Definition</strong></summary>
```sql
CREATE TABLE "agent_budget_applied_call" ("callId" varchar PRIMARY KEY NOT NULL, "createdAt" datetime(3) NOT NULL DEFAULT (STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')))
```
</details>
## Columns
| Name | Type | Default | Nullable | Children | Parents | Comment |
| ---- | ---- | ------- | -------- | -------- | ------- | ------- |
| callId | varchar | | false | | | |
| createdAt | datetime(3) | STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW') | false | | | |
## Constraints
| Name | Type | Definition |
| ---- | ---- | ---------- |
| callId | PRIMARY KEY | PRIMARY KEY (callId) |
| sqlite_autoindex_agent_budget_applied_call_1 | PRIMARY KEY | PRIMARY KEY (callId) |
## Indexes
| Name | Definition |
| ---- | ---------- |
| sqlite_autoindex_agent_budget_applied_call_1 | PRIMARY KEY (callId) |
## Relations
```mermaid
erDiagram
"agent_budget_applied_call" {
varchar callId PK
datetime_3_ createdAt
}
```
---
> Generated by [tbls](https://github.com/k1LoW/tbls)

View File

@ -0,0 +1,53 @@
# agent_budget_spend
## Description
<details>
<summary><strong>Table Definition</strong></summary>
```sql
CREATE TABLE "agent_budget_spend" ("key" varchar(128) PRIMARY KEY NOT NULL, "totalUsd" real NOT NULL DEFAULT (0), "createdAt" datetime(3) NOT NULL DEFAULT (STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')), "updatedAt" datetime(3) NOT NULL DEFAULT (STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')), CONSTRAINT "CHK_agent_budget_spend_total_usd_non_negative" CHECK ("totalUsd" >= 0))
```
</details>
## Columns
| Name | Type | Default | Nullable | Children | Parents | Comment |
| ---- | ---- | ------- | -------- | -------- | ------- | ------- |
| createdAt | datetime(3) | STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW') | false | | | |
| key | varchar(128) | | false | | | |
| totalUsd | REAL | 0 | false | | | |
| updatedAt | datetime(3) | STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW') | false | | | |
## Constraints
| Name | Type | Definition |
| ---- | ---- | ---------- |
| - | CHECK | CHECK ("totalUsd" >= 0) |
| key | PRIMARY KEY | PRIMARY KEY (key) |
| sqlite_autoindex_agent_budget_spend_1 | PRIMARY KEY | PRIMARY KEY (key) |
## Indexes
| Name | Definition |
| ---- | ---------- |
| sqlite_autoindex_agent_budget_spend_1 | PRIMARY KEY (key) |
## Relations
```mermaid
erDiagram
"agent_budget_spend" {
datetime_3_ createdAt
varchar_128_ key PK
REAL totalUsd
datetime_3_ updatedAt
}
```
---
> Generated by [tbls](https://github.com/k1LoW/tbls)

View File

@ -165,11 +165,7 @@ export {
} from './sdk/vector-store-filter';
export type { VectorFilterInput } from './sdk/vector-store-filter';
export { Guardrail } from './sdk/guardrail';
export {
budgetMonthKey,
createBudgetGuardrail,
InMemorySpendLedger,
} from './runtime/guardrails/budget-guardrail';
export { budgetMonthKey, createBudgetGuardrail } from './runtime/guardrails/budget-guardrail';
export type {
BudgetGuardrailOptions,
SpendEntry,

View File

@ -33,7 +33,8 @@ import type {
} from '../../types/sdk/tool';
import type { BuiltTelemetry } from '../../types/telemetry';
import { Workspace, getToolResultRunDirectory } from '../../workspace';
import { createBudgetGuardrail, InMemorySpendLedger } from '../guardrails/budget-guardrail';
import { createBudgetGuardrail } from '../guardrails/budget-guardrail';
import { spendLedger } from '../guardrails/__tests__/spend-ledger';
import { AgentRuntime } from '../loop/agent-runtime';
import { InMemoryMemory } from '../memory/memory-store';
import { OBSERVATION_CONTINUATION_REMINDER } from '../model/message-list';
@ -1630,7 +1631,7 @@ describe('AgentRuntime — guardrails', () => {
});
const runtime = createRuntimeWithEchoTool(async () => ({ ok: true }));
const hook = createBudgetGuardrail({
ledger: new InMemorySpendLedger(),
ledger: spendLedger(),
sessionId: 'session-1',
sessionCostCapUsd: 1,
});

View File

@ -1,5 +1,6 @@
import type { GuardrailModelCallContext, TokenUsage } from '../../../types';
import { budgetMonthKey, createBudgetGuardrail, InMemorySpendLedger } from '../budget-guardrail';
import { budgetMonthKey, createBudgetGuardrail } from '../budget-guardrail';
import { spendLedger } from './spend-ledger';
const usage = (cost?: number): TokenUsage => ({
promptTokens: 1,
@ -16,31 +17,9 @@ function monthKey(agentId: string): string {
return budgetMonthKey(agentId);
}
describe('InMemorySpendLedger', () => {
it('does not increase totals when add repeats a callId', async () => {
const ledger = new InMemorySpendLedger();
await ledger.add('call-1', [
{ key: 'session', usd: 1 },
{ key: 'month', usd: 1 },
]);
const replay = await ledger.add('call-1', [
{ key: 'session', usd: 1 },
{ key: 'month', usd: 1 },
]);
expect(await ledger.read('session')).toBe(1);
expect(await ledger.read('month')).toBe(1);
expect(replay).toEqual([
{ key: 'session', totalUsd: 1, previousUsd: 1 },
{ key: 'month', totalUsd: 1, previousUsd: 1 },
]);
});
});
describe('createBudgetGuardrail', () => {
it('records usage.cost on the session key and the month key when the call is under the cap', async () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
const guardrail = createBudgetGuardrail({
ledger,
sessionId: 'session-1',
@ -57,7 +36,7 @@ describe('createBudgetGuardrail', () => {
});
it('records the crossing call and stops the next before', async () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
const guardrail = createBudgetGuardrail({
ledger,
sessionId: 'session-1',
@ -76,7 +55,7 @@ describe('createBudgetGuardrail', () => {
});
it('stops with budget.monthly when the month is spent and the session cap is not', async () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
const guardrail = createBudgetGuardrail({
ledger,
sessionId: 'session-1',
@ -94,7 +73,7 @@ describe('createBudgetGuardrail', () => {
});
it('leaves the totals unchanged when usage.cost is missing', async () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
const guardrail = createBudgetGuardrail({
ledger,
sessionId: 'session-1',
@ -111,7 +90,7 @@ describe('createBudgetGuardrail', () => {
});
it('fires onNotice once when the month total crosses the alert line', async () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
const onNotice = vi.fn();
const guardrail = createBudgetGuardrail({
ledger,
@ -132,7 +111,7 @@ describe('createBudgetGuardrail', () => {
});
it('does not stop or record when a budget is 0', async () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
const guardrail = createBudgetGuardrail({
ledger,
sessionId: 'session-1',
@ -148,7 +127,7 @@ describe('createBudgetGuardrail', () => {
});
it('stops with budget.misconfigured when a limit has no id', async () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
const missingSession = createBudgetGuardrail({ ledger, sessionCostCapUsd: 1 });
const missingAgent = createBudgetGuardrail({ ledger, monthlyBudgetUsd: 1 });

View File

@ -0,0 +1,29 @@
import type { SpendLedger } from '../budget-guardrail';
/** Ledger for guardrail tests. `add` ignores a repeated `callId`. */
export function spendLedger(): SpendLedger {
const totals = new Map<string, number>();
const appliedCallIds = new Set<string>();
return {
async add(callId, entries) {
if (appliedCallIds.has(callId)) {
return entries.map((entry) => {
const totalUsd = totals.get(entry.key) ?? 0;
return { key: entry.key, totalUsd, previousUsd: totalUsd };
});
}
appliedCallIds.add(callId);
return entries.map((entry) => {
const previousUsd = totals.get(entry.key) ?? 0;
const totalUsd = previousUsd + entry.usd;
totals.set(entry.key, totalUsd);
return { key: entry.key, totalUsd, previousUsd };
});
},
async read(key) {
return totals.get(key) ?? 0;
},
};
}

View File

@ -11,39 +11,12 @@ export interface SpendTotal {
previousUsd: number;
}
/** Process-local spend. `add` is atomic and idempotent on `callId`. */
/** Spend totals. `add` is idempotent on `callId`. */
export interface SpendLedger {
add(callId: string, entries: SpendEntry[]): Promise<SpendTotal[]>;
read(key: string): Promise<number>;
}
export class InMemorySpendLedger implements SpendLedger {
private readonly totals = new Map<string, number>();
private readonly appliedCallIds = new Set<string>();
async add(callId: string, entries: SpendEntry[]): Promise<SpendTotal[]> {
if (this.appliedCallIds.has(callId)) {
return entries.map((entry) => {
const totalUsd = this.totals.get(entry.key) ?? 0;
return { key: entry.key, totalUsd, previousUsd: totalUsd };
});
}
const totals = entries.map((entry) => {
const previousUsd = this.totals.get(entry.key) ?? 0;
const totalUsd = previousUsd + entry.usd;
this.totals.set(entry.key, totalUsd);
return { key: entry.key, totalUsd, previousUsd };
});
this.appliedCallIds.add(callId);
return totals;
}
async read(key: string): Promise<number> {
return this.totals.get(key) ?? 0;
}
}
export interface BudgetGuardrailOptions {
ledger: SpendLedger;
sessionId?: string;

View File

@ -0,0 +1,38 @@
import type { MigrationContext, ReversibleMigration } from '../migration-types';
/**
* Running budget totals and the model call ids that already added spend.
* A restart and a second main read the same rows.
*/
export class CreateAgentBudgetSpendTables1791209478293 implements ReversibleMigration {
async up({ schemaBuilder: { createTable, column }, tablePrefix, escape }: MigrationContext) {
await createTable('agent_budget_spend')
.withColumns(
column('key')
.varchar(128)
.primary.comment('Session key is the root thread id. Month key is {agentId}:{YYYY-MM}.'),
column('totalUsd')
.double.notNull.default(0)
.comment('Running USD estimate for this ledger key. A check rejects a negative total.'),
)
.withTimestamps.withCheck(
`CHK_${tablePrefix}agent_budget_spend_total_usd_non_negative`,
`${escape.columnName('totalUsd')} >= 0`,
);
await createTable('agent_budget_applied_call').withColumns(
column('callId').uuid.primary.comment(
'Model call id. A second insert of the same id does not add spend.',
),
column('createdAt')
.timestampTimezone()
.notNull.default('NOW()')
.comment('Time this call id was stored. A later prune can delete old rows.'),
);
}
async down({ schemaBuilder: { dropTable } }: MigrationContext) {
await dropTable('agent_budget_applied_call');
await dropTable('agent_budget_spend');
}
}

View File

@ -0,0 +1,151 @@
import { randomUUID } from 'node:crypto';
import { testDb, testModules } from '@n8n/backend-test-utils';
import type { Logger } from '@n8n/backend-common';
import { Container } from '@n8n/di';
import { mock } from 'vitest-mock-extended';
import { AgentBudgetAppliedCall } from '../entities/agent-budget-applied-call.entity';
import { AgentSpendLedger } from '../budget-guardrail';
import { AgentBudgetSpendRepository } from '../repositories/agent-budget-spend.repository';
describe('AgentBudgetSpendRepository', () => {
let repository: AgentBudgetSpendRepository;
beforeAll(async () => {
await testModules.loadModules(['agents']);
await testDb.init();
repository = Container.get(AgentBudgetSpendRepository);
});
beforeEach(async () => {
await repository.manager.delete(AgentBudgetAppliedCall, {});
await repository.delete({});
});
afterAll(async () => {
await testDb.terminate();
});
it('reads a missing key as 0', async () => {
await expect(repository.readTotal('missing-key')).resolves.toBe(0);
});
it('sums two calls on the same key', async () => {
const key = 'thread-1';
const first = await repository.applySpend(randomUUID(), [{ key, usd: 1.25 }]);
const second = await repository.applySpend(randomUUID(), [{ key, usd: 0.75 }]);
expect(first).toEqual([{ key, totalUsd: 1.25, previousUsd: 0 }]);
expect(second).toEqual([{ key, totalUsd: 2, previousUsd: 1.25 }]);
await expect(repository.readTotal(key)).resolves.toBe(2);
});
it('returns the exact pre-increment total at a float boundary', async () => {
const key = 'thread-1';
await repository.applySpend(randomUUID(), [{ key, usd: 8 }]);
const second = await repository.applySpend(randomUUID(), [{ key, usd: 0.2 }]);
// 8.2 - 0.2 is 7.999999999999999 in binary float. The alert check compares
// previousUsd with the alert line, so the value must be exactly 8.
expect(second).toEqual([{ key, totalUsd: 8.2, previousUsd: 8 }]);
});
it('adds repeated keys in one call in entry order', async () => {
const key = 'thread-1';
// The ledger contract allows repeated keys. Postgres rejects two rows with
// the same conflict target in one statement, so the upsert coalesces them.
const totals = await repository.applySpend(randomUUID(), [
{ key, usd: 1.25 },
{ key, usd: 0.75 },
]);
expect(totals).toEqual([
{ key, totalUsd: 1.25, previousUsd: 0 },
{ key, totalUsd: 2, previousUsd: 1.25 },
]);
await expect(repository.readTotal(key)).resolves.toBe(2);
});
it('counts both calls when two writers add to the same keys at the same time', async () => {
const sessionKey = 'thread-1';
const monthKey = 'agent-1:2026-10';
// Two mains can add spend for the same keys at the same time. The atomic
// upsert must count both. SQLite serializes the two transactions through
// the write connection; on Postgres they run as real concurrent transactions.
const [first, second] = await Promise.all([
repository.applySpend(randomUUID(), [
{ key: sessionKey, usd: 1.25 },
{ key: monthKey, usd: 1.25 },
]),
repository.applySpend(randomUUID(), [
{ key: sessionKey, usd: 0.75 },
{ key: monthKey, usd: 0.75 },
]),
]);
for (const key of [sessionKey, monthKey]) {
const results = [...first, ...second].filter((r) => r.key === key);
expect(results.filter((r) => r.previousUsd === 0)).toHaveLength(1);
expect(results.filter((r) => r.totalUsd === 2)).toHaveLength(1);
await expect(repository.readTotal(key)).resolves.toBe(2);
}
});
it('does not add spend again when callId repeats', async () => {
const callId = randomUUID();
const key = 'thread-1';
await repository.applySpend(callId, [{ key, usd: 1.5 }]);
const replay = await repository.applySpend(callId, [{ key, usd: 1.5 }]);
expect(replay).toEqual([{ key, totalUsd: 1.5, previousUsd: 1.5 }]);
await expect(repository.readTotal(key)).resolves.toBe(1.5);
});
it('updates the session key and the month key in entry order', async () => {
const sessionKey = 'thread-1';
const monthKey = 'agent-1:2026-10';
const totals = await repository.applySpend(randomUUID(), [
{ key: sessionKey, usd: 0.5 },
{ key: monthKey, usd: 0.5 },
]);
expect(totals).toEqual([
{ key: sessionKey, totalUsd: 0.5, previousUsd: 0 },
{ key: monthKey, totalUsd: 0.5, previousUsd: 0 },
]);
await expect(repository.readTotal(sessionKey)).resolves.toBe(0.5);
await expect(repository.readTotal(monthKey)).resolves.toBe(0.5);
});
it('keeps both totals when a new ledger reads the same database', async () => {
const callId = randomUUID();
const sessionKey = 'thread-1';
const monthKey = 'agent-1:2026-10';
const first = new AgentSpendLedger(repository, mock<Logger>());
await first.add(callId, [
{ key: sessionKey, usd: 1.5 },
{ key: monthKey, usd: 1.5 },
]);
const restarted = new AgentSpendLedger(repository, mock<Logger>());
await expect(restarted.read(sessionKey)).resolves.toBe(1.5);
await expect(restarted.read(monthKey)).resolves.toBe(1.5);
const replay = await restarted.add(callId, [
{ key: sessionKey, usd: 1.5 },
{ key: monthKey, usd: 1.5 },
]);
expect(replay).toEqual([
{ key: sessionKey, totalUsd: 1.5, previousUsd: 1.5 },
{ key: monthKey, totalUsd: 1.5, previousUsd: 1.5 },
]);
});
});

View File

@ -52,6 +52,7 @@ import { AgentTurnExecutionService } from '../agent-turn-execution.service';
import type { AgentExecutionStreamChunk } from '../types/agent-steering';
import type { AgentToolApprovalService } from '../agent-tool-approval.service';
import type { AgentValidationService } from '../agent-validation.service';
import { AgentSpendLedger } from '../budget-guardrail';
import type { Agent } from '../entities/agent.entity';
import type { AgentBackgroundJob } from '../entities/agent-background-job.entity';
import type { AgentExecutionThread } from '../entities/agent-execution-thread.entity';
@ -239,6 +240,16 @@ function makeService(sandboxEnabled = false) {
const wakeService = mock<AgentWakeService>();
Container.set(AgentWakeService, wakeService);
// The budget guardrail resolves its ledger from the container. Stub it with
// an empty ledger: reads stay under every cap, and `add` reports each entry's
// cost as the new total so alert-crossing tests can pick a crossing amount.
const spendLedger = mock<AgentSpendLedger>();
spendLedger.read.mockResolvedValue(0);
spendLedger.add.mockImplementation(async (_callId, entries) =>
entries.map((entry) => ({ key: entry.key, totalUsd: entry.usd, previousUsd: 0 })),
);
Container.set(AgentSpendLedger, spendLedger);
executionService.canUseDraftThread.mockResolvedValue(true);
executionService.findThreadById.mockResolvedValue({
id: 'thread-1',

View File

@ -0,0 +1,135 @@
import type { SpendEntry } from '@n8n/agents';
import type { Logger } from '@n8n/backend-common';
import { mock } from 'vitest-mock-extended';
import { AgentSpendLedger } from '../budget-guardrail';
import type { AgentBudgetSpendRepository } from '../repositories/agent-budget-spend.repository';
const entry = (key: string, usd: number): SpendEntry => ({ key, usd });
/** Repository whose writes and reads fail while `down` is true. */
function flakyRepository() {
const repository = mock<AgentBudgetSpendRepository>();
const state = { down: true };
repository.applySpend.mockImplementation(async (_callId: string, entries: SpendEntry[]) => {
if (state.down) throw new Error('database is down');
return entries.map((e) => ({ key: e.key, totalUsd: e.usd, previousUsd: 0 }));
});
repository.readTotal.mockImplementation(async (key: string) => {
if (state.down) throw new Error('database is down');
return key === 'recovered-key' ? 42 : 0;
});
return { repository, state };
}
describe('AgentSpendLedger', () => {
it('records spend when the database is up', async () => {
const repository = mock<AgentBudgetSpendRepository>();
repository.applySpend.mockResolvedValue([{ key: 'k', totalUsd: 1, previousUsd: 0 }]);
const ledger = new AgentSpendLedger(repository, mock<Logger>());
await expect(ledger.add('call-1', [entry('k', 1)])).resolves.toEqual([
{ key: 'k', totalUsd: 1, previousUsd: 0 },
]);
});
it('buffers the call and does not throw when the write fails', async () => {
const { repository } = flakyRepository();
const ledger = new AgentSpendLedger(repository, mock<Logger>());
await expect(ledger.add('call-1', [entry('k', 1.5)])).resolves.toEqual([
{ key: 'k', totalUsd: 1.5, previousUsd: 0 },
]);
});
it('reads the buffered totals when the read fails', async () => {
const { repository } = flakyRepository();
const ledger = new AgentSpendLedger(repository, mock<Logger>());
await ledger.add('call-1', [entry('session-1', 1), entry('agent-1:2026-10', 1)]);
await ledger.add('call-2', [entry('session-1', 0.5)]);
await expect(ledger.read('session-1')).resolves.toBe(1.5);
await expect(ledger.read('agent-1:2026-10')).resolves.toBe(1);
await expect(ledger.read('unknown')).resolves.toBe(0);
});
it('includes buffered spend when a read succeeds after a failed flush', async () => {
const repository = mock<AgentBudgetSpendRepository>();
// Writes keep failing (write-lock timeout); reads succeed.
repository.applySpend.mockRejectedValue(new Error('write lock timeout'));
repository.readTotal.mockResolvedValue(0.25);
const ledger = new AgentSpendLedger(repository, mock<Logger>());
await ledger.add('call-1', [entry('session-1', 1.5)]);
// 0.25 stored + 1.50 still buffered: a 1 USD session cap must stop the next call.
await expect(ledger.read('session-1')).resolves.toBe(1.75);
});
it('reconstructs the buffered previous total without float noise', async () => {
const { repository } = flakyRepository();
const ledger = new AgentSpendLedger(repository, mock<Logger>());
await ledger.add('call-1', [entry('k', 8)]);
// 8.2 - 0.2 is 7.999999999999999 in binary float; the alert check needs exactly 8.
await expect(ledger.add('call-2', [entry('k', 0.2)])).resolves.toEqual([
{ key: 'k', totalUsd: 8.2, previousUsd: 8 },
]);
});
it('flushes buffered calls in arrival order on the next add after recovery', async () => {
const { repository, state } = flakyRepository();
const ledger = new AgentSpendLedger(repository, mock<Logger>());
await ledger.add('call-1', [entry('k', 1)]);
await ledger.add('call-2', [entry('k', 2)]);
state.down = false;
await ledger.add('call-3', [entry('k', 3)]);
expect(repository.applySpend.mock.calls.map(([callId]) => callId)).toEqual([
'call-1', // first add, fails
'call-1', // flush attempt during the second add, fails
'call-2', // second add, fails
'call-1', // flush after recovery
'call-2',
'call-3',
]);
});
it('flushes the buffer on read once the database recovers', async () => {
const { repository, state } = flakyRepository();
const ledger = new AgentSpendLedger(repository, mock<Logger>());
await ledger.add('call-1', [entry('recovered-key', 1)]);
state.down = false;
await expect(ledger.read('recovered-key')).resolves.toBe(42);
expect(repository.applySpend.mock.calls.map(([callId]) => callId)).toEqual([
'call-1', // add, fails
'call-1', // flush during the read
]);
// The buffer is empty now: a second read must not replay the call.
await ledger.read('recovered-key');
expect(repository.applySpend).toHaveBeenCalledTimes(2);
});
it('drops the oldest buffered call when the buffer is full', async () => {
const { repository } = flakyRepository();
const logger = mock<Logger>();
const ledger = new AgentSpendLedger(repository, logger);
for (let i = 0; i < 10_001; i++) {
await ledger.add(`call-${i}`, [entry('k', 1)]);
}
// call-0 dropped, call-1..call-10000 buffered at 1 USD each
await expect(ledger.read('k')).resolves.toBe(10_000);
expect(logger.warn).toHaveBeenCalledWith(
'Budget spend buffer is full; dropped the oldest buffered call',
{ callId: 'call-0' },
);
});
});

View File

@ -1,5 +1,4 @@
import { budgetMonthKey, type SpendLedger } from '@n8n/agents';
import { vi } from 'vitest';
import { budgetMonthKey } from '@n8n/agents';
import { mock } from 'vitest-mock-extended';
import type { CollaborationService } from '@/collaboration/collaboration.service';
@ -61,14 +60,9 @@ describe('AgentsConfigController getValidation', () => {
describe('AgentsConfigController getBudget', () => {
it('returns the ledger total for the current month key', async () => {
const agentRepository = mock<AgentRepository>();
agentRepository.findByIdAndProjectId.mockResolvedValue({ id: 'agent-1' } as never);
const read = vi.fn<SpendLedger['read']>().mockResolvedValue(12.5);
const agentSpendLedger = {
ledger: {
read,
add: vi.fn<SpendLedger['add']>(),
},
};
agentRepository.existsByIdAndProjectId.mockResolvedValue(true);
const agentSpendLedger = mock<AgentSpendLedger>();
agentSpendLedger.read.mockResolvedValue(12.5);
const controller = new AgentsConfigController(
mock<AgentConfigService>(),
mock<AgentCustomToolsService>(),
@ -84,6 +78,6 @@ describe('AgentsConfigController getBudget', () => {
params: { projectId: 'project-1', agentId: 'agent-1' },
} as never),
).resolves.toEqual({ spentUsd: 12.5 });
expect(read).toHaveBeenCalledWith(budgetMonthKey('agent-1'));
expect(agentSpendLedger.read).toHaveBeenCalledWith(budgetMonthKey('agent-1'));
});
});

View File

@ -1,7 +1,35 @@
import { InMemorySpendLedger, type ExecutionOptions, type RunOptions } from '@n8n/agents';
import type { ExecutionOptions, RunOptions, SpendLedger } from '@n8n/agents';
import { withBudgetGuardrail } from '../budget-guardrail';
/** Ledger for these tests. `add` ignores a repeated `callId`. */
function spendLedger(): SpendLedger {
const totals = new Map<string, number>();
const appliedCallIds = new Set<string>();
return {
async add(callId, entries) {
if (appliedCallIds.has(callId)) {
return entries.map((entry) => {
const totalUsd = totals.get(entry.key) ?? 0;
return { key: entry.key, totalUsd, previousUsd: totalUsd };
});
}
appliedCallIds.add(callId);
return entries.map((entry) => {
const previousUsd = totals.get(entry.key) ?? 0;
const totalUsd = previousUsd + entry.usd;
totals.set(entry.key, totalUsd);
return { key: entry.key, totalUsd, previousUsd };
});
},
async read(key) {
return totals.get(key) ?? 0;
},
};
}
const base: RunOptions & ExecutionOptions = {
persistence: { threadId: 'thread-1', resourceId: 'user-1' },
};
@ -13,7 +41,7 @@ const saved = {
};
describe('withBudgetGuardrail', () => {
const ledger = new InMemorySpendLedger();
const ledger = spendLedger();
it('attaches nothing when budget is missing or turned off', () => {
expect(
@ -54,7 +82,7 @@ describe('withBudgetGuardrail', () => {
});
it('fires the attached onNotice once when the month total crosses the alert line', async () => {
const freshLedger = new InMemorySpendLedger();
const freshLedger = spendLedger();
const onNotice = vi.fn();
const attached = withBudgetGuardrail(base, {
ledger: freshLedger,

View File

@ -37,11 +37,11 @@ export class AgentsConfigController {
req: AuthenticatedRequest<{ projectId: string; agentId: string }>,
): Promise<AgentBudgetSpend> {
const { projectId, agentId } = req.params;
const agent = await this.agentRepository.findByIdAndProjectId(agentId, projectId);
if (!agent) {
const exists = await this.agentRepository.existsByIdAndProjectId(agentId, projectId);
if (!exists) {
throw new NotFoundError('Agent not found');
}
const spentUsd = await this.agentSpendLedger.ledger.read(budgetMonthKey(agentId));
const spentUsd = await this.agentSpendLedger.read(budgetMonthKey(agentId));
return { spentUsd };
}

View File

@ -73,6 +73,11 @@ export class AgentsModule implements ModuleInterface {
const { AgentHistoryRepository } = await import('./repositories/agent-history.repository.js');
Container.get(AgentHistoryRepository);
const { AgentBudgetSpendRepository } = await import(
'./repositories/agent-budget-spend.repository.js'
);
Container.get(AgentBudgetSpendRepository);
// Register the sandboxed runtime service (lazy — the V8 isolate is only
// created on first use, so this import has negligible startup cost).
const { AgentSecureRuntime } = await import('./runtime/agent-secure-runtime.js');
@ -257,6 +262,10 @@ export class AgentsModule implements ModuleInterface {
);
const { AgentTask } = await import('./entities/agent-task.entity.js');
const { AgentTaskRunLock } = await import('./entities/agent-task-run-lock.entity.js');
const { AgentBudgetSpend } = await import('./entities/agent-budget-spend.entity.js');
const { AgentBudgetAppliedCall } = await import(
'./entities/agent-budget-applied-call.entity.js'
);
const { AgentTaskSnapshot } = await import('./entities/agent-task-snapshot.entity.js');
const { AgentObservationEntity } = await import('./entities/agent-observation.entity.js');
const { AgentObservationCursorEntity } = await import(
@ -299,6 +308,8 @@ export class AgentsModule implements ModuleInterface {
AgentWorkflowDependency,
AgentTask,
AgentTaskRunLock,
AgentBudgetSpend,
AgentBudgetAppliedCall,
AgentTaskSnapshot,
AgentObservationEntity,
AgentObservationCursorEntity,

View File

@ -1,17 +1,113 @@
import {
createBudgetGuardrail,
InMemorySpendLedger,
type ExecutionOptions,
type RunOptions,
type SpendEntry,
type SpendLedger,
type SpendTotal,
} from '@n8n/agents';
import type { BudgetGuardrailConfig } from '@n8n/api-types';
import { Logger } from '@n8n/backend-common';
import { Container, Service } from '@n8n/di';
/** One ledger for this process. Queue mode and multi-main need a later adapter. */
import {
AgentBudgetSpendRepository,
previousTotalUsd,
} from './repositories/agent-budget-spend.repository';
/** Past this many buffered calls the oldest drops, so a long outage cannot grow memory. */
const PENDING_CALLS_CAP = 10_000;
/**
* Budget spend shared by every main. Reads and writes the database rows.
* A database error does not stop a turn: the spend is buffered in memory and
* retried on the next ledger call. Buffered spend is lost on a process crash,
* and an alert crossing during an outage can go unnoticed.
*/
@Service()
export class AgentSpendLedger {
readonly ledger: SpendLedger = new InMemorySpendLedger();
export class AgentSpendLedger implements SpendLedger {
/** Calls not yet recorded, in arrival order. */
private readonly pending = new Map<string, SpendEntry[]>();
constructor(
private readonly budgetSpendRepository: AgentBudgetSpendRepository,
private readonly logger: Logger,
) {}
async add(callId: string, entries: SpendEntry[]): Promise<SpendTotal[]> {
await this.flushPending();
try {
return await this.budgetSpendRepository.applySpend(callId, entries);
} catch (error) {
this.logger.warn('Failed to record budget spend; buffered the call for a later retry', {
callId,
error: error instanceof Error ? error.message : String(error),
});
this.buffer(callId, entries);
return this.bufferedTotals(entries);
}
}
async read(key: string): Promise<number> {
await this.flushPending();
try {
const stored = await this.budgetSpendRepository.readTotal(key);
return stored + this.bufferedTotalFor(key);
} catch (error) {
this.logger.warn('Failed to read budget spend; using the buffered totals', {
key,
error: error instanceof Error ? error.message : String(error),
});
return this.bufferedTotalFor(key);
}
}
/**
* Replays buffered calls in arrival order. `applySpend` is idempotent on
* `callId`, so a retry after an ambiguous failure cannot add spend twice.
* Stops at the first failure: the database is still down.
*/
private async flushPending(): Promise<void> {
for (const [callId, entries] of this.pending) {
try {
await this.budgetSpendRepository.applySpend(callId, entries);
} catch {
return;
}
this.pending.delete(callId);
}
}
private buffer(callId: string, entries: SpendEntry[]): void {
if (this.pending.size >= PENDING_CALLS_CAP && !this.pending.has(callId)) {
const oldest = this.pending.keys().next().value;
if (oldest !== undefined) {
this.pending.delete(oldest);
this.logger.warn('Budget spend buffer is full; dropped the oldest buffered call', {
callId: oldest,
});
}
}
this.pending.set(callId, entries);
}
/** Totals from the buffer alone. The database share is unknown during an outage. */
private bufferedTotals(entries: SpendEntry[]): SpendTotal[] {
return entries.map((entry) => {
const totalUsd = this.bufferedTotalFor(entry.key);
return { key: entry.key, totalUsd, previousUsd: previousTotalUsd(totalUsd, entry.usd) };
});
}
private bufferedTotalFor(key: string): number {
let totalUsd = 0;
for (const entries of this.pending.values()) {
for (const entry of entries) {
if (entry.key === key) totalUsd += entry.usd;
}
}
return totalUsd;
}
}
export interface BudgetAttachInput {
@ -81,7 +177,7 @@ export function withBudgetGuardrail<T extends RunOptions & ExecutionOptions>(
const limits = resolveBudgetLimits(input);
if (!limits) return options;
const ledger = input.ledger ?? Container.get(AgentSpendLedger).ledger;
const ledger = input.ledger ?? Container.get(AgentSpendLedger);
const hook = createBudgetGuardrail({
ledger,
...limits,

View File

@ -0,0 +1,16 @@
import { WithCreatedAt } from '@n8n/db';
import { Entity, PrimaryColumn } from '@n8n/typeorm';
/**
* One model call that already added budget spend.
* `createdAt` is the handle for a later prune. Totals stay on `agent_budget_spend`.
*/
@Entity({ name: 'agent_budget_applied_call' })
export class AgentBudgetAppliedCall extends WithCreatedAt {
/** Model call id. A second insert of the same id does not add spend. */
@PrimaryColumn({
type: 'uuid',
comment: 'Model call id. A second insert of the same id does not add spend.',
})
callId: string;
}

View File

@ -0,0 +1,22 @@
import { WithTimestamps } from '@n8n/db';
import { Column, Entity, PrimaryColumn } from '@n8n/typeorm';
/** Running USD total for one budget ledger key. */
@Entity({ name: 'agent_budget_spend' })
export class AgentBudgetSpend extends WithTimestamps {
/** Session key is the root thread id. Month key is {agentId}:{YYYY-MM}. */
@PrimaryColumn({
type: 'varchar',
length: 128,
comment: 'Session key is the root thread id. Month key is {agentId}:{YYYY-MM}.',
})
key: string;
/** Running USD estimate for this ledger key. A check rejects a negative total. */
@Column({
type: 'double precision',
default: 0,
comment: 'Running USD estimate for this ledger key. A check rejects a negative total.',
})
totalUsd: number;
}

View File

@ -0,0 +1,183 @@
import { BaseRepository, dbNowLiteral, TransactionRunner } from '@n8n/db';
import { Service } from '@n8n/di';
import { DataSource, In, type EntityManager } from '@n8n/typeorm';
import { isRecord } from '@n8n/utils/is-record';
import { UnexpectedError } from 'n8n-workflow';
import { AgentBudgetAppliedCall } from '../entities/agent-budget-applied-call.entity';
import { AgentBudgetSpend } from '../entities/agent-budget-spend.entity';
/** One ledger key and the USD amount to add. */
export interface AgentBudgetSpendEntry {
key: string;
usd: number;
}
/** Totals for one key after `applySpend`. */
export interface AgentBudgetSpendTotal {
key: string;
totalUsd: number;
previousUsd: number;
}
/**
* Reconstructs the pre-increment total from the new total. The subtraction
* keeps the as-applied value under concurrent writers, but binary float adds
* noise (8.2 - 0.2 = 7.999999999999999) that can re-fire an alert crossing.
* Twelve significant digits drop the noise and keep real cost precision.
*/
export function previousTotalUsd(totalUsd: number, usd: number): number {
return Number((totalUsd - usd).toPrecision(12));
}
/**
* Stores agent budget totals and the call ids that already added spend.
* `applySpend` inserts the call id and increments the totals in one transaction.
*/
@Service()
export class AgentBudgetSpendRepository extends BaseRepository<AgentBudgetSpend> {
constructor(dataSource: DataSource, transactionRunner: TransactionRunner) {
super(AgentBudgetSpend, dataSource.manager, transactionRunner);
}
/** Returns the stored total for a ledger key. A missing key is 0. */
async readTotal(key: string): Promise<number> {
const totals = await this.readTotalsWith(this.managerFor({}), [key]);
return totals.get(key) ?? 0;
}
/**
* Adds spend for one model call.
* A repeated callId returns the current totals and does not add again.
* The result list follows `entries` order.
*/
async applySpend(
callId: string,
entries: AgentBudgetSpendEntry[],
): Promise<AgentBudgetSpendTotal[]> {
return await this.runInTransaction({}, async (manager) => {
const inserted = await this.insertAppliedCall(manager, callId);
if (inserted) await this.incrementKeys(manager, entries);
const stored = await this.readTotalsWith(
manager,
entries.map((entry) => entry.key),
);
// Repeated keys share one stored total: walk it back in entry order.
const unapplied = new Map<string, number>();
for (const entry of entries) {
unapplied.set(entry.key, (unapplied.get(entry.key) ?? 0) + entry.usd);
}
return entries.map((entry) => {
const totalUsd = stored.get(entry.key);
if (totalUsd === undefined) {
if (inserted) throw new UnexpectedError('Budget spend row is missing after insert');
// Replay of a call whose spend never landed: the key has no row yet.
return { key: entry.key, totalUsd: 0, previousUsd: 0 };
}
if (!inserted) return { key: entry.key, totalUsd, previousUsd: totalUsd };
const left = unapplied.get(entry.key) ?? entry.usd;
const stillLeft = previousTotalUsd(left, entry.usd);
unapplied.set(entry.key, stillLeft);
return {
key: entry.key,
totalUsd: previousTotalUsd(totalUsd, stillLeft),
previousUsd: previousTotalUsd(totalUsd, left),
};
});
});
}
/** Returns true when this call id is new. A conflict means the call already added spend. */
private async insertAppliedCall(manager: EntityManager, callId: string): Promise<boolean> {
const table = this.quote(manager, this.tableName(manager, AgentBudgetAppliedCall));
const callIdColumn = this.quote(manager, 'callId');
const createdAtColumn = this.quote(manager, 'createdAt');
const now = dbNowLiteral(this.isPostgres(manager));
const sql =
`INSERT INTO ${table} (${callIdColumn}, ${createdAtColumn}) VALUES (:callId, ${now}) ` +
`ON CONFLICT (${callIdColumn}) DO NOTHING`;
const affected = await this.execute(manager, sql, { callId });
return affected > 0;
}
/**
* Adds each entry's `usd` to its key in one statement.
* The addition runs in SQL so two mains cannot drop each other's update.
* The SQLite driver discards RETURNING rows on INSERT, so the new totals are
* read back in `applySpend`.
*/
private async incrementKeys(
manager: EntityManager,
entries: AgentBudgetSpendEntry[],
): Promise<void> {
if (entries.length === 0) return;
const table = this.quote(manager, this.tableName(manager, AgentBudgetSpend));
const keyColumn = this.quote(manager, 'key');
const totalColumn = this.quote(manager, 'totalUsd');
const createdAtColumn = this.quote(manager, 'createdAt');
const updatedAtColumn = this.quote(manager, 'updatedAt');
const now = dbNowLiteral(this.isPostgres(manager));
// Coalesce repeated keys: Postgres rejects two rows with the same conflict
// target in one statement. Sort so concurrent writers lock rows in one order.
const byKey = new Map<string, number>();
for (const entry of entries) byKey.set(entry.key, (byKey.get(entry.key) ?? 0) + entry.usd);
const rows = [...byKey.entries()].sort(([a], [b]) => a.localeCompare(b));
const values = rows.map((_, i) => `(:key${i}, :usd${i}, ${now}, ${now})`).join(', ');
const sql =
`INSERT INTO ${table} (${keyColumn}, ${totalColumn}, ${createdAtColumn}, ${updatedAtColumn}) ` +
`VALUES ${values} ` +
`ON CONFLICT (${keyColumn}) DO UPDATE SET ${totalColumn} = ` +
`${table}.${totalColumn} + EXCLUDED.${totalColumn}, ${updatedAtColumn} = ${now}`;
const parameters: Record<string, unknown> = {};
rows.forEach(([key, usd], i) => {
parameters[`key${i}`] = key;
parameters[`usd${i}`] = usd;
});
await this.execute(manager, sql, parameters);
}
private async readTotalsWith(
manager: EntityManager,
keys: string[],
): Promise<Map<string, number>> {
if (keys.length === 0) return new Map();
const rows = await manager.findBy(AgentBudgetSpend, { key: In(keys) });
return new Map(rows.map((row) => [row.key, Number(row.totalUsd)]));
}
private async execute(
manager: EntityManager,
sql: string,
parameters: Record<string, unknown>,
): Promise<number> {
const queryRunner = manager.queryRunner;
if (!queryRunner) {
throw new UnexpectedError('Budget spend write has no query runner');
}
const [query, queryParameters] = manager.connection.driver.escapeQueryWithParameters(
sql,
parameters,
{},
);
const result: unknown = await queryRunner.query(query, queryParameters, true);
if (!isRecord(result)) {
throw new UnexpectedError('Budget spend write returned an unexpected result');
}
return typeof result.affected === 'number' ? result.affected : 0;
}
private tableName(
manager: EntityManager,
target: typeof AgentBudgetSpend | typeof AgentBudgetAppliedCall,
): string {
return manager.connection.getMetadata(target).tableName;
}
private quote(manager: EntityManager, identifier: string): string {
return manager.connection.driver.escape(identifier);
}
private isPostgres(manager: EntityManager): boolean {
return manager.connection.options.type === 'postgres';
}
}