claude-code/packages/remote-control-server/src/store.ts

import { v4 as uuid } from "uuid";

// ---------- Types ----------

export interface UserRecord {
  username: string;
  createdAt: Date;
}

export interface EnvironmentRecord {
  id: string;
  secret: string;
  machineName: string | null;
  directory: string | null;
  branch: string | null;
  gitRepoUrl: string | null;
  maxSessions: number;
  workerType: string;
  bridgeId: string | null;
  capabilities: Record<string, unknown> | null;
  status: string;
  username: string | null;
  lastPollAt: Date | null;
  createdAt: Date;
  updatedAt: Date;
}

export interface SessionRecord {
  id: string;
  environmentId: string | null;
  title: string | null;
  status: string;
  source: string;
  permissionMode: string | null;
  workerEpoch: number;
  username: string | null;
  createdAt: Date;
  updatedAt: Date;
}

export interface WorkItemRecord {
  id: string;
  environmentId: string;
  sessionId: string;
  state: string;
  secret: string;
  createdAt: Date;
  updatedAt: Date;
}

export interface SessionWorkerRecord {
  sessionId: string;
  workerStatus: string | null;
  externalMetadata: Record<string, unknown> | null;
  requiresActionDetails: Record<string, unknown> | null;
  lastHeartbeatAt: Date | null;
  createdAt: Date;
  updatedAt: Date;
}

// ---------- Stores (in-memory Maps) ----------

const users = new Map<string, UserRecord>();
const tokenToUser = new Map<string, { username: string; createdAt: Date }>();
const environments = new Map<string, EnvironmentRecord>();
const sessions = new Map<string, SessionRecord>();
const workItems = new Map<string, WorkItemRecord>();
const sessionWorkers = new Map<string, SessionWorkerRecord>();

// UUID → session ownership: sessionId → Set of UUIDs
const sessionOwners = new Map<string, Set<string>>();

// ---------- User ----------

export function storeCreateUser(username: string): UserRecord {
  const existing = users.get(username);
  if (existing) return existing;
  const record: UserRecord = { username, createdAt: new Date() };
  users.set(username, record);
  return record;
}

export function storeGetUser(username: string): UserRecord | undefined {
  return users.get(username);
}

export function storeCreateToken(username: string, token: string): void {
  tokenToUser.set(token, { username, createdAt: new Date() });
}

export function storeGetUserByToken(token: string): { username: string; createdAt: Date } | undefined {
  return tokenToUser.get(token);
}

export function storeDeleteToken(token: string): boolean {
  return tokenToUser.delete(token);
}

// ---------- Environment ----------

/** Find an active or offline environment by machineName (optionally filtered by workerType).
 *  Includes "offline" so ACP agents can be reused on reconnect. */
export function storeFindEnvironmentByMachineName(
  machineName: string,
  workerType?: string,
): EnvironmentRecord | undefined {
  for (const rec of environments.values()) {
    if (rec.machineName === machineName && (rec.status === "active" || rec.status === "offline")) {
      if (!workerType || rec.workerType === workerType) {
        return rec;
      }
    }
  }
  return undefined;
}

export function storeCreateEnvironment(req: {
  secret: string;
  machineName?: string;
  directory?: string;
  branch?: string;
  gitRepoUrl?: string;
  maxSessions?: number;
  workerType?: string;
  bridgeId?: string;
  username?: string;
  capabilities?: Record<string, unknown>;
}): EnvironmentRecord {
  // ACP: reuse existing active record by machineName
  if (req.workerType === "acp" && req.machineName) {
    const existing = storeFindEnvironmentByMachineName(req.machineName, "acp");
    if (existing) {
      Object.assign(existing, {
        status: "active",
        lastPollAt: new Date(),
        updatedAt: new Date(),
        maxSessions: req.maxSessions ?? existing.maxSessions,
        bridgeId: req.bridgeId ?? existing.bridgeId,
        capabilities: req.capabilities ?? existing.capabilities,
        username: req.username ?? existing.username,
      });
      return existing;
    }
  }

  const id = `env_${uuid().replace(/-/g, "")}`;
  const now = new Date();
  const record: EnvironmentRecord = {
    id,
    secret: req.secret,
    machineName: req.machineName ?? null,
    directory: req.directory ?? null,
    branch: req.branch ?? null,
    gitRepoUrl: req.gitRepoUrl ?? null,
    maxSessions: req.maxSessions ?? 1,
    workerType: req.workerType ?? "claude_code",
    bridgeId: req.bridgeId ?? null,
    capabilities: req.capabilities ?? null,
    status: "active",
    username: req.username ?? null,
    lastPollAt: now,
    createdAt: now,
    updatedAt: now,
  };
  environments.set(id, record);
  return record;
}

export function storeGetEnvironment(id: string): EnvironmentRecord | undefined {
  return environments.get(id);
}

export function storeUpdateEnvironment(id: string, patch: Partial<Pick<EnvironmentRecord, "status" | "lastPollAt" | "updatedAt" | "capabilities" | "machineName" | "maxSessions" | "bridgeId">>): boolean {
  const rec = environments.get(id);
  if (!rec) return false;
  Object.assign(rec, patch, { updatedAt: new Date() });
  return true;
}

export function storeListActiveEnvironments(): EnvironmentRecord[] {
  return [...environments.values()].filter((e) => e.status === "active");
}

export function storeListActiveEnvironmentsByUsername(username: string): EnvironmentRecord[] {
  return [...environments.values()].filter((e) => e.status === "active" && e.username === username);
}

// ---------- Session ----------

export function storeCreateSession(req: {
  environmentId?: string | null;
  title?: string | null;
  source?: string;
  permissionMode?: string | null;
  idPrefix?: string;
  username?: string | null;
}): SessionRecord {
  const id = `${req.idPrefix || "session_"}${uuid().replace(/-/g, "")}`;
  const now = new Date();
  const record: SessionRecord = {
    id,
    environmentId: req.environmentId ?? null,
    title: req.title ?? null,
    status: "idle",
    source: req.source ?? "remote-control",
    permissionMode: req.permissionMode ?? null,
    workerEpoch: 0,
    username: req.username ?? null,
    createdAt: now,
    updatedAt: now,
  };
  sessions.set(id, record);
  return record;
}

export function storeGetSession(id: string): SessionRecord | undefined {
  return sessions.get(id);
}

export function storeUpdateSession(id: string, patch: Partial<Pick<SessionRecord, "title" | "status" | "workerEpoch" | "updatedAt">>): boolean {
  const rec = sessions.get(id);
  if (!rec) return false;
  Object.assign(rec, patch, { updatedAt: new Date() });
  return true;
}

export function storeListSessions(): SessionRecord[] {
  return [...sessions.values()];
}

export function storeListSessionsByUsername(username: string): SessionRecord[] {
  return [...sessions.values()].filter((s) => s.username === username);
}

export function storeListSessionsByEnvironment(envId: string): SessionRecord[] {
  return [...sessions.values()].filter((s) => s.environmentId === envId);
}

export function storeDeleteSession(id: string): boolean {
  sessionWorkers.delete(id);
  return sessions.delete(id);
}

// ---------- Session Worker ----------

export function storeGetSessionWorker(sessionId: string): SessionWorkerRecord | undefined {
  return sessionWorkers.get(sessionId);
}

export function storeUpsertSessionWorker(sessionId: string, patch: {
  workerStatus?: string | null;
  externalMetadata?: Record<string, unknown> | null;
  requiresActionDetails?: Record<string, unknown> | null;
  lastHeartbeatAt?: Date | null;
}): SessionWorkerRecord {
  const now = new Date();
  const existing = sessionWorkers.get(sessionId);
  const record: SessionWorkerRecord = existing ?? {
    sessionId,
    workerStatus: null,
    externalMetadata: null,
    requiresActionDetails: null,
    lastHeartbeatAt: null,
    createdAt: now,
    updatedAt: now,
  };

  if (patch.workerStatus !== undefined) {
    record.workerStatus = patch.workerStatus;
  }
  if (patch.externalMetadata !== undefined) {
    if (patch.externalMetadata === null) {
      record.externalMetadata = null;
    } else {
      record.externalMetadata = {
        ...(record.externalMetadata ?? {}),
        ...patch.externalMetadata,
      };
    }
  }
  if (patch.requiresActionDetails !== undefined) {
    record.requiresActionDetails = patch.requiresActionDetails;
  }
  if (patch.lastHeartbeatAt !== undefined) {
    record.lastHeartbeatAt = patch.lastHeartbeatAt;
  }
  record.updatedAt = now;

  sessionWorkers.set(sessionId, record);
  return record;
}

// ---------- Work Items ----------

// ---------- Session Ownership (UUID-based) ----------

export function storeBindSession(sessionId: string, uuid: string): void {
  let owners = sessionOwners.get(sessionId);
  if (!owners) {
    owners = new Set();
    sessionOwners.set(sessionId, owners);
  }
  owners.add(uuid);
}

export function storeIsSessionOwner(sessionId: string, uuid: string): boolean {
  const owners = sessionOwners.get(sessionId);
  return owners ? owners.has(uuid) : false;
}

export function storeGetSessionOwners(sessionId: string): Set<string> | undefined {
  return sessionOwners.get(sessionId);
}

export function storeListSessionsByOwnerUuid(uuid: string): SessionRecord[] {
  const result: SessionRecord[] = [];
  const resultIds = new Set<string>();

  // Collect sessions already owned by this UUID
  for (const [sessionId, owners] of sessionOwners) {
    if (owners.has(uuid)) {
      const session = sessions.get(sessionId);
      if (session) {
        result.push(session);
        resultIds.add(sessionId);
      }
    }
  }

  // Auto-bind orphaned sessions (no owner — typically ACP agent sessions created via REST registration)
  for (const [sessionId, session] of sessions) {
    if (resultIds.has(sessionId)) continue;
    const owners = sessionOwners.get(sessionId);
    // No owners map entry at all, or empty owners set
    const isOrphaned = !owners || owners.size === 0;
    if (isOrphaned) {
      storeBindSession(sessionId, uuid);
      result.push(session);
      resultIds.add(sessionId);
    }
  }

  return result;
}

// ---------- Work Items (cont.) ----------

export function storeCreateWorkItem(req: {
  environmentId: string;
  sessionId: string;
  secret: string;
}): WorkItemRecord {
  const id = `work_${uuid().replace(/-/g, "")}`;
  const now = new Date();
  const record: WorkItemRecord = {
    id,
    environmentId: req.environmentId,
    sessionId: req.sessionId,
    state: "pending",
    secret: req.secret,
    createdAt: now,
    updatedAt: now,
  };
  workItems.set(id, record);
  return record;
}

export function storeGetWorkItem(id: string): WorkItemRecord | undefined {
  return workItems.get(id);
}

export function storeGetPendingWorkItem(environmentId: string): WorkItemRecord | undefined {
  for (const item of workItems.values()) {
    if (item.environmentId === environmentId && item.state === "pending") {
      return item;
    }
  }
  return undefined;
}

export function storeUpdateWorkItem(id: string, patch: Partial<Pick<WorkItemRecord, "state" | "updatedAt">>): boolean {
  const rec = workItems.get(id);
  if (!rec) return false;
  Object.assign(rec, patch, { updatedAt: new Date() });
  return true;
}

// ---------- ACP Agent (reuses EnvironmentRecord with workerType="acp") ----------

/** List all ACP agents (environments with workerType="acp") */
export function storeListAcpAgents(): EnvironmentRecord[] {
  return [...environments.values()].filter((e) => e.workerType === "acp");
}

/** List ACP agents by channel group (stored in bridgeId field) */
export function storeListAcpAgentsByChannelGroup(channelGroupId: string): EnvironmentRecord[] {
  return [...environments.values()].filter(
    (e) => e.workerType === "acp" && e.bridgeId === channelGroupId,
  );
}

/** List online ACP agents */
export function storeListOnlineAcpAgents(): EnvironmentRecord[] {
  return [...environments.values()].filter(
    (e) => e.workerType === "acp" && e.status === "active",
  );
}

/** Mark an ACP agent as offline */
export function storeMarkAcpAgentOffline(id: string): boolean {
  const rec = environments.get(id);
  if (!rec || rec.workerType !== "acp") return false;
  Object.assign(rec, { status: "offline", updatedAt: new Date() });
  return true;
}

/** Mark an ACP agent as online (on reconnect) */
export function storeMarkAcpAgentOnline(id: string): boolean {
  const rec = environments.get(id);
  if (!rec || rec.workerType !== "acp") return false;
  Object.assign(rec, { status: "active", lastPollAt: new Date(), updatedAt: new Date() });
  return true;
}

// ---------- Reset (for tests) ----------

export function storeReset() {
  users.clear();
  tokenToUser.clear();
  environments.clear();
  sessions.clear();
  workItems.clear();
  sessionWorkers.clear();
  sessionOwners.clear();
}