AgentCompany/src/main/database/runs.ts

import { randomUUID } from "node:crypto";
import { writeFileSync, readFileSync, existsSync, statSync, openSync, readSync, closeSync } from "node:fs";
import { appendFile } from "node:fs/promises";
import { join } from "node:path";
import type { DatabaseSync, SQLInputValue } from "node:sqlite";
import log from "electron-log/main.js";
import type { ConnectorId, RunLogChunk, RunRecord } from "@shared/types";
import { type ReadAll, type SqliteRow, type Transaction, nowIso, toBoolean } from "./helpers";

/**
 * Find the last complete UTF-8 character boundary in a buffer.
 * Returns the number of bytes up to (and including) the last complete character.
 * This prevents corrupting multi-byte CJK characters when slicing at chunk edges.
 */
function findUtf8Boundary(buffer: Buffer, bytesRead: number): number {
  if (bytesRead === 0) return 0;
  let pos = bytesRead - 1;
  while (pos >= 0 && (buffer[pos] & 0xc0) === 0x80) {
    pos--;
  }
  if (pos < 0) return bytesRead;
  const lead = buffer[pos];
  let charLen: number;
  if ((lead & 0x80) === 0) charLen = 1;
  else if ((lead & 0xe0) === 0xc0) charLen = 2;
  else if ((lead & 0xf0) === 0xe0) charLen = 3;
  else if ((lead & 0xf8) === 0xf0) charLen = 4;
  else return bytesRead;
  if (pos + charLen <= bytesRead) return bytesRead;
  return pos;
}

export function buildRunMethods(
  getDb: () => DatabaseSync,
  readAll: ReadAll,
  transaction: Transaction,
  deps: {
    logsDir: string;
  },
) {
  return {
    getRun(id: string): RunRecord {
      const row = getDb()
        .prepare(
          "select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId, connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode, signal, model, session_display_id as sessionDisplayId, cost_usd as costUsd, unattributed_cost as unattributedCost, started_at as startedAt, finished_at as finishedAt, log_path as logPath, created_at as createdAt, updated_at as updatedAt from runs where id = ?",
        )
        .get(id) as RunRecord | undefined;
      if (!row) throw new Error(`Run not found: ${id}`);
      row.unattributedCost = Boolean(row.unattributedCost);
      return row;
    },

    getRunAuthState(runId: string): { id: string; companyId: string; agentId: string; status: RunRecord["status"] } | null {
      const row = getDb()
        .prepare("select id, company_id as companyId, agent_id as agentId, status from runs where id = ?")
        .get(runId) as { id: string; companyId: string; agentId: string; status: RunRecord["status"] } | undefined;
      return row ?? null;
    },

    createRun(input: {
      id?: string;
      companyId: string;
      taskId: string | null;
      agentId: string;
      workspaceId: string | null;
      connectorId: ConnectorId;
    }) {
      const id = input.id ?? randomUUID();
      const ts = nowIso();
      const logPath = join(deps.logsDir, `${id}.log`);
      transaction(() => {
        getDb()
          .prepare(
            "insert into runs (id, company_id, task_id, agent_id, workspace_id, connector_id, status, summary, unattributed_cost, log_path, created_at, updated_at) values (?, ?, ?, ?, ?, ?, 'queued', '', 0, ?, ?, ?)",
          )
          .run(id, input.companyId, input.taskId, input.agentId, input.workspaceId, input.connectorId, logPath, ts, ts);
        if (input.taskId !== null) {
          getDb().prepare("update tasks set active_run_id = ?, updated_at = ? where id = ?").run(id, ts, input.taskId);
        }
      });
      writeFileSync(logPath, "", "utf8");
      return { id, logPath };
    },

    appendRunLog(runId: string, stream: "stdout" | "stderr", chunk: string) {
      const run = getDb().prepare("select log_path as logPath from runs where id = ?").get(runId) as { logPath: string | null } | undefined;
      if (!run?.logPath) return;
      return appendFile(run.logPath, `[${new Date().toISOString()}][${stream}] ${chunk}`, "utf8");
    },

    readRunLog(runId: string) {
      const run = getDb().prepare("select log_path as logPath from runs where id = ?").get(runId) as { logPath: string | null } | undefined;
      if (!run?.logPath) return "";
      try {
        return readFileSync(run.logPath, "utf8");
      } catch (err) {
        log.warn(`[database/runs] readFileSync failed for run ${runId} log ${run.logPath}:`, err instanceof Error ? err.message : String(err));
        return "";
      }
    },

    readRunLogChunk(runId: string, offset = 0, limit = 65_536): RunLogChunk {
      const run = getDb().prepare("select log_path as logPath from runs where id = ?").get(runId) as { logPath: string | null } | undefined;
      if (!run?.logPath || !existsSync(run.logPath)) {
        return { content: "", offset: 0, nextOffset: 0, totalBytes: 0, eof: true };
      }

      const totalBytes = statSync(run.logPath).size;
      const safeOffset = Math.max(0, Math.min(offset, totalBytes));
      const safeLimit = Math.max(1_024, Math.min(limit, 262_144));
      const bytesToRead = Math.max(0, Math.min(safeLimit, totalBytes - safeOffset));

      if (bytesToRead === 0) {
        return { content: "", offset: safeOffset, nextOffset: safeOffset, totalBytes, eof: true };
      }

      const fileDescriptor = openSync(run.logPath, "r");
      try {
        const buffer = Buffer.allocUnsafe(bytesToRead);
        const bytesRead = readSync(fileDescriptor, buffer, 0, bytesToRead, safeOffset);
        const safeBytesRead = findUtf8Boundary(buffer, bytesRead);
        const nextOffset = safeOffset + safeBytesRead;
        return {
          content: buffer.subarray(0, safeBytesRead).toString("utf8"),
          offset: safeOffset,
          nextOffset,
          totalBytes,
          eof: nextOffset >= totalBytes,
        };
      } finally {
        closeSync(fileDescriptor);
      }
    },

    markRunStatus(runId: string, status: RunRecord["status"], message: string | null) {
      const ts = nowIso();
      const startedAt = status === "running" ? ts : null;
      transaction(() => {
        getDb()
          .prepare(
            "update runs set status = ?, summary = case when ? is not null then ? else summary end, started_at = coalesce(started_at, ?), updated_at = ? where id = ?",
          )
          .run(status, message, message, startedAt, ts, runId);

        if (status === "running") {
          const run = getDb().prepare("select agent_id from runs where id = ?").get(runId) as { agent_id: string } | undefined;
          if (run) {
            getDb()
              .prepare("update agents set status = 'running', updated_at = ? where id = ? and status in ('idle', 'active')")
              .run(ts, run.agent_id);
          }
        }
      });
    },

    finishRun(input: {
      runId: string;
      status: RunRecord["status"];
      summary: string | null;
      errorMessage: string | null;
      exitCode: number | null;
      signal: string | null;
      model: string | null;
      sessionDisplayId: string | null;
      unattributedCost: boolean;
    }) {
      const run = getDb().prepare("select task_id as taskId, agent_id as agentId, company_id as companyId, connector_id as connectorId from runs where id = ?").get(input.runId) as
        | { taskId: string | null; agentId: string; companyId: string; connectorId: ConnectorId }
        | undefined;
      if (!run) return { budgetHardStopped: false };
      const ts = nowIso();

      // Cost is owned by CostWriter — per-message rows in `cost_entries`.
      // finishRun reads the aggregate instead of writing another row, so
      // there is exactly one source of truth and `cost_entries.run_id`
      // cannot double-count against the previous finishRun INSERT.
      const costRow = getDb()
        .prepare("select coalesce(sum(amount_usd), 0) as total from cost_entries where run_id = ? and unattributed = 0")
        .get(input.runId) as { total: number } | undefined;
      const runCostUsd = costRow?.total ?? 0;

      const agent = getDb().prepare(
        "select budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, status from agents where id = ?",
      ).get(run.agentId) as { budgetMonthlyUsd: number; spentMonthlyUsd: number; status: string } | undefined;
      const nextSpend = (agent?.spentMonthlyUsd ?? 0) + runCostUsd;
      const budgetHardStopped = Boolean(
        agent &&
        agent.budgetMonthlyUsd > 0 &&
        nextSpend >= agent.budgetMonthlyUsd &&
        agent.status !== "paused" &&
        agent.status !== "terminated",
      );
      transaction(() => {
        getDb()
          .prepare(
            "update runs set status = ?, summary = ?, error_message = ?, exit_code = ?, signal = ?, model = ?, session_display_id = ?, cost_usd = ?, unattributed_cost = ?, finished_at = ?, updated_at = ? where id = ?",
          )
          .run(
            input.status, input.summary ?? "", input.errorMessage, input.exitCode,
            input.signal, input.model, input.sessionDisplayId, runCostUsd,
            input.unattributedCost ? 1 : 0, ts, ts, input.runId,
          );
        if (run.taskId !== null) {
          getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, run.taskId);
          getDb()
            .prepare("UPDATE tasks SET checkout_run_id = NULL, execution_locked_at = NULL, updated_at = ? WHERE id = ? AND checkout_run_id = ?")
            .run(ts, run.taskId, input.runId);
        }
        if (runCostUsd > 0) {
          if (run.taskId !== null) {
            getDb()
              .prepare("update tasks set cost_usd = cost_usd + ?, updated_at = ? where id = ?")
              .run(runCostUsd, ts, run.taskId);
          }
          getDb()
            .prepare("update agents set spent_monthly_usd = spent_monthly_usd + ?, updated_at = ? where id = ?")
            .run(runCostUsd, ts, run.agentId);
        }
        if (budgetHardStopped) {
          getDb()
            .prepare("update agents set status = 'paused', updated_at = ? where id = ?")
            .run(ts, run.agentId);
        } else {
          getDb()
            .prepare("update agents set status = 'idle', updated_at = ? where id = ? and status = 'running'")
            .run(ts, run.agentId);
        }
      });
      return { budgetHardStopped };
    },

    cancelRunForDeletion(runId: string, taskId: string | null, workspaceId: string | null, summary: string) {
      const ts = nowIso();
      getDb()
        .prepare(
          "update runs set status = 'cancelled', summary = ?, error_message = null, exit_code = null, signal = null, model = null, session_display_id = null, cost_usd = null, unattributed_cost = 1, finished_at = ?, updated_at = ? where id = ?",
        )
        .run(summary, ts, ts, runId);
      if (taskId !== null) {
        getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, taskId);
        getDb()
          .prepare("UPDATE tasks SET checkout_run_id = NULL, execution_locked_at = NULL, updated_at = ? WHERE id = ? AND checkout_run_id = ?")
          .run(ts, taskId, runId);
      }
      if (workspaceId) {
        getDb()
          .prepare(
            `update workspaces
               set lock_holder_run_id = null, lock_acquired_at = null, state = 'ready', updated_at = ?
             where id = ? and lock_holder_run_id = ?`,
          )
          .run(ts, workspaceId, runId);
      }
    },

    hasActiveRun(filter: { companyId: string; taskId?: string; agentId?: string }) {
      const whereClauses = ["company_id = ?", "status in ('queued', 'running')"];
      const params: SQLInputValue[] = [filter.companyId];
      if (filter.taskId) {
        whereClauses.push("task_id = ?");
        params.push(filter.taskId);
      }
      if (filter.agentId) {
        whereClauses.push("agent_id = ?");
        params.push(filter.agentId);
      }
      const row = getDb().prepare(
        `select 1 as present from runs where ${whereClauses.join(" and ")} limit 1`,
      ).get(...params) as { present: number } | undefined;
      return Boolean(row?.present);
    },

    /**
     * Single writer for `cost_entries`. CostWriter calls this per
     * `message_end`; tests also call it to simulate per-turn cost.
     * finishRun reads SUM(cost_entries) for the run-level aggregate,
     * so this is the authoritative path for cost attribution.
     */
    recordCostEntry(input: {
      runId: string;
      companyId: string;
      agentId: string;
      taskId: string | null;
      connectorId: string;
      amountUsd: number | null;
      unattributed: boolean;
    }): string {
      const id = randomUUID();
      const ts = nowIso();
      const projectRow = input.taskId
        ? getDb().prepare("SELECT project_id FROM tasks WHERE id = ?").get(input.taskId) as { project_id: string | null } | undefined
        : undefined;
      const projectId = projectRow?.project_id ?? null;
      getDb()
        .prepare(
          "insert into cost_entries (id, company_id, run_id, task_id, agent_id, connector_id, amount_usd, unattributed, project_id, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        )
        .run(
          id, input.companyId, input.runId, input.taskId, input.agentId,
          input.connectorId, input.amountUsd,
          input.unattributed ? 1 : 0, projectId, ts,
        );
      return id;
    },

    listQueuedRuns() {
      return readAll<RunRecord>(
        `select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId,
                connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode,
                signal, model, session_display_id as sessionDisplayId, cost_usd as costUsd,
                unattributed_cost as unattributedCost, started_at as startedAt, finished_at as finishedAt,
                log_path as logPath, created_at as createdAt, updated_at as updatedAt
           from runs
          where status = 'queued'
          order by created_at asc`,
      ).map((record) => ({ ...record, unattributedCost: toBoolean(record.unattributedCost) })) as RunRecord[];
    },

    getInterruptedRunsForRecovery(): Array<{
      runId: string; taskId: string; taskTitle: string; agentId: string;
      agentName: string; companyId: string; connectorId: string; interruptedAt: string;
    }> {
      return readAll<{
        runId: string; taskId: string; taskTitle: string; agentId: string;
        agentName: string; companyId: string; connectorId: string; interruptedAt: string;
      }>(
        `select r.id as runId, r.task_id as taskId, t.title as taskTitle,
                r.agent_id as agentId, a.name as agentName,
                r.company_id as companyId,
                r.connector_id as connectorId, r.finished_at as interruptedAt
         from runs r
         join tasks t on t.id = r.task_id
         join agents a on a.id = r.agent_id
         where r.status = 'interrupted'
         and t.status = 'blocked'
         order by r.finished_at desc`
      );
    },

    appendRunEvent(input: {
      companyId: string; runId: string; agentId: string;
      eventType: string; stream?: string | null; level?: string | null;
      message?: string | null; payload?: Record<string, unknown> | null;
    }): void {
      const ts = nowIso();
      transaction(() => {
        const lastSeq = getDb().prepare(
          "SELECT MAX(seq) as maxSeq FROM run_events WHERE run_id = ?"
        ).get(input.runId) as { maxSeq: number | null } | undefined;
        const seq = (lastSeq?.maxSeq ?? 0) + 1;
        getDb().prepare(
          "INSERT INTO run_events (company_id, run_id, agent_id, seq, event_type, stream, level, message, payload, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
        ).run(
          input.companyId, input.runId, input.agentId, seq,
          input.eventType, input.stream ?? null, input.level ?? null,
          input.message ?? null, input.payload ? JSON.stringify(input.payload) : null, ts
        );
      });
    },

    listRunEvents(runId: string, companyId: string, opts?: { offset?: number; limit?: number }): Array<{
      id: number; runId: string; agentId: string; seq: number; eventType: string;
      stream: string | null; level: string | null; message: string | null;
      payload: Record<string, unknown> | null; createdAt: string;
    }> {
      const offset = opts?.offset ?? 0;
      const limit = opts?.limit ?? 500;
      return readAll<SqliteRow>(
        "SELECT id, run_id, agent_id, seq, event_type, stream, level, message, payload, created_at FROM run_events WHERE run_id = ? AND company_id = ? ORDER BY seq ASC LIMIT ? OFFSET ?",
        runId, companyId, limit, offset
      ).map((row: SqliteRow) => ({
        id: row.id,
        runId: row.run_id,
        agentId: row.agent_id,
        seq: row.seq,
        eventType: row.event_type,
        stream: row.stream,
        level: row.level,
        message: row.message,
        payload: row.payload ? JSON.parse(row.payload) : null,
        createdAt: row.created_at,
      }));
    },

    listRunEventsAfter(runId: string, companyId: string, afterSeq: number, limit: number): Array<{
      id: number; runId: string; agentId: string; seq: number; eventType: string;
      stream: string | null; level: string | null; message: string | null;
      payload: Record<string, unknown> | null; createdAt: string;
    }> {
      return readAll<SqliteRow>(
        "SELECT id, run_id, agent_id, seq, event_type, stream, level, message, payload, created_at FROM run_events WHERE run_id = ? AND company_id = ? AND seq > ? ORDER BY seq ASC LIMIT ?",
        runId, companyId, afterSeq, limit
      ).map((row: SqliteRow) => ({
        id: row.id,
        runId: row.run_id,
        agentId: row.agent_id,
        seq: row.seq,
        eventType: row.event_type,
        stream: row.stream,
        level: row.level,
        message: row.message,
        payload: row.payload ? JSON.parse(row.payload) : null,
        createdAt: row.created_at,
      }));
    },

    updateRunActor(runId: string, actorType: string, actorId: string | null): void {
      getDb().prepare(
        "UPDATE runs SET actor_type = ?, actor_id = ?, updated_at = ? WHERE id = ?"
      ).run(actorType, actorId, nowIso(), runId);
    },
  };
}