oc/plugins/memory/recall.ts

// Retrieval: tokenization, scoring, and the block that gets injected.
//
// Tokenization note — this is the part that had to be measured, not guessed.
//
// FTS5's two built-in tokenizers both fail on the memory corpus this plugin
// exists to hold. Measured on SQLite 3.53.4 with the string
// "部署脚本用 pnpm 而不是 npm，因为 lockfile 不兼容":
//
//   unicode61 (FTS5 default)  query "部署" -> 0 hits
//                            query "部署脚本" -> 0 hits
//       It emits 7 terms for the whole sentence, treating 部署脚本用 as ONE
//       token, so no substring can ever match.
//
//   trigram                  query "部署" -> 0 hits   (needs >= 3 chars)
//                            query "部署脚" -> 1 hit
//       Its vocabulary is sliding 3-char windows (不兼 / 而不 / m，因), so every
//       two-character Chinese word — 部署, 配置, 依赖, 接口 — is unreachable.
//
// So tokenization happens here, not in SQL: CJK becomes character bigrams and
// latin becomes words. Two-character Chinese then matches exactly, and the
// extra noise bigrams introduce is harmless because rarer terms score lower.

import fs from "node:fs"
import path from "node:path"

const CJK =
  /[\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff\u3040-\u30ff\uac00-\ud7af]/
const LATIN_WORD = /[A-Za-z0-9][A-Za-z0-9_+.#-]*/g

/** Stopwords that would otherwise dominate a scoring pass. */
const STOPWORDS = new Set([
  "the", "a", "an", "and", "or", "of", "to", "in", "is", "are", "was", "were",
  "be", "for", "on", "at", "by", "with", "as", "it", "this", "that", "from",
  "我", "的", "了", "是", "在", "和", "有", "就", "不", "也", "你",
])

export function tokenize(text: string): string[] {
  if (!text) return []
  const tokens: string[] = []

  // Latin / digits: whole words, lowercased.
  for (const match of text.toLowerCase().matchAll(LATIN_WORD)) {
    const word = match[0]
    if (word.length < 2 || STOPWORDS.has(word)) continue
    tokens.push(word)
  }

  // CJK: unigrams plus bigrams. Bigrams carry the matching power; unigrams
  // rescue one-character queries that have no bigram at all.
  const runs = text.match(new RegExp(CJK.source, "g")) ?? []
  for (const run of runs) {
    const chars = [...run]
    for (let i = 0; i < chars.length - 1; i++) {
      const bigram = chars[i] + chars[i + 1]
      if (!STOPWORDS.has(bigram)) tokens.push(bigram)
    }
    if (chars.length === 1 && !STOPWORDS.has(chars[0])) tokens.push(chars[0])
  }

  return tokens
}

export interface ScoredFile {
  file: import("./store.ts").MemoryFile
  score: number
  reasons: string[]
}

/**
 * Score memories against a query.
 *
 * Ranking is scope-first on purpose: a project-specific fact about how *this*
 * repo deploys should beat a global preference when both match, because the
 * narrower claim is the more useful one. Ties inside one root fall back to
 * decay and hit count.
 */
export function scoreFiles(
  files: readonly import("./store.ts").MemoryFile[],
  query: string,
  states: Record<string, import("./store.ts").FileState>,
  options: { now?: number; halfLifeDays?: number; pinnedBoost?: boolean } = {},
): ScoredFile[] {
  const now = options.now ?? Date.now()
  const halfLife = (options.halfLifeDays ?? 14) * 86_400_000
  const queryTokens = tokenize(query)
  if (queryTokens.length === 0) return []
  if (files.length === 0) return []

  const querySet = new Map<string, number>()
  for (const token of queryTokens) querySet.set(token, (querySet.get(token) ?? 0) + 1)

  // Document frequency for idf: how many files contain each term.
  // The previous version passed a per-file match count into a total-files
  // denominator, which gave every file the same idf and reduced ranking to
  // raw term frequency.
  const docFreq = new Map<string, number>()
  const docTokenList: Set<string>[] = []
  for (const file of files) {
    const haystack = [file.title, file.summary, file.highlight, file.loadWhen.join(" "), file.body].join("\n")
    const docTokens = new Set(tokenize(haystack))
    docTokenList.push(docTokens)
    for (const token of docTokens) docFreq.set(token, (docFreq.get(token) ?? 0) + 1)
  }

  const results: ScoredFile[] = []

  for (let i = 0; i < files.length; i++) {
    const file = files[i]
    // A superseded memory stays on disk for audit but must never be injected.
    if (file.fm.superseded_by) continue

    const state = states[file.rel]
    const docTokens = docTokenList[i]

    let overlap = 0
    const matched: string[] = []
    for (const [token, weight] of querySet) {
      if (!docTokens.has(token)) continue
      // Diminishing returns: a term repeated in the query should not dominate.
      overlap += weight * (1 / (1 + Math.log2(weight)))
      if (matched.length < 4) matched.push(token)
    }
    if (overlap <= 0) continue

    // Inverse document frequency: a term appearing in every file is useless.
    let idfSum = 0
    let idfCount = 0
    for (const token of querySet.keys()) {
      if (!docTokens.has(token)) continue
      const df = docFreq.get(token) ?? 1
      idfSum += Math.log(1 + files.length / df)
      idfCount++
    }
    if (idfCount === 0) continue
    const idf = idfSum / idfCount

    let score = overlap * idf

    // Scope prior.
    score *= file.root.rank === 20 ? 1 : 0.5

    // Recency decay, unless pinned or a hard rule.
    const hard = file.pinned || String(file.fm.kind ?? "").toLowerCase() === "rule" || file.category === "rules"
    if (!hard && halfLife > 0) {
      const age = Math.max(0, now - (state?.lastHitAt ?? file.updated))
      score *= Math.pow(0.5, age / halfLife)
    }

    // Reinforcement: memories that keep proving useful float up.
    score *= 1 + Math.log1p(state?.hitCount ?? 0) * 0.25

    if (options.pinnedBoost !== false && hard) score *= 1.15

    const reasons = [...matched]
    if (hard) reasons.push("铁规则")
    if ((state?.hitCount ?? 0) > 0) reasons.push(`命中 ${state?.hitCount} 次`)

    results.push({ file, score, reasons })
  }

  return results.sort((a, b) => b.score - a.score)
}

export interface InjectionOptions {
  /** Hard cap on the injected block, in characters. */
  maxChars?: number
  /** Max number of memory bodies to inline. */
  maxEntries?: number
  /**
   * Half-life in days for the recency decay. 0 disables decay.
   * Declared here because the caller passes it: it was silently dropped when
   * this option was missing, so injection ranked with the default 14-day
   * half-life regardless of configuration.
   */
  halfLifeDays?: number
  now?: number
}

export interface InjectionInput {
  roots: readonly import("./store.ts").MemoryRoot[]
  filesByRoot: Map<string, import("./store.ts").MemoryFile[]>
  statesByRoot: Map<string, Record<string, import("./store.ts").FileState>>
  query: string
  options?: InjectionOptions
}

/**
 * Build the text injected into the agent loop.
 *
 * Two tiers on purpose:
 *   - outlines are cheap and always accurate, so they always ride along
 *   - full bodies are expensive, so only the best few ride along
 * A `memory_recall` call reaches the rest.
 */
export function buildInjection(input: InjectionInput): string {
  const { roots, filesByRoot, statesByRoot, query } = input
  const maxChars = input.options?.maxChars ?? 600
  const maxEntries = input.options?.maxEntries ?? 0
  const now = input.options?.now ?? Date.now()

  const blocks: string[] = []

  // Tier 1 — hard rules. Resident regardless of query: these are the user's
  // non-negotiables and must never be outranked by relevance.
  // Body first line only: the title usually restates it, and this block is
  // re-sent on every model request.
  const ruleLines: string[] = []
  for (const root of roots) {
    for (const file of filesByRoot.get(root.key) ?? []) {
      const isRule = file.category === "rules" || String(file.fm.kind ?? "").toLowerCase() === "rule"
      if (!isRule || file.fm.superseded_by) continue
      const firstLine = file.body.split(/\r?\n/).find((l) => l.trim())?.trim() ?? ""
      if (firstLine) ruleLines.push(`- [${root.label}] ${firstLine}`)
    }
  }

  // Tier 2 — the outline: what exists, and when each part is worth loading.
  // This is the ROOT.md idea: one small file answers "what do I know and when
  // do I need it", instead of dumping every memory into every request.
  //
  // Goals only. The slug list was the single most expensive part of the block
  // and the model does not need it: `memory_recall` discovers files by query.
  const outline: string[] = []
  for (const root of roots) {
    const goals = readGoal(root.dir)
    if (goals) outline.push(`目标(${root.label}): ${goals}`)
  }

  // Tier 3 — full bodies for the top matches.
  const bodyLines: string[] = []
  let budget = maxChars - ruleLines.join("\n").length - outline.join("\n").length
  const all: ScoredFile[] = []
  for (const root of roots) {
    const scored = scoreFiles(filesByRoot.get(root.key) ?? [], query, statesByRoot.get(root.key) ?? {}, { now })
    for (const entry of scored) {
      // Project wins over global on the same file path.
      if (all.some((a) => a.file.rel === entry.file.rel && a.file.root.rank < entry.file.root.rank)) continue
      all.push(entry)
    }
  }
  for (const entry of all) {
    if (bodyLines.length >= maxEntries) break
    if (entry.score <= 0) continue
    const body = entry.file.body.slice(0, 400)
    const line = `\n### ${entry.file.root.label}/${entry.file.rel}\n${body}`
    if (budget - line.length < 0) break
    budget -= line.length
    bodyLines.push(line)
  }

  const sections: string[] = []
  if (ruleLines.length) {
    sections.push(`## 铁规则（必须遵守）\n${ruleLines.join("\n")}`)
  }
  if (outline.length) {
    sections.push(`## 记忆大纲（按需深挖，用 memory_recall 取全文）\n${outline.join("\n")}`)
  }
  if (bodyLines.length) {
    sections.push(`## 与当前任务相关的记忆\n${bodyLines.join("\n")}`)
  }
  if (sections.length === 0) return ""

  const header = "# 记忆上下文（来自 memory 插件；如需更多记忆用 memory_recall 工具）"
  let text = `${header}\n\n${sections.join("\n\n")}`
  if (text.length > maxChars) {
    // Truncate from the least important tier: keep rules, drop the tail of the
    // outline, then the bodies.
    const keep = sections.slice(0, Math.max(1, sections.length - 1)).join("\n\n")
    text = `${header}\n\n${keep}`.slice(0, maxChars)
  }
  return text
}

function groupCategories(files: readonly import("./store.ts").MemoryFile[]) {
  const map = new Map<string, import("./store.ts").MemoryFile[]>()
  for (const file of files) {
    const list = map.get(file.category) ?? []
    list.push(file)
    map.set(file.category, list)
  }
  return [...map.entries()].sort((a, b) => a[0].localeCompare(b[0]))
}

/** Read the hand-written "当前目标" block from ROOT.md. */
export function readGoal(rootDir: string): string | undefined {
  try {
    const text = fs.readFileSync(path.join(rootDir, "ROOT.md"), "utf8")
    // Match until the next section header or auto marker. The previous regex
    // stopped at any `##` or `<!--`, which truncated content containing
    // sub-headings or HTML comments inside the goals section.
    const m = /^##\s*当前目标\s*\n([\s\S]*?)(?=\n##\s|\n<!-- auto:begin|$)/m.exec(text)
    const body = m?.[1]
      ?.replace(/<!--(?!auto:)[\s\S]*?-->/g, "")
      .split(/\r?\n/)
      .map((l) => l.trim())
      .filter((l) => l && !l.startsWith("<!--"))
      .join(" ")
      .trim()
    return body || undefined
  } catch {
    // A root nobody has written to yet has no goals. This runs on every
    // injection, so a missing ROOT.md is a quiet outcome, not a failure.
    return undefined
  }
}