diff --git a/.gitignore b/.gitignore index 9957289..3501e6f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,5 +4,6 @@ dist/ .betterdb_context.md .mcp.json *.log +*.bun-build bun.lockb .idea/ diff --git a/CLAUDE.md b/CLAUDE.md index 70a04fa..47b72a6 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -48,6 +48,24 @@ bunx @betterdb/memory install ## Setup (for development) +### macOS + +**Never run Ollama in Docker on macOS.** Docker on macOS has no GPU +passthrough, so models run CPU-only inside the VM — measured 388s to first +token vs 10s native. Use the native Ollama.app and the Mac compose file, +which runs only Valkey (with AOF persistence enabled): + +```bash +docker compose -f docker-compose.mac.yml up -d # Valkey only +open -a Ollama # Native Ollama on :11434 +ollama pull mxbai-embed-large +ollama pull qwen2.5:3b +bun install +bun run setup-index +``` + +### Linux + ```bash docker compose up -d # Start Valkey (with search module) + Ollama bun install # Install dependencies diff --git a/docker-compose.mac.yml b/docker-compose.mac.yml new file mode 100644 index 0000000..6e0f642 --- /dev/null +++ b/docker-compose.mac.yml @@ -0,0 +1,21 @@ +services: + valkey: + image: valkey/valkey-bundle:9.1-alpine + command: + - valkey-server + - --dir + - /data + - --appendonly + - "yes" + - --save + - "60 1" + container_name: betterdb-valkey + ports: + - "6390:6379" + volumes: + - valkey-data:/data + restart: unless-stopped + +volumes: + valkey-data: + name: betterdb-valkey-data diff --git a/package.json b/package.json index deaf6de..0fac555 100644 --- a/package.json +++ b/package.json @@ -38,6 +38,7 @@ "test": "bun test", "test:unit": "bun test tests/unit", "test:integration": "bun test tests/integration", + "build:hooks": "bun run scripts/build-hooks.ts", "validate": "bash scripts/validate-pack.sh", "prepublishOnly": "bash scripts/validate-pack.sh", "setup-index": "bun run scripts/setup-index.ts", diff --git a/scripts/build-hooks.ts b/scripts/build-hooks.ts new file mode 100644 index 0000000..c678cf4 --- /dev/null +++ b/scripts/build-hooks.ts @@ -0,0 +1,57 @@ +#!/usr/bin/env bun + +/** + * Compile each hook entry point in HOOK_SPECS, plus the companion binaries + * hooks spawn at runtime, to standalone binaries. + * + * Usage: + * bun run scripts/build-hooks.ts + * + * Sources live in src/hooks/; binaries are written to + * dist/hooks/, the paths install-hooks.sh registers. + */ + +import { mkdirSync, readdirSync, rmSync } from "node:fs"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; +import { COMPANION_BINARIES, HOOK_SPECS } from "../src/hook-spec.js"; + +const projectRoot = resolve(dirname(fileURLToPath(import.meta.url)), ".."); +const hooksDir = join(projectRoot, "src", "hooks"); +const outDir = join(projectRoot, "dist", "hooks"); + +const OPTIONAL_PROVIDERS = ["openai"]; + +mkdirSync(outDir, { recursive: true }); + +const specs = [...HOOK_SPECS, ...COMPANION_BINARIES]; + +for (const spec of specs) { + const source = join(hooksDir, spec.source); + const outfile = join(outDir, spec.binary); + console.log(`Compiling ${spec.source} → dist/hooks/${spec.binary}`); + const result = Bun.spawnSync( + [ + "bun", + "build", + "--compile", + ...OPTIONAL_PROVIDERS.flatMap((pkg) => ["--external", pkg]), + source, + "--outfile", + outfile, + ], + { stdout: "inherit", stderr: "inherit" }, + ); + if (!result.success) { + console.error(`ERROR: failed to compile ${spec.source}`); + process.exit(1); + } +} + +for (const name of readdirSync(projectRoot)) { + if (name.endsWith(".bun-build")) { + rmSync(join(projectRoot, name), { force: true }); + } +} + +console.log(`Compiled ${specs.length} binaries to dist/hooks/`); diff --git a/scripts/install-hooks.sh b/scripts/install-hooks.sh index 64eec2b..b2ed482 100755 --- a/scripts/install-hooks.sh +++ b/scripts/install-hooks.sh @@ -42,7 +42,8 @@ if [ ! -f "$GLOBAL_SETTINGS" ]; then echo '{}' > "$GLOBAL_SETTINGS" fi -# Merge hooks into existing settings using Bun (preserves other fields, overwrites hooks block) +# Merge hooks into existing settings using Bun — replaces our own entries per +# event and preserves every other field, including third-party hooks # Each hook command sources the .env file first so compiled binaries get the right env vars # (bun build --compile binaries don't auto-load .env like `bun run` does) DIST_DIR="$PROJECT_DIR/dist/hooks" @@ -50,6 +51,7 @@ ENV_FILE="$PROJECT_DIR/.env" bun -e " const fs = require('fs'); +const { buildHookMap, isOwnHookEntry } = require('$PROJECT_DIR/src/hook-spec.ts'); const settingsPath = '$GLOBAL_SETTINGS'; const envFile = '$ENV_FILE'; const distDir = '$DIST_DIR'; @@ -59,12 +61,12 @@ const existing = JSON.parse(fs.readFileSync(settingsPath, 'utf8')); const wrap = (bin) => 'bash -c ' + JSON.stringify('set -a; [ -f ' + envFile + ' ] && . ' + envFile + '; set +a; ' + distDir + '/' + bin); -const hooks = { - SessionStart: [{ hooks: [{ type: 'command', command: wrap('session-start') }] }], - PreToolUse: [{ matcher: '', hooks: [{ type: 'command', command: wrap('pre-tool') }] }], - PostToolUse: [{ matcher: '', hooks: [{ type: 'command', command: wrap('post-tool') }] }], - SessionEnd: [{ hooks: [{ type: 'command', command: wrap('session-end') }] }], -}; +const hooks = { ...(existing.hooks || {}) }; +for (const [event, entries] of Object.entries(buildHookMap((spec) => wrap(spec.binary)))) { + const prev = Array.isArray(hooks[event]) ? hooks[event] : []; + const kept = prev.filter((entry) => !isOwnHookEntry(entry, [distDir, 'betterdb'])); + hooks[event] = [...kept, ...entries]; +} fs.writeFileSync(settingsPath, JSON.stringify({ ...existing, hooks }, null, 2)); console.log('Hook configuration written to: ' + settingsPath); " @@ -80,9 +82,17 @@ echo "" echo "Verifying global settings..." bun -e " const fs = require('fs'); +const { HOOK_SPECS } = require('$PROJECT_DIR/src/hook-spec.ts'); const settings = JSON.parse(fs.readFileSync('$HOME/.claude/settings.json', 'utf8')); -const hookCount = Object.keys(settings.hooks || {}).length; -console.log('Hooks registered: ' + hookCount + ' lifecycle events'); +const missing = HOOK_SPECS.filter((spec) => { + const entries = (settings.hooks || {})[spec.event] || []; + return !JSON.stringify(entries).includes(spec.binary); +}); +console.log('Hooks registered: ' + (HOOK_SPECS.length - missing.length) + '/' + HOOK_SPECS.length + ' lifecycle events'); +if (missing.length > 0) { + console.error('ERROR: not registered: ' + missing.map((spec) => spec.event).join(', ')); + process.exit(1); +} " # Summary @@ -90,10 +100,13 @@ echo "" echo "=== Installation Complete ===" echo "" echo "Hooks written to: ~/.claude/settings.json" -echo " SessionStart → $DIST_DIR/session-start" -echo " SessionEnd → $DIST_DIR/session-end" -echo " PreToolUse → $DIST_DIR/pre-tool" -echo " PostToolUse → $DIST_DIR/post-tool" +bun -e " +const { formatHookSummary } = require('$PROJECT_DIR/src/hook-spec.ts'); +const distDir = '$DIST_DIR'; +for (const line of formatHookSummary((spec) => distDir + '/' + spec.binary)) { + console.log(line); +} +" echo "" echo "MCP server: betterdb-memory (stdio)" echo "" diff --git a/scripts/register-hooks.ts b/scripts/register-hooks.ts index eb01dfd..1e3fd63 100644 --- a/scripts/register-hooks.ts +++ b/scripts/register-hooks.ts @@ -12,7 +12,13 @@ import { existsSync, readFileSync, writeFileSync, mkdirSync } from "node:fs"; import { join, resolve } from "node:path"; -import { stripLegacyBetterdbHooks } from "../src/hook-migration.js"; +import { + HOOK_SPECS, + buildHookMap, + formatHookSummary, + isOwnHookEntry, + type HookSpec, +} from "../src/hook-spec.js"; const pluginRoot = process.argv[2]; if (!pluginRoot) { @@ -24,15 +30,11 @@ const resolvedRoot = resolve(pluginRoot); const hooksDir = join(resolvedRoot, "src", "hooks"); // Verify hook source files exist -const hookFiles = [ - "session-start.ts", - "pre-tool.ts", - "post-tool.ts", - "session-end.ts", -]; -for (const file of hookFiles) { - if (!existsSync(join(hooksDir, file))) { - console.error(`ERROR: Hook source not found: ${join(hooksDir, file)}`); +for (const spec of HOOK_SPECS) { + if (!existsSync(join(hooksDir, spec.source))) { + console.error( + `ERROR: Hook source not found: ${join(hooksDir, spec.source)}`, + ); process.exit(1); } } @@ -60,37 +62,18 @@ if (existsSync(settingsPath)) { } } -function cmd(hookFile: string): string { - return `bash -c 'bun run "${join(hooksDir, hookFile)}"'`; +function cmd(spec: HookSpec): string { + return `bash -c 'bun run "${join(hooksDir, spec.source)}"'`; } // Merge hooks — replaces BetterDB entries per event, preserves all others. -// The loop below only visits events in betterdbHooks, so the legacy Stop -// registration must be stripped explicitly or it survives forever. -const existingHooks = stripLegacyBetterdbHooks( - (settings["hooks"] ?? {}) as Record, - ["betterdb", hooksDir], -); -const betterdbHooks: Record = { - SessionStart: [ - { hooks: [{ type: "command", command: cmd("session-start.ts") }] }, - ], - PreToolUse: [ - { matcher: "", hooks: [{ type: "command", command: cmd("pre-tool.ts") }] }, - ], - PostToolUse: [ - { matcher: "", hooks: [{ type: "command", command: cmd("post-tool.ts") }] }, - ], - SessionEnd: [ - { hooks: [{ type: "command", command: cmd("session-end.ts") }] }, - ], -}; +const existingHooks = (settings["hooks"] ?? {}) as Record; +const betterdbHooks = buildHookMap(cmd); for (const [event, entries] of Object.entries(betterdbHooks)) { const prev = Array.isArray(existingHooks[event]) ? existingHooks[event] : []; const filtered = prev.filter((entry) => { - const json = JSON.stringify(entry); - return !json.includes("betterdb") && !json.includes(hooksDir); + return !isOwnHookEntry(entry, [hooksDir, "betterdb"]); }); existingHooks[event] = [...filtered, ...entries]; } @@ -99,9 +82,8 @@ settings["hooks"] = existingHooks; writeFileSync(settingsPath, JSON.stringify(settings, null, 2) + "\n"); console.log("BetterDB Memory — Hooks registered in ~/.claude/settings.json\n"); -console.log(" SessionStart → session-start.ts"); -console.log(" PreToolUse → pre-tool.ts"); -console.log(" PostToolUse → post-tool.ts"); -console.log(" SessionEnd → session-end.ts"); +for (const line of formatHookSummary()) { + console.log(line); +} console.log(`\n Plugin root: ${resolvedRoot}`); console.log("\n Restart Claude Code for hooks to take effect."); diff --git a/src/client/providers/ollama.ts b/src/client/providers/ollama.ts index dd2b81f..dceda15 100644 --- a/src/client/providers/ollama.ts +++ b/src/client/providers/ollama.ts @@ -4,13 +4,32 @@ import { SessionSummarySchema, type SessionSummary } from "../../memory/schema.j import type { ModelClient, ModelPreset } from "../model.js"; import { buildSummarizePrompt, stripCodeFences } from "./_prompt.js"; +export interface OllamaTransport { + chat(request: { + model: string; + messages: { role: string; content: string }[]; + format: string; + stream: true; + keep_alive: string; + }): Promise>; + embed(request: { + model: string; + input: string; + }): Promise<{ embeddings: number[][] }>; +} + export class OllamaModelClient implements ModelClient { - private ollama: Ollama; + private ollama: OllamaTransport; readonly preset: ModelPreset; readonly embedDim: number; - constructor(preset: ModelPreset, ollamaUrl?: string) { - this.ollama = new Ollama({ host: ollamaUrl ?? config.ollama.url }); + constructor( + preset: ModelPreset, + ollamaUrl?: string, + transport?: OllamaTransport, + ) { + this.ollama = + transport ?? new Ollama({ host: ollamaUrl ?? config.ollama.url }); this.preset = preset; this.embedDim = preset.embedDim; } @@ -27,17 +46,30 @@ export class OllamaModelClient implements ModelClient { return first; } + /** + * Streamed rather than awaited whole: Bun's fetch enforces an idle timeout, + * and a non-streaming generate sends no bytes until the entire summary is + * done — long generations died as TimeoutError. Chunks keep the connection + * alive for as long as the model keeps producing. + * + * A cold model load is silent for longer than the idle timeout allows, but + * the load keeps going server-side after the client gives up — so one + * timed-out attempt is retried against the by-then warm model, and + * keep_alive holds the model in memory between calls. + */ async summarize(transcript: string): Promise { - const response = await this.ollama.chat({ - model: this.preset.summarizeModel, - messages: [ - { role: "user", content: buildSummarizePrompt(transcript) }, - ], - format: "json", - }); + let content: string; + try { + content = await this.chatSummary(transcript); + } catch (err) { + if (!(err instanceof DOMException && err.name === "TimeoutError")) { + throw err; + } + content = await this.chatSummary(transcript); + } const parsed = SessionSummarySchema.safeParse( - JSON.parse(stripCodeFences(response.message.content)), + JSON.parse(stripCodeFences(content)), ); if (!parsed.success) { @@ -50,4 +82,22 @@ export class OllamaModelClient implements ModelClient { return parsed.data; } + + private async chatSummary(transcript: string): Promise { + const stream = await this.ollama.chat({ + model: this.preset.summarizeModel, + messages: [ + { role: "user", content: buildSummarizePrompt(transcript) }, + ], + format: "json", + stream: true, + keep_alive: config.ollama.keepAlive, + }); + + let content = ""; + for await (const chunk of stream) { + content += chunk.message.content; + } + return content; + } } diff --git a/src/client/valkey.ts b/src/client/valkey.ts index bf3826b..f0d999f 100644 --- a/src/client/valkey.ts +++ b/src/client/valkey.ts @@ -429,21 +429,20 @@ let clientInstance: ValkeyClient | null = null; export async function getValkeyClient(): Promise { if (clientInstance) return clientInstance; - const retryDelays = [100, 500, 2000]; + // Hooks run on every tool call, so a dead or unresponsive Valkey must cost + // milliseconds, not minutes. + const retryDelays = [100, 300]; let lastError: Error | null = null; - for (const delay of retryDelays) { + for (const delay of [...retryDelays, 0]) { try { - const redis = new Redis(config.valkey.url, { - maxRetriesPerRequest: 3, - lazyConnect: true, - }); - await redis.connect(); - clientInstance = new ValkeyClient(redis); + clientInstance = new ValkeyClient(await connectValkey(config.valkey.url)); return clientInstance; } catch (err) { lastError = err instanceof Error ? err : new Error(String(err)); - await new Promise((resolve) => setTimeout(resolve, delay)); + if (delay > 0) { + await new Promise((resolve) => setTimeout(resolve, delay)); + } } } @@ -455,3 +454,40 @@ export async function getValkeyClient(): Promise { export function resetValkeyClient(): void { clientInstance = null; } + +/** + * One connection attempt under a hard deadline. connectTimeout only bounds + * the TCP connect — a proxy that accepts for a missing backend passes it and + * then hangs the ready check — so the whole attempt races the deadline, and + * the null retryStrategy stops ioredis from retrying forever underneath the + * caller's retry loop. + */ +export async function connectValkey( + url: string, + deadlineMs = 1000, +): Promise { + const redis = new Redis(url, { + maxRetriesPerRequest: 3, + lazyConnect: true, + connectTimeout: deadlineMs, + retryStrategy: () => null, + }); + let deadline: ReturnType | undefined; + try { + await Promise.race([ + redis.connect(), + new Promise((_, reject) => { + deadline = setTimeout( + () => reject(new Error("connection attempt deadline exceeded")), + deadlineMs, + ); + }), + ]); + return redis; + } catch (err) { + redis.disconnect(); + throw err; + } finally { + clearTimeout(deadline); + } +} diff --git a/src/config.ts b/src/config.ts index fb486a0..14bf05a 100644 --- a/src/config.ts +++ b/src/config.ts @@ -43,6 +43,7 @@ export const config = { embedModel: env("BETTERDB_EMBED_MODEL") ?? "mxbai-embed-large", summarizeModel: env("BETTERDB_SUMMARIZE_MODEL") ?? "mistral:7b", embedDim: Number(env("BETTERDB_EMBED_DIM") ?? 1024), + keepAlive: env("BETTERDB_OLLAMA_KEEP_ALIVE") ?? "5m", }, memory: { maxContextMemories: Number(env("BETTERDB_MAX_CONTEXT_MEMORIES") ?? 5), diff --git a/src/hook-migration.ts b/src/hook-migration.ts deleted file mode 100644 index 84c8d1c..0000000 --- a/src/hook-migration.ts +++ /dev/null @@ -1,38 +0,0 @@ -// Events this plugin used to register on and no longer does. mergeHooks only -// touches events present in the current hook map, so without an explicit strip -// a stale registration would survive upgrades forever in a user's -// ~/.claude/settings.json. -const LEGACY_EVENTS = ["Stop"]; - -/** - * `markers` identify an entry as ours. The default catches installed binaries - * under ~/.betterdb; dev registrations point at a plugin checkout whose path - * may not contain "betterdb", so callers pass that path as an extra marker. - */ -export function stripLegacyBetterdbHooks( - hooks: Record, - markers: string[] = ["betterdb"], -): Record { - const out: Record = { ...hooks }; - - for (const event of LEGACY_EVENTS) { - const entries = out[event]; - if (!Array.isArray(entries)) { - continue; - } - // Only ours — a third party's hook on the same event must survive. - const kept = entries.filter((entry) => { - const json = JSON.stringify(entry); - return !markers.some((m) => { - return m.length > 0 && json.includes(m); - }); - }); - if (kept.length > 0) { - out[event] = kept; - } else { - delete out[event]; - } - } - - return out; -} diff --git a/src/hook-spec.ts b/src/hook-spec.ts new file mode 100644 index 0000000..547ec4b --- /dev/null +++ b/src/hook-spec.ts @@ -0,0 +1,141 @@ +export type HookEvent = + | "SessionStart" + | "PreToolUse" + | "PostToolUse" + | "SessionEnd" + | "Stop"; + +export interface HookSpec { + readonly event: HookEvent; + readonly source: string; + readonly binary: string; + readonly matcher?: string; + readonly timeout: number; +} + +export interface HookCommand { + readonly type: "command"; + readonly command: string; + readonly timeout: number; +} + +export interface HookEntry { + readonly matcher?: string; + readonly hooks: readonly HookCommand[]; +} + +/** + * Timeouts are seconds, enforced by Claude Code per hook invocation. They are + * a backstop over the client-side fail-fast connect: pre/post fire on every + * tool call and must never make a tool call feel slow; the session-boundary + * hooks may parse large transcripts and get more room. + */ +export const HOOK_SPECS: readonly HookSpec[] = [ + { + event: "SessionStart", + source: "session-start.ts", + binary: "session-start", + timeout: 30, + }, + { + event: "PreToolUse", + source: "pre-tool.ts", + binary: "pre-tool", + matcher: "", + timeout: 10, + }, + { + event: "PostToolUse", + source: "post-tool.ts", + binary: "post-tool", + matcher: "", + timeout: 10, + }, + { + event: "SessionEnd", + source: "session-end.ts", + binary: "session-end", + timeout: 60, + }, + { + event: "Stop", + source: "stop-checkpoint.ts", + binary: "stop-checkpoint", + timeout: 30, + }, +]; + +export interface BinarySpec { + readonly source: string; + readonly binary: string; +} + +/** + * Non-hook binaries that must ship next to the hook binaries: the SessionEnd + * hook resolves the drainer as a sibling of its own executable, so every + * shape that compiles hooks must compile these too. + */ +export const COMPANION_BINARIES: readonly BinarySpec[] = [ + { source: "drain.ts", binary: "drain" }, +]; + +export const HOOK_COUNT = HOOK_SPECS.length; + +const HOOK_SOURCES: readonly string[] = HOOK_SPECS.map((spec) => spec.source); + +const HOOK_BINARY_PATHS: readonly string[] = HOOK_SPECS.map( + (spec) => `dist/hooks/${spec.binary}`, +); + +/** + * Whether a hook entry already in ~/.claude/settings.json is one this plugin + * wrote, and so may be replaced. + * + * Matching an install path alone is not enough: register-hooks.ts points hooks + * at a checkout whose path need not contain "betterdb", so such an entry + * survived a later install and kept firing alongside the new registration. The + * source filenames and dist/hooks binary paths come from HOOK_SPECS, so they + * identify our entries wherever the checkout lives — install-hooks.sh + * registers compiled binaries that contain no source filename. `markers` adds + * the caller's own install location. + * + * Deliberately conservative — a false positive deletes a third party's hook, + * which is worse than leaving one of ours behind. + */ +export function isOwnHookEntry( + entry: unknown, + markers: readonly string[] = [], +): boolean { + const json = JSON.stringify(entry) ?? ""; + if (HOOK_SOURCES.some((source) => json.includes(source))) { + return true; + } + if (HOOK_BINARY_PATHS.some((path) => json.includes(path))) { + return true; + } + return markers.some((marker) => marker.length > 0 && json.includes(marker)); +} + +export function buildHookMap( + toCommand: (spec: HookSpec) => string, +): Record { + const map: Record = {}; + for (const spec of HOOK_SPECS) { + const hooks: HookCommand[] = [ + { type: "command", command: toCommand(spec), timeout: spec.timeout }, + ]; + const entry: HookEntry = + spec.matcher === undefined ? { hooks } : { matcher: spec.matcher, hooks }; + (map[spec.event] ??= []).push(entry); + } + return map; +} + +export function formatHookSummary( + describe: (spec: HookSpec) => string = (spec) => spec.source, +): string[] { + const width = Math.max(...HOOK_SPECS.map((spec) => spec.event.length)); + return HOOK_SPECS.map( + (spec) => ` ${spec.event.padEnd(width)} → ${describe(spec)}`, + ); +} diff --git a/src/hooks/drain.ts b/src/hooks/drain.ts index 69eff47..7b02161 100644 --- a/src/hooks/drain.ts +++ b/src/hooks/drain.ts @@ -17,7 +17,7 @@ async function main(): Promise { const store = await getPluginMemoryStore((t) => modelClient.embed(t)); const pipeline = new AgingPipeline(valkeyClient, store, modelClient); - const { processed, skipped } = await pipeline.processIngestQueue(); + const { processed, skipped } = await pipeline.drainIngestQueue(); console.error(`[betterdb] drain: processed=${processed} skipped=${skipped}`); await store.close(); diff --git a/src/hooks/flush-segments.ts b/src/hooks/flush-segments.ts new file mode 100644 index 0000000..bfdcd6a --- /dev/null +++ b/src/hooks/flush-segments.ts @@ -0,0 +1,44 @@ +export interface SegmentSink { + pushIngestQueue(transcript: string, meta: object): Promise; +} + +export interface FlushMeta { + readonly project: string; + readonly branch: string; + readonly sessionId: string; + readonly baseSegment: number; +} + +/** + * Push tail segments to the ingest queue, numbering them after the segments + * the Stop hook already queued. Swallows a mid-loop failure and reports how + * far it got: SessionEnd fires exactly once, so an escaped exception here + * would skip both the drain spawn and cleanup, stranding everything already + * queued this session. + */ +export async function flushSegments( + client: SegmentSink, + segments: readonly string[], + meta: FlushMeta, +): Promise<{ pushed: number; failed: boolean }> { + let pushed = 0; + for (const segment of segments) { + try { + await client.pushIngestQueue(segment, { + project: meta.project, + branch: meta.branch, + timestamp: new Date().toISOString(), + sessionId: meta.sessionId, + segment: meta.baseSegment + pushed, + }); + } catch (err) { + console.error( + `[betterdb] flush failed after ${pushed}/${segments.length} segments:`, + err instanceof Error ? err.message : String(err), + ); + return { pushed, failed: true }; + } + pushed++; + } + return { pushed, failed: false }; +} diff --git a/src/hooks/locate-drain.ts b/src/hooks/locate-drain.ts new file mode 100644 index 0000000..fb56cac --- /dev/null +++ b/src/hooks/locate-drain.ts @@ -0,0 +1,36 @@ +import { join } from "node:path"; + +export interface DrainLocations { + readonly execDir: string; + readonly installBinDir: string; + readonly sourceDir: string; +} + +/** + * Resolve the drainer to spawn, in order of preference: the binary compiled + * next to the running hook (dev builds via build:hooks), the binary the CLI + * install placed in ~/.betterdb/bin, then bun-running the source (the + * register-hooks.ts shape). Inside a compiled executable sourceDir is a bunfs + * virtual path where drain.ts never exists, which is why the sibling binary + * must be checked first. + */ +export async function locateDrainCommand( + locations: DrainLocations, +): Promise { + const siblingBin = join(locations.execDir, "drain"); + if (await Bun.file(siblingBin).exists()) { + return [siblingBin]; + } + + const installedBin = join(locations.installBinDir, "drain"); + if (await Bun.file(installedBin).exists()) { + return [installedBin]; + } + + const source = join(locations.sourceDir, "drain.ts"); + if (await Bun.file(source).exists()) { + return ["bun", "run", source]; + } + + return null; +} diff --git a/src/hooks/session-end.ts b/src/hooks/session-end.ts index ccf21cf..c044c94 100644 --- a/src/hooks/session-end.ts +++ b/src/hooks/session-end.ts @@ -6,10 +6,19 @@ import { getCwdProject, } from "../memory/capture.js"; import { SessionEventSchema } from "../memory/schema.js"; -import { selectTranscript, type TranscriptTurn } from "../memory/transcript.js"; +import { + CHECKPOINT_THRESHOLD, + parseTurnsFrom, + planTailSegments, + readCheckpoint, + removeCheckpoint, + type OffsetTurn, +} from "../memory/checkpoint.js"; import { config, isConfigured } from "../config.js"; +import { flushSegments } from "./flush-segments.js"; +import { locateDrainCommand } from "./locate-drain.js"; import { unlink } from "node:fs/promises"; -import { join } from "node:path"; +import { dirname, join } from "node:path"; /** * SessionEnd hook: captures the session transcript and queues it. @@ -39,17 +48,20 @@ runHook(async () => { } const eventFilePath = `/tmp/betterdb-${sessionId}.jsonl`; - let turns: TranscriptTurn[] = []; + const checkpoint = await readCheckpoint(sessionId); + let turns: OffsetTurn[] = []; - // Prefer transcript_path — contains the full conversation including user messages if (transcriptPath) { - turns = await parseTranscriptTurns(transcriptPath); + turns = await parseTurnsFrom(transcriptPath, checkpoint.byteOffset); } - // Fall back to JSONL event file (tool calls captured by PostToolUse hook). - // Event lines rank as tool turns; with no user turns present the selector - // keeps them, so a tool-only fallback transcript is never emptied. - if (turns.length === 0) { + // Fall back to JSONL event file (tool calls captured by PostToolUse hook) + // only when no checkpoint ever ran. Event lines rank as tool turns; with no + // user turns present the selector keeps them, so a tool-only fallback + // transcript is never emptied. After a checkpoint an empty tail means the + // transcript is already fully queued, and the event file covers the whole + // session — re-queueing it would duplicate earlier segments. + if (turns.length === 0 && checkpoint.segment === 0) { const eventFile = Bun.file(eventFilePath); if (await eventFile.exists()) { const raw = await eventFile.text(); @@ -66,19 +78,29 @@ runHook(async () => { .buildTranscript() .split("\n") .filter(Boolean) - .map((text) => ({ role: "tool" as const, text })); + .map((text) => ({ role: "tool" as const, text, endByte: 0 })); } } - // Cap to ~8K chars for the summarizer via priority-based selection (user - // turns > assistant turns near user turns > tool lines) instead of the old - // head+tail slice, which dropped the middle of long sessions wholesale. + // Chunk the tail the way Stop does rather than capping it once: a session + // whose Stop hook never checkpointed arrives here whole, and a single cap + // would discard everything past it. Only the sub-threshold remainder is + // selected down. const MAX_TRANSCRIPT = 8000; - const transcript = selectTranscript(turns, MAX_TRANSCRIPT); - - // Nothing to store - if (!transcript || transcript.length < 20) { - await cleanup(eventFilePath); + const segments = planTailSegments( + turns, + CHECKPOINT_THRESHOLD, + MAX_TRANSCRIPT, + ).filter((text) => text.length >= 20); + + // Nothing new to store. The Stop hook may still have queued segments this + // session, and it never spawns a drainer — so returning here without one + // would leave them unsummarized until some later session happened to drain. + if (segments.length === 0) { + if (checkpoint.segment > 0) { + await spawnDrain(); + } + await cleanup(eventFilePath, sessionId); return; } @@ -86,129 +108,59 @@ runHook(async () => { try { valkeyClient = await getValkeyClient(); } catch { - await cleanup(eventFilePath); + await cleanup(eventFilePath, sessionId); return; // Valkey unreachable — skip silently } - const project = getCwdProject(); - const branch = getGitBranch(); - - await valkeyClient.pushIngestQueue(transcript, { - project, - branch, - timestamp: new Date().toISOString(), + const { pushed } = await flushSegments(valkeyClient, segments, { + project: getCwdProject(), + branch: getGitBranch(), sessionId, + baseSegment: checkpoint.segment, }); - await spawnDrain(); + // Drain even after a partial flush: the Stop hook's segments and whatever + // was pushed before the failure are already queued and nothing else will + // summarize them this session. + if (pushed > 0 || checkpoint.segment > 0) { + await spawnDrain(); + } - await valkeyClient.quit(); - await cleanup(eventFilePath); + await valkeyClient.quit().catch(() => {}); + await cleanup(eventFilePath, sessionId); }); /** * Spawn the detached drainer. unref() releases it from this process's event * loop, so the hook exits immediately while summarization continues. - * - * Both install shapes must work: `install` compiles binaries into - * ~/.betterdb/bin, while register-hooks.ts registers `bun run ` and - * compiles nothing. Resolving only the compiled path left that second shape - * with no drainer at all — and the exists() guard made it silent. */ async function spawnDrain(): Promise { // HOME is unset on Windows, where install and config both fall back to // USERPROFILE. const home = process.env["HOME"] ?? process.env["USERPROFILE"] ?? ""; - const drainBin = join(home, ".betterdb", "bin", "drain"); - if (await Bun.file(drainBin).exists()) { - Bun.spawn([drainBin], { - stdin: "ignore", - stdout: "ignore", - stderr: "ignore", - }).unref(); - return; - } - - const drainSrc = join(import.meta.dir, "drain.ts"); - if (await Bun.file(drainSrc).exists()) { - Bun.spawn(["bun", "run", drainSrc], { - stdin: "ignore", - stdout: "ignore", - stderr: "ignore", - }).unref(); + const cmd = await locateDrainCommand({ + execDir: dirname(process.execPath), + installBinDir: join(home, ".betterdb", "bin"), + sourceDir: import.meta.dir, + }); + if (!cmd) { + console.error( + "[betterdb] no drain binary or source found — queued transcripts stay queued until `betterdb-memory drain` runs", + ); return; } - - console.error( - "[betterdb] no drain binary or source found — queued transcripts stay queued until `betterdb-memory drain` runs", - ); -} - -/** - * Parse Claude Code's transcript JSONL into role-tagged turns. - * The JSONL contains objects with type: "user" | "assistant" and message content. - * We extract user/assistant/tool turns so selectTranscript can rank them. - */ -async function parseTranscriptTurns(path: string): Promise { - const file = Bun.file(path); - if (!(await file.exists())) return []; - - const raw = await file.text(); - const turns: TranscriptTurn[] = []; - - for (const line of raw.split("\n").filter(Boolean)) { - try { - const entry = JSON.parse(line); - if (entry.type === "user" && entry.message?.content) { - const content = - typeof entry.message.content === "string" - ? entry.message.content - : Array.isArray(entry.message.content) - ? entry.message.content - .filter((b: { type: string }) => b.type === "text") - .map((b: { text: string }) => b.text) - .join("\n") - : ""; - // Skip system-generated messages (commands, caveats) - if ( - content && - !content.includes("") - ) { - turns.push({ role: "user", text: `User: ${content}` }); - } - } else if (entry.type === "assistant" && entry.message?.content) { - const content = - typeof entry.message.content === "string" - ? entry.message.content - : Array.isArray(entry.message.content) - ? entry.message.content - .filter((b: { type: string }) => b.type === "text") - .map((b: { text: string }) => b.text) - .join("\n") - : ""; - if (content) { - turns.push({ - role: "assistant", - text: `Assistant: ${content.slice(0, 2000)}`, - }); - } - } else if (entry.type === "tool_use" || entry.type === "tool_result") { - // Include tool names for context but keep it brief - const toolName = entry.tool_name ?? entry.name ?? ""; - if (toolName) { - turns.push({ role: "tool", text: `Tool: ${toolName}` }); - } - } - } catch { - // Skip malformed lines - } - } - - return turns; + Bun.spawn(cmd, { + stdin: "ignore", + stdout: "ignore", + stderr: "ignore", + }).unref(); } -async function cleanup(eventFilePath: string): Promise { +async function cleanup( + eventFilePath: string, + sessionId: string, +): Promise { + await removeCheckpoint(sessionId); try { await unlink(eventFilePath); } catch { diff --git a/src/hooks/stop-checkpoint.ts b/src/hooks/stop-checkpoint.ts new file mode 100644 index 0000000..bca5d42 --- /dev/null +++ b/src/hooks/stop-checkpoint.ts @@ -0,0 +1,68 @@ +import { readRawPayload, runHook } from "./_utils.js"; +import { getValkeyClient } from "../client/valkey.js"; +import { getCwdProject, getGitBranch } from "../memory/capture.js"; +import { isConfigured } from "../config.js"; +import { + CHECKPOINT_THRESHOLD, + nextChunk, + parseTurnsFrom, + readCheckpoint, + writeCheckpoint, +} from "../memory/checkpoint.js"; + +runHook(async () => { + if (!isConfigured()) { + return; + } + + const payload = await readRawPayload(); + const sessionId = payload["session_id"] as string; + const cwd = payload["cwd"] as string | undefined; + const transcriptPath = payload["transcript_path"] as string | undefined; + + if (!sessionId || !transcriptPath) { + return; + } + if (cwd) { + process.chdir(cwd); + } + + const transcript = Bun.file(transcriptPath); + if (!(await transcript.exists())) { + return; + } + + const checkpoint = await readCheckpoint(sessionId); + + if (transcript.size <= checkpoint.byteOffset) { + return; + } + + const turns = await parseTurnsFrom(transcriptPath, checkpoint.byteOffset); + const result = nextChunk(turns, CHECKPOINT_THRESHOLD); + if (result === null) { + return; + } + + let valkeyClient; + try { + valkeyClient = await getValkeyClient(); + } catch { + return; + } + + await valkeyClient.pushIngestQueue(result.chunk, { + project: getCwdProject(), + branch: getGitBranch(), + timestamp: new Date().toISOString(), + sessionId, + segment: checkpoint.segment, + }); + + await writeCheckpoint(sessionId, { + byteOffset: result.endByte, + segment: checkpoint.segment + 1, + }); + + await valkeyClient.quit(); +}); diff --git a/src/index.ts b/src/index.ts index 4ffbc16..7ea92d5 100644 --- a/src/index.ts +++ b/src/index.ts @@ -19,7 +19,13 @@ import { rmSync, } from "node:fs"; import { join, resolve } from "node:path"; -import { stripLegacyBetterdbHooks } from "./hook-migration.js"; +import { + COMPANION_BINARIES, + HOOK_COUNT, + HOOK_SPECS, + buildHookMap, + isOwnHookEntry, +} from "./hook-spec.js"; const VERSION = "0.5.0"; const HOME = process.env["HOME"] ?? process.env["USERPROFILE"] ?? ""; @@ -29,14 +35,13 @@ const CONFIG_PATH = join(BETTERDB_DIR, "memory.json"); const MANIFEST_PATH = join(BETTERDB_DIR, "install-manifest.json"); const PKG_ROOT = resolve(import.meta.dir, ".."); -const BINARIES = [ - { src: "src/hooks/session-start.ts", out: "session-start" }, - { src: "src/hooks/session-end.ts", out: "session-end" }, - { src: "src/hooks/pre-tool.ts", out: "pre-tool" }, - { src: "src/hooks/post-tool.ts", out: "post-tool" }, - { src: "src/hooks/drain.ts", out: "drain" }, +const BINARIES: readonly { src: string; out: string }[] = [ + ...[...HOOK_SPECS, ...COMPANION_BINARIES].map((spec) => ({ + src: `src/hooks/${spec.source}`, + out: spec.binary, + })), { src: "src/mcp/server.ts", out: "mcp-server" }, -] as const; +]; const USAGE = ` BetterDB Memory for Claude Code v${VERSION} @@ -235,35 +240,12 @@ async function runInstall() { } } - // mergeHooks only touches events present in betterdbHooks, so the legacy - // Stop registration must be stripped explicitly or it survives upgrades. - const existingHooks = stripLegacyBetterdbHooks( - (settings["hooks"] ?? {}) as Record, - ); - const betterdbHooks: Record = { - SessionStart: [ - { hooks: [{ type: "command", command: join(BIN_DIR, "session-start") }] }, - ], - PreToolUse: [ - { - matcher: "", - hooks: [{ type: "command", command: join(BIN_DIR, "pre-tool") }], - }, - ], - PostToolUse: [ - { - matcher: "", - hooks: [{ type: "command", command: join(BIN_DIR, "post-tool") }], - }, - ], - SessionEnd: [ - { hooks: [{ type: "command", command: join(BIN_DIR, "session-end") }] }, - ], - }; + const existingHooks = (settings["hooks"] ?? {}) as Record; + const betterdbHooks = buildHookMap((spec) => join(BIN_DIR, spec.binary)); settings["hooks"] = mergeHooks(existingHooks, betterdbHooks); writeFileSync(settingsPath, JSON.stringify(settings, null, 2) + "\n"); - console.log(" Registered 4 hooks in ~/.claude/settings.json"); + console.log(` Registered ${HOOK_COUNT} hooks in ~/.claude/settings.json`); // Register MCP server globally (-s user) so it's available in all projects const mcpBin = join(BIN_DIR, "mcp-server"); @@ -359,7 +341,7 @@ async function runInstall() { // 7. PRINT SUMMARY console.log("\n=== Installation Complete ===\n"); console.log(` ✅ Compiled ${BINARIES.length} binaries to ${BIN_DIR}/`); - console.log(" ✅ Registered 4 hooks with Claude Code"); + console.log(` ✅ Registered ${HOOK_COUNT} hooks with Claude Code`); console.log(" ✅ Registered MCP server: betterdb-memory"); console.log(" ✅ Valkey index ready"); console.log(` ✅ Config saved to ${CONFIG_PATH}`); @@ -572,7 +554,7 @@ async function runDrain() { const store = await getPluginMemoryStore((t) => modelClient.embed(t)); const pipeline = new AgingPipeline(valkeyClient, store, modelClient); - const { processed, skipped } = await pipeline.processIngestQueue(); + const { processed, skipped } = await pipeline.drainIngestQueue(); console.log( `Processed ${processed} queued transcript(s), skipped ${skipped} empty.`, ); @@ -920,10 +902,8 @@ function mergeHooks( const merged = { ...existing }; for (const [event, entries] of Object.entries(ours)) { const prev = Array.isArray(merged[event]) ? merged[event] : []; - // Filter out previous BetterDB entries (contain our BIN_DIR or betterdb path) const filtered = prev.filter((entry) => { - const json = JSON.stringify(entry); - return !json.includes(BIN_DIR) && !json.includes("betterdb"); + return !isOwnHookEntry(entry, [BIN_DIR, "betterdb"]); }); merged[event] = [...filtered, ...entries]; } diff --git a/src/memory/aging.ts b/src/memory/aging.ts index 5eeaa01..e372f6b 100644 --- a/src/memory/aging.ts +++ b/src/memory/aging.ts @@ -132,12 +132,18 @@ export class AgingPipeline { // --- Ingest Queue Processing --- - async processIngestQueue(): Promise<{ processed: number; skipped: number }> { + async processIngestQueue(): Promise<{ + processed: number; + skipped: number; + halted: boolean; + }> { const items = await this.valkeyClient.popIngestQueue(20); let processed = 0; let skipped = 0; + let halted = false; - for (const item of items) { + for (let i = 0; i < items.length; i++) { + const item = items[i]!; try { const summary = await this.modelClient.summarize(item.transcript); @@ -169,12 +175,43 @@ export class AgingPipeline { processed++; } catch (err) { console.error("[betterdb] Failed to process queued transcript:", err); - // Re-queue on failure - await this.valkeyClient.pushIngestQueue(item.transcript, item.meta); + // The pop was destructive: everything not yet processed must go back, + // not just the item that failed, or the rest of the batch is lost + // when this process exits. + for (const unprocessed of items.slice(i)) { + await this.valkeyClient.pushIngestQueue( + unprocessed.transcript, + unprocessed.meta, + ); + } + halted = true; break; } } + return { processed, skipped, halted }; + } + + /** + * Drain the ingest queue completely. Each pop is capped, but a + * checkpointing session can enqueue far more segments than one cap; the + * drainer runs detached with nothing waiting on it, so it loops until a + * round comes back empty. A halted round means the failed batch was + * re-queued — looping again would spin against a dead summarizer. + */ + async drainIngestQueue( + maxRounds = 50, + ): Promise<{ processed: number; skipped: number }> { + let processed = 0; + let skipped = 0; + for (let round = 0; round < maxRounds; round++) { + const result = await this.processIngestQueue(); + processed += result.processed; + skipped += result.skipped; + if (result.halted || result.processed + result.skipped === 0) { + break; + } + } return { processed, skipped }; } diff --git a/src/memory/checkpoint.ts b/src/memory/checkpoint.ts new file mode 100644 index 0000000..0309269 --- /dev/null +++ b/src/memory/checkpoint.ts @@ -0,0 +1,145 @@ +import { rename, unlink } from "node:fs/promises"; +import { parseTranscriptLine, selectTranscript } from "./transcript.js"; + +export interface OffsetTurn { + role: "user" | "assistant" | "tool"; + text: string; + endByte: number; +} + +export const CHECKPOINT_THRESHOLD = 8000; + +export function nextChunk( + turns: OffsetTurn[], + threshold: number, +): { chunk: string; endByte: number; consumedTurns: number } | null { + let length = 0; + for (let i = 0; i < turns.length; i++) { + length += (i > 0 ? 1 : 0) + turns[i]!.text.length; + if (length >= threshold) { + const consumed = turns.slice(0, i + 1); + return { + chunk: consumed.map((t) => t.text).join("\n"), + endByte: turns[i]!.endByte, + consumedTurns: i + 1, + }; + } + } + return null; +} + +/** + * Split a session-end tail into the segments to queue, in order. + * + * The Stop hook is an optimization, not a correctness requirement: when it + * never ran (fresh install before a restart, Valkey unreachable) the whole + * transcript arrives here. Capping that with a single `maxChars` selection + * would silently discard everything past the cap — the truncation checkpoint + * capture exists to remove. So chunk the tail exactly as Stop would, and let + * only the final sub-threshold remainder go through `selectTranscript`. + */ +export function planTailSegments( + turns: OffsetTurn[], + threshold: number, + maxChars: number, +): string[] { + const segments: string[] = []; + let remaining = turns; + + for (;;) { + const result = nextChunk(remaining, threshold); + if (result === null) { + break; + } + segments.push(result.chunk); + remaining = remaining.slice(result.consumedTurns); + } + + const tail = selectTranscript(remaining, maxChars); + if (tail.length > 0) { + segments.push(tail); + } + + return segments; +} + +export interface Checkpoint { + byteOffset: number; + segment: number; +} + +const ZERO_CHECKPOINT: Checkpoint = { byteOffset: 0, segment: 0 }; + +export function checkpointPath(sessionId: string): string { + return `/tmp/betterdb-${sessionId}.checkpoint.json`; +} + +export async function readCheckpoint(sessionId: string): Promise { + const file = Bun.file(checkpointPath(sessionId)); + if (!(await file.exists())) { + return { ...ZERO_CHECKPOINT }; + } + try { + const parsed = (await file.json()) as Partial; + if ( + typeof parsed.byteOffset !== "number" || + typeof parsed.segment !== "number" + ) { + return { ...ZERO_CHECKPOINT }; + } + return { byteOffset: parsed.byteOffset, segment: parsed.segment }; + } catch { + return { ...ZERO_CHECKPOINT }; + } +} + +export async function writeCheckpoint( + sessionId: string, + cp: Checkpoint, +): Promise { + const target = checkpointPath(sessionId); + const tmp = `${target}.tmp`; + await Bun.write(tmp, JSON.stringify(cp)); + await rename(tmp, target); +} + +export async function removeCheckpoint(sessionId: string): Promise { + await unlink(checkpointPath(sessionId)).catch(() => {}); +} + +export async function parseTurnsFrom( + path: string, + fromByte: number, +): Promise { + const file = Bun.file(path); + if (!(await file.exists())) { + return []; + } + + const bytes = new Uint8Array(await file.slice(fromByte).arrayBuffer()); + const decoder = new TextDecoder(); + const turns: OffsetTurn[] = []; + let lineStart = 0; + + for (let i = 0; i <= bytes.length; i++) { + const atEnd = i === bytes.length; + if (!atEnd && bytes[i] !== 0x0a) { + continue; + } + const lineBytes = bytes.subarray(lineStart, i); + const endByte = fromByte + (atEnd ? i : i + 1); + lineStart = i + 1; + if (lineBytes.length === 0) { + continue; + } + const line = decoder.decode(lineBytes).trim(); + if (!line) { + continue; + } + for (const turn of parseTranscriptLine(line)) { + turns.push({ ...turn, endByte }); + } + } + + return turns; +} diff --git a/src/memory/transcript.ts b/src/memory/transcript.ts index 26721bc..a8856da 100644 --- a/src/memory/transcript.ts +++ b/src/memory/transcript.ts @@ -93,3 +93,68 @@ export function selectTranscript( if (prev < turns.length - 1) out.push(GAP_MARKER); return out.join("\n"); } + +interface RawTranscriptEntry { + type?: string; + message?: { content?: unknown }; + tool_name?: string; + name?: string; +} + +function extractContent(content: unknown): string { + if (typeof content === "string") { + return content; + } + if (Array.isArray(content)) { + return content + .filter((b): b is { type: string; text: string } => { + return ( + typeof b === "object" && + b !== null && + (b as { type?: unknown }).type === "text" && + typeof (b as { text?: unknown }).text === "string" + ); + }) + .map((b) => b.text) + .join("\n"); + } + return ""; +} + +export function parseTranscriptLine(line: string): TranscriptTurn[] { + let entry: RawTranscriptEntry; + try { + entry = JSON.parse(line) as RawTranscriptEntry; + } catch { + return []; + } + + if (entry.type === "user" && entry.message?.content !== undefined) { + const content = extractContent(entry.message.content); + if ( + content && + !content.includes("") + ) { + return [{ role: "user", text: `User: ${content}` }]; + } + return []; + } + + if (entry.type === "assistant" && entry.message?.content !== undefined) { + const content = extractContent(entry.message.content); + if (content) { + return [{ role: "assistant", text: `Assistant: ${content.slice(0, 2000)}` }]; + } + return []; + } + + if (entry.type === "tool_use" || entry.type === "tool_result") { + const toolName = entry.tool_name ?? entry.name ?? ""; + if (toolName) { + return [{ role: "tool", text: `Tool: ${toolName}` }]; + } + } + + return []; +} diff --git a/tests/integration/checkpoint-capture.test.ts b/tests/integration/checkpoint-capture.test.ts new file mode 100644 index 0000000..1b1bc01 --- /dev/null +++ b/tests/integration/checkpoint-capture.test.ts @@ -0,0 +1,93 @@ +import { describe, expect, test, afterAll } from "bun:test"; +import { unlink } from "node:fs/promises"; +import { getValkeyClient, resetValkeyClient } from "../../src/client/valkey.js"; +import { + CHECKPOINT_THRESHOLD, + nextChunk, + parseTurnsFrom, + readCheckpoint, + writeCheckpoint, + removeCheckpoint, + checkpointPath, +} from "../../src/memory/checkpoint.js"; + +const SKIP = Bun.env.BETTERDB_SKIP_INTEGRATION === "true"; +const SID = "checkpoint-capture-it"; +const FIXTURE = `/tmp/betterdb-${SID}.transcript.jsonl`; + +const META = { + project: "it", + branch: "it", + timestamp: "2026-07-17T00:00:00.000Z", + sessionId: SID, +}; + +async function seedTranscript(): Promise { + const lines: string[] = []; + for (let i = 0; i < 14; i++) { + lines.push( + JSON.stringify({ + type: "user", + message: { content: `turn ${i} ` + "x".repeat(2500) }, + }), + ); + } + await Bun.write(FIXTURE, lines.join("\n") + "\n"); +} + +describe.skipIf(SKIP)("checkpoint capture integration", () => { + afterAll(async () => { + await unlink(FIXTURE).catch(() => {}); + await unlink(checkpointPath(SID)).catch(() => {}); + await resetValkeyClient(); + }); + + test("produces sequential segments plus a tail", async () => { + await seedTranscript(); + await removeCheckpoint(SID); + + const vk = await getValkeyClient(); + // Start from a clean queue without assuming a raw() accessor. + await vk.popIngestQueue(10000); + + let loopSegments = 0; + let exitedViaNull = false; + for (let guard = 0; guard < 20; guard++) { + const cp = await readCheckpoint(SID); + const turns = await parseTurnsFrom(FIXTURE, cp.byteOffset); + const result = nextChunk(turns, CHECKPOINT_THRESHOLD); + if (result === null) { + exitedViaNull = true; + break; + } + await vk.pushIngestQueue(result.chunk, { ...META, segment: cp.segment }); + await writeCheckpoint(SID, { + byteOffset: result.endByte, + segment: cp.segment + 1, + }); + loopSegments++; + } + + const cp = await readCheckpoint(SID); + const tail = await parseTurnsFrom(FIXTURE, cp.byteOffset); + const tailText = tail.map((t) => t.text).join("\n"); + let tailPushed = false; + if (tailText.length >= 20) { + await vk.pushIngestQueue(tailText, { ...META, segment: cp.segment }); + tailPushed = true; + } + + const items = await vk.popIngestQueue(50); + await vk.quit(); + + expect(exitedViaNull).toBe(true); + expect(loopSegments).toBeGreaterThanOrEqual(2); + expect(tailPushed).toBe(true); + expect(items.length).toBe(loopSegments + 1); + const segments = items.map((i) => (i.meta as { segment: number }).segment); + expect(segments).toEqual([...Array(items.length).keys()]); + for (const item of items) { + expect(item.transcript.length).toBeGreaterThan(0); + } + }); +}); diff --git a/tests/unit/aging-guard.test.ts b/tests/unit/aging-guard.test.ts index 9ef6b89..00c2b95 100644 --- a/tests/unit/aging-guard.test.ts +++ b/tests/unit/aging-guard.test.ts @@ -37,6 +37,14 @@ function fakeModel(summary: SessionSummary) { }; } +function failingModel() { + return { + summarize: async () => { + throw new Error("summarizer unavailable"); + }, + }; +} + describe("processIngestQueue empty-summary guard", () => { test("does not store a husk", async () => { const husk = SessionSummarySchema.parse({}); @@ -87,3 +95,116 @@ describe("processIngestQueue empty-summary guard", () => { expect(result.skipped).toBe(0); }); }); + +function countedValkey(items: Array<{ transcript: string; meta: object }>) { + const requeued: string[] = []; + let pops = 0; + return { + client: { + popIngestQueue: async (count: number) => { + pops++; + return items.splice(0, count); + }, + pushIngestQueue: async (t: string) => { + requeued.push(t); + }, + }, + requeued, + pops: () => pops, + }; +} + +describe("drainIngestQueue", () => { + test("drains the whole queue across multiple pops", async () => { + // A checkpointing session can enqueue more segments than one pop cap; + // the drainer is detached and unclocked, so it keeps going until dry. + const real = SessionSummarySchema.parse({ oneLineSummary: "Did a thing" }); + const items = Array.from({ length: 45 }, (_, i) => ({ + transcript: `segment ${i}`, + meta: {}, + })); + const stored: SessionSummary[] = []; + const vk = countedValkey(items); + + const pipeline = new AgingPipeline( + vk.client as never, + fakeStore(stored) as never, + fakeModel(real) as never, + ); + const result = await pipeline.drainIngestQueue(); + + expect(result.processed).toBe(45); + expect(stored.length).toBe(45); + expect(vk.pops()).toBeGreaterThan(2); + }); + + test("stops after a failing round instead of spinning on a dead summarizer", async () => { + const items = Array.from({ length: 25 }, (_, i) => ({ + transcript: `segment ${i}`, + meta: {}, + })); + const vk = countedValkey(items); + + const pipeline = new AgingPipeline( + vk.client as never, + fakeStore([]) as never, + failingModel() as never, + ); + const result = await pipeline.drainIngestQueue(); + + expect(result.processed).toBe(0); + expect(vk.pops()).toBe(1); + expect(vk.requeued.length).toBe(20); + }); +}); + +describe("processIngestQueue failure handling", () => { + test("re-queues the failed item and every remaining popped item", async () => { + // The pop is destructive: anything popped but neither stored nor + // re-queued is lost forever when the drain process exits. + const vk = fakeValkey([ + { transcript: "first", meta: {} }, + { transcript: "second", meta: {} }, + { transcript: "third", meta: {} }, + ]); + + const pipeline = new AgingPipeline( + vk.client as never, + fakeStore([]) as never, + failingModel() as never, + ); + const result = await pipeline.processIngestQueue(); + + expect(result.processed).toBe(0); + expect(vk.requeued.sort()).toEqual(["first", "second", "third"]); + }); + + test("re-queues only the unprocessed tail when a later item fails", async () => { + const real = SessionSummarySchema.parse({ oneLineSummary: "Did a thing" }); + let calls = 0; + const flaky = { + summarize: async () => { + calls++; + if (calls > 1) throw new Error("summarizer died mid-batch"); + return real; + }, + }; + const stored: SessionSummary[] = []; + const vk = fakeValkey([ + { transcript: "first", meta: {} }, + { transcript: "second", meta: {} }, + { transcript: "third", meta: {} }, + ]); + + const pipeline = new AgingPipeline( + vk.client as never, + fakeStore(stored) as never, + flaky as never, + ); + const result = await pipeline.processIngestQueue(); + + expect(result.processed).toBe(1); + expect(stored.length).toBe(1); + expect(vk.requeued.sort()).toEqual(["second", "third"]); + }); +}); diff --git a/tests/unit/checkpoint-file.test.ts b/tests/unit/checkpoint-file.test.ts new file mode 100644 index 0000000..4bf1ccf --- /dev/null +++ b/tests/unit/checkpoint-file.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, test, afterEach } from "bun:test"; +import { unlink } from "node:fs/promises"; +import { + checkpointPath, + readCheckpoint, + writeCheckpoint, + removeCheckpoint, +} from "../../src/memory/checkpoint.js"; + +const SID = "checkpoint-file-test"; + +afterEach(async () => { + await unlink(checkpointPath(SID)).catch(() => {}); +}); + +describe("checkpoint file", () => { + test("missing file reads as the zero checkpoint", async () => { + expect(await readCheckpoint(SID)).toEqual({ byteOffset: 0, segment: 0 }); + }); + + test("round-trips a written checkpoint", async () => { + await writeCheckpoint(SID, { byteOffset: 8002, segment: 1 }); + expect(await readCheckpoint(SID)).toEqual({ byteOffset: 8002, segment: 1 }); + }); + + test("overwrites an existing checkpoint", async () => { + await writeCheckpoint(SID, { byteOffset: 100, segment: 1 }); + await writeCheckpoint(SID, { byteOffset: 200, segment: 2 }); + expect(await readCheckpoint(SID)).toEqual({ byteOffset: 200, segment: 2 }); + }); + + test("a corrupt file reads as the zero checkpoint", async () => { + await Bun.write(checkpointPath(SID), "{not json"); + expect(await readCheckpoint(SID)).toEqual({ byteOffset: 0, segment: 0 }); + }); + + test("removeCheckpoint deletes the file", async () => { + await writeCheckpoint(SID, { byteOffset: 5, segment: 1 }); + await removeCheckpoint(SID); + expect(await readCheckpoint(SID)).toEqual({ byteOffset: 0, segment: 0 }); + }); +}); diff --git a/tests/unit/checkpoint-nextchunk.test.ts b/tests/unit/checkpoint-nextchunk.test.ts new file mode 100644 index 0000000..b60151d --- /dev/null +++ b/tests/unit/checkpoint-nextchunk.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, test } from "bun:test"; +import { nextChunk, type OffsetTurn } from "../../src/memory/checkpoint.js"; + +function turn(text: string, endByte: number): OffsetTurn { + return { role: "user", text, endByte }; +} + +describe("nextChunk", () => { + test("returns null when accumulated turns are below threshold", () => { + const turns = [turn("a".repeat(10), 11), turn("b".repeat(10), 22)]; + expect(nextChunk(turns, 8000)).toBeNull(); + }); + + test("emits the turn that crosses the threshold, whole", () => { + const turns = [ + turn("a".repeat(5000), 5001), + turn("b".repeat(4000), 9002), + turn("c".repeat(10), 9013), + ]; + const result = nextChunk(turns, 8000); + expect(result).not.toBeNull(); + expect(result!.consumedTurns).toBe(2); + expect(result!.endByte).toBe(9002); + expect(result!.chunk).toBe("a".repeat(5000) + "\n" + "b".repeat(4000)); + }); + + test("emits at exactly the threshold", () => { + const turns = [turn("a".repeat(8000), 8001)]; + const result = nextChunk(turns, 8000); + expect(result).not.toBeNull(); + expect(result!.consumedTurns).toBe(1); + }); + + test("emits a single over-threshold turn alone, never split", () => { + const turns = [turn("a".repeat(20000), 20001), turn("b", 20003)]; + const result = nextChunk(turns, 8000); + expect(result!.consumedTurns).toBe(1); + expect(result!.chunk).toBe("a".repeat(20000)); + expect(result!.endByte).toBe(20001); + }); + + test("charges one newline between consumed turns", () => { + const turns = [turn("a".repeat(4000), 4001), turn("b".repeat(3999), 8001)]; + const result = nextChunk(turns, 8000); + expect(result).not.toBeNull(); + expect(result!.consumedTurns).toBe(2); + }); +}); diff --git a/tests/unit/checkpoint-parse.test.ts b/tests/unit/checkpoint-parse.test.ts new file mode 100644 index 0000000..fda79c6 --- /dev/null +++ b/tests/unit/checkpoint-parse.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, test, afterEach } from "bun:test"; +import { unlink } from "node:fs/promises"; +import { parseTurnsFrom } from "../../src/memory/checkpoint.js"; + +const FIXTURE = "/tmp/betterdb-parse-test.jsonl"; + +afterEach(async () => { + await unlink(FIXTURE).catch(() => {}); +}); + +async function writeLines(lines: object[]): Promise { + await Bun.write(FIXTURE, lines.map((l) => JSON.stringify(l)).join("\n") + "\n"); +} + +describe("parseTurnsFrom", () => { + test("returns [] for a missing file", async () => { + expect(await parseTurnsFrom("/tmp/does-not-exist.jsonl", 0)).toEqual([]); + }); + + test("parses all turns from offset 0 with end-byte offsets", async () => { + await writeLines([ + { type: "user", message: { content: "hello" } }, + { type: "assistant", message: { content: "hi" } }, + ]); + const turns = await parseTurnsFrom(FIXTURE, 0); + expect(turns.map((t) => t.text)).toEqual(["User: hello", "Assistant: hi"]); + expect(turns[0]!.endByte).toBeGreaterThan(0); + expect(turns[1]!.endByte).toBeGreaterThan(turns[0]!.endByte); + }); + + test("resuming from a prior endByte yields only later turns", async () => { + await writeLines([ + { type: "user", message: { content: "first" } }, + { type: "user", message: { content: "second" } }, + ]); + const all = await parseTurnsFrom(FIXTURE, 0); + const resumed = await parseTurnsFrom(FIXTURE, all[0]!.endByte); + expect(resumed.map((t) => t.text)).toEqual(["User: second"]); + }); + + test("the final endByte equals the file size", async () => { + await writeLines([{ type: "user", message: { content: "only" } }]); + const size = Bun.file(FIXTURE).size; + const turns = await parseTurnsFrom(FIXTURE, 0); + expect(turns[turns.length - 1]!.endByte).toBe(size); + }); + + test("final endByte equals file size when no trailing newline", async () => { + await Bun.write( + FIXTURE, + JSON.stringify({ type: "user", message: { content: "no newline" } }), + ); + const size = Bun.file(FIXTURE).size; + const turns = await parseTurnsFrom(FIXTURE, 0); + expect(turns).toHaveLength(1); + expect(turns[turns.length - 1]!.endByte).toBe(size); + }); +}); diff --git a/tests/unit/checkpoint-tail.test.ts b/tests/unit/checkpoint-tail.test.ts new file mode 100644 index 0000000..1b550c5 --- /dev/null +++ b/tests/unit/checkpoint-tail.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, test } from "bun:test"; +import { planTailSegments, type OffsetTurn } from "../../src/memory/checkpoint.js"; + +const THRESHOLD = 8000; +const MAX_CHARS = 8000; + +function turn(text: string): OffsetTurn { + return { role: "user", text, endByte: 0 }; +} + +describe("planTailSegments", () => { + test("returns nothing for no turns", () => { + expect(planTailSegments([], THRESHOLD, MAX_CHARS)).toEqual([]); + }); + + test("emits a sub-threshold tail as a single segment", () => { + const segments = planTailSegments([turn("a".repeat(100))], THRESHOLD, MAX_CHARS); + expect(segments).toEqual(["a".repeat(100)]); + }); + + test("chunks a large tail instead of truncating it", () => { + // 10 turns x 2500 chars = ~25K: far past a single 8K cap. The old + // single-selectTranscript path kept 8K and dropped the rest. + const turns = Array.from({ length: 10 }, (_, i) => turn(`${i} ` + "x".repeat(2500))); + const segments = planTailSegments(turns, THRESHOLD, MAX_CHARS); + + expect(segments.length).toBeGreaterThan(1); + const totalKept = segments.reduce((n, s) => n + s.length, 0); + const totalInput = turns.reduce((n, t) => n + t.text.length, 0); + // Only newline joins separate them, so essentially nothing is discarded. + expect(totalKept).toBeGreaterThanOrEqual(totalInput); + }); + + test("every full chunk except the last reaches the threshold", () => { + const turns = Array.from({ length: 10 }, (_, i) => turn(`${i} ` + "x".repeat(2500))); + const segments = planTailSegments(turns, THRESHOLD, MAX_CHARS); + for (const segment of segments.slice(0, -1)) { + expect(segment.length).toBeGreaterThanOrEqual(THRESHOLD); + } + }); + + test("emits a single over-threshold turn as its own segment, never split", () => { + const huge = "x".repeat(20000); + const segments = planTailSegments([turn(huge)], THRESHOLD, MAX_CHARS); + expect(segments).toEqual([huge]); + }); + + test("does not append an empty trailing segment when the tail divides evenly", () => { + const turns = [turn("x".repeat(8000)), turn("y".repeat(8000))]; + const segments = planTailSegments(turns, THRESHOLD, MAX_CHARS); + expect(segments).toEqual(["x".repeat(8000), "y".repeat(8000)]); + expect(segments.every((s) => s.length > 0)).toBe(true); + }); +}); diff --git a/tests/unit/flush-segments.test.ts b/tests/unit/flush-segments.test.ts new file mode 100644 index 0000000..6089d50 --- /dev/null +++ b/tests/unit/flush-segments.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, test } from "bun:test"; +import { flushSegments } from "../../src/hooks/flush-segments.js"; + +const META = { + project: "p", + branch: "b", + sessionId: "s", + baseSegment: 3, +}; + +describe("flushSegments", () => { + test("pushes every segment with sequential segment numbers", async () => { + const pushed: Array<{ text: string; segment: number }> = []; + const client = { + pushIngestQueue: async (text: string, meta: { segment: number }) => { + pushed.push({ text, segment: meta.segment }); + }, + }; + + const result = await flushSegments(client, ["one", "two"], META); + + expect(result).toEqual({ pushed: 2, failed: false }); + expect(pushed.map((p) => p.segment)).toEqual([3, 4]); + }); + + test("a mid-loop failure reports partial progress instead of throwing", async () => { + // SessionEnd runs once; an escaped exception here skipped both the drain + // spawn and cleanup, stranding the already-queued segments unsummarized. + let calls = 0; + const client = { + pushIngestQueue: async () => { + calls++; + if (calls > 1) throw new Error("connection dropped"); + }, + }; + + const result = await flushSegments(client, ["one", "two", "three"], META); + + expect(result).toEqual({ pushed: 1, failed: true }); + }); +}); diff --git a/tests/unit/hook-migration.test.ts b/tests/unit/hook-migration.test.ts deleted file mode 100644 index c2895fd..0000000 --- a/tests/unit/hook-migration.test.ts +++ /dev/null @@ -1,89 +0,0 @@ -import { describe, expect, test } from "bun:test"; -import { stripLegacyBetterdbHooks } from "../../src/hook-migration.js"; - -describe("stripLegacyBetterdbHooks", () => { - test("removes a betterdb Stop entry", () => { - const hooks = { - Stop: [ - { - hooks: [ - { type: "command", command: "/Users/x/.betterdb/bin/session-end" }, - ], - }, - ], - }; - const result = stripLegacyBetterdbHooks(hooks); - expect(result["Stop"]).toBeUndefined(); - }); - - test("preserves a third-party Stop hook", () => { - const hooks = { - Stop: [ - { - hooks: [ - { type: "command", command: "/Users/x/.betterdb/bin/session-end" }, - ], - }, - { hooks: [{ type: "command", command: "/other/tool/hook" }] }, - ], - }; - const result = stripLegacyBetterdbHooks(hooks); - expect(result["Stop"]).toHaveLength(1); - expect(JSON.stringify(result["Stop"])).toContain("/other/tool/hook"); - }); - - test("leaves unrelated events untouched", () => { - const hooks = { - PreCompact: [{ hooks: [{ type: "command", command: "/other/hook" }] }], - }; - const result = stripLegacyBetterdbHooks(hooks); - expect(result["PreCompact"]).toHaveLength(1); - }); - - test("is a no-op on settings with no Stop event", () => { - const hooks = { SessionStart: [{ hooks: [] }] }; - expect(stripLegacyBetterdbHooks(hooks)).toEqual(hooks); - }); - - test("strips a dev registration via an extra marker (path lacks 'betterdb')", () => { - const hooksDir = "/Users/x/dev/memory/src/hooks"; - const hooks = { - Stop: [ - { - hooks: [ - { - type: "command", - command: `bash -c 'bun run "${hooksDir}/session-end.ts"'`, - }, - ], - }, - { hooks: [{ type: "command", command: "/other/tool/hook" }] }, - ], - }; - const result = stripLegacyBetterdbHooks(hooks, ["betterdb", hooksDir]); - expect(result["Stop"]).toHaveLength(1); - expect(JSON.stringify(result["Stop"])).toContain("/other/tool/hook"); - }); - - test("ignores empty markers so they cannot match everything", () => { - const hooks = { - Stop: [{ hooks: [{ type: "command", command: "/other/tool/hook" }] }], - }; - const result = stripLegacyBetterdbHooks(hooks, [""]); - expect(result["Stop"]).toHaveLength(1); - }); - - test("does not mutate its input", () => { - const hooks = { - Stop: [ - { - hooks: [ - { type: "command", command: "/Users/x/.betterdb/bin/session-end" }, - ], - }, - ], - }; - stripLegacyBetterdbHooks(hooks); - expect(hooks["Stop"]).toHaveLength(1); - }); -}); diff --git a/tests/unit/hook-spec.test.ts b/tests/unit/hook-spec.test.ts new file mode 100644 index 0000000..602d6ae --- /dev/null +++ b/tests/unit/hook-spec.test.ts @@ -0,0 +1,96 @@ +import { describe, expect, test } from "bun:test"; +import { isOwnHookEntry, HOOK_SPECS, buildHookMap } from "../../src/hook-spec.js"; + +const BIN_DIR = "/Users/x/.betterdb/bin"; + +function entry(command: string): unknown { + return { hooks: [{ type: "command", command }] }; +} + +describe("isOwnHookEntry", () => { + test("recognizes an installed binary entry via the BIN_DIR marker", () => { + expect(isOwnHookEntry(entry(`${BIN_DIR}/session-end`), [BIN_DIR])).toBe(true); + }); + + test("recognizes a dev entry whose checkout path has no betterdb in it", () => { + // The gap this predicate exists to close: register-hooks.ts points at a + // checkout that need not be named betterdb, so a path-substring match + // missed it and the stale entry survived install. + const dev = entry(`bash -c 'bun run "/work/memory/src/hooks/session-end.ts"'`); + expect(isOwnHookEntry(dev, [BIN_DIR, "betterdb"])).toBe(true); + }); + + test("recognizes a legacy Stop entry pointing at session-end from any path", () => { + const legacy = entry(`bash -c 'bun run "/srv/clone/src/hooks/session-end.ts"'`); + expect(isOwnHookEntry(legacy, [BIN_DIR, "betterdb"])).toBe(true); + }); + + test("recognizes a compiled binary entry from a foreign checkout", () => { + // install-hooks.sh registers dist/hooks/ wrapped in a bash -c + // env loader — no .ts source name anywhere. A later installer passes its + // own markers, which need not cover the old checkout's path. + const bin = entry( + `bash -c "set -a; [ -f /work/memory/.env ] && . /work/memory/.env; set +a; /work/memory/dist/hooks/stop-checkpoint"`, + ); + expect(isOwnHookEntry(bin, [BIN_DIR, "betterdb"])).toBe(true); + }); + + test("recognizes every compiled hook binary this plugin registers", () => { + const map = buildHookMap((spec) => `/work/memory/dist/hooks/${spec.binary}`); + for (const entries of Object.values(map)) { + for (const e of entries) { + expect(isOwnHookEntry(e, [BIN_DIR, "betterdb"])).toBe(true); + } + } + }); + + test("does not claim a third party binary under its own dist/hooks", () => { + expect( + isOwnHookEntry(entry("/opt/tool/dist/hooks/lint"), [BIN_DIR, "betterdb"]), + ).toBe(false); + }); + + test("does not claim a third party's hook", () => { + expect(isOwnHookEntry(entry("/opt/other-tool/bin/their-hook"), [BIN_DIR])).toBe( + false, + ); + }); + + test("does not claim a third party hook merely near our install", () => { + expect(isOwnHookEntry(entry("/usr/local/bin/lint"), [BIN_DIR, "betterdb"])).toBe( + false, + ); + }); + + test("ignores empty markers rather than matching everything", () => { + expect(isOwnHookEntry(entry("/opt/other/hook"), [""])).toBe(false); + }); + + test("every registered hook carries an explicit timeout", () => { + // Without one, Claude Code waits 60s per hook — a wedged backend held + // every tool call hostage for minutes. Pre/post fire on each tool call + // and must be the tightest. + const map = buildHookMap((spec) => spec.binary); + for (const [event, entries] of Object.entries(map)) { + for (const entry of entries) { + for (const cmd of entry.hooks) { + expect(cmd.timeout).toBeGreaterThan(0); + if (event === "PreToolUse" || event === "PostToolUse") { + expect(cmd.timeout).toBeLessThanOrEqual(10); + } + } + } + } + }); + + test("recognizes every hook this plugin registers", () => { + const map = buildHookMap((spec) => `bash -c 'bun run "/work/memory/src/hooks/${spec.source}"'`); + for (const spec of HOOK_SPECS) { + const entries = map[spec.event] ?? []; + expect(entries.length).toBeGreaterThan(0); + for (const e of entries) { + expect(isOwnHookEntry(e, [BIN_DIR, "betterdb"])).toBe(true); + } + } + }); +}); diff --git a/tests/unit/locate-drain.test.ts b/tests/unit/locate-drain.test.ts new file mode 100644 index 0000000..ba59633 --- /dev/null +++ b/tests/unit/locate-drain.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, test } from "bun:test"; +import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterAll } from "bun:test"; +import { locateDrainCommand } from "../../src/hooks/locate-drain.js"; + +const root = mkdtempSync(join(tmpdir(), "locate-drain-")); + +function dir(name: string, files: string[] = []): string { + const d = join(root, name); + mkdirSync(d, { recursive: true }); + for (const f of files) { + writeFileSync(join(d, f), ""); + } + return d; +} + +afterAll(() => { + rmSync(root, { recursive: true, force: true }); +}); + +describe("locateDrainCommand", () => { + test("prefers the drain binary sitting next to the running executable", async () => { + const execDir = dir("exec-with-drain", ["drain"]); + const installBinDir = dir("bin-with-drain", ["drain"]); + const sourceDir = dir("src-with-drain", ["drain.ts"]); + + const cmd = await locateDrainCommand({ execDir, installBinDir, sourceDir }); + + expect(cmd).toEqual([join(execDir, "drain")]); + }); + + test("falls back to the install bin dir when the executable has no sibling", async () => { + const execDir = dir("exec-empty-1"); + const installBinDir = dir("bin-with-drain-2", ["drain"]); + const sourceDir = dir("src-with-drain-2", ["drain.ts"]); + + const cmd = await locateDrainCommand({ execDir, installBinDir, sourceDir }); + + expect(cmd).toEqual([join(installBinDir, "drain")]); + }); + + test("falls back to bun-running the source in the dev shape", async () => { + const execDir = dir("exec-empty-2"); + const installBinDir = dir("bin-empty-2"); + const sourceDir = dir("src-with-drain-3", ["drain.ts"]); + + const cmd = await locateDrainCommand({ execDir, installBinDir, sourceDir }); + + expect(cmd).toEqual(["bun", "run", join(sourceDir, "drain.ts")]); + }); + + test("returns null when no shape provides a drain", async () => { + // The compiled-binary trap this function exists to close: sourceDir is a + // bunfs virtual path inside the executable, so drain.ts never exists there. + const execDir = dir("exec-empty-3"); + const installBinDir = dir("bin-empty-3"); + const sourceDir = join(root, "does-not-exist"); + + const cmd = await locateDrainCommand({ execDir, installBinDir, sourceDir }); + + expect(cmd).toBeNull(); + }); +}); diff --git a/tests/unit/ollama-stream.test.ts b/tests/unit/ollama-stream.test.ts new file mode 100644 index 0000000..2ddb079 --- /dev/null +++ b/tests/unit/ollama-stream.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, test } from "bun:test"; +import { OllamaModelClient } from "../../src/client/providers/ollama.js"; + +const PRESET = { + embedModel: "mxbai-embed-large", + summarizeModel: "mistral:7b", + embedDim: 1024, +}; + +function chunked(pieces: string[]) { + return { + chat: async (request: { stream?: boolean }) => { + expect(request.stream).toBe(true); + async function* stream() { + for (const piece of pieces) { + yield { message: { content: piece } }; + } + } + return stream(); + }, + embed: async () => ({ embeddings: [[0]] }), + }; +} + +function timeoutOnce(pieces: string[]) { + let calls = 0; + const requests: { keep_alive?: string }[] = []; + return { + requests, + transport: { + chat: async (request: { stream?: boolean; keep_alive?: string }) => { + requests.push(request); + calls++; + if (calls === 1) { + throw new DOMException("The operation timed out.", "TimeoutError"); + } + async function* stream() { + for (const piece of pieces) { + yield { message: { content: piece } }; + } + } + return stream(); + }, + embed: async () => ({ embeddings: [[0]] }), + }, + }; +} + +describe("OllamaModelClient.summarize", () => { + test("retries once when the cold model load outlives the fetch timeout", async () => { + // A cold Ollama model load can exceed the client idle timeout; the load + // continues server-side, so a single retry lands on a warm model. + const json = JSON.stringify({ oneLineSummary: "Recovered" }); + const fake = timeoutOnce([json]); + + const client = new OllamaModelClient(PRESET, undefined, fake.transport); + const summary = await client.summarize("User: transcript"); + + expect(summary.oneLineSummary).toBe("Recovered"); + expect(fake.requests.length).toBe(2); + }); + + test("passes the configured keep_alive on every chat call", async () => { + // Default is a modest 5m: native Ollama reloads in seconds, so pinning + // the model in RAM for longer trades gigabytes for nothing. + const json = JSON.stringify({ oneLineSummary: "Recovered" }); + const fake = timeoutOnce([json]); + + const client = new OllamaModelClient(PRESET, undefined, fake.transport); + await client.summarize("User: transcript"); + + for (const request of fake.requests) { + expect(request.keep_alive).toBe("5m"); + } + }); + + test("assembles the summary from streamed chunks", async () => { + // Streaming keeps bytes flowing during generation so Bun's fetch idle + // timeout cannot fire while the model is still producing output. + const json = JSON.stringify({ oneLineSummary: "Did a thing" }); + const mid = Math.floor(json.length / 2); + const transport = chunked([json.slice(0, mid), json.slice(mid)]); + + const client = new OllamaModelClient(PRESET, undefined, transport); + const summary = await client.summarize("User: did we do a thing?"); + + expect(summary.oneLineSummary).toBe("Did a thing"); + }); +}); diff --git a/tests/unit/transcript-line.test.ts b/tests/unit/transcript-line.test.ts new file mode 100644 index 0000000..268ebe4 --- /dev/null +++ b/tests/unit/transcript-line.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, test } from "bun:test"; +import { parseTranscriptLine } from "../../src/memory/transcript.js"; + +describe("parseTranscriptLine", () => { + test("extracts a user turn from string content", () => { + const line = JSON.stringify({ type: "user", message: { content: "hello" } }); + expect(parseTranscriptLine(line)).toEqual([{ role: "user", text: "User: hello" }]); + }); + + test("extracts a user turn from block content", () => { + const line = JSON.stringify({ + type: "user", + message: { content: [{ type: "text", text: "hi there" }] }, + }); + expect(parseTranscriptLine(line)).toEqual([{ role: "user", text: "User: hi there" }]); + }); + + test("drops system-generated command messages", () => { + const line = JSON.stringify({ + type: "user", + message: { content: "/foo" }, + }); + expect(parseTranscriptLine(line)).toEqual([]); + }); + + test("caps an assistant turn at 2000 chars", () => { + const line = JSON.stringify({ + type: "assistant", + message: { content: "x".repeat(5000) }, + }); + const turns = parseTranscriptLine(line); + expect(turns).toHaveLength(1); + expect(turns[0]!.text).toBe("Assistant: " + "x".repeat(2000)); + }); + + test("extracts a tool turn by name", () => { + const line = JSON.stringify({ type: "tool_use", tool_name: "Edit" }); + expect(parseTranscriptLine(line)).toEqual([{ role: "tool", text: "Tool: Edit" }]); + }); + + test("returns [] for malformed JSON", () => { + expect(parseTranscriptLine("{not json")).toEqual([]); + }); + + test("returns [] for an unrecognized entry type", () => { + expect(parseTranscriptLine(JSON.stringify({ type: "summary" }))).toEqual([]); + }); +}); diff --git a/tests/unit/valkey-connect.test.ts b/tests/unit/valkey-connect.test.ts new file mode 100644 index 0000000..0f8f221 --- /dev/null +++ b/tests/unit/valkey-connect.test.ts @@ -0,0 +1,50 @@ +import { afterAll, describe, expect, test } from "bun:test"; +import { connectValkey } from "../../src/client/valkey.js"; + +// A listener that accepts connections and never speaks RESP — the same shape +// as Docker's proxy accepting for a backend that is not there. Hooks run on +// every tool call, so a connect against this must fail fast, not hang until +// the hook harness kills the process. +const blackhole = Bun.listen({ + hostname: "127.0.0.1", + port: 0, + socket: { + data() {}, + }, +}); + +afterAll(() => { + blackhole.stop(true); +}); + +describe("connectValkey against an unresponsive server", () => { + test("fails within the deadline instead of hanging", async () => { + const started = Date.now(); + + await expect( + connectValkey(`redis://127.0.0.1:${blackhole.port}`, 1000), + ).rejects.toThrow(/deadline exceeded/); + + expect(Date.now() - started).toBeLessThan(3000); + }, 30_000); + + test("does not leak an unhandled rejection from the losing connect", async () => { + // disconnect() rejects the still-pending connect() after the deadline + // already won the race; unabsorbed, that surfaces as an unhandled + // rejection a tick later. + const rejections: unknown[] = []; + const handler = (err: unknown) => { + rejections.push(err); + }; + process.on("unhandledRejection", handler); + try { + await connectValkey(`redis://127.0.0.1:${blackhole.port}`, 50).catch( + () => {}, + ); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(rejections).toEqual([]); + } finally { + process.off("unhandledRejection", handler); + } + }, 30_000); +});