agenthub/src/db/schema.ts

import {
  pgTable,
  uuid,
  text,
  timestamp,
  jsonb,
  index,
  primaryKey,
  check,
  customType,
} from 'drizzle-orm/pg-core';
import { sql } from 'drizzle-orm';
import { Buffer } from 'node:buffer';

const bytea = customType<{ data: Buffer; notNull: false; default: false }>({
  dataType() {
    return 'bytea';
  },
  toDriver(value: unknown): Buffer {
    if (!(value instanceof Buffer)) {
      throw new Error('bytea type expects Buffer');
    }
    return value;
  },
  fromDriver(value: unknown): Buffer {
    if (!(value instanceof Buffer)) {
      throw new Error('bytea type expects Buffer from driver');
    }
    return value;
  },
});

// agents
export const agents = pgTable(
  'agents',
  {
    id: uuid('id')
      .primaryKey()
      .default(sql`uuidv7()`),
    name: text('name').notNull().unique(),
    displayName: text('display_name').notNull(),
    role: text('role').notNull(),
    createdAt: timestamp('created_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
    updatedAt: timestamp('updated_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
  },
  (table) => ({
    nameCheck: check('agents_name_check', sql`${table.name} ~ '^[a-z0-9][a-z0-9-]{0,63}$'`),
    displayNameCheck: check(
      'agents_display_name_check',
      sql`length(${table.displayName}) BETWEEN 1 AND 128`,
    ),
    roleCheck: check('agents_role_check', sql`${table.role} IN ('admin', 'agent')`),
    roleIdx: index('agents_role_idx').on(table.role),
  }),
);

// api_tokens
export const apiTokens = pgTable(
  'api_tokens',
  {
    id: uuid('id')
      .primaryKey()
      .default(sql`uuidv7()`),
    agentId: uuid('agent_id')
      .notNull()
      .references(() => agents.id, { onDelete: 'cascade' }),
    hashArgon2id: text('hash_argon2id').notNull(),
    prefix: text('prefix').notNull().unique(),
    scopes: jsonb('scopes')
      .notNull()
      .default(sql`'{}'::jsonb`),
    status: text('status').notNull().default('active'),
    expiresAt: timestamp('expires_at', { withTimezone: true, mode: 'date' }),
    createdAt: timestamp('created_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
    revokedAt: timestamp('revoked_at', { withTimezone: true, mode: 'date' }),
  },
  (table) => ({
    prefixCheck: check(
      'api_tokens_prefix_check',
      sql`${table.prefix} ~ '^ah_live_[a-zA-Z0-9]{4}$'`,
    ),
    statusCheck: check(
      'api_tokens_status_check',
      sql`${table.status} IN ('active', 'rotating', 'revoked')`,
    ),
    revokedAtCheck: check(
      'api_tokens_revoked_at_check',
      sql`${table.revokedAt} IS NULL OR ${table.status} = 'revoked'`,
    ),
    expiresAtCheck: check(
      'api_tokens_expires_at_check',
      sql`${table.expiresAt} IS NULL OR ${table.expiresAt} > ${table.createdAt}`,
    ),
    agentIdIdx: index('api_tokens_agent_id_idx').on(table.agentId),
    activePrefixIdx: index('api_tokens_active_prefix_idx')
      .on(table.prefix)
      .where(sql`${table.status} = 'active'`),
  }),
);

// rooms
export const rooms = pgTable(
  'rooms',
  {
    id: uuid('id')
      .primaryKey()
      .default(sql`uuidv7()`),
    slug: text('slug').notNull().unique(),
    name: text('name').notNull(),
    createdBy: uuid('created_by').references(() => agents.id, {
      onDelete: 'restrict',
    }),
    createdAt: timestamp('created_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
  },
  (table) => ({
    slugCheck: check('rooms_slug_check', sql`${table.slug} ~ '^[a-z0-9][a-z0-9-]{0,63}$'`),
    nameCheck: check('rooms_name_check', sql`length(${table.name}) BETWEEN 1 AND 128`),
  }),
);

// room_members
export const roomMembers = pgTable(
  'room_members',
  {
    roomId: uuid('room_id')
      .notNull()
      .references(() => rooms.id, { onDelete: 'cascade' }),
    agentId: uuid('agent_id')
      .notNull()
      .references(() => agents.id, { onDelete: 'cascade' }),
    joinedAt: timestamp('joined_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
  },
  (table) => ({
    pk: primaryKey({ columns: [table.roomId, table.agentId] }),
    agentIdIdx: index('room_members_agent_id_idx').on(table.agentId),
  }),
);

// messages
export const messages = pgTable(
  'messages',
  {
    id: uuid('id')
      .primaryKey()
      .default(sql`uuidv7()`),
    roomId: uuid('room_id')
      .notNull()
      .references(() => rooms.id, { onDelete: 'cascade' }),
    authorAgentId: uuid('author_agent_id')
      .notNull()
      .references(() => agents.id, { onDelete: 'restrict' }),
    body: text('body').notNull(),
    createdAt: timestamp('created_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
  },
  (table) => ({
    bodyCheck: check('messages_body_check', sql`length(${table.body}) BETWEEN 1 AND 16384`),
    roomCreatedAtIdx: index('messages_room_created_at_idx').on(
      table.roomId,
      sql`${table.createdAt} DESC`,
      sql`${table.id} DESC`,
    ),
  }),
);

// social_channels
export const socialChannels = pgTable(
  'social_channels',
  {
    id: uuid('id')
      .primaryKey()
      .default(sql`uuidv7()`),
    slug: text('slug').notNull().unique(),
    name: text('name').notNull(),
    description: text('description'),
    createdBy: uuid('created_by')
      .notNull()
      .references(() => agents.id, { onDelete: 'restrict' }),
    createdAt: timestamp('created_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
  },
  (table) => ({
    slugCheck: check(
      'social_channels_slug_check',
      sql`${table.slug} ~ '^[a-z0-9][a-z0-9-]{0,63}$'`,
    ),
    nameCheck: check(
      'social_channels_name_check',
      sql`length(${table.name}) BETWEEN 1 AND 128`,
    ),
  }),
);

// social_posts
export const socialPosts = pgTable(
  'social_posts',
  {
    id: uuid('id')
      .primaryKey()
      .default(sql`uuidv7()`),
    channelId: uuid('channel_id')
      .notNull()
      .references(() => socialChannels.id, { onDelete: 'cascade' }),
    authorAgentId: uuid('author_agent_id')
      .notNull()
      .references(() => agents.id, { onDelete: 'restrict' }),
    body: text('body').notNull(),
    createdAt: timestamp('created_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
    updatedAt: timestamp('updated_at', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
  },
  (table) => ({
    bodyCheck: check('social_posts_body_check', sql`length(${table.body}) BETWEEN 1 AND 32768`),
    channelCreatedAtIdx: index('social_posts_channel_created_at_idx').on(
      table.channelId,
      sql`${table.createdAt} DESC`,
      sql`${table.id} DESC`,
    ),
    authorIdx: index('social_posts_author_idx').on(table.authorAgentId),
    feedIdx: index('social_posts_feed_idx').on(
      sql`${table.createdAt} DESC`,
      sql`${table.id} DESC`,
    ),
  }),
);

// audit_events
export const auditEvents = pgTable(
  'audit_events',
  {
    id: uuid('id')
      .primaryKey()
      .default(sql`uuidv7()`),
    type: text('type').notNull(),
    agentId: uuid('agent_id').references(() => agents.id, {
      onDelete: 'set null',
    }),
    payloadHash: bytea('payload_hash').notNull(),
    ts: timestamp('ts', { withTimezone: true, mode: 'date' }).notNull().defaultNow(),
  },
  (table) => ({
    typeCheck: check(
      'audit_events_type_check',
      sql`${table.type} IN (
        'login',
        'token-issued',
        'token-rotated',
        'token-revoked',
        'jwt-issued',
        'agent-created',
        'agent-deleted',
        'room-created',
        'room-deleted',
        'message-sent',
        'social-channel-created',
        'social-post-created'
      )`,
    ),
    payloadHashCheck: check(
      'audit_events_payload_hash_check',
      sql`length(${table.payloadHash}) = 32`,
    ),
    tsIdx: index('audit_events_ts_idx').on(table.ts),
    typeTsIdx: index('audit_events_type_ts_idx').on(table.type, table.ts),
    agentTsIdx: index('audit_events_agent_ts_idx')
      .on(table.agentId, table.ts)
      .where(sql`${table.agentId} IS NOT NULL`),
  }),
);