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";

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

  const createMockAgent = (overrides?: Partial<AgentInfo>): AgentInfo => ({
    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]).toEqual(agent);
    });

    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]).toEqual(agent1);
      expect(agents[1]).toEqual(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]).toEqual(agent1);
      expect(agents2[0]).toEqual(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();
    });
  });

  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
    });
  });

  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]).toEqual(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);
    });
  });
});