hikari-desktop/src/lib/utils/stateMapper.test.ts

import { describe, it, expect } from "vitest";
import { mapMessageToState, extractTextFromMessage, extractToolInfo } from "./stateMapper";
import type { ClaudeStreamMessage } from "$lib/types/messages";

describe("stateMapper", () => {
  describe("mapMessageToState", () => {
    it("returns idle for system init message", () => {
      const message: ClaudeStreamMessage = {
        type: "system",
        subtype: "init",
        session_id: "test-session",
        cwd: "/home/test",
        tools: ["Read", "Write", "Edit"],
      };
      expect(mapMessageToState(message)).toBe("idle");
    });

    it("returns null for non-init system messages", () => {
      const message: ClaudeStreamMessage = {
        type: "system",
        subtype: "compact_boundary",
      };
      expect(mapMessageToState(message)).toBeNull();
    });

    it("returns searching for Read tool", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            {
              type: "tool_use",
              id: "tool-1",
              name: "Read",
              input: { file_path: "/test/file.txt" },
            },
          ],
          model: "claude-3",
          stop_reason: "tool_use",
        },
      };
      expect(mapMessageToState(message)).toBe("searching");
    });

    it("returns coding for Edit tool", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            {
              type: "tool_use",
              id: "tool-1",
              name: "Edit",
              input: { file_path: "/test/file.txt", old_string: "foo", new_string: "bar" },
            },
          ],
          model: "claude-3",
          stop_reason: "tool_use",
        },
      };
      expect(mapMessageToState(message)).toBe("coding");
    });

    it("returns mcp for mcp__ prefixed tools", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            {
              type: "tool_use",
              id: "tool-1",
              name: "mcp__github__list_repos",
              input: {},
            },
          ],
          model: "claude-3",
          stop_reason: "tool_use",
        },
      };
      expect(mapMessageToState(message)).toBe("mcp");
    });

    it("returns thinking for Task tool", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            {
              type: "tool_use",
              id: "tool-1",
              name: "Task",
              input: { prompt: "test task" },
            },
          ],
          model: "claude-3",
          stop_reason: "tool_use",
        },
      };
      expect(mapMessageToState(message)).toBe("thinking");
    });

    it("returns typing for text content", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [{ type: "text", text: "Hello, Naomi!" }],
          model: "claude-3",
          stop_reason: "end_turn",
        },
      };
      expect(mapMessageToState(message)).toBe("typing");
    });

    it("returns success for result success message", () => {
      const message: ClaudeStreamMessage = {
        type: "result",
        subtype: "success",
        result: "Task completed",
      };
      expect(mapMessageToState(message)).toBe("success");
    });

    it("returns error for result error message", () => {
      const message: ClaudeStreamMessage = {
        type: "result",
        subtype: "error_max_turns",
      };
      expect(mapMessageToState(message)).toBe("error");
    });

    it("returns null for user messages", () => {
      const message: ClaudeStreamMessage = {
        type: "user",
        message: { content: [{ type: "text", text: "Hello" }] },
      };
      expect(mapMessageToState(message)).toBeNull();
    });

    it("returns typing for unknown tool", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            {
              type: "tool_use",
              id: "tool-1",
              name: "SomeUnknownTool",
              input: {},
            },
          ],
          model: "claude-3",
          stop_reason: "tool_use",
        },
      };
      expect(mapMessageToState(message)).toBe("typing");
    });

    it("returns thinking for thinking content block", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [{ type: "thinking", thinking: "Analyzing the problem..." }],
          model: "claude-3",
          stop_reason: "end_turn",
        },
      };
      expect(mapMessageToState(message)).toBe("thinking");
    });

    it("returns null for assistant message with no recognizable content", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [],
          model: "claude-3",
          stop_reason: "end_turn",
        },
      };
      expect(mapMessageToState(message)).toBeNull();
    });

    it("returns thinking for thinking_delta stream event", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "content_block_delta",
          index: 0,
          delta: {
            type: "thinking_delta",
            thinking: "Thinking...",
          },
        },
      };
      expect(mapMessageToState(message)).toBe("thinking");
    });

    it("returns typing for text_delta stream event", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "content_block_delta",
          index: 0,
          delta: {
            type: "text_delta",
            text: "Hello",
          },
        },
      };
      expect(mapMessageToState(message)).toBe("typing");
    });

    it("returns thinking for thinking content_block_start", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "content_block_start",
          index: 0,
          content_block: {
            type: "thinking",
            thinking: "",
          },
        },
      };
      expect(mapMessageToState(message)).toBe("thinking");
    });

    it("returns typing for text content_block_start", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "content_block_start",
          index: 0,
          content_block: {
            type: "text",
            text: "",
          },
        },
      };
      expect(mapMessageToState(message)).toBe("typing");
    });

    it("returns correct state for tool_use content_block_start", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "content_block_start",
          index: 0,
          content_block: {
            type: "tool_use",
            id: "tool-1",
            name: "Read",
            input: {},
          },
        },
      };
      expect(mapMessageToState(message)).toBe("searching");
    });

    it("returns null for stream_event with unrecognized type", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "message_start",
        },
      };
      expect(mapMessageToState(message)).toBeNull();
    });

    it("returns null for result with unknown subtype", () => {
      const message = {
        type: "result",
        subtype: "unknown_type",
      } as unknown as ClaudeStreamMessage;
      expect(mapMessageToState(message)).toBeNull();
    });

    it("returns null for unknown message type", () => {
      const message = {
        type: "unknown_type",
      } as unknown as ClaudeStreamMessage;
      expect(mapMessageToState(message)).toBeNull();
    });
  });

  describe("extractTextFromMessage", () => {
    it("extracts text from assistant message", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            { type: "text", text: "Hello!" },
            { type: "text", text: "How are you?" },
          ],
          model: "claude-3",
          stop_reason: "end_turn",
        },
      };
      expect(extractTextFromMessage(message)).toBe("Hello!\nHow are you?");
    });

    it("returns null for assistant message without text", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            {
              type: "tool_use",
              id: "tool-1",
              name: "Read",
              input: { file_path: "/test/file.txt" },
            },
          ],
          model: "claude-3",
          stop_reason: "tool_use",
        },
      };
      expect(extractTextFromMessage(message)).toBeNull();
    });

    it("extracts text from stream_event delta", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "content_block_delta",
          index: 0,
          delta: { type: "text_delta", text: "streaming text" },
        },
      };
      expect(extractTextFromMessage(message)).toBe("streaming text");
    });

    it("extracts result from result message", () => {
      const message: ClaudeStreamMessage = {
        type: "result",
        subtype: "success",
        result: "Completed successfully",
      };
      expect(extractTextFromMessage(message)).toBe("Completed successfully");
    });

    it("returns null for result without result field", () => {
      const message: ClaudeStreamMessage = {
        type: "result",
        subtype: "success",
      };
      expect(extractTextFromMessage(message)).toBeNull();
    });

    it("returns null for stream_event without delta text", () => {
      const message: ClaudeStreamMessage = {
        type: "stream_event",
        event: {
          type: "content_block_start",
          index: 0,
          content_block: {
            type: "text",
            text: "",
          },
        },
      };
      expect(extractTextFromMessage(message)).toBeNull();
    });

    it("returns null for unknown message type", () => {
      const message = {
        type: "unknown",
      } as unknown as ClaudeStreamMessage;
      expect(extractTextFromMessage(message)).toBeNull();
    });
  });

  describe("extractToolInfo", () => {
    it("extracts tool info from assistant message", () => {
      const message: ClaudeStreamMessage = {
        type: "assistant",
        message: {
          content: [
            {
              type: "tool_use",
              id: "tool-1",
              name: "Read",
              input: { file_path: "/test/file.txt" },
            },
            {
              type: "tool_use",
              id: "tool-2",
              name: "Edit",
              input: { file_path: "/test/file.txt", old_string: "a", new_string: "b" },
            },
          ],
          model: "claude-3",
          stop_reason: "tool_use",
        },
      };
      const tools = extractToolInfo(message);
      expect(tools).toHaveLength(2);
      expect(tools[0]).toEqual({
        name: "Read",
        input: { file_path: "/test/file.txt" },
      });
      expect(tools[1]).toEqual({
        name: "Edit",
        input: { file_path: "/test/file.txt", old_string: "a", new_string: "b" },
      });
    });

    it("returns empty array for non-assistant messages", () => {
      const message: ClaudeStreamMessage = {
        type: "user",
        message: { content: [{ type: "text", text: "Hello" }] },
      };
      expect(extractToolInfo(message)).toEqual([]);
    });
  });
});