typeclaw 0.36.7 → 0.37.0

package/src/bundled-plugins/memory/vector/doctor.ts

@@ -0,0 +1,124 @@
+import { existsSync } from 'node:fs'
+import { rm } from 'node:fs/promises'
+import { join } from 'node:path'
+
+import type { PluginCheckResult } from '@/plugin'
+
+import { inspectVectorIndex, type VectorIndexProblem } from './inspect'
+import { collectPassages, findMissingPassages } from './passages'
+import { VectorStore } from './store'
+
+export const VECTOR_INDEX_REL_PATH = join('memory', '.vectors', 'index.db')
+
+export async function runVectorIndexDoctor(agentDir: string): Promise<PluginCheckResult> {
+  const dbPath = join(agentDir, VECTOR_INDEX_REL_PATH)
+
+  if (!existsSync(dbPath)) {
+    return {
+      status: 'warning',
+      message: 'vector memory is enabled but the index DB is missing; it rebuilds on the next startup',
+    }
+  }
+
+  const finding = inspectVectorIndex(dbPath)
+
+  if (finding.kind === 'unreadable' || finding.kind === 'corrupt' || finding.kind === 'schema-missing') {
+    return corruptionResult(dbPath, finding)
+  }
+
+  const passages = await collectPassages(agentDir)
+  const wantedIds = new Set(passages.map((passage) => passage.id))
+  const orphans = finding.rowIds.filter((id) => !wantedIds.has(id))
+  const backfillCount = countBackfill(dbPath, passages)
+
+  return summarize(dbPath, {
+    rowCount: finding.rowCount,
+    orphans,
+    modelMismatch: finding.modelMismatch,
+    malformed: finding.malformed,
+    backfillCount,
+    indexedCount: passages.length - backfillCount,
+    wantedCount: passages.length,
+  })
+}
+
+function countBackfill(dbPath: string, passages: Awaited<ReturnType<typeof collectPassages>>): number {
+  const store = VectorStore.open(dbPath)
+  try {
+    return findMissingPassages(store, passages).length
+  } finally {
+    store.close()
+  }
+}
+
+type Summary = {
+  rowCount: number
+  orphans: string[]
+  modelMismatch: string[]
+  malformed: string[]
+  backfillCount: number
+  indexedCount: number
+  wantedCount: number
+}
+
+function summarize(dbPath: string, s: Summary): PluginCheckResult {
+  // Dedupe: a row can be both orphaned and malformed/variant, so the union by
+  // id keeps the count and the deletion list honest.
+  const repairable = [...new Set([...s.orphans, ...s.modelMismatch, ...s.malformed])]
+  const details: string[] = []
+  if (s.orphans.length > 0) details.push(`${s.orphans.length} orphaned row(s) for deleted topics/fragments`)
+  if (s.modelMismatch.length > 0) {
+    details.push(`${s.modelMismatch.length} row(s) from a different embedding model/dims`)
+  }
+  if (s.malformed.length > 0) details.push(`${s.malformed.length} row(s) with a malformed embedding blob`)
+  if (s.backfillCount > 0) details.push(`${s.backfillCount} memory passage(s) need (re)indexing`)
+
+  if (details.length === 0) {
+    return { status: 'ok', message: `vector index healthy: ${s.indexedCount}/${s.wantedCount} memory passages indexed` }
+  }
+
+  const result: PluginCheckResult = {
+    status: 'warning',
+    message: `vector index has ${details.length} issue(s); ${s.rowCount} row(s) stored`,
+    details,
+  }
+
+  if (repairable.length > 0) {
+    result.fix = {
+      description: `Delete ${repairable.length} orphaned/incompatible vector row(s); backfill happens on the next startup`,
+      apply: async () => {
+        const store = VectorStore.open(dbPath)
+        try {
+          store.deleteMany(repairable)
+        } finally {
+          store.close()
+        }
+        return {
+          summary: `pruned ${repairable.length} stale vector row(s) from ${VECTOR_INDEX_REL_PATH}`,
+          changedPaths: [],
+        }
+      },
+    }
+  }
+
+  return result
+}
+
+function corruptionResult(dbPath: string, finding: VectorIndexProblem): PluginCheckResult {
+  const details = finding.kind === 'corrupt' ? finding.detail : [finding.detail]
+  return {
+    status: 'error',
+    message:
+      finding.kind === 'schema-missing'
+        ? 'vector index DB has an invalid schema'
+        : 'vector index DB is unreadable or corrupted',
+    details,
+    fix: {
+      description: `Delete the corrupted ${VECTOR_INDEX_REL_PATH}; it rebuilds from memory on the next startup`,
+      apply: async () => {
+        await rm(dbPath, { force: true })
+        return { summary: `deleted corrupted ${VECTOR_INDEX_REL_PATH}`, changedPaths: [] }
+      },
+    },
+  }
+}

package/src/bundled-plugins/memory/vector/embedder.ts

@@ -0,0 +1,246 @@
+import { join } from 'node:path'
+
+// Type-only import: erased at runtime, so it does NOT evaluate
+// @huggingface/transformers (which eagerly `import sharp`s, crashing the
+// container at startup when sharp's linux binary is missing). The runtime
+// values are pulled lazily via `loadTransformers()` below.
+import type { env as TransformersEnvValue, pipeline as TransformersPipeline } from '@huggingface/transformers'
+
+import {
+  assertModelCacheCompatible,
+  EMBEDDING_DIMS,
+  EMBEDDING_MODEL_DTYPE,
+  EMBEDDING_MODEL_ID,
+  EMBEDDING_MODEL_NAME,
+} from '@/models/embedding-model'
+import { getResolvedTransformersVersion } from '@/models/transformers-version'
+
+import { homeRoot } from '../../../hostd/paths'
+import { type BoundedText, boundEmbeddableText, MAX_MODEL_TOKENS } from './truncation'
+
+// Re-exported for the vector subsystem's existing imports. The canonical
+// definitions live in @/models/embedding-model (shared host + container).
+export const MODEL_NAME = EMBEDDING_MODEL_NAME
+export const DIMS = EMBEDDING_DIMS
+export const MODEL_DTYPE = EMBEDDING_MODEL_DTYPE
+export { EMBEDDING_MODEL_ID }
+
+export type EmbedType = 'query' | 'passage'
+
+// Passages per onnxruntime forward pass. The whole-array embed (a startup index
+// build over thousands of shards+fragments) otherwise allocates activation
+// tensors for every input at once, and that single spike OOM-kills the
+// container mid-build — the agent boots, then dies with only a SIGKILL. Chunking
+// caps peak memory at one batch's worth regardless of corpus size; the model is
+// loaded once and reused across chunks, so the only cost is sequential passes.
+const EMBED_BATCH_SIZE = 64
+
+type TransformersEnv = typeof TransformersEnvValue
+type FeatureExtractor = Awaited<ReturnType<typeof TransformersPipeline<'feature-extraction'>>>
+
+// Defer the transformers (and thus sharp/onnxruntime) module load until an
+// embedding is actually requested. typeclaw's memory plugin is always loaded
+// and `vector.enabled` defaults to false, so a top-level static import would
+// drag the heavy native stack onto every container boot — and crash it when
+// sharp can't resolve its platform binary. Memoized so the module evaluates
+// at most once.
+type TransformersModule = { env: TransformersEnv; pipeline: typeof TransformersPipeline }
+
+let transformersModulePromise: Promise<TransformersModule> | undefined
+
+const realTransformersImport = (): Promise<TransformersModule> =>
+  import('@huggingface/transformers').then((mod) => ({ env: mod.env, pipeline: mod.pipeline }))
+
+// Injectable importer seam. Defaults to the real dynamic import; a test can
+// swap it to drive the module-load layer (e.g. fail once, then succeed) without
+// fighting Bun's mock.module namespace snapshotting. Bun freezes the mocked
+// namespace at registration, so a runtime-toggled failure can't be expressed
+// through mock.module — this seam is the supported way to exercise it.
+let importTransformers: () => Promise<TransformersModule> = realTransformersImport
+
+export function __setTransformersImporterForTests(importer: (() => Promise<TransformersModule>) | undefined): void {
+  importTransformers = importer ?? realTransformersImport
+  transformersModulePromise = undefined
+}
+
+// Injectable cache-compatibility check, mirroring the importer seam above. In
+// production it asserts the host-stamped sentinel matches this container before
+// the local_files_only load. The embedder's own mechanics tests (batching,
+// lazy-load, warm-up) mock transformers and run against the default cache path
+// with no sentinel, so they swap in a no-op — the sentinel guard has its own
+// dedicated coverage in embedding-model.test.ts.
+const realModelCacheCheck = (): Promise<void> =>
+  assertModelCacheCompatible(modelCachePath(), { transformers: getResolvedTransformersVersion() })
+
+let verifyModelCache: () => Promise<void> = realModelCacheCheck
+
+export function __setModelCacheCheckForTests(check: (() => Promise<void>) | undefined): void {
+  verifyModelCache = check ?? realModelCacheCheck
+}
+
+function loadTransformers(): Promise<TransformersModule> {
+  // Clear the memo on rejection (mirroring getEmbedder) so a transient failure
+  // of the dynamic import / native module load doesn't cache the rejected
+  // promise — otherwise every later getEmbedder() awaits the same dead promise
+  // and per-turn embedding stays poisoned for the life of the process.
+  transformersModulePromise ??= importTransformers().catch((err) => {
+    transformersModulePromise = undefined
+    throw err
+  })
+  return transformersModulePromise
+}
+
+export class Embedder {
+  private constructor(private readonly extractor: FeatureExtractor) {}
+
+  static async load(): Promise<Embedder> {
+    const { env, pipeline } = await loadTransformers()
+    // Guard the cache BEFORE local_files_only load: a host/container transformers
+    // drift (or a hand-copied cache) otherwise surfaces as a cryptic missing-file
+    // miss, or silently loads a stale variant. Fails loudly with a refresh hint.
+    await verifyModelCache()
+    configureTransformers(env)
+    const extractor = await pipeline('feature-extraction', MODEL_NAME, { local_files_only: true, dtype: MODEL_DTYPE })
+    return new Embedder(extractor)
+  }
+
+  async embed(texts: string[], type: EmbedType): Promise<Float32Array[]> {
+    if (texts.length === 0) return []
+
+    // Bound every input to the model's token budget BEFORE the tokenizer sees it.
+    // The tokenizer would otherwise truncate silently at 512 tokens; bounding
+    // here makes the cut deterministic and owned by us (the leading heading /
+    // belief sentence — the load-bearing retrieval signal — always survives
+    // because it comes first). The dreaming subagent separately compacts the
+    // topic shards that trip this, but bounding guarantees no silent loss even
+    // for inputs dreaming never rewrites — queries and stream fragments — which
+    // the dreaming over_budget table does not cover, so this is their only
+    // observability path.
+    const results = texts.map((text) => boundEmbeddableText(text))
+    warnIfBounded(results, type)
+
+    // Log the total embed size up front so a process that dies mid-build still
+    // leaves a line naming how much it was embedding. The work is chunked below
+    // (EMBED_BATCH_SIZE per onnxruntime pass) so peak memory no longer scales
+    // with this total — but the count remains the useful breadcrumb for a slow
+    // or wedged build.
+    logEmbedBatch(texts.length, type)
+
+    const prefixed = prefixTexts(
+      results.map((r) => r.text),
+      type,
+    )
+
+    // Gate per-chunk progress on the same LARGE_EMBED threshold as the up-front
+    // line: only the startup index migration runs long enough for "wedged or
+    // just slow?" to be a real question. Smaller embeds (queries, per-write
+    // upserts) finish in a pass or two and would only spam the logs.
+    const reportProgress = prefixed.length >= LARGE_EMBED
+
+    const embeddings: Float32Array[] = []
+    for (let start = 0; start < prefixed.length; start += EMBED_BATCH_SIZE) {
+      const batch = prefixed.slice(start, start + EMBED_BATCH_SIZE)
+      const output = await this.extractor(batch, { pooling: 'mean', normalize: true })
+      embeddings.push(...toEmbeddings(output.data, batch.length))
+      if (reportProgress) logEmbedProgress(embeddings.length, prefixed.length, type)
+    }
+    return embeddings
+  }
+}
+
+let embedderInstance: Promise<Embedder> | null = null
+
+export function getEmbedder(): Promise<Embedder> {
+  // Clear the memo on rejection so a transient load failure (e.g. boot warm-up
+  // racing the host model mount) degrades to a retry on the next call instead
+  // of caching the rejected promise and poisoning every later per-turn embed.
+  embedderInstance ??= Embedder.load().catch((err) => {
+    embedderInstance = null
+    throw err
+  })
+  return embedderInstance
+}
+
+// Boot-time readiness step: force the lazy embedder to load now so the first
+// per-turn query embed doesn't pay the ~2-5s ONNX init on the critical path.
+// Only called on the vector-enabled boot path (see src/run/index.ts), which
+// preserves embedder.ts's lazy-import guarantee for vector-off boots.
+export async function warmEmbedder(): Promise<void> {
+  await getEmbedder()
+}
+
+export async function embed(texts: string[], type: EmbedType): Promise<Float32Array[]> {
+  return (await getEmbedder()).embed(texts, type)
+}
+
+// Structured, content-free signal when any input was bounded, so a truncation
+// is observable in logs (the dreaming over_budget table only covers topic
+// shards — this is the only path for queries and stream fragments). Logs counts
+// and the worst estimate only, never the text, so memory content can't leak.
+function warnIfBounded(results: readonly BoundedText[], type: EmbedType): void {
+  const trimmed = results.filter((r) => r.bounded)
+  if (trimmed.length === 0) return
+  const worst = trimmed.reduce((max, r) => Math.max(max, r.estimatedTokens), 0)
+  console.warn(
+    `[memory] vector embedding: bounded ${trimmed.length}/${results.length} ${type} input(s) to the ` +
+      `${MAX_MODEL_TOKENS}-token model limit (worst ~${worst} est. tokens); their tail is not embedded`,
+  )
+}
+
+// A large embed (a startup index build over thousands of passages) used to be
+// the prime suspect for a container that boots, logs, then dies without an
+// error — the onnxruntime activation tensors spiked the OOM killer. The embed
+// is now chunked at EMBED_BATCH_SIZE, so this is no longer a fatal threshold;
+// it just marks where a build is large enough that its duration is worth noting
+// up front (the count remains the breadcrumb if the build wedges or is slow).
+const LARGE_EMBED = 256
+
+function logEmbedBatch(count: number, type: EmbedType): void {
+  const line = `[memory] vector embedding: ${count} ${type} input(s) (chunked at ${EMBED_BATCH_SIZE}/pass)`
+  if (count >= LARGE_EMBED) {
+    console.info(`${line} — large build, this may take a while`)
+  } else {
+    console.info(line)
+  }
+}
+
+function logEmbedProgress(done: number, total: number, type: EmbedType): void {
+  const pct = Math.floor((done / total) * 100)
+  console.info(`[memory] vector embedding: ${done}/${total} ${type} input(s) embedded (${pct}%)`)
+}
+
+function configureTransformers(env: TransformersEnv): void {
+  env.localModelPath = modelCachePath()
+  env.allowRemoteModels = false
+}
+
+function modelCachePath(): string {
+  const override = process.env.TYPECLAW_MODEL_CACHE
+  if (override && override.length > 0) return override
+  return join(homeRoot(), 'models')
+}
+
+function prefixTexts(texts: string[], type: EmbedType): string[] {
+  const prefix = type === 'query' ? 'query: ' : 'passage: '
+  return texts.map((text) => `${prefix}${text}`)
+}
+
+function toEmbeddings(data: unknown, count: number): Float32Array[] {
+  const values = toFloat32Array(data)
+  if (values.length !== count * DIMS) {
+    throw new Error(`unexpected ${MODEL_NAME} embedding size: got ${values.length}, expected ${count * DIMS}`)
+  }
+
+  return Array.from({ length: count }, (_, index) => values.slice(index * DIMS, (index + 1) * DIMS))
+}
+
+function toFloat32Array(data: unknown): Float32Array {
+  if (data instanceof Float32Array) return data
+  if (!isNumericArrayLike(data)) throw new Error(`${MODEL_NAME} returned non-numeric embeddings`)
+  return Float32Array.from(data)
+}
+
+function isNumericArrayLike(data: unknown): data is ArrayLike<number> {
+  if (!ArrayBuffer.isView(data) || !('length' in data)) return false
+  return !(data instanceof BigInt64Array) && !(data instanceof BigUint64Array)
+}