MagicCode/packages/core/src/agent.ts

import { estimateTextTokens, estimateTokens } 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;
  /** Taille du contexte (tokens) après la dernière réponse du LLM : valeur renvoyée par le serveur, sinon estimation. */
  contextTokens: number;
  /** true : contextTokens vient du serveur (exact) ; false : estimation. */
  contextExact: boolean;
  /** Facteur de calibration de l'estimation (tokens réels / tokens estimés), mesuré à chaque réponse du serveur. */
  tokenScale: 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;
  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,
    contextTokens: 0,
    contextExact: false,
    tokenScale: 1,
    pendingContext: [],
    ui: opts.ui,
    cwd: opts.cwd,
  };
}

const MAX_ITERATIONS = 25;
const MIN_CALIBRATION_TOKENS = 200; // en dessous, le rapport réel/estimé est trop instable pour calibrer
const MIN_SCALE = 0.25;
const MAX_SCALE = 4;


/** 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;
    const sentEstimate = rawSessionEstimate(session); // taille estimée de la requête, pour calibrer l'estimation

    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;
          if (sentEstimate >= MIN_CALIBRATION_TOKENS && ev.inputTokens > 0) {
            session.tokenScale = Math.min(MAX_SCALE, Math.max(MIN_SCALE, ev.inputTokens / sentEstimate));
          }
        }
        // "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);
    session.contextExact = usedTokens !== null;
    ui.onContextUsage?.(session.contextTokens, session.contextExact);

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

    // 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 ${MAX_ITERATIONS} itérations`, true);
}

type Estimable = Pick<AgentSession, "system" | "messages" | "tools">;

function rawSessionEstimate(session: Estimable): number {
  return estimateTokens(session.system, session.messages, [...session.tools.values()].map((t) => t.spec));
}

/** Estimation calibrée du contexte de la conversation (prompt système, outils, messages), sans le contexte en attente. */
export function estimateSessionTokens(session: Estimable & Pick<AgentSession, "tokenScale">): number {
  return Math.round(rawSessionEstimate(session) * session.tokenScale);
}

/**
 * Taille du contexte de la session en cours : dernier compte connu (serveur ou estimation) + contexte en attente
 * (résultats de commandes « ! », fichiers cités) qui sera joint à la prochaine demande.
 */
export function sessionContextTokens(
  session: Pick<AgentSession, "contextTokens" | "pendingContext" | "tokenScale">,
): number {
  const pending = session.pendingContext.reduce((sum, text) => sum + estimateTextTokens(text), 0);
  return session.contextTokens + Math.round(pending * session.tokenScale);
}