context-compactor

import {
  type Chat,
  type PredictionLoopHandlerController,
  type PredictionProcessContentBlockController,
} from "@lmstudio/sdk";
import { archive, loadRecord, saveRecord } from "./archive";
import { createCache } from "./cache";
import {
  buildChat,
  countLeadingSystemMessages,
  findSplitIndex,
  fingerprint,
  folderName,
  isLLM,
  measure,
  memoryNote,
  renderTranscript,
  summarize,
  type TokenSource,
} from "./compaction";
import { configSchematics } from "./config";
import { capTools } from "./toolCap";

// The SDK declares these but does not export them, so derive them from the controller instead.
type ToolUseSession = Awaited<ReturnType<PredictionLoopHandlerController["startToolUseSession"]>>;
type SessionTool = ToolUseSession["tools"][number];

export async function handlePredictionLoop(ctl: PredictionLoopHandlerController) {
  const history = await ctl.pullHistory();
  const source = await ctl.tokenSource();
  const chatToSend = await prepareChat(ctl, source, history);

  let session: ToolUseSession | undefined;
  let tools: Array<SessionTool> = [];
  let sessionOpened = false;
  try {
    session = await ctl.startToolUseSession();
    tools = session.tools;
    sessionOpened = true;
  } catch (error) {
    // The tool session API is still experimental. Degrading to a text-only reply is survivable;
    // doing it silently is not, because tools vanishing would look like the model getting dumber.
    ctl.debug("Tool use session unavailable:", error);
    ctl.createStatus({
      status: "error",
      text: "Tools are unavailable for this reply (the tool session was refused).",
    });
  }

  ctl.debug(`Tool use session: opened=${sessionOpened}, tools=${tools.length}`);

  // A session that opens but hands back nothing is the failure that looks like success: the model
  // gets no tools, cannot say so, and narrates tool calls it never made. Only worth flagging when
  // the session itself worked — a refused session already reported above.
  if (sessionOpened && tools.length === 0) {
    ctl.createStatus({
      status: "error",
      text:
        "No tools were offered for this reply, so the model has none — anything it says about " +
        "calling one is invented. Check that your tool plugins are enabled for this chat.",
    });
  }

  const maxToolResultTokens = ctl.getPluginConfig(configSchematics).get("maxToolResultTokens");
  const budget = await toolBudget(source, chatToSend);

  try {
    await runPrediction(
      ctl,
      source,
      chatToSend,
      capTools(tools, { maxPerResult: maxToolResultTokens, budget }),
    );
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    const name = error instanceof Error ? error.name : "";

    // A malformed tool call is not a plugin failure — it is the model, usually one running out of
    // room to finish the call and truncating it mid-JSON. LM Studio's own loop absorbs this and lets
    // the model try again; if it propagates here, it kills a reply that was otherwise fine. Report it
    // and swallow it. The wider reserve below makes the underlying truncation rarer.
    if (name === "ToolCallRequestError" || /ToolCallRequest/.test(message)) {
      ctl.debug("Malformed tool call from the model:", message);
      ctl.createStatus({
        status: "error",
        text:
          "The model produced a malformed tool call — often a sign it ran out of room to finish it. " +
          "This reply was stopped; send another message to continue.",
      });
      return;
    }

    if (/exceed_context_size|exceeds the available context size/.test(message)) {
      ctl.createStatus({
        status: "error",
        text:
          "The context overflowed during this reply. Tool results accumulate while the model works, " +
          "after compaction has already run — so several large ones in a single turn can still " +
          `fill the window. Lower "Maximum size of a tool result" (currently ${maxToolResultTokens} ` +
          "tokens) or ask for less at a time.",
      });
      return;
    }

    // Anything else is unexpected: surface it rather than swallow it.
    throw error;
  } finally {
    session?.[Symbol.dispose]();
  }
}

/**
 * How many tokens of tool output this reply can afford: whatever the window has left once the prompt
 * is in it, minus room for the model's own answer and for the tool definitions.
 *
 * Tool definitions are the awkward part. They are injected downstream of `applyPromptTemplate`, so
 * the measurement never sees them, and with several MCP servers loaded they are worth thousands of
 * tokens on every single request. The reserve below is a guess at their weight, and a deliberately
 * generous one: overestimating costs a little verbatim tool output, underestimating costs the reply.
 */
async function toolBudget(source: TokenSource, chat: Chat): Promise<number> {
  if (!isLLM(source)) {
    return 4000;
  }
  // A quarter of the window, floor 8k. The reserve has to cover the model's whole reply, and an
  // agentic reply is many rounds of thinking plus tool calls, not one paragraph. Too thin and the
  // model runs out of room mid-tool-call, truncating it into a ToolCallRequestError. Overestimating
  // only costs some verbatim tool output; underestimating costs the reply.
  const { used, limit } = await measure(source, chat);
  const reserve = Math.max(8000, Math.floor(limit * 0.25));
  return Math.max(0, limit - used - reserve);
}

/**
 * Decides what the model actually sees. The chat displayed to the user is never touched: pullHistory
 * hands back a copy, so shrinking it here shortens the prompt while leaving every message on screen.
 */
async function prepareChat(
  ctl: PredictionLoopHandlerController,
  source: TokenSource,
  history: Chat,
): Promise<Chat> {
  const config = ctl.getPluginConfig(configSchematics);
  const triggerPercent = config.get("triggerPercent");
  const keepRecentTokens = config.get("keepRecentTokens");
  const vaultPath = config.get("vaultPath").trim();
  const showStatus = config.get("showStatus");
  const archiving = vaultPath !== "";

  if (!isLLM(source)) {
    return history;
  }

  const messages = history.getMessagesArray();
  const systemCount = countLeadingSystemMessages(messages);
  const key = fingerprint(messages);
  const folder = folderName(messages);

  // Announced on every turn, not just after a compaction: the model has to know its memory exists
  // before it has one, or it will never think to read anyone else's.
  const note = archiving ? memoryNote(vaultPath, folder) : undefined;
  const saved = await loadRecord(vaultPath, key);

  // Reuse the previous compaction verbatim while it still fits: an identical prefix is what lets the
  // KV cache survive between messages. Re-summarizing every turn would reprocess the whole prompt.
  let candidate = buildChat(history, messages, systemCount, {
    memoryNote: note,
    ...(saved !== undefined && saved.splitIndex < messages.length
      ? { splitIndex: saved.splitIndex, summary: saved.summary }
      : {}),
  });

  const before = await measure(source, candidate);
  if (before.used <= before.limit * triggerPercent) {
    return candidate;
  }

  const splitIndex = await findSplitIndex(messages, systemCount, keepRecentTokens, message =>
    source.countTokens(renderTranscript([message])),
  );
  if (splitIndex < 0) {
    // Nothing can be cut without orphaning a tool exchange; let the backend deal with it.
    ctl.debug("Compaction skipped: no safe split point found.");
    return candidate;
  }

  // A content block rather than a status step: only these carry a colour, and compaction should read
  // as part of the conversation in the app's own blue, the way tool calls do. includeInContext keeps
  // it off the model's plate — it is a note to the user, not to the model.
  const announce = showStatus
    ? ctl.createContentBlock({
        includeInContext: false,
        style: { type: "customLabel", label: "Context", color: "blue" },
      })
    : undefined;
  announce?.appendText("Compacting context…");

  try {
    const toCompact = messages.slice(systemCount, splitIndex);
    const result = await summarize(source, toCompact, {
      signal: ctl.abortSignal,
      cache: createCache(vaultPath),
      chunkTokens: config.get("chunkTokens"),
      onProgress: text => announce?.replaceText(text),
    });

    let archivedTo: string | undefined;
    if (archiving) {
      archivedTo = (await archive(vaultPath, folder, result.summary, toCompact)).stateFile;
    }
    // Remembered either way: without this the next message would re-summarize from scratch.
    await saveRecord(vaultPath, key, {
      splitIndex,
      summary: result.summary,
      compactedAt: new Date().toISOString(),
    });

    candidate = buildChat(history, messages, systemCount, {
      memoryNote: note,
      splitIndex,
      summary: result.summary,
    });
    const after = await measure(source, candidate);
    const reused = result.chunksTotal - result.chunksSummarized;

    // Compacting is pointless if the result still does not fit, and saying "compacted" while the
    // prediction is about to die of overflow is worse than pointless. The tail is the only thing
    // that can still be too big — everything before it is now a summary.
    if (after.used >= after.limit) {
      announce?.setStyle({ type: "customLabel", label: "Context", color: "red" });
      announce?.replaceText(
        `Compacted to ${after.used} tokens, which still exceeds the ${after.limit} the model ` +
          'holds. The recent messages kept verbatim are too large on their own — lower "Recent ' +
          'context kept verbatim", or ask the tools for less output.',
      );
      return candidate;
    }

    announce?.replaceText(
      `Context compacted: ${before.used} → ${after.used} tokens ` +
        `(${toCompact.length} messages summarized` +
        (reused > 0 ? `, ${reused}/${result.chunksTotal} parts reused from cache` : "") +
        ")" +
        (archivedTo === undefined ? "" : ` · archived to ${archivedTo}`),
    );
    return candidate;
  } catch (error) {
    // A failed compaction must not cost the user their message. Sending the full history may
    // overflow, but that is the behaviour they had before this plugin existed.
    ctl.debug("Compaction failed:", error);
    announce?.setStyle({ type: "customLabel", label: "Context", color: "red" });
    announce?.replaceText("Could not compact the context — sending the full history instead.");
    return history;
  }
}

async function runPrediction(
  ctl: PredictionLoopHandlerController,
  source: TokenSource,
  chat: Chat,
  tools: Array<SessionTool>,
): Promise<void> {
  if (tools.length === 0) {
    const block = ctl.createContentBlock();
    if (isLLM(source)) {
      await block.pipeFrom(source.respond(chat, { signal: ctl.abortSignal }));
    } else {
      // A generator plugin as token source cannot be piped; fall back to appending the result.
      const result = await source.respond(chat, { signal: ctl.abortSignal });
      block.appendText(result.content);
    }
    return;
  }

  let content: PredictionProcessContentBlockController | undefined;
  let reasoning: PredictionProcessContentBlockController | undefined;
  const contentBlock = () => (content ??= ctl.createContentBlock());

  // callId must be a number, while a tool call's id is a model-supplied string ("call_a1b2c3").
  // Coercing it with Number() yields NaN and the SDK rejects the whole update, which aborts the
  // tool loop mid-flight. Hand out our own numbers instead, keyed by the string so that a request
  // and its result land on the same id and the UI can pair them.
  const callIds = new Map<string, number>();
  const callIdFor = (id: string | undefined): number => {
    const key = id ?? `anonymous-${callIds.size}`;
    let assigned = callIds.get(key);
    if (assigned === undefined) {
      assigned = callIds.size;
      callIds.set(key, assigned);
    }
    return assigned;
  };

  await source.act(chat, tools, {
    signal: ctl.abortSignal,
    onRoundStart: () => {
      content = undefined;
      reasoning = undefined;
    },
    onPredictionFragment: fragment => {
      switch (fragment.reasoningType) {
        case "reasoningStartTag":
          // Delimiters, not content: their text is LM Studio's internal separator. Appending them
          // is what leaks LM_STUDIO_INTERNAL_LSEP_… into the chat — the default loop consumes them
          // to open and close a collapsible thinking block instead.
          return;
        case "reasoningEndTag":
          reasoning?.setStyle({ type: "thinking", ended: true });
          return;
        case "reasoning":
          reasoning ??= ctl.createContentBlock({ style: { type: "thinking" } });
          reasoning.appendText(fragment.content);
          return;
        default:
          contentBlock().appendText(fragment.content);
      }
    },
    onMessage: message => {
      // Text already arrived through fragments; only the tool traffic still needs rendering.
      for (const request of message.getToolCallRequests()) {
        contentBlock().appendToolRequest({
          callId: callIdFor(request.id),
          toolCallRequestId: request.id,
          name: request.name,
          parameters: request.arguments ?? {},
        });
      }

      // Results belong to a tool-role block; the SDK refuses to take them on an assistant one.
      const results = message.getToolCallResults();
      if (results.length > 0) {
        const toolBlock = ctl.createContentBlock({ roleOverride: "tool" });
        for (const result of results) {
          toolBlock.appendToolResult({
            callId: callIdFor(result.toolCallId),
            toolCallRequestId: result.toolCallId,
            content: String(result.content),
          });
        }
      }
    },
  });
}