zn-ai/electron/utils/agent-config.ts

import * as fs from 'fs';
import * as path from 'path';
import type { ProviderAccount } from '@runtime/lib/providers';
import { DEFAULT_AGENT_ID, DEFAULT_MAIN_SESSION_SUFFIX, type AgentSummary, type AgentsSnapshot } from '@runtime/lib/agents';
import { buildMainSessionKey, normalizeAgentId } from '@runtime/lib/models';
import { getUserDataDir } from './paths';

interface StoredAgentEntry {
  id: string;
  name: string;
  providerAccountId?: string | null;
  modelRef?: string | null;
  workspace?: string | null;
  agentDir?: string | null;
  channelTypes?: string[];
  createdAt?: string;
  updatedAt?: string;
}

interface StoredAgentsConfig {
  agents: StoredAgentEntry[];
  channelOwners?: Record<string, string>;
  channelAccountOwners?: Record<string, string>;
  mainSessionSuffix?: string;
}

const STORE_FILE_NAME = 'agents.json';

function getStorePath(): string {
  return path.join(getUserDataDir(), STORE_FILE_NAME);
}

function getAgentsRootDir(): string {
  return path.join(getUserDataDir(), 'agents');
}

function getMainWorkspacePath(): string {
  return path.join(getAgentsRootDir(), DEFAULT_AGENT_ID, 'workspace');
}

function getMainAgentDirPath(): string {
  return path.join(getAgentsRootDir(), DEFAULT_AGENT_ID, 'agent');
}

function getAgentWorkspacePath(agentId: string): string {
  return path.join(getAgentsRootDir(), agentId, 'workspace');
}

function getAgentDirPath(agentId: string): string {
  return path.join(getAgentsRootDir(), agentId, 'agent');
}

function ensureDir(dirPath: string | null | undefined): void {
  if (!dirPath) return;
  fs.mkdirSync(dirPath, { recursive: true });
}

function readStore(): StoredAgentsConfig {
  try {
    const filePath = getStorePath();
    if (fs.existsSync(filePath)) {
      const parsed = JSON.parse(fs.readFileSync(filePath, 'utf-8')) as Partial<StoredAgentsConfig>;
      return {
        agents: Array.isArray(parsed.agents) ? parsed.agents : [],
        channelOwners: parsed.channelOwners && typeof parsed.channelOwners === 'object' ? parsed.channelOwners : {},
        channelAccountOwners: parsed.channelAccountOwners && typeof parsed.channelAccountOwners === 'object'
          ? parsed.channelAccountOwners
          : {},
        mainSessionSuffix: typeof parsed.mainSessionSuffix === 'string' ? parsed.mainSessionSuffix : DEFAULT_MAIN_SESSION_SUFFIX,
      };
    }
  } catch {
    // Fall back to an empty store on malformed JSON.
  }

  return {
    agents: [],
    channelOwners: {},
    channelAccountOwners: {},
    mainSessionSuffix: DEFAULT_MAIN_SESSION_SUFFIX,
  };
}

function writeStore(store: StoredAgentsConfig): void {
  const filePath = getStorePath();
  fs.mkdirSync(path.dirname(filePath), { recursive: true });
  fs.writeFileSync(filePath, JSON.stringify(store, null, 2), 'utf-8');
}

function formatModelDisplay(modelRef: string | null | undefined, fallbackLabel: string): string {
  const trimmed = String(modelRef ?? '').trim();
  if (!trimmed) return fallbackLabel;

  const parts = trimmed.split('/');
  return parts[parts.length - 1] || trimmed;
}

function normalizeAgentName(name: string): string {
  return name.trim() || 'Agent';
}

function slugifyAgentId(name: string): string {
  const normalized = name
    .normalize('NFKD')
    .replace(/[^\w\s-]/g, '')
    .toLowerCase()
    .replace(/[_\s]+/g, '-')
    .replace(/-+/g, '-')
    .replace(/^-|-$/g, '');

  if (!normalized || /^\d+$/.test(normalized)) return 'agent';
  if (normalized === DEFAULT_AGENT_ID) return 'agent';
  return normalized;
}

function buildUniqueAgentId(existingIds: Set<string>, name: string): string {
  const baseId = slugifyAgentId(name);
  let nextId = baseId;
  let index = 2;

  while (existingIds.has(nextId)) {
    nextId = `${baseId}-${index}`;
    index += 1;
  }

  return nextId;
}

function getDefaultAccount(accounts: ProviderAccount[], defaultAccountId: string | null): ProviderAccount | null {
  return accounts.find((account) => account.id === defaultAccountId) ?? accounts[0] ?? null;
}

function buildMainAgent(defaultAccount: ProviderAccount | null, mainSessionSuffix: string): AgentSummary {
  ensureDir(getMainWorkspacePath());
  ensureDir(getMainAgentDirPath());

  return {
    id: DEFAULT_AGENT_ID,
    name: 'Main Agent',
    isDefault: true,
    providerAccountId: defaultAccount?.id ?? null,
    modelRef: defaultAccount?.model ?? null,
    modelDisplay: formatModelDisplay(defaultAccount?.model, defaultAccount?.label || 'Unassigned'),
    mainSessionKey: buildMainSessionKey(DEFAULT_AGENT_ID, mainSessionSuffix),
    vendorId: defaultAccount?.vendorId ?? null,
    source: 'synthetic-main',
    overrideModelRef: null,
    inheritedModel: true,
    workspace: getMainWorkspacePath(),
    agentDir: getMainAgentDirPath(),
    channelTypes: [],
  };
}

function collectChannelOwners(
  store: StoredAgentsConfig,
  entries: StoredAgentEntry[],
): { channelOwners: Record<string, string>; channelAccountOwners: Record<string, string> } {
  const explicitChannelOwners = store.channelOwners && typeof store.channelOwners === 'object'
    ? { ...store.channelOwners }
    : {};
  const explicitChannelAccountOwners = store.channelAccountOwners && typeof store.channelAccountOwners === 'object'
    ? { ...store.channelAccountOwners }
    : {};

  for (const entry of entries) {
    const normalizedId = normalizeAgentId(entry.id);
    const channels = Array.isArray(entry.channelTypes) ? entry.channelTypes.filter(Boolean) : [];
    for (const channelType of channels) {
      if (!explicitChannelOwners[channelType]) {
        explicitChannelOwners[channelType] = normalizedId;
      }
    }
  }

  return {
    channelOwners: explicitChannelOwners,
    channelAccountOwners: explicitChannelAccountOwners,
  };
}

function mapStoredAgentToSummary(
  entry: StoredAgentEntry,
  accounts: ProviderAccount[],
  defaultAccount: ProviderAccount | null,
  mainSessionSuffix: string,
): AgentSummary {
  const normalizedId = normalizeAgentId(entry.id);
  const configuredAccount = entry.providerAccountId
    ? accounts.find((account) => account.id === entry.providerAccountId) ?? null
    : null;
  const effectiveAccount = configuredAccount ?? defaultAccount;
  const effectiveModelRef = entry.modelRef ?? effectiveAccount?.model ?? defaultAccount?.model ?? null;
  const workspace = entry.workspace || getAgentWorkspacePath(normalizedId);
  const agentDir = entry.agentDir || getAgentDirPath(normalizedId);

  ensureDir(agentDir);
  ensureDir(workspace);

  return {
    id: normalizedId,
    name: normalizeAgentName(entry.name),
    isDefault: false,
    providerAccountId: entry.providerAccountId ?? null,
    modelRef: effectiveModelRef,
    modelDisplay: formatModelDisplay(effectiveModelRef, normalizeAgentName(entry.name)),
    mainSessionKey: buildMainSessionKey(normalizedId, mainSessionSuffix),
    vendorId: configuredAccount?.vendorId ?? effectiveAccount?.vendorId ?? null,
    source: 'agent-config',
    overrideModelRef: entry.modelRef ?? null,
    inheritedModel: !entry.modelRef,
    workspace,
    agentDir,
    channelTypes: Array.isArray(entry.channelTypes) ? entry.channelTypes.filter(Boolean) : [],
  };
}

function buildSnapshotFromStore(
  store: StoredAgentsConfig,
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  const defaultAccount = getDefaultAccount(accounts, defaultAccountId);
  const mainSessionSuffix = store.mainSessionSuffix || DEFAULT_MAIN_SESSION_SUFFIX;
  const normalizedEntries = Array.from(
    new Map(
      store.agents
        .filter((entry) => Boolean(entry) && typeof entry.id === 'string' && entry.id.trim().length > 0)
        .map((entry) => [normalizeAgentId(entry.id), { ...entry, id: normalizeAgentId(entry.id) }]),
    ).values(),
  ).filter((entry) => entry.id !== DEFAULT_AGENT_ID);

  const { channelOwners, channelAccountOwners } = collectChannelOwners(store, normalizedEntries);
  const agents = [
    buildMainAgent(defaultAccount, mainSessionSuffix),
    ...normalizedEntries.map((entry) => mapStoredAgentToSummary(entry, accounts, defaultAccount, mainSessionSuffix)),
  ];

  const configuredChannelTypes = Array.from(new Set([
    ...Object.keys(channelOwners),
    ...Object.keys(channelAccountOwners).map((key) => key.split(':')[0]).filter(Boolean),
  ]));

  return {
    agents,
    models: agents,
    defaultAgentId: DEFAULT_AGENT_ID,
    defaultProviderAccountId: defaultAccount?.id ?? null,
    defaultModelRef: defaultAccount?.model ?? null,
    mainSessionSuffix,
    configuredChannelTypes,
    channelOwners,
    channelAccountOwners,
  };
}

function syncAgentChannelMembership(store: StoredAgentsConfig, channelType: string): void {
  const normalizedChannelType = String(channelType ?? '').trim();
  if (!normalizedChannelType) return;

  const ownerIds = new Set<string>();
  const channelOwnerId = store.channelOwners?.[normalizedChannelType];
  if (channelOwnerId) {
    ownerIds.add(normalizeAgentId(channelOwnerId));
  }

  for (const [key, ownerId] of Object.entries(store.channelAccountOwners ?? {})) {
    if (!key.startsWith(`${normalizedChannelType}:`) || !ownerId) continue;
    ownerIds.add(normalizeAgentId(ownerId));
  }

  store.agents = store.agents.map((entry) => {
    const normalizedId = normalizeAgentId(entry.id);
    const currentChannelTypes = Array.isArray(entry.channelTypes) ? entry.channelTypes.filter(Boolean) : [];
    const hasChannel = currentChannelTypes.includes(normalizedChannelType);
    const shouldHaveChannel = ownerIds.has(normalizedId);

    if (hasChannel === shouldHaveChannel) {
      return entry;
    }

    return {
      ...entry,
      channelTypes: shouldHaveChannel
        ? [...currentChannelTypes, normalizedChannelType]
        : currentChannelTypes.filter((value) => value !== normalizedChannelType),
      updatedAt: new Date().toISOString(),
    };
  });
}

export function listAgentsSnapshot(
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  return buildSnapshotFromStore(readStore(), accounts, defaultAccountId);
}

export function createAgentConfig(
  name: string,
  options: { inheritWorkspace?: boolean } | undefined,
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  const trimmedName = normalizeAgentName(name);
  const store = readStore();
  const existingIds = new Set(store.agents.map((entry) => normalizeAgentId(entry.id)));
  const agentId = buildUniqueAgentId(existingIds, trimmedName);
  const now = new Date().toISOString();
  const workspace = options?.inheritWorkspace ? getMainWorkspacePath() : getAgentWorkspacePath(agentId);
  const agentDir = getAgentDirPath(agentId);

  ensureDir(workspace);
  ensureDir(agentDir);

  store.agents = [
    ...store.agents,
    {
      id: agentId,
      name: trimmedName,
      providerAccountId: null,
      modelRef: null,
      workspace,
      agentDir,
      channelTypes: [],
      createdAt: now,
      updatedAt: now,
    },
  ];

  writeStore(store);
  return buildSnapshotFromStore(store, accounts, defaultAccountId);
}

export function updateAgentName(
  agentId: string,
  name: string,
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  const normalizedId = normalizeAgentId(agentId);
  if (normalizedId === DEFAULT_AGENT_ID) {
    throw new Error('Main Agent cannot be renamed');
  }

  const store = readStore();
  const index = store.agents.findIndex((entry) => normalizeAgentId(entry.id) === normalizedId);
  if (index === -1) {
    throw new Error(`Agent "${normalizedId}" not found`);
  }

  store.agents[index] = {
    ...store.agents[index],
    name: normalizeAgentName(name),
    updatedAt: new Date().toISOString(),
  };
  writeStore(store);
  return buildSnapshotFromStore(store, accounts, defaultAccountId);
}

export function updateAgentModelConfig(
  agentId: string,
  modelRef: string | null,
  providerAccountId: string | null | undefined,
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  const normalizedId = normalizeAgentId(agentId);
  if (normalizedId === DEFAULT_AGENT_ID) {
    throw new Error('Main Agent model is managed from Models');
  }

  const store = readStore();
  const index = store.agents.findIndex((entry) => normalizeAgentId(entry.id) === normalizedId);
  if (index === -1) {
    throw new Error(`Agent "${normalizedId}" not found`);
  }

  const trimmedModelRef = typeof modelRef === 'string' ? modelRef.trim() : '';
  const inferredAccountId = providerAccountId
    ?? (
      trimmedModelRef
        ? accounts.find((account) => account.model === trimmedModelRef)?.id ?? store.agents[index].providerAccountId ?? null
        : null
    );

  store.agents[index] = {
    ...store.agents[index],
    providerAccountId: inferredAccountId ?? null,
    modelRef: trimmedModelRef || null,
    updatedAt: new Date().toISOString(),
  };
  writeStore(store);
  return buildSnapshotFromStore(store, accounts, defaultAccountId);
}

export function deleteAgentConfig(
  agentId: string,
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  const normalizedId = normalizeAgentId(agentId);
  if (normalizedId === DEFAULT_AGENT_ID) {
    throw new Error('Main Agent cannot be deleted');
  }

  const store = readStore();
  const existingAgent = store.agents.find((entry) => normalizeAgentId(entry.id) === normalizedId);
  if (!existingAgent) {
    throw new Error(`Agent "${normalizedId}" not found`);
  }

  store.agents = store.agents.filter((entry) => normalizeAgentId(entry.id) !== normalizedId);
  store.channelOwners = Object.fromEntries(
    Object.entries(store.channelOwners ?? {}).filter(([, ownerId]) => normalizeAgentId(ownerId) !== normalizedId),
  );
  store.channelAccountOwners = Object.fromEntries(
    Object.entries(store.channelAccountOwners ?? {}).filter(([, ownerId]) => normalizeAgentId(ownerId) !== normalizedId),
  );
  writeStore(store);

  const agentRootDir = path.join(getAgentsRootDir(), normalizedId);
  try {
    fs.rmSync(agentRootDir, { recursive: true, force: true });
  } catch {
    // Best effort cleanup only.
  }

  return buildSnapshotFromStore(store, accounts, defaultAccountId);
}

export function assignChannelToAgent(
  agentId: string,
  channelType: string,
  accountId: string | null | undefined,
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  const normalizedId = normalizeAgentId(agentId);
  const normalizedChannelType = String(channelType ?? '').trim();
  if (!normalizedChannelType) {
    throw new Error('channelType is required');
  }

  const store = readStore();
  if (normalizedId !== DEFAULT_AGENT_ID && !store.agents.some((entry) => normalizeAgentId(entry.id) === normalizedId)) {
    throw new Error(`Agent "${normalizedId}" not found`);
  }

  store.channelOwners = {
    ...(store.channelOwners ?? {}),
    ...(accountId ? {} : { [normalizedChannelType]: normalizedId }),
  };

  if (accountId) {
    store.channelAccountOwners = {
      ...(store.channelAccountOwners ?? {}),
      [`${normalizedChannelType}:${accountId}`]: normalizedId,
    };
  }

  syncAgentChannelMembership(store, normalizedChannelType);

  writeStore(store);
  return buildSnapshotFromStore(store, accounts, defaultAccountId);
}

export function clearChannelBinding(
  channelType: string,
  accountId: string | null | undefined,
  accounts: ProviderAccount[],
  defaultAccountId: string | null,
): AgentsSnapshot {
  const normalizedChannelType = String(channelType ?? '').trim();
  if (!normalizedChannelType) {
    throw new Error('channelType is required');
  }

  const store = readStore();
  const ownerId = accountId
    ? store.channelAccountOwners?.[`${normalizedChannelType}:${accountId}`]
    : store.channelOwners?.[normalizedChannelType];

  if (accountId) {
    const nextAccountOwners = { ...(store.channelAccountOwners ?? {}) };
    delete nextAccountOwners[`${normalizedChannelType}:${accountId}`];
    store.channelAccountOwners = nextAccountOwners;
  } else {
    const nextChannelOwners = { ...(store.channelOwners ?? {}) };
    delete nextChannelOwners[normalizedChannelType];
    store.channelOwners = nextChannelOwners;
  }

  if (ownerId) {
    syncAgentChannelMembership(store, normalizedChannelType);
  }

  writeStore(store);
  return buildSnapshotFromStore(store, accounts, defaultAccountId);
}