diff --git a/packages/agentswarm/src/budget.test.ts b/packages/agentswarm/src/budget.test.ts new file mode 100644 index 0000000..0955ef9 --- /dev/null +++ b/packages/agentswarm/src/budget.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, it, vi } from "vitest"; +import { InMemoryBudgetLedger, SwarmError, withBudget } from "./index.js"; +import type { SwarmRunner } from "./index.js"; + +const DID = "did:coinpay:agent:abc"; +const ask = { messages: [{ role: "user" as const, content: "hi" }] }; + +function okRunner(): SwarmRunner { + return { + run: vi.fn(async (input) => ({ + threadId: "t", + messages: [...input.messages, { role: "assistant" as const, content: "done" }], + output: "done" + })) + }; +} + +describe("InMemoryBudgetLedger", () => { + it("credits and reports balance", async () => { + const ledger = new InMemoryBudgetLedger(); + expect(await ledger.balance(DID)).toBe(0); + await ledger.credit(DID, 100); + expect(await ledger.balance(DID)).toBe(100); + }); + + it("reserve holds funds and settle refunds the unused remainder", async () => { + const ledger = new InMemoryBudgetLedger(); + await ledger.credit(DID, 100); + const hold = await ledger.reserve(DID, 40); + expect(await ledger.balance(DID)).toBe(60); + await ledger.settle(hold, 10); + expect(await ledger.balance(DID)).toBe(90); // 60 + 30 refunded + }); + + it("reserve throws SwarmError 402 when short", async () => { + const ledger = new InMemoryBudgetLedger(); + await ledger.credit(DID, 5); + await expect(ledger.reserve(DID, 10)).rejects.toMatchObject({ status: 402 }); + }); + + it("release returns the full hold", async () => { + const ledger = new InMemoryBudgetLedger(); + await ledger.credit(DID, 50); + const hold = await ledger.reserve(DID, 50); + expect(await ledger.balance(DID)).toBe(0); + await ledger.release(hold); + expect(await ledger.balance(DID)).toBe(50); + }); +}); + +describe("withBudget", () => { + it("charges the actual cost on success", async () => { + const ledger = new InMemoryBudgetLedger(); + await ledger.credit(DID, 100); + const runner = withBudget({ runner: okRunner(), ledger, agentDid: DID, maxCost: 30, costOf: () => 12 }); + + const res = await runner.run(ask); + expect(res.output).toBe("done"); + expect(await ledger.balance(DID)).toBe(88); // 100 - 12 + }); + + it("rejects with 402 and never runs when out of budget", async () => { + const ledger = new InMemoryBudgetLedger(); + await ledger.credit(DID, 5); + const inner = okRunner(); + const runner = withBudget({ runner: inner, ledger, agentDid: DID, maxCost: 30 }); + + await expect(runner.run(ask)).rejects.toMatchObject({ status: 402 }); + expect(inner.run).not.toHaveBeenCalled(); + }); + + it("releases the hold (no charge) when the run throws", async () => { + const ledger = new InMemoryBudgetLedger(); + await ledger.credit(DID, 100); + const failing: SwarmRunner = { + run: vi.fn(async () => { + throw new Error("boom"); + }) + }; + const runner = withBudget({ runner: failing, ledger, agentDid: DID, maxCost: 30 }); + + await expect(runner.run(ask)).rejects.toThrow("boom"); + expect(await ledger.balance(DID)).toBe(100); // fully restored + }); + + it("re-exports SwarmError for status mapping", () => { + expect(new SwarmError("x", 402).status).toBe(402); + }); +}); diff --git a/packages/agentswarm/src/budget.ts b/packages/agentswarm/src/budget.ts new file mode 100644 index 0000000..a45e9d3 --- /dev/null +++ b/packages/agentswarm/src/budget.ts @@ -0,0 +1,111 @@ +import type { SwarmRunResult, SwarmRunner } from "./types.js"; +import { SwarmError } from "./types.js"; + +/** + * A swarm agent's identity. The DID mirrors @logicsrc/agentstack + * (`did:coinpay:agent:`), so an agent can carry an agentgit/CoinPay identity. + */ +export interface AgentIdentity { + did: string; + name?: string; +} + +/** + * Spend ledger for agents. Amounts are integer minor units (e.g. cents). A run + * reserves up to its budget, then settles to the actual cost — releasing the + * remainder. Mirrors the b1dz budget+ledger concept; back it with a real store + * for production. + */ +export interface BudgetLedger { + balance(agentDid: string): Promise; + credit(agentDid: string, amount: number): Promise; + /** Hold funds for an in-flight run; throws {@link SwarmError} 402 if short. */ + reserve(agentDid: string, amount: number): Promise; + /** Settle a hold to a final cost (clamped to the reserved amount). */ + settle(holdId: string, finalCost: number): Promise; + /** Release a hold without charging (e.g. the run failed). */ + release(holdId: string): Promise; +} + +/** In-memory reference {@link BudgetLedger} — good for tests and single-process hosts. */ +export class InMemoryBudgetLedger implements BudgetLedger { + private readonly balances = new Map(); + private readonly holds = new Map(); + private seq = 0; + + async balance(agentDid: string): Promise { + return this.balances.get(agentDid) ?? 0; + } + + async credit(agentDid: string, amount: number): Promise { + if (amount < 0) throw new Error("credit amount must be >= 0"); + const next = (this.balances.get(agentDid) ?? 0) + amount; + this.balances.set(agentDid, next); + return next; + } + + async reserve(agentDid: string, amount: number): Promise { + if (amount < 0) throw new Error("reserve amount must be >= 0"); + const available = this.balances.get(agentDid) ?? 0; + if (available < amount) { + throw new SwarmError( + `insufficient budget for ${agentDid}: need ${amount}, have ${available}`, + 402 + ); + } + this.balances.set(agentDid, available - amount); + const holdId = `hold_${(this.seq += 1)}`; + this.holds.set(holdId, { agentDid, amount }); + return holdId; + } + + async settle(holdId: string, finalCost: number): Promise { + const hold = this.holds.get(holdId); + if (!hold) throw new Error(`unknown hold: ${holdId}`); + this.holds.delete(holdId); + const cost = Math.max(0, Math.min(finalCost, hold.amount)); + const refund = hold.amount - cost; + if (refund > 0) await this.credit(hold.agentDid, refund); + } + + async release(holdId: string): Promise { + const hold = this.holds.get(holdId); + if (!hold) return; + this.holds.delete(holdId); + await this.credit(hold.agentDid, hold.amount); + } +} + +export interface BudgetRunnerOptions { + runner: SwarmRunner; + ledger: BudgetLedger; + /** DID charged for runs. */ + agentDid: string; + /** Funds reserved per run, in minor units. */ + maxCost: number; + /** Final cost from the result (e.g. derived from token usage). Default `maxCost`. */ + costOf?: (result: SwarmRunResult) => number; +} + +/** + * Wrap a {@link SwarmRunner} so each run reserves budget up front (throwing 402 + * when the agent is out of funds), then settles the actual cost afterward and + * releases the hold if the run throws. This is the per-agent spend guard for + * multi-step / swarm runs. + */ +export function withBudget(options: BudgetRunnerOptions): SwarmRunner { + const { runner, ledger, agentDid, maxCost, costOf } = options; + return { + async run(input) { + const holdId = await ledger.reserve(agentDid, maxCost); + try { + const result = await runner.run(input); + await ledger.settle(holdId, costOf ? costOf(result) : maxCost); + return result; + } catch (error) { + await ledger.release(holdId); + throw error; + } + } + }; +} diff --git a/packages/agentswarm/src/index.ts b/packages/agentswarm/src/index.ts index d986f93..699031f 100644 --- a/packages/agentswarm/src/index.ts +++ b/packages/agentswarm/src/index.ts @@ -17,6 +17,8 @@ export type { SwarmOptions, LLMRouterOptions } from "./swarm.js"; +export { InMemoryBudgetLedger, withBudget } from "./budget.js"; +export type { AgentIdentity, BudgetLedger, BudgetRunnerOptions } from "./budget.js"; export { createRubricRunner, createLLMJudge } from "./rubric.js"; export type { RubricEvaluation,