context-compactor

import { type ChatMessage } from "@lmstudio/sdk";
import { promises as fs } from "fs";
import * as path from "path";
import { renderTranscript } from "./compaction";

export interface CompactionRecord {
  splitIndex: number;
  summary: string;
  compactedAt: string;
}

/** Persisted so that restarting LM Studio does not force a re-summarization on the next message. */
type StateFile = Record<string, CompactionRecord>;

function statePath(vaultPath: string): string {
  return path.join(vaultPath, ".context-compactor-state.json");
}

export async function loadState(vaultPath: string): Promise<StateFile> {
  try {
    return JSON.parse(await fs.readFile(statePath(vaultPath), "utf8")) as StateFile;
  } catch {
    return {};
  }
}

export async function saveState(
  vaultPath: string,
  key: string,
  record: CompactionRecord,
): Promise<void> {
  const state = await loadState(vaultPath);
  state[key] = record;
  await fs.mkdir(vaultPath, { recursive: true });
  await fs.writeFile(statePath(vaultPath), JSON.stringify(state, null, 2), "utf8");
}

// Module-level, for the same reason as the summary cache: without it, an unconfigured archive folder
// means no compaction is ever remembered, and every single message re-summarizes the whole history.
const memoryState = new Map<string, CompactionRecord>();

/** Reads the last compaction for a conversation, from disk when archiving, from memory otherwise. */
export async function loadRecord(
  vaultPath: string,
  key: string,
): Promise<CompactionRecord | undefined> {
  if (vaultPath === "") {
    return memoryState.get(key);
  }
  return (await loadState(vaultPath))[key];
}

export async function saveRecord(
  vaultPath: string,
  key: string,
  record: CompactionRecord,
): Promise<void> {
  if (vaultPath === "") {
    memoryState.set(key, record);
    return;
  }
  await saveState(vaultPath, key, record);
}

function stamp(date: Date): string {
  const pad = (n: number) => String(n).padStart(2, "0");
  return (
    `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}` +
    `-${pad(date.getHours())}${pad(date.getMinutes())}${pad(date.getSeconds())}`
  );
}

export interface ArchiveResult {
  transcriptFile: string;
  stateFile: string;
}

/**
 * Writes both halves of the memory: the verbatim transcript of everything leaving the context, and
 * the consolidated state note. The transcript is what makes compaction non-destructive — nothing is
 * lost, it just stops being in the prompt.
 *
 * Each conversation gets its own folder. A new chat is not a continuation of an old one, so its
 * memory does not belong in a shared pile: mixing several projects into flat Logs/ and Projects/
 * directories makes the vault unreadable and invites the wrong state being picked up. Resuming an
 * earlier chat stays a deliberate act — the user points the model at the folder.
 */
export async function archive(
  vaultPath: string,
  folder: string,
  summary: string,
  compactedMessages: Array<ChatMessage>,
): Promise<ArchiveResult> {
  const now = new Date();
  const conversationDir = path.join(vaultPath, folder);
  await fs.mkdir(conversationDir, { recursive: true });

  const transcriptName = `transcript-${stamp(now)}`;
  const transcriptFile = path.join(conversationDir, `${transcriptName}.md`);
  const stateFile = path.join(conversationDir, "state.md");

  const transcript = [
    "---",
    `created: ${now.toISOString()}`,
    `conversation: ${folder}`,
    `messages: ${compactedMessages.length}`,
    "tags: [lmstudio, transcript]",
    "---",
    "",
    `# Compacted transcript — ${now.toLocaleString()}`,
    "",
    `These ${compactedMessages.length} messages were removed from the model's context and replaced`,
    "by the consolidated state in [[state]]. Nothing is lost: it is all below.",
    "",
    renderTranscript(compactedMessages),
  ].join("\n");

  const state = [
    "---",
    `updated: ${now.toISOString()}`,
    `conversation: ${folder}`,
    "tags: [lmstudio, state]",
    "---",
    "",
    "# Consolidated state",
    "",
    `Last compaction: ${now.toLocaleString()}`,
    `Latest transcript: [[${transcriptName}]]`,
    "",
    "To resume this work in a new chat, tell the model to read this note.",
    "",
    "---",
    "",
    summary,
  ].join("\n");

  await fs.writeFile(transcriptFile, transcript, "utf8");
  await fs.writeFile(stateFile, state, "utf8");

  return { transcriptFile, stateFile };
}