Homelab MCP Server

import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { z } from "zod";
import { ServiceContainer } from "../services/container.js";
import { type ToolDefinition, ToolRegistry } from "./registry.js";

describe("ToolRegistry", () => {
  let container: ServiceContainer;
  let registry: ToolRegistry;

  beforeEach(() => {
    container = new ServiceContainer();
    registry = new ToolRegistry(container);
  });

  const mockTool: ToolDefinition = {
    name: "test-tool",
    title: "Test Tool",
    description: "A test tool",
    inputSchema: z.object({ input: z.string() }),
    handler: async () => "test output",
  };

  describe("register", () => {
    it("should register a tool", () => {
      registry.register(mockTool);
      expect(registry.has("test-tool")).toBe(true);
    });

    it("should throw on duplicate registration", () => {
      registry.register(mockTool);
      expect(() => registry.register(mockTool)).toThrow('Tool "test-tool" is already registered');
    });

    it("should register multiple tools", () => {
      const tool2 = { ...mockTool, name: "tool2" };
      registry.register(mockTool);
      registry.register(tool2);

      expect(registry.has("test-tool")).toBe(true);
      expect(registry.has("tool2")).toBe(true);
    });
  });

  describe("has", () => {
    it("should return true for registered tool", () => {
      registry.register(mockTool);
      expect(registry.has("test-tool")).toBe(true);
    });

    it("should return false for unregistered tool", () => {
      expect(registry.has("nonexistent")).toBe(false);
    });
  });

  describe("get", () => {
    it("should return registered tool definition", () => {
      registry.register(mockTool);
      const tool = registry.get("test-tool");
      expect(tool).toBe(mockTool);
    });

    it("should return undefined for unregistered tool", () => {
      const tool = registry.get("nonexistent");
      expect(tool).toBeUndefined();
    });
  });

  describe("getToolNames", () => {
    it("should return empty array when no tools registered", () => {
      expect(registry.getToolNames()).toEqual([]);
    });

    it("should return all registered tool names", () => {
      const tool2 = { ...mockTool, name: "tool2" };
      registry.register(mockTool);
      registry.register(tool2);

      const names = registry.getToolNames();
      expect(names).toContain("test-tool");
      expect(names).toContain("tool2");
      expect(names).toHaveLength(2);
    });
  });

  describe("unregister", () => {
    it("should remove registered tool", () => {
      registry.register(mockTool);
      expect(registry.has("test-tool")).toBe(true);

      const result = registry.unregister("test-tool");
      expect(result).toBe(true);
      expect(registry.has("test-tool")).toBe(false);
    });

    it("should return false when removing non-existent tool", () => {
      const result = registry.unregister("nonexistent");
      expect(result).toBe(false);
    });
  });

  describe("clear", () => {
    it("should remove all tools", () => {
      const tool2 = { ...mockTool, name: "tool2" };
      registry.register(mockTool);
      registry.register(tool2);

      registry.clear();
      expect(registry.getToolNames()).toEqual([]);
      expect(registry.has("test-tool")).toBe(false);
      expect(registry.has("tool2")).toBe(false);
    });
  });

  describe("registerAll", () => {
    it("should register all tools with MCP server", async () => {
      const mockServer = {
        registerTool: vi.fn(),
      } as unknown as McpServer;

      const tool1: ToolDefinition = {
        name: "tool1",
        title: "Tool 1",
        inputSchema: z.object({ a: z.string() }),
        handler: async () => "output1",
      };

      const tool2: ToolDefinition = {
        name: "tool2",
        title: "Tool 2",
        description: "Custom description",
        inputSchema: z.object({ b: z.number() }),
        annotations: {
          readOnlyHint: true,
          destructiveHint: false,
          idempotentHint: true,
          openWorldHint: false,
        },
        handler: async () => "output2",
      };

      registry.register(tool1);
      registry.register(tool2);
      registry.registerAll(mockServer);

      expect(mockServer.registerTool).toHaveBeenCalledTimes(2);

      // Verify first tool registration
      const calls = vi.mocked(mockServer.registerTool).mock.calls;
      const [name1, config1] = calls[0];
      expect(name1).toBe("tool1");
      expect(config1.title).toBe("Tool 1");
      expect(config1.inputSchema).toBe(tool1.inputSchema);

      // Verify second tool registration with custom annotations
      const [name2, config2] = calls[1];
      expect(name2).toBe("tool2");
      expect(config2.title).toBe("Tool 2");
      expect(config2.description).toBe("Custom description");
      expect(config2.annotations).toEqual({
        readOnlyHint: true,
        destructiveHint: false,
        idempotentHint: true,
        openWorldHint: false,
      });
    });

    it("should prefer mcpInputSchema over inputSchema when provided", () => {
      const mockServer = {
        registerTool: vi.fn(),
      } as unknown as McpServer;

      const strictSchema = z.object({ id: z.string().uuid() });
      const advertisedSchema = z.object({ id: z.string() }).passthrough();

      const tool: ToolDefinition = {
        name: "schema-tool",
        title: "Schema Tool",
        inputSchema: strictSchema,
        mcpInputSchema: advertisedSchema,
        handler: async () => "ok",
      };

      registry.register(tool);
      registry.registerAll(mockServer);

      const config = vi.mocked(mockServer.registerTool).mock.calls[0][1];
      expect(config.inputSchema).toBe(advertisedSchema);
    });

    it("should wrap handler with error logging", async () => {
      const mockServer = {
        registerTool: vi.fn(),
      } as unknown as McpServer;

      const handlerSpy = vi.fn().mockResolvedValue("success");
      const tool: ToolDefinition = {
        name: "test",
        title: "Test",
        inputSchema: z.object({}),
        handler: handlerSpy,
      };

      registry.register(tool);
      registry.registerAll(mockServer);

      // Extract registered handler from mock
      const registeredHandler = vi.mocked(mockServer.registerTool).mock.calls[0][2];

      // Call the wrapped handler
      const result = await registeredHandler({ input: "test" });

      expect(handlerSpy).toHaveBeenCalledWith({ input: "test" }, container);
      expect(result).toEqual({ content: [{ type: "text", text: "success" }] });
    });

    it("should log errors and re-throw", async () => {
      const mockServer = {
        registerTool: vi.fn(),
      } as unknown as McpServer;

      const error = new Error("Handler failed");
      const handlerSpy = vi.fn().mockRejectedValue(error);
      const tool: ToolDefinition = {
        name: "test",
        title: "Test Tool",
        inputSchema: z.object({}),
        handler: handlerSpy,
      };

      registry.register(tool);
      registry.registerAll(mockServer);

      // Extract registered handler from mock
      const registeredHandler = vi.mocked(mockServer.registerTool).mock.calls[0][2];

      // Expect wrapped handler to throw
      await expect(registeredHandler({ input: "test" })).rejects.toThrow("Handler failed");

      expect(handlerSpy).toHaveBeenCalled();
    });
  });

  describe("tool definitions", () => {
    it("should support minimal tool definition", () => {
      const minimalTool: ToolDefinition = {
        name: "minimal",
        title: "Minimal Tool",
        inputSchema: z.object({}),
        handler: async () => "output",
      };

      expect(() => registry.register(minimalTool)).not.toThrow();
      expect(registry.has("minimal")).toBe(true);
    });

    it("should support full tool definition with all fields", () => {
      const fullTool: ToolDefinition = {
        name: "full",
        title: "Full Tool",
        description: "A fully configured tool",
        inputSchema: z.object({ param: z.string() }),
        annotations: {
          readOnlyHint: true,
          destructiveHint: false,
          idempotentHint: true,
          openWorldHint: false,
        },
        handler: async (params) => `Received: ${JSON.stringify(params)}`,
      };

      expect(() => registry.register(fullTool)).not.toThrow();
      expect(registry.has("full")).toBe(true);

      const retrieved = registry.get("full");
      expect(retrieved?.description).toBe("A fully configured tool");
      expect(retrieved?.annotations?.readOnlyHint).toBe(true);
    });
  });
});