MagicCode/packages/core/src/agent.ts

import { countContextTokens } from "./context";
import { runToolCallsDetailed, type ToolUseBlock } from "./tool-runner";
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[];
  /** Résumé compacté qui sert de contexte de départ (présent dans `system`) ; absent : contexte MAIA.md / CLAUDE.md. */
  startContext?: string;
  /** 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;
  startContext?: 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: [],
    startContext: opts.startContext,
    pendingContext: [],
    ui: opts.ui,
    cwd: opts.cwd,
  };
}

/** Nombre d'itérations par défaut avant de s'arrêter (ou de demander s'il faut continuer). */
export const MAX_ITERATIONS = 50;

/** Fin d'un tour : réponse finale du modèle, ou arrêt faute d'itérations (le dernier message est alors un résultat d'outil). */
/** done : réponse finale · max_iterations : limite atteinte · stopped : l'utilisateur a répondu « stop » (rien n'est annulé). */
export type TurnEnd = "done" | "max_iterations" | "stopped";


/** Un tour complet : message utilisateur → (LLM → outils)* → réponse finale. */
export async function runTurn(
  session: AgentSession,
  userText: string,
  signal?: AbortSignal,
  opts: {
    maxIterations?: number;
    /**
     * Appelée quand la limite d'itérations est atteinte sans réponse finale : true → continuer pour `maxIterations`
     * itérations de plus, false → s'arrêter. Absente : arrêt avec un message d'erreur.
     */
    onLimit?: (iterations: number) => Promise<boolean>;
  } = {},
): Promise<TurnEnd> {
  const maxIterations = opts.maxIterations ?? MAX_ITERATIONS;
  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;
  // Après un arrêt à la limite d'itérations, le dernier message est déjà un message « utilisateur » (résultats d'outils) :
  // la nouvelle demande s'y ajoute (l'alternance user/assistant est exigée par certains serveurs)
  const last = session.messages.at(-1);
  if (last?.role === "user") {
    session.messages[session.messages.length - 1] = { ...last, content: [...last.content, { type: "text", text }] };
  } else {
    session.messages.push({ role: "user", content: [{ type: "text", text }] });
  }
  ui.onContextUsage?.(sessionContextTokens(session)); // la demande (et le contexte joint) occupe déjà des tokens

  let iterations = 0;
  let limit = maxIterations;
  for (;;) {
    if (iterations >= limit) {
      if (opts.onLimit && (await opts.onLimit(iterations))) limit += maxIterations;
      else break;
    }
    iterations++;
    const assistant: ContentBlock[] = [];
    let text = "";
    let truncated = false; // réponse coupée par la limite de tokens de sortie du modèle
    let reportedOutput = 0; // tokens de sortie annoncés par le serveur

    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") {
          reportedOutput = ev.outputTokens;
        } else if (ev.type === "done" && ev.stopReason === "max_tokens") {
          truncated = true;
        }
        // La consommation du serveur ne sert pas à la taille du contexte (comptée localement), seulement à repérer une coupure.
        // "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 });
    ui.onContextUsage?.(sessionContextTokens(session));
    // Certains serveurs (LiteLLM derrière Bedrock/Anthropic…) « réparent » un appel d'outil coupé par la limite de sortie et
    // annoncent « tool_calls » au lieu de « length » : une sortie qui tombe pile sur une limite usuelle (4096…) est suspecte.
    if (!truncated && assistant.some((b) => b.type === "tool_use") && isUsualOutputLimit(reportedOutput)) truncated = true;
    if (truncated) {
      ui.onToolResult(
        "agent",
        "Réponse probablement coupée par la limite de tokens de sortie du modèle : un texte ou un appel d'outil (ex. le contenu d'un fichier à écrire) peut être incomplet. " +
          "Demandez un contenu plus court ou, si la limite du serveur est basse, fixez « MaxOutputTokens » dans ~/.maia/settings.json.",
        true,
      );
    }

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

    // 1 appel : premier plan ; 2 appels ou plus : une tâche de fond par outil, on attend toutes les fins
    const { results, stopped } = await runToolCallsDetailed(session, calls, signal);
    session.messages.push({ role: "user", content: results });
    ui.onContextUsage?.(sessionContextTokens(session)); // les résultats d'outils pèsent avant l'itération suivante
    // « stop » : le contexte est conservé tel quel ; le prochain message de l'utilisateur s'ajoute à ces résultats
    if (stopped) return "stopped";
  }
  session.ui.onToolResult("agent", `Arrêt : plus de ${iterations} itérations`, true);
  return "max_iterations";
}

/**
 * Taille du contexte de la session en cours, comptée avec js-tiktoken : prompt système, outils, messages, plus le
 * contexte en attente (résultats de commandes « ! », fichiers cités) qui sera joint à la prochaine demande.
 */
export function sessionContextTokens(
  session: Pick<AgentSession, "system" | "messages" | "tools" | "pendingContext">,
): number {
  return countContextTokens({
    system: session.system,
    messages: session.messages,
    tools: [...session.tools.values()].map((t) => t.spec),
    pending: session.pendingContext,
  });
}

const USUAL_OUTPUT_LIMITS = new Set([1024, 2048, 4096, 8192, 16384, 32768, 65536]);

/** Vrai si le nombre de tokens de sortie est exactement une limite usuelle (réponse très probablement coupée). */
export function isUsualOutputLimit(outputTokens: number): boolean {
  return USUAL_OUTPUT_LIMITS.has(outputTokens);
}