diff --git a/docs/generated/postgres-schema/README.md b/docs/generated/postgres-schema/README.md index 04f070b7a20..abf8d2c0dce 100644 --- a/docs/generated/postgres-schema/README.md +++ b/docs/generated/postgres-schema/README.md @@ -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 diff --git a/docs/generated/postgres-schema/public.agent_budget_applied_call.md b/docs/generated/postgres-schema/public.agent_budget_applied_call.md new file mode 100644 index 00000000000..8094627205e --- /dev/null +++ b/docs/generated/postgres-schema/public.agent_budget_applied_call.md @@ -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) diff --git a/docs/generated/postgres-schema/public.agent_budget_spend.md b/docs/generated/postgres-schema/public.agent_budget_spend.md new file mode 100644 index 00000000000..8884d60713e --- /dev/null +++ b/docs/generated/postgres-schema/public.agent_budget_spend.md @@ -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) diff --git a/docs/generated/sqlite-schema/README.md b/docs/generated/sqlite-schema/README.md index b12a6b13672..d935caa4bc5 100644 --- a/docs/generated/sqlite-schema/README.md +++ b/docs/generated/sqlite-schema/README.md @@ -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 diff --git a/docs/generated/sqlite-schema/agent_budget_applied_call.md b/docs/generated/sqlite-schema/agent_budget_applied_call.md new file mode 100644 index 00000000000..3aebbe0ccb2 --- /dev/null +++ b/docs/generated/sqlite-schema/agent_budget_applied_call.md @@ -0,0 +1,48 @@ +# agent_budget_applied_call + +## Description + +
+Table Definition + +```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'))) +``` + +
+ +## 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) diff --git a/docs/generated/sqlite-schema/agent_budget_spend.md b/docs/generated/sqlite-schema/agent_budget_spend.md new file mode 100644 index 00000000000..03284271705 --- /dev/null +++ b/docs/generated/sqlite-schema/agent_budget_spend.md @@ -0,0 +1,53 @@ +# agent_budget_spend + +## Description + +
+Table Definition + +```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)) +``` + +
+ +## 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) diff --git a/packages/@n8n/agents/src/index.ts b/packages/@n8n/agents/src/index.ts index e347347d184..44bdd8d922b 100644 --- a/packages/@n8n/agents/src/index.ts +++ b/packages/@n8n/agents/src/index.ts @@ -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, diff --git a/packages/@n8n/agents/src/runtime/__tests__/agent-runtime.test.ts b/packages/@n8n/agents/src/runtime/__tests__/agent-runtime.test.ts index 73f70d18ea8..91bdd09977d 100644 --- a/packages/@n8n/agents/src/runtime/__tests__/agent-runtime.test.ts +++ b/packages/@n8n/agents/src/runtime/__tests__/agent-runtime.test.ts @@ -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, }); diff --git a/packages/@n8n/agents/src/runtime/guardrails/__tests__/budget-guardrail.test.ts b/packages/@n8n/agents/src/runtime/guardrails/__tests__/budget-guardrail.test.ts index 5b28141bc97..6b8b26a5e4b 100644 --- a/packages/@n8n/agents/src/runtime/guardrails/__tests__/budget-guardrail.test.ts +++ b/packages/@n8n/agents/src/runtime/guardrails/__tests__/budget-guardrail.test.ts @@ -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 }); diff --git a/packages/@n8n/agents/src/runtime/guardrails/__tests__/spend-ledger.ts b/packages/@n8n/agents/src/runtime/guardrails/__tests__/spend-ledger.ts new file mode 100644 index 00000000000..5eb9a9d5e01 --- /dev/null +++ b/packages/@n8n/agents/src/runtime/guardrails/__tests__/spend-ledger.ts @@ -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(); + const appliedCallIds = new Set(); + + 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; + }, + }; +} diff --git a/packages/@n8n/agents/src/runtime/guardrails/budget-guardrail.ts b/packages/@n8n/agents/src/runtime/guardrails/budget-guardrail.ts index e93aca2404b..2d6fa4827a3 100644 --- a/packages/@n8n/agents/src/runtime/guardrails/budget-guardrail.ts +++ b/packages/@n8n/agents/src/runtime/guardrails/budget-guardrail.ts @@ -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; read(key: string): Promise; } -export class InMemorySpendLedger implements SpendLedger { - private readonly totals = new Map(); - private readonly appliedCallIds = new Set(); - - async add(callId: string, entries: SpendEntry[]): Promise { - 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 { - return this.totals.get(key) ?? 0; - } -} - export interface BudgetGuardrailOptions { ledger: SpendLedger; sessionId?: string; diff --git a/packages/@n8n/db/src/migrations/common/1791209478293-CreateAgentBudgetSpendTables.ts b/packages/@n8n/db/src/migrations/common/1791209478293-CreateAgentBudgetSpendTables.ts new file mode 100644 index 00000000000..297eaf2867f --- /dev/null +++ b/packages/@n8n/db/src/migrations/common/1791209478293-CreateAgentBudgetSpendTables.ts @@ -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'); + } +} diff --git a/packages/cli/src/modules/agents/__tests__/agent-budget-spend.repository.integration.test.ts b/packages/cli/src/modules/agents/__tests__/agent-budget-spend.repository.integration.test.ts new file mode 100644 index 00000000000..be7a0a639a2 --- /dev/null +++ b/packages/cli/src/modules/agents/__tests__/agent-budget-spend.repository.integration.test.ts @@ -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()); + + await first.add(callId, [ + { key: sessionKey, usd: 1.5 }, + { key: monthKey, usd: 1.5 }, + ]); + + const restarted = new AgentSpendLedger(repository, mock()); + 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 }, + ]); + }); +}); diff --git a/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts b/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts index 025c6ded9c0..e4c8011d97f 100644 --- a/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts @@ -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(); 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(); + 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', diff --git a/packages/cli/src/modules/agents/__tests__/agent-spend-ledger.test.ts b/packages/cli/src/modules/agents/__tests__/agent-spend-ledger.test.ts new file mode 100644 index 00000000000..6c0a9e9b78a --- /dev/null +++ b/packages/cli/src/modules/agents/__tests__/agent-spend-ledger.test.ts @@ -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(); + 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(); + repository.applySpend.mockResolvedValue([{ key: 'k', totalUsd: 1, previousUsd: 0 }]); + const ledger = new AgentSpendLedger(repository, mock()); + + 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()); + + 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()); + + 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(); + // 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()); + + 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()); + + 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()); + + 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()); + + 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(); + 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' }, + ); + }); +}); diff --git a/packages/cli/src/modules/agents/__tests__/agents-config.controller.test.ts b/packages/cli/src/modules/agents/__tests__/agents-config.controller.test.ts index 3905b2db7cd..62ec902c08a 100644 --- a/packages/cli/src/modules/agents/__tests__/agents-config.controller.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agents-config.controller.test.ts @@ -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.findByIdAndProjectId.mockResolvedValue({ id: 'agent-1' } as never); - const read = vi.fn().mockResolvedValue(12.5); - const agentSpendLedger = { - ledger: { - read, - add: vi.fn(), - }, - }; + agentRepository.existsByIdAndProjectId.mockResolvedValue(true); + const agentSpendLedger = mock(); + agentSpendLedger.read.mockResolvedValue(12.5); const controller = new AgentsConfigController( mock(), mock(), @@ -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')); }); }); diff --git a/packages/cli/src/modules/agents/__tests__/budget-guardrail.test.ts b/packages/cli/src/modules/agents/__tests__/budget-guardrail.test.ts index cf76709ab2d..063d35fb411 100644 --- a/packages/cli/src/modules/agents/__tests__/budget-guardrail.test.ts +++ b/packages/cli/src/modules/agents/__tests__/budget-guardrail.test.ts @@ -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(); + const appliedCallIds = new Set(); + + 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, diff --git a/packages/cli/src/modules/agents/agents-config.controller.ts b/packages/cli/src/modules/agents/agents-config.controller.ts index 1aee052abfa..e74f20a294c 100644 --- a/packages/cli/src/modules/agents/agents-config.controller.ts +++ b/packages/cli/src/modules/agents/agents-config.controller.ts @@ -37,11 +37,11 @@ export class AgentsConfigController { req: AuthenticatedRequest<{ projectId: string; agentId: string }>, ): Promise { 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 }; } diff --git a/packages/cli/src/modules/agents/agents.module.ts b/packages/cli/src/modules/agents/agents.module.ts index 6518008df9c..3792ca16023 100644 --- a/packages/cli/src/modules/agents/agents.module.ts +++ b/packages/cli/src/modules/agents/agents.module.ts @@ -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, diff --git a/packages/cli/src/modules/agents/budget-guardrail.ts b/packages/cli/src/modules/agents/budget-guardrail.ts index 46e9f58c9ce..9638416d85a 100644 --- a/packages/cli/src/modules/agents/budget-guardrail.ts +++ b/packages/cli/src/modules/agents/budget-guardrail.ts @@ -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(); + + constructor( + private readonly budgetSpendRepository: AgentBudgetSpendRepository, + private readonly logger: Logger, + ) {} + + async add(callId: string, entries: SpendEntry[]): Promise { + 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 { + 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 { + 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( 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, diff --git a/packages/cli/src/modules/agents/entities/agent-budget-applied-call.entity.ts b/packages/cli/src/modules/agents/entities/agent-budget-applied-call.entity.ts new file mode 100644 index 00000000000..64c7953d8bf --- /dev/null +++ b/packages/cli/src/modules/agents/entities/agent-budget-applied-call.entity.ts @@ -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; +} diff --git a/packages/cli/src/modules/agents/entities/agent-budget-spend.entity.ts b/packages/cli/src/modules/agents/entities/agent-budget-spend.entity.ts new file mode 100644 index 00000000000..c494af508dc --- /dev/null +++ b/packages/cli/src/modules/agents/entities/agent-budget-spend.entity.ts @@ -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; +} diff --git a/packages/cli/src/modules/agents/repositories/agent-budget-spend.repository.ts b/packages/cli/src/modules/agents/repositories/agent-budget-spend.repository.ts new file mode 100644 index 00000000000..57314436af6 --- /dev/null +++ b/packages/cli/src/modules/agents/repositories/agent-budget-spend.repository.ts @@ -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 { + 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 { + 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 { + 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(); + 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 { + 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 { + 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(); + 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 = {}; + 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> { + 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, + ): Promise { + 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'; + } +}