feat(budgets): add Redis-backed distributed mutex guard for spending limits

src/modules/budgets/budget.module.ts

@@ -6,6 +6,7 @@ import { BudgetReservationService } from './services/budget-reservation.service'
 import { PolicyEvaluatorService } from './services/policy-evaluator.service';
 import { RedisLock } from '../../common/locks/redis-lock.util';
 import { RollingWindowBudgetService } from './services/rolling-window-budget.service';
+import { SpendingLimitGuardService } from './services/spending-limit-guard.service';
 
 /**
  * Budget module. Exports the service so the transactions pipeline can enforce
@@ -20,14 +21,7 @@ import { RollingWindowBudgetService } from './services/rolling-window-budget.ser
  */
 @Module({
   controllers: [BudgetController],
-  providers: [
-    BudgetService,
-    BudgetRepository,
-    BudgetReservationService,
-    RedisLock,
-    PolicyEvaluatorService,
-    RollingWindowBudgetService,
-  ],
-  exports: [BudgetService, PolicyEvaluatorService, RollingWindowBudgetService],
+  providers: [BudgetService, BudgetRepository, PolicyEvaluatorService, RollingWindowBudgetService, SpendingLimitGuardService],
+  exports: [BudgetService, PolicyEvaluatorService, RollingWindowBudgetService, SpendingLimitGuardService],
 })
 export class BudgetModule {}

src/modules/budgets/services/spending-limit-guard.service.spec.ts

@@ -0,0 +1,76 @@
+import { describe, expect, it, vi, beforeEach } from 'vitest';
+import { SpendingLimitGuardService } from './spending-limit-guard.service';
+import { RedisLock } from '../../../common/locks/redis-lock.util';
+import { BudgetService } from '../budget.service';
+
+describe('SpendingLimitGuardService', () => {
+  let service: SpendingLimitGuardService;
+  let redisLock: RedisLock;
+  let budgets: BudgetService;
+
+  beforeEach(() => {
+    redisLock = {
+      withLock: vi.fn(),
+    } as unknown as RedisLock;
+
+    budgets = {
+      assertWithinBudget: vi.fn(),
+      consume: vi.fn(),
+    } as unknown as BudgetService;
+
+    service = new SpendingLimitGuardService(redisLock, budgets);
+  });
+
+  it('should acquire lock and call assertWithinBudget + consume', async () => {
+    const fakeBudget = { id: 'b1', spent: 0, limitAmount: { toFixed: () => '100' } };
+    vi.mocked(redisLock.withLock).mockImplementation(async (_key, fn) => fn());
+    vi.mocked(budgets.assertWithinBudget).mockResolvedValue(undefined as never);
+    vi.mocked(budgets.consume).mockResolvedValue(fakeBudget as never);
+
+    const result = await service.guardAndConsume('org1', 'b1', 50);
+
+    expect(redisLock.withLock).toHaveBeenCalledWith(
+      'budget:guard:b1',
+      expect.any(Function),
+      5000,
+    );
+    expect(budgets.assertWithinBudget).toHaveBeenCalledWith('org1', 'b1', 50);
+    expect(budgets.consume).toHaveBeenCalledWith('org1', 'b1', 50);
+    expect(result).toBe(fakeBudget);
+  });
+
+  it('should not call consume if assertWithinBudget throws', async () => {
+    vi.mocked(redisLock.withLock).mockImplementation(async (_key, fn) => fn());
+    vi.mocked(budgets.assertWithinBudget).mockRejectedValueOnce(
+      new Error('Budget exceeded'),
+    );
+
+    await expect(service.guardAndConsume('org1', 'b1', 200)).rejects.toThrow('Budget exceeded');
+    expect(budgets.consume).not.toHaveBeenCalled();
+  });
+
+  it('should use custom TTL when provided', async () => {
+    vi.mocked(redisLock.withLock).mockImplementation(async (_key, fn) => fn());
+    vi.mocked(budgets.assertWithinBudget).mockResolvedValue(undefined as never);
+    vi.mocked(budgets.consume).mockResolvedValue({} as never);
+
+    await service.guardAndConsume('org1', 'b1', 10, 10_000);
+
+    expect(redisLock.withLock).toHaveBeenCalledWith(
+      'budget:guard:b1',
+      expect.any(Function),
+      10_000,
+    );
+  });
+
+  it('should propagate lock acquisition failure', async () => {
+    vi.mocked(redisLock.withLock).mockRejectedValue(
+      new Error('Lock not acquired'),
+    );
+
+    await expect(service.guardAndConsume('org1', 'b1', 50)).rejects.toThrow(
+      'Lock not acquired',
+    );
+    expect(budgets.assertWithinBudget).not.toHaveBeenCalled();
+  });
+});

src/modules/budgets/services/spending-limit-guard.service.ts

@@ -0,0 +1,53 @@
+import { Injectable } from '@nestjs/common';
+import { RedisLock } from '../../../common/locks/redis-lock.util';
+import { BudgetService } from '../budget.service';
+
+/**
+ * Wraps the budget check + consume cycle in a single Redis distributed lock
+ * to eliminate the TOCTOU race between assertWithinBudget (step 5) and
+ * consume (step 7) in the transaction pipeline.
+ *
+ * Without this guard, concurrent transactions can both pass the headroom
+ * check before either consumes, overshooting the budget limit.
+ *
+ * The lock key is `lock:budget:guard:{budgetId}` — distinct from the
+ * individual check/consume keys in BudgetService.
+ */
+@Injectable()
+export class SpendingLimitGuardService {
+  constructor(
+    private readonly redisLock: RedisLock,
+    private readonly budgets: BudgetService,
+  ) {}
+
+  /**
+   * Atomically validates headroom and debits the budget under a single
+   * distributed lock. Throws BudgetExceededException if the spend would
+   * breach the limit (no state is mutated in that case).
+   *
+   * @param organizationId - Owning organization.
+   * @param budgetId - Budget to check and consume from.
+   * @param amount - Amount to spend.
+   * @param ttlMs - Lock TTL in milliseconds (default 5 000).
+   * @returns The budget after consumption.
+   */
+  async guardAndConsume(
+    organizationId: string,
+    budgetId: string,
+    amount: number,
+    ttlMs = 5_000,
+  ) {
+    const lockKey = `budget:guard:${budgetId}`;
+    return this.redisLock.withLock(
+      lockKey,
+      async () => {
+        // Step 1: Validate headroom (same logic as assertWithinBudget).
+        await this.budgets.assertWithinBudget(organizationId, budgetId, amount);
+
+        // Step 2: Debit the budget (same logic as consume).
+        return this.budgets.consume(organizationId, budgetId, amount);
+      },
+      ttlMs,
+    );
+  }
+}

src/modules/transactions/transaction.service.ts

@@ -16,6 +16,7 @@ import { WalletService, toNetworkName } from '../wallets/wallet.service';
 import { PolicyService } from '../policies/policy.service';
 import { RiskService } from '../risk/risk.service';
 import { BudgetService } from '../budgets/budget.service';
+import { SpendingLimitGuardService } from '../budgets/services/spending-limit-guard.service';
 import { StellarService } from '../stellar/stellar.service';
 import { AgentService } from '../agents/agent.service';
 import { TransactionIntent } from '../policies/policy.types';
@@ -65,6 +66,7 @@ export class TransactionService {
     private readonly policies: PolicyService,
     private readonly risk: RiskService,
     private readonly budgets: BudgetService,
+    private readonly spendingGuard: SpendingLimitGuardService,
     private readonly stellar: StellarService,
     private readonly eventBus: EventBusService,
     private readonly prisma: PrismaService,
@@ -98,46 +100,41 @@ export class TransactionService {
       { actorId },
     );
 
-    // 5. Reserve the budget headroom (authoritative — persisted atomically under
-    //    a lock, so concurrent agent requests cannot overspend a single budget).
-    if (input.budgetId) {
-      await this.budgets.reserve(organizationId, input.budgetId, amount);
-    }
-
     const requiresApproval = policyResult.requiresApproval || !assessment.canAutoExecute;
 
-    // 6. Persist the transaction row. If persistence fails, return the reserved
-    //    headroom so the budget is not silently eaten.
-    let transaction: Transaction;
-    try {
-      transaction = await this.repository.create({
-        organization: { connect: { id: organizationId } },
-        wallet: { connect: { id: wallet.id } },
-        ...(input.agentId ? { agent: { connect: { id: input.agentId } } } : {}),
-        ...(policyResult.matchedPolicyId
-          ? { policy: { connect: { id: policyResult.matchedPolicyId } } }
-          : {}),
-        ...(input.budgetId ? { budget: { connect: { id: input.budgetId } } } : {}),
-        asset: input.asset,
-        amount: new Decimal(input.amount),
-        senderAddress: wallet.stellarAddress,
-        recipientAddress: input.recipientAddress,
-        memo: memoValue,
-        purpose: input.purpose,
-        status: TransactionStatus.DRAFT,
-        riskScore: assessment.score,
-        riskBand: assessment.band,
-        requiresApproval,
-        metadata: input.metadata as Prisma.InputJsonValue,
-      });
-    } catch (error) {
-      if (input.budgetId) {
-        await this.budgets
-          .release(organizationId, input.budgetId, amount)
-          .catch(() => undefined);
+    // 5. Budget check — atomically guard+consume for auto-executable txns,
+    //    or check-only for approval-required txns (consume deferred to execute).
+    let budgetConsumed = false;
+    if (input.budgetId) {
+      if (requiresApproval) {
+        await this.budgets.assertWithinBudget(organizationId, input.budgetId, amount);
+      } else {
+        await this.spendingGuard.guardAndConsume(organizationId, input.budgetId, amount);
+        budgetConsumed = true;
       }
-      throw error;
     }
+
+    // 6. Persist the transaction row.
+    const transaction = await this.repository.create({
+      organization: { connect: { id: organizationId } },
+      wallet: { connect: { id: wallet.id } },
+      ...(input.agentId ? { agent: { connect: { id: input.agentId } } } : {}),
+      ...(policyResult.matchedPolicyId
+        ? { policy: { connect: { id: policyResult.matchedPolicyId } } }
+        : {}),
+      ...(input.budgetId ? { budget: { connect: { id: input.budgetId } } } : {}),
+      asset: input.asset,
+      amount: new Decimal(input.amount),
+      senderAddress: wallet.stellarAddress,
+      recipientAddress: input.recipientAddress,
+      memo: memoValue,
+      purpose: input.purpose,
+      status: TransactionStatus.DRAFT,
+      riskScore: assessment.score,
+      riskBand: assessment.band,
+      requiresApproval,
+      metadata: input.metadata as Prisma.InputJsonValue,
+    });
     await this.eventBus.emit(
       DomainEventName.TransactionCreated,
       { transactionId: transaction.id, amount: input.amount, riskBand: assessment.band },
@@ -156,7 +153,7 @@ export class TransactionService {
       };
     }
 
-    const executed = await this.execute(organizationId, transaction.id, actorId);
+    const executed = await this.execute(organizationId, transaction.id, actorId, budgetConsumed);
     return { transaction: executed, requiresApproval: false, risk: assessment };
   }
 
@@ -165,7 +162,7 @@ export class TransactionService {
    * auto-executable transactions and by the approvals module once a proposal has
    * gathered the required approvals.
    */
-  async execute(organizationId: string, transactionId: string, actorId?: string): Promise<Transaction> {
+  async execute(organizationId: string, transactionId: string, actorId?: string, budgetConsumed = false): Promise<Transaction> {
     const tx = await this.getOrThrow(organizationId, transactionId);
     if (
       tx.status === TransactionStatus.COMPLETED ||
@@ -202,9 +199,7 @@ export class TransactionService {
         confirmationCount: result.ledger ? 1 : 0,
       });
 
-      if (result.successful && tx.budgetId) {
-        // Settles the reservation booked at creation — spend was already
-        // reserved, so this only records the event and warns near the cap.
+      if (result.successful && tx.budgetId && !budgetConsumed) {
         await this.budgets.consume(organizationId, tx.budgetId, Number(tx.amount));
       }