MagicCode/packages/core/src/agent.ts

import { estimateTokens } from "./context";
import type {
  AgentSkill,
  AgentUI,
  ContentBlock,
  LLMProvider,
  Message,
  ModeInitResult,
  PermissionPolicy,
  PermissionSource,
  SlashCommand,
  Tool,
} from "./types";

export interface AgentSession {
  provider: LLMProvider;
  system: string;
  tools: Map<string, Tool>;
  commands: SlashCommand[];
  skills: AgentSkill[];
  permissions?: PermissionPolicy;
  /** Fichiers dont viennent les règles de permissions (affichage avec /permissions). */
  permissionSources: PermissionSource[];
  env?: Record<string, string>;
  messages: Message[];
  /** Taille du contexte (tokens) après la dernière réponse du LLM : valeur renvoyée par le serveur, sinon estimation. */
  contextTokens: number;
  /** Contexte ajouté par l'utilisateur hors conversation (ex. résultat d'une commande "!"), joint à sa prochaine demande. */
  pendingContext: string[];
  ui: AgentUI;
  cwd: string;
}

export function createSession(opts: {
  provider: LLMProvider;
  init: ModeInitResult;
  ui: AgentUI;
  cwd: string;
}): AgentSession {
  return {
    provider: opts.provider,
    system: opts.init.systemPrompt,
    tools: new Map(opts.init.tools.map((t) => [t.spec.name, t])),
    commands: opts.init.commands ?? [],
    skills: opts.init.skills ?? [],
    permissions: opts.init.permissions,
    permissionSources: opts.init.permissionSources ?? [],
    env: opts.init.env,
    messages: [],
    contextTokens: 0,
    pendingContext: [],
    ui: opts.ui,
    cwd: opts.cwd,
  };
}

const MAX_ITERATIONS = 25;

type ToolUseBlock = Extract<ContentBlock, { type: "tool_use" }>;

/** Un tour complet : message utilisateur → (LLM → outils)* → réponse finale. */
export async function runTurn(
  session: AgentSession,
  userText: string,
  signal?: AbortSignal,
): Promise<void> {
  const { ui } = session;
  // Le contexte en attente est joint à la demande (un seul message : certains providers exigent l'alternance user/assistant)
  const context = session.pendingContext.splice(0);
  const text = context.length > 0 ? `${context.join("\n\n")}\n\n${userText}` : userText;
  session.messages.push({ role: "user", content: [{ type: "text", text }] });

  for (let i = 0; i < MAX_ITERATIONS; i++) {
    const assistant: ContentBlock[] = [];
    let text = "";
    let usedTokens: number | null = null;

    ui.onThinking(true);
    let first = true;
    try {
      for await (const ev of session.provider.chat(
        {
          system: session.system,
          messages: session.messages,
          tools: [...session.tools.values()].map((t) => t.spec),
        },
        signal,
      )) {
        if (first) {
          ui.onThinking(false);
          first = false;
        }
        if (ev.type === "text_delta") {
          text += ev.text;
          ui.onText(ev.text);
        } else if (ev.type === "tool_use") {
          assistant.push({ type: "tool_use", id: ev.id, name: ev.name, input: ev.input });
        } else if (ev.type === "usage") {
          // Requête + réponse : c'est la taille du contexte une fois la réponse ajoutée à la conversation
          usedTokens = ev.inputTokens + ev.outputTokens;
        }
        // "done" : la fin du flux suffit, les appels d'outils décident de la suite
      }
    } finally {
      if (first) ui.onThinking(false);
    }

    if (text) {
      ui.onTextEnd();
      assistant.unshift({ type: "text", text });
    }
    session.messages.push({ role: "assistant", content: assistant });
    session.contextTokens = usedTokens ?? estimateSessionTokens(session);
    ui.onContextUsage?.(session.contextTokens);

    const calls = assistant.filter((b): b is ToolUseBlock => b.type === "tool_use");
    if (calls.length === 0) return;

    const results: ContentBlock[] = [];
    for (const call of calls) {
      results.push(await executeToolCall(session, call, signal));
    }
    session.messages.push({ role: "user", content: results });
  }
  session.ui.onToolResult("agent", `Arrêt : plus de ${MAX_ITERATIONS} itérations`, true);
}

async function executeToolCall(
  session: AgentSession,
  call: ToolUseBlock,
  signal?: AbortSignal,
): Promise<ContentBlock> {
  const { ui } = session;
  const tool = session.tools.get(call.name);
  const result = (content: string, isError = false): ContentBlock => ({
    type: "tool_result",
    toolUseId: call.id,
    content,
    isError,
  });

  if (!tool) return result(`Outil inconnu : ${call.name}`, true);

  const summary = tool.describe?.(call.input) ?? JSON.stringify(call.input).slice(0, 120);
  ui.onToolCall(call.name, summary);

  // Les règles de l'agent (deny > ask > allow) priment sur la permission par défaut de l'outil
  const verdict = session.permissions?.check(call.name, call.input, { cwd: session.cwd });
  if (verdict?.decision === "deny") {
    const why = verdict.reason ? ` (${verdict.reason})` : "";
    ui.onToolResult(call.name, `Refusé par la configuration${why}`, true);
    return result(`Action refusée par les règles de permissions${why}. Ne réessaie pas telle quelle.`, true);
  }
  const needsConfirm = verdict ? verdict.decision === "ask" : tool.permission === "ask";

  if (needsConfirm) {
    const ok = await ui.confirm(`Autoriser ${call.name} : ${summary} ?`);
    if (!ok) {
      ui.onToolResult(call.name, "Refusé par l'utilisateur", true);
      return result("L'utilisateur a refusé cette action.", true);
    }
  }

  try {
    const output = await tool.execute(call.input, { cwd: session.cwd, model: session.provider.model, signal, env: session.env });
    ui.onToolResult(call.name, output, false);
    return result(output);
  } catch (err) {
    const msg = (err as Error).message;
    ui.onToolResult(call.name, msg, true);
    return result(msg, true);
  }
}

/** Estimation de la taille du contexte d'une session (sans réponse du serveur) : prompt système, outils et messages. */
export function estimateSessionTokens(session: Pick<AgentSession, "system" | "messages" | "tools">): number {
  return estimateTokens(session.system, session.messages, [...session.tools.values()].map((t) => t.spec));
}