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

import { Inject, Injectable } from '@nestjs/common';
import {
  and,
  connectorLeaseAuditLog,
  eq,
  gt,
  isNull,
  logicalAgentConnectorLeases,
  sql,
  type Db,
} from '@mosaicstack/db';
import { ConnectorLeaseError } from '@mosaicstack/agent';
import type {
  ConnectorLease,
  ConnectorLeaseAcquireMutation,
  ConnectorLeaseAuditEvent,
  ConnectorLeaseHeartbeatMutation,
  ConnectorLeaseRejectReason,
  ConnectorLeaseReleaseMutation,
  ConnectorLeaseStore,
  ConnectorLeaseTakeoverMutation,
  LogicalAgentBinding,
} from '@mosaicstack/types';
import { DB } from '../database/database.module.js';

interface SuccessfulMutation {
  readonly ok: true;
  readonly lease: ConnectorLease;
}

interface FailedMutation {
  readonly ok: false;
  readonly reason: ConnectorLeaseRejectReason;
}

type MutationResult = SuccessfulMutation | FailedMutation;

@Injectable()
export class ConnectorLeaseRepository implements ConnectorLeaseStore {
  constructor(@Inject(DB) private readonly db: Db) {}

  async acquire(input: ConnectorLeaseAcquireMutation): Promise<ConnectorLease> {
    const result: MutationResult = await this.db.transaction(
      async (tx): Promise<MutationResult> => {
        const inserted = await tx
          .insert(logicalAgentConnectorLeases)
          .values({
            leaseId: input.leaseId,
            tenantId: input.identity.tenantId,
            logicalAgentId: input.identity.logicalAgentId,
            bindingId: input.bindingId,
            connectorId: input.connectorId,
            scopes: [...input.scopes],
            leaseEpoch: 1n,
            acquiredAt: new Date(input.now),
            heartbeatAt: new Date(input.now),
            expiresAt: new Date(input.expiresAt),
            updatedAt: new Date(input.now),
          })
          .onConflictDoNothing()
          .returning();
        const row = inserted[0];
        if (row) {
          const lease = toLease(row);
          await insertAudit(tx, lifecycleAudit(input, lease, 'acquire'));
          return { ok: true, lease };
        }

        const current = await findRow(tx, input);
        if (current && current.expiresAt <= new Date(input.now) && !current.releasedAt) {
          await insertAudit(tx, lifecycleAudit(input, toLease(current), 'expiry'));
        }
        const reason: ConnectorLeaseRejectReason =
          current && (current.releasedAt || current.expiresAt <= new Date(input.now))
            ? 'takeover_required'
            : 'lease_held';
        await insertAudit(tx, rejectionAudit(input, current ? toLease(current) : null, reason));
        return { ok: false, reason };
      },
    );
    return unwrap(result);
  }

  async takeover(input: ConnectorLeaseTakeoverMutation): Promise<ConnectorLease> {
    const result: MutationResult = await this.db.transaction(
      async (tx): Promise<MutationResult> => {
        const current = await findRow(tx, input);
        if (!current || current.leaseEpoch.toString(10) !== input.expectedEpoch) {
          await insertAudit(
            tx,
            rejectionAudit(input, current ? toLease(current) : null, 'cas_mismatch'),
          );
          return { ok: false, reason: 'cas_mismatch' };
        }
        if (current.expiresAt <= new Date(input.now) && !current.releasedAt) {
          await insertAudit(tx, lifecycleAudit(input, toLease(current), 'expiry'));
        }
        const updated = await tx
          .update(logicalAgentConnectorLeases)
          .set({
            leaseId: input.leaseId,
            connectorId: input.connectorId,
            scopes: [...input.scopes],
            leaseEpoch: sql`${logicalAgentConnectorLeases.leaseEpoch} + 1`,
            acquiredAt: new Date(input.now),
            heartbeatAt: new Date(input.now),
            expiresAt: new Date(input.expiresAt),
            releasedAt: null,
            updatedAt: new Date(input.now),
          })
          .where(
            and(
              bindingPredicate(input),
              eq(logicalAgentConnectorLeases.leaseId, current.leaseId),
              eq(logicalAgentConnectorLeases.leaseEpoch, BigInt(input.expectedEpoch)),
            ),
          )
          .returning();
        const row = updated[0];
        if (!row) {
          await insertAudit(tx, rejectionAudit(input, toLease(current), 'cas_mismatch'));
          return { ok: false, reason: 'cas_mismatch' };
        }
        const lease = toLease(row);
        await insertAudit(tx, lifecycleAudit(input, lease, 'takeover'));
        return { ok: true, lease };
      },
    );
    return unwrap(result);
  }

  async heartbeat(input: ConnectorLeaseHeartbeatMutation): Promise<ConnectorLease> {
    const result: MutationResult = await this.db.transaction(
      async (tx): Promise<MutationResult> => {
        const updated = await tx
          .update(logicalAgentConnectorLeases)
          .set({
            heartbeatAt: new Date(input.now),
            expiresAt: new Date(input.expiresAt),
            updatedAt: new Date(input.now),
          })
          .where(
            and(
              bindingPredicate(input.lease),
              eq(logicalAgentConnectorLeases.leaseId, input.lease.leaseId),
              eq(logicalAgentConnectorLeases.connectorId, input.lease.connectorId),
              eq(logicalAgentConnectorLeases.leaseEpoch, BigInt(input.lease.leaseEpoch)),
              isNull(logicalAgentConnectorLeases.releasedAt),
              gt(logicalAgentConnectorLeases.expiresAt, new Date(input.now)),
            ),
          )
          .returning();
        const row = updated[0];
        if (row) {
          const lease = toLease(row);
          await insertAudit(tx, lifecycleAudit(input, lease, 'renew'));
          return { ok: true, lease };
        }
        const current = await findRow(tx, input.lease);
        const reason = classifyAuthorityFailure(
          current ? toLease(current) : null,
          input.lease,
          input.now,
        );
        if (reason === 'lease_expired' && current) {
          await insertAudit(tx, lifecycleAudit(input, toLease(current), 'expiry'));
        }
        await insertAudit(
          tx,
          rejectionAudit(
            { ...input.lease, correlationId: input.correlationId, now: input.now },
            current ? toLease(current) : null,
            reason,
          ),
        );
        return { ok: false, reason };
      },
    );
    return unwrap(result);
  }

  async release(input: ConnectorLeaseReleaseMutation): Promise<void> {
    const result: MutationResult = await this.db.transaction(
      async (tx): Promise<MutationResult> => {
        const updated = await tx
          .update(logicalAgentConnectorLeases)
          .set({
            releasedAt: new Date(input.now),
            expiresAt: new Date(input.now),
            updatedAt: new Date(input.now),
          })
          .where(
            and(
              bindingPredicate(input.lease),
              eq(logicalAgentConnectorLeases.leaseId, input.lease.leaseId),
              eq(logicalAgentConnectorLeases.connectorId, input.lease.connectorId),
              eq(logicalAgentConnectorLeases.leaseEpoch, BigInt(input.lease.leaseEpoch)),
              isNull(logicalAgentConnectorLeases.releasedAt),
              gt(logicalAgentConnectorLeases.expiresAt, new Date(input.now)),
            ),
          )
          .returning();
        const row = updated[0];
        if (row) {
          const lease = toLease(row);
          await insertAudit(tx, lifecycleAudit(input, lease, 'release'));
          return { ok: true, lease };
        }
        const current = await findRow(tx, input.lease);
        const reason = classifyAuthorityFailure(
          current ? toLease(current) : null,
          input.lease,
          input.now,
        );
        if (reason === 'lease_expired' && current) {
          await insertAudit(tx, lifecycleAudit(input, toLease(current), 'expiry'));
        }
        await insertAudit(
          tx,
          rejectionAudit(
            { ...input.lease, correlationId: input.correlationId, now: input.now },
            current ? toLease(current) : null,
            reason,
          ),
        );
        return { ok: false, reason };
      },
    );
    unwrap(result);
  }

  async findCurrent(binding: LogicalAgentBinding): Promise<ConnectorLease | null> {
    const row = await findRow(this.db, binding);
    return row ? toLease(row) : null;
  }

  async recordAudit(event: ConnectorLeaseAuditEvent): Promise<void> {
    await insertAudit(this.db, event);
  }
}

function unwrap(result: MutationResult): ConnectorLease {
  if (!result.ok) throw new ConnectorLeaseError(result.reason, safeErrorMessage(result.reason));
  return result.lease;
}

function safeErrorMessage(reason: ConnectorLeaseRejectReason): string {
  return `Connector lease mutation denied: ${reason}`;
}

function bindingPredicate(binding: LogicalAgentBinding) {
  return and(
    eq(logicalAgentConnectorLeases.tenantId, binding.identity.tenantId),
    eq(logicalAgentConnectorLeases.logicalAgentId, binding.identity.logicalAgentId),
    eq(logicalAgentConnectorLeases.bindingId, binding.bindingId),
  );
}

async function findRow(
  db: Pick<Db, 'select'>,
  binding: LogicalAgentBinding,
): Promise<typeof logicalAgentConnectorLeases.$inferSelect | null> {
  const rows = await db
    .select()
    .from(logicalAgentConnectorLeases)
    .where(bindingPredicate(binding))
    .limit(1);
  return rows[0] ?? null;
}

function toLease(row: typeof logicalAgentConnectorLeases.$inferSelect): ConnectorLease {
  return Object.freeze({
    identity: Object.freeze({ tenantId: row.tenantId, logicalAgentId: row.logicalAgentId }),
    bindingId: row.bindingId,
    leaseId: row.leaseId,
    connectorId: row.connectorId,
    scopes: Object.freeze([...row.scopes]),
    leaseEpoch: row.leaseEpoch.toString(10),
    acquiredAt: row.acquiredAt.toISOString(),
    heartbeatAt: row.heartbeatAt.toISOString(),
    expiresAt: row.expiresAt.toISOString(),
    ...(row.releasedAt ? { releasedAt: row.releasedAt.toISOString() } : {}),
  });
}

function classifyAuthorityFailure(
  current: ConnectorLease | null,
  claimed: ConnectorLease,
  now: string,
): ConnectorLeaseRejectReason {
  if (!current) return 'lease_missing';
  if (current.releasedAt) return 'lease_released';
  if (new Date(current.expiresAt) <= new Date(now)) return 'lease_expired';
  if (current.leaseEpoch !== claimed.leaseEpoch) return 'stale_epoch';
  return 'connector_mismatch';
}

function lifecycleAudit(
  input: { readonly correlationId: string; readonly now: string },
  lease: ConnectorLease,
  event: Exclude<ConnectorLeaseAuditEvent['event'], 'reject'>,
): ConnectorLeaseAuditEvent {
  return {
    identity: lease.identity,
    bindingId: lease.bindingId,
    connectorId: lease.connectorId,
    leaseId: lease.leaseId,
    leaseEpoch: lease.leaseEpoch,
    event,
    outcome: 'succeeded',
    correlationId: input.correlationId,
    occurredAt: input.now,
  };
}

function rejectionAudit(
  input: {
    readonly identity: ConnectorLease['identity'];
    readonly bindingId: string;
    readonly connectorId: string;
    readonly correlationId: string;
    readonly now: string;
  },
  current: ConnectorLease | null,
  reason: ConnectorLeaseRejectReason,
): ConnectorLeaseAuditEvent {
  return {
    identity: input.identity,
    bindingId: input.bindingId,
    connectorId: input.connectorId,
    event: 'reject',
    outcome: 'denied',
    correlationId: input.correlationId,
    occurredAt: input.now,
    ...(current ? { leaseId: current.leaseId, leaseEpoch: current.leaseEpoch } : {}),
    reason,
  };
}

async function insertAudit(db: Pick<Db, 'insert'>, event: ConnectorLeaseAuditEvent): Promise<void> {
  await db.insert(connectorLeaseAuditLog).values({
    tenantId: event.identity.tenantId,
    logicalAgentId: event.identity.logicalAgentId,
    bindingId: event.bindingId,
    connectorId: event.connectorId,
    ...(event.leaseId ? { leaseId: event.leaseId } : {}),
    ...(event.leaseEpoch ? { leaseEpoch: BigInt(event.leaseEpoch) } : {}),
    event: event.event,
    outcome: event.outcome,
    ...(event.reason ? { reason: event.reason } : {}),
    correlationId: event.correlationId,
    occurredAt: new Date(event.occurredAt),
  });
}