/** * 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 { const agentId = uuidv4(); const organizationId = data.organizationId ?? 'org_system'; const result: QueryResult = 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 { const result: QueryResult = 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 { const result: QueryResult = 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 = 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 { 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 = 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 { const result: QueryResult = 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 { 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); } }