mirror of
https://github.com/n8n-io/n8n.git
synced 2026-10-17 05:29:44 +02:00
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:
parent
5f73c7cdcb
commit
a2fd4189f1
12
docs/generated/postgres-schema/README.md
generated
12
docs/generated/postgres-schema/README.md
generated
|
|
@ -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
|
||||
|
|
|
|||
38
docs/generated/postgres-schema/public.agent_budget_applied_call.md
generated
Normal file
38
docs/generated/postgres-schema/public.agent_budget_applied_call.md
generated
Normal 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)
|
||||
45
docs/generated/postgres-schema/public.agent_budget_spend.md
generated
Normal file
45
docs/generated/postgres-schema/public.agent_budget_spend.md
generated
Normal 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)
|
||||
12
docs/generated/sqlite-schema/README.md
generated
12
docs/generated/sqlite-schema/README.md
generated
|
|
@ -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
|
||||
|
|
|
|||
48
docs/generated/sqlite-schema/agent_budget_applied_call.md
generated
Normal file
48
docs/generated/sqlite-schema/agent_budget_applied_call.md
generated
Normal 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)
|
||||
53
docs/generated/sqlite-schema/agent_budget_spend.md
generated
Normal file
53
docs/generated/sqlite-schema/agent_budget_spend.md
generated
Normal 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)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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');
|
||||
}
|
||||
}
|
||||
|
|
@ -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 },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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' },
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
@ -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'));
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 };
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
@ -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;
|
||||
}
|
||||
|
|
@ -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';
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user