server/runtime/compact.ts

/**
 * Compaction module — EXACT parity with original claw-code Rust compact.rs.
 *
 * Implements:
 * - CompactionConfig with preserve_recent_messages and max_estimated_tokens
 * - estimate_session_tokens / estimate_message_tokens
 * - should_compact
 * - compact_session (produces CompactionResult)
 * - format_compact_summary with <analysis>/<summary> tag handling
 * - get_compact_continuation_message
 * - merge_compact_summaries (multi-level compaction)
 * - summarize_messages with tool names, recent user requests, pending work, key files, timeline
 * - Helper functions: extract_tag_block, strip_tag_block, collapse_blank_lines, etc.
 */

// ─── Constants (match original compact.rs) ─────────────────────────────

const COMPACT_CONTINUATION_PREAMBLE =
  "This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.\n\n";

const COMPACT_RECENT_MESSAGES_NOTE = "Recent messages are preserved verbatim.";

const COMPACT_DIRECT_RESUME_INSTRUCTION =
  "Continue the conversation from where it left off without asking the user any further questions. Resume directly — do not acknowledge the summary, do not recap what was happening, and do not preface with continuation text.";

// ─── Types ──────────────────────────────────────────────────────────────

export interface CompactionConfig {
  preserveRecentMessages: number;
  maxEstimatedTokens: number;
}

export const DEFAULT_COMPACTION_CONFIG: CompactionConfig = {
  preserveRecentMessages: 4,
  maxEstimatedTokens: 10_000,
};

export type MessageRole = "system" | "user" | "assistant" | "tool";

export interface ContentBlock {
  type: "text" | "tool_use" | "tool_result";
  text?: string;
  name?: string;
  input?: string;
  toolName?: string;
  output?: string;
  isError?: boolean;
}

export interface ConversationMessage {
  role: MessageRole;
  blocks: ContentBlock[];
  usage?: { inputTokens?: number; outputTokens?: number };
}

export interface Session {
  version: number;
  messages: ConversationMessage[];
}

export interface CompactionResult {
  summary: string;
  formattedSummary: string;
  compactedSession: Session;
  removedMessageCount: number;
}

// ─── Token Estimation ───────────────────────────────────────────────────

function estimateBlockTokens(block: ContentBlock): number {
  switch (block.type) {
    case "text":
      return Math.floor((block.text?.length || 0) / 4) + 1;
    case "tool_use":
      return Math.floor(((block.name?.length || 0) + (block.input?.length || 0)) / 4) + 1;
    case "tool_result":
      return Math.floor(((block.toolName?.length || 0) + (block.output?.length || 0)) / 4) + 1;
    default:
      return 1;
  }
}

export function estimateMessageTokens(message: ConversationMessage): number {
  return message.blocks.reduce((sum, block) => sum + estimateBlockTokens(block), 0);
}

export function estimateSessionTokens(session: Session): number {
  return session.messages.reduce((sum, msg) => sum + estimateMessageTokens(msg), 0);
}

// ─── Should Compact ─────────────────────────────────────────────────────

function compactedSummaryPrefixLen(session: Session): number {
  const first = session.messages[0];
  if (!first) return 0;
  return extractExistingCompactedSummary(first) !== null ? 1 : 0;
}

export function shouldCompact(session: Session, config: CompactionConfig = DEFAULT_COMPACTION_CONFIG): boolean {
  const start = compactedSummaryPrefixLen(session);
  const compactable = session.messages.slice(start);
  if (compactable.length <= config.preserveRecentMessages) return false;
  const totalTokens = compactable.reduce((sum, msg) => sum + estimateMessageTokens(msg), 0);
  return totalTokens >= config.maxEstimatedTokens;
}

// ─── Tag Helpers ────────────────────────────────────────────────────────

function extractTagBlock(content: string, tag: string): string | null {
  const startTag = `<${tag}>`;
  const endTag = `</${tag}>`;
  const startIdx = content.indexOf(startTag);
  if (startIdx === -1) return null;
  const contentStart = startIdx + startTag.length;
  const endIdx = content.indexOf(endTag, contentStart);
  if (endIdx === -1) return null;
  return content.substring(contentStart, endIdx);
}

function stripTagBlock(content: string, tag: string): string {
  const startTag = `<${tag}>`;
  const endTag = `</${tag}>`;
  const startIdx = content.indexOf(startTag);
  if (startIdx === -1) return content;
  const endIdx = content.indexOf(endTag);
  if (endIdx === -1) return content;
  return content.substring(0, startIdx) + content.substring(endIdx + endTag.length);
}

function collapseBlankLines(content: string): string {
  const lines = content.split("\n");
  const result: string[] = [];
  let lastBlank = false;
  for (const line of lines) {
    const isBlank = line.trim() === "";
    if (isBlank && lastBlank) continue;
    result.push(line);
    lastBlank = isBlank;
  }
  return result.join("\n");
}

// ─── Format Summary ────────────────────────────────────────────────────

export function formatCompactSummary(summary: string): string {
  const withoutAnalysis = stripTagBlock(summary, "analysis");
  const summaryContent = extractTagBlock(withoutAnalysis, "summary");
  let formatted: string;
  if (summaryContent !== null) {
    formatted = withoutAnalysis.replace(
      `<summary>${summaryContent}</summary>`,
      `Summary:\n${summaryContent.trim()}`
    );
  } else {
    formatted = withoutAnalysis;
  }
  return collapseBlankLines(formatted).trim();
}

// ─── Continuation Message ──────────────────────────────────────────────

export function getCompactContinuationMessage(
  summary: string,
  suppressFollowUpQuestions: boolean,
  recentMessagesPreserved: boolean
): string {
  let base = `${COMPACT_CONTINUATION_PREAMBLE}${formatCompactSummary(summary)}`;
  if (recentMessagesPreserved) {
    base += `\n\n${COMPACT_RECENT_MESSAGES_NOTE}`;
  }
  if (suppressFollowUpQuestions) {
    base += `\n${COMPACT_DIRECT_RESUME_INSTRUCTION}`;
  }
  return base;
}

// ─── Summarize Messages ────────────────────────────────────────────────

function truncateSummary(content: string, maxChars: number): string {
  if (content.length <= maxChars) return content;
  return content.substring(0, maxChars) + "…";
}

function firstTextBlock(message: ConversationMessage): string | null {
  for (const block of message.blocks) {
    if (block.type === "text" && block.text && block.text.trim()) {
      return block.text;
    }
  }
  return null;
}

function summarizeBlock(block: ContentBlock): string {
  let raw: string;
  switch (block.type) {
    case "text":
      raw = block.text || "";
      break;
    case "tool_use":
      raw = `tool_use ${block.name || "?"}(${block.input || ""})`;
      break;
    case "tool_result":
      raw = `tool_result ${block.toolName || "?"}: ${block.isError ? "error " : ""}${block.output || ""}`;
      break;
    default:
      raw = "";
  }
  return truncateSummary(raw, 160);
}

function collectRecentRoleSummaries(
  messages: ConversationMessage[],
  role: MessageRole,
  limit: number
): string[] {
  const filtered = messages.filter((m) => m.role === role);
  const reversed = [...filtered].reverse();
  const summaries: string[] = [];
  for (const msg of reversed) {
    if (summaries.length >= limit) break;
    const text = firstTextBlock(msg);
    if (text) summaries.push(truncateSummary(text, 160));
  }
  return summaries.reverse();
}

function hasInterestingExtension(candidate: string): boolean {
  const ext = candidate.split(".").pop()?.toLowerCase() || "";
  return ["rs", "ts", "tsx", "js", "json", "md", "py", "go", "java", "css", "html"].includes(ext);
}

function extractFileCandidates(content: string): string[] {
  return content
    .split(/\s+/)
    .filter((token) => {
      const candidate = token.replace(/^[,.:;)"'`]+|[,.:;)"'`]+$/g, "");
      return candidate.includes("/") && hasInterestingExtension(candidate);
    })
    .map((token) => token.replace(/^[,.:;)"'`]+|[,.:;)"'`]+$/g, ""));
}

function collectKeyFiles(messages: ConversationMessage[]): string[] {
  const files = new Set<string>();
  for (const msg of messages) {
    for (const block of msg.blocks) {
      let content = "";
      if (block.type === "text") content = block.text || "";
      else if (block.type === "tool_use") content = block.input || "";
      else if (block.type === "tool_result") content = block.output || "";
      for (const file of extractFileCandidates(content)) {
        files.add(file);
      }
    }
  }
  return Array.from(files).sort().slice(0, 8);
}

function inferPendingWork(messages: ConversationMessage[]): string[] {
  const results: string[] = [];
  const reversed = [...messages].reverse();
  for (const msg of reversed) {
    if (results.length >= 3) break;
    const text = firstTextBlock(msg);
    if (text) {
      const lowered = text.toLowerCase();
      if (
        lowered.includes("todo") ||
        lowered.includes("next") ||
        lowered.includes("pending") ||
        lowered.includes("follow up") ||
        lowered.includes("remaining")
      ) {
        results.push(truncateSummary(text, 160));
      }
    }
  }
  return results.reverse();
}

function inferCurrentWork(messages: ConversationMessage[]): string | null {
  const reversed = [...messages].reverse();
  for (const msg of reversed) {
    const text = firstTextBlock(msg);
    if (text && text.trim()) return truncateSummary(text, 200);
  }
  return null;
}

function summarizeMessages(messages: ConversationMessage[]): string {
  const userMessages = messages.filter((m) => m.role === "user").length;
  const assistantMessages = messages.filter((m) => m.role === "assistant").length;
  const toolMessages = messages.filter((m) => m.role === "tool").length;

  const toolNames = new Set<string>();
  for (const msg of messages) {
    for (const block of msg.blocks) {
      if (block.type === "tool_use" && block.name) toolNames.add(block.name);
      if (block.type === "tool_result" && block.toolName) toolNames.add(block.toolName);
    }
  }
  const sortedToolNames = Array.from(toolNames).sort();

  const lines: string[] = [
    "<summary>",
    "Conversation summary:",
    `- Scope: ${messages.length} earlier messages compacted (user=${userMessages}, assistant=${assistantMessages}, tool=${toolMessages}).`,
  ];

  if (sortedToolNames.length > 0) {
    lines.push(`- Tools mentioned: ${sortedToolNames.join(", ")}.`);
  }

  const recentUserRequests = collectRecentRoleSummaries(messages, "user", 3);
  if (recentUserRequests.length > 0) {
    lines.push("- Recent user requests:");
    for (const req of recentUserRequests) {
      lines.push(`  - ${req}`);
    }
  }

  const pendingWork = inferPendingWork(messages);
  if (pendingWork.length > 0) {
    lines.push("- Pending work:");
    for (const item of pendingWork) {
      lines.push(`  - ${item}`);
    }
  }

  const keyFiles = collectKeyFiles(messages);
  if (keyFiles.length > 0) {
    lines.push(`- Key files referenced: ${keyFiles.join(", ")}.`);
  }

  const currentWork = inferCurrentWork(messages);
  if (currentWork) {
    lines.push(`- Current work: ${currentWork}`);
  }

  lines.push("- Key timeline:");
  for (const msg of messages) {
    const content = msg.blocks.map(summarizeBlock).join(" | ");
    lines.push(`  - ${msg.role}: ${content}`);
  }

  lines.push("</summary>");
  return lines.join("\n");
}

// ─── Merge Summaries ───────────────────────────────────────────────────

function extractSummaryHighlights(summary: string): string[] {
  const lines: string[] = [];
  let inTimeline = false;
  for (const line of formatCompactSummary(summary).split("\n")) {
    const trimmed = line.trimEnd();
    if (!trimmed || trimmed === "Summary:" || trimmed === "Conversation summary:") continue;
    if (trimmed === "- Key timeline:") {
      inTimeline = true;
      continue;
    }
    if (inTimeline) continue;
    lines.push(trimmed);
  }
  return lines;
}

function extractSummaryTimeline(summary: string): string[] {
  const lines: string[] = [];
  let inTimeline = false;
  for (const line of formatCompactSummary(summary).split("\n")) {
    const trimmed = line.trimEnd();
    if (trimmed === "- Key timeline:") {
      inTimeline = true;
      continue;
    }
    if (!inTimeline) continue;
    if (!trimmed) break;
    lines.push(trimmed);
  }
  return lines;
}

function mergeCompactSummaries(existingSummary: string | null, newSummary: string): string {
  if (!existingSummary) return newSummary;

  const previousHighlights = extractSummaryHighlights(existingSummary);
  const newFormattedSummary = formatCompactSummary(newSummary);
  const newHighlights = extractSummaryHighlights(newFormattedSummary);
  const newTimeline = extractSummaryTimeline(newFormattedSummary);

  const lines: string[] = ["<summary>", "Conversation summary:"];

  if (previousHighlights.length > 0) {
    lines.push("- Previously compacted context:");
    for (const line of previousHighlights) {
      lines.push(`  ${line}`);
    }
  }

  if (newHighlights.length > 0) {
    lines.push("- Newly compacted context:");
    for (const line of newHighlights) {
      lines.push(`  ${line}`);
    }
  }

  if (newTimeline.length > 0) {
    lines.push("- Key timeline:");
    for (const line of newTimeline) {
      lines.push(`  ${line}`);
    }
  }

  lines.push("</summary>");
  return lines.join("\n");
}

// ─── Extract Existing Summary ──────────────────────────────────────────

function extractExistingCompactedSummary(message: ConversationMessage): string | null {
  if (message.role !== "system") return null;
  const text = firstTextBlock(message);
  if (!text) return null;
  if (!text.startsWith(COMPACT_CONTINUATION_PREAMBLE)) return null;

  let summary = text.substring(COMPACT_CONTINUATION_PREAMBLE.length);

  const recentNoteIdx = summary.indexOf(`\n\n${COMPACT_RECENT_MESSAGES_NOTE}`);
  if (recentNoteIdx !== -1) {
    summary = summary.substring(0, recentNoteIdx);
  }

  const resumeIdx = summary.indexOf(`\n${COMPACT_DIRECT_RESUME_INSTRUCTION}`);
  if (resumeIdx !== -1) {
    summary = summary.substring(0, resumeIdx);
  }

  return summary.trim();
}

// ─── Main Compact Function ─────────────────────────────────────────────

/**
 * LLM-based summarization: calls the configured LLM to produce a real summary
 * instead of heuristic keyword extraction.
 */
export async function compactSessionWithLLM(
  session: Session,
  config: CompactionConfig = DEFAULT_COMPACTION_CONFIG,
  llmFetch?: (messages: Array<{role: string; content: string}>) => Promise<string>
): Promise<CompactionResult> {
  const heuristicResult = compactSessionSync(session, config);
  if (!heuristicResult.removedMessageCount || !llmFetch) return heuristicResult;

  // Build a prompt for the LLM to summarize the removed messages
  const removedText = heuristicResult.summary;
  try {
    const llmSummary = await llmFetch([
      {
        role: "system",
        content: "You are a conversation summarizer. Produce a concise but complete summary of the conversation below. Include: 1) What the user asked for, 2) What actions were taken, 3) What files were created/modified, 4) Current state and pending work. Be specific with file names, function names, and technical details. Output ONLY the summary, no preamble."
      },
      {
        role: "user",
        content: `Summarize this conversation segment:\n\n${removedText}`
      }
    ]);
    if (llmSummary && llmSummary.length > 50) {
      // Replace heuristic summary with LLM summary
      const formattedSummary = formatCompactSummary(llmSummary);
      const continuation = getCompactContinuationMessage(llmSummary, true, heuristicResult.compactedSession.messages.length > 1);
      heuristicResult.summary = llmSummary;
      heuristicResult.formattedSummary = formattedSummary;
      heuristicResult.compactedSession.messages[0] = {
        role: "system",
        blocks: [{ type: "text", text: continuation }],
      };
    }
  } catch (e) {
    console.error("[compact] LLM summarization failed, using heuristic:", e);
  }
  return heuristicResult;
}

/**
 * Synchronous heuristic compaction (original behavior, used as fallback)
 */
export function compactSession(
  session: Session,
  config: CompactionConfig = DEFAULT_COMPACTION_CONFIG
): CompactionResult {
  return compactSessionSync_impl(session, config);
}

export function compactSessionSync(
  session: Session,
  config: CompactionConfig = DEFAULT_COMPACTION_CONFIG
): CompactionResult {
  return compactSessionSync_impl(session, config);
}

function compactSessionSync_impl(
  session: Session,
  config: CompactionConfig = DEFAULT_COMPACTION_CONFIG
): CompactionResult {
  if (!shouldCompact(session, config)) {
    return {
      summary: "",
      formattedSummary: "",
      compactedSession: { ...session, messages: [...session.messages] },
      removedMessageCount: 0,
    };
  }

  const existingSummary = session.messages[0]
    ? extractExistingCompactedSummary(session.messages[0])
    : null;

  const compactedPrefixLen = existingSummary !== null ? 1 : 0;
  const keepFrom = Math.max(compactedPrefixLen, session.messages.length - config.preserveRecentMessages);
  const removed = session.messages.slice(compactedPrefixLen, keepFrom);
  const preserved = session.messages.slice(keepFrom);

  const summary = mergeCompactSummaries(existingSummary, summarizeMessages(removed));
  const formattedSummary = formatCompactSummary(summary);
  const continuation = getCompactContinuationMessage(summary, true, preserved.length > 0);

  const compactedMessages: ConversationMessage[] = [
    {
      role: "system",
      blocks: [{ type: "text", text: continuation }],
    },
    ...preserved,
  ];

  return {
    summary,
    formattedSummary,
    compactedSession: { version: session.version, messages: compactedMessages },
    removedMessageCount: removed.length,
  };
}

// ─── Convert DB messages to Session format ─────────────────────────────

export function dbMessagesToSession(
  dbMessages: Array<{ role: string; content: string; toolName?: string | null; toolCallId?: string | null }>
): Session {
  const messages: ConversationMessage[] = dbMessages.map((msg) => {
    const blocks: ContentBlock[] = [];
    if (msg.role === "tool") {
      blocks.push({
        type: "tool_result",
        toolName: msg.toolName || undefined,
        output: msg.content,
        isError: false,
      });
    } else {
      // Try to parse tool_use from content
      try {
        const parsed = JSON.parse(msg.content);
        if (parsed.tool_calls && Array.isArray(parsed.tool_calls)) {
          for (const tc of parsed.tool_calls) {
            blocks.push({
              type: "tool_use",
              name: tc.function?.name || tc.name,
              input: typeof tc.function?.arguments === "string"
                ? tc.function.arguments
                : JSON.stringify(tc.function?.arguments || tc.arguments || {}),
            });
          }
        } else {
          blocks.push({ type: "text", text: msg.content });
        }
      } catch {
        blocks.push({ type: "text", text: msg.content });
      }
    }
    return {
      role: msg.role as MessageRole,
      blocks,
    };
  });

  return { version: 1, messages };
}