agenthub/src/core/index.ts

import Database from 'better-sqlite3';
import { mkdirSync } from 'fs';
import { dirname } from 'path';
import { getIndexPath } from './paths.js';

export interface IndexEntry {
  id: string;
  type: string;
  title: string;
  content: string;
  filePath: string;
  createdAt: string;
  updatedAt: string;
  status?: string;
  role?: string;
  assignedTo?: string;
  reviewer?: string;
  tags?: string;
  // Handoff-specific routing fields
  fromRole?: string;
  toRole?: string;
  fromAgent?: string;
  toAgent?: string;
  // Handoff → task link
  taskId?: string;
  // Memory → task links (JSON array string, e.g. '["TSK-0001"]')
  relatedTasks?: string;
}

export class Index {
  private db: Database.Database;

  constructor(cwd: string) {
    const path = getIndexPath(cwd);
    mkdirSync(dirname(path), { recursive: true });
    this.db = new Database(path);
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS entities (
        id TEXT PRIMARY KEY,
        type TEXT NOT NULL,
        title TEXT NOT NULL,
        content TEXT NOT NULL,
        filePath TEXT NOT NULL,
        createdAt TEXT NOT NULL,
        updatedAt TEXT NOT NULL,
        status TEXT,
        role TEXT,
        assignedTo TEXT,
        reviewer TEXT,
        tags TEXT,
        fromRole TEXT,
        toRole TEXT,
        fromAgent TEXT,
        toAgent TEXT,
        taskId TEXT,
        relatedTasks TEXT
      );
      CREATE VIRTUAL TABLE IF NOT EXISTS search USING fts5(id, title, content);
    `);

    // Migration: add columns to existing DBs that pre-date this schema.
    const existingCols = new Set(
      (this.db.pragma('table_info(entities)') as Array<{ name: string }>).map((r) => r.name),
    );
    for (const col of ['reviewer', 'fromRole', 'toRole', 'fromAgent', 'toAgent', 'taskId', 'relatedTasks']) {
      if (!existingCols.has(col)) {
        this.db.exec(`ALTER TABLE entities ADD COLUMN ${col} TEXT`);
      }
    }
  }

  upsert(entry: IndexEntry): void {
    const params = {
      status: null,
      role: null,
      assignedTo: null,
      reviewer: null,
      tags: null,
      fromRole: null,
      toRole: null,
      fromAgent: null,
      toAgent: null,
      taskId: null,
      relatedTasks: null,
      ...entry,
    };

    const insert = this.db.prepare(`
      INSERT INTO entities (id, type, title, content, filePath, createdAt, updatedAt, status, role, assignedTo, reviewer, tags, fromRole, toRole, fromAgent, toAgent, taskId, relatedTasks)
      VALUES (@id, @type, @title, @content, @filePath, @createdAt, @updatedAt, @status, @role, @assignedTo, @reviewer, @tags, @fromRole, @toRole, @fromAgent, @toAgent, @taskId, @relatedTasks)
      ON CONFLICT(id) DO UPDATE SET
        type=@type, title=@title, content=@content, filePath=@filePath,
        createdAt=@createdAt, updatedAt=@updatedAt, status=@status, role=@role,
        assignedTo=@assignedTo, reviewer=@reviewer, tags=@tags,
        fromRole=@fromRole, toRole=@toRole, fromAgent=@fromAgent, toAgent=@toAgent,
        taskId=@taskId, relatedTasks=@relatedTasks
    `);
    insert.run(params);

    const search = this.db.prepare(`
      INSERT OR REPLACE INTO search (id, title, content) VALUES (@id, @title, @content)
    `);
    search.run(params);
  }

  /** Permanently drop an entity from both the entities table and the FTS index. */
  remove(id: string): void {
    this.db.prepare('DELETE FROM entities WHERE id = @id').run({ id });
    this.db.prepare('DELETE FROM search WHERE id = @id').run({ id });
  }

  search(query: string): Array<{ id: string; type: string; title: string }> {
    const ftsQuery = toFts5Phrase(query);
    if (!ftsQuery) return [];
    const stmt = this.db.prepare(`
      SELECT e.id, e.type, e.title
      FROM search s
      JOIN entities e ON s.id = e.id
      WHERE search MATCH @query
      ORDER BY rank
    `);
    return stmt.all({ query: ftsQuery }) as Array<{ id: string; type: string; title: string }>;
  }

  list(type?: string, filters?: { status?: string; role?: string; assignedTo?: string }): IndexEntry[] {
    let sql = 'SELECT * FROM entities WHERE 1=1';
    const params: Record<string, string> = {};

    if (type) {
      sql += ' AND type = @type';
      params.type = type;
    }
    if (filters?.status) {
      sql += ' AND status = @status';
      params.status = filters.status;
    }
    if (filters?.role) {
      sql += ' AND role = @role';
      params.role = filters.role;
    }
    if (filters?.assignedTo) {
      sql += ' AND assignedTo = @assignedTo';
      params.assignedTo = filters.assignedTo;
    }

    sql += ' ORDER BY createdAt DESC';

    const stmt = this.db.prepare(sql);
    return stmt.all(params) as IndexEntry[];
  }

  /** Returns all handoff entries whose taskId matches the given task ID. */
  listHandoffsByTask(taskId: string): IndexEntry[] {
    const stmt = this.db.prepare(
      "SELECT * FROM entities WHERE type = 'handoff' AND taskId = @taskId ORDER BY createdAt ASC",
    );
    return stmt.all({ taskId }) as IndexEntry[];
  }

  /**
   * Returns all memory entries whose relatedTasks JSON array contains the given
   * task ID. Uses a LIKE search on the serialised JSON string — safe because
   * task IDs follow the form TSK-NNNN and cannot be substrings of each other.
   */
  listMemoryByTask(taskId: string): IndexEntry[] {
    const stmt = this.db.prepare(
      "SELECT * FROM entities WHERE type = 'memory' AND relatedTasks LIKE @pattern ORDER BY createdAt ASC",
    );
    return stmt.all({ pattern: `%${taskId}%` }) as IndexEntry[];
  }

  close(): void {
    this.db.close();
  }
}

function toFts5Phrase(query: string): string {
  const trimmed = query.trim();
  if (!trimmed) return '';
  return `"${trimmed.replace(/"/g, '""')}"`;
}