hikari-desktop/src/lib/stores/agents.ts

import { writable, derived } from "svelte/store";
import type { AgentInfo } from "$lib/types/agents";
import { assignCharacter } from "$lib/utils/agentCharacters";

// Map of conversation ID -> agents in that conversation
const agentsByConversation = writable<Record<string, AgentInfo[]>>({});

function createAgentStore() {
  return {
    subscribe: agentsByConversation.subscribe,

    addAgent(conversationId: string, agent: Omit<AgentInfo, "characterName" | "characterAvatar">) {
      agentsByConversation.update((state) => {
        const existing = state[conversationId] || [];
        const activeNames = existing.map((a) => a.characterName);
        const character = assignCharacter(activeNames);
        return {
          ...state,
          [conversationId]: [
            ...existing,
            { ...agent, characterName: character.name, characterAvatar: character.avatar },
          ],
        };
      });
    },

    updateAgentId(conversationId: string, toolUseId: string, agentId: string, agentType?: string) {
      agentsByConversation.update((state) => {
        const agents = state[conversationId];
        if (!agents) return state;

        const agentIndex = agents.findIndex((a) => a.toolUseId === toolUseId);
        if (agentIndex === -1) return state;

        const updated = [...agents];
        updated[agentIndex] = {
          ...updated[agentIndex],
          agentId,
          ...(agentType !== undefined ? { agentType } : {}),
        };

        return {
          ...state,
          [conversationId]: updated,
        };
      });
    },

    endAgent(
      conversationId: string,
      toolUseId: string,
      endedAt: number,
      isError: boolean,
      lastAssistantMessage?: string
    ) {
      agentsByConversation.update((state) => {
        const agents = state[conversationId];
        if (!agents) return state;

        const agentIndex = agents.findIndex((a) => a.toolUseId === toolUseId);
        if (agentIndex === -1) return state;

        const updated = [...agents];
        const agent = updated[agentIndex];
        const durationMs = endedAt - agent.startedAt;

        updated[agentIndex] = {
          ...agent,
          endedAt,
          status: isError ? "errored" : "completed",
          durationMs,
          lastAssistantMessage,
        };

        return {
          ...state,
          [conversationId]: updated,
        };
      });
    },

    markAllErrored(conversationId: string) {
      agentsByConversation.update((state) => {
        const agents = state[conversationId];
        if (!agents) return state;

        const now = Date.now();
        const updated = agents.map((agent) =>
          agent.status === "running"
            ? { ...agent, endedAt: now, status: "errored" as const }
            : agent
        );

        return {
          ...state,
          [conversationId]: updated,
        };
      });
    },

    clearCompleted(conversationId: string) {
      agentsByConversation.update((state) => {
        const agents = state[conversationId];
        if (!agents) return state;

        return {
          ...state,
          [conversationId]: agents.filter((a) => a.status === "running"),
        };
      });
    },

    clearConversation(conversationId: string) {
      agentsByConversation.update((state) => {
        // eslint-disable-next-line @typescript-eslint/no-unused-vars -- Unused destructured value
        const { [conversationId]: _, ...rest } = state;
        return rest;
      });
    },
  };
}

export const agentStore = createAgentStore();

export function getAgentsForConversation(conversationId: string) {
  return derived(agentsByConversation, ($state) => $state[conversationId] || []);
}

export const runningAgentCount = derived(agentsByConversation, ($state) => {
  let count = 0;
  for (const agents of Object.values($state)) {
    count += agents.filter((a) => a.status === "running").length;
  }
  return count;
});