coding_eval.ts

import { parseArgs } from "util";
import { Database } from "bun:sqlite";
import { DISTILLATION_SYSTEM, distillationUser } from "../src/prompt";

const BASE_URL = "http://localhost:4096";
const MODEL = { providerID: "anthropic", modelID: "claude-sonnet-4-6" };
const POLL_INTERVAL = 2000;
const MAX_WAIT = 120000;

const { values } = parseArgs({
  args: Bun.argv.slice(2),
  options: {
    data: { type: "string", default: "eval/data/coding_memory_eval.json" },
    out: { type: "string", default: "eval/results/coding_eval.jsonl" },
    mode: { type: "string", default: "all" }, // "default", "lore", or "all"
    concurrency: { type: "string", default: "3" },
  },
});

const concurrency = parseInt(values.concurrency!, 10);
const targetMode = values.mode!;

type Question = {
  session_id: string;
  session_label: string;
  question: string;
  answer: string;
  question_type: string;
  message_index: number;
};

// --- DB access ---
const DB_PATH =
  process.env.NUUM_DB ??
  `${process.env.HOME}/.local/share/opencode-lore/lore.db`;

function getTemporalMessages(sessionID: string): Array<{
  role: string;
  content: string;
  tokens: number;
  created_at: number;
}> {
  const d = new Database(DB_PATH, { readonly: true });
  const rows = d
    .query(
      "SELECT role, content, tokens, created_at FROM temporal_messages WHERE session_id = ? ORDER BY created_at ASC",
    )
    .all(sessionID) as Array<{
    role: string;
    content: string;
    tokens: number;
    created_at: number;
  }>;
  d.close();
  return rows;
}

function getDistillations(
  sessionID: string,
): Array<{ observations: string; created_at: number }> {
  const d = new Database(DB_PATH, { readonly: true });
  const projectRow = d
    .query(
      "SELECT DISTINCT project_id FROM temporal_messages WHERE session_id = ? LIMIT 1",
    )
    .get(sessionID) as { project_id: string } | null;
  if (!projectRow) {
    d.close();
    return [];
  }
  const rows = d
    .query(
      "SELECT observations, created_at FROM distillations WHERE project_id = ? AND session_id = ? ORDER BY created_at ASC",
    )
    .all(projectRow.project_id, sessionID) as Array<{
    observations: string;
    created_at: number;
  }>;
  d.close();
  return rows;
}


// Purge temporal messages from eval sessions to prevent recall tool contamination.
// Eval sessions are small (≤5 messages, <1000 tokens total) and not in the test dataset.
// The token check guards against accidentally deleting small but meaningful real sessions.
function purgeEvalMessages(testSessionIDs: string[]) {
  const d = new Database(DB_PATH);
  const testSet = new Set(testSessionIDs);

  // Find small low-token sessions (≤5 messages AND <1000 total tokens) that aren't test sessions
  const sessions = d
    .query(
      `SELECT session_id, count(*) as c, SUM(tokens) as total_tokens
       FROM temporal_messages
       GROUP BY session_id
       HAVING c <= 5 AND total_tokens < 1000`,
    )
    .all() as Array<{ session_id: string; c: number; total_tokens: number }>;
  const toDelete = sessions.filter((s) => !testSet.has(s.session_id));
  if (!toDelete.length) {
    d.close();
    return;
  }

  // Delete content table rows in batches, then rebuild FTS index.
  // FTS5 content-sync tables don't support direct DELETE — must rebuild after content changes.
  const batch = 100;
  let deleted = 0;
  const totalMsgs = toDelete.reduce((s, x) => s + x.c, 0);
  for (let i = 0; i < toDelete.length; i += batch) {
    const chunk = toDelete.slice(i, i + batch).map((s) => s.session_id);
    const placeholders = chunk.map(() => '?').join(',');
    d.query(
      `DELETE FROM temporal_messages WHERE session_id IN (${placeholders})`,
    ).run(...chunk);
    deleted += chunk.length;
  }
  // Rebuild FTS index to reflect deleted rows
  d.query("INSERT INTO temporal_fts(temporal_fts) VALUES('rebuild')").run();
  d.close();
  console.log(`Purged ${deleted} eval sessions (${totalMsgs} messages) from temporal storage`);
}

// --- Eval root session (hidden from UI) ---
let evalRoot: string;

async function createEvalRoot(): Promise<string> {
  const res = await fetch(`${BASE_URL}/session`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      title: `coding-eval - ${targetMode} - ${new Date().toISOString()}`,
    }),
  }).then((r) => r.json() as Promise<{ id: string }>);
  return res.id;
}

async function createSession(): Promise<string> {
  const res = await fetch(`${BASE_URL}/session`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ parentID: evalRoot }),
  }).then((r) => r.json() as Promise<{ id: string }>);
  return res.id;
}

type TokenInfo = {
  input: number;
  output: number;
  reasoning: number;
  cache: { read: number; write: number };
};

type MessageInfo = {
  info: {
    id: string;
    role: string;
    time: { created: number; updated: number };
    cost?: number;
    tokens?: TokenInfo;
  };
  parts: Array<{ type: string; text?: string; tool?: string; state?: { status: string } }>;
};

type AggregatedTokens = {
  input: number;       // input + cache.read + cache.write (total model context)
  raw_input: number;   // input only (non-cached)
  cache_read: number;
  cache_write: number;
  output: number;
  reasoning: number;
  cost: number;
  cache_hit_rate: number; // cache.read / input (fraction served from cache)
  api_calls: number;      // number of promptAndWait calls accumulated
};

function emptyTokens(): AggregatedTokens {
  return {
    input: 0, raw_input: 0, cache_read: 0, cache_write: 0,
    output: 0, reasoning: 0, cost: 0, cache_hit_rate: 0, api_calls: 0,
  };
}

function aggregateTokens(msgs: MessageInfo[]): AggregatedTokens {
  let rawInput = 0, cacheRead = 0, cacheWrite = 0, output = 0, reasoning = 0, cost = 0;
  for (const msg of msgs) {
    if (msg.info.role !== "assistant" || !msg.info.tokens) continue;
    const t = msg.info.tokens;
    rawInput += t.input;
    cacheRead += t.cache.read;
    cacheWrite += t.cache.write;
    output += t.output;
    reasoning += t.reasoning;
    cost += msg.info.cost ?? 0;
  }
  const totalInput = rawInput + cacheRead + cacheWrite;
  return {
    input: totalInput,
    raw_input: rawInput,
    cache_read: cacheRead,
    cache_write: cacheWrite,
    output,
    reasoning,
    cost,
    cache_hit_rate: totalInput > 0 ? cacheRead / totalInput : 0,
    api_calls: 1,
  };
}

function addTokens(a: AggregatedTokens, b: AggregatedTokens): AggregatedTokens {
  const totalInput = a.input + b.input;
  const totalCacheRead = a.cache_read + b.cache_read;
  return {
    input: totalInput,
    raw_input: a.raw_input + b.raw_input,
    cache_read: totalCacheRead,
    cache_write: a.cache_write + b.cache_write,
    output: a.output + b.output,
    reasoning: a.reasoning + b.reasoning,
    cost: a.cost + b.cost,
    cache_hit_rate: totalInput > 0 ? totalCacheRead / totalInput : 0,
    api_calls: a.api_calls + b.api_calls,
  };
}

// Wait for the model to finish responding. Handles multi-turn tool use by waiting until
// the last assistant message has stabilized (no new messages and has a text part).
// Returns the response text AND aggregated token usage across all assistant messages.
async function promptAndWait(
  sessionID: string,
  text: string,
  options?: { system?: string; agent?: string },
): Promise<{ text: string; tokens: AggregatedTokens }> {
  const body: Record<string, unknown> = {
    parts: [{ type: "text", text: options?.system ? `${options.system}\n\n${text}` : text }],
    model: MODEL,
    agent: options?.agent ?? "lore-distill",
  };
  await fetch(`${BASE_URL}/session/${sessionID}/prompt_async`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });

  const deadline = Date.now() + MAX_WAIT;
  let stableCount = 0;
  let lastMsgCount = 0;
  let lastMsgs: MessageInfo[] = [];

  while (Date.now() < deadline) {
    await Bun.sleep(POLL_INTERVAL);
    const msgs = await fetch(`${BASE_URL}/session/${sessionID}/message`).then(
      (r) => r.json() as Promise<MessageInfo[]>,
    );

    // Check if the response has settled: assistant message with text, not followed by more activity
    const assistants = msgs.filter((m) => m.info.role === "assistant");
    if (assistants.length > 0) {
      const last = assistants[assistants.length - 1];
      const hasText = last.parts.some((p) => p.type === "text" && p.text?.trim());
      const hasPendingTool = last.parts.some(
        (p) => p.type === "tool" && p.state?.status !== "completed",
      );

      // If the last assistant has text and no pending tools, and message count is stable, we're done
      if (hasText && !hasPendingTool) {
        if (msgs.length === lastMsgCount) stableCount++;
        else stableCount = 0;
        // Wait for 2 stable polls to be sure tool results aren't still arriving
        if (stableCount >= 1) {
          const textPart = last.parts.find((p) => p.type === "text" && p.text?.trim());
          if (textPart?.text) {
            lastMsgs = msgs;
            return { text: textPart.text.trim(), tokens: aggregateTokens(lastMsgs) };
          }
        }
      }
    }
    lastMsgCount = msgs.length;
    lastMsgs = msgs;
  }
  return { text: "[TIMEOUT]", tokens: aggregateTokens(lastMsgs) };
}

// OpenCode's compaction prompt — same wording used in session/compaction.ts
const COMPACTION_PROMPT =
  "Provide a detailed prompt for continuing our conversation above. Focus on information that would be helpful for continuing the conversation, including what we did, what we're doing, which files we're working on, and what we're going to do next considering new session will not have access to our conversation.";

const COMPACTION_SYSTEM = `You are a helpful AI assistant tasked with summarizing conversations.

When asked to summarize, provide a detailed but concise summary of the conversation. 
Focus on information that would be helpful for continuing the conversation, including:
- What was done
- What is currently being worked on
- Which files are being modified
- What needs to be done next
- Key user requests, constraints, or preferences that should persist
- Important technical decisions and why they were made

Your summary should be comprehensive enough to provide context but concise enough to be quickly understood.`;

async function compactSession(
  msgs: Array<{ role: string; content: string; created_at: number }>,
): Promise<string> {
  // Chunk messages into segments that fit within context (~80K tokens, rough char/4 estimate).
  // Each chunk is compacted with the prior summary carried forward, mimicking iterative compaction.
  const CHUNK_TOKEN_LIMIT = 80_000;
  const segments: string[] = [];
  let current: string[] = [];
  let tokens = 0;
  for (const m of msgs) {
    const line = `[${m.role}]: ${m.content}`;
    const est = Math.ceil(line.length / 4);
    if (tokens + est > CHUNK_TOKEN_LIMIT && current.length > 0) {
      segments.push(current.join("\n\n"));
      current = [];
      tokens = 0;
    }
    current.push(line);
    tokens += est;
  }
  if (current.length) segments.push(current.join("\n\n"));

  let summary = "";
  for (let i = 0; i < segments.length; i++) {
    const prior = summary
      ? `Here is a summary of the conversation so far:\n\n${summary}\n\n---\n\nContinuation of the conversation:\n\n`
      : "";
    const prompt = `${prior}${segments[i]}\n\n${COMPACTION_PROMPT}`;
    const sid = await createSession();
    console.log(`    Compacting chunk ${i + 1}/${segments.length}...`);
    summary = (await promptAndWait(sid, prompt, { system: COMPACTION_SYSTEM })).text;
  }
  return summary;
}

function buildLore(distillations: Array<{ observations: string }>): string {
  if (!distillations.length)
    return "[No distilled observations available for this session]";
  return distillations
    .map((d, i) => `## Session segment ${i + 1}\n${d.observations}`)
    .join("\n\n");
}



async function distillOnDemand(
  msgs: Array<{ role: string; content: string; created_at: number }>,
): Promise<string> {
  const segments: string[] = [];
  let current: string[] = [];
  let tokens = 0;
  for (const m of msgs) {
    const est = Math.ceil(m.content.length / 4);
    if (tokens + est > 20000 && current.length > 0) {
      segments.push(current.join("\n\n"));
      current = [];
      tokens = 0;
    }
    const time = new Date(m.created_at);
    const hh = time.getHours().toString().padStart(2, "0");
    const mm = time.getMinutes().toString().padStart(2, "0");
    current.push(`[${m.role}] (${hh}:${mm}) ${m.content}`);
    tokens += est;
  }
  if (current.length) segments.push(current.join("\n\n"));

  const date = msgs[0]
    ? new Date(msgs[0].created_at).toLocaleDateString("en-US", {
        year: "numeric",
        month: "long",
        day: "numeric",
      })
    : "unknown date";

  let allObservations = "";
  for (const segment of segments) {
    const userMsg = distillationUser({
      date,
      messages: segment,
      priorObservations: allObservations || undefined,
    });

    const sid = await createSession();
    const { text: responseText } = await promptAndWait(sid, userMsg, { system: DISTILLATION_SYSTEM });
    const match = responseText.match(/<observations>([\s\S]*?)<\/observations>/i);
    const obs = match ? match[1].trim() : responseText.trim();
    allObservations += (allObservations ? "\n" : "") + obs;
  }
  return allObservations;
}

// --- QA system prompts ---
const QA_SYSTEM = `You are a helpful coding assistant answering questions about past coding sessions. Answer concisely based on the context provided. If the information is not present in the context, say "I don't know."`;

const QA_SYSTEM_WITH_RECALL = `You are a helpful coding assistant answering questions about past coding sessions.

You have two sources of information:
1. Distilled observations provided in the context below
2. A "recall" tool that searches raw message archives and long-term knowledge

IMPORTANT: If the distilled observations don't contain enough detail to answer the question confidently, USE THE RECALL TOOL to search for the specific information. Try different search queries if the first doesn't return useful results.

Answer concisely. If after checking both observations and recall you still can't find the answer, say "I don't know."`;

// --- Process one question ---
async function processQuestion(
  q: Question,
  mode: string,
  loreContext: string,
  msgs: Array<{ role: string; content: string; tokens: number; created_at: number }>
): Promise<{
  question: string;
  answer: string;
  hypothesis: string;
  mode: string;
  tokens: AggregatedTokens;
}> {
  const sid = await createSession();

  if (mode === "lore") {
    // Lore mode: distilled observations as context + recall tool available.
    // Use default agent (not lore-distill) so the recall tool is registered.
    const prompt = `Here are distilled observations from a past coding session:\n\n${loreContext}\n\nQuestion: ${q.question}\n\nAnswer concisely. If the observations don't have enough detail, use the recall tool to search for it.`;
    const { text: hypothesis, tokens } = await promptAndWait(sid, prompt, {
      system: QA_SYSTEM_WITH_RECALL,
    });
    return { question: q.question, answer: q.answer, hypothesis, mode, tokens };
  }

  // Default mode: static tail window of last ~80K tokens (simulating OpenCode's recency-biased context)
  const TAIL_BUDGET = 80_000;
  let tailTokens = 0;
  let cutoff = msgs.length;
  for (let i = msgs.length - 1; i >= 0; i--) {
    tailTokens += msgs[i].tokens;
    if (tailTokens > TAIL_BUDGET) {
      cutoff = i + 1;
      break;
    }
  }
  const tailContext = msgs.slice(cutoff).map((m) => `[${m.role}]: ${m.content}`).join("\n\n");
  const dropped = msgs.length - (msgs.length - cutoff);
  const prefix = dropped > 0 ? `[Note: ${dropped} earlier messages were compacted/lost from context]\n\n` : "";
  const prompt = `Here is context from a past coding session:\n\n${prefix}${tailContext}\n\nQuestion: ${q.question}\n\nAnswer concisely:`;
  const { text: hypothesis, tokens } = await promptAndWait(sid, prompt, {
    system: QA_SYSTEM,
    agent: "lore-distill",
  });
  return { question: q.question, answer: q.answer, hypothesis, mode, tokens };
}

// --- Judge ---
const JUDGE_SYSTEM = `You are evaluating whether a hypothesis correctly answers a question about a coding session. Compare the hypothesis against the reference answer. Say "yes" if the hypothesis contains the key information from the reference (it can have extra detail). Say "no" if critical information is missing or wrong. Respond with ONLY "yes" or "no".`;

async function judge(
  question: string,
  reference: string,
  hypothesis: string,
): Promise<boolean> {
  const prompt = `Question: ${question}\nReference answer: ${reference}\nHypothesis: ${hypothesis}\n\nDoes the hypothesis correctly answer the question?`;
  const sid = await createSession();
  const { text: response } = await promptAndWait(sid, prompt, {
    system: JUDGE_SYSTEM,
  });
  return response.toLowerCase().startsWith("yes");
}

// --- Concurrency pool ---
async function pool<T, R>(
  items: T[],
  fn: (item: T, idx: number) => Promise<R>,
  max: number,
): Promise<R[]> {
  const results: R[] = [];
  let idx = 0;
  async function worker() {
    while (idx < items.length) {
      const i = idx++;
      results[i] = await fn(items[i], i);
    }
  }
  await Promise.all(
    Array.from({ length: Math.min(max, items.length) }, () => worker()),
  );
  return results;
}

// --- Main ---
const questions = (await Bun.file(values.data!).json()) as Question[];
evalRoot = await createEvalRoot();

console.log(`Coding Memory Eval`);
console.log(`Mode: ${targetMode}`);
console.log(`Questions: ${questions.length}`);
console.log(`Concurrency: ${concurrency}`);
console.log(`Output: ${values.out}`);
console.log("");

// Pre-load all session data
const sessionCache = new Map<
  string,
  {
    msgs: Array<{
      role: string;
      content: string;
      tokens: number;
      created_at: number;
    }>;
    lore: string;
  }
>();

const sessionIDs = [...new Set(questions.map((q) => q.session_id))];
purgeEvalMessages(sessionIDs);
const needDefault = targetMode === "all" || targetMode === "default";
const needLore = targetMode === "all" || targetMode === "lore";
for (const sid of sessionIDs) {
  console.log(`Loading session ${sid.substring(0, 16)}...`);
  const msgs = getTemporalMessages(sid);
  console.log(
    `  ${msgs.length} messages, ${msgs.reduce((s, m) => s + m.tokens, 0)} tokens`,
  );

  let lore = "";
  if (needLore) {
    const distillations = getDistillations(sid);
    if (
      distillations.length > 0 &&
      distillations.some((d) => d.observations?.trim())
    ) {
      console.log(`  Using ${distillations.length} existing distillation(s)`);
      lore = buildLore(distillations);
    } else {
      console.log(`  No existing distillations — running on-demand observer...`);
      lore = await distillOnDemand(msgs);
    }
    console.log(`  Lore context: ${lore.length} chars`);
  }

  sessionCache.set(sid, { msgs, lore });
}

console.log("");

// Build work items
type WorkItem = { q: Question; mode: string };
const work: WorkItem[] = [];
const modes =
  targetMode === "all" ? ["default", "lore"] : [targetMode];
for (const q of questions) {
  for (const mode of modes) {
    work.push({ q, mode });
  }
}

console.log(
  `Running ${work.length} evaluations (${questions.length} questions × ${modes.length} modes)...`,
);
console.log("");

const startTime = Date.now();
let completed = 0;
const writer = Bun.file(values.out!).writer();

await pool(
  work,
  async ({ q, mode }) => {
    const session = sessionCache.get(q.session_id)!;
    const result = await processQuestion(q, mode, session.lore, session.msgs);
    const label = await judge(q.question, q.answer, result.hypothesis);
    const entry = {
      session_label: q.session_label,
      question_type: q.question_type,
      question: q.question,
      answer: q.answer,
      hypothesis: result.hypothesis,
      mode: result.mode,
      label,
      tokens: result.tokens,
    };
    writer.write(JSON.stringify(entry) + "\n");
    writer.flush();
    completed++;

    const elapsed = (Date.now() - startTime) / 1000;
    const icon = label ? "✓" : "✗";
    const t = result.tokens;
    const cacheStr = t.input > 0
      ? ` [${(t.cache_hit_rate * 100).toFixed(0)}% cache, ${(t.input / 1000).toFixed(1)}K in, $${t.cost.toFixed(4)}]`
      : "";
    console.log(
      `[${completed}/${work.length}] ${icon} ${mode.padEnd(7)} ${q.session_label.padEnd(12)} "${q.question.substring(0, 50)}"${cacheStr}`,
    );
    return entry;
  },
  concurrency,
);

writer.end();

// --- Summary ---
const results = (await Bun.file(values.out!).text())
  .trim()
  .split("\n")
  .flatMap((l) => { try { return [JSON.parse(l)]; } catch { return []; } });

console.log("\n=== Results ===");
for (const mode of modes) {
  const modeResults = results.filter((r: any) => r.mode === mode);
  const correct = modeResults.filter((r: any) => r.label).length;
  console.log(
    `${mode.padEnd(10)} ${correct}/${modeResults.length} (${((correct / modeResults.length) * 100).toFixed(1)}%)`,
  );
}

if (modes.length > 1) {
  console.log("\n--- By session ---");
  for (const label of [...new Set(results.map((r: any) => r.session_label))]) {
    console.log(`\n${label}:`);
    for (const mode of modes) {
      const subset = results.filter(
        (r: any) => r.mode === mode && r.session_label === label,
      );
      const correct = subset.filter((r: any) => r.label).length;
      console.log(
        `  ${mode.padEnd(10)} ${correct}/${subset.length} (${((correct / subset.length) * 100).toFixed(1)}%)`,
      );
    }
  }
}

// --- Token usage summary ---
console.log("\n=== Token Usage ===");
for (const mode of modes) {
  const modeResults = results.filter((r: any) => r.mode === mode && r.tokens);
  if (!modeResults.length) continue;

  const totals = modeResults.reduce(
    (acc: AggregatedTokens, r: any) => addTokens(acc, r.tokens as AggregatedTokens),
    emptyTokens(),
  );
  const n = modeResults.length;
  const correct = modeResults.filter((r: any) => r.label).length;

  console.log(`\n${mode}:`);
  console.log(`  Avg input/question : ${(totals.input / n / 1000).toFixed(1)}K tokens`);
  console.log(`  Avg output/question: ${(totals.output / n).toFixed(0)} tokens`);
  console.log(`  Cache hit rate     : ${(totals.cache_hit_rate * 100).toFixed(1)}%`);
  console.log(`    cache_read total : ${(totals.cache_read / 1000).toFixed(1)}K`);
  console.log(`    cache_write total: ${(totals.cache_write / 1000).toFixed(1)}K`);
  console.log(`    raw_input total  : ${(totals.raw_input / 1000).toFixed(1)}K`);
  console.log(`  Total cost         : $${totals.cost.toFixed(4)}`);
  console.log(`  Cost/correct answer: $${correct > 0 ? (totals.cost / correct).toFixed(4) : "N/A"}`);
  console.log(`  Total API calls    : ${totals.api_calls}`);
}

const elapsed = (Date.now() - startTime) / 1000;
console.log(`\nDone! ${completed} evaluations in ${elapsed.toFixed(1)}s`);