Kilo-Org_kilocode/packages/opencode/test/tool/registry.test.ts

import { afterEach, describe, expect, test } from "bun:test"
import path from "path"
import fs from "fs/promises"
import { Effect, Layer } from "effect"
import { Instance } from "../../src/project/instance"
import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
import { ToolRegistry } from "@/tool/registry"
import { provideTmpdirInstance, tmpdir } from "../fixture/fixture" // kilocode_change
import { testEffect } from "../lib/effect"

const node = CrossSpawnSpawner.defaultLayer

const it = testEffect(Layer.mergeAll(ToolRegistry.defaultLayer, node))

afterEach(async () => {
  await Instance.disposeAll()
})

describe("tool.registry", () => {
  // kilocode_change start - plan_exit is always registered
  it.live("plan_exit is always registered regardless of client", () =>
    Effect.gen(function* () {
      const original = process.env["KILO_CLIENT"]
      try {
        for (const client of ["cli", "vscode", "desktop", "app"]) {
          process.env["KILO_CLIENT"] = client
          yield* provideTmpdirInstance(
            () =>
              Effect.gen(function* () {
                const registry = yield* ToolRegistry.Service
                const ids = yield* registry.ids()
                expect(ids).toContain("plan_exit")
              }),
            { git: true },
          )
        }
      } finally {
        if (original === undefined) delete process.env["KILO_CLIENT"]
        else process.env["KILO_CLIENT"] = original
      }
    }),
  )
  // kilocode_change end

  // kilocode_change start
  test("suggest is registered for cli and vscode only", async () => {
    const original = process.env["KILO_CLIENT"]
    const originalQuestion = process.env["KILO_ENABLE_QUESTION_TOOL"]
    const originalConfig = process.env["KILO_CONFIG_DIR"]
    try {
      for (const client of ["cli", "vscode", "desktop", "app"]) {
        process.env["KILO_CLIENT"] = client
        process.env["KILO_ENABLE_QUESTION_TOOL"] = client === "vscode" ? "true" : "false"
        await using tmp = await tmpdir({ git: true })
        process.env["KILO_CONFIG_DIR"] = tmp.path
        await Instance.provide({
          directory: tmp.path,
          fn: async () => {
            const ids = await ToolRegistry.ids()
            if (client === "cli" || client === "vscode") expect(ids).toContain("suggest")
            else expect(ids).not.toContain("suggest")
          },
        })
      }
    } finally {
      if (original === undefined) delete process.env["KILO_CLIENT"]
      else process.env["KILO_CLIENT"] = original
      if (originalQuestion === undefined) delete process.env["KILO_ENABLE_QUESTION_TOOL"]
      else process.env["KILO_ENABLE_QUESTION_TOOL"] = originalQuestion
      if (originalConfig === undefined) delete process.env["KILO_CONFIG_DIR"]
      else process.env["KILO_CONFIG_DIR"] = originalConfig
    }
  })
  // kilocode_change end

  it.live("loads tools from .opencode/tool (singular)", () =>
    provideTmpdirInstance((dir) =>
      Effect.gen(function* () {
        const opencode = path.join(dir, ".opencode")
        const tool = path.join(opencode, "tool")
        yield* Effect.promise(() => fs.mkdir(tool, { recursive: true }))
        yield* Effect.promise(() =>
          Bun.write(
            path.join(tool, "hello.ts"),
            [
              "export default {",
              "  description: 'hello tool',",
              "  args: {},",
              "  execute: async () => {",
              "    return 'hello world'",
              "  },",
              "}",
              "",
            ].join("\n"),
          ),
        )
        const registry = yield* ToolRegistry.Service
        const ids = yield* registry.ids()
        expect(ids).toContain("hello")
      }),
    ),
  )

  it.live("loads tools from .opencode/tools (plural)", () =>
    provideTmpdirInstance((dir) =>
      Effect.gen(function* () {
        const opencode = path.join(dir, ".opencode")
        const tools = path.join(opencode, "tools")
        yield* Effect.promise(() => fs.mkdir(tools, { recursive: true }))
        yield* Effect.promise(() =>
          Bun.write(
            path.join(tools, "hello.ts"),
            [
              "export default {",
              "  description: 'hello tool',",
              "  args: {},",
              "  execute: async () => {",
              "    return 'hello world'",
              "  },",
              "}",
              "",
            ].join("\n"),
          ),
        )
        const registry = yield* ToolRegistry.Service
        const ids = yield* registry.ids()
        expect(ids).toContain("hello")
      }),
    ),
  )

  it.live("loads tools with external dependencies without crashing", () =>
    provideTmpdirInstance((dir) =>
      Effect.gen(function* () {
        const opencode = path.join(dir, ".opencode")
        const tools = path.join(opencode, "tools")
        yield* Effect.promise(() => fs.mkdir(tools, { recursive: true }))
        yield* Effect.promise(() =>
          Bun.write(
            path.join(opencode, "package.json"),
            JSON.stringify({
              name: "custom-tools",
              dependencies: {
                "@kilocode/plugin": "^0.0.0",
                cowsay: "^1.6.0",
              },
            }),
          ),
        )
        yield* Effect.promise(() =>
          Bun.write(
            path.join(opencode, "package-lock.json"),
            JSON.stringify({
              name: "custom-tools",
              lockfileVersion: 3,
              packages: {
                "": {
                  dependencies: {
                    "@kilocode/plugin": "^0.0.0",
                    cowsay: "^1.6.0",
                  },
                },
              },
            }),
          ),
        )

        const cowsay = path.join(opencode, "node_modules", "cowsay")
        yield* Effect.promise(() => fs.mkdir(cowsay, { recursive: true }))
        yield* Effect.promise(() =>
          Bun.write(
            path.join(cowsay, "package.json"),
            JSON.stringify({
              name: "cowsay",
              type: "module",
              exports: "./index.js",
            }),
          ),
        )
        yield* Effect.promise(() =>
          Bun.write(
            path.join(cowsay, "index.js"),
            ["export function say({ text }) {", "  return `moo ${text}`", "}", ""].join("\n"),
          ),
        )
        yield* Effect.promise(() =>
          Bun.write(
            path.join(tools, "cowsay.ts"),
            [
              "import { say } from 'cowsay'",
              "export default {",
              "  description: 'tool that imports cowsay at top level',",
              "  args: { text: { type: 'string' } },",
              "  execute: async ({ text }: { text: string }) => {",
              "    return say({ text })",
              "  },",
              "}",
              "",
            ].join("\n"),
          ),
        )
        const registry = yield* ToolRegistry.Service
        const ids = yield* registry.ids()
        expect(ids).toContain("cowsay")
      }),
    ),
  )
})