MagicCode/packages/core/src/agent.ts

import { countContextTokens } from "./context";
import { runToolCalls, 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,
  };
}

const MAX_ITERATIONS = 25;

/** 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). */
export type TurnEnd = "done" | "max_iterations";


/** Un tour complet : message utilisateur → (LLM → outils)* → réponse finale. */
export async function runTurn(
  session: AgentSession,
  userText: string,
  signal?: AbortSignal,
  opts: { maxIterations?: number } = {},
): 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;
  session.messages.push({ role: "user", content: [{ type: "text", text }] });

  for (let i = 0; i < maxIterations; i++) {
    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 = await runToolCalls(session, calls, signal);
    session.messages.push({ role: "user", content: results });
  }
  session.ui.onToolResult("agent", `Arrêt : plus de ${maxIterations} 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);
}