sentryagent-idp/src/repositories/AgentRepository.ts

/**
 * Agent Repository for SentryAgent.ai AgentIdP.
 * All SQL queries for the agents table live exclusively here.
 */

import { Pool, QueryResult } from 'pg';
import { v4 as uuidv4 } from 'uuid';
import {
  IAgent,
  ICreateAgentRequest,
  IUpdateAgentRequest,
  IAgentListFilters,
  AgentStatus,
} from '../types/index.js';

/** Raw database row for an agent. */
interface AgentRow {
  agent_id: string;
  organization_id: string;
  email: string;
  agent_type: string;
  version: string;
  capabilities: string[];
  owner: string;
  deployment_env: string;
  status: string;
  created_at: Date;
  updated_at: Date;
}

/**
 * Maps a raw database row to the IAgent domain model.
 *
 * @param row - Raw row from the agents table.
 * @returns Typed IAgent object.
 */
function mapRowToAgent(row: AgentRow): IAgent {
  return {
    agentId: row.agent_id,
    organizationId: row.organization_id,
    email: row.email,
    agentType: row.agent_type as IAgent['agentType'],
    version: row.version,
    capabilities: row.capabilities,
    owner: row.owner,
    deploymentEnv: row.deployment_env as IAgent['deploymentEnv'],
    status: row.status as AgentStatus,
    createdAt: row.created_at,
    updatedAt: row.updated_at,
  };
}

/**
 * Repository for all agent database operations.
 * Receives a pg.Pool via constructor injection.
 */
export class AgentRepository {
  /**
   * @param pool - The PostgreSQL connection pool.
   */
  constructor(private readonly pool: Pool) {}

  /**
   * Creates a new agent record in the database.
   *
   * @param data - The fields for the new agent.
   * @returns The created agent record.
   */
  async create(data: ICreateAgentRequest): Promise<IAgent> {
    const agentId = uuidv4();
    const organizationId = data.organizationId ?? 'org_system';
    const result: QueryResult<AgentRow> = await this.pool.query(
      `INSERT INTO agents
         (agent_id, organization_id, email, agent_type, version, capabilities, owner, deployment_env, status, created_at, updated_at)
       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'active', NOW(), NOW())
       RETURNING *`,
      [
        agentId,
        organizationId,
        data.email,
        data.agentType,
        data.version,
        data.capabilities,
        data.owner,
        data.deploymentEnv,
      ],
    );
    return mapRowToAgent(result.rows[0]);
  }

  /**
   * Finds an agent by its UUID.
   *
   * @param agentId - The agent UUID.
   * @returns The agent record, or null if not found.
   */
  async findById(agentId: string): Promise<IAgent | null> {
    const result: QueryResult<AgentRow> = await this.pool.query(
      'SELECT * FROM agents WHERE agent_id = $1',
      [agentId],
    );
    if (result.rows.length === 0) return null;
    return mapRowToAgent(result.rows[0]);
  }

  /**
   * Finds an agent by its email address.
   *
   * @param email - The agent email.
   * @returns The agent record, or null if not found.
   */
  async findByEmail(email: string): Promise<IAgent | null> {
    const result: QueryResult<AgentRow> = await this.pool.query(
      'SELECT * FROM agents WHERE email = $1',
      [email],
    );
    if (result.rows.length === 0) return null;
    return mapRowToAgent(result.rows[0]);
  }

  /**
   * Returns a paginated list of agents with optional filters.
   *
   * @param filters - Pagination and filter criteria.
   * @returns Object containing the agent list and total count.
   */
  async findAll(filters: IAgentListFilters): Promise<{ agents: IAgent[]; total: number }> {
    const conditions: string[] = [];
    const params: unknown[] = [];
    let paramIndex = 1;

    if (filters.owner !== undefined) {
      conditions.push(`owner = $${paramIndex++}`);
      params.push(filters.owner);
    }
    if (filters.agentType !== undefined) {
      conditions.push(`agent_type = $${paramIndex++}`);
      params.push(filters.agentType);
    }
    if (filters.status !== undefined) {
      conditions.push(`status = $${paramIndex++}`);
      params.push(filters.status);
    }

    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';

    const countResult: QueryResult<{ count: string }> = await this.pool.query(
      `SELECT COUNT(*) as count FROM agents ${whereClause}`,
      params,
    );
    const total = parseInt(countResult.rows[0].count, 10);

    const offset = (filters.page - 1) * filters.limit;
    const dataParams = [...params, filters.limit, offset];

    const dataResult: QueryResult<AgentRow> = await this.pool.query(
      `SELECT * FROM agents ${whereClause}
       ORDER BY created_at DESC
       LIMIT $${paramIndex++} OFFSET $${paramIndex}`,
      dataParams,
    );

    return {
      agents: dataResult.rows.map(mapRowToAgent),
      total,
    };
  }

  /**
   * Partially updates an agent record.
   *
   * @param agentId - The agent UUID to update.
   * @param data - The fields to update (only provided fields are changed).
   * @returns The updated agent record, or null if not found.
   */
  async update(agentId: string, data: IUpdateAgentRequest): Promise<IAgent | null> {
    const setClauses: string[] = [];
    const params: unknown[] = [];
    let paramIndex = 1;

    if (data.agentType !== undefined) {
      setClauses.push(`agent_type = $${paramIndex++}`);
      params.push(data.agentType);
    }
    if (data.version !== undefined) {
      setClauses.push(`version = $${paramIndex++}`);
      params.push(data.version);
    }
    if (data.capabilities !== undefined) {
      setClauses.push(`capabilities = $${paramIndex++}`);
      params.push(data.capabilities);
    }
    if (data.owner !== undefined) {
      setClauses.push(`owner = $${paramIndex++}`);
      params.push(data.owner);
    }
    if (data.deploymentEnv !== undefined) {
      setClauses.push(`deployment_env = $${paramIndex++}`);
      params.push(data.deploymentEnv);
    }
    if (data.status !== undefined) {
      setClauses.push(`status = $${paramIndex++}`);
      params.push(data.status);
    }

    if (setClauses.length === 0) return null;

    setClauses.push(`updated_at = NOW()`);
    params.push(agentId);

    const result: QueryResult<AgentRow> = await this.pool.query(
      `UPDATE agents SET ${setClauses.join(', ')}
       WHERE agent_id = $${paramIndex}
       RETURNING *`,
      params,
    );

    if (result.rows.length === 0) return null;
    return mapRowToAgent(result.rows[0]);
  }

  /**
   * Sets an agent's status to 'decommissioned'.
   *
   * @param agentId - The agent UUID to decommission.
   * @returns The updated agent record, or null if not found.
   */
  async decommission(agentId: string): Promise<IAgent | null> {
    const result: QueryResult<AgentRow> = await this.pool.query(
      `UPDATE agents
       SET status = 'decommissioned', updated_at = NOW()
       WHERE agent_id = $1
       RETURNING *`,
      [agentId],
    );
    if (result.rows.length === 0) return null;
    return mapRowToAgent(result.rows[0]);
  }

  /**
   * Counts all agents excluding decommissioned ones (for free-tier limit checks).
   *
   * @returns Total count of active and suspended agents.
   */
  async countActive(): Promise<number> {
    const result: QueryResult<{ count: string }> = await this.pool.query(
      `SELECT COUNT(*) as count FROM agents WHERE status != 'decommissioned'`,
    );
    return parseInt(result.rows[0].count, 10);
  }
}