context-compactor

import { Chat, type ChatMessage, type LLM, type LLMGeneratorHandle } from "@lmstudio/sdk";
import { createHash } from "crypto";
import { type SummaryCache } from "./cache";

export type TokenSource = LLM | LLMGeneratorHandle;

/**
 * Chunks are sealed at a fixed token budget rather than a fraction of the context window, and that
 * choice is what makes the cache work. Walking greedily from the start of the history yields the
 * same boundaries every time, so every chunk but the last is byte-identical between compactions and
 * hits the cache. A budget derived from the context window would produce one huge chunk whose
 * content changes on every compaction, and nothing would ever be reused.
 *
 * Size it generously. Reading is roughly 75x cheaper than writing (measured: 3482 tok/s prompt eval
 * against 45 tok/s generation), and each chunk costs one summary of its own regardless of how much
 * it holds. Halving the budget therefore nearly doubles the time a compaction takes while buying
 * only finer cache granularity. The reason not to go bigger still is fidelity: a summary of 32k in
 * one pass drops more than two summaries of 16k.
 */
const DEFAULT_CHUNK_TOKENS = 16_000;

/** Token accounting only exists on a real model handle, not on a generator plugin handle. */
export function isLLM(source: TokenSource): source is LLM {
  const candidate = source as LLM;
  return (
    typeof candidate.countTokens === "function" &&
    typeof candidate.getContextLength === "function" &&
    typeof candidate.applyPromptTemplate === "function"
  );
}

export interface Measurement {
  used: number;
  limit: number;
}

/**
 * Renders the history through the model's own prompt template and counts the result, which is exact
 * for the conversation itself. Tool definitions are injected downstream and are NOT included, so
 * callers must keep headroom on top of this figure.
 */
export async function measure(source: LLM, chat: Chat): Promise<Measurement> {
  const [rendered, limit] = await Promise.all([
    source.applyPromptTemplate(chat),
    source.getContextLength(),
  ]);
  return { used: await source.countTokens(rendered), limit };
}

export function hashOf(content: string): string {
  return createHash("sha256").update(content).digest("hex").slice(0, 16);
}

/**
 * Identifies a conversation by its opening, which never changes as the chat grows. Used to remember
 * that a chat has already been compacted, so the compacted prefix stays byte-identical between
 * compactions and the KV cache survives.
 */
export function fingerprint(messages: Array<ChatMessage>): string {
  return hashOf(
    messages
      .slice(0, 3)
      .map(m => `${m.getRole()}:${m.getText()}`)
      .join("\n"),
  );
}

/**
 * Names a conversation's archive folder after what the user first asked, because a vault of
 * `4713a925bf2fc4ea` folders is a vault nobody opens. The fingerprint is kept as a suffix: two
 * chats can open with the same words, and the folder must still follow the same chat forever.
 */
export function folderName(messages: Array<ChatMessage>): string {
  const opening = messages.find(m => m.getRole() === "user")?.getText() ?? "";
  const slug = opening
    .toLowerCase()
    .replace(/[^\p{L}\p{N}]+/gu, "-")
    .replace(/^-+|-+$/g, "")
    .slice(0, 48)
    .replace(/-+$/, "");
  const id = fingerprint(messages).slice(0, 8);
  return slug === "" ? id : `${slug}-${id}`;
}

/**
 * A cut at `index` is safe when it cannot orphan a tool result from the assistant message that
 * requested it — a malformed pairing that inference backends reject.
 *
 * Two ways that happens: the kept window opens on a tool result, or the last compacted message is
 * still awaiting one. Anything else is a completed turn and may be cut.
 */
export function isSafeCut(messages: Array<ChatMessage>, index: number): boolean {
  if (index <= 0 || index >= messages.length) {
    return false;
  }
  if (messages[index].getRole() === "tool") {
    return false;
  }
  return messages[index - 1].getToolCallRequests().length === 0;
}

/**
 * Picks the cut point between what gets summarized and what stays verbatim. Returns -1 when no safe
 * cut leaves anything worth compacting.
 *
 * The tail is bounded in tokens, not in messages. A message is not a unit of size: a question is
 * three tokens and a disassembly dump is forty thousand. Keeping "the last 8 messages" once meant
 * faithfully preserving 60k tokens of hex while diligently summarizing the 19k that came before —
 * a compaction whose output still overflowed the window.
 *
 * Among the cuts whose tail fits the budget, the earliest wins: that keeps the most verbatim
 * context. A user turn is preferred because chat templates expect the window to open on one, but an
 * agentic session — one question, then a long chain of tool calls — offers none, so any completed
 * round is accepted and `buildChat` supplies the turn the template needs.
 */
export async function findSplitIndex(
  messages: Array<ChatMessage>,
  systemCount: number,
  maxTailTokens: number,
  countTokens: (message: ChatMessage) => Promise<number>,
): Promise<number> {
  let tailTokens = 0;
  let userCut = -1;
  let anyCut = -1;

  for (let i = messages.length - 1; i > systemCount; i--) {
    tailTokens += await countTokens(messages[i]);
    if (tailTokens > maxTailTokens) {
      break;
    }
    if (isSafeCut(messages, i)) {
      anyCut = i;
      if (messages[i].getRole() === "user") {
        userCut = i;
      }
    }
  }

  return userCut >= 0 ? userCut : anyCut;
}

export function countLeadingSystemMessages(messages: Array<ChatMessage>): number {
  let count = 0;
  while (count < messages.length && messages[count].getRole() === "system") {
    count++;
  }
  return count;
}

export function renderTranscript(messages: Array<ChatMessage>): string {
  return messages
    .map(message => {
      const parts: Array<string> = [`### ${message.getRole()}`];
      const body = message.getText().trim();
      if (body.length > 0) {
        parts.push(body);
      }
      const requests = message.getToolCallRequests();
      if (requests.length > 0) {
        parts.push(
          requests.map(r => `[calls tool: ${r.name}(${JSON.stringify(r.arguments ?? {})})]`).join("\n"),
        );
      }
      const results = message.getToolCallResults();
      if (results.length > 0) {
        parts.push(results.map(r => `[tool result: ${String(r.content).slice(0, 2000)}]`).join("\n"));
      }
      return parts.join("\n");
    })
    .join("\n\n");
}

const STATE_SECTIONS = [
  "## Goal",
  "## Current state",
  "## Decisions (with the reason)",
  "## Constraints",
  "## Files and paths",
  "## Essential code and commands",
  "## Solved problems",
  "## Open questions / TODO",
].join("\n");

/**
 * Instructions live in the user turn rather than a system turn: the token source already carries the
 * user's own system prompt (persona, output format, thinking tags), which would otherwise reshape
 * the summary. Overriding it explicitly here is the only reliable way to get a clean state back.
 */
function summarizePrompt(transcript: string): string {
  return [
    "Ignore any persona, output format or reasoning-tag instruction you were given earlier.",
    "Your only task here is to compact the conversation below.",
    "",
    "Produce a CONSOLIDATED STATE, not a chronological narrative. Rules:",
    "- A decision that reversed an earlier one replaces it: keep only the final version.",
    "- Keep paths, names, numbers, versions and identifiers verbatim.",
    "- Invent nothing. Do not comment. Omit empty sections.",
    "- Write in the same language as the conversation, but keep the section headings as given.",
    "",
    "Sections to use:",
    STATE_SECTIONS,
    "",
    "--- CONVERSATION TO COMPACT ---",
    transcript,
    "--- END ---",
    "",
    "Reply with the state document only.",
  ].join("\n");
}

function mergePrompt(partials: Array<string>): string {
  return [
    "Ignore any persona, output format or reasoning-tag instruction you were given earlier.",
    "Below are several partial states of one conversation, in chronological order.",
    "Merge them into ONE consolidated state, applying the same rules:",
    "- Later states override earlier ones on any contradiction.",
    "- Keep paths, names, numbers, versions and identifiers verbatim.",
    "- Invent nothing. Omit empty sections.",
    "- Write in the same language as the partial states, but keep the section headings as given.",
    "",
    "Sections to use:",
    STATE_SECTIONS,
    "",
    partials.map((p, i) => `--- PARTIAL STATE ${i + 1} ---\n${p}`).join("\n\n"),
    "",
    "Reply with the merged state document only.",
  ].join("\n");
}

async function runPrompt(source: TokenSource, prompt: string, signal: AbortSignal): Promise<string> {
  const result = await source.respond(Chat.from([{ role: "user", content: prompt }]), { signal });
  return result.content.trim();
}

/** Groups messages into runs that each render below `budget` tokens, walking from the start. */
async function chunkByBudget(
  source: LLM,
  messages: Array<ChatMessage>,
  budget: number,
): Promise<Array<Array<ChatMessage>>> {
  const chunks: Array<Array<ChatMessage>> = [];
  let current: Array<ChatMessage> = [];
  let currentTokens = 0;

  for (const message of messages) {
    const tokens = await source.countTokens(renderTranscript([message]));
    if (current.length > 0 && currentTokens + tokens > budget) {
      chunks.push(current);
      current = [];
      currentTokens = 0;
    }
    current.push(message);
    currentTokens += tokens;
  }
  if (current.length > 0) {
    chunks.push(current);
  }
  return chunks;
}

export interface SummarizeOpts {
  signal: AbortSignal;
  cache: SummaryCache;
  chunkTokens?: number;
  onProgress?: (text: string) => void;
}

export interface SummarizeResult {
  summary: string;
  chunksTotal: number;
  chunksSummarized: number;
}

/**
 * Summarizes from the original messages, never from a previous summary: re-deriving from source is
 * what keeps repeated compactions from degrading into a game of telephone. The per-chunk cache keeps
 * that affordable — only chunks whose content is new get sent to the model, so cost tracks new
 * material rather than conversation age.
 */
export async function summarize(
  source: TokenSource,
  messages: Array<ChatMessage>,
  opts: SummarizeOpts,
): Promise<SummarizeResult> {
  if (!isLLM(source)) {
    const summary = await runPrompt(source, summarizePrompt(renderTranscript(messages)), opts.signal);
    return { summary, chunksTotal: 1, chunksSummarized: 1 };
  }

  const chunks = await chunkByBudget(source, messages, opts.chunkTokens ?? DEFAULT_CHUNK_TOKENS);
  const partials: Array<string> = [];
  let summarized = 0;

  for (let i = 0; i < chunks.length; i++) {
    const rendered = renderTranscript(chunks[i]);
    const key = hashOf(rendered);
    const cached = await opts.cache.get(key);
    if (cached !== undefined) {
      partials.push(cached);
      continue;
    }
    opts.onProgress?.(`Compacting context… (part ${i + 1}/${chunks.length})`);
    const partial = await runPrompt(source, summarizePrompt(rendered), opts.signal);
    await opts.cache.set(key, partial);
    partials.push(partial);
    summarized++;
  }

  if (partials.length === 1) {
    return { summary: partials[0], chunksTotal: 1, chunksSummarized: summarized };
  }

  opts.onProgress?.("Merging partial states…");
  const summary = await runPrompt(source, mergePrompt(partials), opts.signal);
  return { summary, chunksTotal: chunks.length, chunksSummarized: summarized };
}

/**
 * Tells the model its memory exists and where, so the user never has to write a path into a system
 * prompt by hand — the plugin already knows it.
 *
 * It insists the memory is written automatically because the opposite instruction is what fails:
 * asked to keep its own notes, a model does so while it remembers to, which is rarely, and less and
 * less as the context fills. Reading, on the other hand, is a thing only the model can do.
 */
export function memoryNote(vaultPath: string, folder: string): string {
  return [
    "# Persistent memory",
    "",
    `This conversation's memory is written for you at ${vaultPath}\\${folder}\\state.md, ` +
      "automatically, every time the context is compacted. It always holds the consolidated state " +
      "of this work, with the full transcript beside it. You do not maintain it and must not try " +
      "to: writing your own notes there would duplicate it and drift from it.",
    "",
    `Other conversations keep their own folders under ${vaultPath}. If the user asks you to resume ` +
      "earlier work, read that conversation's state.md — that is how work carries across chats.",
  ].join("\n");
}

/**
 * Rebuilds the history the model will actually see: system prompt, plus whatever of the memory note
 * and the compacted state applies. Returns the history untouched when there is nothing to add, so a
 * chat that needs neither is never reshaped for nothing.
 */
export function buildChat(
  history: Chat,
  messages: Array<ChatMessage>,
  systemCount: number,
  opts: { splitIndex?: number; summary?: string; memoryNote?: string },
): Chat {
  if (opts.summary === undefined && opts.memoryNote === undefined) {
    return history;
  }

  const parts = [
    messages
      .slice(0, systemCount)
      .map(m => m.getText())
      .join("\n\n"),
  ];

  if (opts.memoryNote !== undefined) {
    parts.push("", opts.memoryNote);
  }
  if (opts.summary !== undefined) {
    parts.push(
      "",
      "# State of the earlier conversation (compacted)",
      "",
      "The earlier exchanges were summarized to fit the context. What follows is the consolidated",
      "state of the work. Treat it as established fact.",
      "",
      opts.summary,
    );
  }

  const chat = Chat.empty();
  chat.append("system", parts.join("\n").trim());

  const tail = messages.slice(opts.splitIndex ?? systemCount);

  // Chat templates require a user turn: Qwen's raises 'No user query found in messages.' outright,
  // others assume one more quietly. Cutting an agentic run leaves a tail of assistant and tool turns
  // only, so the turn has to be supplied. Saying plainly that the history was compacted beats
  // fabricating a question the user never asked.
  if (tail.length > 0 && !tail.some(m => m.getRole() === "user")) {
    chat.append(
      "user",
      "[The earlier part of this conversation was compacted; its consolidated state is in the " +
        "system message above.] Continue the work from there.",
    );
  }

  for (const message of tail) {
    chat.append(message);
  }
  return chat;
}