stack/apps/gateway/src/agent/connector-lease.repository.test.ts

import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
import {
  connectorLeaseAuditLog,
  createPgliteDb,
  eq,
  runPgliteMigrations,
  type DbHandle,
} from '@mosaicstack/db';
import { ConnectorLeaseCoordinator, ConnectorLeaseError } from '@mosaicstack/agent';
import { ConnectorLeaseRepository } from './connector-lease.repository.js';

const identity = { tenantId: 'tenant-a', logicalAgentId: 'mos' } as const;

function acquireCommand(connectorId: string, correlationId: string) {
  return {
    identity,
    bindingId: 'operator-chat',
    connectorId,
    scopes: ['runtime.send', 'tool.execute'],
    ttlMs: 60_000,
    correlationId,
  };
}

describe('ConnectorLeaseRepository PostgreSQL semantics', (): void => {
  let dataDir: string;
  let handle: DbHandle;
  let now: Date;
  let coordinator: ConnectorLeaseCoordinator;

  beforeEach(async (): Promise<void> => {
    dataDir = await mkdtemp(join(tmpdir(), 'mosaic-connector-lease-'));
    handle = createPgliteDb(dataDir);
    await runPgliteMigrations(handle);
    now = new Date('2026-07-14T17:00:00.000Z');
    coordinator = new ConnectorLeaseCoordinator(new ConnectorLeaseRepository(handle.db), {
      now: (): Date => now,
    });
  });

  afterEach(async (): Promise<void> => {
    await handle.close();
    await rm(dataDir, { recursive: true, force: true });
  });

  it('allows only one concurrent contender to acquire a binding', async (): Promise<void> => {
    const outcomes = await Promise.allSettled([
      coordinator.acquire(acquireCommand('connector-a', 'correlation-a')),
      coordinator.acquire(acquireCommand('connector-b', 'correlation-b')),
    ]);

    expect(outcomes.filter((result) => result.status === 'fulfilled')).toHaveLength(1);
    const rejected = outcomes.find((result) => result.status === 'rejected');
    expect(rejected).toMatchObject({
      reason: { code: 'lease_held' } satisfies Partial<ConnectorLeaseError>,
    });
  });

  it('uses compare-and-swap takeover and increments the fencing epoch monotonically', async (): Promise<void> => {
    const acquired = await coordinator.acquire(acquireCommand('connector-a', 'correlation-a'));
    const results = await Promise.allSettled([
      coordinator.takeover({
        ...acquireCommand('connector-b', 'correlation-b'),
        expectedEpoch: acquired.leaseEpoch,
      }),
      coordinator.takeover({
        ...acquireCommand('connector-c', 'correlation-c'),
        expectedEpoch: acquired.leaseEpoch,
      }),
    ]);
    const winner = results.find((result) => result.status === 'fulfilled');

    expect(results.filter((result) => result.status === 'fulfilled')).toHaveLength(1);
    expect(winner?.status === 'fulfilled' ? winner.value.leaseEpoch : null).toBe('2');
    expect(results.find((result) => result.status === 'rejected')).toMatchObject({
      reason: { code: 'cas_mismatch' } satisfies Partial<ConnectorLeaseError>,
    });
  });

  it('heartbeats and releases only the current connector epoch', async (): Promise<void> => {
    const acquired = await coordinator.acquire(acquireCommand('connector-a', 'correlation-a'));
    now = new Date('2026-07-14T17:00:30.000Z');
    const renewed = await coordinator.heartbeat({
      lease: acquired,
      ttlMs: 120_000,
      correlationId: 'correlation-renew',
    });
    expect(renewed.expiresAt).toBe('2026-07-14T17:02:30.000Z');

    await coordinator.release({ lease: renewed, correlationId: 'correlation-release' });
    await expect(
      coordinator.heartbeat({
        lease: renewed,
        ttlMs: 120_000,
        correlationId: 'correlation-stale',
      }),
    ).rejects.toMatchObject({ code: 'lease_released' } satisfies Partial<ConnectorLeaseError>);
  });

  it('survives close/reopen and requires CAS takeover to recover an expired lease', async (): Promise<void> => {
    const acquired = await coordinator.acquire(acquireCommand('connector-a', 'correlation-a'));
    await handle.close();

    now = new Date('2026-07-14T17:02:00.000Z');
    handle = createPgliteDb(dataDir);
    await runPgliteMigrations(handle);
    coordinator = new ConnectorLeaseCoordinator(new ConnectorLeaseRepository(handle.db), {
      now: (): Date => now,
    });

    await expect(
      coordinator.acquire(acquireCommand('connector-b', 'correlation-plain-acquire')),
    ).rejects.toMatchObject({ code: 'takeover_required' } satisfies Partial<ConnectorLeaseError>);
    const recovered = await coordinator.takeover({
      ...acquireCommand('connector-b', 'correlation-takeover'),
      expectedEpoch: acquired.leaseEpoch,
    });
    expect(recovered).toMatchObject({ connectorId: 'connector-b', leaseEpoch: '2' });
  });

  it('writes credential-safe lifecycle and rejection audit records', async (): Promise<void> => {
    const acquired = await coordinator.acquire(acquireCommand('connector-a', 'correlation-a'));
    await coordinator.heartbeat({
      lease: acquired,
      ttlMs: 60_000,
      correlationId: 'correlation-renew',
    });
    await expect(
      coordinator.acquire(acquireCommand('connector-b', 'correlation-reject')),
    ).rejects.toBeInstanceOf(ConnectorLeaseError);

    const rows = await handle.db
      .select()
      .from(connectorLeaseAuditLog)
      .where(eq(connectorLeaseAuditLog.tenantId, identity.tenantId));
    expect(rows.map((row) => row.event)).toEqual(
      expect.arrayContaining(['acquire', 'renew', 'reject']),
    );
    const serialized = JSON.stringify(rows, (_key: string, value: unknown): unknown =>
      typeof value === 'bigint' ? value.toString(10) : value,
    );
    expect(serialized).not.toContain('tool.execute');
    expect(serialized).not.toContain('runtime.send');
    expect(serialized).not.toMatch(/token|secret|credential/i);
  });
});