tf_code/packages/tfcode/test/agent/tf-agent.test.ts

import { afterAll, beforeAll, beforeEach, afterEach, test, expect, describe } from "bun:test"
import path from "path"
import fs from "fs"
import os from "os"
import { tmpdir } from "../fixture/fixture"
import { Instance } from "../../src/project/instance"
import { Agent } from "../../src/agent/agent"
import { LLM } from "../../src/session/llm"
import { Provider } from "../../src/provider/provider"
import { Global } from "../../src/global"
import { Filesystem } from "../../src/util/filesystem"
import { ModelsDev } from "../../src/provider/models"
import { ProviderID, ModelID } from "../../src/provider/schema"
import { SessionID, MessageID } from "../../src/session/schema"
import type { MessageV2 } from "../../src/session/message-v2"

const TF_TOOLS_PATH = path.join(os.homedir(), ".tfcode", "tools.json")

// Server for capturing LLM requests
const state = {
  server: null as ReturnType<typeof Bun.serve> | null,
  queue: [] as Array<{ path: string; response: Response; resolve: (value: any) => void }>,
}

function deferred<T>() {
  const result = {} as { promise: Promise<T>; resolve: (value: T) => void }
  result.promise = new Promise((resolve) => {
    result.resolve = resolve
  })
  return result
}

function waitRequest(pathname: string, response: Response) {
  const pending = deferred<{ url: URL; headers: Headers; body: Record<string, unknown> }>()
  state.queue.push({ path: pathname, response, resolve: pending.resolve })
  return pending.promise
}

function createChatStream(text: string) {
  const payload =
    [
      `data: ${JSON.stringify({
        id: "chatcmpl-1",
        object: "chat.completion.chunk",
        choices: [{ delta: { role: "assistant" } }],
      })}`,
      `data: ${JSON.stringify({
        id: "chatcmpl-1",
        object: "chat.completion.chunk",
        choices: [{ delta: { content: text } }],
      })}`,
      `data: ${JSON.stringify({
        id: "chatcmpl-1",
        object: "chat.completion.chunk",
        choices: [{ delta: {}, finish_reason: "stop" }],
      })}`,
      "data: [DONE]",
    ].join("\n\n") + "\n\n"
  return new ReadableStream({
    start(controller) {
      controller.enqueue(new TextEncoder().encode(payload))
      controller.close()
    },
  })
}

async function loadFixture(providerID: string, modelID: string) {
  const fixturePath = path.join(import.meta.dir, "../tool/fixtures/models-api.json")
  const data = await Filesystem.readJson<Record<string, ModelsDev.Provider>>(fixturePath)
  const provider = data[providerID]
  if (!provider) throw new Error(`Missing provider in fixture: ${providerID}`)
  const model = provider.models[modelID]
  if (!model) throw new Error(`Missing model in fixture: ${modelID}`)
  return { provider, model }
}

// Test the flow from tools.json -> Agent.Info -> highlighted instructions

describe("ToothFairyAI Agent Loading", () => {
  let originalDataPath: string
  let originalToolsContent: string | null = null

  beforeEach(async () => {
    originalDataPath = Global.Path.data
    const testDataDir = path.join(path.dirname(originalDataPath), "tf-agent-test-data")
    ;(Global.Path as { data: string }).data = testDataDir

    // Backup existing tools.json if it exists
    try {
      originalToolsContent = await Bun.file(TF_TOOLS_PATH).text()
    } catch {
      originalToolsContent = null
    }
    await fs.promises.mkdir(path.dirname(TF_TOOLS_PATH), { recursive: true })
  })

  afterEach(async () => {
    await Instance.disposeAll()
    ;(Global.Path as { data: string }).data = originalDataPath

    // Restore original tools.json
    if (originalToolsContent !== null) {
      await fs.promises.writeFile(TF_TOOLS_PATH, originalToolsContent)
    } else {
      try {
        await fs.promises.unlink(TF_TOOLS_PATH)
      } catch {}
    }
  })

  describe("loadTFCoderAgents", () => {
    test("parses tools.json with full agent data", async () => {
      const toolsData = {
        success: true,
        tools: [
          {
            id: "coder-agent-1",
            name: "Code Reviewer",
            description: "Reviews code for quality",
            tool_type: "coder_agent",
            request_type: null,
            url: null,
            auth_via: "tf_agent",
            interpolation_string: "You are a code reviewer. Review code thoroughly.",
            goals: "Identify bugs. Suggest improvements.",
            temperature: 0.3,
            max_tokens: 4096,
            llm_base_model: "claude-3-5-sonnet",
            llm_provider: "toothfairyai",
          },
          {
            id: "coder-agent-2",
            name: "Test Writer",
            description: "Writes tests",
            tool_type: "coder_agent",
            request_type: null,
            url: null,
            auth_via: "tf_agent",
            interpolation_string: "You are a test writer.",
            goals: "Write comprehensive tests.",
            temperature: 0.5,
            max_tokens: null,
            llm_base_model: "gpt-4",
            llm_provider: null,
          },
        ],
        by_type: { coder_agent: 2 },
      }

      const toolsPath = TF_TOOLS_PATH
      await fs.promises.writeFile(toolsPath, JSON.stringify(toolsData, null, 2))

      await using tmp = await tmpdir()

      await Instance.provide({
        directory: tmp.path,
        fn: async () => {
          const agents = await Agent.list()
          const codeReviewer = agents.find((a) => a.name === "Code Reviewer")
          const testWriter = agents.find((a) => a.name === "Test Writer")

          expect(codeReviewer).toBeDefined()
          expect(codeReviewer?.description).toBe("Reviews code for quality")
          expect(codeReviewer?.prompt).toBe("You are a code reviewer. Review code thoroughly.")
          expect(codeReviewer?.goals).toBe("Identify bugs. Suggest improvements.")
          expect(codeReviewer?.temperature).toBe(0.3)
          expect(codeReviewer?.native).toBe(false)
          expect(codeReviewer?.options?.tf_agent_id).toBe("coder-agent-1")
          expect(codeReviewer?.options?.tf_auth_via).toBe("tf_agent")
          expect(codeReviewer?.options?.tf_max_tokens).toBe(4096)
          expect(String(codeReviewer?.model?.providerID)).toBe("toothfairyai")
          expect(String(codeReviewer?.model?.modelID)).toBe("claude-3-5-sonnet")

          expect(testWriter).toBeDefined()
          expect(String(testWriter?.model?.providerID)).toBe("toothfairyai")
          expect(String(testWriter?.model?.modelID)).toBe("gpt-4")
        },
      })
    })

    test("maps tf provider to toothfairyai", async () => {
      const toolsData = {
        success: true,
        tools: [
          {
            id: "tf-provider-agent",
            name: "TF Provider Agent",
            description: "Test",
            tool_type: "coder_agent",
            interpolation_string: "Test",
            goals: "Test",
            temperature: 0.7,
            max_tokens: 2048,
            llm_base_model: "test-model",
            llm_provider: "tf",
          },
        ],
      }

      const toolsPath = TF_TOOLS_PATH
      await fs.promises.writeFile(toolsPath, JSON.stringify(toolsData))

      await using tmp = await tmpdir()

      await Instance.provide({
        directory: tmp.path,
        fn: async () => {
          const agent = await Agent.get("TF Provider Agent")
          expect(String(agent?.model?.providerID)).toBe("toothfairyai")
          expect(String(agent?.model?.modelID)).toBe("test-model")
        },
      })
    })

    test("does not map external providers", async () => {
      const toolsData = {
        success: true,
        tools: [
          {
            id: "external-agent",
            name: "External Agent",
            description: "Test",
            tool_type: "coder_agent",
            interpolation_string: "Test",
            goals: "Test",
            temperature: 0.7,
            max_tokens: 2048,
            llm_base_model: "claude-3-5-sonnet",
            llm_provider: "anthropic",
          },
        ],
      }

      const toolsPath = TF_TOOLS_PATH
      await fs.promises.writeFile(toolsPath, JSON.stringify(toolsData))

      await using tmp = await tmpdir()

      await Instance.provide({
        directory: tmp.path,
        fn: async () => {
          const agent = await Agent.get("External Agent")
          expect(agent?.model).toBeUndefined()
        },
      })
    })

    test("handles agent without interpolation_string or goals", async () => {
      const toolsData = {
        success: true,
        tools: [
          {
            id: "minimal-agent",
            name: "Minimal Agent",
            description: "Test",
            tool_type: "coder_agent",
            interpolation_string: null,
            goals: null,
            temperature: null,
            max_tokens: null,
            llm_base_model: null,
            llm_provider: null,
          },
        ],
      }

      const toolsPath = TF_TOOLS_PATH
      await fs.promises.writeFile(toolsPath, JSON.stringify(toolsData))

      await using tmp = await tmpdir()

      await Instance.provide({
        directory: tmp.path,
        fn: async () => {
          const agent = await Agent.get("Minimal Agent")
          expect(agent).toBeDefined()
          expect(agent?.prompt).toBeUndefined()
          expect(agent?.goals).toBeUndefined()
          expect(agent?.model).toBeUndefined()
        },
      })
    })
  })
})

// Separate describe block for LLM stream tests
describe("ToothFairyAI Agent Instructions in LLM Stream", () => {
  let originalDataPath: string
  let originalToolsContent: string | null = null

  beforeAll(() => {
    state.server = Bun.serve({
      port: 0,
      async fetch(req) {
        const next = state.queue.shift()
        if (!next) return new Response("unexpected request", { status: 500 })
        const url = new URL(req.url)
        const body = (await req.json()) as Record<string, unknown>
        next.resolve({ url, headers: req.headers, body })
        if (!url.pathname.endsWith(next.path)) return new Response("not found", { status: 404 })
        return next.response
      },
    })
  })

  afterAll(() => {
    state.server?.stop()
  })

  beforeEach(async () => {
    state.queue.length = 0
    originalDataPath = Global.Path.data
    const testDataDir = path.join(path.dirname(originalDataPath), "tf-agent-test-data")
    ;(Global.Path as { data: string }).data = testDataDir

    // Backup existing tools.json if it exists
    try {
      originalToolsContent = await Bun.file(TF_TOOLS_PATH).text()
    } catch {
      originalToolsContent = null
    }
    await fs.promises.mkdir(path.dirname(TF_TOOLS_PATH), { recursive: true })
  })

  afterEach(async () => {
    await Instance.disposeAll()
    ;(Global.Path as { data: string }).data = originalDataPath

    // Restore original tools.json
    if (originalToolsContent !== null) {
      await fs.promises.writeFile(TF_TOOLS_PATH, originalToolsContent)
    } else {
      try {
        await fs.promises.unlink(TF_TOOLS_PATH)
      } catch {}
    }
  })

  test("includes highlighted TF agent instructions in system prompt", async () => {
    const server = state.server
    if (!server) throw new Error("Server not initialized")

    const providerID = "alibaba"
    const modelID = "qwen-plus"
    const fixture = await loadFixture(providerID, modelID)

    // Setup TF agent with this model
    const toolsData = {
      success: true,
      tools: [
        {
          id: "code-reviewer-123",
          name: "Code Reviewer",
          description: "Reviews code for quality and best practices",
          tool_type: "coder_agent",
          auth_via: "tf_agent",
          interpolation_string:
            "You are a code reviewer. Always check for bugs, security issues, and suggest improvements.",
          goals: "Review all code thoroughly. Provide actionable feedback. Ensure code quality standards.",
          temperature: 0.3,
          max_tokens: 4096,
          llm_base_model: modelID,
          llm_provider: providerID,
        },
      ],
    }

    await fs.promises.writeFile(TF_TOOLS_PATH, JSON.stringify(toolsData, null, 2))

    const request = waitRequest(
      "/chat/completions",
      new Response(createChatStream("I'll review your code."), {
        status: 200,
        headers: { "Content-Type": "text/event-stream" },
      }),
    )

    await using tmp = await tmpdir({
      init: async (dir) => {
        await Bun.write(
          path.join(dir, "opencode.json"),
          JSON.stringify({
            $schema: "https://opencode.ai/config.json",
            enabled_providers: [providerID],
            provider: {
              [providerID]: {
                options: {
                  apiKey: "test-key",
                  baseURL: `${server.url.origin}/v1`,
                },
              },
            },
          }),
        )
      },
    })

    await Instance.provide({
      directory: tmp.path,
      fn: async () => {
        const agent = await Agent.get("Code Reviewer")
        expect(agent).toBeDefined()

        const resolved = await Provider.getModel(ProviderID.make(providerID), ModelID.make(modelID))
        const sessionID = SessionID.make("test-session")
        const user: MessageV2.User = {
          id: MessageID.make("user-1"),
          sessionID,
          role: "user",
          time: { created: Date.now() },
          agent: "Code Reviewer",
          model: { providerID: ProviderID.make(providerID), modelID: ModelID.make(modelID) },
        }

        const stream = await LLM.stream({
          user,
          sessionID,
          model: resolved,
          agent: agent!,
          system: [],
          abort: new AbortController().signal,
          messages: [{ role: "user", content: "Hello" }],
          tools: {},
        })

        for await (const _ of stream.fullStream) {
        }

        const capture = await request
        const body = capture.body
        const messages = body.messages as Array<{ role: string; content: string }>

        const systemMessage = messages.find((m) => m.role === "system")
        expect(systemMessage).toBeDefined()

        const systemContent = systemMessage!.content

        expect(systemContent).toContain("ULTRA IMPORTANT - AGENT CONFIGURATION")
        expect(systemContent).toContain('You are acting as the agent: "Code Reviewer"')
        expect(systemContent).toContain("Reviews code for quality and best practices")
        expect(systemContent).toContain('AGENT "Code Reviewer" INSTRUCTIONS')
        expect(systemContent).toContain(
          "You are a code reviewer. Always check for bugs, security issues, and suggest improvements.",
        )
        expect(systemContent).toContain('AGENT "Code Reviewer" GOALS')
        expect(systemContent).toContain(
          "Review all code thoroughly. Provide actionable feedback. Ensure code quality standards.",
        )
      },
    })
  })

  test("does NOT include highlighted instructions for native agents", async () => {
    const server = state.server
    if (!server) throw new Error("Server not initialized")

    const providerID = "alibaba"
    const modelID = "qwen-plus"
    const fixture = await loadFixture(providerID, modelID)

    const request = waitRequest(
      "/chat/completions",
      new Response(createChatStream("Hello"), {
        status: 200,
        headers: { "Content-Type": "text/event-stream" },
      }),
    )

    await using tmp = await tmpdir({
      init: async (dir) => {
        await Bun.write(
          path.join(dir, "opencode.json"),
          JSON.stringify({
            $schema: "https://opencode.ai/config.json",
            enabled_providers: [providerID],
            provider: {
              [providerID]: {
                options: {
                  apiKey: "test-key",
                  baseURL: `${server.url.origin}/v1`,
                },
              },
            },
          }),
        )
      },
    })

    await Instance.provide({
      directory: tmp.path,
      fn: async () => {
        const agent = await Agent.get("build")
        expect(agent).toBeDefined()
        expect(agent?.native).toBe(true)

        const resolved = await Provider.getModel(ProviderID.make(providerID), ModelID.make(modelID))
        const sessionID = SessionID.make("test-session")
        const user: MessageV2.User = {
          id: MessageID.make("user-1"),
          sessionID,
          role: "user",
          time: { created: Date.now() },
          agent: "build",
          model: { providerID: ProviderID.make(providerID), modelID: ModelID.make(modelID) },
        }

        const stream = await LLM.stream({
          user,
          sessionID,
          model: resolved,
          agent: agent!,
          system: [],
          abort: new AbortController().signal,
          messages: [{ role: "user", content: "Hello" }],
          tools: {},
        })

        for await (const _ of stream.fullStream) {
        }

        const capture = await request
        const body = capture.body
        const messages = body.messages as Array<{ role: string; content: string }>

        const systemMessage = messages.find((m) => m.role === "system")
        expect(systemMessage).toBeDefined()

        const systemContent = systemMessage!.content

        expect(systemContent).not.toContain("ULTRA IMPORTANT - AGENT CONFIGURATION")
        expect(systemContent).not.toContain("You are acting as the agent:")
      },
    })
  })
})