MagicCode/packages/core/src/agent.test.ts

import { describe, expect, it } from "vitest";
import { createSession, estimateSessionTokens, runTurn, sessionContextTokens } from "./agent";
import type { AgentUI, LLMProvider, PermissionPolicy, StreamEvent, Tool } from "./types";

const silentUI: AgentUI = {
  onThinking() {},
  onText() {},
  onTextEnd() {},
  onToolCall() {},
  onToolResult() {},
  confirm: async () => true,
};

/** Provider factice : 1er appel → demande d'outil, 2e appel → réponse texte. */
function fakeProvider(toolName = "echo"): LLMProvider {
  let call = 0;
  return {
    id: "fake",
    model: "fake",
    capabilities: { tools: true },
    async *chat(): AsyncIterable<StreamEvent> {
      call++;
      if (call === 1) {
        yield { type: "tool_use", id: "t1", name: toolName, input: { value: "salut" } };
        yield { type: "done", stopReason: "tool_use" };
      } else {
        yield { type: "text_delta", text: "terminé" };
        yield { type: "done", stopReason: "end_turn" };
      }
    },
  };
}

describe("runTurn", () => {
  it("exécute l'outil demandé puis renvoie la réponse finale", async () => {
    const echo: Tool = {
      spec: { name: "echo", description: "echo", inputSchema: { type: "object" } },
      permission: "auto",
      execute: async (input) => String(input.value),
    };
    const session = createSession({
      provider: fakeProvider(),
      init: { systemPrompt: "test", tools: [echo] },
      ui: silentUI,
      cwd: process.cwd(),
    });

    await runTurn(session, "bonjour");

    // user, assistant(tool_use), user(tool_result), assistant(texte)
    expect(session.messages).toHaveLength(4);
    expect(session.messages[2]?.content[0]).toMatchObject({
      type: "tool_result",
      content: "salut",
    });
  });
});

describe("politique de permissions", () => {
  const makeTool = (ran: string[]): Tool => ({
    spec: { name: "danger", description: "d", inputSchema: { type: "object" } },
    permission: "ask",
    execute: async () => {
      ran.push("x");
      return "ok";
    },
  });
  const policy = (decision: "allow" | "ask" | "deny"): PermissionPolicy => ({
    check: () => ({ decision, reason: "règle de test" }),
  });

  async function run(decision: "allow" | "ask" | "deny", answer: boolean) {
    const ran: string[] = [];
    const confirms: string[] = [];
    const session = createSession({
      provider: fakeProvider("danger"),
      init: { systemPrompt: "t", tools: [makeTool(ran)], permissions: policy(decision) },
      ui: { ...silentUI, confirm: async (m) => (confirms.push(m), answer) },
      cwd: process.cwd(),
    });
    await runTurn(session, "go");
    return { ran, confirms, session };
  }

  it("allow : exécute sans confirmation", async () => {
    const { ran, confirms } = await run("allow", false);
    expect(ran).toHaveLength(1);
    expect(confirms).toHaveLength(0);
  });

  it("deny : refuse sans demander", async () => {
    const { ran, confirms, session } = await run("deny", true);
    expect(ran).toHaveLength(0);
    expect(confirms).toHaveLength(0);
    expect(session.messages[2]?.content[0]).toMatchObject({ type: "tool_result", isError: true });
  });

  it("ask : demande confirmation même pour un outil automatique", async () => {
    const { ran, confirms } = await run("ask", true);
    expect(confirms).toHaveLength(1);
    expect(ran).toHaveLength(1);
  });
});

describe("contexte en attente", () => {
  it("est joint à la prochaine demande puis vidé", async () => {
    const session = createSession({
      provider: fakeProvider("echo"),
      init: {
        systemPrompt: "t",
        tools: [{ spec: { name: "echo", description: "e", inputSchema: { type: "object" } }, permission: "auto", execute: async () => "ok" }],
      },
      ui: silentUI,
      cwd: process.cwd(),
    });
    session.pendingContext.push("CONTEXTE-A", "CONTEXTE-B");

    await runTurn(session, "ma demande");

    const first = session.messages[0]?.content[0];
    expect(first).toMatchObject({ type: "text", text: "CONTEXTE-A\n\nCONTEXTE-B\n\nma demande" });
    expect(session.pendingContext).toEqual([]);
  });
});

describe("taille du contexte", () => {
  const echo: Tool = {
    spec: { name: "echo", description: "echo", inputSchema: { type: "object" } },
    permission: "auto",
    execute: async () => "ok",
  };

  it("utilise la consommation renvoyée par le serveur après chaque réponse", async () => {
    const seen: number[] = [];
    let call = 0;
    const provider: LLMProvider = {
      id: "fake",
      model: "fake",
      capabilities: { tools: true },
      async *chat(): AsyncIterable<StreamEvent> {
        call++;
        if (call === 1) {
          yield { type: "tool_use", id: "t1", name: "echo", input: {} };
          yield { type: "usage", inputTokens: 1000, outputTokens: 50 };
        } else {
          yield { type: "text_delta", text: "fin" };
          yield { type: "usage", inputTokens: 1200, outputTokens: 30 };
        }
        yield { type: "done", stopReason: call === 1 ? "tool_use" : "end_turn" };
      },
    };
    const session = createSession({
      provider,
      init: { systemPrompt: "t", tools: [echo] },
      ui: { ...silentUI, onContextUsage: (n) => seen.push(n) },
      cwd: process.cwd(),
    });
    await runTurn(session, "go");
    expect(seen).toEqual([1050, 1230]);
    expect(session.contextTokens).toBe(1230);
  });

  it("sans consommation renvoyée : estimation à partir de la conversation", async () => {
    const seen: number[] = [];
    const session = createSession({
      provider: fakeProvider("echo"),
      init: { systemPrompt: "x".repeat(4000), tools: [echo] },
      ui: { ...silentUI, onContextUsage: (n) => seen.push(n) },
      cwd: process.cwd(),
    });
    await runTurn(session, "go");
    expect(seen).toHaveLength(2);
    expect(seen[0]).toBeGreaterThanOrEqual(1000); // prompt système de 4000 caractères
    expect(seen[1]).toBeGreaterThan(seen[0]!);
    expect(session.contextTokens).toBe(seen[1]);
  });

  it("valeur du serveur : marquée exacte ; estimation : marquée estimée", async () => {
    const flags: (boolean | undefined)[] = [];
    let call = 0;
    const withUsage: LLMProvider = {
      id: "f",
      model: "f",
      capabilities: { tools: false },
      async *chat(): AsyncIterable<StreamEvent> {
        call++;
        yield { type: "text_delta", text: "ok" };
        yield { type: "usage", inputTokens: 500, outputTokens: 5 };
        yield { type: "done", stopReason: "end_turn" };
      },
    };
    const a = createSession({ provider: withUsage, init: { systemPrompt: "t", tools: [] }, ui: { ...silentUI, onContextUsage: (_n, exact) => flags.push(exact) }, cwd: "." });
    await runTurn(a, "x");
    const b = createSession({ provider: fakeProvider("echo"), init: { systemPrompt: "t", tools: [echo] }, ui: { ...silentUI, onContextUsage: (_n, exact) => flags.push(exact) }, cwd: "." });
    await runTurn(b, "x");
    expect(call).toBe(1);
    expect(flags).toEqual([true, false, false]);
    expect(a.contextExact).toBe(true);
    expect(b.contextExact).toBe(false);
  });

  it("calibre l'estimation sur la consommation réelle du serveur", async () => {
    // Le serveur compte 2× moins de tokens que l'estimation brute : l'estimation suivante doit en tenir compte
    const provider: LLMProvider = {
      id: "f",
      model: "f",
      capabilities: { tools: false },
      async *chat(req): AsyncIterable<StreamEvent> {
        const raw = Math.ceil(req.system.length / 4);
        yield { type: "text_delta", text: "ok" };
        yield { type: "usage", inputTokens: Math.round(raw / 2) + 10, outputTokens: 1 };
        yield { type: "done", stopReason: "end_turn" };
      },
    };
    const session = createSession({ provider, init: { systemPrompt: "x".repeat(8000), tools: [] }, ui: silentUI, cwd: "." });
    expect(session.tokenScale).toBe(1);
    await runTurn(session, "salut");
    expect(session.tokenScale).toBeGreaterThan(0.4);
    expect(session.tokenScale).toBeLessThan(0.6);
    expect(estimateSessionTokens(session)).toBeLessThan(1200); // 8000 caractères ≈ 2000 tokens bruts, ≈ 1000 calibrés
  });

  it("le contexte en attente (commandes « ! ») est compté dans la taille de la session", () => {
    const session = createSession({ provider: fakeProvider("echo"), init: { systemPrompt: "t", tools: [] }, ui: silentUI, cwd: "." });
    session.contextTokens = 1000;
    expect(sessionContextTokens(session)).toBe(1000);
    session.pendingContext.push("y".repeat(4000)); // ≈ 1000 tokens
    expect(sessionContextTokens(session)).toBe(2000);
    session.tokenScale = 0.5;
    expect(sessionContextTokens(session)).toBe(1500);
  });
});