diff --git a/.gitignore b/.gitignore index 241a389..10bd78c 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,9 @@ Thumbs.db .claude/ .gstack/ +# vibecop context optimization cache +.vibecop/ + # Environment .env .env.* diff --git a/CHANGELOG.md b/CHANGELOG.md index 068ecd2..5420f4e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,30 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.4.1] - 2026-04-07 + +### Added + +- **Context optimization** — intercepts Claude Code Read tool re-reads to reduce token consumption + - AST skeleton caching via `bun:sqlite` (WAL mode, zero-dep under bun) + - Smart-limits unchanged re-reads to 30 lines + skeleton injected via `additionalContext` + - `vibecop init --context` configures PreToolUse/PostToolUse/PostCompact hooks with conflict detection + - `vibecop context benchmark` shows projected savings for any project (runs under node, no bun required) + - `vibecop context stats` shows actual token savings after sessions (requires bun) +- **MCP tool: `vibecop_context_benchmark`** — benchmark context optimization potential through the MCP server +- **Shared AST utilities** — new `src/ast-utils.ts` with `findImports()`, `findFunctions()`, `findClasses()`, `findExports()`, `extractJsPackageName()` +- Build target: `bun run build:context` produces `dist/context.js` (26KB, bun runtime) + +### Changed + +- Refactored 4 detectors (god-component, god-function, mixed-concerns, undeclared-import) to use shared AST utilities — reduced code duplication, same behavior +- `god-function` removed 68 lines of local helper functions now shared via ast-utils + +### Internal + +- 42 new tests (610 total): ast-utils (28), skeleton (12), cache (15), integration (11), init --context (4) +- `.vibecop/` added to `.gitignore` (SQLite cache is project-local state) + ## [0.4.0] - 2026-04-06 ### Added diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1659764..d777569 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -22,21 +22,26 @@ src/ cli.ts # CLI entry point (commander) config.ts # Config loading + Zod validation engine.ts # File discovery, AST parsing, detector orchestration + ast-utils.ts # Shared AST traversal (findImports, findFunctions, findClasses, findExports) + context.ts # Context optimization entry point (hook handlers, requires bun) + context/ + benchmark.ts # Benchmark command (pure ast-grep, no bun:sqlite — runs under node) + cache.ts # SQLite persistence (bun:sqlite, WAL mode) + session.ts # File hashing, token estimation, extension checks + skeleton.ts # AST skeleton extraction from source code + stats.ts # Token savings reporting project.ts # Package manifest parsing (dependencies, devDependencies) types.ts # All shared types (Detector, Finding, DetectionContext, etc.) + init.ts # Agent tool detection + config generation (vibecop init) detectors/ index.ts # builtinDetectors array — register new detectors here - empty-error-handler.ts - god-function.ts - sql-injection.ts - ... # 22 detectors total + utils.ts # Detector-specific helpers (makeFinding, test function detection) + ... # 35 detectors across quality, security, correctness, testing formatters/ - index.ts # Formatter registry - text.ts # Default human-readable output - json.ts # JSON output - html.ts # HTML report - sarif.ts # SARIF for IDE/CI integration - github.ts # GitHub Actions annotations + index.ts # Formatter registry (text, json, html, sarif, github, agent, gcc) + mcp/ + index.ts # MCP server setup + tool registration + server.ts # Tool handlers (scan, check, explain, context_benchmark) action/ main.ts # GitHub Action entry point diff.ts # PR diff parsing @@ -46,11 +51,10 @@ src/ test/ detectors/ # One test file per detector formatters/ # Formatter tests + context/ # Context optimization tests (skeleton, cache, integration) + mcp/ # MCP server tests fixtures/ # Sample source files for testing - cli.test.ts - config.test.ts - engine.test.ts - project.test.ts + ast-utils.test.ts # Shared AST utility tests ``` ## How to Contribute @@ -140,10 +144,11 @@ Add the detector to the README detector table with its id, name, category, and s Run all checks locally: ```bash -bun run lint # Biome linter — fix all errors -bun run typecheck # TypeScript strict — no type errors -bun test # All tests must pass -bun run build # Build must succeed +bun run lint # Biome linter — fix all errors +bun run typecheck # TypeScript strict — no type errors +bun test # All tests must pass +bun run build # CLI build must succeed +bun run build:context # Context optimization build must succeed ``` ## Review Process diff --git a/README.md b/README.md index 9aacbc8..cc395bb 100644 --- a/README.md +++ b/README.md @@ -64,7 +64,26 @@ Agent writes code → vibecop hook fires → Findings? Agent fixes → Clean? Co } ``` -Three tools: `vibecop_scan`, `vibecop_check`, `vibecop_explain`. +Four tools: `vibecop_scan`, `vibecop_check`, `vibecop_explain`, `vibecop_context_benchmark`. + +## Context Optimization + +Reduce token consumption by ~35% on Read tool re-reads. When Claude Code reads a file it's already seen, vibecop intercepts the Read and serves a compact AST skeleton instead of the full file. Unchanged files get smart-limited to 30 lines + skeleton context. + +**Requires bun runtime** (uses `bun:sqlite` for zero-dependency caching). + +```bash +vibecop context benchmark # See projected savings for your project +vibecop init --context # Configure hooks (Claude Code only) +vibecop context stats # View actual token savings after sessions +``` + +How it works: +1. **First read** — full file passes through, skeleton is cached +2. **Re-read (unchanged)** — smart-limited to 30 lines + skeleton injected via `additionalContext` +3. **Re-read (changed)** — full file passes through with "file changed" note + +Skeletons include imports, function signatures, class outlines, and exports — enough for Claude to understand file structure without re-reading the full implementation. ## Benchmarks @@ -114,7 +133,7 @@ Catches: god functions, N+1 queries, unsafe shell exec, SQL injection, hardcoded - [x] **Phase 2.5**: Agent integration (7 tools), 6 LLM/agent detectors, `vibecop init` - [x] **Phase 3**: Test quality detectors, custom YAML rules (28 → 35) - [x] **Phase 3.5**: MCP server with scan/check/explain tools -- [ ] **Phase 4**: Context optimization (Read tool interception, AST skeleton caching) +- [x] **Phase 4**: Context optimization (Read tool interception, AST skeleton caching) - [ ] **Phase 5**: VS Code extension, cross-file analysis ## Links diff --git a/package.json b/package.json index 0bbe4cd..8d57f6a 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "vibecop", - "version": "0.4.0", + "version": "0.4.1", "description": "AI code quality toolkit — deterministic linter for the AI coding era", "type": "module", "bin": { @@ -8,6 +8,7 @@ }, "scripts": { "build": "bun build src/cli.ts --outdir dist --target node --external @ast-grep/napi --external @ast-grep/lang-python", + "build:context": "bun build src/context.ts --outdir dist --target bun --external @ast-grep/napi --external @ast-grep/lang-python", "build:action": "npx ncc build src/action/main.ts -o dist/action --source-map --license licenses.txt", "test": "bun test", "lint": "bunx biome check", @@ -17,6 +18,7 @@ }, "files": [ "dist/cli.js", + "dist/context.js", "README.md", "LICENSE" ], diff --git a/product-docs/src/content/docs/agent-integration/claude-code.mdx b/product-docs/src/content/docs/agent-integration/claude-code.mdx index 5b48b4b..e4b0599 100644 --- a/product-docs/src/content/docs/agent-integration/claude-code.mdx +++ b/product-docs/src/content/docs/agent-integration/claude-code.mdx @@ -68,6 +68,54 @@ This scans only files with uncommitted changes, making the hook fast even in lar If `.claude/settings.json` already exists, `vibecop init` will skip it to avoid overwriting your configuration. In that case, add the `PostToolUse` hook manually by merging the JSON above into your existing settings. +## Context Optimization (v0.4.1+) + +vibecop can intercept Read tool calls to reduce token consumption on re-reads. When Claude Code reads a file it has already seen in the current session, vibecop smart-limits the read to 30 lines and injects an AST skeleton (imports, function signatures, class outlines) via `additionalContext`. + +**Requires bun runtime** for SQLite caching. + +### Setup + +```bash +npx vibecop init --context +``` + +This adds three hooks to `.claude/settings.json`: + +- **PreToolUse Read** -- checks if the file was read before, smart-limits re-reads +- **PostToolUse Read** -- caches the AST skeleton after each read +- **PostCompact** -- prunes old session data + +### Benchmark First + +Before enabling, see projected savings for your project: + +```bash +npx vibecop context benchmark +``` + +This scans your codebase and shows per-file skeleton compression ratios and projected token savings. No bun required for the benchmark. + +### View Stats + +After working with context optimization enabled: + +```bash +npx vibecop context stats +``` + +Shows total reads, cache hits, hit rate, and tokens saved across sessions. + +### How It Works + +1. **First read** -- full file passes through, skeleton is cached in `.vibecop/.vibecop-context.db` +2. **Re-read (unchanged)** -- smart-limited to 30 lines + skeleton injected via `additionalContext` +3. **Re-read (changed)** -- full file passes through with "file changed" note + +### Conflict Detection + +`vibecop init --context` checks for existing PreToolUse Read hooks before adding its own. If another tool already uses `updatedInput` on Read, vibecop skips to avoid conflicts (only one hook can modify `updatedInput` per tool call). + ## Troubleshooting ### vibecop not found diff --git a/product-docs/src/content/docs/agent-integration/mcp-server.mdx b/product-docs/src/content/docs/agent-integration/mcp-server.mdx index 244df3b..39e38cb 100644 --- a/product-docs/src/content/docs/agent-integration/mcp-server.mdx +++ b/product-docs/src/content/docs/agent-integration/mcp-server.mdx @@ -3,7 +3,7 @@ title: MCP Server description: Use vibecop as an MCP server for Continue.dev, Amazon Q, Zed, and other MCP-compatible tools --- -vibecop includes a built-in MCP (Model Context Protocol) server that exposes scan, check, and explain tools over stdio transport. This enables integration with any MCP-compatible AI coding tool. +vibecop includes a built-in MCP (Model Context Protocol) server that exposes scan, check, explain, and context benchmark tools over stdio transport. This enables integration with any MCP-compatible AI coding tool. ## Starting the Server @@ -30,7 +30,7 @@ Configure your MCP client to connect to vibecop: ## Available Tools -The MCP server exposes three tools: +The MCP server exposes four tools: ### vibecop_scan @@ -72,6 +72,18 @@ Get detailed information about what a specific detector checks for. If the detector ID is not found, returns an error with the list of all available detector IDs. +### vibecop_context_benchmark + +Benchmark context optimization potential for a project. Shows per-file skeleton compression ratios and projected token savings at different re-read rates. Does not require bun. + +**Parameters:** + +| Parameter | Type | Required | Default | Description | +|-----------|------|:--------:|---------|-------------| +| `path` | string | No | Current working directory | Directory to benchmark | + +**Response:** JSON object with `totalFiles`, `totalTokens`, `topFiles` (top 10 by token count with skeleton compression), `projections` (savings at 20%/40%/60% re-read rates), and `enableCommand`. + ## Tool-Specific Setup ### Continue.dev @@ -128,9 +140,10 @@ The MCP server is built on the `@modelcontextprotocol/sdk` package and uses stdi MCP Client (Continue.dev / Amazon Q / Zed) ↕ JSON-RPC over stdio vibecop MCP Server - → vibecop_scan → engine.scan() - → vibecop_check → engine.checkFile() - → vibecop_explain → detector metadata lookup + → vibecop_scan → engine.scan() + → vibecop_check → engine.checkFile() + → vibecop_explain → detector metadata lookup + → vibecop_context_benchmark → benchmark() (token savings projection) ``` The server handles graceful shutdown on SIGINT and SIGTERM. All logging goes to stderr so it does not interfere with the MCP protocol on stdout. diff --git a/product-docs/src/content/docs/architecture.mdx b/product-docs/src/content/docs/architecture.mdx index 4e60a2e..1e749e8 100644 --- a/product-docs/src/content/docs/architecture.mdx +++ b/product-docs/src/content/docs/architecture.mdx @@ -8,10 +8,12 @@ description: How vibecop works internally — engine data flow, detector pattern ``` vibecop CLI (Commander.js) ├── Scan Engine — discovers files, loads AST, runs detectors, dedup by priority -├── MCP Server — stdio transport, 3 tools (scan, check, explain) +├── MCP Server — stdio transport, 4 tools (scan, check, explain, context_benchmark) +├── Context Optimization — Read tool interception, AST skeleton caching (bun:sqlite) ├── Init Wizard — auto-detects AI tools, generates hook/rule configs ├── Custom Rules Engine — loads .vibecop/rules/*.yaml, validates with Zod, runs via ast-grep ├── Config Loader (Zod) — validates .vibecop.yml, merges defaults, per-rule config +├── AST Utilities — shared findImports, findFunctions, findClasses, findExports ├── Detectors (35) — AST pattern matching via ast-grep (@ast-grep/napi) ├── Formatters (7) — text, json, html, sarif, github, agent, gcc output ├── Project Analyzer — parses package.json, requirements.txt, lockfiles @@ -40,9 +42,17 @@ vibecop/ │ │ ├── html.ts — Single-file HTML report │ │ ├── agent.ts — Token-efficient one-per-line for AI hooks │ │ └── gcc.ts — GCC-style for editor integration +│ ├── ast-utils.ts — Shared AST: findImports, findFunctions, findClasses, findExports +│ ├── context.ts — Context optimization entry (hook handlers, requires bun) +│ ├── context/ +│ │ ├── benchmark.ts — Benchmark command (pure ast-grep, runs under node) +│ │ ├── cache.ts — SQLite persistence (bun:sqlite, WAL mode) +│ │ ├── session.ts — File hashing, token estimation +│ │ ├── skeleton.ts — AST skeleton extraction +│ │ └── stats.ts — Token savings reporting │ ├── mcp/ │ │ ├── index.ts — MCP module entry, server creation + transport -│ │ └── server.ts — Tool handlers: scan, check, explain +│ │ └── server.ts — Tool handlers: scan, check, explain, context_benchmark │ ├── detectors/ │ │ ├── index.ts — Detector registry (all 35 built-in detectors) │ │ ├── utils.ts — makeFinding/makeLineFinding helpers @@ -163,9 +173,36 @@ The engine handles these error conditions explicitly: MCP Client (Continue.dev / Amazon Q / Zed) ↕ JSON-RPC over stdio vibecop MCP Server (@modelcontextprotocol/sdk) - → vibecop_scan → engine.scan() - → vibecop_check → engine.checkFile() - → vibecop_explain → detector metadata lookup + → vibecop_scan → engine.scan() + → vibecop_check → engine.checkFile() + → vibecop_explain → detector metadata lookup + → vibecop_context_benchmark → benchmark() (project token savings projection) ``` The MCP server reuses the same `scan()` and `checkFile()` functions as the CLI. Tool input schemas are defined with Zod. The server handles graceful shutdown on SIGINT/SIGTERM. + +## Context Optimization Architecture + +Context optimization intercepts Claude Code Read tool calls to reduce token consumption on re-reads. It has two build targets with different runtimes: + +| Build | Runtime | Contents | +|-------|---------|----------| +| `dist/cli.js` | node | CLI + MCP + benchmark (no SQLite) | +| `dist/context.js` | bun | Hook handlers + stats (bun:sqlite for caching) | + +``` +Read tool call + ├── PreToolUse fires → bun dist/context.js --pre + │ ├── First read? → passthrough (+ skeleton if cached) + │ ├── Re-read, unchanged? → smart-limit to 30 lines + skeleton + │ └── Re-read, changed? → passthrough + "file changed" note + │ + ├── Read tool executes (full or smart-limited) + │ + └── PostToolUse fires → bun dist/context.js --post + ├── Record session read + ├── Extract AST skeleton (imports, functions, classes, exports) + └── Cache skeleton + token counts in SQLite +``` + +Skeletons are extracted using the same `@ast-grep/napi` parser as the detectors, via shared utilities in `src/ast-utils.ts`. The SQLite database uses WAL mode for concurrent access and is stored at `.vibecop/.vibecop-context.db` (gitignored). diff --git a/product-docs/src/content/docs/configuration/cli-reference.mdx b/product-docs/src/content/docs/configuration/cli-reference.mdx index 61826dd..f213a5c 100644 --- a/product-docs/src/content/docs/configuration/cli-reference.mdx +++ b/product-docs/src/content/docs/configuration/cli-reference.mdx @@ -95,10 +95,16 @@ vibecop check src/auth.ts --format json Auto-detect AI coding tools and generate integration config files. ```bash -vibecop init +vibecop init [--context] ``` -No arguments or options. Detects tools by checking for: +**Options:** + +| Flag | Description | Default | +|------|-------------|---------| +| `--context` | Enable context optimization hooks (requires bun) | `false` | + +Detects tools by checking for: - `.claude/` directory (Claude Code) - `.cursor/` directory (Cursor) @@ -118,7 +124,34 @@ Start the MCP server on stdio transport. vibecop serve ``` -No arguments or options. The server exposes three tools (`vibecop_scan`, `vibecop_check`, `vibecop_explain`) via the Model Context Protocol. See [MCP Server](/vibecop/agent-integration/mcp-server/) for client configuration. +No arguments or options. The server exposes four tools (`vibecop_scan`, `vibecop_check`, `vibecop_explain`, `vibecop_context_benchmark`) via the Model Context Protocol. See [MCP Server](/vibecop/agent-integration/mcp-server/) for client configuration. + +### vibecop context + +Context optimization commands. + +```bash +vibecop context [benchmark|stats] +``` + +**Subcommands:** + +| Subcommand | Description | Requires bun | +|------------|-------------|:------------:| +| `benchmark` | Show projected token savings for the current project | No | +| `stats` | Show actual token savings from past sessions | Yes | + +**Examples:** + +```bash +# See projected savings (runs under node, no bun needed) +vibecop context benchmark + +# View actual savings after sessions (requires bun for SQLite) +vibecop context stats +``` + +The `benchmark` command scans all supported files, extracts AST skeletons, and projects token savings at 20%, 40%, and 60% re-read rates. The `stats` command reads from the SQLite database populated by the context optimization hooks. ### vibecop test-rules diff --git a/scripts/baseline-read-counter.ts b/scripts/baseline-read-counter.ts new file mode 100644 index 0000000..a1b4f47 --- /dev/null +++ b/scripts/baseline-read-counter.ts @@ -0,0 +1,101 @@ +#!/usr/bin/env bun +/** + * Baseline Read counter — logs read activity WITHOUT optimizing anything. + * Use this as a PreToolUse Read hook in a session WITHOUT context optimization + * to establish a baseline for comparison. + * + * Setup: Add to .claude/settings.json PreToolUse Read hook + * Stats: bun scripts/baseline-read-counter.ts --stats + * Reset: bun scripts/baseline-read-counter.ts --reset + * + * Always passes through (outputs {}) — zero impact on the session. + */ + +import { appendFileSync, existsSync, readFileSync, unlinkSync } from "node:fs"; +import { extname } from "node:path"; + +const LOG = "/tmp/vibecop-baseline-reads.jsonl"; +const SUPPORTED = new Set([".js", ".jsx", ".mjs", ".cjs", ".ts", ".tsx", ".py"]); + +if (process.argv.includes("--reset")) { + if (existsSync(LOG)) unlinkSync(LOG); + console.log("Baseline log cleared."); + process.exit(0); +} + +if (process.argv.includes("--stats")) { + if (!existsSync(LOG)) { + console.log("No baseline data. Run a session with the baseline hook first."); + process.exit(0); + } + + const lines = readFileSync(LOG, "utf-8").trim().split("\n").filter(Boolean); + const perFile: Record = {}; + let totalReads = 0; + let supportedReads = 0; + let reReads = 0; + + for (const line of lines) { + try { + const entry = JSON.parse(line); + const path = entry.path ?? ""; + perFile[path] = (perFile[path] ?? 0) + 1; + totalReads++; + if (entry.supported) supportedReads++; + } catch {} + } + + for (const count of Object.values(perFile)) { + if (count > 1) reReads += count - 1; + } + + const uniqueFiles = Object.keys(perFile).length; + const reReadPercent = totalReads > 0 ? ((reReads / totalReads) * 100).toFixed(1) : "0"; + + // Estimate token savings if context optimization were enabled + // Avg file ~300 lines, ~5100 tokens. Skeleton ~900 tokens. Savings ~4200 per re-read. + const estimatedSavings = reReads * 4200; + + console.log("Baseline Session Stats (no optimization)"); + console.log("========================================="); + console.log(`Total Read calls: ${totalReads}`); + console.log(`Unique files: ${uniqueFiles}`); + console.log(`Supported files: ${supportedReads} (${SUPPORTED.size} extensions)`); + console.log(`Re-reads: ${reReads} (${reReadPercent}% of total)`); + console.log(`Avg reads per file: ${(totalReads / Math.max(uniqueFiles, 1)).toFixed(1)}`); + console.log(); + console.log(`Estimated token savings with context optimization:`); + console.log(` ~${estimatedSavings.toLocaleString()} tokens (${reReads} re-reads x ~4,200 tokens each)`); + console.log(); + + // Show top re-read files + const sorted = Object.entries(perFile).sort((a, b) => b[1] - a[1]).slice(0, 10); + if (sorted.length > 0 && sorted[0][1] > 1) { + console.log("Most re-read files:"); + for (const [path, count] of sorted) { + if (count <= 1) break; + console.log(` ${count}x ${path}`); + } + } + + process.exit(0); +} + +// Hook mode — log and passthrough +try { + const stdin = readFileSync("/dev/stdin", "utf-8"); + const input = JSON.parse(stdin); + const path = input.tool_input?.file_path ?? ""; + const ext = extname(path); + + appendFileSync(LOG, JSON.stringify({ + ts: new Date().toISOString(), + path, + ext, + supported: SUPPORTED.has(ext), + session_id: input.session_id ?? null, + }) + "\n"); +} catch {} + +// Always passthrough +console.log("{}"); diff --git a/scripts/test-hook-dedup.ts b/scripts/test-hook-dedup.ts new file mode 100755 index 0000000..a97b5db --- /dev/null +++ b/scripts/test-hook-dedup.ts @@ -0,0 +1,103 @@ +#!/usr/bin/env bun +/** + * Empirical test for Claude Code PreToolUse hook behavior with Read tool dedup. + * + * Verifies: + * 1. Whether PreToolUse hooks fire BEFORE built-in read deduplication + * 2. The exact JSON format of hook stdin for Read tool calls + * 3. Whether updatedInput (limit) and additionalContext work as expected + * + * Usage: + * As hook: Configured via .claude/settings.json PreToolUse matcher "Read" + * Reset: bun scripts/test-hook-dedup.ts --reset + * Stats: bun scripts/test-hook-dedup.ts --stats + */ + +import { appendFileSync, existsSync, readFileSync, unlinkSync, writeFileSync } from "node:fs"; + +const LOG_FILE = "/tmp/vibecop-hook-test.jsonl"; +const READS_FILE = "/tmp/vibecop-hook-reads.json"; + +// ── CLI flags ─────────────────────────────────────────────────────────────── + +if (process.argv.includes("--reset")) { + for (const f of [LOG_FILE, READS_FILE]) { + if (existsSync(f)) unlinkSync(f); + } + console.log("Cleared tracking and log files."); + process.exit(0); +} + +if (process.argv.includes("--stats")) { + if (!existsSync(LOG_FILE)) { + console.log("No log file found. Nothing recorded yet."); + process.exit(0); + } + const lines = readFileSync(LOG_FILE, "utf-8").trim().split("\n").filter(Boolean); + const perFile: Record = {}; + for (const line of lines) { + try { + const entry = JSON.parse(line); + const path = entry.tool_input?.file_path ?? ""; + perFile[path] ??= { total: 0, rereads: 0 }; + perFile[path].total++; + if (entry.is_reread) perFile[path].rereads++; + } catch {} + } + console.log(`Total invocations: ${lines.length}`); + for (const [path, counts] of Object.entries(perFile)) { + console.log(` ${path} → ${counts.total} reads (${counts.rereads} re-reads)`); + } + process.exit(0); +} + +// ── Hook entrypoint ───────────────────────────────────────────────────────── + +try { + const stdin = readFileSync("/dev/stdin", "utf-8"); + const hook = JSON.parse(stdin) as { + tool_name: string; + tool_input: { file_path?: string; offset?: number; limit?: number }; + session_id?: string; + }; + + // Load previously-seen reads + let seen: Record = {}; + if (existsSync(READS_FILE)) { + try { seen = JSON.parse(readFileSync(READS_FILE, "utf-8")); } catch {} + } + + const filePath = hook.tool_input?.file_path ?? ""; + const isReread = filePath in seen; + + // Log every invocation + appendFileSync( + LOG_FILE, + JSON.stringify({ + ts: new Date().toISOString(), + tool_name: hook.tool_name, + tool_input: hook.tool_input, + session_id: hook.session_id, + is_reread: isReread, + seen_count: seen[filePath] ?? 0, + }) + "\n", + ); + + // Update tracking + seen[filePath] = (seen[filePath] ?? 0) + 1; + writeFileSync(READS_FILE, JSON.stringify(seen, null, 2)); + + if (!isReread) { + // First read — passthrough + console.log("{}"); + } else { + // Re-read — test updatedInput + additionalContext + console.log(JSON.stringify({ + updatedInput: { file_path: filePath, limit: 30 }, + additionalContext: `[vibecop-test] Re-read #${seen[filePath]} of ${filePath}. Skeleton would go here.`, + })); + } +} catch { + // Never block the agent + console.log("{}"); +} diff --git a/scripts/validate-context.ts b/scripts/validate-context.ts new file mode 100644 index 0000000..a18a7c3 --- /dev/null +++ b/scripts/validate-context.ts @@ -0,0 +1,223 @@ +#!/usr/bin/env bun +/** + * Validation script for context optimization. + * Tests latency, skeleton quality, and end-to-end hook flow. + * + * Usage: bun scripts/validate-context.ts + */ + +import { mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; + +const PASS = "\x1b[32m✓\x1b[0m"; +const FAIL = "\x1b[31m✗\x1b[0m"; +const BOLD = "\x1b[1m"; +const RESET = "\x1b[0m"; + +let passed = 0; +let failed = 0; + +function check(name: string, ok: boolean, detail?: string) { + if (ok) { + console.log(` ${PASS} ${name}${detail ? ` (${detail})` : ""}`); + passed++; + } else { + console.log(` ${FAIL} ${name}${detail ? ` — ${detail}` : ""}`); + failed++; + } +} + +// ── Test 1: Skeleton quality ──────────────────────────────────────────────── + +console.log(`\n${BOLD}1. Skeleton Quality${RESET}\n`); + +const { extractSkeleton } = await import("../src/context/skeleton.js"); + +// Test on a real file from this repo +const engineSource = readFileSync(join(import.meta.dir, "../src/engine.ts"), "utf-8"); +const engineSkeleton = extractSkeleton(engineSource, "typescript"); +const engineLines = engineSource.split("\n").length; +const skeletonLines = engineSkeleton.split("\n").length; +const compressionRatio = ((1 - engineSkeleton.length / engineSource.length) * 100).toFixed(1); + +check("Skeleton is shorter than source", engineSkeleton.length < engineSource.length, `${compressionRatio}% reduction`); +check("Skeleton has imports", engineSkeleton.includes("import")); +check("Skeleton has function names", engineSkeleton.includes("scan") || engineSkeleton.includes("discoverFiles")); +check("Skeleton preserves exports", engineSkeleton.includes("export") || engineSkeleton.includes("scan")); +check(`Source: ${engineLines} lines → Skeleton: ${skeletonLines} lines`, skeletonLines < engineLines / 2); + +// Show the skeleton for manual inspection +console.log(`\n ${BOLD}engine.ts skeleton (${skeletonLines} lines):${RESET}`); +for (const line of engineSkeleton.split("\n").slice(0, 15)) { + console.log(` ${line}`); +} +if (skeletonLines > 15) console.log(` ... (${skeletonLines - 15} more lines)`); + +// Test Python skeleton +const pySource = `import os\nfrom pathlib import Path\n\ndef process(path: str) -> str:\n with open(path) as f:\n content = f.read()\n lines = content.split("\\n")\n return "\\n".join(lines)\n\nclass FileHandler:\n def __init__(self, root: str):\n self.root = root\n def read(self, name: str) -> str:\n return (Path(self.root) / name).read_text()\n`; +const pySkeleton = extractSkeleton(pySource, "python"); +check("Python skeleton works", pySkeleton.includes("import os") && pySkeleton.includes("process") && pySkeleton.includes("FileHandler")); + +// Token savings estimate +const { estimateTokens } = await import("../src/context/session.js"); +const fullTokens = estimateTokens(engineSource); +const skeletonTokens = estimateTokens(engineSkeleton); +const savedTokens = fullTokens - skeletonTokens; +const savingsPercent = ((savedTokens / fullTokens) * 100).toFixed(1); +check(`Token savings: ${savedTokens} tokens saved (${savingsPercent}%)`, Number(savingsPercent) > 50); + +// ── Test 2: Latency benchmarks ────────────────────────────────────────────── + +console.log(`\n${BOLD}2. Latency Benchmarks${RESET}\n`); + +const testDir = join(tmpdir(), `vibecop-validate-${Date.now()}`); +mkdirSync(testDir, { recursive: true }); +mkdirSync(join(testDir, ".git")); + +const contextScript = join(import.meta.dir, "../src/context.ts"); +const testFile = join(testDir, "test-file.ts"); +writeFileSync(testFile, engineSource); + +async function timeHook(command: string, stdin: string): Promise<{ ms: number; stdout: string }> { + const start = performance.now(); + const proc = Bun.spawn(["bun", contextScript, command], { + stdin: new Blob([stdin]), + cwd: testDir, + stdout: "pipe", + stderr: "pipe", + }); + const stdout = await new Response(proc.stdout).text(); + await proc.exited; + const ms = performance.now() - start; + return { ms: Math.round(ms), stdout: stdout.trim() }; +} + +const hookInput = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: testFile }, + session_id: "validate-session", +}); + +// Cold start (first post-hook — cache miss, AST parse) +const postCold = await timeHook("--post", hookInput); +check(`Post-hook cold (cache miss): ${postCold.ms}ms`, postCold.ms < 500, "budget: <500ms"); + +// Pre-hook first read (no session history) +const preCold = await timeHook("--pre", hookInput); +check(`Pre-hook first read: ${preCold.ms}ms`, preCold.ms < 200, "budget: <200ms"); + +// Pre-hook re-read (cache hit, smart limit) +const preWarm = await timeHook("--pre", hookInput); +check(`Pre-hook re-read (cache hit): ${preWarm.ms}ms`, preWarm.ms < 200, "budget: <200ms"); + +// Parse the re-read response +try { + const response = JSON.parse(preWarm.stdout); + check("Re-read sets updatedInput.limit=30", response.updatedInput?.limit === 30); + check("Re-read has additionalContext", typeof response.additionalContext === "string" && response.additionalContext.length > 0); + check("additionalContext under 10K chars", (response.additionalContext?.length ?? 0) < 10000, `${response.additionalContext?.length ?? 0} chars`); +} catch { + check("Re-read response is valid JSON", false, preWarm.stdout.slice(0, 100)); +} + +// Post-hook warm (cache hit — skeleton already exists) +const postWarm = await timeHook("--post", hookInput); +check(`Post-hook warm (cache hit): ${postWarm.ms}ms`, postWarm.ms < 200, "budget: <200ms"); + +// Run 5 iterations to get a stable P90 +const latencies: number[] = []; +for (let i = 0; i < 5; i++) { + const { ms } = await timeHook("--pre", hookInput); + latencies.push(ms); +} +latencies.sort((a, b) => a - b); +const p90 = latencies[Math.floor(latencies.length * 0.9)]; +check(`P90 pre-hook latency: ${p90}ms`, p90 < 200, "budget: <200ms"); + +// ── Test 3: End-to-end hook protocol ──────────────────────────────────────── + +console.log(`\n${BOLD}3. End-to-End Hook Protocol${RESET}\n`); + +// Non-supported file → passthrough +const rustInput = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: "/foo/bar.rs" }, + session_id: "validate-session", +}); +const rustResult = await timeHook("--pre", rustInput); +check("Non-supported extension passes through", rustResult.stdout === "{}"); + +// Partial read (offset specified) → passthrough +const partialInput = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: testFile, offset: 10, limit: 20 }, + session_id: "validate-session", +}); +const partialResult = await timeHook("--pre", partialInput); +check("Partial read (offset>0) passes through", partialResult.stdout === "{}"); + +// No session_id → passthrough +const noSessionInput = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: testFile }, +}); +const noSessionResult = await timeHook("--pre", noSessionInput); +check("Missing session_id passes through", noSessionResult.stdout === "{}"); + +// Invalid JSON → passthrough (never crash) +const proc = Bun.spawn(["bun", contextScript, "--pre"], { + stdin: new Blob(["not json at all"]), + cwd: testDir, + stdout: "pipe", + stderr: "pipe", +}); +const badStdout = (await new Response(proc.stdout).text()).trim(); +const badExit = await proc.exited; +check("Invalid JSON → passthrough, no crash", badStdout === "{}" && badExit === 0); + +// File changed between reads → allows full read +const changedFile = join(testDir, "changing.ts"); +writeFileSync(changedFile, "const v1 = 1;\n"); +const changeInput1 = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: changedFile }, + session_id: "change-session", +}); +await timeHook("--post", changeInput1); // cache v1 +writeFileSync(changedFile, "const v2 = 2;\nexport function changed() { return v2; }\n"); +const changeResult = await timeHook("--pre", changeInput1); +const changeParsed = JSON.parse(changeResult.stdout); +check("Changed file: no limit override", changeParsed.updatedInput === undefined); +check("Changed file: notes the change", changeParsed.additionalContext?.includes("changed")); + +// Stats command +const statsProc = Bun.spawn(["bun", contextScript, "stats"], { + cwd: testDir, + stdout: "pipe", + stderr: "pipe", +}); +const statsOut = await new Response(statsProc.stdout).text(); +await statsProc.exited; +check("Stats command shows data", statsOut.includes("Sessions tracked") || statsOut.includes("Total reads")); + +// ── Cleanup ───────────────────────────────────────────────────────────────── + +try { rmSync(testDir, { recursive: true }); } catch {} + +// ── Summary ───────────────────────────────────────────────────────────────── + +console.log(`\n${BOLD}═══════════════════════════════════════${RESET}`); +console.log(`${BOLD}Results: ${passed} passed, ${failed} failed${RESET}`); + +if (failed > 0) { + console.log(`\n${FAIL} NOT READY TO SHIP — fix failures above`); + process.exit(1); +} else { + console.log(`\n${PASS} ALL CHECKS PASSED — ready for manual A/B testing`); + console.log(`\nNext: enable in a real session and compare token usage:`); + console.log(` vibecop init --context`); + console.log(` # work for 15 min, then:`); + console.log(` vibecop context stats`); +} +console.log(); diff --git a/src/ast-utils.ts b/src/ast-utils.ts new file mode 100644 index 0000000..6982169 --- /dev/null +++ b/src/ast-utils.ts @@ -0,0 +1,282 @@ +import type { SgNode } from "@ast-grep/napi"; +import type { Lang } from "./types.js"; + +export interface ImportInfo { + node: SgNode; + /** The module/package being imported (e.g., "react", "os.path") */ + source: string; + /** Full import statement text */ + text: string; +} + +export interface FunctionInfo { + node: SgNode; + /** Function name, or "" for unnamed arrow functions */ + name: string; + /** Number of parameters (excludes self/cls for Python) */ + params: number; + /** The function body AST node */ + body: SgNode | null; + /** AST node kind (e.g., "function_declaration", "arrow_function") */ + kind: string; +} + +export interface ClassInfo { + node: SgNode; + name: string; + methods: string[]; +} + +export interface ExportInfo { + node: SgNode; + name: string; + kind: "function" | "class" | "variable" | "type" | "default"; +} + +/** + * Extract the npm package name from an import specifier. + * `lodash/merge` → `lodash`, `@scope/pkg/sub` → `@scope/pkg` + */ +export function extractJsPackageName(specifier: string): string | null { + if (!specifier) return null; + if (specifier.startsWith("@")) { + const parts = specifier.split("/"); + return parts.length >= 2 ? `${parts[0]}/${parts[1]}` : null; + } + const slashIdx = specifier.indexOf("/"); + return slashIdx === -1 ? specifier : specifier.slice(0, slashIdx); +} + +/** Find all import declarations in a file. */ +export function findImports(root: SgNode, language: Lang): ImportInfo[] { + if (language === "python") return findPythonImports(root); + return findJsImports(root); +} + +function findJsImports(root: SgNode): ImportInfo[] { + const results: ImportInfo[] = []; + const nodes = root.findAll({ rule: { kind: "import_statement" } }); + for (const node of nodes) { + const sourceNode = node.children().find((ch) => ch.kind() === "string"); + if (!sourceNode) continue; + const source = sourceNode.text().slice(1, -1); + results.push({ node, source, text: node.text() }); + } + return results; +} + +function findPythonImports(root: SgNode): ImportInfo[] { + const results: ImportInfo[] = []; + + // `import X` or `import X.Y` + for (const node of root.findAll({ rule: { kind: "import_statement" } })) { + const nameNode = node.children().find( + (ch) => ch.kind() === "dotted_name" || ch.kind() === "aliased_import", + ); + if (!nameNode) continue; + let source: string; + if (nameNode.kind() === "aliased_import") { + const dotted = nameNode.children().find((ch) => ch.kind() === "dotted_name"); + source = dotted ? dotted.text() : nameNode.text(); + } else { + source = nameNode.text(); + } + results.push({ node, source, text: node.text() }); + } + + // `from X import Y` + for (const node of root.findAll({ rule: { kind: "import_from_statement" } })) { + const nameNode = node.children().find((ch) => ch.kind() === "dotted_name"); + if (!nameNode) continue; + results.push({ node, source: nameNode.text(), text: node.text() }); + } + + return results; +} + +/** Find all function/method declarations in a file. */ +export function findFunctions(root: SgNode, language: Lang): FunctionInfo[] { + if (language === "python") return findPythonFunctions(root); + return findJsFunctions(root); +} + +function findJsFunctions(root: SgNode): FunctionInfo[] { + const results: FunctionInfo[] = []; + + for (const kind of ["function_declaration", "method_definition", "arrow_function"] as const) { + for (const node of root.findAll({ rule: { kind } })) { + const name = getJsFunctionName(node); + const params = countJsParams(node); + const body = node.children().find((c) => c.kind() === "statement_block") ?? null; + results.push({ node, name, params, body, kind }); + } + } + + return results; +} + +function getJsFunctionName(node: SgNode): string { + const kind = node.kind(); + + if (kind === "function_declaration") { + return node.children().find((ch) => ch.kind() === "identifier")?.text() ?? ""; + } + + if (kind === "method_definition") { + const nameNode = node.children().find( + (ch) => ch.kind() === "property_identifier" || ch.kind() === "identifier", + ); + return nameNode?.text() ?? ""; + } + + if (kind === "arrow_function") { + const parent = node.parent(); + if (parent?.kind() === "variable_declarator") { + return parent.children().find((ch) => ch.kind() === "identifier")?.text() ?? ""; + } + if (parent?.kind() === "pair") { + const nameNode = parent.children().find( + (ch) => ch.kind() === "property_identifier" || ch.kind() === "string", + ); + return nameNode?.text() ?? ""; + } + return ""; + } + + return ""; +} + +function countJsParams(node: SgNode): number { + const params = node.children().find((ch) => ch.kind() === "formal_parameters"); + if (!params) return 0; + return params.children().filter((ch) => { + const k = ch.kind(); + return k !== "(" && k !== ")" && k !== ","; + }).length; +} + +function findPythonFunctions(root: SgNode): FunctionInfo[] { + const results: FunctionInfo[] = []; + + for (const node of root.findAll({ rule: { kind: "function_definition" } })) { + const nameNode = node.children().find((ch) => ch.kind() === "identifier"); + const name = nameNode?.text() ?? ""; + const body = node.children().find((ch) => ch.kind() === "block") ?? null; + const params = countPyParams(node); + results.push({ node, name, params, body, kind: "function_definition" }); + } + + return results; +} + +function countPyParams(node: SgNode): number { + const params = node.children().find((ch) => ch.kind() === "parameters"); + if (!params) return 0; + return params.children().filter((ch) => { + const k = ch.kind(); + if (k === "(" || k === ")" || k === ",") return false; + const text = ch.text().split(":")[0].split("=")[0].trim(); + return text !== "self" && text !== "cls"; + }).length; +} + +/** Find all class declarations in a file. */ +export function findClasses(root: SgNode, language: Lang): ClassInfo[] { + if (language === "python") return findPythonClasses(root); + return findJsClasses(root); +} + +function findJsClasses(root: SgNode): ClassInfo[] { + const results: ClassInfo[] = []; + + for (const node of root.findAll({ rule: { kind: "class_declaration" } })) { + const nameNode = node.children().find( + (ch) => ch.kind() === "type_identifier" || ch.kind() === "identifier", + ); + const name = nameNode?.text() ?? ""; + const methods: string[] = []; + const classBody = node.children().find((ch) => ch.kind() === "class_body"); + if (classBody) { + for (const member of classBody.findAll({ rule: { kind: "method_definition" } })) { + const methodName = member.children().find( + (ch) => ch.kind() === "property_identifier" || ch.kind() === "identifier", + ); + if (methodName) methods.push(methodName.text()); + } + } + results.push({ node, name, methods }); + } + + return results; +} + +function findPythonClasses(root: SgNode): ClassInfo[] { + const results: ClassInfo[] = []; + + for (const node of root.findAll({ rule: { kind: "class_definition" } })) { + const nameNode = node.children().find((ch) => ch.kind() === "identifier"); + const name = nameNode?.text() ?? ""; + const methods: string[] = []; + for (const method of node.findAll({ rule: { kind: "function_definition" } })) { + const methodName = method.children().find((ch) => ch.kind() === "identifier"); + if (methodName) methods.push(methodName.text()); + } + results.push({ node, name, methods }); + } + + return results; +} + +/** Find all export declarations (JS/TS only). */ +export function findExports(root: SgNode, language: Lang): ExportInfo[] { + if (language === "python") return []; + + const results: ExportInfo[] = []; + + for (const node of root.findAll({ rule: { kind: "export_statement" } })) { + const children = node.children(); + const hasDefault = children.some((ch) => ch.kind() === "default"); + + if (hasDefault) { + results.push({ node, name: "default", kind: "default" }); + continue; + } + + const funcDecl = children.find((ch) => ch.kind() === "function_declaration"); + if (funcDecl) { + const name = funcDecl.children().find((ch) => ch.kind() === "identifier")?.text() ?? ""; + results.push({ node, name, kind: "function" }); + continue; + } + + const classDecl = children.find((ch) => ch.kind() === "class_declaration"); + if (classDecl) { + const nameNode = classDecl.children().find( + (ch) => ch.kind() === "type_identifier" || ch.kind() === "identifier", + ); + results.push({ node, name: nameNode?.text() ?? "", kind: "class" }); + continue; + } + + const typeDecl = children.find((ch) => + ch.kind() === "type_alias_declaration" || ch.kind() === "interface_declaration", + ); + if (typeDecl) { + const name = typeDecl.children().find((ch) => ch.kind() === "type_identifier")?.text() ?? ""; + results.push({ node, name, kind: "type" }); + continue; + } + + const lexDecl = children.find((ch) => ch.kind() === "lexical_declaration"); + if (lexDecl) { + const declarator = lexDecl.children().find((ch) => ch.kind() === "variable_declarator"); + const name = declarator?.children().find((ch) => ch.kind() === "identifier")?.text() ?? ""; + results.push({ node, name, kind: "variable" }); + continue; + } + + results.push({ node, name: "", kind: "variable" }); + } + + return results; +} diff --git a/src/cli.ts b/src/cli.ts index 39f78ef..9812199 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -237,9 +237,10 @@ function main(): void { program .command("init") .description("Set up vibecop integration with AI coding tools") - .action(async () => { + .option("--context", "Enable context optimization (requires bun)", false) + .action(async (options: { context: boolean }) => { const { runInit } = await import("./init.js"); - await runInit(); + await runInit(undefined, { context: options.context }); }); program @@ -250,6 +251,47 @@ function main(): void { await startServer(); }); + program + .command("context") + .description("Context optimization — run as hook handler or view stats") + .argument("[mode]", "Mode: stats | benchmark (default: stats)") + .option("--pre", "PreToolUse handler (reads stdin)", false) + .option("--post", "PostToolUse handler (reads stdin)", false) + .option("--compact", "PostCompact handler (reads stdin)", false) + .action(async (mode: string | undefined, options: { pre: boolean; post: boolean; compact: boolean }) => { + // benchmark runs directly under node (no bun:sqlite needed) + if (mode === "benchmark") { + const { benchmark, formatBenchmark } = await import("./context/benchmark.js"); + console.log(formatBenchmark(benchmark(process.cwd()))); + return; + } + + // stats and hooks need bun:sqlite — shell out to bun + const args: string[] = []; + if (options.pre) args.push("--pre"); + else if (options.post) args.push("--post"); + else if (options.compact) args.push("--compact"); + else args.push(mode ?? "stats"); + + // context.js uses bun:sqlite — always requires bun runtime + const { execSync } = await import("node:child_process"); + const { existsSync } = await import("node:fs"); + const distPath = new URL("../dist/context.js", import.meta.url).pathname; + const srcPath = new URL("./context.ts", import.meta.url).pathname; + const contextPath = existsSync(distPath) ? distPath : srcPath; + + try { + execSync(`bun ${contextPath} ${args.join(" ")}`, { + stdio: "inherit", + cwd: process.cwd(), + }); + } catch (err: unknown) { + if (err && typeof err === "object" && "status" in err) { + process.exit((err as { status: number }).status); + } + } + }); + program .command("test-rules") .description("Validate custom rules against their inline examples") diff --git a/src/context.ts b/src/context.ts new file mode 100644 index 0000000..2150e27 --- /dev/null +++ b/src/context.ts @@ -0,0 +1,268 @@ +#!/usr/bin/env bun +/** + * Context optimization entry point. + * Called by Claude Code hooks to intercept Read tool calls and provide + * AST skeleton caching for token reduction on re-reads. + * + * Usage (from hooks): + * bun dist/context.js --pre # PreToolUse Read handler + * bun dist/context.js --post # PostToolUse Read handler + * bun dist/context.js --compact # PostCompact handler + * bun dist/context.js stats # Print stats + */ + +import { readFileSync, readdirSync } from "node:fs"; +import { dirname, extname, resolve } from "node:path"; +import { + openDb, + getSkeleton, + upsertSkeleton, + hasSessionRead, + recordSessionRead, + incrementStats, + pruneOldSessions, +} from "./context/cache.js"; +import { hashFile, estimateTokens, isSupportedExtension, isFileEligible } from "./context/session.js"; +import { extractSkeleton, languageForExtension } from "./context/skeleton.js"; +import { benchmark, formatBenchmark } from "./context/benchmark.js"; +import { printStats } from "./context/stats.js"; + +interface HookInput { + tool_name: string; + tool_input: { + file_path?: string; + offset?: number; + limit?: number; + }; + session_id?: string; +} + +interface PreHookResponse { + updatedInput?: { + file_path: string; + limit?: number; + offset?: number; + }; + additionalContext?: string; +} + +function findProjectRoot(): string { + let dir = process.cwd(); + while (true) { + try { + const entries = new Set(readdirSync(dir)); + if (entries.has(".vibecop.yml") || entries.has(".git")) return dir; + } catch {} + const parent = dirname(dir); + if (parent === dir) return process.cwd(); + dir = parent; + } +} + +// ── Handlers ──────────────────────────────────────────────────────────────── + +function handlePre(input: HookInput): void { + const filePath = input.tool_input?.file_path; + if (!filePath) { + console.log("{}"); + return; + } + + // Skip non-supported file types + if (!isSupportedExtension(filePath)) { + console.log("{}"); + return; + } + + // Skip partial reads (offset/limit already specified by user/agent) + if (input.tool_input.offset !== undefined && input.tool_input.offset > 0) { + console.log("{}"); + return; + } + + const sessionId = input.session_id; + if (!sessionId) { + console.log("{}"); + return; + } + + const projectRoot = findProjectRoot(); + + try { + const db = openDb(projectRoot); + try { + const resolvedPath = resolve(filePath); + const currentHash = hashFile(resolvedPath); + if (!currentHash) { + console.log("{}"); + return; + } + + const previousRead = hasSessionRead(db, sessionId, resolvedPath); + + if (!previousRead) { + // First read — check if we have a cached skeleton to inject as context + const cached = getSkeleton(db, resolvedPath, currentHash); + if (cached) { + const response: PreHookResponse = { + additionalContext: `[vibecop] File structure:\n${cached.skeleton}`, + }; + console.log(JSON.stringify(response)); + } else { + console.log("{}"); + } + return; + } + + // Re-read — check if file has changed + if (previousRead.hash === currentHash) { + // Unchanged: smart-limit to 30 lines + inject skeleton + const cached = getSkeleton(db, resolvedPath, currentHash); + const response: PreHookResponse = { + updatedInput: { + file_path: filePath, + limit: 30, + }, + }; + if (cached) { + response.additionalContext = + `[vibecop] File unchanged since last read. Structure:\n${cached.skeleton}`; + // Tokens saved = full file - (30 lines + skeleton) + // 30 lines ≈ 30 * 17 chars ≈ 510 chars ≈ 128 tokens + const limitedTokens = 128 + cached.skeletonTokens; + const saved = Math.max(0, cached.fullTokens - limitedTokens); + incrementStats(db, sessionId, { cacheHits: 1, tokensSaved: saved }); + } else { + incrementStats(db, sessionId, { cacheHits: 1 }); + } + console.log(JSON.stringify(response)); + } else { + // Changed: allow full read, note the change + const response: PreHookResponse = { + additionalContext: "[vibecop] File has changed since last read.", + }; + console.log(JSON.stringify(response)); + } + } finally { + db.close(); + } + } catch { + // Never block the agent + console.log("{}"); + } +} + +function handlePost(input: HookInput): void { + const filePath = input.tool_input?.file_path; + if (!filePath) return; + if (!isSupportedExtension(filePath)) return; + + const sessionId = input.session_id; + if (!sessionId) return; + + const projectRoot = findProjectRoot(); + + try { + const db = openDb(projectRoot); + try { + const resolvedPath = resolve(filePath); + + if (!isFileEligible(resolvedPath)) return; + + const currentHash = hashFile(resolvedPath); + if (!currentHash) return; + + // Record that this file was read in this session + recordSessionRead(db, sessionId, resolvedPath, currentHash); + incrementStats(db, sessionId, { totalReads: 1 }); + + // Extract and cache skeleton if not already cached with this hash + const existing = getSkeleton(db, resolvedPath, currentHash); + if (!existing) { + const ext = extname(resolvedPath); + const lang = languageForExtension(ext); + if (!lang) return; + + const source = readFileSync(resolvedPath, "utf-8"); + const skeleton = extractSkeleton(source, lang); + if (skeleton) { + const fullTokens = estimateTokens(source); + const skeletonTokens = estimateTokens(skeleton); + upsertSkeleton(db, resolvedPath, currentHash, skeleton, lang, fullTokens, skeletonTokens); + } + } + } finally { + db.close(); + } + } catch { + // Silent failure — never affect the agent + } +} + +function handleCompact(input: HookInput): void { + const sessionId = input.session_id; + if (!sessionId) return; + + const projectRoot = findProjectRoot(); + + try { + const db = openDb(projectRoot); + try { + // On compaction, prune old sessions + pruneOldSessions(db, 7); + } finally { + db.close(); + } + } catch { + // Silent + } +} + +function handleStats(sessionId?: string): void { + const projectRoot = findProjectRoot(); + try { + const db = openDb(projectRoot); + try { + printStats(db, sessionId); + } finally { + db.close(); + } + } catch (err) { + console.error(`Failed to read stats: ${err instanceof Error ? err.message : String(err)}`); + process.exit(1); + } +} + +// ── Main ──────────────────────────────────────────────────────────────────── + +const args = process.argv.slice(2); +const command = args[0]; + +if (command === "stats") { + handleStats(args[1]); +} else if (command === "benchmark") { + console.log(formatBenchmark(benchmark(findProjectRoot()))); +} else if (command === "--pre" || command === "--post" || command === "--compact") { + try { + const stdin = readFileSync("/dev/stdin", "utf-8"); + const input = JSON.parse(stdin) as HookInput; + + switch (command) { + case "--pre": + handlePre(input); + break; + case "--post": + handlePost(input); + break; + case "--compact": + handleCompact(input); + break; + } + } catch { + // Never block the agent — output passthrough on any error + if (command === "--pre") console.log("{}"); + } +} else { + console.error("Usage: vibecop context [--pre|--post|--compact|stats]"); + process.exit(1); +} diff --git a/src/context/benchmark.ts b/src/context/benchmark.ts new file mode 100644 index 0000000..212a5ab --- /dev/null +++ b/src/context/benchmark.ts @@ -0,0 +1,125 @@ +/** + * Context optimization benchmark. Zero bun:sqlite dependency. + * Works under both node and bun. Can be imported by MCP server and CLI directly. + */ + +import { readdirSync, readFileSync, statSync } from "node:fs"; +import { extname, join, relative, resolve } from "node:path"; +import { extractSkeleton, languageForExtension } from "./skeleton.js"; +import { estimateTokens, isSupportedExtension } from "./session.js"; + +export interface FileMetrics { + path: string; + fullTokens: number; + skeletonTokens: number; + reductionPercent: number; +} + +export interface BenchmarkResult { + files: FileMetrics[]; + totalFiles: number; + totalTokens: number; + projections: Array<{ + rereadPercent: number; + tokensSaved: number; + percentOfTotal: number; + }>; +} + +function walkSupported(dir: string, root: string, results: FileMetrics[]): void { + let entries; + try { entries = readdirSync(dir, { withFileTypes: true }); } catch { return; } + + for (const entry of entries) { + if (entry.name.startsWith(".")) continue; + if (entry.name === "node_modules" || entry.name === "dist" || entry.name === "build") continue; + + const fullPath = join(dir, entry.name); + if (entry.isDirectory()) { + walkSupported(fullPath, root, results); + } else if (entry.isFile() && isSupportedExtension(fullPath)) { + try { + const stat = statSync(fullPath); + if (stat.size > 500_000 || stat.size === 0) continue; + + const source = readFileSync(fullPath, "utf-8"); + const lang = languageForExtension(extname(fullPath)); + if (!lang) continue; + + const fullTokens = estimateTokens(source); + const skeleton = extractSkeleton(source, lang); + const skeletonTokens = estimateTokens(skeleton); + const reductionPercent = fullTokens > 0 + ? Math.round((1 - (skeletonTokens + 128) / fullTokens) * 100) + : 0; + + results.push({ + path: relative(root, fullPath), + fullTokens, + skeletonTokens, + reductionPercent: Math.max(0, reductionPercent), + }); + } catch {} + } + } +} + +/** Run benchmark on a project. Returns structured data (no console output). */ +export function benchmark(projectRoot: string): BenchmarkResult { + const root = resolve(projectRoot); + const files: FileMetrics[] = []; + walkSupported(root, root, files); + + const totalTokens = files.reduce((sum, f) => sum + f.fullTokens, 0); + const sorted = [...files].sort((a, b) => b.fullTokens - a.fullTokens); + + const projections = [20, 40, 60].map((rereadPercent) => { + const rereadCount = Math.round(files.length * rereadPercent / 100); + const rereadFiles = sorted.slice(0, rereadCount); + const tokensSaved = rereadFiles.reduce((sum, f) => { + const limited = 128 + f.skeletonTokens; + return sum + Math.max(0, f.fullTokens - limited); + }, 0); + return { + rereadPercent, + tokensSaved, + percentOfTotal: totalTokens > 0 ? Math.round(tokensSaved / totalTokens * 100) : 0, + }; + }); + + return { files: sorted, totalFiles: files.length, totalTokens, projections }; +} + +/** Format benchmark result as human-readable text. */ +export function formatBenchmark(result: BenchmarkResult): string { + if (result.totalFiles === 0) { + return "No supported files found (.js, .ts, .tsx, .py)."; + } + + const lines: string[] = []; + lines.push("vibecop context benchmark"); + lines.push("═════════════════════════"); + lines.push(""); + lines.push(`Files: ${result.totalFiles} supported`); + lines.push(`Total tokens: ~${result.totalTokens.toLocaleString()}`); + lines.push(""); + + const top = result.files.slice(0, 10); + lines.push("Largest files (most savings potential):"); + const maxPath = Math.max(...top.map(f => f.path.length), 10); + for (const f of top) { + const pathPad = f.path.padEnd(maxPath); + lines.push(` ${pathPad} ${f.fullTokens.toLocaleString().padStart(6)} tokens → skeleton: ${f.skeletonTokens.toLocaleString().padStart(5)} (${f.reductionPercent}% reduction)`); + } + lines.push(""); + + lines.push("Projected savings per session:"); + for (const p of result.projections) { + lines.push(` ${p.rereadPercent}% re-read rate: ~${p.tokensSaved.toLocaleString()} tokens saved (${p.percentOfTotal}% of total Read usage)`); + } + + lines.push(""); + lines.push("To enable: vibecop init --context"); + + return lines.join("\n"); +} diff --git a/src/context/cache.ts b/src/context/cache.ts new file mode 100644 index 0000000..4bf58b9 --- /dev/null +++ b/src/context/cache.ts @@ -0,0 +1,174 @@ +import { Database } from "bun:sqlite"; +import { existsSync, mkdirSync } from "node:fs"; +import { dirname, join } from "node:path"; + +const DB_FILENAME = ".vibecop-context.db"; + +/** Get the database path for a project root. */ +export function getDbPath(projectRoot: string): string { + return join(projectRoot, ".vibecop", DB_FILENAME); +} + +/** Open or create the SQLite database with WAL mode. */ +export function openDb(projectRoot: string): Database { + const dbPath = getDbPath(projectRoot); + const dir = dirname(dbPath); + if (!existsSync(dir)) mkdirSync(dir, { recursive: true }); + + const db = new Database(dbPath); + db.run("PRAGMA journal_mode = WAL"); + db.run("PRAGMA busy_timeout = 1000"); + initSchema(db); + return db; +} + +function initSchema(db: Database): void { + db.run(` + CREATE TABLE IF NOT EXISTS skeletons ( + path TEXT PRIMARY KEY, + hash TEXT NOT NULL, + skeleton TEXT NOT NULL, + language TEXT NOT NULL, + full_tokens INTEGER NOT NULL DEFAULT 0, + skeleton_tokens INTEGER NOT NULL DEFAULT 0, + updated_at INTEGER NOT NULL DEFAULT (unixepoch()) + ) + `); + + db.run(` + CREATE TABLE IF NOT EXISTS session_reads ( + session_id TEXT NOT NULL, + path TEXT NOT NULL, + hash TEXT NOT NULL, + read_at INTEGER NOT NULL DEFAULT (unixepoch()), + PRIMARY KEY (session_id, path) + ) + `); + + db.run(` + CREATE TABLE IF NOT EXISTS stats ( + session_id TEXT NOT NULL, + total_reads INTEGER NOT NULL DEFAULT 0, + cache_hits INTEGER NOT NULL DEFAULT 0, + tokens_saved INTEGER NOT NULL DEFAULT 0, + updated_at INTEGER NOT NULL DEFAULT (unixepoch()), + PRIMARY KEY (session_id) + ) + `); +} + +// ── Skeleton operations ───────────────────────────────────────────────────── + +export interface SkeletonRecord { + skeleton: string; + fullTokens: number; + skeletonTokens: number; +} + +export function getSkeleton(db: Database, path: string, hash: string): SkeletonRecord | null { + const row = db.query<{ skeleton: string; full_tokens: number; skeleton_tokens: number }, [string, string]>( + "SELECT skeleton, full_tokens, skeleton_tokens FROM skeletons WHERE path = ? AND hash = ?", + ).get(path, hash); + if (!row) return null; + return { skeleton: row.skeleton, fullTokens: row.full_tokens, skeletonTokens: row.skeleton_tokens }; +} + +export function upsertSkeleton( + db: Database, + path: string, + hash: string, + skeleton: string, + language: string, + fullTokens: number, + skeletonTokens: number, +): void { + db.run( + `INSERT INTO skeletons (path, hash, skeleton, language, full_tokens, skeleton_tokens, updated_at) + VALUES (?, ?, ?, ?, ?, ?, unixepoch()) + ON CONFLICT(path) DO UPDATE SET hash=excluded.hash, skeleton=excluded.skeleton, language=excluded.language, full_tokens=excluded.full_tokens, skeleton_tokens=excluded.skeleton_tokens, updated_at=excluded.updated_at`, + [path, hash, skeleton, language, fullTokens, skeletonTokens], + ); +} + +// ── Session read tracking ─────────────────────────────────────────────────── + +export function hasSessionRead(db: Database, sessionId: string, path: string): { hash: string } | null { + const row = db.query<{ hash: string }, [string, string]>( + "SELECT hash FROM session_reads WHERE session_id = ? AND path = ?", + ).get(sessionId, path); + return row ?? null; +} + +export function recordSessionRead(db: Database, sessionId: string, path: string, hash: string): void { + db.run( + `INSERT INTO session_reads (session_id, path, hash, read_at) + VALUES (?, ?, ?, unixepoch()) + ON CONFLICT(session_id, path) DO UPDATE SET hash=excluded.hash, read_at=excluded.read_at`, + [sessionId, path, hash], + ); +} + +export function clearSession(db: Database, sessionId: string): void { + db.run("DELETE FROM session_reads WHERE session_id = ?", [sessionId]); +} + +// ── Stats ─────────────────────────────────────────────────────────────────── + +export function incrementStats( + db: Database, + sessionId: string, + fields: { totalReads?: number; cacheHits?: number; tokensSaved?: number }, +): void { + db.run( + `INSERT INTO stats (session_id, total_reads, cache_hits, tokens_saved, updated_at) + VALUES (?, ?, ?, ?, unixepoch()) + ON CONFLICT(session_id) DO UPDATE SET + total_reads = stats.total_reads + excluded.total_reads, + cache_hits = stats.cache_hits + excluded.cache_hits, + tokens_saved = stats.tokens_saved + excluded.tokens_saved, + updated_at = excluded.updated_at`, + [sessionId, fields.totalReads ?? 0, fields.cacheHits ?? 0, fields.tokensSaved ?? 0], + ); +} + +export interface SessionStats { + sessionId: string; + totalReads: number; + cacheHits: number; + tokensSaved: number; +} + +export function getSessionStats(db: Database, sessionId: string): SessionStats | null { + const row = db.query< + { session_id: string; total_reads: number; cache_hits: number; tokens_saved: number }, + [string] + >("SELECT * FROM stats WHERE session_id = ?").get(sessionId); + if (!row) return null; + return { + sessionId: row.session_id, + totalReads: row.total_reads, + cacheHits: row.cache_hits, + tokensSaved: row.tokens_saved, + }; +} + +export function getAllStats(db: Database): SessionStats[] { + const rows = db.query< + { session_id: string; total_reads: number; cache_hits: number; tokens_saved: number }, + [] + >("SELECT * FROM stats ORDER BY updated_at DESC").all(); + return rows.map((r) => ({ + sessionId: r.session_id, + totalReads: r.total_reads, + cacheHits: r.cache_hits, + tokensSaved: r.tokens_saved, + })); +} + +/** Prune old session data older than N days. */ +export function pruneOldSessions(db: Database, days: number = 7): number { + const cutoff = Math.floor(Date.now() / 1000) - days * 86400; + const readsResult = db.run("DELETE FROM session_reads WHERE read_at < ?", [cutoff]); + const statsResult = db.run("DELETE FROM stats WHERE updated_at < ?", [cutoff]); + return (readsResult.changes ?? 0) + (statsResult.changes ?? 0); +} diff --git a/src/context/session.ts b/src/context/session.ts new file mode 100644 index 0000000..75dd9ce --- /dev/null +++ b/src/context/session.ts @@ -0,0 +1,43 @@ +import { createHash } from "node:crypto"; +import { readFileSync, statSync } from "node:fs"; + +/** Compute SHA-256 hash of file content. */ +export function hashFile(filePath: string): string | null { + try { + const content = readFileSync(filePath); + return createHash("sha256").update(content).digest("hex"); + } catch { + return null; + } +} + +/** Compute SHA-256 hash of a string. */ +export function hashString(content: string): string { + return createHash("sha256").update(content).digest("hex"); +} + +/** Estimate token count from string length (rough ~4 chars per token). */ +export function estimateTokens(text: string): number { + return Math.ceil(text.length / 4); +} + +/** Check if a file extension is supported for context optimization. */ +const SUPPORTED_EXTENSIONS = new Set([".js", ".jsx", ".mjs", ".cjs", ".ts", ".tsx", ".py"]); + +export function isSupportedExtension(filePath: string): boolean { + const lastDot = filePath.lastIndexOf("."); + if (lastDot === -1) return false; + return SUPPORTED_EXTENSIONS.has(filePath.slice(lastDot)); +} + +/** Check if file exists and is not too large for skeleton extraction. */ +const MAX_FILE_SIZE = 500_000; // 500KB + +export function isFileEligible(filePath: string): boolean { + try { + const stat = statSync(filePath); + return stat.isFile() && stat.size <= MAX_FILE_SIZE; + } catch { + return false; + } +} diff --git a/src/context/skeleton.ts b/src/context/skeleton.ts new file mode 100644 index 0000000..8e392a5 --- /dev/null +++ b/src/context/skeleton.ts @@ -0,0 +1,139 @@ +import { parse, Lang as SgLang, registerDynamicLanguage } from "@ast-grep/napi"; +import { createRequire } from "node:module"; +import { findImports, findFunctions, findClasses, findExports } from "../ast-utils.js"; +import type { Lang } from "../types.js"; + +const EXTENSION_TO_LANG: Record = { + ".js": "javascript", + ".jsx": "javascript", + ".mjs": "javascript", + ".cjs": "javascript", + ".ts": "typescript", + ".tsx": "tsx", + ".py": "python", +}; + +const LANG_TO_SG: Record = { + javascript: SgLang.JavaScript, + typescript: SgLang.TypeScript, + tsx: SgLang.Tsx, + python: "python", +}; + +let pythonRegistered = false; +function ensurePython(): void { + if (pythonRegistered) return; + try { + const req = createRequire(import.meta.url); + const pythonLang = req("@ast-grep/lang-python") as { + libraryPath: string; + extensions: string[]; + languageSymbol?: string; + expandoChar?: string; + }; + registerDynamicLanguage({ python: pythonLang }); + pythonRegistered = true; + } catch { + // Python support unavailable + } +} + +export function languageForExtension(ext: string): Lang | null { + return EXTENSION_TO_LANG[ext] ?? null; +} + +/** + * Extract a structural skeleton from source code. + * Returns a compact representation of imports, function signatures, + * class outlines, and exports — without implementation bodies. + * + * Used by context optimization to provide Claude with file structure + * on re-reads instead of full file content. + */ +export function extractSkeleton(source: string, language: Lang): string { + if (language === "python") ensurePython(); + + const sgLang = LANG_TO_SG[language]; + let root; + try { + root = parse(sgLang as SgLang, source).root(); + } catch { + return fallbackSkeleton(source); + } + + const lines: string[] = []; + + // Imports + const imports = findImports(root, language); + if (imports.length > 0) { + for (const imp of imports) { + lines.push(imp.text); + } + lines.push(""); + } + + // Classes with method signatures + const classes = findClasses(root, language); + for (const cls of classes) { + if (language === "python") { + lines.push(`class ${cls.name}:`); + } else { + lines.push(`class ${cls.name} {`); + } + for (const method of cls.methods) { + lines.push(` ${method}(...)`); + } + if (language !== "python") lines.push("}"); + lines.push(""); + } + + // Standalone functions (not inside classes) + const functions = findFunctions(root, language); + const classMethodNodes = new Set( + classes.flatMap((cls) => + cls.node.findAll({ rule: { kind: language === "python" ? "function_definition" : "method_definition" } }), + ).map((n) => n.range().start.line), + ); + + for (const fn of functions) { + // Skip class methods (already shown above) + if (classMethodNodes.has(fn.node.range().start.line)) continue; + + const range = fn.node.range(); + const sig = source.split("\n").slice(range.start.line, range.start.line + 1)[0]?.trim() ?? ""; + if (sig) lines.push(sig.replace(/\{[\s\S]*$/, "{ ... }")); + } + + if (functions.length > 0 && classes.length === 0) lines.push(""); + + // Exports (JS/TS only) + const exports = findExports(root, language); + for (const exp of exports) { + if (exp.kind === "function" || exp.kind === "class") continue; // already shown + lines.push(`export ${exp.kind === "default" ? "default" : ""} ${exp.name}`.trim()); + } + + const skeleton = lines.join("\n").trim(); + return skeleton || fallbackSkeleton(source); +} + +/** Regex-based fallback when AST parsing fails. */ +function fallbackSkeleton(source: string): string { + const lines = source.split("\n"); + const skeleton: string[] = []; + for (const line of lines) { + const trimmed = line.trim(); + if ( + trimmed.startsWith("import ") || + trimmed.startsWith("from ") || + trimmed.startsWith("export ") || + trimmed.startsWith("class ") || + /^(async\s+)?function\s/.test(trimmed) || + /^(const|let|var)\s+\w+\s*=\s*(async\s+)?\(/.test(trimmed) || + /^def\s/.test(trimmed) + ) { + skeleton.push(trimmed); + } + } + return skeleton.join("\n"); +} diff --git a/src/context/stats.ts b/src/context/stats.ts new file mode 100644 index 0000000..7fc4788 --- /dev/null +++ b/src/context/stats.ts @@ -0,0 +1,62 @@ +import type { Database } from "bun:sqlite"; +import { getAllStats, getSessionStats, type SessionStats } from "./cache.js"; + +export function formatStats(stats: SessionStats[]): string { + if (stats.length === 0) return "No context optimization data recorded yet.\nRun `vibecop context benchmark` to see projected savings for this project."; + + const lines: string[] = []; + + let totalReads = 0; + let totalHits = 0; + let totalSaved = 0; + + for (const s of stats) { + totalReads += s.totalReads; + totalHits += s.cacheHits; + totalSaved += s.tokensSaved; + } + + const hitRate = totalReads > 0 ? ((totalHits / totalReads) * 100).toFixed(1) : "0.0"; + + lines.push("vibecop context optimization"); + lines.push("════════════════════════════"); + lines.push(""); + lines.push(`Sessions: ${stats.length}`); + lines.push(`Total reads: ${totalReads}`); + lines.push(`Cache hits: ${totalHits} (${hitRate}% of reads)`); + lines.push(`Tokens saved: ~${totalSaved.toLocaleString()}`); + + if (totalSaved > 0 && totalHits > 0) { + const avgSavedPerHit = Math.round(totalSaved / totalHits); + lines.push(`Avg savings: ~${avgSavedPerHit.toLocaleString()} tokens per re-read`); + } + + lines.push(""); + + if (stats.length <= 10 && stats.length > 0) { + lines.push("Per-session:"); + for (const s of stats) { + const rate = s.totalReads > 0 + ? ((s.cacheHits / s.totalReads) * 100).toFixed(0) + : "0"; + lines.push( + ` ${s.sessionId.slice(0, 8)}… ${s.totalReads} reads, ${s.cacheHits} hits (${rate}%), ~${s.tokensSaved.toLocaleString()} saved`, + ); + } + } + + return lines.join("\n"); +} + +export function printStats(db: Database, sessionId?: string): void { + if (sessionId) { + const stats = getSessionStats(db, sessionId); + if (stats) { + console.log(formatStats([stats])); + } else { + console.log(`No stats found for session ${sessionId}`); + } + } else { + console.log(formatStats(getAllStats(db))); + } +} diff --git a/src/detectors/god-component.ts b/src/detectors/god-component.ts index b5e0380..9be9695 100644 --- a/src/detectors/god-component.ts +++ b/src/detectors/god-component.ts @@ -1,3 +1,4 @@ +import { findImports } from "../ast-utils.js"; import type { Detector, DetectionContext, Finding } from "../types.js"; import { makeLineFinding } from "./utils.js"; @@ -73,8 +74,7 @@ function detect(ctx: DetectionContext): Finding[] { // --- Count imports via AST --- const root = ctx.root.root(); - const importStatements = root.findAll({ rule: { kind: "import_statement" } }); - const importCount = importStatements.length; + const importCount = findImports(root, ctx.file.language).length; // --- Count total lines --- const lineCount = source.split("\n").length; diff --git a/src/detectors/god-function.ts b/src/detectors/god-function.ts index 213086d..fc282b3 100644 --- a/src/detectors/god-function.ts +++ b/src/detectors/god-function.ts @@ -1,4 +1,5 @@ import type { SgNode } from "@ast-grep/napi"; +import { findFunctions } from "../ast-utils.js"; import type { Detector, DetectionContext, Finding } from "../types.js"; import { makeLineFinding } from "./utils.js"; @@ -74,74 +75,6 @@ function countComplexity(node: SgNode, isPython: boolean): number { return count; } -/** Count formal parameters for a JS/TS function node */ -function countJsParams(funcNode: SgNode): number { - const params = funcNode.children().find( - ch => ch.kind() === "formal_parameters", - ); - if (!params) return 0; - - return params.children().filter(ch => { - const k = ch.kind(); - return k !== "(" && k !== ")" && k !== ","; - }).length; -} - -/** Count parameters for a Python function node */ -function countPyParams(funcNode: SgNode): number { - const params = funcNode.children().find( - ch => ch.kind() === "parameters", - ); - if (!params) return 0; - - const paramNodes = params.children().filter(ch => { - const k = ch.kind(); - return k !== "(" && k !== ")" && k !== ","; - }); - - // Exclude `self` and `cls` as they are implicit - return paramNodes.filter(ch => { - const text = ch.text().split(":")[0].split("=")[0].trim(); - return text !== "self" && text !== "cls"; - }).length; -} - -/** Get function name from a JS/TS function node or its parent context */ -function getJsFunctionName(node: SgNode): string { - const kind = node.kind(); - - if (kind === "function_declaration") { - const nameNode = node.children().find(ch => ch.kind() === "identifier"); - return nameNode?.text() ?? ""; - } - - if (kind === "method_definition") { - const nameNode = node.children().find( - ch => ch.kind() === "property_identifier" || ch.kind() === "identifier", - ); - return nameNode?.text() ?? ""; - } - - if (kind === "arrow_function") { - // Walk up to find variable_declarator parent - const parent = node.parent(); - if (parent?.kind() === "variable_declarator") { - const nameNode = parent.children().find(ch => ch.kind() === "identifier"); - return nameNode?.text() ?? ""; - } - // Property assignment: { key: () => {} } - if (parent?.kind() === "pair") { - const nameNode = parent.children().find( - ch => ch.kind() === "property_identifier" || ch.kind() === "string", - ); - return nameNode?.text() ?? ""; - } - return ""; - } - - return ""; -} - function buildFinding( ctx: DetectionContext, m: FunctionMetrics, @@ -168,17 +101,15 @@ function detectJavaScript(ctx: DetectionContext): Finding[] { const maxComplexity = (ctx.config as Record)?.maxComplexity as number ?? 15; const maxParams = (ctx.config as Record)?.maxParams as number ?? 5; - const funcKinds = ["function_declaration", "method_definition", "arrow_function"]; + const allFunctions = findFunctions(root, ctx.file.language); - for (const kind of funcKinds) { - const nodes = root.findAll({ rule: { kind } }); - - for (const node of nodes) { + for (const fn of allFunctions) { + const node = fn.node; const range = node.range(); const lines = range.end.line - range.start.line + 1; // For arrow functions, skip short inline callbacks - if (kind === "arrow_function") { + if (fn.kind === "arrow_function") { if (lines <= 10) continue; // Only flag if assigned to a variable or property const parent = node.parent(); @@ -189,9 +120,9 @@ function detectJavaScript(ctx: DetectionContext): Finding[] { } } - const name = getJsFunctionName(node); + const name = fn.name; const complexity = 1 + countComplexity(node, false); - const params = countJsParams(node); + const params = fn.params; const linesExceeded = lines > maxLines; const complexityExceeded = complexity > maxComplexity; @@ -220,7 +151,6 @@ function detectJavaScript(ctx: DetectionContext): Finding[] { }; findings.push(buildFinding(ctx, metrics, severity)); - } } return findings; @@ -234,15 +164,15 @@ function detectPython(ctx: DetectionContext): Finding[] { const maxComplexity = (ctx.config as Record)?.maxComplexity as number ?? 15; const maxParams = (ctx.config as Record)?.maxParams as number ?? 5; - const funcNodes = root.findAll({ rule: { kind: "function_definition" } }); + const allFunctions = findFunctions(root, ctx.file.language); - for (const node of funcNodes) { + for (const fn of allFunctions) { + const node = fn.node; const range = node.range(); const lines = range.end.line - range.start.line + 1; - const nameNode = node.children().find(ch => ch.kind() === "identifier"); - const name = nameNode?.text() ?? ""; + const name = fn.name; const complexity = 1 + countComplexity(node, true); - const params = countPyParams(node); + const params = fn.params; const linesExceeded = lines > maxLines; const complexityExceeded = complexity > maxComplexity; diff --git a/src/detectors/mixed-concerns.ts b/src/detectors/mixed-concerns.ts index 5eac456..2195f30 100644 --- a/src/detectors/mixed-concerns.ts +++ b/src/detectors/mixed-concerns.ts @@ -1,3 +1,4 @@ +import { findImports } from "../ast-utils.js"; import type { Detector, DetectionContext, Finding } from "../types.js"; import { makeLineFinding } from "./utils.js"; @@ -35,7 +36,7 @@ function detect(ctx: DetectionContext): Finding[] { if (ctx.source.includes('"use server"') || ctx.source.includes("'use server'")) return findings; const root = ctx.root.root(); - const imports = root.findAll({ rule: { kind: "import_statement" } }); + const imports = findImports(root, ctx.file.language); let hasUIImport = false; let hasDBImport = false; @@ -44,10 +45,7 @@ function detect(ctx: DetectionContext): Finding[] { let dbImportName = ""; for (const imp of imports) { - const sourceNode = imp.children().find(ch => ch.kind() === "string"); - if (!sourceNode) continue; - - const specifier = sourceNode.text().slice(1, -1); + const specifier = imp.source; const pkgName = specifier.startsWith("@") ? specifier.split("/").slice(0, 2).join("/") : specifier.split("/")[0]; diff --git a/src/detectors/undeclared-import.ts b/src/detectors/undeclared-import.ts index a4e0d47..b7f423c 100644 --- a/src/detectors/undeclared-import.ts +++ b/src/detectors/undeclared-import.ts @@ -1,5 +1,6 @@ import { existsSync, readFileSync, readdirSync } from "node:fs"; import { basename, dirname, join } from "node:path"; +import { extractJsPackageName } from "../ast-utils.js"; import type { Detector, DetectionContext, Finding } from "../types.js"; import { makeFinding } from "./utils.js"; @@ -68,30 +69,6 @@ const PYTHON_BUILTINS = new Set([ "formatter", "rlcompleter", ]); -/** - * Extract the package name from a JS/TS import specifier. - * - `lodash` -> `lodash` - * - `lodash/merge` -> `lodash` - * - `@scope/pkg` -> `@scope/pkg` - * - `@scope/pkg/sub` -> `@scope/pkg` - */ -function extractJsPackageName(specifier: string): string | null { - if (!specifier) return null; - - // Scoped package - if (specifier.startsWith("@")) { - const parts = specifier.split("/"); - if (parts.length >= 2) { - return `${parts[0]}/${parts[1]}`; - } - return null; - } - - // Unscoped: take first segment - const slashIdx = specifier.indexOf("/"); - return slashIdx === -1 ? specifier : specifier.slice(0, slashIdx); -} - function isRelativeImport(specifier: string): boolean { return specifier.startsWith("./") || specifier.startsWith("../") || specifier === "." || specifier === ".."; } diff --git a/src/init.ts b/src/init.ts index ef7a599..dbb9671 100644 --- a/src/init.ts +++ b/src/init.ts @@ -2,6 +2,29 @@ import { execSync } from "node:child_process"; import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; import { join } from "node:path"; +/** Resolve the absolute path to dist/context.js from the vibecop package. */ +function resolveContextScript(): string { + // When installed via npm: this file is at /dist/cli.js or /src/init.ts + // dist/context.js sits next to dist/cli.js + const fromDist = new URL("../dist/context.js", import.meta.url); + try { + const { existsSync } = require("node:fs") as typeof import("node:fs"); + if (existsSync(fromDist)) return fromDist.pathname; + } catch {} + // Fallback: try relative to cwd (local dev) + const { resolve } = require("node:path") as typeof import("node:path"); + return resolve("dist/context.js"); +} + +function contextCommands(): { pre: string; post: string; compact: string } { + const script = resolveContextScript(); + return { + pre: `bun ${script} --pre`, + post: `bun ${script} --post`, + compact: `bun ${script} --compact`, + }; +} + interface DetectedTool { name: string; detected: boolean; @@ -270,13 +293,138 @@ function padEnd(str: string, len: number): string { return str + " ".repeat(Math.max(0, len - str.length)); } -export async function runInit(cwd?: string): Promise { +function isBunAvailable(): boolean { + try { + execSync("bun --version", { stdio: "pipe" }); + return true; + } catch { + return false; + } +} + +function generateContextHooks(root: string): GeneratedFile[] { + const generated: GeneratedFile[] = []; + + // Only Claude Code supports the hooks needed for context optimization + if (!existsSync(join(root, ".claude"))) { + console.log(" Context optimization requires Claude Code (.claude/ directory)."); + return generated; + } + + const settingsPath = join(root, ".claude", "settings.json"); + let settings: Record = {}; + + if (existsSync(settingsPath)) { + try { + settings = JSON.parse(readFileSync(settingsPath, "utf-8")); + } catch { + console.log(" Warning: Could not parse .claude/settings.json"); + return generated; + } + } + + const hooks = (settings.hooks ?? {}) as Record; + + // Check for existing context hooks to avoid conflicts + const preHooks = (hooks.PreToolUse ?? []) as Array<{ matcher?: string }>; + const hasExistingReadHook = preHooks.some( + (h) => h.matcher && /\bRead\b/.test(h.matcher), + ); + + if (hasExistingReadHook) { + console.log(" Warning: Existing PreToolUse Read hook detected."); + console.log(" Context optimization uses updatedInput which is single-consumer."); + console.log(" Skipping to avoid conflicts. See docs/agent-integration.md."); + generated.push({ + path: ".claude/settings.json", + description: "context hooks skipped (existing Read hook)", + }); + return generated; + } + + // Resolve absolute paths to context.js + const cmds = contextCommands(); + + // Add context hooks + hooks.PreToolUse = [ + ...(hooks.PreToolUse ?? []), + { + matcher: "Read", + hooks: [{ type: "command", command: cmds.pre }], + }, + ]; + + hooks.PostToolUse = [ + ...(hooks.PostToolUse ?? []), + { + matcher: "Read", + hooks: [{ type: "command", command: cmds.post }], + }, + ]; + + // PostCompact is a session-level event, no matcher + const postCompact = (hooks.PostCompact ?? []) as unknown[]; + postCompact.push({ + hooks: [{ type: "command", command: cmds.compact }], + }); + hooks.PostCompact = postCompact; + + settings.hooks = hooks; + mkdirSync(join(root, ".claude"), { recursive: true }); + writeFileSync(settingsPath, `${JSON.stringify(settings, null, 2)}\n`); + + generated.push({ + path: ".claude/settings.json", + description: "context optimization hooks (Pre/Post Read, PostCompact)", + }); + + return generated; +} + +export interface InitOptions { + context?: boolean; +} + +export async function runInit(cwd?: string, options?: InitOptions): Promise { const root = cwd ?? process.cwd(); + const enableContext = options?.context ?? false; console.log(""); console.log(" vibecop — agent integration setup"); console.log(""); + if (enableContext) { + // Context optimization mode + if (!isBunAvailable()) { + console.log(" Error: Context optimization requires the bun runtime."); + console.log(" Install bun: https://bun.sh"); + console.log(""); + return; + } + + console.log(" Setting up context optimization..."); + console.log(""); + + const generated = generateContextHooks(root); + + if (generated.length > 0) { + const maxPath = Math.max(...generated.map((g) => g.path.length)); + console.log(" Generated:"); + for (const file of generated) { + console.log( + ` ${padEnd(file.path, maxPath)} — ${file.description}`, + ); + } + console.log(""); + } + + console.log(" Context optimization configured."); + console.log(" Run 'vibecop context stats' to see token savings."); + console.log(""); + return; + } + + // Standard init — detect tools and generate configs const tools = detectTools(root); const anyDetected = tools.some((t) => t.detected); diff --git a/src/mcp/index.ts b/src/mcp/index.ts index a334683..dbdc46e 100644 --- a/src/mcp/index.ts +++ b/src/mcp/index.ts @@ -8,6 +8,8 @@ import { handleCheck, explainInputSchema, handleExplain, + contextBenchmarkInputSchema, + handleContextBenchmark, } from "./server.js"; /** Read version from package.json */ @@ -49,6 +51,12 @@ export function createServer(): McpServer { inputSchema: explainInputSchema, }, handleExplain); + server.registerTool("vibecop_context_benchmark", { + description: + "Benchmark context optimization potential for a project. Shows per-file skeleton compression ratios and projected token savings at different re-read rates. No bun required.", + inputSchema: contextBenchmarkInputSchema, + }, handleContextBenchmark); + return server; } diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 9e054f6..0d1099a 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1,6 +1,7 @@ import { z } from "zod"; import { scan, checkFile } from "../engine.js"; import { builtinDetectors } from "../detectors/index.js"; +import { benchmark } from "../context/benchmark.js"; import type { ScanResult } from "../types.js"; /** Format a ScanResult as a JSON text content block for MCP */ @@ -117,3 +118,56 @@ export async function handleExplain(args: { detector_id: string }) { ], }; } + +/** Input schema for vibecop_context_benchmark */ +export const contextBenchmarkInputSchema = { + path: z.string().optional().describe("Directory to benchmark. Defaults to current working directory."), +}; + +/** Handler for vibecop_context_benchmark tool */ +export async function handleContextBenchmark(args: { path?: string }) { + try { + const result = benchmark(args.path ?? "."); + if (result.totalFiles === 0) { + return { + content: [{ type: "text" as const, text: "No supported files found (.js, .ts, .tsx, .py)." }], + }; + } + + const top10 = result.files.slice(0, 10).map((f) => ({ + file: f.path, + tokens: f.fullTokens, + skeletonTokens: f.skeletonTokens, + reduction: `${f.reductionPercent}%`, + })); + + return { + content: [ + { + type: "text" as const, + text: JSON.stringify( + { + totalFiles: result.totalFiles, + totalTokens: result.totalTokens, + topFiles: top10, + projections: result.projections.map((p) => ({ + rereadRate: `${p.rereadPercent}%`, + tokensSaved: p.tokensSaved, + percentOfTotal: `${p.percentOfTotal}%`, + })), + enableCommand: "vibecop init --context", + }, + null, + 2, + ), + }, + ], + }; + } catch (err: unknown) { + const message = err instanceof Error ? err.message : String(err); + return { + content: [{ type: "text" as const, text: `Error running benchmark: ${message}` }], + isError: true, + }; + } +} diff --git a/test/ast-utils.test.ts b/test/ast-utils.test.ts new file mode 100644 index 0000000..336f18d --- /dev/null +++ b/test/ast-utils.test.ts @@ -0,0 +1,287 @@ +import { describe, expect, test } from "bun:test"; +import { parse, Lang, registerDynamicLanguage } from "@ast-grep/napi"; +import { createRequire } from "node:module"; +import { + extractJsPackageName, + findImports, + findFunctions, + findClasses, + findExports, +} from "../src/ast-utils.js"; + +// Register Python for tests +try { + const req = createRequire(import.meta.url); + const pythonLang = req("@ast-grep/lang-python") as { + libraryPath: string; + extensions: string[]; + languageSymbol?: string; + expandoChar?: string; + }; + registerDynamicLanguage({ python: pythonLang }); +} catch { + // Python support may not be available +} + +function parseJs(source: string) { + return parse(Lang.JavaScript, source).root(); +} + +function parseTs(source: string) { + return parse(Lang.TypeScript, source).root(); +} + +function parseTsx(source: string) { + return parse(Lang.Tsx, source).root(); +} + +function parsePy(source: string) { + return parse("python" as Lang, source).root(); +} + +// ── extractJsPackageName ──────────────────────────────────────────────────── + +describe("extractJsPackageName", () => { + test("unscoped package", () => { + expect(extractJsPackageName("lodash")).toBe("lodash"); + }); + + test("unscoped package with subpath", () => { + expect(extractJsPackageName("lodash/merge")).toBe("lodash"); + }); + + test("scoped package", () => { + expect(extractJsPackageName("@scope/pkg")).toBe("@scope/pkg"); + }); + + test("scoped package with subpath", () => { + expect(extractJsPackageName("@scope/pkg/sub")).toBe("@scope/pkg"); + }); + + test("empty string returns null", () => { + expect(extractJsPackageName("")).toBeNull(); + }); + + test("bare scope returns null", () => { + expect(extractJsPackageName("@scope")).toBeNull(); + }); +}); + +// ── findImports ───────────────────────────────────────────────────────────── + +describe("findImports", () => { + describe("JavaScript/TypeScript", () => { + test("finds ES module imports", () => { + const root = parseTs(` + import React from "react"; + import { useState } from "react"; + import * as path from "node:path"; + `); + const imports = findImports(root, "typescript"); + expect(imports).toHaveLength(3); + expect(imports[0].source).toBe("react"); + expect(imports[1].source).toBe("react"); + expect(imports[2].source).toBe("node:path"); + }); + + test("returns full import text", () => { + const root = parseTs(`import { foo } from "bar";`); + const imports = findImports(root, "typescript"); + expect(imports[0].text).toContain("bar"); + }); + + test("returns empty array for no imports", () => { + const root = parseTs("const x = 1;"); + expect(findImports(root, "typescript")).toHaveLength(0); + }); + }); + + describe("Python", () => { + test("finds import statements", () => { + const root = parsePy("import os\nimport sys"); + const imports = findImports(root, "python"); + expect(imports).toHaveLength(2); + expect(imports[0].source).toBe("os"); + expect(imports[1].source).toBe("sys"); + }); + + test("finds from-import statements", () => { + const root = parsePy("from os.path import join\nfrom collections import defaultdict"); + const imports = findImports(root, "python"); + expect(imports).toHaveLength(2); + expect(imports[0].source).toBe("os.path"); + expect(imports[1].source).toBe("collections"); + }); + + test("handles aliased imports", () => { + const root = parsePy("import numpy as np"); + const imports = findImports(root, "python"); + expect(imports).toHaveLength(1); + // Source should be the original module name + expect(imports[0].source).toBe("numpy"); + }); + }); +}); + +// ── findFunctions ─────────────────────────────────────────────────────────── + +describe("findFunctions", () => { + describe("JavaScript/TypeScript", () => { + test("finds function declarations", () => { + const root = parseTs("function foo(a: string, b: number) { return a; }"); + const fns = findFunctions(root, "typescript"); + expect(fns).toHaveLength(1); + expect(fns[0].name).toBe("foo"); + expect(fns[0].params).toBe(2); + expect(fns[0].kind).toBe("function_declaration"); + expect(fns[0].body).not.toBeNull(); + }); + + test("finds arrow functions assigned to variables", () => { + const root = parseTs("const add = (a: number, b: number) => { return a + b; };"); + const fns = findFunctions(root, "typescript"); + const arrow = fns.find((f) => f.kind === "arrow_function"); + expect(arrow).toBeDefined(); + expect(arrow!.name).toBe("add"); + expect(arrow!.params).toBe(2); + }); + + test("finds method definitions", () => { + const root = parseTs(` + class Foo { + bar(x: number) { return x; } + static baz() {} + } + `); + const fns = findFunctions(root, "typescript"); + const methods = fns.filter((f) => f.kind === "method_definition"); + expect(methods.length).toBeGreaterThanOrEqual(2); + const names = methods.map((m) => m.name); + expect(names).toContain("bar"); + expect(names).toContain("baz"); + }); + + test("unnamed arrow function gets ", () => { + const root = parseJs("[1, 2].map((x) => { return x * 2; });"); + const fns = findFunctions(root, "javascript"); + const arrow = fns.find((f) => f.kind === "arrow_function"); + expect(arrow?.name).toBe(""); + }); + }); + + describe("Python", () => { + test("finds function definitions", () => { + const root = parsePy("def greet(name, greeting='hello'):\n return f'{greeting} {name}'"); + const fns = findFunctions(root, "python"); + expect(fns).toHaveLength(1); + expect(fns[0].name).toBe("greet"); + expect(fns[0].params).toBe(2); + expect(fns[0].kind).toBe("function_definition"); + }); + + test("excludes self/cls from param count", () => { + const root = parsePy("class Foo:\n def bar(self, x):\n pass"); + const fns = findFunctions(root, "python"); + const bar = fns.find((f) => f.name === "bar"); + expect(bar).toBeDefined(); + expect(bar!.params).toBe(1); // self excluded + }); + + test("finds nested functions", () => { + const root = parsePy("def outer():\n def inner():\n pass\n inner()"); + const fns = findFunctions(root, "python"); + expect(fns.length).toBeGreaterThanOrEqual(2); + const names = fns.map((f) => f.name); + expect(names).toContain("outer"); + expect(names).toContain("inner"); + }); + }); +}); + +// ── findClasses ───────────────────────────────────────────────────────────── + +describe("findClasses", () => { + describe("JavaScript/TypeScript", () => { + test("finds class with methods", () => { + const root = parseTs(` + class UserService { + constructor() {} + getUser(id: string) { return id; } + deleteUser(id: string) {} + } + `); + const classes = findClasses(root, "typescript"); + expect(classes).toHaveLength(1); + expect(classes[0].name).toBe("UserService"); + expect(classes[0].methods).toContain("constructor"); + expect(classes[0].methods).toContain("getUser"); + expect(classes[0].methods).toContain("deleteUser"); + }); + + test("returns empty array for no classes", () => { + const root = parseTs("const x = 1;"); + expect(findClasses(root, "typescript")).toHaveLength(0); + }); + }); + + describe("Python", () => { + test("finds class with methods", () => { + const root = parsePy(`class Dog: + def __init__(self, name): + self.name = name + def bark(self): + print("woof")`); + const classes = findClasses(root, "python"); + expect(classes).toHaveLength(1); + expect(classes[0].name).toBe("Dog"); + expect(classes[0].methods).toContain("__init__"); + expect(classes[0].methods).toContain("bark"); + }); + }); +}); + +// ── findExports ───────────────────────────────────────────────────────────── + +describe("findExports", () => { + test("finds exported function", () => { + const root = parseTs("export function hello() {}"); + const exports = findExports(root, "typescript"); + expect(exports).toHaveLength(1); + expect(exports[0].name).toBe("hello"); + expect(exports[0].kind).toBe("function"); + }); + + test("finds exported class", () => { + const root = parseTs("export class Foo {}"); + const exports = findExports(root, "typescript"); + expect(exports).toHaveLength(1); + expect(exports[0].name).toBe("Foo"); + expect(exports[0].kind).toBe("class"); + }); + + test("finds exported const", () => { + const root = parseTs("export const bar = 42;"); + const exports = findExports(root, "typescript"); + expect(exports).toHaveLength(1); + expect(exports[0].name).toBe("bar"); + expect(exports[0].kind).toBe("variable"); + }); + + test("finds default export", () => { + const root = parseTs("export default function() {}"); + const exports = findExports(root, "typescript"); + expect(exports.some((e) => e.kind === "default")).toBe(true); + }); + + test("finds exported type", () => { + const root = parseTs("export type Foo = string;"); + const exports = findExports(root, "typescript"); + expect(exports).toHaveLength(1); + expect(exports[0].name).toBe("Foo"); + expect(exports[0].kind).toBe("type"); + }); + + test("returns empty for Python", () => { + expect(findExports(parsePy("x = 1"), "python")).toHaveLength(0); + }); +}); diff --git a/test/context/cache.test.ts b/test/context/cache.test.ts new file mode 100644 index 0000000..30c96b7 --- /dev/null +++ b/test/context/cache.test.ts @@ -0,0 +1,132 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { rmSync, existsSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; +import { + openDb, + getSkeleton, + upsertSkeleton, + hasSessionRead, + recordSessionRead, + clearSession, + incrementStats, + getSessionStats, + getAllStats, + pruneOldSessions, + getDbPath, +} from "../../src/context/cache.js"; +import type { Database } from "bun:sqlite"; + +let testDir: string; +let db: Database; + +beforeEach(() => { + testDir = join(tmpdir(), `vibecop-cache-test-${Date.now()}`); + db = openDb(testDir); +}); + +afterEach(() => { + db.close(); + try { rmSync(join(testDir, ".vibecop"), { recursive: true }); } catch {} +}); + +describe("database setup", () => { + test("creates database file", () => { + expect(existsSync(getDbPath(testDir))).toBe(true); + }); + + test("uses WAL mode", () => { + const result = db.query<{ journal_mode: string }, []>("PRAGMA journal_mode").get(); + expect(result?.journal_mode).toBe("wal"); + }); +}); + +describe("skeleton operations", () => { + test("upsert and get skeleton", () => { + upsertSkeleton(db, "src/foo.ts", "hash1", "import { x } from 'y';", "typescript", 500, 50); + const result = getSkeleton(db, "src/foo.ts", "hash1"); + expect(result).not.toBeNull(); + expect(result!.skeleton).toBe("import { x } from 'y';"); + expect(result!.fullTokens).toBe(500); + expect(result!.skeletonTokens).toBe(50); + }); + + test("returns null for wrong hash", () => { + upsertSkeleton(db, "src/foo.ts", "hash1", "skeleton", "typescript", 500, 50); + expect(getSkeleton(db, "src/foo.ts", "wrong-hash")).toBeNull(); + }); + + test("returns null for non-existent path", () => { + expect(getSkeleton(db, "nope.ts", "hash1")).toBeNull(); + }); + + test("upsert overwrites existing entry", () => { + upsertSkeleton(db, "src/foo.ts", "hash1", "old", "typescript", 500, 50); + upsertSkeleton(db, "src/foo.ts", "hash2", "new", "typescript", 600, 60); + expect(getSkeleton(db, "src/foo.ts", "hash2")?.skeleton).toBe("new"); + expect(getSkeleton(db, "src/foo.ts", "hash1")).toBeNull(); + }); +}); + +describe("session read tracking", () => { + test("records and checks session read", () => { + expect(hasSessionRead(db, "sess1", "src/foo.ts")).toBeNull(); + recordSessionRead(db, "sess1", "src/foo.ts", "hash1"); + const read = hasSessionRead(db, "sess1", "src/foo.ts"); + expect(read).not.toBeNull(); + expect(read!.hash).toBe("hash1"); + }); + + test("different sessions are independent", () => { + recordSessionRead(db, "sess1", "src/foo.ts", "hash1"); + expect(hasSessionRead(db, "sess2", "src/foo.ts")).toBeNull(); + }); + + test("updates hash on re-read", () => { + recordSessionRead(db, "sess1", "src/foo.ts", "hash1"); + recordSessionRead(db, "sess1", "src/foo.ts", "hash2"); + expect(hasSessionRead(db, "sess1", "src/foo.ts")?.hash).toBe("hash2"); + }); + + test("clears session data", () => { + recordSessionRead(db, "sess1", "src/foo.ts", "hash1"); + recordSessionRead(db, "sess1", "src/bar.ts", "hash2"); + clearSession(db, "sess1"); + expect(hasSessionRead(db, "sess1", "src/foo.ts")).toBeNull(); + expect(hasSessionRead(db, "sess1", "src/bar.ts")).toBeNull(); + }); +}); + +describe("stats", () => { + test("increments stats", () => { + incrementStats(db, "sess1", { totalReads: 1, cacheHits: 0, tokensSaved: 0 }); + incrementStats(db, "sess1", { totalReads: 1, cacheHits: 1, tokensSaved: 500 }); + const stats = getSessionStats(db, "sess1"); + expect(stats).not.toBeNull(); + expect(stats!.totalReads).toBe(2); + expect(stats!.cacheHits).toBe(1); + expect(stats!.tokensSaved).toBe(500); + }); + + test("returns null for unknown session", () => { + expect(getSessionStats(db, "nope")).toBeNull(); + }); + + test("getAllStats returns all sessions", () => { + incrementStats(db, "sess1", { totalReads: 1 }); + incrementStats(db, "sess2", { totalReads: 2 }); + const all = getAllStats(db); + expect(all.length).toBe(2); + }); +}); + +describe("pruning", () => { + test("prunes old session data", () => { + recordSessionRead(db, "old-sess", "src/foo.ts", "hash1"); + // Manually backdate the read_at + db.run("UPDATE session_reads SET read_at = unixepoch() - 864000 WHERE session_id = 'old-sess'"); + const pruned = pruneOldSessions(db, 7); + expect(pruned).toBeGreaterThan(0); + expect(hasSessionRead(db, "old-sess", "src/foo.ts")).toBeNull(); + }); +}); diff --git a/test/context/integration.test.ts b/test/context/integration.test.ts new file mode 100644 index 0000000..aab2e5b --- /dev/null +++ b/test/context/integration.test.ts @@ -0,0 +1,192 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdirSync, rmSync, writeFileSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; + +let testDir: string; +const contextScript = join(import.meta.dir, "../../src/context.ts"); + +function makeHookInput(filePath: string, sessionId: string = "test-session") { + return JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: filePath }, + session_id: sessionId, + }); +} + +async function runHook( + command: string, + stdin: string, + cwd?: string, +): Promise<{ stdout: string; stderr: string; exitCode: number }> { + const proc = Bun.spawn(["bun", contextScript, command], { + stdin: new Blob([stdin]), + cwd: cwd ?? testDir, + stdout: "pipe", + stderr: "pipe", + }); + + const [stdout, stderr] = await Promise.all([ + new Response(proc.stdout).text(), + new Response(proc.stderr).text(), + ]); + const exitCode = await proc.exited; + return { stdout: stdout.trim(), stderr: stderr.trim(), exitCode }; +} + +beforeEach(() => { + testDir = join(tmpdir(), `vibecop-integration-${Date.now()}`); + mkdirSync(testDir, { recursive: true }); + // Create a .git directory so project root detection works + mkdirSync(join(testDir, ".git")); +}); + +afterEach(() => { + try { rmSync(testDir, { recursive: true }); } catch {} +}); + +describe("pre-hook", () => { + test("passes through for non-supported file", async () => { + const input = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: "/some/file.rs" }, + session_id: "sess1", + }); + const result = await runHook("--pre", input); + expect(result.stdout).toBe("{}"); + }); + + test("passes through for first read of supported file", async () => { + const tsFile = join(testDir, "app.ts"); + writeFileSync(tsFile, 'const x = 1;\nexport function foo() { return x; }\n'); + + const result = await runHook("--pre", makeHookInput(tsFile)); + // First read with no cached skeleton — should passthrough + expect(result.exitCode).toBe(0); + const parsed = JSON.parse(result.stdout); + // Either {} or has additionalContext (if skeleton was cached in a prior post-hook) + expect(typeof parsed).toBe("object"); + }); + + test("passes through when no session_id", async () => { + const input = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: join(testDir, "foo.ts") }, + }); + const result = await runHook("--pre", input); + expect(result.stdout).toBe("{}"); + }); + + test("passes through for partial reads with offset", async () => { + const input = JSON.stringify({ + tool_name: "Read", + tool_input: { file_path: join(testDir, "foo.ts"), offset: 10, limit: 20 }, + session_id: "sess1", + }); + const result = await runHook("--pre", input); + expect(result.stdout).toBe("{}"); + }); +}); + +describe("post-hook", () => { + test("caches skeleton after first read", async () => { + const tsFile = join(testDir, "service.ts"); + writeFileSync( + tsFile, + `import { db } from "./db";\n\nexport function getUser(id: string) {\n return db.find(id);\n}\n`, + ); + + // Post-hook should cache the skeleton + const result = await runHook("--post", makeHookInput(tsFile)); + expect(result.exitCode).toBe(0); + + // Now a pre-hook re-read should find the cached skeleton + const preResult = await runHook("--pre", makeHookInput(tsFile)); + const parsed = JSON.parse(preResult.stdout); + // Should have smart-limited (updatedInput.limit = 30) since it's a re-read + expect(parsed.updatedInput?.limit).toBe(30); + expect(parsed.additionalContext).toContain("vibecop"); + }); + + test("tracks session reads", async () => { + const tsFile = join(testDir, "tracker.ts"); + writeFileSync(tsFile, "const x = 1;\n"); + + // First post-hook + await runHook("--post", makeHookInput(tsFile)); + + // Pre-hook should know this was read before + const preResult = await runHook("--pre", makeHookInput(tsFile)); + const parsed = JSON.parse(preResult.stdout); + // Re-read → should smart-limit + expect(parsed.updatedInput?.limit).toBe(30); + }); +}); + +describe("stats", () => { + test("shows empty stats when no data", async () => { + const proc = Bun.spawn(["bun", contextScript, "stats"], { + cwd: testDir, + stdout: "pipe", + stderr: "pipe", + }); + const stdout = await new Response(proc.stdout).text(); + await proc.exited; + expect(stdout).toContain("No context optimization data"); + }); + + test("shows stats after activity", async () => { + const tsFile = join(testDir, "activity.ts"); + writeFileSync(tsFile, "export const x = 1;\n"); + + // Generate some activity + await runHook("--post", makeHookInput(tsFile, "sess-stats")); + await runHook("--pre", makeHookInput(tsFile, "sess-stats")); + + const proc = Bun.spawn(["bun", contextScript, "stats"], { + cwd: testDir, + stdout: "pipe", + stderr: "pipe", + }); + const stdout = await new Response(proc.stdout).text(); + await proc.exited; + expect(stdout).toContain("Sessions:"); + }); +}); + +describe("compact", () => { + test("runs without error", async () => { + const input = JSON.stringify({ + tool_name: "Compact", + tool_input: {}, + session_id: "sess-compact", + }); + const result = await runHook("--compact", input); + expect(result.exitCode).toBe(0); + }); +}); + +describe("error handling", () => { + test("pre-hook outputs {} on invalid JSON stdin", async () => { + const proc = Bun.spawn(["bun", contextScript, "--pre"], { + stdin: new Blob(["not json"]), + cwd: testDir, + stdout: "pipe", + stderr: "pipe", + }); + const stdout = (await new Response(proc.stdout).text()).trim(); + await proc.exited; + expect(stdout).toBe("{}"); + }); + + test("post-hook doesn't crash on invalid JSON stdin", async () => { + const proc = Bun.spawn(["bun", contextScript, "--post"], { + stdin: new Blob(["not json"]), + cwd: testDir, + stdout: "pipe", + stderr: "pipe", + }); + await proc.exited; + // Should exit without crashing + }); +}); diff --git a/test/context/skeleton.test.ts b/test/context/skeleton.test.ts new file mode 100644 index 0000000..5c7c89e --- /dev/null +++ b/test/context/skeleton.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, test } from "bun:test"; +import { extractSkeleton, languageForExtension } from "../../src/context/skeleton.js"; + +describe("languageForExtension", () => { + test("maps .ts to typescript", () => { + expect(languageForExtension(".ts")).toBe("typescript"); + }); + + test("maps .py to python", () => { + expect(languageForExtension(".py")).toBe("python"); + }); + + test("returns null for unknown extension", () => { + expect(languageForExtension(".rs")).toBeNull(); + }); +}); + +describe("extractSkeleton", () => { + describe("TypeScript", () => { + test("extracts imports", () => { + const source = `import { useState } from "react";\nimport path from "node:path";\n\nconst x = 1;`; + const skeleton = extractSkeleton(source, "typescript"); + expect(skeleton).toContain("react"); + expect(skeleton).toContain("node:path"); + }); + + test("extracts function signatures", () => { + const source = ` +function greet(name: string): string { + return "hello " + name; +} + +const add = (a: number, b: number) => { + return a + b; +}; +`; + const skeleton = extractSkeleton(source, "typescript"); + expect(skeleton).toContain("greet"); + }); + + test("extracts class with methods", () => { + const source = ` +class UserService { + constructor(private db: Database) {} + getUser(id: string) { return this.db.find(id); } + deleteUser(id: string) { this.db.delete(id); } +} +`; + const skeleton = extractSkeleton(source, "typescript"); + expect(skeleton).toContain("UserService"); + expect(skeleton).toContain("constructor"); + expect(skeleton).toContain("getUser"); + expect(skeleton).toContain("deleteUser"); + }); + + test("extracts exports", () => { + const source = `export const VERSION = "1.0";\nexport type Config = { verbose: boolean };`; + const skeleton = extractSkeleton(source, "typescript"); + expect(skeleton).toContain("VERSION"); + expect(skeleton).toContain("Config"); + }); + + test("is much shorter than original", () => { + const source = ` +import { readFileSync } from "node:fs"; + +function processFile(path: string): string { + const content = readFileSync(path, "utf-8"); + const lines = content.split("\\n"); + const filtered = lines.filter(l => l.trim().length > 0); + const numbered = filtered.map((l, i) => \`\${i + 1}: \${l}\`); + const result = numbered.join("\\n"); + console.log(\`Processed \${filtered.length} lines\`); + return result; +} + +export function main() { + const result = processFile("./input.txt"); + console.log(result); +} +`; + const skeleton = extractSkeleton(source, "typescript"); + expect(skeleton.length).toBeLessThan(source.length * 0.6); + }); + }); + + describe("JavaScript", () => { + test("extracts require-style code", () => { + const source = ` +const fs = require("fs"); + +function readConfig(path) { + return JSON.parse(fs.readFileSync(path, "utf-8")); +} + +module.exports = { readConfig }; +`; + const skeleton = extractSkeleton(source, "javascript"); + expect(skeleton).toContain("readConfig"); + }); + }); + + describe("Python", () => { + test("extracts imports and functions", () => { + const source = `import os\nfrom pathlib import Path\n\ndef process(path: str) -> str:\n with open(path) as f:\n return f.read()\n`; + const skeleton = extractSkeleton(source, "python"); + expect(skeleton).toContain("import os"); + expect(skeleton).toContain("pathlib"); + expect(skeleton).toContain("process"); + }); + + test("extracts class with methods", () => { + const source = `class Dog:\n def __init__(self, name):\n self.name = name\n def bark(self):\n print("woof")\n`; + const skeleton = extractSkeleton(source, "python"); + expect(skeleton).toContain("Dog"); + expect(skeleton).toContain("__init__"); + expect(skeleton).toContain("bark"); + }); + }); + + describe("edge cases", () => { + test("handles empty source", () => { + const skeleton = extractSkeleton("", "typescript"); + expect(typeof skeleton).toBe("string"); + }); + + test("handles source with only comments", () => { + const skeleton = extractSkeleton("// just a comment\n/* block */", "typescript"); + expect(typeof skeleton).toBe("string"); + }); + }); +}); diff --git a/test/init.test.ts b/test/init.test.ts index 24b90d1..70cad4b 100644 --- a/test/init.test.ts +++ b/test/init.test.ts @@ -173,4 +173,76 @@ describe("vibecop init", () => { ); expect(hooks.hooks.PostToolUse).toBeDefined(); }); + + describe("--context flag", () => { + test("generates context hooks when .claude/ exists", async () => { + mkdirSync(join(tempDir, ".claude"), { recursive: true }); + + await runInit(tempDir, { context: true }); + + const settingsPath = join(tempDir, ".claude", "settings.json"); + expect(existsSync(settingsPath)).toBe(true); + + const settings = JSON.parse(readFileSync(settingsPath, "utf-8")); + // Should have PreToolUse Read hook + const preHooks = settings.hooks?.PreToolUse ?? []; + expect(preHooks.some((h: { matcher?: string }) => h.matcher === "Read")).toBe(true); + // Should have PostToolUse Read hook + const postHooks = settings.hooks?.PostToolUse ?? []; + expect(postHooks.some((h: { matcher?: string }) => h.matcher === "Read")).toBe(true); + // Should have PostCompact hook + expect(settings.hooks?.PostCompact).toBeDefined(); + }); + + test("merges with existing hooks instead of overwriting", async () => { + mkdirSync(join(tempDir, ".claude"), { recursive: true }); + const existing = { + hooks: { + PostToolUse: [ + { matcher: "Edit|Write", hooks: [{ type: "command", command: "npx vibecop scan" }] }, + ], + }, + }; + writeFileSync(join(tempDir, ".claude", "settings.json"), JSON.stringify(existing)); + + await runInit(tempDir, { context: true }); + + const settings = JSON.parse( + readFileSync(join(tempDir, ".claude", "settings.json"), "utf-8"), + ); + // Original PostToolUse hook should still exist + const postHooks = settings.hooks.PostToolUse; + expect(postHooks.some((h: { matcher?: string }) => h.matcher === "Edit|Write")).toBe(true); + // New Read PostToolUse hook should be added + expect(postHooks.some((h: { matcher?: string }) => h.matcher === "Read")).toBe(true); + }); + + test("skips when existing Read hook would conflict", async () => { + mkdirSync(join(tempDir, ".claude"), { recursive: true }); + const existing = { + hooks: { + PreToolUse: [ + { matcher: "Read", hooks: [{ type: "command", command: "other-tool" }] }, + ], + }, + }; + writeFileSync(join(tempDir, ".claude", "settings.json"), JSON.stringify(existing)); + + await runInit(tempDir, { context: true }); + + const settings = JSON.parse( + readFileSync(join(tempDir, ".claude", "settings.json"), "utf-8"), + ); + // Should NOT add another PreToolUse Read hook + const preHooks = settings.hooks.PreToolUse; + expect(preHooks.length).toBe(1); + expect(preHooks[0].hooks[0].command).toBe("other-tool"); + }); + + test("requires .claude/ directory", async () => { + // No .claude/ directory + await runInit(tempDir, { context: true }); + expect(existsSync(join(tempDir, ".claude", "settings.json"))).toBe(false); + }); + }); });