Feat/prem keys

.env.example

@@ -34,6 +34,9 @@ DATABASE_URL=
 # Redis
 REDIS_URL=
 
+# Warn a device to replenish when its unconsumed one-time prekeys drop below
+# this count (emits `prekeys_low` over the socket). Defaults to 20.
+PREKEY_LOW_THRESHOLD=
 # Rate limits and quotas (see docs/security/rate-limits.md).
 # Any bucket in apps/backend/src/config/rateLimits.ts can be overridden with
 # RATE_LIMIT_<BUCKET>=<limit>[/<windowSeconds>]. Defaults apply when unset.

apps/backend/docs/e2ee-onboarding.md

@@ -535,6 +535,50 @@ Required guarantees for this path:
 - signed prekey must still be present if the device remains reachable for session bootstrap
 - client must treat this as lower-entropy/fallback first-contact establishment and should trigger recipient prekey replenishment UX when possible
 
+### C) Low-prekey warning before exhaustion
+
+Waiting for exhaustion means every sender in the meantime is downgraded to
+3-DH, so the backend warns the owning device *before* it runs dry.
+
+Two surfaces expose this:
+
+1. `GET /devices` returns `oneTimePreKeysRemaining` per device — a count of
+   unconsumed one-time prekeys, `0` when the device has none. Poll-free clients
+   can read this at startup to decide whether to top up.
+2. A `prekeys_low` Socket.IO event is emitted after a bundle fetch drops a
+   device below the threshold (default `20`, overridable with the
+   `PREKEY_LOW_THRESHOLD` env var).
+
+Event payload:
+
+```json
+{
+  "deviceId": "uuid",
+  "oneTimePreKeysRemaining": 19,
+  "threshold": 20
+}
+```
+
+Delivery and debounce semantics:
+
+- emitted only to the `device:{deviceId}` room — the owning device, on whichever
+  gateway holds its socket. No other device on the account sees it, since only
+  the owner can generate replacement prekeys.
+- fired **at most once per threshold crossing**. A device that keeps serving
+  bundles while below the threshold is told once, not once per fetch.
+- the signal re-arms when `POST /devices/:id/prekeys` brings the device back to
+  or above the threshold, so a later crossing warns again. Revoking a device
+  also re-arms it (its prekeys are deleted).
+- a device that is offline when the threshold is crossed misses the event; it
+  should read `oneTimePreKeysRemaining` from `GET /devices` on reconnect.
+
+Client behavior:
+
+- on `prekeys_low`, generate and upload a fresh batch via
+  `POST /devices/:id/prekeys`, respecting the `200` cap
+- the upload response echoes `oneTimePreKeysRemaining` so the client can confirm
+  it is back above the threshold without a follow-up `GET /devices`
+
 ## End-to-end ordering contract
 
 For compatibility with the current implementation, clients should rely on this ordering:
@@ -564,8 +608,10 @@ For compatibility with the current implementation, clients should rely on this o
 - device registration/listing/revocation/prekey upload: `apps/backend/src/routes/devices.ts`
 - recipient key-bundle fetch: `apps/backend/src/routes/users.ts` (`GET /users/:userId/devices/:deviceId/key-bundle`)
 - E2EE-related schema: `apps/backend/src/db/schema.ts` (`devices`, `devicePrekeys`)
+- low-prekey signal + debounce latch: `apps/backend/src/services/prekeyLowSignal.ts`
 - prekey route tests: `apps/backend/src/__tests__/devices.prekeys.test.ts`
 - key-bundle route tests: `apps/backend/src/__tests__/users.bundle.test.ts`
+- low-prekey signal tests: `apps/backend/src/__tests__/prekeysLow.test.ts`
 
 ## Gaps to close for full first-DM support
 

apps/backend/src/__tests__/devices.prekeys.test.ts

@@ -35,6 +35,14 @@ vi.mock('../db/schema.js', () => ({
 
 vi.mock('../lib/redis.js', () => ({ redis: null }));
 
+// Real threshold + count logic, but the latch release is spied so we can assert
+// the upload route re-arms the signal only once the device is back in range.
+const mockReleaseLatch = vi.fn().mockResolvedValue(undefined);
+vi.mock('../services/prekeyLowSignal.js', async (importOriginal) => ({
+  ...(await importOriginal<typeof import('../services/prekeyLowSignal.js')>()),
+  releasePrekeysLowLatch: mockReleaseLatch,
+}));
+
 vi.mock('drizzle-orm', () => ({
   eq: vi.fn((col: unknown, val: unknown) => ({ col, val })),
   and: vi.fn((...args: unknown[]) => args),
@@ -190,6 +198,30 @@ describe('POST /devices/:id/prekeys', () => {
     expect(mockInsert).toHaveBeenCalledTimes(2); // signed + OTP
   });
 
+  it('re-arms the low-prekey signal once replenished to the threshold', async () => {
+    mockDeviceFindFirst.mockResolvedValue(ACTIVE_DEVICE);
+    setupOtpCount(50); // post-upload recount lands at or above the threshold
+    setupInsertChain();
+
+    const res = await request(makeApp()).post('/devices/device-1/prekeys').send(VALID_BODY);
+
+    expect(res.status).toBe(200);
+    expect(res.body.oneTimePreKeysRemaining).toBe(50);
+    expect(mockReleaseLatch).toHaveBeenCalledWith('device-1');
+  });
+
+  it('leaves the signal armed when still below the threshold after upload', async () => {
+    mockDeviceFindFirst.mockResolvedValue(ACTIVE_DEVICE);
+    setupOtpCount(3);
+    setupInsertChain();
+
+    const res = await request(makeApp()).post('/devices/device-1/prekeys').send(VALID_BODY);
+
+    expect(res.status).toBe(200);
+    expect(res.body.oneTimePreKeysRemaining).toBe(3);
+    expect(mockReleaseLatch).not.toHaveBeenCalled();
+  });
+
   it('trims the OTP batch to the remaining cap space', async () => {
     mockDeviceFindFirst.mockResolvedValue(ACTIVE_DEVICE);
     setupOtpCount(199); // 1 slot left

apps/backend/src/__tests__/prekeysLow.test.ts

@@ -0,0 +1,136 @@
+/**
+ * Tests for the low one-time-prekey signal.
+ *
+ * Covers the debounce contract: `prekeys_low` fires at most once per threshold
+ * crossing, and replenishing back to the threshold re-arms it.
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+// ── Mocks ─────────────────────────────────────────────────────────────────────
+
+const mockSelect = vi.fn();
+const mockEmit = vi.fn();
+const mockTo = vi.fn(() => ({ emit: mockEmit }));
+
+vi.mock('../db/index.js', () => ({ db: { select: mockSelect } }));
+
+vi.mock('../db/schema.js', () => ({
+  devicePrekeys: { deviceId: 'deviceId', keyType: 'keyType', consumed: 'consumed' },
+}));
+
+vi.mock('drizzle-orm', () => ({
+  eq: vi.fn((col: unknown, val: unknown) => ({ col, val })),
+  and: vi.fn((...args: unknown[]) => args),
+  count: vi.fn(() => 'count(*)'),
+}));
+
+vi.mock('../lib/redis.js', () => ({ redis: null }));
+vi.mock('../lib/socket.js', () => ({ getSocketServer: vi.fn(() => ({ to: mockTo })) }));
+
+const {
+  PREKEY_LOW_THRESHOLD,
+  signalPrekeysLowIfNeeded,
+  releasePrekeysLowLatch,
+  countAvailableOneTimePreKeys,
+  __resetPrekeyLowLatches,
+} = await import('../services/prekeyLowSignal.js');
+
+const DEVICE = 'device-1';
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  __resetPrekeyLowLatches();
+});
+
+describe('signalPrekeysLowIfNeeded', () => {
+  it('defaults the threshold to 20', () => {
+    expect(PREKEY_LOW_THRESHOLD).toBe(20);
+  });
+
+  it('emits prekeys_low to the device room when below the threshold', async () => {
+    await signalPrekeysLowIfNeeded(DEVICE, 19);
+
+    expect(mockTo).toHaveBeenCalledWith(`device:${DEVICE}`);
+    expect(mockEmit).toHaveBeenCalledTimes(1);
+    expect(mockEmit).toHaveBeenCalledWith('prekeys_low', {
+      deviceId: DEVICE,
+      oneTimePreKeysRemaining: 19,
+      threshold: 20,
+    });
+  });
+
+  it('does not emit while at or above the threshold', async () => {
+    await signalPrekeysLowIfNeeded(DEVICE, 20);
+    await signalPrekeysLowIfNeeded(DEVICE, 50);
+
+    expect(mockEmit).not.toHaveBeenCalled();
+  });
+
+  it('emits at most once per threshold crossing', async () => {
+    for (let remaining = 19; remaining >= 0; remaining--) {
+      await signalPrekeysLowIfNeeded(DEVICE, remaining);
+    }
+
+    expect(mockEmit).toHaveBeenCalledTimes(1);
+    expect(mockEmit).toHaveBeenCalledWith(
+      'prekeys_low',
+      expect.objectContaining({ oneTimePreKeysRemaining: 19 }),
+    );
+  });
+
+  it('latches per device — one device going low does not silence another', async () => {
+    await signalPrekeysLowIfNeeded('device-a', 5);
+    await signalPrekeysLowIfNeeded('device-b', 5);
+    await signalPrekeysLowIfNeeded('device-a', 4);
+
+    expect(mockEmit).toHaveBeenCalledTimes(2);
+  });
+
+  it('stops firing after replenishment, then fires again on the next crossing', async () => {
+    await signalPrekeysLowIfNeeded(DEVICE, 3);
+    expect(mockEmit).toHaveBeenCalledTimes(1);
+
+    // Replenished above the threshold — this re-arms the latch and is itself
+    // silent.
+    await signalPrekeysLowIfNeeded(DEVICE, 100);
+    expect(mockEmit).toHaveBeenCalledTimes(1);
+
+    // Draining back down below the threshold signals once more.
+    await signalPrekeysLowIfNeeded(DEVICE, 19);
+    await signalPrekeysLowIfNeeded(DEVICE, 18);
+    expect(mockEmit).toHaveBeenCalledTimes(2);
+  });
+
+  it('re-arms after an explicit latch release (e.g. prekey upload)', async () => {
+    await signalPrekeysLowIfNeeded(DEVICE, 2);
+    await releasePrekeysLowLatch(DEVICE);
+    await signalPrekeysLowIfNeeded(DEVICE, 2);
+
+    expect(mockEmit).toHaveBeenCalledTimes(2);
+  });
+
+  it('swallows socket-layer failures', async () => {
+    mockTo.mockImplementationOnce(() => {
+      throw new Error('socket server exploded');
+    });
+
+    await expect(signalPrekeysLowIfNeeded(DEVICE, 1)).resolves.toBeUndefined();
+  });
+});
+
+describe('countAvailableOneTimePreKeys', () => {
+  it('returns the unconsumed one-time prekey count', async () => {
+    const where = vi.fn().mockResolvedValue([{ total: 7 }]);
+    mockSelect.mockReturnValue({ from: vi.fn().mockReturnValue({ where }) });
+
+    expect(await countAvailableOneTimePreKeys(DEVICE)).toBe(7);
+  });
+
+  it('returns 0 when the device has none', async () => {
+    const where = vi.fn().mockResolvedValue([]);
+    mockSelect.mockReturnValue({ from: vi.fn().mockReturnValue({ where }) });
+
+    expect(await countAvailableOneTimePreKeys(DEVICE)).toBe(0);
+  });
+});

apps/backend/src/__tests__/users.bundle.test.ts

@@ -54,9 +54,15 @@ vi.mock('drizzle-orm', () => ({
   ilike: vi.fn(),
   exists: vi.fn(),
   isNull: vi.fn((col: unknown) => ({ op: 'isNull', col })),
+  count: vi.fn(() => 'count(*)'),
   sql: vi.fn(),
 }));
 
+const mockSignalPrekeysLow = vi.fn().mockResolvedValue(undefined);
+vi.mock('../services/prekeyLowSignal.js', () => ({
+  signalPrekeysLowIfNeeded: mockSignalPrekeysLow,
+}));
+
 vi.mock('../lib/redis.js', () => ({ redis: null }));
 vi.mock('../services/presence.js', () => ({ isOnline: vi.fn().mockResolvedValue(false) }));
 vi.mock('../lib/socket.js', () => ({ getSocketServer: vi.fn(() => null) }));
@@ -135,17 +141,26 @@ describe('GET /users/:userId/devices/:deviceId/key-bundle', () => {
     const updateSet = vi.fn().mockReturnValue({ where: updateWhere });
     const lockFor = vi.fn().mockResolvedValue([claimed]);
     const tx = {
-      select: vi.fn().mockReturnValue({
-        from: vi.fn().mockReturnValue({
-          where: vi.fn().mockReturnValue({
-            orderBy: vi.fn().mockReturnValue({
-              limit: vi.fn().mockReturnValue({
-                for: lockFor,
+      // Two selects run inside the claiming transaction: the row-locked claim,
+      // then the post-consumption remaining count that drives `prekeys_low`.
+      select: vi
+        .fn()
+        .mockReturnValueOnce({
+          from: vi.fn().mockReturnValue({
+            where: vi.fn().mockReturnValue({
+              orderBy: vi.fn().mockReturnValue({
+                limit: vi.fn().mockReturnValue({
+                  for: lockFor,
+                }),
               }),
             }),
           }),
+        })
+        .mockReturnValueOnce({
+          from: vi.fn().mockReturnValue({
+            where: vi.fn().mockResolvedValue([{ total: 4 }]),
+          }),
         }),
-      }),
       update: vi.fn().mockReturnValue({ set: updateSet }),
     };
     mockTransaction.mockImplementation(async (cb: (txArg: typeof tx) => unknown) => cb(tx));
@@ -167,6 +182,9 @@ describe('GET /users/:userId/devices/:deviceId/key-bundle', () => {
     // history stays auditable.
     expect(updateSet).toHaveBeenCalledWith({ consumed: true });
     expect(updateWhere).toHaveBeenCalled();
+    // The post-consumption count feeds the low-prekey signal, which decides on
+    // its own whether this crossed the threshold.
+    expect(mockSignalPrekeysLow).toHaveBeenCalledWith('device-2', 4);
   });
 
   it('uses skipLocked so concurrent bundle reads only consume one OTP', async () => {
@@ -234,5 +252,8 @@ describe('GET /users/:userId/devices/:deviceId/key-bundle', () => {
     expect(res.status).toBe(200);
     expect(res.body.oneTimePreKey).toBeNull();
     expect(tx.update).not.toHaveBeenCalled();
+    // Nothing was consumed, so the count did not move — the crossing already
+    // signalled on the fetch that drained the last key.
+    expect(mockSignalPrekeysLow).not.toHaveBeenCalled();
   });
 });

apps/backend/src/routes/devices.ts

@@ -25,6 +25,11 @@ import { SignedPreKeyEntrySchema, PreKeyEntrySchema, verifyEd25519Signature } fr
 import { conversationRoom } from '../services/roomManager.js';
 import { redis } from '../lib/redis.js';
 import { markDeviceRevoked } from '../services/deviceRevocation.js';
+import {
+  PREKEY_LOW_THRESHOLD,
+  countAvailableOneTimePreKeys,
+  releasePrekeysLowLatch,
+} from '../services/prekeyLowSignal.js';
 import { actorFromRequest, recordAuditEvent } from '../services/auditLog.js';
 
 export const devicesRouter: RouterType = Router();
@@ -128,6 +133,10 @@ async function revokeDeviceRow(deviceId: string): Promise<Date> {
   await db.update(devices).set({ revokedAt, updatedAt: revokedAt }).where(eq(devices.id, deviceId));
   await db.delete(devicePrekeys).where(eq(devicePrekeys.deviceId, deviceId));
 
+  // The device's prekeys are gone; leaving a latch behind would suppress the
+  // first low signal if this identity key is later re-registered.
+  await releasePrekeysLowLatch(deviceId);
+
   // Force-disconnect any live socket for this device, on this node or any
   // other (cross-node via Redis pub/sub — see services/deviceRevocation.ts).
   markDeviceRevoked(deviceId);
@@ -347,10 +356,20 @@ devicesRouter.post('/:id/prekeys', validate(UploadPreKeysSchema), async (req: Au
       });
   }
 
+  // Re-arm the low-prekey signal once the device is back at or above the
+  // threshold, so a future crossing fires again. Recounted rather than derived
+  // from `currentCount + trimmedBatch.length` because duplicate keyIds are
+  // dropped by ON CONFLICT DO NOTHING and would inflate the derived figure.
+  const replenishedCount = await countAvailableOneTimePreKeys(deviceId);
+  if (replenishedCount >= PREKEY_LOW_THRESHOLD) {
+    await releasePrekeysLowLatch(deviceId);
+  }
+
   res.status(200).json({
     uploadedSignedPreKey: true,
     uploadedOneTimePreKeys: trimmedBatch.length,
     capped: trimmedBatch.length < otpBatch.length,
+    oneTimePreKeysRemaining: replenishedCount,
   });
 });