pikasTech-agentrun/src/mgr/store.ts

import type { BackendProfile, BackendTurnResult, CommandRecord, CreateCommandInput, CreateRunInput, FailureKind, JsonRecord, RunEvent, RunnerRecord, RunRecord, TerminalStatus } from "../common/types.js";
import { AgentRunError } from "../common/errors.js";
import { newId, nowIso, stableHash } from "../common/validation.js";
import { redactJson } from "../common/redaction.js";
import { backendCapabilities } from "../common/backend-profiles.js";

export type MaybePromise<T> = T | Promise<T>;

export interface StoreHealth extends JsonRecord {
  adapter: "memory-self-test" | "postgres";
  ready: boolean;
  reachable: boolean;
  migrationReady: boolean;
  migrationId: string | null;
  failureKind: FailureKind | null;
  message: string | null;
  credentialValuesPrinted: false;
}

export interface AgentRunStore {
  health(): MaybePromise<StoreHealth>;
  createRun(input: CreateRunInput): MaybePromise<RunRecord>;
  getRun(runId: string): MaybePromise<RunRecord>;
  listEvents(runId: string, afterSeq: number, limit: number): MaybePromise<RunEvent[]>;
  createCommand(runId: string, input: CreateCommandInput): MaybePromise<CommandRecord>;
  getCommand(commandId: string): MaybePromise<CommandRecord>;
  listCommands(runId: string, afterSeq: number, limit: number): MaybePromise<CommandRecord[]>;
  registerRunner(input: Partial<RunnerRecord>): MaybePromise<RunnerRecord>;
  claimRun(runId: string, runnerId: string, leaseMs: number): MaybePromise<RunRecord>;
  heartbeat(runId: string, runnerId: string, leaseMs: number): MaybePromise<RunRecord>;
  ackCommand(commandId: string): MaybePromise<CommandRecord>;
  finishCommand(commandId: string, result: Pick<BackendTurnResult, "terminalStatus" | "failureKind" | "failureMessage">): MaybePromise<CommandRecord>;
  appendEvent(runId: string, type: RunEvent["type"], payload: JsonRecord): MaybePromise<RunEvent>;
  finishRun(runId: string, result: Pick<BackendTurnResult, "terminalStatus" | "failureKind" | "failureMessage">): MaybePromise<RunRecord>;
  backends(): MaybePromise<JsonRecord[]>;
  close?(): MaybePromise<void>;
}

export async function openAgentRunStoreFromEnv(env: NodeJS.ProcessEnv = process.env): Promise<AgentRunStore> {
  const databaseUrl = env.DATABASE_URL?.trim();
  if (databaseUrl) {
    const { createPostgresAgentRunStore } = await import("./postgres-store.js");
    return createPostgresAgentRunStore({ connectionString: databaseUrl });
  }
  const storeMode = env.AGENTRUN_STORE ?? env.AGENTRUN_MGR_STORE;
  if (storeMode === "memory") return new MemoryAgentRunStore();
  throw new AgentRunError("infra-failed", "DATABASE_URL is required for agentrun-mgr live runtime; set AGENTRUN_STORE=memory only for explicit self-test/dev mode", { httpStatus: 503, details: { adapter: "postgres", databaseUrl: "missing", memoryFallback: "disabled" } });
}

export class MemoryAgentRunStore implements AgentRunStore {
  private readonly runs = new Map<string, RunRecord>();
  private readonly commands = new Map<string, CommandRecord>();
  private readonly eventsByRun = new Map<string, RunEvent[]>();
  private readonly runners = new Map<string, RunnerRecord>();

  health(): StoreHealth {
    return { adapter: "memory-self-test", ready: true, reachable: true, migrationReady: true, migrationId: "memory-self-test", failureKind: null, message: null, credentialValuesPrinted: false };
  }

  createRun(input: CreateRunInput): RunRecord {
    const at = nowIso();
    const run: RunRecord = { ...input, id: newId("run"), status: "pending", terminalStatus: null, failureKind: null, failureMessage: null, createdAt: at, updatedAt: at, claimedBy: null, leaseExpiresAt: null };
    this.runs.set(run.id, run);
    this.eventsByRun.set(run.id, []);
    this.appendEvent(run.id, "backend_status", { phase: "run-created", backendProfile: run.backendProfile });
    return run;
  }

  getRun(runId: string): RunRecord {
    const run = this.runs.get(runId);
    if (!run) throw new AgentRunError("schema-invalid", `run ${runId} was not found`, { httpStatus: 404 });
    return run;
  }

  listEvents(runId: string, afterSeq: number, limit: number): RunEvent[] {
    this.getRun(runId);
    return (this.eventsByRun.get(runId) ?? []).filter((event) => event.seq > afterSeq).slice(0, Math.max(1, Math.min(limit, 500)));
  }

  createCommand(runId: string, input: CreateCommandInput): CommandRecord {
    this.getRun(runId);
    const payloadHash = stableHash(input.payload);
    if (input.idempotencyKey) {
      const existing = Array.from(this.commands.values()).find((command) => command.runId === runId && command.idempotencyKey === input.idempotencyKey);
      if (existing) {
        if (existing.payloadHash !== payloadHash) throw new AgentRunError("schema-invalid", "idempotency key reused with different payload", { httpStatus: 409 });
        return existing;
      }
    }
    const at = nowIso();
    const seq = Array.from(this.commands.values()).filter((command) => command.runId === runId).length + 1;
    const command: CommandRecord = { ...input, id: newId("cmd"), runId, seq, state: "pending", payloadHash, createdAt: at, updatedAt: at, acknowledgedAt: null };
    this.commands.set(command.id, command);
    this.appendEvent(runId, "backend_status", { phase: "command-created", commandId: command.id, commandType: command.type });
    return command;
  }

  getCommand(commandId: string): CommandRecord {
    const command = this.commands.get(commandId);
    if (!command) throw new AgentRunError("schema-invalid", `command ${commandId} was not found`, { httpStatus: 404 });
    return command;
  }

  listCommands(runId: string, afterSeq: number, limit: number): CommandRecord[] {
    this.getRun(runId);
    return Array.from(this.commands.values()).filter((command) => command.runId === runId && command.seq > afterSeq).slice(0, Math.max(1, Math.min(limit, 100)));
  }

  registerRunner(input: Partial<RunnerRecord>): RunnerRecord {
    const at = nowIso();
    const runner: RunnerRecord = { id: input.id ?? newId("runner"), registeredAt: at, heartbeatAt: at, ...input };
    this.runners.set(runner.id, runner);
    return runner;
  }

  claimRun(runId: string, runnerId: string, leaseMs: number): RunRecord {
    const run = this.getRun(runId);
    if (run.claimedBy && run.claimedBy !== runnerId && run.status !== "completed" && run.status !== "failed" && run.status !== "cancelled") throw new AgentRunError("runner-lease-conflict", `run ${runId} is already claimed`, { httpStatus: 409 });
    const next = this.updateRun(runId, { status: "claimed", claimedBy: runnerId, leaseExpiresAt: new Date(Date.now() + leaseMs).toISOString() });
    this.appendEvent(runId, "backend_status", { phase: "run-claimed", runnerId });
    return next;
  }

  heartbeat(runId: string, runnerId: string, leaseMs: number): RunRecord {
    const run = this.getRun(runId);
    if (run.claimedBy !== runnerId) throw new AgentRunError("runner-lease-conflict", `run ${runId} is not claimed by ${runnerId}`, { httpStatus: 409 });
    return this.updateRun(runId, { leaseExpiresAt: new Date(Date.now() + leaseMs).toISOString() });
  }

  ackCommand(commandId: string): CommandRecord {
    const command = this.getCommand(commandId);
    const next = { ...command, state: "acknowledged" as const, acknowledgedAt: nowIso(), updatedAt: nowIso() };
    this.commands.set(commandId, next);
    return next;
  }

  finishCommand(commandId: string, result: Pick<BackendTurnResult, "terminalStatus" | "failureKind" | "failureMessage">): CommandRecord {
    const command = this.getCommand(commandId);
    const next = { ...command, state: commandStateFromTerminal(result.terminalStatus), updatedAt: nowIso() };
    this.commands.set(commandId, next);
    this.appendEvent(command.runId, "backend_status", { phase: "command-terminal", commandId, state: next.state, terminalStatus: result.terminalStatus, failureKind: result.failureKind });
    return next;
  }

  appendEvent(runId: string, type: RunEvent["type"], payload: JsonRecord): RunEvent {
    this.getRun(runId);
    const events = this.eventsByRun.get(runId) ?? [];
    const event: RunEvent = { id: newId("evt"), runId, seq: events.length + 1, type, payload: redactJson(payload), createdAt: nowIso() };
    events.push(event);
    this.eventsByRun.set(runId, events);
    return event;
  }

  finishRun(runId: string, result: Pick<BackendTurnResult, "terminalStatus" | "failureKind" | "failureMessage">): RunRecord {
    const status = statusFromTerminal(result.terminalStatus);
    const next = this.updateRun(runId, { status, terminalStatus: result.terminalStatus, failureKind: result.failureKind, failureMessage: result.failureMessage });
    this.appendEvent(runId, "terminal_status", { terminalStatus: result.terminalStatus, failureKind: result.failureKind, message: result.failureMessage });
    return next;
  }

  backends(): JsonRecord[] {
    return backendCapabilities();
  }

  private updateRun(runId: string, patch: Partial<RunRecord>): RunRecord {
    const run = this.getRun(runId);
    const next = { ...run, ...patch, updatedAt: nowIso() };
    this.runs.set(runId, next);
    return next;
  }
}

export function statusFromTerminal(terminalStatus: TerminalStatus): RunRecord["status"] {
  if (terminalStatus === "completed") return "completed";
  if (terminalStatus === "cancelled") return "cancelled";
  if (terminalStatus === "blocked") return "blocked";
  return "failed";
}

export function commandStateFromTerminal(terminalStatus: TerminalStatus): CommandRecord["state"] {
  if (terminalStatus === "completed") return "completed";
  if (terminalStatus === "cancelled") return "cancelled";
  return "failed";
}