branch:
extension-manager.test.ts
15441 bytesRaw
/**
 * Tests for the Extension system: ExtensionManager, HostBridgeLoopback, and
 * the extension Worker contract.
 *
 * Uses vitest-pool-workers with a real WorkerLoader binding.
 */
import { describe, it, expect, beforeEach } from "vitest";
import { env } from "cloudflare:workers";
import { ExtensionManager, sanitizeName } from "../extensions/manager";
import type { ExtensionManifest } from "../extensions/types";

// Simple extension source: a greet tool that returns a greeting
const GREET_EXTENSION_SOURCE = `{
  greet: {
    description: "Greet someone by name",
    parameters: { name: { type: "string", description: "Name to greet" } },
    required: ["name"],
    execute: async (args) => "Hello, " + args.name + "!"
  }
}`;

// Multi-tool extension
const MULTI_TOOL_SOURCE = `{
  add: {
    description: "Add two numbers",
    parameters: {
      a: { type: "number", description: "First number" },
      b: { type: "number", description: "Second number" }
    },
    required: ["a", "b"],
    execute: async (args) => args.a + args.b
  },
  multiply: {
    description: "Multiply two numbers",
    parameters: {
      a: { type: "number", description: "First number" },
      b: { type: "number", description: "Second number" }
    },
    required: ["a", "b"],
    execute: async (args) => args.a * args.b
  }
}`;

// Extension that uses the host bridge for workspace access
const _WORKSPACE_EXTENSION_SOURCE = `{
  readFile: {
    description: "Read a file via host bridge",
    parameters: { path: { type: "string" } },
    required: ["path"],
    execute: async (args, host) => {
      const content = await host.readFile(args.path);
      return content;
    }
  },
  writeFile: {
    description: "Write a file via host bridge",
    parameters: {
      path: { type: "string" },
      content: { type: "string" }
    },
    required: ["path", "content"],
    execute: async (args, host) => {
      await host.writeFile(args.path, args.content);
      return "written";
    }
  }
}`;

// Extension that throws errors
const ERROR_EXTENSION_SOURCE = `{
  fail: {
    description: "Always fails",
    parameters: {},
    execute: async () => { throw new Error("intentional failure"); }
  }
}`;

function makeManifest(
  overrides?: Partial<ExtensionManifest>
): ExtensionManifest {
  return {
    name: "test-ext",
    version: "1.0.0",
    description: "Test extension",
    ...overrides
  };
}

describe("ExtensionManager", () => {
  let manager: ExtensionManager;

  beforeEach(() => {
    manager = new ExtensionManager({ loader: env.LOADER });
  });

  describe("load and discover", () => {
    it("should load an extension and discover its tools", async () => {
      const info = await manager.load(
        makeManifest({ name: "greeter" }),
        GREET_EXTENSION_SOURCE
      );

      expect(info.name).toBe("greeter");
      expect(info.version).toBe("1.0.0");
      expect(info.tools).toEqual(["greeter_greet"]);
    });

    it("should load a multi-tool extension", async () => {
      const info = await manager.load(
        makeManifest({ name: "math" }),
        MULTI_TOOL_SOURCE
      );

      expect(info.tools).toContain("math_add");
      expect(info.tools).toContain("math_multiply");
      expect(info.tools).toHaveLength(2);
    });

    it("should reject duplicate extension names", async () => {
      await manager.load(makeManifest({ name: "dup" }), GREET_EXTENSION_SOURCE);

      await expect(
        manager.load(makeManifest({ name: "dup" }), GREET_EXTENSION_SOURCE)
      ).rejects.toThrow("already loaded");
    });
  });

  describe("unload", () => {
    it("should unload an extension", async () => {
      await manager.load(
        makeManifest({ name: "temp" }),
        GREET_EXTENSION_SOURCE
      );
      expect(manager.list()).toHaveLength(1);

      const removed = await manager.unload("temp");
      expect(removed).toBe(true);
      expect(manager.list()).toHaveLength(0);
    });

    it("should return false for unknown extension", async () => {
      expect(await manager.unload("nonexistent")).toBe(false);
    });
  });

  describe("list", () => {
    it("should list all loaded extensions", async () => {
      await manager.load(
        makeManifest({ name: "ext-a", version: "1.0.0" }),
        GREET_EXTENSION_SOURCE
      );
      await manager.load(
        makeManifest({ name: "ext-b", version: "2.0.0" }),
        MULTI_TOOL_SOURCE
      );

      const list = manager.list();
      expect(list).toHaveLength(2);

      const names = list.map((e) => e.name);
      expect(names).toContain("ext-a");
      expect(names).toContain("ext-b");
    });

    it("should return empty array when no extensions loaded", () => {
      expect(manager.list()).toEqual([]);
    });
  });

  describe("getTools", () => {
    it("should return AI SDK tools from loaded extensions", async () => {
      await manager.load(
        makeManifest({ name: "greeter" }),
        GREET_EXTENSION_SOURCE
      );

      const tools = manager.getTools();
      expect(tools).toHaveProperty("greeter_greet");
    });

    it("should prefix tool names with extension name", async () => {
      await manager.load(makeManifest({ name: "math" }), MULTI_TOOL_SOURCE);

      const tools = manager.getTools();
      expect(Object.keys(tools)).toContain("math_add");
      expect(Object.keys(tools)).toContain("math_multiply");
    });

    it("should merge tools from multiple extensions", async () => {
      await manager.load(
        makeManifest({ name: "greeter" }),
        GREET_EXTENSION_SOURCE
      );
      await manager.load(makeManifest({ name: "math" }), MULTI_TOOL_SOURCE);

      const tools = manager.getTools();
      expect(Object.keys(tools)).toHaveLength(3);
      expect(tools).toHaveProperty("greeter_greet");
      expect(tools).toHaveProperty("math_add");
      expect(tools).toHaveProperty("math_multiply");
    });

    it("should return empty object when no extensions loaded", () => {
      expect(manager.getTools()).toEqual({});
    });
  });

  describe("tool execution", () => {
    it("should execute a simple tool", async () => {
      await manager.load(
        makeManifest({ name: "greeter" }),
        GREET_EXTENSION_SOURCE
      );

      const tools = manager.getTools();
      const greet = tools.greeter_greet;
      const result = await greet.execute!(
        { name: "World" },
        { toolCallId: "tc1", messages: [], abortSignal: undefined as never }
      );

      expect(result).toBe("Hello, World!");
    });

    it("should execute math tools correctly", async () => {
      await manager.load(makeManifest({ name: "math" }), MULTI_TOOL_SOURCE);

      const tools = manager.getTools();

      const sum = await tools.math_add.execute!(
        { a: 3, b: 4 },
        { toolCallId: "tc1", messages: [], abortSignal: undefined as never }
      );
      expect(sum).toBe(7);

      const product = await tools.math_multiply.execute!(
        { a: 5, b: 6 },
        { toolCallId: "tc2", messages: [], abortSignal: undefined as never }
      );
      expect(product).toBe(30);
    });

    it("should propagate tool execution errors", async () => {
      await manager.load(makeManifest({ name: "bad" }), ERROR_EXTENSION_SOURCE);

      const tools = manager.getTools();

      await expect(
        tools.bad_fail.execute!(
          {},
          { toolCallId: "tc1", messages: [], abortSignal: undefined as never }
        )
      ).rejects.toThrow("intentional failure");
    });
  });

  describe("network isolation", () => {
    it("should block network by default (no network permission)", async () => {
      const source = `{
        fetchUrl: {
          description: "Try to fetch a URL",
          parameters: { url: { type: "string" } },
          required: ["url"],
          execute: async (args) => {
            const res = await fetch(args.url);
            return res.status;
          }
        }
      }`;

      await manager.load(
        makeManifest({ name: "fetcher", permissions: {} }),
        source
      );

      const tools = manager.getTools();

      await expect(
        tools.fetcher_fetchUrl.execute!(
          { url: "https://example.com" },
          { toolCallId: "tc1", messages: [], abortSignal: undefined as never }
        )
      ).rejects.toThrow();
    });
  });

  describe("lifecycle after unload", () => {
    it("should not include tools from unloaded extensions", async () => {
      await manager.load(
        makeManifest({ name: "temp" }),
        GREET_EXTENSION_SOURCE
      );
      expect(Object.keys(manager.getTools())).toHaveLength(1);

      await manager.unload("temp");
      expect(Object.keys(manager.getTools())).toHaveLength(0);
    });

    it("should throw when executing a tool from an unloaded extension", async () => {
      await manager.load(
        makeManifest({ name: "temp" }),
        GREET_EXTENSION_SOURCE
      );

      // Capture tools while extension is still loaded
      const tools = manager.getTools();
      expect(tools).toHaveProperty("temp_greet");

      // Unload the extension
      await manager.unload("temp");

      // Executing the captured tool should throw
      await expect(
        tools.temp_greet.execute!(
          { name: "world" },
          { toolCallId: "tc1", messages: [], abortSignal: undefined as never }
        )
      ).rejects.toThrow(/has been unloaded/);
    });

    it("should allow reloading after unload", async () => {
      await manager.load(
        makeManifest({ name: "temp" }),
        GREET_EXTENSION_SOURCE
      );
      await manager.unload("temp");

      const info = await manager.load(
        makeManifest({ name: "temp" }),
        GREET_EXTENSION_SOURCE
      );
      expect(info.tools).toEqual(["temp_greet"]);
    });
  });

  describe("persistence and restore", () => {
    it("should restore extensions from storage after re-creation", async () => {
      // Simulate DO storage with a simple Map-backed mock
      const store = new Map<string, unknown>();
      const mockStorage = {
        get: async (key: string) => store.get(key),
        put: async (key: string, value: unknown) => {
          store.set(key, value);
        },
        delete: async (key: string) => store.delete(key),
        list: async (opts?: { prefix?: string }) => {
          const result = new Map<string, unknown>();
          for (const [k, v] of store) {
            if (!opts?.prefix || k.startsWith(opts.prefix)) {
              result.set(k, v);
            }
          }
          return result;
        }
      } as unknown as DurableObjectStorage;

      // First manager — load an extension
      const manager1 = new ExtensionManager({
        loader: env.LOADER,
        storage: mockStorage
      });
      await manager1.load(makeManifest({ name: "math" }), MULTI_TOOL_SOURCE);
      expect(manager1.list()).toHaveLength(1);

      // Second manager — simulates DO wake from hibernation
      const manager2 = new ExtensionManager({
        loader: env.LOADER,
        storage: mockStorage
      });
      expect(manager2.list()).toHaveLength(0); // empty before restore

      await manager2.restore();
      expect(manager2.list()).toHaveLength(1);
      expect(manager2.list()[0].name).toBe("math");

      // Tools should work after restore
      const tools = manager2.getTools();
      expect(tools).toHaveProperty("math_add");
      const sum = await tools.math_add.execute!(
        { a: 10, b: 20 },
        { toolCallId: "tc1", messages: [], abortSignal: undefined as never }
      );
      expect(sum).toBe(30);
    });

    it("should remove from storage on unload", async () => {
      const store = new Map<string, unknown>();
      const mockStorage = {
        get: async (key: string) => store.get(key),
        put: async (key: string, value: unknown) => {
          store.set(key, value);
        },
        delete: async (key: string) => store.delete(key),
        list: async (opts?: { prefix?: string }) => {
          const result = new Map<string, unknown>();
          for (const [k, v] of store) {
            if (!opts?.prefix || k.startsWith(opts.prefix)) {
              result.set(k, v);
            }
          }
          return result;
        }
      } as unknown as DurableObjectStorage;

      const mgr = new ExtensionManager({
        loader: env.LOADER,
        storage: mockStorage
      });
      await mgr.load(makeManifest({ name: "temp" }), GREET_EXTENSION_SOURCE);
      expect(store.size).toBe(1);

      await mgr.unload("temp");
      expect(store.size).toBe(0);

      // New manager should have nothing to restore
      const mgr2 = new ExtensionManager({
        loader: env.LOADER,
        storage: mockStorage
      });
      await mgr2.restore();
      expect(mgr2.list()).toHaveLength(0);
    });

    it("restore should be idempotent", async () => {
      const store = new Map<string, unknown>();
      const mockStorage = {
        get: async (key: string) => store.get(key),
        put: async (key: string, value: unknown) => {
          store.set(key, value);
        },
        delete: async (key: string) => store.delete(key),
        list: async (opts?: { prefix?: string }) => {
          const result = new Map<string, unknown>();
          for (const [k, v] of store) {
            if (!opts?.prefix || k.startsWith(opts.prefix)) {
              result.set(k, v);
            }
          }
          return result;
        }
      } as unknown as DurableObjectStorage;

      const mgr = new ExtensionManager({
        loader: env.LOADER,
        storage: mockStorage
      });
      await mgr.load(makeManifest({ name: "greeter" }), GREET_EXTENSION_SOURCE);

      const mgr2 = new ExtensionManager({
        loader: env.LOADER,
        storage: mockStorage
      });
      await mgr2.restore();
      await mgr2.restore(); // second call should be a no-op
      expect(mgr2.list()).toHaveLength(1);
    });
  });

  describe("name sanitization", () => {
    it("should sanitize hyphens in extension names for tool names", async () => {
      const info = await manager.load(
        makeManifest({ name: "my-cool-ext" }),
        GREET_EXTENSION_SOURCE
      );

      // ExtensionInfo.tools should have sanitized prefix
      expect(info.tools).toEqual(["my_cool_ext_greet"]);

      // getTools() should use sanitized prefix
      const tools = manager.getTools();
      expect(tools).toHaveProperty("my_cool_ext_greet");
      expect(tools).not.toHaveProperty("my-cool-ext_greet");
    });

    it("sanitizeName should replace non-alphanumeric chars", () => {
      expect(sanitizeName("hello-world")).toBe("hello_world");
      expect(sanitizeName("foo.bar.baz")).toBe("foo_bar_baz");
      expect(sanitizeName("a--b")).toBe("a_b");
      expect(sanitizeName("simple")).toBe("simple");
      expect(sanitizeName("with spaces")).toBe("with_spaces");
    });

    it("sanitizeName should strip leading and trailing underscores", () => {
      expect(sanitizeName("-leading")).toBe("leading");
      expect(sanitizeName("trailing-")).toBe("trailing");
      expect(sanitizeName("--both--")).toBe("both");
      expect(sanitizeName("...dots...")).toBe("dots");
    });

    it("sanitizeName should throw for empty or whitespace-only names", () => {
      expect(() => sanitizeName("")).toThrow("must not be empty");
      expect(() => sanitizeName("   ")).toThrow("must not be empty");
    });
  });
});