feat: implement mail bucket entries and usage tracking

migrations/20260623120000-create-mail-bucket-entries.js

@@ -0,0 +1,53 @@
+'use strict';
+
+const TABLE_NAME = 'mail_bucket_entries';
+
+/** @type {import('sequelize-cli').Migration} */
+module.exports = {
+  async up(queryInterface, Sequelize) {
+    await queryInterface.createTable(TABLE_NAME, {
+      id: {
+        type: Sequelize.UUID,
+        defaultValue: Sequelize.UUIDV4,
+        primaryKey: true,
+        allowNull: false,
+      },
+      mail_account_id: {
+        type: Sequelize.UUID,
+        allowNull: false,
+        references: { model: 'mail_accounts', key: 'id' },
+        onUpdate: 'CASCADE',
+        onDelete: 'CASCADE',
+      },
+      entry_key: {
+        type: Sequelize.STRING(255),
+        allowNull: false,
+        unique: true,
+      },
+      bridge_entry_id: {
+        type: Sequelize.STRING(24),
+        allowNull: false,
+      },
+      size: {
+        type: Sequelize.BIGINT,
+        allowNull: false,
+      },
+      created_at: {
+        type: Sequelize.DATE,
+        allowNull: false,
+        defaultValue: Sequelize.fn('now'),
+      },
+      updated_at: {
+        type: Sequelize.DATE,
+        allowNull: false,
+        defaultValue: Sequelize.fn('now'),
+      },
+    });
+
+    await queryInterface.addIndex(TABLE_NAME, ['mail_account_id']);
+  },
+
+  async down(queryInterface) {
+    await queryInterface.dropTable(TABLE_NAME);
+  },
+};

src/modules/account/repositories/address.repository.ts

@@ -13,6 +13,7 @@ import { MailProviderAccountModel } from '../models/mail-provider-account.model.
 const MAX_BATCH_LOOKUP = 50;
 
 export interface ProviderAccountBucketContext {
+  mailAccountId: string;
   userUuid: string;
   networkBucketId: string | null;
 }
@@ -110,6 +111,7 @@ export class AddressRepository {
     if (!model?.account) return null;
 
     return {
+      mailAccountId: model.account.id,
       userUuid: model.account.userId,
       networkBucketId: model.networkBucketId,
     };
@@ -132,6 +134,7 @@ export class AddressRepository {
     if (!link?.address?.account) return null;
 
     return {
+      mailAccountId: link.address.account.id,
       userUuid: link.address.account.userId,
       networkBucketId: link.address.networkBucketId,
     };

src/modules/email/email.module.ts

@@ -1,5 +1,5 @@
 import { Module } from '@nestjs/common';
-import { BridgeModule } from '../infrastructure/bridge/bridge.module.js';
+import { MailUsageModule } from '../usage/mail-usage.module.js';
 import { JmapModule } from '../infrastructure/jmap/jmap.module.js';
 import { SmtpModule } from '../infrastructure/smtp/smtp.module.js';
 import { ProvisioningModule } from '../provisioning/provisioning.module.js';
@@ -8,7 +8,7 @@ import { EmailService } from './email.service.js';
 import { Reflector } from '@nestjs/core';
 
 @Module({
-  imports: [JmapModule, SmtpModule, ProvisioningModule, BridgeModule],
+  imports: [JmapModule, SmtpModule, ProvisioningModule, MailUsageModule],
   controllers: [EmailController],
   providers: [EmailService, Reflector],
 })

src/modules/email/email.service.ts

@@ -6,7 +6,7 @@ import {
 } from '@nestjs/common';
 import { ConfigService } from '@nestjs/config';
 import { AccountService } from '../account/account.service.js';
-import { BridgeClient } from '../infrastructure/bridge/bridge.service.js';
+import { MailUsageService } from '../usage/mail-usage.service.js';
 import { MailProvider } from './mail-provider.port.js';
 import type {
   DraftEmailDto,
@@ -60,7 +60,7 @@ export class EmailService {
     private readonly accountService: AccountService,
     private readonly smtp: StalwartSmtpService,
     private readonly configService: ConfigService,
-    private readonly bridge: BridgeClient,
+    private readonly usage: MailUsageService,
   ) {}
 
   getMailboxes(userEmail: string): Promise<Mailbox[]> {
@@ -276,11 +276,11 @@ export class EmailService {
     }
 
     try {
-      await this.bridge.deleteBucketEntry(
-        context.userUuid,
-        context.networkBucketId,
+      await this.usage.releaseStoredMessage({
+        userUuid: context.userUuid,
+        bucketId: context.networkBucketId,
         entryKey,
-      );
+      });
     } catch (error) {
       this.logger.warn(
         `Failed to release quota entry '${entryKey}' for '${userEmail}': ${(error as Error).message}`,

src/modules/infrastructure/bridge/bridge.service.ts

@@ -7,7 +7,11 @@ import {
 import { ConfigService } from '@nestjs/config';
 import { JwtService } from '@nestjs/jwt';
 import { Client } from 'undici';
-import type { BucketEntry, MailBucket } from './bridge.types.js';
+import type {
+  BucketEntry,
+  MailBucket,
+  UserSpaceSnapshot,
+} from './bridge.types.js';
 
 @Injectable()
 export class BridgeClient implements OnModuleInit, OnModuleDestroy {
@@ -101,7 +105,6 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy {
   async createBucketEntry(
     userUuid: string,
     bucketId: string,
-    key: string,
     size: number,
   ): Promise<BucketEntry> {
     const token = this.signGatewayToken(userUuid);
@@ -114,14 +117,14 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy {
         accept: 'application/json',
         authorization: `Bearer ${token}`,
       },
-      body: JSON.stringify({ key, size }),
+      body: JSON.stringify({ size }),
     });
 
     const text = await body.text();
 
     if (statusCode !== 200) {
       throw new BridgeApiError(
-        `Failed to create bucket entry '${key}' on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`,
+        `Failed to create bucket entry on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`,
         statusCode,
         text,
       );
@@ -133,13 +136,13 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy {
   async deleteBucketEntry(
     userUuid: string,
     bucketId: string,
-    key: string,
-  ): Promise<void> {
+    entryId: string,
+  ): Promise<UserSpaceSnapshot> {
     const token = this.signGatewayToken(userUuid);
 
     const { statusCode, body } = await this.httpClient.request({
       method: 'DELETE',
-      path: `${this.basePath}/v2/gateway/users/${encodeURIComponent(userUuid)}/buckets/${encodeURIComponent(bucketId)}/entries/${encodeURIComponent(key)}`,
+      path: `${this.basePath}/v2/gateway/users/${encodeURIComponent(userUuid)}/buckets/${encodeURIComponent(bucketId)}/entries/${encodeURIComponent(entryId)}`,
       headers: {
         accept: 'application/json',
         authorization: `Bearer ${token}`,
@@ -150,11 +153,13 @@ export class BridgeClient implements OnModuleInit, OnModuleDestroy {
 
     if (statusCode !== 200) {
       throw new BridgeApiError(
-        `Failed to delete bucket entry '${key}' on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`,
+        `Failed to delete bucket entry '${entryId}' on bucket '${bucketId}' for user '${userUuid}': HTTP ${statusCode}`,
         statusCode,
         text,
       );
     }
+
+    return JSON.parse(text) as UserSpaceSnapshot;
   }
 
   private signGatewayToken(userUuid: string): string {

src/modules/stalwart-events/stalwart-events.module.ts

@@ -1,12 +1,12 @@
 import { Module } from '@nestjs/common';
 import { AccountModule } from '../account/account.module.js';
-import { BridgeModule } from '../infrastructure/bridge/bridge.module.js';
+import { MailUsageModule } from '../usage/mail-usage.module.js';
 import { StalwartEventsController } from './stalwart-events.controller.js';
 import { StalwartEventsAuthGuard } from './stalwart-events-auth.guard.js';
 import { StalwartEventsService } from './stalwart-events.service.js';
 
 @Module({
-  imports: [AccountModule, BridgeModule],
+  imports: [AccountModule, MailUsageModule],
   controllers: [StalwartEventsController],
   providers: [StalwartEventsAuthGuard, StalwartEventsService],
 })

src/modules/stalwart-events/stalwart-events.service.ts

@@ -1,6 +1,6 @@
 import { Injectable, Logger } from '@nestjs/common';
 import { AccountService } from '../account/account.service.js';
-import { BridgeClient } from '../infrastructure/bridge/bridge.service.js';
+import { MailUsageService } from '../usage/mail-usage.service.js';
 import type {
   StalwartEvent,
   StalwartWebhookPayload,
@@ -12,7 +12,7 @@ export class StalwartEventsService {
 
   constructor(
     private readonly accounts: AccountService,
-    private readonly bridge: BridgeClient,
+    private readonly usage: MailUsageService,
   ) {}
 
   async handleBatch(payload: StalwartWebhookPayload): Promise<void> {
@@ -57,16 +57,12 @@ export class StalwartEventsService {
       return;
     }
 
-    const { totalUsedSpaceBytes } = await this.bridge.createBucketEntry(
-      context.userUuid,
-      context.networkBucketId,
+    await this.usage.trackStoredMessage({
+      mailAccountId: context.mailAccountId,
+      userUuid: context.userUuid,
+      bucketId: context.networkBucketId,
       entryKey,
       size,
-    );
-
-    this.logger.log(
-      { entryKey, size, totalUsedSpaceBytes, type: event.type },
-      'Created bucket entry for ingested message',
-    );
+    });
   }
 }

src/modules/usage/domain/mail-bucket-entry.domain.ts

@@ -0,0 +1,27 @@
+export interface MailBucketEntryAttributes {
+  id: string;
+  mailAccountId: string;
+  entryKey: string;
+  bridgeEntryId: string;
+  size: number;
+  createdAt: Date;
+  updatedAt: Date;
+}
+
+export class MailBucketEntry {
+  readonly id!: string;
+  readonly mailAccountId!: string;
+  readonly entryKey!: string;
+  readonly bridgeEntryId!: string;
+  readonly size!: number;
+  readonly createdAt!: Date;
+  readonly updatedAt!: Date;
+
+  private constructor(attributes: MailBucketEntryAttributes) {
+    Object.assign(this, attributes);
+  }
+
+  static build(attributes: MailBucketEntryAttributes): MailBucketEntry {
+    return new MailBucketEntry(attributes);
+  }
+}

src/modules/usage/mail-usage.module.ts

@@ -0,0 +1,13 @@
+import { Module } from '@nestjs/common';
+import { SequelizeModule } from '@nestjs/sequelize';
+import { BridgeModule } from '../infrastructure/bridge/bridge.module.js';
+import { MailBucketEntryModel } from './models/mail-bucket-entry.model.js';
+import { MailBucketEntryRepository } from './repositories/mail-bucket-entry.repository.js';
+import { MailUsageService } from './mail-usage.service.js';
+
+@Module({
+  imports: [SequelizeModule.forFeature([MailBucketEntryModel]), BridgeModule],
+  providers: [MailBucketEntryRepository, MailUsageService],
+  exports: [MailUsageService],
+})
+export class MailUsageModule {}

src/modules/usage/mail-usage.service.ts

@@ -0,0 +1,94 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { BridgeClient } from '../infrastructure/bridge/bridge.service.js';
+import {
+  DuplicateEntryKeyError,
+  MailBucketEntryRepository,
+} from './repositories/mail-bucket-entry.repository.js';
+
+export interface TrackStoredMessageParams {
+  mailAccountId: string;
+  userUuid: string;
+  bucketId: string;
+  entryKey: string;
+  size: number;
+}
+
+export interface ReleaseStoredMessageParams {
+  userUuid: string;
+  bucketId: string;
+  entryKey: string;
+}
+
+@Injectable()
+export class MailUsageService {
+  private readonly logger = new Logger(MailUsageService.name);
+
+  constructor(
+    private readonly entries: MailBucketEntryRepository,
+    private readonly bridge: BridgeClient,
+  ) {}
+
+  async trackStoredMessage(params: TrackStoredMessageParams): Promise<void> {
+    const { mailAccountId, userUuid, bucketId, entryKey, size } = params;
+
+    const existing = await this.entries.findByEntryKey(entryKey);
+    if (existing) {
+      this.logger.debug({ entryKey }, 'Message already tracked; skipping');
+      return;
+    }
+
+    const { id: bridgeEntryId, totalUsedSpaceBytes } =
+      await this.bridge.createBucketEntry(userUuid, bucketId, size);
+
+    try {
+      await this.entries.create({
+        mailAccountId,
+        entryKey,
+        bridgeEntryId,
+        size,
+      });
+    } catch (error) {
+      if (error instanceof DuplicateEntryKeyError) {
+        this.logger.debug(
+          { entryKey, bridgeEntryId },
+          'Concurrent tracking detected; rolling back minted bucket entry',
+        );
+        await this.bridge.deleteBucketEntry(userUuid, bucketId, bridgeEntryId);
+        return;
+      }
+      throw error;
+    }
+
+    this.logger.log(
+      { entryKey, bridgeEntryId, size, totalUsedSpaceBytes },
+      'Tracked stored message',
+    );
+  }
+
+  async releaseStoredMessage(
+    params: ReleaseStoredMessageParams,
+  ): Promise<void> {
+    const { userUuid, bucketId, entryKey } = params;
+
+    const existing = await this.entries.findByEntryKey(entryKey);
+    if (!existing) {
+      this.logger.debug(
+        { entryKey },
+        'No tracked entry for message; skipping release',
+      );
+      return;
+    }
+
+    await this.bridge.deleteBucketEntry(
+      userUuid,
+      bucketId,
+      existing.bridgeEntryId,
+    );
+    await this.entries.deleteByEntryKey(entryKey);
+
+    this.logger.log(
+      { entryKey, bridgeEntryId: existing.bridgeEntryId },
+      'Released stored message',
+    );
+  }
+}