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

import { describe, it, expect, beforeEach } from "vitest";
import { agentStore, getAgentsForConversation, runningAgentCount } from "./agents";
import { get } from "svelte/store";
import type { AgentInfo } from "$lib/types/agents";
import { CHARACTER_POOL } from "$lib/utils/agentCharacters";

describe("agents store", () => {
  const conversationId = "test-conversation-1";
  const otherConversationId = "test-conversation-2";

  type AgentInput = Omit<AgentInfo, "characterName" | "characterAvatar">;

  const createMockAgent = (overrides?: Partial<AgentInput>): AgentInput => ({
    toolUseId: "toolu_test123",
    description: "Test agent",
    subagentType: "Explore",
    startedAt: Date.now(),
    status: "running",
    ...overrides,
  });

  beforeEach(() => {
    // Clear all conversations by subscribing and getting state
    let state: Record<string, AgentInfo[]> = {};
    const unsub = agentStore.subscribe((s) => {
      state = s;
    });
    unsub();

    // Clear each conversation
    for (const convId of Object.keys(state)) {
      agentStore.clearConversation(convId);
    }
  });

  describe("addAgent", () => {
    it("adds an agent to a conversation", () => {
      const agent = createMockAgent();
      agentStore.addAgent(conversationId, agent);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents).toHaveLength(1);
      expect(agents[0]).toMatchObject(agent);
    });

    it("preserves model field when provided", () => {
      const agent = createMockAgent({ model: "claude-opus-4-6" });
      agentStore.addAgent(conversationId, agent);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].model).toBe("claude-opus-4-6");
    });

    it("leaves model undefined when not provided", () => {
      const agent = createMockAgent();
      agentStore.addAgent(conversationId, agent);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].model).toBeUndefined();
    });

    it("assigns a character name and avatar to added agents", () => {
      const agent = createMockAgent();
      agentStore.addAgent(conversationId, agent);

      const agents = get(getAgentsForConversation(conversationId));
      const validNames = CHARACTER_POOL.map((c) => c.name);
      expect(validNames).toContain(agents[0].characterName);
      expect(agents[0].characterAvatar).toMatch(/^https:\/\//u);
    });

    it("avoids duplicate character names across agents when possible", () => {
      // Add 6 agents - each should ideally get a unique character
      for (let i = 0; i < 6; i++) {
        agentStore.addAgent(conversationId, createMockAgent({ toolUseId: `tool${i.toString()}` }));
      }

      const agents = get(getAgentsForConversation(conversationId));
      const names = agents.map((a) => a.characterName);
      const uniqueNames = new Set(names);
      expect(uniqueNames.size).toBe(6);
    });

    it("adds multiple agents to the same conversation", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1" });
      const agent2 = createMockAgent({ toolUseId: "tool2" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(conversationId, agent2);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents).toHaveLength(2);
      expect(agents[0]).toMatchObject(agent1);
      expect(agents[1]).toMatchObject(agent2);
    });

    it("keeps agents in different conversations separate", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1" });
      const agent2 = createMockAgent({ toolUseId: "tool2" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(otherConversationId, agent2);

      const agents1 = get(getAgentsForConversation(conversationId));
      const agents2 = get(getAgentsForConversation(otherConversationId));

      expect(agents1).toHaveLength(1);
      expect(agents2).toHaveLength(1);
      expect(agents1[0]).toMatchObject(agent1);
      expect(agents2[0]).toMatchObject(agent2);
    });
  });

  describe("updateAgentId", () => {
    it("updates the agent_id for a specific agent", () => {
      const agent = createMockAgent({ agentId: undefined });
      agentStore.addAgent(conversationId, agent);

      agentStore.updateAgentId(conversationId, agent.toolUseId, "agent-abc123");

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].agentId).toBe("agent-abc123");
    });

    it("does nothing if conversation doesn't exist", () => {
      agentStore.updateAgentId("nonexistent", "tool1", "agent1");
      // Should not throw
      expect(true).toBe(true);
    });

    it("does nothing if tool_use_id doesn't exist", () => {
      const agent = createMockAgent();
      agentStore.addAgent(conversationId, agent);

      agentStore.updateAgentId(conversationId, "nonexistent-tool", "agent1");

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].agentId).toBeUndefined();
    });

    it("updates agentType when provided alongside agentId", () => {
      const agent = createMockAgent({ agentId: undefined });
      agentStore.addAgent(conversationId, agent);

      agentStore.updateAgentId(conversationId, agent.toolUseId, "agent-abc123", "general-purpose");

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].agentId).toBe("agent-abc123");
      expect(agents[0].agentType).toBe("general-purpose");
    });

    it("does not set agentType when not provided", () => {
      const agent = createMockAgent({ agentId: undefined });
      agentStore.addAgent(conversationId, agent);

      agentStore.updateAgentId(conversationId, agent.toolUseId, "agent-abc123");

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].agentId).toBe("agent-abc123");
      expect(agents[0].agentType).toBeUndefined();
    });
  });

  describe("endAgent", () => {
    it("marks an agent as completed", () => {
      const agent = createMockAgent({ status: "running" });
      agentStore.addAgent(conversationId, agent);

      const endTime = Date.now();
      agentStore.endAgent(conversationId, agent.toolUseId, endTime, false);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].status).toBe("completed");
      expect(agents[0].endedAt).toBe(endTime);
      expect(agents[0].durationMs).toBeGreaterThanOrEqual(0); // Duration can be 0 if timestamps are the same
    });

    it("marks an agent as errored", () => {
      const agent = createMockAgent({ status: "running" });
      agentStore.addAgent(conversationId, agent);

      const endTime = Date.now();
      agentStore.endAgent(conversationId, agent.toolUseId, endTime, true);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].status).toBe("errored");
      expect(agents[0].endedAt).toBe(endTime);
    });

    it("calculates duration correctly", () => {
      const startTime = Date.now() - 5000; // 5 seconds ago
      const agent = createMockAgent({ startedAt: startTime, status: "running" });
      agentStore.addAgent(conversationId, agent);

      const endTime = Date.now();
      agentStore.endAgent(conversationId, agent.toolUseId, endTime, false);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].durationMs).toBeGreaterThanOrEqual(5000);
      expect(agents[0].durationMs).toBeLessThanOrEqual(6000); // Allow some buffer
    });

    it("does nothing if conversation doesn't exist", () => {
      agentStore.endAgent("nonexistent", "tool1", Date.now(), false);
      // Should not throw
      expect(true).toBe(true);
    });

    it("does nothing if agent doesn't exist", () => {
      const agent = createMockAgent();
      agentStore.addAgent(conversationId, agent);

      agentStore.endAgent(conversationId, "nonexistent-tool", Date.now(), false);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].status).toBe("running"); // Status unchanged
    });

    it("stores lastAssistantMessage when provided", () => {
      const agent = createMockAgent({ status: "running" });
      agentStore.addAgent(conversationId, agent);

      agentStore.endAgent(
        conversationId,
        agent.toolUseId,
        Date.now(),
        false,
        "Task completed successfully."
      );

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].lastAssistantMessage).toBe("Task completed successfully.");
    });

    it("leaves lastAssistantMessage undefined when not provided", () => {
      const agent = createMockAgent({ status: "running" });
      agentStore.addAgent(conversationId, agent);

      agentStore.endAgent(conversationId, agent.toolUseId, Date.now(), false);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].lastAssistantMessage).toBeUndefined();
    });
  });

  describe("markAllErrored", () => {
    it("marks all running agents as errored", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1", status: "running" });
      const agent2 = createMockAgent({ toolUseId: "tool2", status: "running" });
      const agent3 = createMockAgent({ toolUseId: "tool3", status: "completed" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(conversationId, agent2);
      agentStore.addAgent(conversationId, agent3);

      agentStore.markAllErrored(conversationId);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].status).toBe("errored");
      expect(agents[0].endedAt).toBeGreaterThan(0);
      expect(agents[1].status).toBe("errored");
      expect(agents[1].endedAt).toBeGreaterThan(0);
      expect(agents[2].status).toBe("completed"); // Already completed, unchanged
    });

    it("does nothing if conversation doesn't exist", () => {
      agentStore.markAllErrored("nonexistent");
      // Should not throw
      expect(true).toBe(true);
    });

    it("does nothing if conversation has no running agents", () => {
      const agent = createMockAgent({ status: "completed" });
      agentStore.addAgent(conversationId, agent);

      agentStore.markAllErrored(conversationId);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents[0].status).toBe("completed"); // Unchanged
    });
  });

  describe("clearCompleted", () => {
    it("removes completed and errored agents", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1", status: "running" });
      const agent2 = createMockAgent({ toolUseId: "tool2", status: "completed" });
      const agent3 = createMockAgent({ toolUseId: "tool3", status: "errored" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(conversationId, agent2);
      agentStore.addAgent(conversationId, agent3);

      agentStore.clearCompleted(conversationId);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents).toHaveLength(1);
      expect(agents[0].toolUseId).toBe("tool1"); // Only running agent remains
    });

    it("does nothing if conversation doesn't exist", () => {
      agentStore.clearCompleted("nonexistent");
      // Should not throw
      expect(true).toBe(true);
    });

    it("clears all agents if all are completed", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1", status: "completed" });
      const agent2 = createMockAgent({ toolUseId: "tool2", status: "errored" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(conversationId, agent2);

      agentStore.clearCompleted(conversationId);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents).toHaveLength(0);
    });
  });

  describe("clearConversation", () => {
    it("removes all agents from a conversation", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1" });
      const agent2 = createMockAgent({ toolUseId: "tool2" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(conversationId, agent2);

      agentStore.clearConversation(conversationId);

      const agents = get(getAgentsForConversation(conversationId));
      expect(agents).toHaveLength(0);
    });

    it("only removes agents from the specified conversation", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1" });
      const agent2 = createMockAgent({ toolUseId: "tool2" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(otherConversationId, agent2);

      agentStore.clearConversation(conversationId);

      const agents1 = get(getAgentsForConversation(conversationId));
      const agents2 = get(getAgentsForConversation(otherConversationId));

      expect(agents1).toHaveLength(0);
      expect(agents2).toHaveLength(1);
      expect(agents2[0]).toMatchObject(agent2);
    });

    it("does nothing if conversation doesn't exist", () => {
      agentStore.clearConversation("nonexistent");
      // Should not throw
      expect(true).toBe(true);
    });
  });

  describe("runningAgentCount", () => {
    it("counts running agents across all conversations", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1", status: "running" });
      const agent2 = createMockAgent({ toolUseId: "tool2", status: "running" });
      const agent3 = createMockAgent({ toolUseId: "tool3", status: "completed" });
      const agent4 = createMockAgent({ toolUseId: "tool4", status: "running" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(conversationId, agent2);
      agentStore.addAgent(conversationId, agent3);
      agentStore.addAgent(otherConversationId, agent4);

      const count = get(runningAgentCount);
      expect(count).toBe(3); // 2 from first conversation + 1 from second
    });

    it("returns 0 when no agents are running", () => {
      const agent1 = createMockAgent({ status: "completed" });
      const agent2 = createMockAgent({ status: "errored" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(otherConversationId, agent2);

      const count = get(runningAgentCount);
      expect(count).toBe(0);
    });

    it("updates when agents complete", () => {
      const agent = createMockAgent({ status: "running" });
      agentStore.addAgent(conversationId, agent);

      let count = get(runningAgentCount);
      expect(count).toBe(1);

      agentStore.endAgent(conversationId, agent.toolUseId, Date.now(), false);

      count = get(runningAgentCount);
      expect(count).toBe(0);
    });

    it("updates when conversation is cleared", () => {
      const agent1 = createMockAgent({ toolUseId: "tool1", status: "running" });
      const agent2 = createMockAgent({ toolUseId: "tool2", status: "running" });

      agentStore.addAgent(conversationId, agent1);
      agentStore.addAgent(conversationId, agent2);

      let count = get(runningAgentCount);
      expect(count).toBe(2);

      agentStore.clearConversation(conversationId);

      count = get(runningAgentCount);
      expect(count).toBe(0);
    });
  });
});