diff --git a/src/common/locks/budget-lock.decorator.spec.ts b/src/common/locks/budget-lock.decorator.spec.ts new file mode 100644 index 0000000..e7b257b --- /dev/null +++ b/src/common/locks/budget-lock.decorator.spec.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from 'vitest'; +import { Reflector } from '@nestjs/core'; +import { BUDGET_LOCK_KEY, UseBudgetLock } from './budget-lock.decorator'; + +describe('UseBudgetLock', () => { + it('sets the budget-lock metadata with default options', () => { + class Controller { + @UseBudgetLock() + allocate(): void {} + } + + const reflector = new Reflector(); + expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.allocate)).toEqual({}); + }); + + it('stores custom key resolver and ttl options', () => { + const keyResolver = () => 'budget:custom'; + + class Controller { + @UseBudgetLock({ key: keyResolver, ttl: 1000 }) + allocate(): void {} + } + + const reflector = new Reflector(); + expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.allocate)).toEqual({ + key: keyResolver, + ttl: 1000, + }); + }); + + it('is only applied to the decorated method', () => { + class Controller { + @UseBudgetLock() + allocate(): void {} + + list(): void {} + } + + const reflector = new Reflector(); + expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.allocate)).toEqual({}); + expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.list)).toBeUndefined(); + }); +}); diff --git a/src/common/locks/budget-lock.decorator.ts b/src/common/locks/budget-lock.decorator.ts new file mode 100644 index 0000000..0da467d --- /dev/null +++ b/src/common/locks/budget-lock.decorator.ts @@ -0,0 +1,37 @@ +import { applyDecorators, SetMetadata } from '@nestjs/common'; +import { Request } from 'express'; + +/** Reflector metadata key used by {@link BudgetLockInterceptor}. */ +export const BUDGET_LOCK_KEY = 'budgetLock'; + +export interface BudgetLockOptions { + /** + * Static lock key or a resolver producing the lock key from the incoming + * request. Defaults to `budget:{params.id}`. + */ + key?: string | ((request: Request) => string); + /** Lock time-to-live in milliseconds (defaults to 5000). */ + ttl?: number; +} + +/** + * Decorator that serializes concurrent mutations on the same budget resource + * by acquiring a Redis distributed lock around the handler. + * + * This prevents race conditions when multiple agents or requests simultaneously + * attempt to allocate, deduct, or modify budgets — ensuring atomic + * checks-and-balances for the same budget resource. + * + * Usage: + * ```ts + * @UseBudgetLock() + * @Post(':id/allocate') + * allocate(...) { ... } + * ``` + * + * When the lock cannot be acquired (another request is already mutating the + * same budget), the request fails with a `409 LOCK_ACQUISITION_FAILED` error. + */ +export function UseBudgetLock(options: BudgetLockOptions = {}): MethodDecorator { + return applyDecorators(SetMetadata(BUDGET_LOCK_KEY, options)); +} diff --git a/src/common/locks/budget-lock.interceptor.spec.ts b/src/common/locks/budget-lock.interceptor.spec.ts new file mode 100644 index 0000000..7b80642 --- /dev/null +++ b/src/common/locks/budget-lock.interceptor.spec.ts @@ -0,0 +1,136 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { ExecutionContext, CallHandler } from '@nestjs/common'; +import { Reflector } from '@nestjs/core'; +import { of, throwError } from 'rxjs'; +import { BudgetLockInterceptor } from './budget-lock.interceptor'; +import { RedisLock } from './redis-lock.util'; +import { BUDGET_LOCK_KEY, BudgetLockOptions } from './budget-lock.decorator'; +import { LockNotAcquiredException } from '../exceptions/domain.exception'; + +function buildContext(req: Record): ExecutionContext { + return { + switchToHttp: () => ({ getRequest: () => req }), + getHandler: () => handler, + getClass: () => Controller, + } as unknown as ExecutionContext; +} + +const handler = (): void => {}; +class Controller {} + +function decorate(options: BudgetLockOptions | undefined): void { + if (options === undefined) { + Reflect.deleteMetadata(BUDGET_LOCK_KEY, handler); + return; + } + Reflect.defineMetadata(BUDGET_LOCK_KEY, options, handler); +} + +describe('BudgetLockInterceptor', () => { + let reflector: Reflector; + let redisLock: { withLock: ReturnType }; + let interceptor: BudgetLockInterceptor; + + beforeEach(() => { + vi.clearAllMocks(); + Reflect.deleteMetadata(BUDGET_LOCK_KEY, handler); + reflector = new Reflector(); + redisLock = { + withLock: vi.fn().mockImplementation(async (_key: string, fn: () => Promise) => fn()), + }; + interceptor = new BudgetLockInterceptor(reflector, redisLock as unknown as RedisLock); + }); + + it('passes through when the handler is not decorated with @UseBudgetLock()', async () => { + decorate(undefined); + const ctx = buildContext({ params: { id: 'budget-1' } }); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + const result = await interceptor.intercept(ctx, next).toPromise(); + + expect(result).toBe('ok'); + expect(redisLock.withLock).not.toHaveBeenCalled(); + }); + + it('acquires a lock on the default budget key for decorated handlers', async () => { + decorate({}); + const ctx = buildContext({ params: { id: 'budget-1' } }); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + await interceptor.intercept(ctx, next).toPromise(); + + expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-1', expect.any(Function), 5000); + }); + + it('supports a custom static lock key', async () => { + decorate({ key: 'budget:custom' }); + const ctx = buildContext({ params: { id: 'budget-1' } }); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + await interceptor.intercept(ctx, next).toPromise(); + + expect(redisLock.withLock).toHaveBeenCalledWith('budget:custom', expect.any(Function), 5000); + }); + + it('supports a custom key resolver function receiving the request', async () => { + const resolver = vi.fn().mockReturnValue('budget:resolved'); + decorate({ key: resolver }); + const req = { params: { id: 'budget-1' } }; + const ctx = buildContext(req); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + await interceptor.intercept(ctx, next).toPromise(); + + expect(resolver).toHaveBeenCalledWith(req); + expect(redisLock.withLock).toHaveBeenCalledWith('budget:resolved', expect.any(Function), 5000); + }); + + it('falls back to body.budgetId when no route param is present', async () => { + decorate({}); + const ctx = buildContext({ params: {}, body: { budgetId: 'budget-9' } }); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + await interceptor.intercept(ctx, next).toPromise(); + + expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-9', expect.any(Function), 5000); + }); + + it('uses the configured ttl when provided', async () => { + decorate({ ttl: 250 }); + const ctx = buildContext({ params: { id: 'budget-1' } }); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + await interceptor.intercept(ctx, next).toPromise(); + + expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-1', expect.any(Function), 250); + }); + + it('rejects the request when the budget id cannot be resolved', async () => { + decorate({}); + const ctx = buildContext({ params: {}, body: {} }); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + await expect(interceptor.intercept(ctx, next).toPromise()).rejects.toThrow('budget id'); + expect(redisLock.withLock).not.toHaveBeenCalled(); + }); + + it('propagates lock acquisition failures to the caller', async () => { + decorate({}); + redisLock.withLock.mockRejectedValue(new LockNotAcquiredException('budget-1')); + const ctx = buildContext({ params: { id: 'budget-1' } }); + const next = { handle: () => of('ok') } as unknown as CallHandler; + + await expect(interceptor.intercept(ctx, next).toPromise()).rejects.toBeInstanceOf( + LockNotAcquiredException, + ); + }); + + it('propagates handler errors through the lock boundary', async () => { + decorate({}); + const ctx = buildContext({ params: { id: 'budget-1' } }); + const next = { handle: () => throwError(() => new Error('boom')) } as unknown as CallHandler; + + await expect(interceptor.intercept(ctx, next).toPromise()).rejects.toThrow('boom'); + expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-1', expect.any(Function), 5000); + }); +}); diff --git a/src/common/locks/budget-lock.interceptor.ts b/src/common/locks/budget-lock.interceptor.ts new file mode 100644 index 0000000..55f90fc --- /dev/null +++ b/src/common/locks/budget-lock.interceptor.ts @@ -0,0 +1,66 @@ +import { + BadRequestException, + CallHandler, + ExecutionContext, + Injectable, + NestInterceptor, +} from '@nestjs/common'; +import { Reflector } from '@nestjs/core'; +import { from, lastValueFrom, Observable, throwError } from 'rxjs'; +import { Request } from 'express'; +import { RedisLock, DEFAULT_LOCK_TTL_MS } from './redis-lock.util'; +import { BUDGET_LOCK_KEY, BudgetLockOptions } from './budget-lock.decorator'; + +/** + * Resolves the budget resource id from the request: `req.params.id` with a + * fallback to `req.body.budgetId`. + */ +function defaultBudgetKeyResolver(request: Request): string { + const params = request.params as { id?: string } | undefined; + const body = request.body as { budgetId?: string } | undefined; + const budgetId = params?.id ?? body?.budgetId; + if (!budgetId) { + throw new BadRequestException('A budget id is required to acquire the budget lock'); + } + return `budget:${budgetId}`; +} + +/** + * Global interceptor enforcing `@UseBudgetLock()`. For handlers decorated with + * the decorator it acquires a Redis distributed lock for the budget resource + * before invoking the handler and releases it afterwards (including on error), + * so concurrent mutations of the same budget are serialized across instances. + * + * Handlers without the decorator pass straight through untouched. + */ +@Injectable() +export class BudgetLockInterceptor implements NestInterceptor { + constructor( + private readonly reflector: Reflector, + private readonly redisLock: RedisLock, + ) {} + + intercept(context: ExecutionContext, next: CallHandler): Observable { + const options = this.reflector.getAllAndOverride(BUDGET_LOCK_KEY, [ + context.getHandler(), + context.getClass(), + ]); + + if (!options) { + return next.handle(); + } + + try { + const request = context.switchToHttp().getRequest(); + const resourceKey = + typeof options.key === 'function' + ? options.key(request) + : options.key ?? defaultBudgetKeyResolver(request); + const ttl = options.ttl ?? DEFAULT_LOCK_TTL_MS; + + return from(this.redisLock.withLock(resourceKey, () => lastValueFrom(next.handle()), ttl)); + } catch (error) { + return throwError(() => error); + } + } +} diff --git a/src/common/locks/locks.module.ts b/src/common/locks/locks.module.ts index 2110e16..7ea0d0a 100644 --- a/src/common/locks/locks.module.ts +++ b/src/common/locks/locks.module.ts @@ -6,13 +6,15 @@ import { RedisConfig } from '../../config/redis.config'; import { REDIS_CLIENT } from './locks.constants'; import { RedisLock } from './redis-lock.util'; import { AgentLockInterceptor } from './agent-lock.interceptor'; +import { BudgetLockInterceptor } from './budget-lock.interceptor'; /** * Global distributed-locking infrastructure. * * Provides a single shared ioredis client and the {@link RedisLock} service to - * every module, and registers the {@link AgentLockInterceptor} that enforces - * `@UseAgentLock()` on any decorated controller method. + * every module, and registers the {@link AgentLockInterceptor} and + * {@link BudgetLockInterceptor} that enforce `@UseAgentLock()` and + * `@UseBudgetLock()` on any decorated controller method. */ @Global() @Module({ @@ -27,6 +29,7 @@ import { AgentLockInterceptor } from './agent-lock.interceptor'; }, RedisLock, { provide: APP_INTERCEPTOR, useClass: AgentLockInterceptor }, + { provide: APP_INTERCEPTOR, useClass: BudgetLockInterceptor }, ], exports: [REDIS_CLIENT, RedisLock], }) diff --git a/src/modules/budgets/budget.controller.ts b/src/modules/budgets/budget.controller.ts index 5625fdc..3ad6704 100644 --- a/src/modules/budgets/budget.controller.ts +++ b/src/modules/budgets/budget.controller.ts @@ -19,6 +19,7 @@ import { } from '@nestjs/swagger'; import { UserRole } from '@prisma/client'; import { BudgetService } from './budget.service'; +import { UseBudgetLock } from '../../common/locks/budget-lock.decorator'; import { allocateBudgetSchema, AllocateBudgetInput, @@ -95,6 +96,7 @@ export class BudgetController { } @Patch(':id') + @UseBudgetLock() @Roles(UserRole.OWNER, UserRole.ADMIN, UserRole.FINANCE) @ApiOperation({ summary: 'Update a budget', @@ -116,6 +118,7 @@ export class BudgetController { } @Post(':id/allocate') + @UseBudgetLock() @Roles(UserRole.OWNER, UserRole.ADMIN, UserRole.FINANCE) @ApiOperation({ summary: 'Allocate funds from the parent budget to this child', @@ -139,6 +142,7 @@ export class BudgetController { } @Delete(':id') + @UseBudgetLock() @Roles(UserRole.OWNER, UserRole.ADMIN, UserRole.FINANCE) @ApiOperation({ summary: 'Delete (soft) a budget', diff --git a/src/modules/budgets/budget.service-lock.spec.ts b/src/modules/budgets/budget.service-lock.spec.ts new file mode 100644 index 0000000..170fad2 --- /dev/null +++ b/src/modules/budgets/budget.service-lock.spec.ts @@ -0,0 +1,299 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { Budget, Prisma } from '@prisma/client'; +import { BudgetService } from './budget.service'; +import { BudgetRepository } from './budget.repository'; +import { RedisLock } from '../../common/locks/redis-lock.util'; +import { EventBusService } from '../../events/event-bus.service'; +import { BudgetExceededException } from '../../common/exceptions/domain.exception'; + +const Decimal = Prisma.Decimal; + +/** Minimal mock Budget row builder. */ +function mockBudget( + overrides: Partial<{ + id: string; + organizationId: string; + parentBudgetId: string | null; + limitAmount: number; + spent: number; + name: string; + enabled: boolean; + }> = {}, +): Budget { + return { + id: overrides.id ?? 'budget-child-1', + organizationId: overrides.organizationId ?? 'org-1', + parentBudgetId: overrides.parentBudgetId ?? 'budget-parent-1', + agentId: null, + name: overrides.name ?? 'Child Budget', + currency: 'USDC', + limitAmount: new Decimal(overrides.limitAmount ?? 5000), + spent: new Decimal(overrides.spent ?? 0), + period: 'MONTHLY', + periodStart: new Date('2026-08-01T00:00:00.000Z'), + rollover: false, + enabled: overrides.enabled ?? true, + createdAt: new Date('2026-08-01T00:00:00.000Z'), + updatedAt: new Date('2026-08-01T00:00:00.000Z'), + deletedAt: null, + } as unknown as Budget; +} + +describe('BudgetService — distributed locking', () => { + let repository: BudgetRepository; + let redisLock: RedisLock; + let eventBus: EventBusService; + let service: BudgetService; + + beforeEach(() => { + vi.clearAllMocks(); + + repository = { + findById: vi.fn(), + update: vi.fn(), + incrementSpent: vi.fn(), + create: vi.fn(), + findManyAndCount: vi.fn(), + findChildren: vi.fn(), + softDelete: vi.fn(), + findEnabledByAgentId: vi.fn(), + } as unknown as BudgetRepository; + + redisLock = { + withLock: vi.fn((_key: unknown, fn: () => Promise) => fn()), + acquire: vi.fn(), + onModuleDestroy: vi.fn(), + } as unknown as RedisLock; + + eventBus = { + emit: vi.fn().mockResolvedValue(undefined), + } as unknown as EventBusService; + + service = new BudgetService(repository, eventBus, redisLock); + }); + + // ── allocate ── + + describe('allocate', () => { + it('acquires a lock on the parent budget key during allocation', async () => { + const parent = mockBudget({ + id: 'budget-parent-1', + limitAmount: 10000, + spent: 0, + parentBudgetId: null, + }); + const child = mockBudget({ id: 'budget-child-1', parentBudgetId: 'budget-parent-1' }); + + vi.mocked(repository.findById) + .mockResolvedValueOnce(child) // getOrThrow for child + .mockResolvedValueOnce(parent) // getOrThrow for parent + .mockResolvedValueOnce(parent); // re-read inside lock + vi.mocked(repository.update).mockResolvedValue({ ...child, limitAmount: new Decimal(6000) } as Budget); + + await service.allocate('org-1', 'user-1', 'budget-child-1', { amount: '1000' }); + + expect(redisLock.withLock).toHaveBeenCalledWith( + 'budget:budget-parent-1', + expect.any(Function), + ); + }); + + it('prevents race conditions when multiple agents allocate from the same parent', async () => { + const parent = mockBudget({ + id: 'budget-parent-1', + limitAmount: 10000, + spent: 0, + parentBudgetId: null, + }); + const child1 = mockBudget({ id: 'budget-child-1', parentBudgetId: 'budget-parent-1' }); + const child2 = mockBudget({ id: 'budget-child-2', parentBudgetId: 'budget-parent-1' }); + + vi.mocked(repository.findById).mockImplementation(async (_orgId: string, id: string) => { + if (id === 'budget-child-1') return child1; + if (id === 'budget-child-2') return child2; + return parent; + }); + + vi.mocked(repository.update).mockImplementation(async (id: string) => { + return { ...(id === 'budget-child-1' ? child1 : child2), limitAmount: new Decimal(5000) } as Budget; + }); + + const promise1 = service.allocate('org-1', 'user-1', 'budget-child-1', { amount: '4000' }); + const promise2 = service.allocate('org-1', 'user-2', 'budget-child-2', { amount: '4000' }); + + await Promise.all([promise1, promise2]); + + // Both allocations should succeed: 4000 + 4000 = 8000 ≤ 10000 + expect(repository.update).toHaveBeenCalledTimes(2); + }); + + it('rejects the second concurrent allocation when it would exceed the parent limit', async () => { + const parent = mockBudget({ + id: 'budget-parent-1', + limitAmount: 5000, + spent: 0, + parentBudgetId: null, + }); + const child = mockBudget({ id: 'budget-child-1', parentBudgetId: 'budget-parent-1' }); + + // First allocation uses the full parent remaining. + vi.mocked(repository.findById) + .mockResolvedValueOnce(child) // child getOrThrow + .mockResolvedValueOnce(parent) // parent getOrThrow + .mockResolvedValueOnce(parent); // re-read inside lock + + vi.mocked(repository.update).mockResolvedValue({ ...child, limitAmount: new Decimal(4500) } as Budget); + + await service.allocate('org-1', 'user-1', 'budget-child-1', { amount: '4500' }); + + // Now try to allocate more than the remaining 500. + const parentAfterFirst = { ...parent, spent: new Decimal(4500) }; + vi.mocked(repository.findById) + .mockResolvedValueOnce(child) // child getOrThrow + .mockResolvedValueOnce(parent) // parent getOrThrow + .mockResolvedValueOnce(parentAfterFirst); // re-read inside lock + + await expect( + service.allocate('org-1', 'user-2', 'budget-child-1', { amount: '1000' }), + ).rejects.toThrow(BudgetExceededException); + }); + + it('uses the lock with default TTL for allocate', async () => { + const parent = mockBudget({ id: 'p1', parentBudgetId: null, limitAmount: 10000 }); + const child = mockBudget({ id: 'c1', parentBudgetId: 'p1' }); + + vi.mocked(repository.findById) + .mockResolvedValueOnce(child) + .mockResolvedValueOnce(parent) + .mockResolvedValueOnce(parent); + vi.mocked(repository.update).mockResolvedValue({ ...child } as Budget); + + await service.allocate('org-1', 'user-1', 'c1', { amount: '100' }); + + expect(redisLock.withLock).toHaveBeenCalledWith( + 'budget:p1', + expect.any(Function), + ); + }); + }); + + // ── consume ── + + describe('consume', () => { + it('acquires a lock on the budget consume key during spend increment', async () => { + const budget = mockBudget({ id: 'budget-1', spent: 0, limitAmount: 10000 }); + vi.mocked(repository.incrementSpent).mockResolvedValue({ ...budget, spent: new Decimal(100) } as Budget); + + await service.consume('org-1', 'budget-1', 100); + + expect(redisLock.withLock).toHaveBeenCalledWith( + 'budget:consume:budget-1', + expect.any(Function), + ); + }); + + it('serializes concurrent consume calls for the same budget', async () => { + const budget = mockBudget({ id: 'budget-1', spent: 0, limitAmount: 10000 }); + let callCount = 0; + + vi.mocked(repository.incrementSpent).mockImplementation(async () => { + callCount++; + return { ...budget, spent: new Decimal(callCount * 100) } as Budget; + }); + + const promises = [ + service.consume('org-1', 'budget-1', 100), + service.consume('org-1', 'budget-1', 200), + service.consume('org-1', 'budget-1', 50), + ]; + + await Promise.all(promises); + + expect(redisLock.withLock).toHaveBeenCalledTimes(3); + expect(repository.incrementSpent).toHaveBeenCalledTimes(3); + }); + + it('uses separate lock keys for different budgets', async () => { + const budget1 = mockBudget({ id: 'b1', spent: 0 }); + const budget2 = mockBudget({ id: 'b2', spent: 0 }); + + vi.mocked(repository.incrementSpent) + .mockResolvedValueOnce({ ...budget1, spent: new Decimal(100) } as Budget) + .mockResolvedValueOnce({ ...budget2, spent: new Decimal(200) } as Budget); + + await service.consume('org-1', 'b1', 100); + await service.consume('org-1', 'b2', 200); + + expect(redisLock.withLock).toHaveBeenNthCalledWith(1, 'budget:consume:b1', expect.any(Function)); + expect(redisLock.withLock).toHaveBeenNthCalledWith(2, 'budget:consume:b2', expect.any(Function)); + }); + + it('emits budget consumed and warning events after locked spend increment', async () => { + const budget = mockBudget({ id: 'budget-1', spent: 7500, limitAmount: 10000 }); + vi.mocked(repository.incrementSpent).mockResolvedValue({ ...budget, spent: new Decimal(8500) } as Budget); + + await service.consume('org-1', 'budget-1', 1000); + + expect(eventBus.emit).toHaveBeenCalledWith( + 'budget.consumed', + expect.objectContaining({ budgetId: 'budget-1', amount: 1000 }), + expect.any(Object), + ); + // 85% utilisation → should also emit budget warning + expect(eventBus.emit).toHaveBeenCalledWith( + 'budget.warning', + expect.objectContaining({ budgetId: 'budget-1', utilisation: '0.8500' }), + expect.any(Object), + ); + }); + }); + + // ── assertWithinBudget ── + + describe('assertWithinBudget', () => { + it('acquires a lock on the budget check key during the pre-flight check', async () => { + const budget = mockBudget({ id: 'budget-1', spent: 1000, limitAmount: 10000 }); + vi.mocked(repository.findById).mockResolvedValue(budget); + + await service.assertWithinBudget('org-1', 'budget-1', 500); + + expect(redisLock.withLock).toHaveBeenCalledWith( + 'budget:check:budget-1', + expect.any(Function), + ); + }); + + it('throws BudgetExceededException within the lock when the spend would exceed the limit', async () => { + const budget = mockBudget({ id: 'budget-1', spent: 9500, limitAmount: 10000 }); + vi.mocked(repository.findById).mockResolvedValue(budget); + + await expect(service.assertWithinBudget('org-1', 'budget-1', 1000)).rejects.toThrow( + BudgetExceededException, + ); + + // Lock should still have been called (no leak) + expect(redisLock.withLock).toHaveBeenCalledWith( + 'budget:check:budget-1', + expect.any(Function), + ); + }); + + it('releases the lock even when BudgetExceededException is thrown', async () => { + const budget = mockBudget({ id: 'budget-1', spent: 9500, limitAmount: 10000 }); + vi.mocked(repository.findById).mockResolvedValue(budget); + + let lockReleased = false; + vi.mocked(redisLock.withLock).mockImplementation(async (_key: unknown, fn: () => Promise) => { + try { + return await fn(); + } finally { + lockReleased = true; + } + }); + + await expect(service.assertWithinBudget('org-1', 'budget-1', 1000)).rejects.toThrow(); + + expect(lockReleased).toBe(true); + }); + }); +}); diff --git a/src/modules/budgets/budget.service.ts b/src/modules/budgets/budget.service.ts index c67c8a7..8c06dc3 100644 --- a/src/modules/budgets/budget.service.ts +++ b/src/modules/budgets/budget.service.ts @@ -15,6 +15,7 @@ import { import { Paginated } from '../../common/interfaces/api-response.interface'; import { EventBusService } from '../../events/event-bus.service'; import { DomainEventName } from '../../events/event-names'; +import { RedisLock } from '../../common/locks/redis-lock.util'; const SORTABLE = ['createdAt', 'name', 'limitAmount', 'spent', 'period']; const Decimal = Prisma.Decimal; @@ -36,6 +37,7 @@ export class BudgetService { constructor( private readonly repository: BudgetRepository, private readonly eventBus: EventBusService, + private readonly redisLock: RedisLock, ) {} async create(organizationId: string, actorId: string, input: CreateBudgetInput): Promise { @@ -120,27 +122,35 @@ export class BudgetService { childId: string, input: AllocateBudgetInput, ) { + // Acquire a distributed lock keyed on the parent budget to serialize + // concurrent allocations from the same parent, preventing double-spend. const child = await this.getOrThrow(organizationId, childId); if (!child.parentBudgetId) { throw new ConflictException('Only a child budget can receive an allocation'); } const parent = await this.getOrThrow(organizationId, child.parentBudgetId); + const lockKey = `budget:${parent.id}`; const amount = new Decimal(input.amount); - if (amount.greaterThan(remaining(parent))) { - throw new BudgetExceededException('Allocation exceeds the parent budget remaining balance', { - parentRemaining: remaining(parent).toFixed(7), - requested: amount.toFixed(7), + + return this.redisLock.withLock(lockKey, async () => { + // Re-read the parent inside the lock to get a fresh snapshot. + const freshParent = await this.getOrThrow(organizationId, parent.id); + if (amount.greaterThan(remaining(freshParent))) { + throw new BudgetExceededException('Allocation exceeds the parent budget remaining balance', { + parentRemaining: remaining(freshParent).toFixed(7), + requested: amount.toFixed(7), + }); + } + const updated = await this.repository.update(childId, { + limitAmount: new Decimal(child.limitAmount).plus(amount), }); - } - const updated = await this.repository.update(childId, { - limitAmount: new Decimal(child.limitAmount).plus(amount), + await this.eventBus.emit( + DomainEventName.BudgetAllocated, + { budgetId: childId, parentBudgetId: freshParent.id, amount: input.amount }, + { organizationId, actorId, aggregateType: 'budget', aggregateId: childId }, + ); + return updated; }); - await this.eventBus.emit( - DomainEventName.BudgetAllocated, - { budgetId: childId, parentBudgetId: parent.id, amount: input.amount }, - { organizationId, actorId, aggregateType: 'budget', aggregateId: childId }, - ); - return updated; } /** @@ -149,43 +159,53 @@ export class BudgetService { * mutate state — call {@link consume} after the payment succeeds. */ async assertWithinBudget(organizationId: string, budgetId: string, amount: number) { - const budget = await this.getOrThrow(organizationId, budgetId); - const spendAfter = new Decimal(budget.spent).plus(amount); - if (spendAfter.greaterThan(budget.limitAmount)) { - await this.eventBus.emit( - DomainEventName.BudgetExceeded, - { budgetId, limit: budget.limitAmount.toFixed(7), attempted: spendAfter.toFixed(7) }, - { organizationId, aggregateType: 'budget', aggregateId: budgetId }, - ); - throw new BudgetExceededException('Transaction would exceed the budget limit', { - budgetId, - limit: budget.limitAmount.toFixed(7), - spent: budget.spent.toFixed(7), - attempted: amount, - }); - } - return budget; + // Lock the budget for the duration of the check-and-consume cycle so + // concurrent callers see a consistent view of the remaining headroom. + const lockKey = `budget:check:${budgetId}`; + return this.redisLock.withLock(lockKey, async () => { + const budget = await this.getOrThrow(organizationId, budgetId); + const spendAfter = new Decimal(budget.spent).plus(amount); + if (spendAfter.greaterThan(budget.limitAmount)) { + await this.eventBus.emit( + DomainEventName.BudgetExceeded, + { budgetId, limit: budget.limitAmount.toFixed(7), attempted: spendAfter.toFixed(7) }, + { organizationId, aggregateType: 'budget', aggregateId: budgetId }, + ); + throw new BudgetExceededException('Transaction would exceed the budget limit', { + budgetId, + limit: budget.limitAmount.toFixed(7), + spent: budget.spent.toFixed(7), + attempted: amount, + }); + } + return budget; + }); } /** Records realised spend after a payment completes; emits warnings near cap. */ async consume(organizationId: string, budgetId: string, amount: number) { - const budget = await this.repository.incrementSpent(budgetId, new Decimal(amount)); - const utilisation = new Decimal(budget.spent).dividedBy( - budget.limitAmount.isZero() ? new Decimal(1) : budget.limitAmount, - ); - await this.eventBus.emit( - DomainEventName.BudgetConsumed, - { budgetId, amount, spent: budget.spent.toFixed(7) }, - { organizationId, aggregateType: 'budget', aggregateId: budgetId }, - ); - if (utilisation.greaterThanOrEqualTo(0.8)) { + // Serialize spend increments on the same budget to prevent concurrent + // agents from overshooting the limit together. + const lockKey = `budget:consume:${budgetId}`; + return this.redisLock.withLock(lockKey, async () => { + const budget = await this.repository.incrementSpent(budgetId, new Decimal(amount)); + const utilisation = new Decimal(budget.spent).dividedBy( + budget.limitAmount.isZero() ? new Decimal(1) : budget.limitAmount, + ); await this.eventBus.emit( - DomainEventName.BudgetWarning, - { budgetId, utilisation: utilisation.toFixed(4) }, + DomainEventName.BudgetConsumed, + { budgetId, amount, spent: budget.spent.toFixed(7) }, { organizationId, aggregateType: 'budget', aggregateId: budgetId }, ); - } - return budget; + if (utilisation.greaterThanOrEqualTo(0.8)) { + await this.eventBus.emit( + DomainEventName.BudgetWarning, + { budgetId, utilisation: utilisation.toFixed(4) }, + { organizationId, aggregateType: 'budget', aggregateId: budgetId }, + ); + } + return budget; + }); } async remove(organizationId: string, actorId: string, id: string) { diff --git a/src/modules/health/health.controller.ts b/src/modules/health/health.controller.ts index d332a36..a08efaf 100644 --- a/src/modules/health/health.controller.ts +++ b/src/modules/health/health.controller.ts @@ -1,4 +1,4 @@ -import { Controller, Get, UseGuards } from '@nestjs/common'; +import { Controller, Get, HttpException, HttpStatus, UseGuards } from '@nestjs/common'; import { ApiOperation, ApiTags, @@ -6,6 +6,10 @@ import { ApiResponse, } from '@nestjs/swagger'; import { StellarHealthIndicator, StellarHealthReport } from './indicators/stellar.health'; +import { + DatabaseMigrationHealthIndicator, + MigrationHealthReport, +} from './indicators/database-migration.health'; import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard'; import { RolesGuard } from '../../common/guards/roles.guard'; import { Roles } from '../../common/decorators/roles.decorator'; @@ -14,7 +18,10 @@ import { UserRole } from '@prisma/client'; @ApiTags('health') @Controller('health') export class HealthController { - constructor(private readonly stellarHealthIndicator: StellarHealthIndicator) {} + constructor( + private readonly stellarHealthIndicator: StellarHealthIndicator, + private readonly databaseMigrationIndicator: DatabaseMigrationHealthIndicator, + ) {} @Get('stellar') @UseGuards(JwtAuthGuard, RolesGuard) @@ -33,4 +40,39 @@ export class HealthController { async checkStellarHealth(): Promise { return this.stellarHealthIndicator.checkHealth(); } + + @Get('database') + @UseGuards(JwtAuthGuard, RolesGuard) + @Roles(UserRole.OWNER, UserRole.ADMIN, UserRole.DEVELOPER, UserRole.AUDITOR) + @ApiBearerAuth('access-token') + @ApiOperation({ + summary: 'Database migration health check', + description: + 'Verifies that all Prisma migrations are applied and the database schema is up to date. ' + + 'Returns 503 when pending migrations are detected, which is critical for Kubernetes ' + + 'liveness and readiness probes.', + }) + @ApiResponse({ status: 200, description: 'All migrations applied, database schema is current' }) + @ApiResponse({ status: 401, description: 'Not authenticated' }) + @ApiResponse({ status: 403, description: 'Insufficient permissions' }) + @ApiResponse({ status: 503, description: 'Pending migrations detected or database unreachable' }) + async checkDatabaseMigrationHealth(): Promise { + const report = await this.databaseMigrationIndicator.checkHealth(); + + if (report.status === 'down' || report.status === 'degraded') { + // Throw to return 503 — Kubernetes probes will mark the pod as unhealthy. + throw new HttpException( + { + statusCode: 503, + message: report.status === 'down' + ? 'Database unreachable during migration health check' + : `${report.pendingMigrations} pending migration(s) detected — run 'prisma migrate deploy'`, + report, + }, + HttpStatus.SERVICE_UNAVAILABLE, + ); + } + + return report; + } } diff --git a/src/modules/health/health.module.ts b/src/modules/health/health.module.ts index f87f852..c0b143e 100644 --- a/src/modules/health/health.module.ts +++ b/src/modules/health/health.module.ts @@ -1,10 +1,11 @@ -import { Module } from '@nestjs/common'; -import { StellarHealthIndicator } from './indicators/stellar.health'; -import { HealthController } from './health.controller'; - -@Module({ - controllers: [HealthController], - providers: [StellarHealthIndicator], - exports: [StellarHealthIndicator], -}) +import { Module } from '@nestjs/common'; +import { StellarHealthIndicator } from './indicators/stellar.health'; +import { DatabaseMigrationHealthIndicator } from './indicators/database-migration.health'; +import { HealthController } from './health.controller'; + +@Module({ + controllers: [HealthController], + providers: [StellarHealthIndicator, DatabaseMigrationHealthIndicator], + exports: [StellarHealthIndicator, DatabaseMigrationHealthIndicator], +}) export class HealthModule {} diff --git a/src/modules/health/index.ts b/src/modules/health/index.ts index 0931ca3..b4f17c6 100644 --- a/src/modules/health/index.ts +++ b/src/modules/health/index.ts @@ -1,3 +1,4 @@ -export * from './health.module'; -export * from './health.controller'; +export * from './health.module'; +export * from './health.controller'; export * from './indicators/stellar.health'; +export * from './indicators/database-migration.health'; diff --git a/src/modules/health/indicators/database-migration.health.spec.ts b/src/modules/health/indicators/database-migration.health.spec.ts new file mode 100644 index 0000000..47623e3 --- /dev/null +++ b/src/modules/health/indicators/database-migration.health.spec.ts @@ -0,0 +1,167 @@ +import { beforeEach, describe, expect, it, vi, afterEach } from 'vitest'; +import { DatabaseMigrationHealthIndicator } from './database-migration.health'; +import { PrismaService } from '../../../database/prisma.service'; + +describe('DatabaseMigrationHealthIndicator', () => { + let prisma: { $queryRaw: ReturnType }; + let indicator: DatabaseMigrationHealthIndicator; + + beforeEach(() => { + vi.clearAllMocks(); + vi.unstubAllEnvs(); + + prisma = { + $queryRaw: vi.fn(), + }; + + indicator = new DatabaseMigrationHealthIndicator( + prisma as unknown as PrismaService, + ); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + }); + + describe('isEnabled', () => { + it('defaults to true when the env var is not set', () => { + expect(indicator.isEnabled).toBe(true); + }); + + it('returns true when env var is set to "true"', () => { + vi.stubEnv('HEALTH_CHECK_MIGRATIONS_ENABLED', 'true'); + expect(indicator.isEnabled).toBe(true); + }); + + it('returns false when env var is set to "false"', () => { + vi.stubEnv('HEALTH_CHECK_MIGRATIONS_ENABLED', 'false'); + expect(indicator.isEnabled).toBe(false); + }); + }); + + describe('checkHealth', () => { + it('returns UP when all migrations are applied', async () => { + prisma.$queryRaw.mockResolvedValue([ + { + migration_name: '20260801_add_budgets', + started_at: new Date('2026-08-01T00:00:00.000Z'), + finished_at: new Date('2026-08-01T00:00:05.000Z'), + applied_steps_count: 3, + checksum: 'abc123', + rolled_back_at: null, + }, + { + migration_name: '20260701_init', + started_at: new Date('2026-07-01T00:00:00.000Z'), + finished_at: new Date('2026-07-01T00:00:10.000Z'), + applied_steps_count: 5, + checksum: 'def456', + rolled_back_at: null, + }, + ]); + + const report = await indicator.checkHealth(); + + expect(report.status).toBe('up'); + expect(report.pendingMigrations).toBe(0); + expect(report.lastMigrationName).toBe('20260801_add_budgets'); + expect(report.lastMigrationApplied).toBe('2026-08-01T00:00:05.000Z'); + expect(report.error).toBeUndefined(); + }); + + it('returns DEGRADED when there are pending (unfinished) migrations', async () => { + prisma.$queryRaw.mockResolvedValue([ + { + migration_name: '20260815_add_locks', + started_at: new Date('2026-08-15T00:00:00.000Z'), + finished_at: null, // pending! + applied_steps_count: 0, + checksum: 'ghi789', + rolled_back_at: null, + }, + { + migration_name: '20260801_add_budgets', + started_at: new Date('2026-08-01T00:00:00.000Z'), + finished_at: new Date('2026-08-01T00:00:05.000Z'), + applied_steps_count: 3, + checksum: 'abc123', + rolled_back_at: null, + }, + ]); + + const report = await indicator.checkHealth(); + + expect(report.status).toBe('degraded'); + expect(report.pendingMigrations).toBe(1); + expect(report.lastMigrationName).toBe('20260801_add_budgets'); + }); + + it('counts multiple pending migrations correctly', async () => { + prisma.$queryRaw.mockResolvedValue([ + { + migration_name: '20260901_pending_a', + started_at: new Date(), + finished_at: null, + applied_steps_count: 0, + checksum: 'aaa', + rolled_back_at: null, + }, + { + migration_name: '20260915_pending_b', + started_at: new Date(), + finished_at: null, + applied_steps_count: 0, + checksum: 'bbb', + rolled_back_at: null, + }, + { + migration_name: '20260801_done', + started_at: new Date(), + finished_at: new Date(), + applied_steps_count: 2, + checksum: 'ccc', + rolled_back_at: null, + }, + ]); + + const report = await indicator.checkHealth(); + + expect(report.status).toBe('degraded'); + expect(report.pendingMigrations).toBe(2); + expect(report.lastMigrationName).toBe('20260801_done'); + }); + + it('returns DOWN when the database query fails', async () => { + prisma.$queryRaw.mockRejectedValue(new Error('Connection refused')); + + const report = await indicator.checkHealth(); + + expect(report.status).toBe('down'); + expect(report.pendingMigrations).toBe(-1); + expect(report.error).toContain('Connection refused'); + expect(report.lastMigrationName).toBeNull(); + expect(report.lastMigrationApplied).toBeNull(); + }); + + it('returns UP with null last migration when no migrations exist', async () => { + prisma.$queryRaw.mockResolvedValue([]); + + const report = await indicator.checkHealth(); + + expect(report.status).toBe('up'); + expect(report.pendingMigrations).toBe(0); + expect(report.lastMigrationName).toBeNull(); + expect(report.lastMigrationApplied).toBeNull(); + }); + + it('includes a timestamp in the report', async () => { + prisma.$queryRaw.mockResolvedValue([]); + + const report = await indicator.checkHealth(); + + expect(report.timestamp).toBeDefined(); + // Should be a valid ISO string + expect(new Date(report.timestamp).toISOString()).toBe(report.timestamp); + }); + }); +}); diff --git a/src/modules/health/indicators/database-migration.health.ts b/src/modules/health/indicators/database-migration.health.ts new file mode 100644 index 0000000..7ecb3ed --- /dev/null +++ b/src/modules/health/indicators/database-migration.health.ts @@ -0,0 +1,113 @@ +import { Injectable, Logger } from '@nestjs/common'; +import { PrismaService } from '../../../database/prisma.service'; + +export interface MigrationHealthReport { + status: 'up' | 'down' | 'degraded'; + timestamp: string; + pendingMigrations: number; + lastMigrationName: string | null; + lastMigrationApplied: string | null; + error?: string; +} + +/** + * A Prisma migration row as returned by the raw query against + * `_prisma_migrations`. + */ +interface PrismaMigrationRow { + migration_name: string; + started_at: Date; + finished_at: Date | null; + applied_steps_count: number; + checksum: string; + rolled_back_at: Date | null; +} + +/** + * Health indicator that verifies whether the database schema is fully up to + * date with the Prisma migration history. Queries the `_prisma_migrations` + * table to detect pending (unapplied) migrations, returning a 503-style + * response when the schema is out of date. + * + * This is critical for Kubernetes liveness/readiness probes — an application + * running against an unmigrated database will produce runtime errors on + * agent request execution. + * + * The check is gated behind a configurable toggle + * (`HEALTH_CHECK_MIGRATIONS_ENABLED`) so it can be disabled in local + * development while remaining active in production. + */ +@Injectable() +export class DatabaseMigrationHealthIndicator { + private readonly logger = new Logger(DatabaseMigrationHealthIndicator.name); + + constructor(private readonly prisma: PrismaService) {} + + /** + * Returns the configured migration check toggle. When `false`, the health + * endpoint should skip this indicator entirely. + */ + get isEnabled(): boolean { + // Default to true — migrations should be checked in production. + // Set HEALTH_CHECK_MIGRATIONS_ENABLED=false to disable in dev. + const raw = process.env.HEALTH_CHECK_MIGRATIONS_ENABLED; + if (raw === undefined) return true; + return raw !== 'false'; + } + + /** + * Queries the `_prisma_migrations` table to determine whether any + * migrations have not yet been applied. Returns a structured report + * suitable for inclusion in the health endpoint response. + */ + async checkHealth(): Promise { + try { + // Raw query is safe here — _prisma_migrations is a Prisma-managed + // internal table with a fixed schema. + const rows = await this.prisma.$queryRaw` + SELECT + migration_name, + started_at, + finished_at, + applied_steps_count, + checksum, + rolled_back_at + FROM _prisma_migrations + WHERE rolled_back_at IS NULL + ORDER BY started_at DESC + `; + + const pendingMigrations = rows.filter((r) => r.finished_at === null).length; + const completedMigrations = rows.filter((r) => r.finished_at !== null); + const lastApplied = completedMigrations[0] ?? null; + + let status: 'up' | 'down' | 'degraded' = 'up'; + if (pendingMigrations > 0) { + status = 'degraded'; + this.logger.warn( + `Database has ${pendingMigrations} pending migration(s) — schema may be out of date`, + ); + } + + return { + status, + timestamp: new Date().toISOString(), + pendingMigrations, + lastMigrationName: lastApplied?.migration_name ?? null, + lastMigrationApplied: lastApplied?.finished_at?.toISOString() ?? null, + }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + this.logger.error(`Migration health check failed: ${message}`); + + return { + status: 'down', + timestamp: new Date().toISOString(), + pendingMigrations: -1, + lastMigrationName: null, + lastMigrationApplied: null, + error: message, + }; + } + } +}