diff --git a/.gitattributes b/.gitattributes index 3ae0c73af..861194900 100644 --- a/.gitattributes +++ b/.gitattributes @@ -59,6 +59,7 @@ CODE_OF_CONDUCT* text *.xml text *.csv text *.sql text +*.scm text eol=lf # Graphics (binary) *.png binary diff --git a/README.md b/README.md index 5c2e9b0d3..8dd673191 100644 --- a/README.md +++ b/README.md @@ -265,6 +265,7 @@ Advanced and source-build guides: - **Images**: URL and base64 image inputs for providers that support vision - **Provider profiles**: Guided setup plus saved user-level provider profile support - **Local and remote model backends**: Cloud APIs, local servers, and Apple Silicon local inference +- **Codebase intelligence (repo map)**: Structural map of the repository ranked by PageRank importance, auto-injected into context when the `REPO_MAP` flag is enabled or the `REPO_MAP` environment variable is set. Inspect with `/repomap` (2048-token default). See [docs/repo-map.md](docs/repo-map.md) for details. ## Provider Notes diff --git a/bun.lock b/bun.lock index 719109cf1..1b0d1620c 100644 --- a/bun.lock +++ b/bun.lock @@ -37,9 +37,12 @@ "fuse.js": "7.1.0", "get-east-asian-width": "1.5.0", "google-auth-library": "10.6.2", + "graphology": "0.26.0", + "graphology-metrics": "2.4.0", "https-proxy-agent": "7.0.6", "ignore": "7.0.5", "indent-string": "5.0.0", + "js-tiktoken": "1.0.21", "jsonc-parser": "3.3.1", "lodash-es": "4.18.1", "lru-cache": "11.2.7", @@ -57,11 +60,13 @@ "signal-exit": "4.1.0", "supports-hyperlinks": "3.2.0", "tree-kill": "1.2.2", + "tree-sitter-wasms": "0.1.13", "turndown": "7.2.2", "type-fest": "4.41.0", "undici": "7.28.0", "usehooks-ts": "3.1.1", "vscode-languageserver-protocol": "3.17.5", + "web-tree-sitter": "0.25.10", "wrap-ansi": "9.0.2", "ws": "8.21.0", "xss": "1.0.15", @@ -448,6 +453,8 @@ "@vscode/ripgrep-win32-x64": ["@vscode/ripgrep-win32-x64@1.18.0", "", { "os": "win32", "cpu": "x64" }, "sha512-KNPvtElldqILHdnAetujPaowkNbpqJy3ssIGGN6F6Kve9Qi+nNLI2DN01O83JjCEVQbCzl8Ov3QZ9Eov3BR8Dg=="], + "@yomguithereal/helpers": ["@yomguithereal/helpers@1.1.1", "", {}, "sha512-UYvAq/XCA7xoh1juWDYsq3W0WywOB+pz8cgVnE1b45ZfdMhBvHDrgmSFG3jXeZSr2tMTYLGHFHON+ekG05Jebg=="], + "accepts": ["accepts@2.0.0", "", { "dependencies": { "mime-types": "^3.0.0", "negotiator": "^1.0.0" } }, "sha512-5cvg6CtKwfgdmVqY1WIiXKc3Q1bkRqGLi+2W/6ao+6Y7gu/RCwRuAhGEzh5B4KlszSuTLgZYuqFqo5bImjNKng=="], "agent-base": ["agent-base@7.1.4", "", {}, "sha512-MnA+YT8fwfJPgBx3m60MNqakm30XOkyIoH1y6huTQvC0PwZG7ki8NacLBcrPbNoo8vEZy7Jpuk7+jMO+CUovTQ=="], @@ -570,6 +577,8 @@ "etag": ["etag@1.8.1", "", {}, "sha512-aIL5Fx7mawVa300al2BnEE4iNvo1qETxLrPI/o05L7z6go7fCw1J6EQmbK4FmJ2AS7kgVF/KEZWufBfdClMcPg=="], + "events": ["events@3.3.0", "", {}, "sha512-mQw+2fkQbALzQ7V0MY0IqdnXNOeTtP4r0lN9z7AAawCXgqea7bDii20AYrIBrFd/Hx0M2Ocz6S111CaFkUcb0Q=="], + "eventsource": ["eventsource@3.0.7", "", { "dependencies": { "eventsource-parser": "^3.0.1" } }, "sha512-CRT1WTyuQoD771GW56XEZFQ/ZoSfWid1alKGDYMmkt2yl8UXrVR4pspqWNEcqKvVIzg6PAltWjxcSSPrboA4iA=="], "eventsource-parser": ["eventsource-parser@3.0.8", "", {}, "sha512-70QWGkr4snxr0OXLRWsFLeRBIRPuQOvt4s8QYjmUlmlkyTZkRqS7EDVRZtzU3TiyDbXSzaOeF0XUKy8PchzukQ=="], @@ -644,6 +653,18 @@ "graceful-fs": ["graceful-fs@4.2.11", "", {}, "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ=="], + "graphology": ["graphology@0.26.0", "", { "dependencies": { "events": "^3.3.0" }, "peerDependencies": { "graphology-types": ">=0.24.0" } }, "sha512-8SSImzgUUYC89Z042s+0r/vMibY7GX/Emz4LDO5e7jYXhuoWfHISPFJYjpRLUSJGq6UQ6xlenvX1p/hJdfXuXg=="], + + "graphology-indices": ["graphology-indices@0.17.0", "", { "dependencies": { "graphology-utils": "^2.4.2", "mnemonist": "^0.39.0" }, "peerDependencies": { "graphology-types": ">=0.20.0" } }, "sha512-A7RXuKQvdqSWOpn7ZVQo4S33O0vCfPBnUSf7FwE0zNCasqwZVUaCXePuWo5HBpWw68KJcwObZDHpFk6HKH6MYQ=="], + + "graphology-metrics": ["graphology-metrics@2.4.0", "", { "dependencies": { "graphology-indices": "^0.17.0", "graphology-shortest-path": "^2.0.0", "graphology-utils": "^2.4.4", "mnemonist": "^0.39.0", "pandemonium": "2.4.1" }, "peerDependencies": { "graphology-types": ">=0.20.0" } }, "sha512-7WOfOP+mFLCaTJx55Qg4eY+211vr1/b3D/R3biz3SXGhAaCVcWYkfabnmO4O4WBNWANEHtVnFrGgJ0kj6MM6xw=="], + + "graphology-shortest-path": ["graphology-shortest-path@2.1.0", "", { "dependencies": { "@yomguithereal/helpers": "^1.1.1", "graphology-indices": "^0.17.0", "graphology-utils": "^2.4.3", "mnemonist": "^0.39.0" }, "peerDependencies": { "graphology-types": ">=0.20.0" } }, "sha512-KbT9CTkP/u72vGEJzyRr24xFC7usI9Es3LMmCPHGwQ1KTsoZjxwA9lMKxfU0syvT/w+7fZUdB/Hu2wWYcJBm6Q=="], + + "graphology-types": ["graphology-types@0.24.8", "", {}, "sha512-hDRKYXa8TsoZHjgEaysSRyPdT6uB78Ci8WnjgbStlQysz7xR52PInxNsmnB7IBOM1BhikxkNyCVEFgmPKnpx3Q=="], + + "graphology-utils": ["graphology-utils@2.5.2", "", { "peerDependencies": { "graphology-types": ">=0.23.0" } }, "sha512-ckHg8MXrXJkOARk56ZaSCM1g1Wihe2d6iTmz1enGOz4W/l831MBCKSayeFQfowgF8wd+PQ4rlch/56Vs/VZLDQ=="], + "has-flag": ["has-flag@4.0.0", "", {}, "sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ=="], "has-symbols": ["has-symbols@1.1.0", "", {}, "sha512-1cDNdwJ2Jaohmb3sg4OmKaMBwuC48sYni5HUw2DvsC8LjGTLK9h+eb1X6RyuOHe4hT0ULCW68iomhjUoKUqlPQ=="], @@ -692,6 +713,8 @@ "jose": ["jose@6.2.3", "", {}, "sha512-YYVDInQKFJfR/xa3ojUTl8c2KoTwiL1R5Wg9YCydwH0x0B9grbzlg5HC7mMjCtUJjbQ/YnGEZIhI5tCgfTb4Hw=="], + "js-tiktoken": ["js-tiktoken@1.0.21", "", { "dependencies": { "base64-js": "^1.5.1" } }, "sha512-biOj/6M5qdgx5TKjDnFT1ymSpM5tbd3ylwDtrQvFQSu0Z7bBYko2dF+W/aUkXUPuk6IVpRxk/3Q2sHOzGlS36g=="], + "json-bigint": ["json-bigint@1.0.0", "", { "dependencies": { "bignumber.js": "^9.0.0" } }, "sha512-SiPv/8VpZuWbvLSMtTDU8hEfrZWg/mH/nV/b4o0CYbSxu1UIQPLdwKOCIyLQX+VIPO5vrLX3i8qtqFyhdPSUSQ=="], "json-schema-to-ts": ["json-schema-to-ts@3.1.1", "", { "dependencies": { "@babel/runtime": "^7.18.3", "ts-algebra": "^2.0.0" } }, "sha512-+DWg8jCJG2TEnpy7kOm/7/AxaYoaRbjVB4LFZLySZlWn8exGs3A4OLJR966cVvU26N7X9TWxl+Jsw7dzAqKT6g=="], @@ -732,6 +755,8 @@ "mime-types": ["mime-types@3.0.2", "", { "dependencies": { "mime-db": "^1.54.0" } }, "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A=="], + "mnemonist": ["mnemonist@0.39.8", "", { "dependencies": { "obliterator": "^2.0.1" } }, "sha512-vyWo2K3fjrUw8YeeZ1zF0fy6Mu59RHokURlld8ymdUPjMlD9EC9ov1/YPqTgqRvUN9nTr3Gqfz29LYAmu0PHPQ=="], + "ms": ["ms@2.1.3", "", {}, "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA=="], "mz": ["mz@2.7.0", "", { "dependencies": { "any-promise": "^1.0.0", "object-assign": "^4.0.1", "thenify-all": "^1.0.0" } }, "sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q=="], @@ -750,6 +775,8 @@ "object-inspect": ["object-inspect@1.13.4", "", {}, "sha512-W67iLl4J2EXEGTbfeHCffrjDfitvLANg0UlX3wFUUSTx92KXRFegMHUVgSqE+wvhAbi4WqjGg9czysTV2Epbew=="], + "obliterator": ["obliterator@2.0.5", "", {}, "sha512-42CPE9AhahZRsMNslczq0ctAEtqk8Eka26QofnqC346BZdHDySk3LWka23LI7ULIw11NmltpiLagIq8gBozxTw=="], + "on-finished": ["on-finished@2.4.1", "", { "dependencies": { "ee-first": "1.1.1" } }, "sha512-oVlzkg3ENAhCk2zdv7IJwd/QUD4z2RxRwpkcGY8psCVcCYZNq4wYnVWALHM+brtuJjePWiYF/ClmuDr8Ch5+kg=="], "once": ["once@1.4.0", "", { "dependencies": { "wrappy": "1" } }, "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w=="], @@ -766,6 +793,8 @@ "p-try": ["p-try@2.2.0", "", {}, "sha512-R4nPAVTAU0B9D35/Gk3uJf/7XYbQcyohSKdvAxIRSNghFl4e71hVoGnBNQz9cWaXxO2I10KTC+3jMdvvoKw6dQ=="], + "pandemonium": ["pandemonium@2.4.1", "", { "dependencies": { "mnemonist": "^0.39.2" } }, "sha512-wRqjisUyiUfXowgm7MFH2rwJzKIr20rca5FsHXCMNm1W5YPP1hCtrZfgmQ62kP7OZ7Xt+cR858aB28lu5NX55g=="], + "parse-ms": ["parse-ms@4.0.0", "", {}, "sha512-TXfryirbmq34y8QBwgqCVLi+8oA3oWx2eAnSn62ITyEhEYaWRlVZ2DvMM9eZbMs/RfxPu/PK/aBLyGj4IrqMHw=="], "parse5": ["parse5@5.1.1", "", {}, "sha512-ugq4DFI0Ptb+WWjAdOK16+u/nHfiIrcE+sh8kZMaM0WllQKLI9rOUq6c2b7cwPkXdzfQESqvoqK6ug7U/Yyzug=="], @@ -892,6 +921,8 @@ "tree-kill": ["tree-kill@1.2.2", "", { "bin": { "tree-kill": "cli.js" } }, "sha512-L0Orpi8qGpRG//Nd+H90vFB+3iHnue1zSSGmNOOCh1GLJ7rUKVwV2HvijphGQS2UmhUZewS9VgvxYIdgr+fG1A=="], + "tree-sitter-wasms": ["tree-sitter-wasms@0.1.13", "", { "dependencies": { "tree-sitter-wasms": "^0.1.11" } }, "sha512-wT+cR6DwaIz80/vho3AvSF0N4txuNx/5bcRKoXouOfClpxh/qqrF4URNLQXbbt8MaAxeksZcZd1j8gcGjc+QxQ=="], + "ts-algebra": ["ts-algebra@2.0.0", "", {}, "sha512-FPAhNPFMrkwz76P7cdjdmiShwMynZYN6SgOujD1urY4oNm80Ou9oMdmbR45LotcKOXoy7wSmHkRFE6Mxbrhefw=="], "tslib": ["tslib@1.14.1", "", {}, "sha512-Xni35NKzjgMrwevysHTCArtLDpPvye8zV/0E4EyYn43P7/7qvQwPh9BGkHewbMulVntbigmcT7rdX3BNo9wRJg=="], @@ -928,6 +959,8 @@ "web-streams-polyfill": ["web-streams-polyfill@3.3.3", "", {}, "sha512-d2JWLCivmZYTSIoge9MsgFCZrt571BikcWGYkjC1khllbTeDlGqZ2D8vD8E/lJa8WGWbb7Plm8/XJYV7IJHZZw=="], + "web-tree-sitter": ["web-tree-sitter@0.25.10", "", { "peerDependencies": { "@types/emscripten": "^1.40.0" }, "optionalPeers": ["@types/emscripten"] }, "sha512-Y09sF44/13XvgVKgO2cNDw5rGk6s26MgoZPXLESvMXeefBf7i6/73eFurre0IsTW6E14Y0ArIzhUMmjoc7xyzA=="], + "which": ["which@2.0.2", "", { "dependencies": { "isexe": "^2.0.0" }, "bin": { "node-which": "./bin/node-which" } }, "sha512-BLI3Tl1TW3Pvl70l3yq3Y64i+awpwXqsGBYWkkqMtnbXgrMD+yj7rhW0kuEDxzJaYXGjEW5ogapKNMEKNMjibA=="], "which-module": ["which-module@2.0.1", "", {}, "sha512-iBdZ57RDvnOR9AGBhML2vFZf7h8vmBjhoaZqODJBFWHVtKkDmKuHai3cx5PgVMrX5YDNp27AofYbAwctSS+vhQ=="], diff --git a/docs/repo-map.md b/docs/repo-map.md new file mode 100644 index 000000000..c66ae37f4 --- /dev/null +++ b/docs/repo-map.md @@ -0,0 +1,68 @@ +# Codebase Intelligence — Repo Map + +The repo map feature gives the AI model structural awareness of your codebase at the start of each session. Instead of the model needing to explore the repository with `Grep`, `Glob`, and `Read` calls, it starts with a ranked summary of the most important files and their key signatures. + +## How it works + +1. **File enumeration** — Lists tracked files plus untracked, unignored files via `git ls-files --cached --others --exclude-standard` (falls back to a manual directory walk when not in a git repo) +2. **Symbol extraction** — Parses each supported source file with tree-sitter to extract function, class, type, and interface definitions, plus cross-file references +3. **Reference graph** — Builds a directed graph where an edge from file A to file B means A references a symbol defined in B. Edges are weighted by reference count multiplied by the IDF (inverse document frequency) of the symbol name — common names like `get`, `set`, `value` contribute less +4. **PageRank** — Ranks files by structural importance using PageRank. Files imported by many others rank highest +5. **Rendering** — Walks ranked files top-down, emitting file paths and definition signatures, stopping when the token budget is reached + +Results are cached to disk (`~/.openclaude/repomap-cache/`) keyed by file path, mtime, and size. Only changed files are re-parsed on subsequent runs. + +## Supported languages + +- TypeScript (`.ts`, `.tsx`) +- JavaScript (`.js`, `.jsx`, `.mjs`, `.cjs`) +- Python (`.py`) + +Additional language grammars will be added in future releases. + +## Enabling auto-injection + +The repo map is gated behind the `REPO_MAP` feature flag, **off by default**. To enable auto-injection into the session context: + +Set the environment variable before launching: + +```bash +REPO_MAP=1 openclaude +``` + +Or add it to your shell profile for persistent use. + +When enabled, the map is built once per session and prepended to the system context alongside git status and CLAUDE.md content. The auto-injected map uses a 1024-token budget. + +Auto-injection is skipped in: +- Bare mode (`--bare`) +- Remote sessions (`CLAUDE_CODE_REMOTE`) + +## The /repomap slash command + +The `/repomap` command is always available regardless of the feature flag. It lets you inspect and tune the map interactively. + +```text +/repomap # Show the map with default settings (2048 tokens) +/repomap --tokens 4096 # Increase the token budget for a larger map +/repomap --focus src/tools/ # Boost specific paths in the ranking +/repomap --focus src/context.ts # Can use multiple --focus flags +/repomap --focus-symbols buildTool # Boost files that define specific symbols +/repomap --stats # Show cache statistics +/repomap --invalidate # Clear cache and rebuild from scratch +``` + +## The RepoMap tool + +The model can also call the `RepoMap` tool on demand during a session. This is useful when: +- The model needs structural context mid-conversation +- The user asks about specific areas (the model can pass `focus_files` or `focus_symbols`) +- A larger token budget is needed than the auto-injected default + +## Known limitations + +- **Signatures only** — The map shows function/class/type declarations, not implementations. The model still needs `Read` to see function bodies. +- **Cold build time** — First build on large repos (2000+ files) can take 20-30 seconds due to WASM-based parsing. Subsequent builds use the disk cache and complete in under 100ms. +- **Language coverage** — Only TypeScript, JavaScript, and Python are supported. Files in other languages are skipped. +- **TypeScript references** — The TypeScript tree-sitter query captures type annotations and `new` expressions as references, but not plain function calls. This means the ranking slightly favors type-heavy hub files. +- **Git dependency** — File enumeration uses `git ls-files --cached --others --exclude-standard` by default, so untracked files that are not ignored can appear in the map. Non-git repos fall back to a directory walk with hardcoded exclusions. diff --git a/knip.json b/knip.json index 41bca0aa9..95ca84bc4 100644 --- a/knip.json +++ b/knip.json @@ -16,6 +16,7 @@ "project": ["src/**/*.{ts,tsx}", "scripts/**/*.ts"], "ignore": [ "src/integrations/generated/**", + "src/context/repoMap/__fixtures__/**", "src/services/internalLogging.ts", "src/types/generated/**", "src/test/fixtures/**", diff --git a/package.json b/package.json index 88d9c6eba..24daf0eb7 100644 --- a/package.json +++ b/package.json @@ -105,11 +105,14 @@ "fflate": "0.8.2", "figures": "6.1.0", "fuse.js": "7.1.0", + "graphology": "0.26.0", "get-east-asian-width": "1.5.0", "google-auth-library": "10.6.2", "https-proxy-agent": "7.0.6", "ignore": "7.0.5", + "graphology-metrics": "2.4.0", "indent-string": "5.0.0", + "js-tiktoken": "1.0.21", "jsonc-parser": "3.3.1", "lodash-es": "4.18.1", "lru-cache": "11.2.7", @@ -127,10 +130,12 @@ "signal-exit": "4.1.0", "supports-hyperlinks": "3.2.0", "tree-kill": "1.2.2", + "tree-sitter-wasms": "0.1.13", "turndown": "7.2.2", "type-fest": "4.41.0", "undici": "7.28.0", "usehooks-ts": "3.1.1", + "web-tree-sitter": "0.25.10", "vscode-languageserver-protocol": "3.17.5", "wrap-ansi": "9.0.2", "ws": "8.21.0", diff --git a/scripts/build.ts b/scripts/build.ts index 9af3373db..f5e6acb90 100644 --- a/scripts/build.ts +++ b/scripts/build.ts @@ -100,6 +100,9 @@ const featureFlags: Record = { COWORKER_TYPE_TELEMETRY: false, // Telemetry for agent/coworker type classification MCP_SKILLS: true, // Dynamic MCP skill discovery via skill:// resources + // ── Disabled by default, opt-in via runtime env var ───────────────── + REPO_MAP: false, // Auto-injected codebase intelligence repo-map; users opt in with REPO_MAP=1 (the runtime gate in src/context.ts honors the env var even when this flag is false) + // ── Enabled: upstream defaults ────────────────────────────────────── COORDINATOR_MODE: true, // Multi-agent coordinator with worker delegation BUILTIN_EXPLORE_PLAN_AGENTS: true, // Built-in Explore/Plan specialized subagents diff --git a/scripts/externals.ts b/scripts/externals.ts index 27aa8d0f0..b2551c455 100644 --- a/scripts/externals.ts +++ b/scripts/externals.ts @@ -25,6 +25,13 @@ export const COMMON_EXTERNALS: string[] = [ // Orama search engine '@orama/orama', '@orama/plugin-data-persistence', + // web-tree-sitter ships a WASM file alongside its JS and resolves the + // path via require.resolve at runtime; bundling would freeze the build + // host's absolute path, so keep it external. + 'web-tree-sitter', + // tree-sitter-wasms ships per-language .wasm files resolved via + // require.resolve at runtime — same bundling concern as web-tree-sitter. + 'tree-sitter-wasms', ] // Additional packages external only in the SDK bundle (TUI + heavy deps) @@ -124,4 +131,9 @@ export const INTENTIONALLY_BUNDLED: string[] = [ 'vscode-languageserver-protocol', // File watching 'chokidar', + // Graph algorithms (repo map PageRank) + 'graphology', + 'graphology-metrics', + // Tokenizer for repo map token budgeting + 'js-tiktoken', ] diff --git a/src/commands.ts b/src/commands.ts index 1dc952962..5b73f7268 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -27,6 +27,7 @@ import doctor from './commands/doctor/index.js' import onboardGithub from './commands/onboard-github/index.js' import knowledge from './commands/knowledge/index.js' import memory from './commands/memory/index.js' +import repomap from './commands/repomap/index.js' import help from './commands/help/index.js' import ide from './commands/ide/index.js' import init from './commands/init.js' @@ -333,6 +334,7 @@ const COMMANDS = memoize((): Command[] => [ releaseNotes, reloadPlugins, rename, + repomap, replay, requestSize, requestSizeNonInteractive, diff --git a/src/commands/repomap/index.ts b/src/commands/repomap/index.ts new file mode 100644 index 000000000..da09fa9c9 --- /dev/null +++ b/src/commands/repomap/index.ts @@ -0,0 +1,17 @@ +/** + * /repomap command - minimal metadata only. + * Implementation is lazy-loaded from repomap.ts to reduce startup time. + */ +import type { Command } from '../../commands.js' + +const repomap = { + type: 'local', + name: 'repomap', + description: + 'Show or configure the repository structural map (codebase intelligence)', + isHidden: false, + supportsNonInteractive: true, + load: () => import('./repomap.js'), +} satisfies Command + +export default repomap diff --git a/src/commands/repomap/repomap.test.ts b/src/commands/repomap/repomap.test.ts new file mode 100644 index 000000000..094bce98f --- /dev/null +++ b/src/commands/repomap/repomap.test.ts @@ -0,0 +1,299 @@ +import { describe, expect, test } from 'bun:test' +import { mkdtempSync, rmSync, writeFileSync } from 'fs' +import { tmpdir } from 'os' +import { join } from 'path' +import { invalidateCache } from '../../context/repoMap/index.js' +import { getCwdState, setCwdState } from '../../bootstrap/state.js' +import { call, parseArgs, runRepoMapCommand } from './repomap.js' +import type { CacheStats, RepoMapResult } from '../../context/repoMap/index.js' + +const SAMPLE_RESULT: RepoMapResult = { + map: 'main.ts:\n export function main(): string', + cacheHit: false, + buildTimeMs: 12, + fileCount: 1, + totalFileCount: 1, + tokenCount: 42, +} + +function fakeDeps(overrides: Partial<{ + buildRepoMap: (options: { + root: string + maxTokens: number + focusFiles?: string[] + focusSymbols?: string[] + }) => Promise + invalidateCache: () => void + getCacheStats: () => CacheStats +}> = {}) { + return Promise.resolve({ + buildRepoMap: async () => SAMPLE_RESULT, + invalidateCache: () => {}, + getCacheStats: () => ({ + cacheDir: '/tmp/cache', + cacheFile: null, + entryCount: 0, + exists: false, + }), + ...overrides, + }) +} + +async function runTextCommand(args: string, root: string): Promise { + const result = await runRepoMapCommand(args, root) + return textValue(result) +} + +function textValue(result: Awaited>): string { + if (result.type !== 'text') { + throw new Error(`/repomap must return type:'text', got ${result.type}`) + } + return result.value +} + +describe('/repomap argument parsing', () => { + test('defaults to 2048 tokens with no flags', () => { + const result = parseArgs('') + expect(result.tokens).toBe(2048) + expect(result.focus).toEqual([]) + expect(result.invalidate).toBe(false) + expect(result.stats).toBe(false) + }) + + test('parses --tokens flag', () => { + const result = parseArgs('--tokens 4096') + expect(result.tokens).toBe(4096) + }) + + test('rejects --tokens below 256', () => { + const result = parseArgs('--tokens 100') + expect(result.tokens).toBe(2048) // falls back to default + }) + + test('rejects --tokens above 16384', () => { + const result = parseArgs('--tokens 20000') + expect(result.tokens).toBe(2048) // falls back to default + }) + + test('parses --focus flag', () => { + const result = parseArgs('--focus src/tools/') + expect(result.focus).toEqual(['src/tools/']) + }) + + test('parses quoted --focus paths with spaces', () => { + const result = parseArgs('--focus "src/my dir" --tokens 4096') + expect(result.tokens).toBe(4096) + expect(result.focus).toEqual(['src/my dir']) + }) + + test('parses unquoted --focus glob values without shifting flags', () => { + const result = parseArgs('--focus src/*.ts --tokens 4096') + expect(result.tokens).toBe(4096) + expect(result.focus).toEqual(['src/*.ts']) + }) + + test('does not treat shell operators as flag values', () => { + const result = parseArgs('--focus && --tokens 4096') + expect(result.tokens).toBe(4096) + expect(result.focus).toEqual([]) + }) + + test('parses multiple --focus flags', () => { + const result = parseArgs('--focus src/tools/ --focus src/context.ts') + expect(result.focus).toEqual(['src/tools/', 'src/context.ts']) + }) + + test('parses --focus-symbols flag', () => { + const result = parseArgs('--focus-symbols buildTool') + expect(result.focusSymbols).toEqual(['buildTool']) + }) + + test('parses multiple --focus-symbols flags', () => { + const result = parseArgs('--focus-symbols buildTool --focus-symbols ToolUseContext') + expect(result.focusSymbols).toEqual(['buildTool', 'ToolUseContext']) + }) + + test('parses combined --focus and --focus-symbols flags', () => { + const result = parseArgs('--focus src/tools/ --focus-symbols buildTool --tokens 4096') + expect(result.tokens).toBe(4096) + expect(result.focus).toEqual(['src/tools/']) + expect(result.focusSymbols).toEqual(['buildTool']) + }) + + test('parses --invalidate flag', () => { + const result = parseArgs('--invalidate') + expect(result.invalidate).toBe(true) + expect(result.stats).toBe(false) + }) + + test('parses --stats flag', () => { + const result = parseArgs('--stats') + expect(result.stats).toBe(true) + expect(result.invalidate).toBe(false) + }) + + test('parses combined flags', () => { + const result = parseArgs('--tokens 2048 --focus src/tools/ --invalidate') + expect(result.tokens).toBe(2048) + expect(result.focus).toEqual(['src/tools/']) + expect(result.invalidate).toBe(true) + }) +}) + +describe('/repomap command', () => { + test('builds a repository map using the default token budget', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-command-')) + try { + writeFileSync( + join(tempDir, 'main.ts'), + 'export function main(): string { return "hello" }\n', + ) + + const value = await runTextCommand('', tempDir) + + expect(value).toContain('Repository map:') + expect(value).toContain('main.ts:') + expect(value).toContain('Tokens:') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('call wrapper uses the current cwd state', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-command-call-')) + const previousCwd = getCwdState() + try { + writeFileSync( + join(tempDir, 'wrapper.ts'), + 'export function commandWrapperRoot(): string { return "ok" }\n', + ) + setCwdState(tempDir) + + const result = await call('', {} as Parameters[1]) + + expect(textValue(result)).toContain('wrapper.ts:') + expect(textValue(result)).toContain('commandWrapperRoot') + } finally { + setCwdState(previousCwd) + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('reports cache stats without building a map', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-command-')) + try { + const value = await runTextCommand('--stats', tempDir) + + expect(value).toContain('Repository map cache stats:') + expect(value).toContain('Cached entries:') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('invalidates and rebuilds the cache', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-command-')) + try { + writeFileSync( + join(tempDir, 'main.ts'), + 'export function value(): number { return 1 }\n', + ) + + const value = await runTextCommand('--invalidate --tokens 512', tempDir) + + expect(value).toContain('Cache invalidated and rebuilt.') + expect(value).toContain('main.ts:') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('forwards focus-symbols to buildRepoMap', async () => { + const captured: { focusSymbols?: string[] } = {} + const result = await runRepoMapCommand( + '--focus-symbols buildTool --focus-symbols ToolUseContext', + '/tmp/repo', + fakeDeps({ + buildRepoMap: async (options) => { + captured.focusSymbols = options.focusSymbols + return SAMPLE_RESULT + }, + }), + ) + + expect(textValue(result)).toContain('Repository map:') + expect(captured.focusSymbols).toEqual(['buildTool', 'ToolUseContext']) + }) + + test('reports build failures without throwing', async () => { + const result = await runRepoMapCommand( + '', + '/tmp/repo', + fakeDeps({ + buildRepoMap: async () => { + throw new Error('parser unavailable') + }, + }), + ) + + expect(textValue(result)).toContain('Failed to build repository map: parser unavailable') + }) + + test('distinguishes invalidate failures from rebuild failures', async () => { + const invalidateResult = await runRepoMapCommand( + '--invalidate', + '/tmp/repo', + fakeDeps({ + invalidateCache: () => { + throw new Error('cache denied') + }, + }), + ) + expect(textValue(invalidateResult)).toContain( + 'Failed to invalidate repository map cache: cache denied', + ) + + const rebuildResult = await runRepoMapCommand( + '--invalidate', + '/tmp/repo', + fakeDeps({ + buildRepoMap: async () => { + throw new Error('wasm missing') + }, + }), + ) + expect(textValue(rebuildResult)).toContain( + 'Cache invalidated, but rebuilding the repository map failed: wasm missing', + ) + }) + + test('reports stats failures without throwing', async () => { + const result = await runRepoMapCommand( + '--stats', + '/tmp/repo', + fakeDeps({ + getCacheStats: () => { + throw new Error('stats denied') + }, + }), + ) + + expect(textValue(result)).toContain( + 'Failed to read repository map cache stats: stats denied', + ) + }) + + test('reports lazy import failures without throwing', async () => { + const result = await runRepoMapCommand( + '', + '/tmp/repo', + Promise.reject(new Error('module missing')), + ) + + expect(textValue(result)).toContain('Failed to load repo map module: module missing') + }) +}) diff --git a/src/commands/repomap/repomap.ts b/src/commands/repomap/repomap.ts new file mode 100644 index 000000000..5fd35e999 --- /dev/null +++ b/src/commands/repomap/repomap.ts @@ -0,0 +1,169 @@ +import type { + LocalCommandCall, + LocalCommandResult, +} from '../../types/command.js' +import type { CacheStats, RepoMapResult } from '../../context/repoMap/index.js' +import { getCwd } from '../../utils/cwd.js' +import type { ParseEntry } from '../../utils/bash/shellQuote.js' +import { tryParseShellCommand } from '../../utils/bash/shellQuote.js' + +type ArgPart = string | null + +/** Parse CLI-style arguments from the command string. */ +export function parseArgs(args: string): { + tokens: number + focus: string[] + focusSymbols: string[] + invalidate: boolean + stats: boolean +} { + const parsed = tryParseShellCommand(args) + const parts: ArgPart[] = parsed.success + ? parsed.tokens.map(normalizeParsedToken) + : args.trim().split(/\s+/).filter(Boolean) + let tokens = 2048 + const focus: string[] = [] + const focusSymbols: string[] = [] + let invalidate = false + let stats = false + + for (let i = 0; i < parts.length; i++) { + const part = parts[i]! + const next = parts[i + 1] + if (part === '--tokens' && typeof next === 'string') { + const n = parseInt(next, 10) + if (!isNaN(n) && n >= 256 && n <= 16384) { + tokens = n + } + i++ + } else if (part === '--focus' && typeof next === 'string') { + focus.push(next) + i++ + } else if (part === '--focus-symbols' && typeof next === 'string') { + focusSymbols.push(next) + i++ + } else if (part === '--invalidate') { + invalidate = true + } else if (part === '--stats') { + stats = true + } + } + + return { tokens, focus, focusSymbols, invalidate, stats } +} + +function normalizeParsedToken(part: ParseEntry): ArgPart { + if (typeof part === 'string') return part + if ( + 'op' in part && + part.op === 'glob' && + 'pattern' in part && + typeof part.pattern === 'string' + ) { + return part.pattern + } + return null +} + +export const call: LocalCommandCall = async (args) => { + const root = getCwd() + return runRepoMapCommand(args ?? '', root) +} + +type RepoMapCommandDeps = { + buildRepoMap: (options: { + root: string + maxTokens: number + focusFiles?: string[] + focusSymbols?: string[] + }) => Promise + invalidateCache: (root?: string) => void + getCacheStats: (root?: string) => CacheStats +} + +async function loadRepoMapDeps(): Promise { + return import('../../context/repoMap/index.js') +} + +export async function runRepoMapCommand( + args: string, + root: string, + depsPromise: Promise = loadRepoMapDeps(), +): Promise { + const { tokens, focus, focusSymbols, invalidate, stats } = parseArgs(args) + + let deps: RepoMapCommandDeps + try { + deps = await depsPromise + } catch (err) { + return renderError('Failed to load repo map module', err) + } + + if (stats) { + try { + const cacheStats = deps.getCacheStats(root) + const lines = [ + `Repository map cache stats:`, + ` Cache directory: ${cacheStats.cacheDir}`, + ` Cache file: ${cacheStats.cacheFile ?? '(none)'}`, + ` Cached entries: ${cacheStats.entryCount}`, + ` Cache exists: ${cacheStats.exists}`, + ] + return { type: 'text', value: lines.join('\n') } + } catch (err) { + return renderError('Failed to read repository map cache stats', err) + } + } + + if (invalidate) { + try { + deps.invalidateCache(root) + } catch (err) { + return renderError('Failed to invalidate repository map cache', err) + } + + try { + const result = await deps.buildRepoMap({ + root, + maxTokens: tokens, + focusFiles: focus.length > 0 ? focus : undefined, + focusSymbols: focusSymbols.length > 0 ? focusSymbols : undefined, + }) + return formatRepoMapResult('Cache invalidated and rebuilt.', result) + } catch (err) { + return renderError('Cache invalidated, but rebuilding the repository map failed', err) + } + } + + try { + const result = await deps.buildRepoMap({ + root, + maxTokens: tokens, + focusFiles: focus.length > 0 ? focus : undefined, + focusSymbols: focusSymbols.length > 0 ? focusSymbols : undefined, + }) + + return formatRepoMapResult('Repository map:', result) + } catch (err) { + return renderError('Failed to build repository map', err) + } +} + +function formatRepoMapResult(prefix: string, result: RepoMapResult) { + return { + type: 'text' as const, + value: [ + `${prefix} ${result.fileCount} files ranked (${result.totalFileCount} total) | Tokens: ${result.tokenCount} | Time: ${result.buildTimeMs}ms | Cache hit: ${result.cacheHit}`, + '', + result.map, + ].join('\n'), + } +} + +function renderError(prefix: string, err: unknown) { + const detail = err instanceof Error ? err.message : String(err) + return { + type: 'text' as const, + value: `${prefix}: ${detail}`, + } +} diff --git a/src/context.repoMap.test.ts b/src/context.repoMap.test.ts new file mode 100644 index 000000000..0f2cacf70 --- /dev/null +++ b/src/context.repoMap.test.ts @@ -0,0 +1,236 @@ +import { afterEach, describe, expect, test } from 'bun:test' +import { getCwdState, setCwdState } from './bootstrap/state.js' + +const originalCwdState = getCwdState() +const originalRepoMapEnv = process.env.REPO_MAP +const originalRemoteEnv = process.env.CLAUDE_CODE_REMOTE + +afterEach(async () => { + const { getRepoMapContext, getSystemContext } = await import('./context.js') + if (originalRepoMapEnv === undefined) delete process.env.REPO_MAP + else process.env.REPO_MAP = originalRepoMapEnv + if (originalRemoteEnv === undefined) delete process.env.CLAUDE_CODE_REMOTE + else process.env.CLAUDE_CODE_REMOTE = originalRemoteEnv + setCwdState(originalCwdState) + getRepoMapContext.cache.clear?.() + getSystemContext.cache.clear?.() +}) + +// The feature() function from bun:bundle is shimmed at build time. +// In tests, it's not available, so we test the getRepoMapContext logic +// by importing and calling it directly — the function checks feature('REPO_MAP') +// which in the test environment (no bun:bundle shim) will throw or return false. +// We test the actual logic paths through integration-style tests. + +describe('getRepoMapContext', () => { + test( + 'returns null when REPO_MAP flag is off (default)', + async () => { + const { getRepoMapContext } = await import('./context.js') + const previous = process.env.REPO_MAP + delete process.env.REPO_MAP + getRepoMapContext.cache.clear?.() + + try { + await expect(getRepoMapContext()).resolves.toBeNull() + } finally { + if (previous === undefined) delete process.env.REPO_MAP + else process.env.REPO_MAP = previous + getRepoMapContext.cache.clear?.() + } + }, + { timeout: 10000 }, + ) + + test('buildRepoMap produces valid output for context injection', async () => { + const { mkdtempSync, writeFileSync, rmSync } = await import('fs') + const { tmpdir } = await import('os') + const { join } = await import('path') + const { buildRepoMap } = await import('./context/repoMap/index.js') + + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-ctx-')) + try { + writeFileSync( + join(tempDir, 'main.ts'), + 'export function main(): void { console.log("hello") }\n', + ) + writeFileSync( + join(tempDir, 'utils.ts'), + 'import { main } from "./main"\nexport function helper(): void { main() }\n', + ) + + const result = await buildRepoMap({ + root: tempDir, + maxTokens: 1024, + }) + + // Valid map that could be injected + expect(result.map.length).toBeGreaterThan(0) + expect(result.tokenCount).toBeGreaterThan(0) + expect(result.tokenCount).toBeLessThanOrEqual(1024) + expect(typeof result.cacheHit).toBe('boolean') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + const { invalidateCache } = await import('./context/repoMap/index.js') + invalidateCache(tempDir) + } + }) + + test('auto-injection builds from cwd state instead of process cwd', async () => { + const { mkdtempSync, writeFileSync, rmSync } = await import('fs') + const { tmpdir } = await import('os') + const { join } = await import('path') + const { getRepoMapContext } = await import('./context.js') + const { invalidateCache } = await import('./context/repoMap/index.js') + + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-cwd-state-')) + process.env.REPO_MAP = '1' + setCwdState(tempDir) + getRepoMapContext.cache.clear?.() + + try { + writeFileSync( + join(tempDir, 'state-root.ts'), + 'export function AutoInjectedStateRoot(): void {}\n', + ) + + const context = await getRepoMapContext() + + expect(context).toContain('state-root.ts:') + expect(context).toContain('AutoInjectedStateRoot') + } finally { + invalidateCache(tempDir) + rmSync(tempDir, { recursive: true, force: true }) + } + }) + + test('auto-injection does not leave the timeout timer active after a fast build', async () => { + const { mkdtempSync, writeFileSync, rmSync } = await import('fs') + const { tmpdir } = await import('os') + const { join } = await import('path') + const { getRepoMapContext } = await import('./context.js') + const { invalidateCache } = await import('./context/repoMap/index.js') + + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-timer-')) + process.env.REPO_MAP = '1' + setCwdState(tempDir) + getRepoMapContext.cache.clear?.() + + try { + writeFileSync( + join(tempDir, 'timer-root.ts'), + 'export function TimerRoot(): void {}\n', + ) + + const getActiveHandles = (process as unknown as { + _getActiveHandles: () => unknown[] + })._getActiveHandles + const before = getActiveHandles().length + await getRepoMapContext() + const after = getActiveHandles().length + + expect(after).toBeLessThanOrEqual(before) + } finally { + invalidateCache(tempDir) + rmSync(tempDir, { recursive: true, force: true }) + } + }) + + test('timeout cancellation is observed by the repo map build hook', async () => { + const { runRepoMapBuildWithTimeout } = await import('./context.js') + let resolveCancellation: (value: boolean) => void = () => {} + const cancellationObserved = new Promise(resolve => { + resolveCancellation = resolve + }) + + const result = await runRepoMapBuildWithTimeout({ + root: '/tmp/repo', + maxTokens: 1024, + timeoutMs: 1, + buildRepoMap: async ({ shouldContinue }) => { + await new Promise(resolve => setTimeout(resolve, 10)) + try { + shouldContinue?.() + } catch (err) { + resolveCancellation(true) + throw err + } + resolveCancellation(false) + return { + map: 'late.ts:\n export function late(): void {}', + cacheHit: false, + buildTimeMs: 10, + fileCount: 1, + totalFileCount: 1, + tokenCount: 10, + } + }, + }) + + expect(result).toEqual({ result: null, timedOut: true }) + await expect(cancellationObserved).resolves.toBe(true) + }) + + test('getSystemContext does not include repoMap key when flag is off', async () => { + const { getRepoMapContext, getSystemContext } = await import('./context.js') + const previousRepoMap = process.env.REPO_MAP + const previousRemote = process.env.CLAUDE_CODE_REMOTE + delete process.env.REPO_MAP + process.env.CLAUDE_CODE_REMOTE = '1' + getRepoMapContext.cache.clear?.() + getSystemContext.cache.clear?.() + + try { + const context = await getSystemContext() + expect(context).not.toHaveProperty('repoMap') + } finally { + if (previousRepoMap === undefined) delete process.env.REPO_MAP + else process.env.REPO_MAP = previousRepoMap + if (previousRemote === undefined) delete process.env.CLAUDE_CODE_REMOTE + else process.env.CLAUDE_CODE_REMOTE = previousRemote + getRepoMapContext.cache.clear?.() + getSystemContext.cache.clear?.() + } + }) +}) + +describe('REPO_MAP feature flag', () => { + test('flag defaults to false in build config', async () => { + const { readFileSync } = await import('fs') + const ts = await import('typescript') + const source = ts.createSourceFile( + 'scripts/build.ts', + readFileSync('scripts/build.ts', 'utf-8'), + ts.ScriptTarget.Latest, + true, + ) + + let repoMapFlag: boolean | null = null + source.forEachChild(node => { + if (!ts.isVariableStatement(node)) return + for (const declaration of node.declarationList.declarations) { + if ( + ts.isIdentifier(declaration.name) && + declaration.name.text === 'featureFlags' && + declaration.initializer && + ts.isObjectLiteralExpression(declaration.initializer) + ) { + const property = declaration.initializer.properties.find(prop => + ts.isPropertyAssignment(prop) && + ts.isIdentifier(prop.name) && + prop.name.text === 'REPO_MAP', + ) + if ( + property && + ts.isPropertyAssignment(property) && + property.initializer.kind === ts.SyntaxKind.FalseKeyword + ) { + repoMapFlag = false + } + } + } + }) + + expect(repoMapFlag as boolean | null).toBe(false) + }) +}) diff --git a/src/context.ts b/src/context.ts index 290118bc2..9ed4ed668 100644 --- a/src/context.ts +++ b/src/context.ts @@ -16,8 +16,62 @@ import { execFileNoThrow } from './utils/execFileNoThrow.js' import { getBranch, getDefaultBranch, getIsGit, gitExe } from './utils/git.js' import { shouldIncludeGitInstructions } from './utils/gitSettings.js' import { logError } from './utils/log.js' +import { getCwd } from './utils/cwd.js' +import type { RepoMapResult } from './context/repoMap/index.js' const MAX_STATUS_CHARS = 2000 +const REPO_MAP_CONTEXT_TIMEOUT_MS = 5000 +const REPO_MAP_TIMEOUT = Symbol('repo_map_timeout') +const REPO_MAP_CANCELLED = Symbol('repo_map_cancelled') + +type RepoMapBuildFn = (options: { + root: string + maxTokens: number + shouldContinue?: () => void +}) => Promise + +export async function runRepoMapBuildWithTimeout({ + buildRepoMap, + root, + maxTokens, + timeoutMs, +}: { + buildRepoMap: RepoMapBuildFn + root: string + maxTokens: number + timeoutMs: number +}): Promise<{ result: RepoMapResult | null; timedOut: boolean }> { + let timeout: ReturnType | undefined + let timedOut = false + const result = await Promise.race([ + buildRepoMap({ + root, + maxTokens, + shouldContinue: () => { + if (timedOut) throw REPO_MAP_CANCELLED + }, + }).catch(err => { + if (err === REPO_MAP_CANCELLED) return REPO_MAP_TIMEOUT + throw err + }), + new Promise(resolve => { + timeout = setTimeout( + () => { + timedOut = true + resolve(REPO_MAP_TIMEOUT) + }, + timeoutMs, + ) + }), + ]).finally(() => { + if (timeout) clearTimeout(timeout) + }) as RepoMapResult | typeof REPO_MAP_TIMEOUT + + if (result === REPO_MAP_TIMEOUT) { + return { result: null, timedOut: true } + } + return { result, timedOut: false } +} // System prompt injection for cache breaking (internal-only, ephemeral debugging state) let systemPromptInjection: string | null = null @@ -31,6 +85,7 @@ export function setSystemPromptInjection(value: string | null): void { // Clear context caches immediately when injection changes getUserContext.cache.clear?.() getSystemContext.cache.clear?.() + getRepoMapContext.cache.clear?.() } export const getGitStatus = memoize(async (): Promise => { @@ -110,6 +165,55 @@ export const getGitStatus = memoize(async (): Promise => { } }) +export const getRepoMapContext = memoize( + async (): Promise => { + // Enable via compile-time feature flag OR runtime env var. + // The runtime env var lets users enable auto-injection without rebuilding. + const runtimeEnabled = isEnvTruthy(process.env.REPO_MAP) + if (!feature('REPO_MAP') && !runtimeEnabled) return null + if (isBareMode()) return null + if (isEnvTruthy(process.env.CLAUDE_CODE_REMOTE)) return null + + try { + const startTime = Date.now() + logForDiagnosticsNoPII('info', 'repo_map_started') + const { buildRepoMap } = await import('./context/repoMap/index.js') + const { result, timedOut } = await runRepoMapBuildWithTimeout({ + buildRepoMap, + root: getCwd(), + maxTokens: 1024, + timeoutMs: REPO_MAP_CONTEXT_TIMEOUT_MS, + }) + if (timedOut) { + logForDiagnosticsNoPII('warn', 'repo_map_timeout', { + duration_ms: Date.now() - startTime, + }) + getRepoMapContext.cache.clear?.() + getSystemContext.cache.clear?.() + return null + } + if (result === null) return null + if (!result.map || result.map.length === 0) { + return null + } + logForDiagnosticsNoPII('info', 'repo_map_completed', { + duration_ms: Date.now() - startTime, + token_count: result.tokenCount, + file_count: result.fileCount, + cache_hit: result.cacheHit, + }) + return `This is a structural map of the repository, ranked by importance. Use it to understand the codebase architecture.\n\n${result.map}` + } catch (err) { + logForDiagnosticsNoPII('warn', 'repo_map_failed', { + error: String(err), + }) + getRepoMapContext.cache.clear?.() + getSystemContext.cache.clear?.() + return null + } + }, +) + /** * This context is prepended to each conversation, and cached for the duration of the conversation. */ @@ -120,12 +224,16 @@ export const getSystemContext = memoize( const startTime = Date.now() logForDiagnosticsNoPII('info', 'system_context_started') - // Skip git status in CCR (unnecessary overhead on resume) or when git instructions are disabled - const gitStatus = + const gitStatusPromise = isEnvTruthy(process.env.CLAUDE_CODE_REMOTE) || !shouldIncludeGitInstructions() - ? null - : await getGitStatus() + ? Promise.resolve(null) + : getGitStatus() + + const [gitStatus, repoMap] = await Promise.all([ + gitStatusPromise, + getRepoMapContext(), + ]) // Include system prompt injection if set (for cache breaking, internal-only) const injection = feature('BREAK_CACHE_COMMAND') @@ -135,11 +243,13 @@ export const getSystemContext = memoize( logForDiagnosticsNoPII('info', 'system_context_completed', { duration_ms: Date.now() - startTime, has_git_status: gitStatus !== null, + has_repo_map: repoMap !== null, has_injection: injection !== null, }) return { ...(gitStatus && { gitStatus }), + ...(repoMap && { repoMap }), ...(feature('BREAK_CACHE_COMMAND') && injection ? { cacheBreaker: `[CACHE_BREAKER: ${injection}]`, diff --git a/src/context/repoMap/__fixtures__/mini-repo/fileA.ts b/src/context/repoMap/__fixtures__/mini-repo/fileA.ts new file mode 100644 index 000000000..f0ceb962e --- /dev/null +++ b/src/context/repoMap/__fixtures__/mini-repo/fileA.ts @@ -0,0 +1,29 @@ +// fileA — imports from fileB and fileC + +import { CacheLayer, buildCache } from './fileB' +import { createStore, type StoreConfig } from './fileC' + +export class AppController { + private cache: CacheLayer + private config: StoreConfig + + constructor(config: StoreConfig) { + this.cache = buildCache() + this.config = config + } + + initialize(): void { + const store = createStore() + this.cache.cacheSet('primary', store) + } + + getFromCache(key: string): unknown { + return this.cache.cacheGet(key) + } +} + +export function startApp(config: StoreConfig): AppController { + const app = new AppController(config) + app.initialize() + return app +} diff --git a/src/context/repoMap/__fixtures__/mini-repo/fileB.ts b/src/context/repoMap/__fixtures__/mini-repo/fileB.ts new file mode 100644 index 000000000..a063c01f8 --- /dev/null +++ b/src/context/repoMap/__fixtures__/mini-repo/fileB.ts @@ -0,0 +1,23 @@ +// fileB — imports from fileC + +import { DataStore, createStore } from './fileC' + +export class CacheLayer { + private store: DataStore + + constructor() { + this.store = createStore() + } + + cacheGet(key: string): unknown | undefined { + return this.store.lookup(key) + } + + cacheSet(key: string, value: unknown): void { + this.store.add(key, value) + } +} + +export function buildCache(): CacheLayer { + return new CacheLayer() +} diff --git a/src/context/repoMap/__fixtures__/mini-repo/fileC.ts b/src/context/repoMap/__fixtures__/mini-repo/fileC.ts new file mode 100644 index 000000000..8433452f6 --- /dev/null +++ b/src/context/repoMap/__fixtures__/mini-repo/fileC.ts @@ -0,0 +1,22 @@ +// fileC — the most imported module (imported by fileA and fileB) + +export class DataStore { + private items: Map = new Map() + + add(key: string, value: unknown): void { + this.items.set(key, value) + } + + lookup(key: string): unknown | undefined { + return this.items.get(key) + } +} + +export function createStore(): DataStore { + return new DataStore() +} + +export interface StoreConfig { + maxSize: number + ttl: number +} diff --git a/src/context/repoMap/__fixtures__/mini-repo/fileD.ts b/src/context/repoMap/__fixtures__/mini-repo/fileD.ts new file mode 100644 index 000000000..4f94c5dfa --- /dev/null +++ b/src/context/repoMap/__fixtures__/mini-repo/fileD.ts @@ -0,0 +1,9 @@ +// fileD — imports from fileA + +import { AppController, startApp } from './fileA' + +export function runApp(): void { + const controller: AppController = startApp({ maxSize: 100, ttl: 3600 }) + const result = controller.getFromCache('test') + console.log(result) +} diff --git a/src/context/repoMap/__fixtures__/mini-repo/fileE.ts b/src/context/repoMap/__fixtures__/mini-repo/fileE.ts new file mode 100644 index 000000000..487083c04 --- /dev/null +++ b/src/context/repoMap/__fixtures__/mini-repo/fileE.ts @@ -0,0 +1,25 @@ +// fileE — isolated, no imports from other fixture files + +export interface Logger { + log(message: string): void + warn(message: string): void + error(message: string): void +} + +export class ConsoleLogger implements Logger { + log(message: string): void { + console.log(`[LOG] ${message}`) + } + + warn(message: string): void { + console.warn(`[WARN] ${message}`) + } + + error(message: string): void { + console.error(`[ERROR] ${message}`) + } +} + +export function createLogger(): Logger { + return new ConsoleLogger() +} diff --git a/src/context/repoMap/cache.ts b/src/context/repoMap/cache.ts new file mode 100644 index 000000000..4c6be1069 --- /dev/null +++ b/src/context/repoMap/cache.ts @@ -0,0 +1,235 @@ +import { createHash } from 'crypto' +import { + existsSync, + mkdirSync, + readFileSync, + rmSync, + statSync, + writeFileSync, +} from 'fs' +import { join } from 'path' +import { getClaudeConfigHomeDir } from '../../utils/envUtils.js' +import type { + CacheData, + CacheStats, + FileStatFingerprint, + RenderedCacheEntry, + Tag, +} from './types.js' + +const CACHE_VERSION = 2 +const MAX_RENDERED_ENTRIES = 20 + +function getCacheDir(): string { + return join(getClaudeConfigHomeDir(), 'repomap-cache') +} + +function getCacheFilePath(root: string): string { + const hash = createHash('sha256').update(root).digest('hex') + return join(getCacheDir(), `${hash}.json`) +} + +function emptyCache(): CacheData { + return { version: CACHE_VERSION, entries: {}, renderedEntries: {} } +} + +function ensureCacheDir(): void { + const cacheDir = getCacheDir() + if (!existsSync(cacheDir)) { + mkdirSync(cacheDir, { recursive: true }) + } +} + +/** Load cache from disk. Returns empty cache if not found or invalid. */ +export function loadCache(root: string): CacheData { + const path = getCacheFilePath(root) + try { + const raw = readFileSync(path, 'utf-8') + const data = JSON.parse(raw) as CacheData + if (data.version !== CACHE_VERSION) { + return emptyCache() + } + data.renderedEntries ??= {} + return data + } catch { + return emptyCache() + } +} + +/** Save cache to disk. */ +export function saveCache(root: string, cache: CacheData): void { + try { + ensureCacheDir() + const path = getCacheFilePath(root) + writeFileSync(path, JSON.stringify(cache), 'utf-8') + } catch { + // Cache persistence is an optimization; repo-map results should still be + // usable in read-only home directories or sandboxed runtimes. + } +} + +/** + * Check if a file's cached entry is still valid based on mtime and size. + * Returns the cached tags if valid, null otherwise. + */ +export function getCachedTags( + cache: CacheData, + filePath: string, + root: string, + stat?: FileStatFingerprint, +): Tag[] | null { + const entry = cache.entries[filePath] + if (!entry) return null + + try { + const fingerprint = stat ?? statFile(root, filePath) + if ( + fingerprint && + fingerprint.mtimeMs === entry.mtimeMs && + fingerprint.size === entry.size + ) { + return entry.tags + } + } catch { + // File may have been deleted + } + return null +} + +/** Update the cache entry for a file. */ +export function setCachedTags( + cache: CacheData, + filePath: string, + root: string, + tags: Tag[], + stat?: FileStatFingerprint, +): void { + try { + const fingerprint = stat ?? statFile(root, filePath) + if (!fingerprint) return + cache.entries[filePath] = { + tags, + mtimeMs: fingerprint.mtimeMs, + size: fingerprint.size, + } + } catch { + // If we can't stat, don't cache + } +} + +/** + * Compute a hash of the inputs that affect the rendered map. + * Used to cache the final rendered output. + * + * The fingerprint includes per-file mtime+size so that editing a file + * (without changing the file list) invalidates the rendered cache. Without + * this, `buildRepoMap` would return a stale rendered map until manual + * `invalidateCache()`. + */ +export function computeMapHash( + files: string[], + maxTokens: number, + focusFiles: string[], + focusSymbols: string[], + root: string, + fileStats?: Map, +): string { + const sorted = [...files].sort() + const fingerprint = sorted.map(file => { + const stat = fileStats?.get(file) ?? statFile(root, file) + return stat ? `${file}:${stat.mtimeMs}:${stat.size}` : `${file}:missing` + }) + const input = JSON.stringify({ + files: sorted, + fingerprint, + maxTokens, + focusFiles: [...focusFiles].sort(), + focusSymbols: [...focusSymbols].sort(), + }) + return createHash('sha256').update(input).digest('hex') +} + +export function statFile( + root: string, + filePath: string, +): FileStatFingerprint | null { + try { + const stat = statSync(join(root, filePath)) + return { mtimeMs: stat.mtimeMs, size: stat.size } + } catch { + return null + } +} + +export function getRenderedCache( + cache: CacheData, + mapHash: string, +): RenderedCacheEntry | null { + return cache.renderedEntries[mapHash] ?? null +} + +export function setRenderedCache( + cache: CacheData, + mapHash: string, + entry: Omit, +): void { + cache.renderedEntries[mapHash] = { ...entry, createdAt: Date.now() } + pruneRenderedEntries(cache) +} + +export function pruneCache(cache: CacheData, currentFiles: string[]): void { + const current = new Set(currentFiles) + for (const key of Object.keys(cache.entries)) { + if (!current.has(key)) { + delete cache.entries[key] + } + } + pruneRenderedEntries(cache) +} + +function pruneRenderedEntries(cache: CacheData): void { + const entries = Object.entries(cache.renderedEntries) + if (entries.length <= MAX_RENDERED_ENTRIES) return + + entries + .sort((a, b) => b[1].createdAt - a[1].createdAt) + .slice(MAX_RENDERED_ENTRIES) + .forEach(([key]) => { + delete cache.renderedEntries[key] + }) +} + +/** Get cache statistics. */ +export function getCacheStats(root: string): CacheStats { + const cacheFile = getCacheFilePath(root) + const exists = existsSync(cacheFile) + let entryCount = 0 + + if (exists) { + try { + const data = JSON.parse(readFileSync(cacheFile, 'utf-8')) as CacheData + entryCount = + Object.keys(data.entries ?? {}).length + + Object.keys(data.renderedEntries ?? {}).length + } catch { + // corrupted cache + } + } + + return { + cacheDir: getCacheDir(), + cacheFile: exists ? cacheFile : null, + entryCount, + exists, + } +} + +/** Delete the cache for a repo root. */ +export function invalidateCache(root: string): void { + const path = getCacheFilePath(root) + try { + rmSync(path, { force: true }) + } catch { + // File may not exist + } +} diff --git a/src/context/repoMap/gitFiles.ts b/src/context/repoMap/gitFiles.ts new file mode 100644 index 000000000..fd30b335a --- /dev/null +++ b/src/context/repoMap/gitFiles.ts @@ -0,0 +1,121 @@ +import { execFile } from 'child_process' +import { readdirSync } from 'fs' +import { join, relative } from 'path' +import type { SupportedLanguage } from './types.js' + +const SUPPORTED_EXTENSIONS: Record = { + '.ts': 'typescript', + '.tsx': 'tsx', + '.js': 'javascript', + '.jsx': 'javascript', + '.mjs': 'javascript', + '.cjs': 'javascript', + '.py': 'python', +} + +const EXCLUDED_DIRS = new Set([ + 'node_modules', + 'dist', + '.git', + '.hg', + '.svn', + 'build', + 'out', + 'coverage', + '__pycache__', + '.next', + '.nuxt', + 'vendor', + '.worktrees', +]) + +const EXCLUDED_FILES = new Set([ + 'bun.lock', + 'bun.lockb', + 'package-lock.json', + 'yarn.lock', + 'pnpm-lock.yaml', +]) + +export function getLanguageForFile(filePath: string): SupportedLanguage | null { + const ext = filePath.substring(filePath.lastIndexOf('.')) + return SUPPORTED_EXTENSIONS[ext] ?? null +} + +export function isSupportedFile(filePath: string): boolean { + return getLanguageForFile(filePath) !== null +} + +function normalizeRepoPath(filePath: string): string { + return filePath.replace(/\\/g, '/') +} + +function gitChildEnv(): NodeJS.ProcessEnv { + const env = { ...process.env } + delete env.GIT_DIR + delete env.GIT_WORK_TREE + delete env.GIT_INDEX_FILE + return env +} + +/** List files using git ls-files. Returns relative paths. */ +function gitLsFiles(root: string): Promise { + return new Promise((resolve, reject) => { + execFile( + 'git', + ['ls-files', '-z', '--cached', '--others', '--exclude-standard'], + { cwd: root, env: gitChildEnv(), maxBuffer: 10 * 1024 * 1024 }, + (error, stdout) => { + if (error) { + reject(error) + return + } + const files = stdout + .split('\0') + .map(normalizeRepoPath) + .filter(f => f.length > 0) + resolve(files) + }, + ) + }) +} + +/** Walk directory tree manually as fallback when git is unavailable. */ +function walkDirectory(root: string, currentDir: string = root): string[] { + const results: string[] = [] + let entries + try { + entries = readdirSync(currentDir, { withFileTypes: true }) + } catch { + return results + } + + for (const entry of entries) { + const name = entry.name + if (entry.isDirectory()) { + if (!EXCLUDED_DIRS.has(name) && !name.startsWith('.')) { + results.push(...walkDirectory(root, join(currentDir, name))) + } + } else if (entry.isFile()) { + if (!EXCLUDED_FILES.has(name)) { + results.push(normalizeRepoPath(relative(root, join(currentDir, name)))) + } + } + } + return results +} + +/** + * Enumerate all supported source files in the repo. + * Tries git ls-files first, falls back to manual walk. + */ +export async function getRepoFiles(root: string): Promise { + let files: string[] + try { + files = await gitLsFiles(root) + } catch { + files = walkDirectory(root) + } + + return files.filter(isSupportedFile) +} diff --git a/src/context/repoMap/graph.ts b/src/context/repoMap/graph.ts new file mode 100644 index 000000000..6efdf50f8 --- /dev/null +++ b/src/context/repoMap/graph.ts @@ -0,0 +1,92 @@ +import Graph from 'graphology' +import type { FileTags } from './types.js' + +// Common identifiers that should contribute less weight (high IDF penalty). +const COMMON_NAMES = new Set([ + 'map', 'get', 'set', 'value', 'key', 'data', 'result', 'error', + 'name', 'type', 'id', 'index', 'item', 'items', 'list', 'options', + 'config', 'args', 'params', 'props', 'state', 'event', 'callback', + 'handler', 'fn', 'func', 'self', 'this', 'ctx', 'context', 'req', + 'res', 'next', 'err', 'msg', 'obj', 'arr', 'str', 'num', 'val', + 'init', 'start', 'stop', 'run', 'main', 'test', 'setup', 'teardown', + 'constructor', 'toString', 'valueOf', 'length', 'size', 'count', + 'push', 'pop', 'shift', 'filter', 'reduce', 'forEach', 'find', + 'log', 'warn', 'info', 'debug', 'trace', +]) +const MAX_DEFINITION_FANOUT = 100 + +/** + * Build a directed graph from file tags. + * + * Nodes are file paths. An edge from A to B means file A references + * a symbol defined in file B. Edge weight = refCount * idf(symbolName). + */ +export function buildGraph(allFileTags: FileTags[]): Graph { + const graph = new Graph({ multi: false, type: 'directed' }) + + // Build a map from symbol name → files that define it + const defIndex = new Map>() + for (const ft of allFileTags) { + for (const tag of ft.tags) { + if (tag.kind === 'def') { + let files = defIndex.get(tag.name) + if (!files) { + files = new Set() + defIndex.set(tag.name, files) + } + files.add(ft.path) + } + } + } + + // Compute IDF: log(totalFiles / filesDefiningSymbol) + // Common names get an extra penalty + const totalFiles = allFileTags.length + function idf(symbolName: string): number { + const defFiles = defIndex.get(symbolName) + const docFreq = defFiles ? defFiles.size : 1 + const rawIdf = Math.log(totalFiles / docFreq) + return COMMON_NAMES.has(symbolName) ? rawIdf * 0.1 : rawIdf + } + + // Add all files as nodes + for (const ft of allFileTags) { + if (!graph.hasNode(ft.path)) { + graph.addNode(ft.path) + } + } + + // Build edges: for each ref in a file, find where it's defined + for (const ft of allFileTags) { + // Count refs per target file + const edgeWeights = new Map() + + for (const tag of ft.tags) { + if (tag.kind !== 'ref') continue + + const defFiles = defIndex.get(tag.name) + if (!defFiles) continue + if (defFiles.size > MAX_DEFINITION_FANOUT) continue + + const weight = idf(tag.name) + if (!Number.isFinite(weight) || weight <= 0) continue + for (const defFile of defFiles) { + if (defFile === ft.path) continue // skip self-references + const current = edgeWeights.get(defFile) ?? 0 + edgeWeights.set(defFile, current + weight) + } + } + + for (const [target, weight] of edgeWeights) { + if (!Number.isFinite(weight) || weight <= 0) continue + if (graph.hasEdge(ft.path, target)) { + graph.setEdgeAttribute(ft.path, target, 'weight', + graph.getEdgeAttribute(ft.path, target, 'weight') + weight) + } else { + graph.addEdge(ft.path, target, { weight }) + } + } + } + + return graph +} diff --git a/src/context/repoMap/index.ts b/src/context/repoMap/index.ts new file mode 100644 index 000000000..0573956dd --- /dev/null +++ b/src/context/repoMap/index.ts @@ -0,0 +1,217 @@ +import { + computeMapHash, + getCachedTags, + getCacheStats as getCacheStatsImpl, + getRenderedCache, + invalidateCache as invalidateCacheImpl, + loadCache, + pruneCache, + saveCache, + setCachedTags, + setRenderedCache, + statFile, +} from './cache.js' +import { getRepoFiles } from './gitFiles.js' +import { buildGraph } from './graph.js' +import { rankFiles } from './pagerank.js' +import { renderMap } from './renderer.js' +import { extractTags } from './symbolExtractor.js' +import type { + CacheData, + CacheStats, + FileStatFingerprint, + FileTags, + RepoMapOptions, + RepoMapResult, +} from './types.js' + +const DEFAULT_MAX_TOKENS = 2048 +const TAG_EXTRACTION_BATCH_SIZE = 50 + +export async function extractTagsWithCache({ + files, + root, + cache, + fileStats, + shouldContinue, +}: { + files: string[] + root: string + cache: CacheData + fileStats?: Map + shouldContinue?: () => void +}): Promise { + const fileTagsByPath = new Map() + const uncachedFiles: string[] = [] + + for (const file of files) { + shouldContinue?.() + const cachedTags = getCachedTags( + cache, + file, + root, + fileStats?.get(file) ?? undefined, + ) + if (cachedTags) { + fileTagsByPath.set(file, { path: file, tags: cachedTags }) + } else { + uncachedFiles.push(file) + } + } + + for (let i = 0; i < uncachedFiles.length; i += TAG_EXTRACTION_BATCH_SIZE) { + shouldContinue?.() + const batch = uncachedFiles.slice(i, i + TAG_EXTRACTION_BATCH_SIZE) + const results = await Promise.all( + batch.map(file => extractTags(file, root).catch(() => null)), + ) + for (const fileTags of results) { + if (!fileTags) continue + fileTagsByPath.set(fileTags.path, fileTags) + setCachedTags( + cache, + fileTags.path, + root, + fileTags.tags, + fileStats?.get(fileTags.path) ?? undefined, + ) + } + } + + return files.flatMap(file => { + const fileTags = fileTagsByPath.get(file) + return fileTags ? [fileTags] : [] + }) +} + +/** + * Build a structural summary of a code repository. + * + * Walks the repo, extracts symbols via tree-sitter, builds an IDF-weighted + * reference graph, ranks files with PageRank, and renders a token-budgeted + * structural summary. + */ +export async function buildRepoMap(options: RepoMapOptions = {}): Promise { + const startTime = Date.now() + const root = options.root ?? process.cwd() + const maxTokens = options.maxTokens ?? DEFAULT_MAX_TOKENS + const focusFiles = options.focusFiles ?? [] + const focusSymbols = options.focusSymbols ?? [] + const shouldContinue = options.shouldContinue + + // Get files + shouldContinue?.() + const files = options.files ?? await getRepoFiles(root) + shouldContinue?.() + const totalFileCount = files.length + + const fileStats = new Map(files.map(file => [file, statFile(root, file)])) + const existingFileStats = new Map( + [...fileStats.entries()].filter( + (entry): entry is [string, NonNullable] => + entry[1] !== null, + ), + ) + const mapHash = computeMapHash( + files, + maxTokens, + focusFiles, + focusSymbols, + root, + existingFileStats, + ) + const cache = loadCache(root) + pruneCache(cache, files) + + const renderedEntry = getRenderedCache(cache, mapHash) + if (renderedEntry) { + return { + map: renderedEntry.map, + cacheHit: true, + buildTimeMs: Date.now() - startTime, + fileCount: renderedEntry.fileCount, + totalFileCount, + tokenCount: renderedEntry.tokenCount, + } + } + + const allFileTags = await extractTagsWithCache({ + files, + root, + cache, + fileStats, + shouldContinue, + }) + + const resolvedFocusFiles = resolveFocusFiles({ + focusFiles, + focusSymbols, + allFileTags, + }) + shouldContinue?.() + + // Build graph and rank + const graph = buildGraph(allFileTags) + const ranked = rankFiles(graph, resolvedFocusFiles) + + // Build a lookup map + const fileTagsMap = new Map() + for (const ft of allFileTags) { + fileTagsMap.set(ft.path, ft) + } + + // Render + const { map, tokenCount, fileCount } = renderMap(ranked, fileTagsMap, maxTokens) + + setRenderedCache(cache, mapHash, { map, fileCount, tokenCount }) + + saveCache(root, cache) + + return { + map, + cacheHit: false, + buildTimeMs: Date.now() - startTime, + fileCount, + totalFileCount, + tokenCount, + } +} + +/** Invalidate the disk cache for a given repo root. */ +export function invalidateCache(root?: string): void { + invalidateCacheImpl(root ?? process.cwd()) +} + +/** Get cache statistics for a given repo root. */ +export function getCacheStats(root?: string): CacheStats { + return getCacheStatsImpl(root ?? process.cwd()) +} + +function resolveFocusFiles({ + focusFiles, + focusSymbols, + allFileTags, +}: { + focusFiles: string[] + focusSymbols: string[] + allFileTags: FileTags[] +}): string[] { + if (focusSymbols.length === 0) return focusFiles + + const symbolSet = new Set(focusSymbols) + const symbolFiles: string[] = [] + + for (const result of allFileTags) { + const hasMatch = result.tags.some( + tag => tag.kind === 'def' && symbolSet.has(tag.name), + ) + if (hasMatch) { + symbolFiles.push(result.path) + } + } + + return [...focusFiles, ...symbolFiles] +} + +// Re-export types for convenience +export type { RepoMapOptions, RepoMapResult, CacheStats } from './types.js' diff --git a/src/context/repoMap/pagerank.ts b/src/context/repoMap/pagerank.ts new file mode 100644 index 000000000..f01c70133 --- /dev/null +++ b/src/context/repoMap/pagerank.ts @@ -0,0 +1,84 @@ +import type Graph from 'graphology' +import pagerank from 'graphology-metrics/centrality/pagerank' + +export interface RankedFile { + path: string + score: number +} + +function normalizeFocusPath(path: string): string { + return path.trim().replace(/\\/g, '/').replace(/^\.\//, '') +} + +function expandFocusFiles(graph: Graph, focusFiles: string[]): string[] { + const expanded = new Set() + + for (const rawFocus of focusFiles) { + const focus = normalizeFocusPath(rawFocus) + if (!focus) continue + + if (graph.hasNode(focus)) { + expanded.add(focus) + } + + const prefix = focus.endsWith('/') ? focus : `${focus}/` + graph.forEachNode((node) => { + if (node.startsWith(prefix)) { + expanded.add(node) + } + }) + } + + return [...expanded] +} + +/** + * Run PageRank on the file reference graph. + * + * PageRank runs on the full graph, then focusFiles and their neighbors get + * a post-processing boost so they rank higher in the rendered map. + * + * Returns files sorted by score descending. + */ +export function rankFiles( + graph: Graph, + focusFiles: string[] = [], +): RankedFile[] { + if (graph.order === 0) return [] + + const expandedFocusFiles = expandFocusFiles(graph, focusFiles) + const hasPersonalization = expandedFocusFiles.length > 0 + + // graphology-metrics PageRank accepts getEdgeWeight option + const scores: Record = pagerank(graph, { + alpha: 0.85, + maxIterations: 100, + tolerance: 1e-6, + getEdgeWeight: 'weight', + }) + + // Apply focus boost post-hoc if focus files are specified + if (hasPersonalization) { + for (const file of expandedFocusFiles) { + if (scores[file] !== undefined) { + scores[file] *= 100 + } + } + + // Also boost direct neighbors of focus files + for (const file of expandedFocusFiles) { + if (!graph.hasNode(file)) continue + graph.forEachNeighbor(file, (neighbor) => { + if (scores[neighbor] !== undefined) { + scores[neighbor] *= 10 + } + }) + } + } + + const ranked: RankedFile[] = Object.entries(scores) + .map(([path, score]) => ({ path, score })) + .sort((a, b) => b.score - a.score) + + return ranked +} diff --git a/src/context/repoMap/parser.ts b/src/context/repoMap/parser.ts new file mode 100644 index 000000000..a8308d77c --- /dev/null +++ b/src/context/repoMap/parser.ts @@ -0,0 +1,168 @@ +import { join, posix, resolve, win32 } from 'path' +import { fileURLToPath } from 'url' +import { createRequire } from 'module' +import { getBundledQuery } from './queries.js' +import type { SupportedLanguage } from './types.js' +import type { Language } from 'web-tree-sitter' + +const require = createRequire(import.meta.url) +const __filename = fileURLToPath(import.meta.url) +const __projectRoot = resolveProjectRoot(__filename) + +export function resolveProjectRoot(filePath: string): string { + const pathApi = filePath.includes('\\') ? win32 : posix + const normalized = filePath.replace(/\\/g, '/') + const isSourcePath = normalized.includes('/src/context/repoMap/') + return isSourcePath + ? pathApi.join(pathApi.dirname(filePath), '../../../') + : pathApi.join(pathApi.dirname(filePath), '../') +} + +// web-tree-sitter types +type TreeSitterParser = { + parse(input: string): { rootNode: unknown } + setLanguage(lang: unknown): void + delete(): void +} + +// The actual module exports { Parser, Language } as named exports +let ParserClass: (new () => TreeSitterParser) & { + init(opts?: { locateFile?: (file: string) => string }): Promise +} | null = null +let LanguageLoader: { + load(path: string | Uint8Array): Promise +} | null = null + +let initialized = false +let initPromise: Promise | null = null +const languageCache = new Map() +const queryCache = new Map() + +/** Resolve the path to the tree-sitter WASM file. */ +function getTreeSitterWasmPath(): string { + // Try require.resolve first (works in source mode with node_modules) + try { + const webTsDir = resolve( + require.resolve('web-tree-sitter/package.json'), + '..', + ) + return join(webTsDir, 'tree-sitter.wasm') + } catch { + // Fallback: relative to project root + return join(__projectRoot, 'node_modules', 'web-tree-sitter', 'tree-sitter.wasm') + } +} + +/** Resolve the path to a language WASM grammar file. */ +function getLanguageWasmPath(language: SupportedLanguage): string { + const wasmName = language === 'typescript' ? 'tree-sitter-typescript' : + language === 'tsx' ? 'tree-sitter-tsx' : + language === 'javascript' ? 'tree-sitter-javascript' : + `tree-sitter-${language}` + + try { + const wasmDir = resolve( + require.resolve('tree-sitter-wasms/package.json'), + '..', + 'out', + ) + return join(wasmDir, `${wasmName}.wasm`) + } catch { + return join(__projectRoot, 'node_modules', 'tree-sitter-wasms', 'out', `${wasmName}.wasm`) + } +} + +/** Initialize the tree-sitter WASM module. */ +export async function initParser(): Promise { + if (initialized) return + if (initPromise) return initPromise + + initPromise = (async () => { + const mod = await import('web-tree-sitter') + ParserClass = mod.Parser as typeof ParserClass + LanguageLoader = mod.Language as typeof LanguageLoader + + const wasmPath = getTreeSitterWasmPath() + await ParserClass!.init({ + locateFile: () => wasmPath, + }) + initialized = true + })() + + try { + await initPromise + } catch (err) { + // eslint-disable-next-line no-console + console.error('[repoMap] Failed to initialize tree-sitter:', err) + initPromise = null + throw err + } +} + +/** Load a language grammar. Cached after first load. */ +export async function loadLanguage(language: SupportedLanguage): Promise { + if (languageCache.has(language)) { + return languageCache.get(language)! + } + + if (!initialized) { + await initParser() + } + + try { + const wasmPath = getLanguageWasmPath(language) + const lang = await LanguageLoader!.load(wasmPath) + languageCache.set(language, lang) + return lang + } catch (err) { + // eslint-disable-next-line no-console + console.error(`[repoMap] Failed to load ${language} grammar:`, err) + return null + } +} + +/** Load the tag query for a language. Cached after first load. + * + * Reads from bundled string constants so the queries ship inside dist/cli.mjs + * — the .scm files in ./queries/ are kept as canonical source only, and are + * not part of the published npm package. + */ +export function loadQuery(language: SupportedLanguage): string | null { + if (queryCache.has(language)) { + return queryCache.get(language)! + } + + const content = getBundledQuery(language) + if (content === null) return null + + queryCache.set(language, content) + return content +} + +/** Create a new parser instance with the given language set. */ +export async function createParser(language: SupportedLanguage): Promise { + if (!initialized) { + await initParser() + } + + const lang = await loadLanguage(language) + if (!lang) return null + + try { + const parser = new ParserClass!() + parser.setLanguage(lang) + return parser + } catch { + return null + } +} + +/** Clear all caches (useful for testing). */ +export function clearParserCaches(): void { + languageCache.clear() + queryCache.clear() + initialized = false + initPromise = null + ParserClass = null + LanguageLoader = null +} diff --git a/src/context/repoMap/queries.test.ts b/src/context/repoMap/queries.test.ts new file mode 100644 index 000000000..d51396609 --- /dev/null +++ b/src/context/repoMap/queries.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, test } from 'bun:test' +import { readFileSync } from 'fs' +import { join } from 'path' +import { fileURLToPath } from 'url' +import { getBundledQuery } from './queries.js' +import type { SupportedLanguage } from './types.js' + +const __dirname = join(fileURLToPath(import.meta.url), '..') + +describe('bundled query drift guard', () => { + test.each(['typescript', 'javascript', 'python'])( + '%s: bundled query matches the .scm source file byte-for-byte', + (language) => { + // Normalize CRLF → LF so Windows checkouts (where Git may convert line + // endings) still pass the byte-for-byte drift guard against the + // LF-only TypeScript string constants in queries.ts. + const fromFile = readFileSync( + join(__dirname, 'queries', `${language}-tags.scm`), + 'utf-8', + ).replace(/\r\n/g, '\n') + const bundled = getBundledQuery(language) + expect(bundled).not.toBeNull() + expect(bundled).toBe(fromFile) + }, + ) + + test('tsx reuses the TypeScript query with the TSX grammar', () => { + expect(getBundledQuery('tsx')).toBe(getBundledQuery('typescript')) + }) + + test('returns null for unknown language', () => { + expect(getBundledQuery('unknown' as SupportedLanguage)).toBeNull() + }) +}) diff --git a/src/context/repoMap/queries.ts b/src/context/repoMap/queries.ts new file mode 100644 index 000000000..f4767dbce --- /dev/null +++ b/src/context/repoMap/queries.ts @@ -0,0 +1,186 @@ +/** + * Bundled tree-sitter tag queries. + * + * The .scm files in ./queries/ are the canonical source-of-truth (kept for + * readability and Aider attribution), but the runtime reads from these inlined + * string constants so the queries are bundled into dist/cli.mjs and ship with + * the published npm package — the .scm files themselves are not in the + * package.json `files` allowlist and would otherwise be missing post-install. + * + * If you edit a .scm file, mirror the change here. A unit test guards drift. + */ + +import type { SupportedLanguage } from './types.js' + +const TYPESCRIPT_TAGS = `; Source: https://github.com/Aider-AI/aider/blob/main/aider/queries/tree-sitter-languages/typescript-tags.scm +; License: MIT (Apache-2.0 dual) — see https://github.com/Aider-AI/aider/blob/main/LICENSE +; Copied for use in openclaude's repo-map feature. + +(function_signature + name: (identifier) @name.definition.function) @definition.function + +(method_signature + name: (property_identifier) @name.definition.method) @definition.method + +(abstract_method_signature + name: (property_identifier) @name.definition.method) @definition.method + +(abstract_class_declaration + name: (type_identifier) @name.definition.class) @definition.class + +(module + name: (identifier) @name.definition.module) @definition.module + +(interface_declaration + name: (type_identifier) @name.definition.interface) @definition.interface + +(type_annotation + (type_identifier) @name.reference.type) @reference.type + +(new_expression + constructor: (identifier) @name.reference.class) @reference.class + +(function_declaration + name: (identifier) @name.definition.function) @definition.function + +(lexical_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + +(variable_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + +(method_definition + name: (property_identifier) @name.definition.method) @definition.method + +(class_declaration + name: (type_identifier) @name.definition.class) @definition.class + +(type_alias_declaration + name: (type_identifier) @name.definition.type) @definition.type + +(enum_declaration + name: (identifier) @name.definition.enum) @definition.enum +` + +const JAVASCRIPT_TAGS = `; Source: https://github.com/Aider-AI/aider/blob/main/aider/queries/tree-sitter-languages/javascript-tags.scm +; License: MIT (Apache-2.0 dual) — see https://github.com/Aider-AI/aider/blob/main/LICENSE +; Copied for use in openclaude's repo-map feature. + +( + (comment)* @doc + . + (method_definition + name: (property_identifier) @name.definition.method) @definition.method + (#not-eq? @name.definition.method "constructor") + (#strip! @doc "^[\\\\s\\\\*/]+|^[\\\\s\\\\*/]$") + (#select-adjacent! @doc @definition.method) +) + +( + (comment)* @doc + . + [ + (class + name: (_) @name.definition.class) + (class_declaration + name: (_) @name.definition.class) + ] @definition.class + (#strip! @doc "^[\\\\s\\\\*/]+|^[\\\\s\\\\*/]$") + (#select-adjacent! @doc @definition.class) +) + +( + (comment)* @doc + . + [ + (function_declaration + name: (identifier) @name.definition.function) + (generator_function_declaration + name: (identifier) @name.definition.function) + ] @definition.function + (#strip! @doc "^[\\\\s\\\\*/]+|^[\\\\s\\\\*/]$") + (#select-adjacent! @doc @definition.function) +) + +( + (comment)* @doc + . + (lexical_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + (#strip! @doc "^[\\\\s\\\\*/]+|^[\\\\s\\\\*/]$") + (#select-adjacent! @doc @definition.function) +) + +( + (comment)* @doc + . + (variable_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + (#strip! @doc "^[\\\\s\\\\*/]+|^[\\\\s\\\\*/]$") + (#select-adjacent! @doc @definition.function) +) + +(assignment_expression + left: [ + (identifier) @name.definition.function + (member_expression + property: (property_identifier) @name.definition.function) + ] + right: [(arrow_function) (function_expression)] +) @definition.function + +(pair + key: (property_identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function + +( + (call_expression + function: (identifier) @name.reference.call) @reference.call + (#not-match? @name.reference.call "^(require)$") +) + +(call_expression + function: (member_expression + property: (property_identifier) @name.reference.call) + arguments: (_) @reference.call) + +(new_expression + constructor: (_) @name.reference.class) @reference.class +` + +const PYTHON_TAGS = `; Source: https://github.com/Aider-AI/aider/blob/main/aider/queries/tree-sitter-languages/python-tags.scm +; License: MIT (Apache-2.0 dual) — see https://github.com/Aider-AI/aider/blob/main/LICENSE +; Copied for use in openclaude's repo-map feature. + +(class_definition + name: (identifier) @name.definition.class) @definition.class + +(function_definition + name: (identifier) @name.definition.function) @definition.function + +(call + function: [ + (identifier) @name.reference.call + (attribute + attribute: (identifier) @name.reference.call) + ]) @reference.call +` + +const QUERIES: Record = { + typescript: TYPESCRIPT_TAGS, + tsx: TYPESCRIPT_TAGS, + javascript: JAVASCRIPT_TAGS, + python: PYTHON_TAGS, +} + +export function getBundledQuery(language: SupportedLanguage): string | null { + return QUERIES[language] ?? null +} diff --git a/src/context/repoMap/queries/javascript-tags.scm b/src/context/repoMap/queries/javascript-tags.scm new file mode 100644 index 000000000..6c80a9bd4 --- /dev/null +++ b/src/context/repoMap/queries/javascript-tags.scm @@ -0,0 +1,88 @@ +; Source: https://github.com/Aider-AI/aider/blob/main/aider/queries/tree-sitter-languages/javascript-tags.scm +; License: MIT (Apache-2.0 dual) — see https://github.com/Aider-AI/aider/blob/main/LICENSE +; Copied for use in openclaude's repo-map feature. + +( + (comment)* @doc + . + (method_definition + name: (property_identifier) @name.definition.method) @definition.method + (#not-eq? @name.definition.method "constructor") + (#strip! @doc "^[\\s\\*/]+|^[\\s\\*/]$") + (#select-adjacent! @doc @definition.method) +) + +( + (comment)* @doc + . + [ + (class + name: (_) @name.definition.class) + (class_declaration + name: (_) @name.definition.class) + ] @definition.class + (#strip! @doc "^[\\s\\*/]+|^[\\s\\*/]$") + (#select-adjacent! @doc @definition.class) +) + +( + (comment)* @doc + . + [ + (function_declaration + name: (identifier) @name.definition.function) + (generator_function_declaration + name: (identifier) @name.definition.function) + ] @definition.function + (#strip! @doc "^[\\s\\*/]+|^[\\s\\*/]$") + (#select-adjacent! @doc @definition.function) +) + +( + (comment)* @doc + . + (lexical_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + (#strip! @doc "^[\\s\\*/]+|^[\\s\\*/]$") + (#select-adjacent! @doc @definition.function) +) + +( + (comment)* @doc + . + (variable_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + (#strip! @doc "^[\\s\\*/]+|^[\\s\\*/]$") + (#select-adjacent! @doc @definition.function) +) + +(assignment_expression + left: [ + (identifier) @name.definition.function + (member_expression + property: (property_identifier) @name.definition.function) + ] + right: [(arrow_function) (function_expression)] +) @definition.function + +(pair + key: (property_identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function + +( + (call_expression + function: (identifier) @name.reference.call) @reference.call + (#not-match? @name.reference.call "^(require)$") +) + +(call_expression + function: (member_expression + property: (property_identifier) @name.reference.call) + arguments: (_) @reference.call) + +(new_expression + constructor: (_) @name.reference.class) @reference.class diff --git a/src/context/repoMap/queries/python-tags.scm b/src/context/repoMap/queries/python-tags.scm new file mode 100644 index 000000000..30959d631 --- /dev/null +++ b/src/context/repoMap/queries/python-tags.scm @@ -0,0 +1,16 @@ +; Source: https://github.com/Aider-AI/aider/blob/main/aider/queries/tree-sitter-languages/python-tags.scm +; License: MIT (Apache-2.0 dual) — see https://github.com/Aider-AI/aider/blob/main/LICENSE +; Copied for use in openclaude's repo-map feature. + +(class_definition + name: (identifier) @name.definition.class) @definition.class + +(function_definition + name: (identifier) @name.definition.function) @definition.function + +(call + function: [ + (identifier) @name.reference.call + (attribute + attribute: (identifier) @name.reference.call) + ]) @reference.call diff --git a/src/context/repoMap/queries/typescript-tags.scm b/src/context/repoMap/queries/typescript-tags.scm new file mode 100644 index 000000000..b2cb24da5 --- /dev/null +++ b/src/context/repoMap/queries/typescript-tags.scm @@ -0,0 +1,52 @@ +; Source: https://github.com/Aider-AI/aider/blob/main/aider/queries/tree-sitter-languages/typescript-tags.scm +; License: MIT (Apache-2.0 dual) — see https://github.com/Aider-AI/aider/blob/main/LICENSE +; Copied for use in openclaude's repo-map feature. + +(function_signature + name: (identifier) @name.definition.function) @definition.function + +(method_signature + name: (property_identifier) @name.definition.method) @definition.method + +(abstract_method_signature + name: (property_identifier) @name.definition.method) @definition.method + +(abstract_class_declaration + name: (type_identifier) @name.definition.class) @definition.class + +(module + name: (identifier) @name.definition.module) @definition.module + +(interface_declaration + name: (type_identifier) @name.definition.interface) @definition.interface + +(type_annotation + (type_identifier) @name.reference.type) @reference.type + +(new_expression + constructor: (identifier) @name.reference.class) @reference.class + +(function_declaration + name: (identifier) @name.definition.function) @definition.function + +(lexical_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + +(variable_declaration + (variable_declarator + name: (identifier) @name.definition.function + value: [(arrow_function) (function_expression)]) @definition.function) + +(method_definition + name: (property_identifier) @name.definition.method) @definition.method + +(class_declaration + name: (type_identifier) @name.definition.class) @definition.class + +(type_alias_declaration + name: (type_identifier) @name.definition.type) @definition.type + +(enum_declaration + name: (identifier) @name.definition.enum) @definition.enum diff --git a/src/context/repoMap/renderer.ts b/src/context/repoMap/renderer.ts new file mode 100644 index 000000000..4e1784b5d --- /dev/null +++ b/src/context/repoMap/renderer.ts @@ -0,0 +1,72 @@ +import type { FileTags, Tag } from './types.js' +import type { RankedFile } from './pagerank.js' +import { countTokens } from './tokenize.js' + +/** + * Render a token-budgeted repo map from ranked files and their tags. + * + * Format per file: + * path/to/file.ts: + * ⋮ + * signature line for def 1 + * ⋮ + * signature line for def 2 + * ⋮ + * + * Files that don't fit within the budget are dropped entirely. + */ +export function renderMap( + rankedFiles: RankedFile[], + fileTagsMap: Map, + maxTokens: number, +): { map: string; tokenCount: number; fileCount: number } { + const sections: string[] = [] + let currentTokens = 0 + let fileCount = 0 + + for (const { path } of rankedFiles) { + const ft = fileTagsMap.get(path) + if (!ft) continue + + // Only include definitions in the rendered output + const defs = ft.tags + .filter(t => t.kind === 'def') + .sort((a, b) => a.line - b.line) + + if (defs.length === 0) continue + + const section = renderFileSection(path, defs) + const sectionTokens = countTokens(section) + + // Would this section bust the budget? + if (currentTokens + sectionTokens > maxTokens) { + // Don't include partial files, but keep trying smaller lower-ranked files. + continue + } + + sections.push(section) + currentTokens += sectionTokens + fileCount++ + } + + const map = sections.join('\n') + return { map, tokenCount: currentTokens, fileCount } +} + +function renderFileSection(path: string, defs: Tag[]): string { + const lines: string[] = [`${path}:`] + let lastLine = 0 + + for (const def of defs) { + // Add elision marker if there's a gap + if (def.line > lastLine + 1) { + lines.push('⋮') + } + lines.push(` ${def.signature}`) + lastLine = def.line + } + + // Trailing elision marker + lines.push('⋮') + return lines.join('\n') +} diff --git a/src/context/repoMap/repoMap.test.ts b/src/context/repoMap/repoMap.test.ts new file mode 100644 index 000000000..18506cb79 --- /dev/null +++ b/src/context/repoMap/repoMap.test.ts @@ -0,0 +1,866 @@ +import { afterEach, beforeAll, describe, expect, test } from 'bun:test' +import { execFileSync, spawnSync } from 'child_process' +import { cpSync, mkdtempSync, rmSync, utimesSync, writeFileSync } from 'fs' +import { tmpdir } from 'os' +import { join, win32 } from 'path' +import { fileURLToPath } from 'url' +import { + getClaudeConfigHomeDir, + getClaudeConfigHomeDirOverrideForTesting, + setClaudeConfigHomeDirForTesting, +} from '../../utils/envUtils.js' +import { + acquireSharedMutationLock, + releaseSharedMutationLock, +} from '../../test/sharedMutationLock.js' +import { + invalidateCache, + buildRepoMap, + extractTagsWithCache, +} from './index.js' +import { + getCachedTags, + loadCache, + saveCache, + setCachedTags, + statFile, +} from './cache.js' +import { clearSymbolExtractorCaches, extractTags } from './symbolExtractor.js' +import { buildGraph } from './graph.js' +import { getRepoFiles } from './gitFiles.js' +import { rankFiles } from './pagerank.js' +import { initParser, resolveProjectRoot } from './parser.js' +import { renderMap } from './renderer.js' +import { countTokens } from './tokenize.js' +import type { FileTags } from './types.js' + +const FIXTURE_ROOT = join(import.meta.dir, '__fixtures__', 'mini-repo') +const FIXTURE_FILES = ['fileA.ts', 'fileB.ts', 'fileC.ts', 'fileD.ts', 'fileE.ts'] +const REPO_ROOT = fileURLToPath(new URL('../../..', import.meta.url)) + +beforeAll(async () => { + await initParser() +}) + +async function withWritableConfigHome( + callback: (configDir: string) => Promise, +): Promise { + await acquireSharedMutationLock('context/repoMap/repoMap.test.ts config home') + const previousConfigHomeOverride = getClaudeConfigHomeDirOverrideForTesting() + let configDir: string | undefined + + try { + configDir = mkdtempSync(join(tmpdir(), 'repomap-test-config-')) + setClaudeConfigHomeDirForTesting(configDir) + getClaudeConfigHomeDir.cache?.clear?.() + return await callback(configDir) + } finally { + setClaudeConfigHomeDirForTesting(previousConfigHomeOverride) + getClaudeConfigHomeDir.cache?.clear?.() + if (configDir) { + rmSync(configDir, { recursive: true, force: true }) + } + releaseSharedMutationLock() + } +} + +// Clean up cache between tests to avoid cross-test interference +afterEach(() => { + invalidateCache(FIXTURE_ROOT) + clearSymbolExtractorCaches() +}) + +describe('symbol extraction', () => { + test('extracts function and class defs from a TypeScript file', async () => { + const result = await extractTags('fileC.ts', FIXTURE_ROOT) + expect(result).not.toBeNull() + + const defs = result!.tags.filter(t => t.kind === 'def') + const defNames = defs.map(t => t.name) + + expect(defNames).toContain('DataStore') + expect(defNames).toContain('createStore') + expect(defNames).toContain('StoreConfig') + expect(defs.filter(d => d.name === 'StoreConfig')).toHaveLength(1) + + // All defs should have kind='def' + for (const d of defs) { + expect(d.kind).toBe('def') + } + }) + + test('extracts TypeScript const arrow and function expression definitions', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-ts-const-functions-')) + try { + writeFileSync( + join(tempDir, 'actions.ts'), + [ + 'export const makeUser = () => ({ id: 1 })', + 'export const named = function named() { return 1 }', + '', + ].join('\n'), + ) + + const result = await extractTags('actions.ts', tempDir) + expect(result).not.toBeNull() + + const defs = result!.tags + .filter(tag => tag.kind === 'def') + .map(tag => tag.name) + + expect(defs).toContain('makeUser') + expect(defs).toContain('named') + + const map = await buildRepoMap({ + root: tempDir, + files: ['actions.ts'], + maxTokens: 1024, + }) + expect(map.map).toContain('makeUser') + expect(map.map).toContain('named') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('extracts references to imported symbols', async () => { + const result = await extractTags('fileA.ts', FIXTURE_ROOT) + expect(result).not.toBeNull() + + const refs = result!.tags.filter(t => t.kind === 'ref') + const refNames = refs.map(t => t.name) + + // fileA imports CacheLayer from fileB and StoreConfig from fileC + expect(refNames).toContain('CacheLayer') + expect(refNames).toContain('StoreConfig') + }) + + test('parses TSX files with the TSX grammar', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-tsx-')) + try { + writeFileSync( + join(tempDir, 'Widget.tsx'), + [ + 'export function Widget(): JSX.Element {', + ' return
Hello
', + '}', + '', + ].join('\n'), + ) + + const result = await extractTags('Widget.tsx', tempDir) + expect(result).not.toBeNull() + expect(result!.tags.some(tag => + tag.kind === 'def' && tag.name === 'Widget', + )).toBe(true) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + } + }) + + test('extracts definitions and references from JavaScript files', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-js-')) + try { + writeFileSync( + join(tempDir, 'main.js'), + [ + 'class Widget {', + ' render() { return helper() }', + '}', + 'function helper() { return new Widget() }', + 'const makeWidget = () => new Widget()', + 'exports.fromCommonJs = function() { return makeWidget() }', + '', + ].join('\n'), + ) + + const result = await extractTags('main.js', tempDir) + expect(result).not.toBeNull() + + const defs = result!.tags + .filter(tag => tag.kind === 'def') + .map(tag => tag.name) + const refs = result!.tags + .filter(tag => tag.kind === 'ref') + .map(tag => tag.name) + + expect(defs).toContain('Widget') + expect(defs).toContain('helper') + expect(defs).toContain('makeWidget') + expect(defs).toContain('fromCommonJs') + expect(refs).toContain('Widget') + expect(refs).toContain('helper') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('builds a non-empty map for JavaScript-only repos', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-js-only-')) + try { + writeFileSync( + join(tempDir, 'main.js'), + [ + 'export class JavaScriptOnly {}', + 'export function createJavaScriptOnly() {', + ' return new JavaScriptOnly()', + '}', + '', + ].join('\n'), + ) + + const result = await buildRepoMap({ + root: tempDir, + maxTokens: 1024, + files: ['main.js'], + }) + + expect(result.map).toContain('main.js:') + expect(result.map).toContain('JavaScriptOnly') + expect(result.fileCount).toBe(1) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('shares concurrent query loads for the same language', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-js-concurrent-')) + try { + const files = Array.from({ length: 20 }, (_, i) => `file${i}.js`) + for (const file of files) { + writeFileSync( + join(tempDir, file), + `export function ${file.replace('.js', '')}() { return 1 }\n`, + ) + } + + const results = await Promise.all( + files.map(file => extractTags(file, tempDir)), + ) + + expect(results.every(result => result !== null)).toBe(true) + expect( + results.every(result => + result!.tags.some(tag => tag.kind === 'def'), + ), + ).toBe(true) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) +}) + +describe('graph', () => { + test('builds edges between files that reference each other\'s symbols', async () => { + const allTags: FileTags[] = [] + for (const f of FIXTURE_FILES) { + const tags = await extractTags(f, FIXTURE_ROOT) + if (tags) allTags.push(tags) + } + + const graph = buildGraph(allTags) + + // fileA imports from fileB (references CacheLayer defined in fileB) + expect(graph.hasEdge('fileA.ts', 'fileB.ts')).toBe(true) + + // fileA imports from fileC (references StoreConfig, DataStore defined in fileC) + expect(graph.hasEdge('fileA.ts', 'fileC.ts')).toBe(true) + + // fileB imports from fileC (references DataStore defined in fileC) + expect(graph.hasEdge('fileB.ts', 'fileC.ts')).toBe(true) + + // fileD imports from fileA + expect(graph.hasEdge('fileD.ts', 'fileA.ts')).toBe(true) + + // fileE is isolated — no edges to/from it + expect(graph.degree('fileE.ts')).toBe(0) + }) + + test('skips edge creation when one symbol is defined in too many files', () => { + const tags: FileTags[] = [ + { + path: 'source.ts', + tags: [ + { + kind: 'ref', + name: 'SharedName', + line: 1, + signature: 'const value = SharedName', + }, + ], + }, + ] + + for (let i = 0; i < 101; i++) { + tags.push({ + path: `defs/file${i}.ts`, + tags: [ + { + kind: 'def', + name: 'SharedName', + line: 1, + signature: `export const SharedName = ${i}`, + }, + ], + }) + } + + const graph = buildGraph(tags) + expect(graph.outDegree('source.ts')).toBe(0) + }) + + test('does not create zero-weight edges for symbols defined in every file', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-zero-weight-')) + try { + for (const file of ['a.ts', 'b.ts']) { + writeFileSync( + join(tempDir, file), + [ + 'export class Config {}', + 'export function useConfig(input: Config): Config { return input }', + '', + ].join('\n'), + ) + } + + const result = await buildRepoMap({ + root: tempDir, + files: ['a.ts', 'b.ts'], + maxTokens: 2048, + }) + + expect(result.fileCount).toBe(2) + expect(result.map).toContain('a.ts:') + expect(result.map).toContain('b.ts:') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) +}) + +describe('tag cache extraction', () => { + test('preserves input file order when mixing cached and uncached tags', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-cache-order-')) + try { + const files = ['b.ts', 'a.ts', 'c.ts'] + for (const file of files) { + const symbolName = file.replace('.ts', '').toUpperCase() + writeFileSync(join(tempDir, file), `export const ${symbolName} = 1\n`) + } + + const fileStats = new Map(files.map(file => [file, statFile(tempDir, file)])) + const cache = { version: 2, entries: {}, renderedEntries: {} } + const cachedTags = await extractTags('b.ts', tempDir) + expect(cachedTags).not.toBeNull() + setCachedTags( + cache, + cachedTags!.path, + tempDir, + cachedTags!.tags, + fileStats.get(cachedTags!.path) ?? undefined, + ) + + const result = await extractTagsWithCache({ + files, + root: tempDir, + cache, + fileStats, + }) + + expect(result.map(fileTags => fileTags.path)).toEqual(files) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) +}) + +describe('parser project root resolution', () => { + test('resolves source checkout roots for Windows paths', () => { + const filePath = 'C:\\repo\\src\\context\\repoMap\\parser.ts' + expect(resolveProjectRoot(filePath)).toBe('C:\\repo\\') + }) + + test('resolves bundled roots from dist paths', () => { + const filePath = win32.join('C:\\repo', 'dist', 'cli.mjs') + expect(resolveProjectRoot(filePath)).toBe('C:\\repo\\') + }) +}) + +describe('pagerank', () => { + test('ranks the most-imported file highest', async () => { + const result = await buildRepoMap({ + root: FIXTURE_ROOT, + maxTokens: 2048, + files: FIXTURE_FILES, + }) + + // The map starts with the highest-ranked file + const firstFile = result.map.split('\n')[0] + expect(firstFile).toBe('fileC.ts:') + + // fileE should be ranked lowest (or near last) + const lines = result.map.split('\n') + const filePositions = FIXTURE_FILES.map(f => { + const idx = lines.findIndex(l => l === `${f}:`) + return { file: f, position: idx } + }).filter(x => x.position >= 0) + .sort((a, b) => a.position - b.position) + + // fileC should be first + expect(filePositions[0]!.file).toBe('fileC.ts') + + // fileE should be last (or among the last) + const lastFile = filePositions[filePositions.length - 1]!.file + expect(['fileD.ts', 'fileE.ts']).toContain(lastFile) + }) + + test('directory focus boosts matching file nodes', () => { + const graph = buildGraph([ + { + path: 'src/tools/isolated.ts', + tags: [ + { + kind: 'def', + name: 'IsolatedTool', + line: 1, + signature: 'export class IsolatedTool {}', + }, + ], + }, + { + path: 'src/core.ts', + tags: [ + { + kind: 'def', + name: 'Core', + line: 1, + signature: 'export class Core {}', + }, + { + kind: 'ref', + name: 'Helper', + line: 2, + signature: 'new Helper()', + }, + ], + }, + { + path: 'src/helper.ts', + tags: [ + { + kind: 'def', + name: 'Helper', + line: 1, + signature: 'export class Helper {}', + }, + ], + }, + ]) + + const noFocus = rankFiles(graph).map(file => file.path) + const dirFocus = rankFiles(graph, ['src/tools/']).map(file => file.path) + + expect(noFocus[0]).not.toBe('src/tools/isolated.ts') + expect(dirFocus[0]).toBe('src/tools/isolated.ts') + }) + + test('symbol focus boosts matching files in the core build path', async () => { + const result = await buildRepoMap({ + root: FIXTURE_ROOT, + maxTokens: 2048, + files: FIXTURE_FILES, + focusSymbols: ['ConsoleLogger'], + }) + + const firstFile = result.map.split('\n')[0] + expect(firstFile).toBe('fileE.ts:') + expect(result.map).toContain('export class ConsoleLogger implements Logger') + }) +}) + +describe('renderer', () => { + test('respects the token budget within 5%', async () => { + const maxTokens = 500 + const result = await buildRepoMap({ + root: FIXTURE_ROOT, + maxTokens, + files: FIXTURE_FILES, + }) + + const actualTokens = countTokens(result.map) + expect(actualTokens).toBeLessThanOrEqual(maxTokens * 1.05) + expect(result.tokenCount).toBeLessThanOrEqual(maxTokens * 1.05) + }) + + test('drops files that don\'t fit rather than listing their names', async () => { + // Very tight budget — should only fit 1-2 files + const result = await buildRepoMap({ + root: FIXTURE_ROOT, + maxTokens: 100, + files: FIXTURE_FILES, + }) + + // Count how many files appear as headers in the output + const fileHeaders = result.map.split('\n').filter(l => l.endsWith(':') && !l.startsWith(' ')) + + // Every file header in the output should have its signatures listed + for (const header of fileHeaders) { + // The file must have at least one signature line after it + const headerIdx = result.map.indexOf(header) + const afterHeader = result.map.slice(headerIdx + header.length) + // Should have content (signatures), not just the filename + expect(afterHeader.trim().length).toBeGreaterThan(0) + } + + // Should have fewer files than total + expect(fileHeaders.length).toBeLessThan(FIXTURE_FILES.length) + }) + + test('skips an oversized ranked file and keeps later files that fit', () => { + const hugeSignature = `export function enormous(${Array.from( + { length: 120 }, + (_, i) => `arg${i}: string`, + ).join(', ')}): void {}` + const ranked = [ + { path: 'big.ts', score: 10 }, + { path: 'small.ts', score: 1 }, + ] + const fileTagsMap = new Map([ + [ + 'big.ts', + { + path: 'big.ts', + tags: [ + { + kind: 'def', + name: 'enormous', + line: 1, + signature: hugeSignature, + }, + ], + }, + ], + [ + 'small.ts', + { + path: 'small.ts', + tags: [ + { + kind: 'def', + name: 'small', + line: 1, + signature: 'export function small(): void {}', + }, + ], + }, + ], + ]) + + const result = renderMap(ranked, fileTagsMap, 80) + + expect(result.map).toContain('small.ts:') + expect(result.map).not.toContain('big.ts:') + expect(result.fileCount).toBe(1) + expect(result.tokenCount).toBeLessThanOrEqual(80) + }) +}) + +describe('cache', () => { + test('saveCache does not throw when persistence is unavailable', () => { + expect(() => { + saveCache('\0invalid-root', { + version: 2, + entries: { + bad: { + tags: [{ value: BigInt(1) }], + mtimeMs: 0, + size: 0, + }, + }, + renderedEntries: {}, + } as never) + }).not.toThrow() + }) + + test('stores repo map cache under OPENCLAUDE_CONFIG_DIR when configured', async () => { + const configDir = mkdtempSync(join(tmpdir(), 'repomap-config-home-')) + const expectedCacheDir = join(configDir, 'repomap-cache') + const { + CLAUDE_CONFIG_DIR: _legacyConfigDir, + OPENCLAUDE_CONFIG_DIR: _openClaudeConfigDir, + ...env + } = process.env + + try { + const result = spawnSync( + process.execPath, + [ + '--eval', + [ + "import { getCacheStats } from './src/context/repoMap/index.ts'", + 'const stats = getCacheStats(process.argv[1])', + 'if (stats.cacheDir !== process.argv[2]) {', + ' console.error(`Expected ${process.argv[2]}, got ${stats.cacheDir}`)', + ' process.exit(1)', + '}', + ].join('\n'), + FIXTURE_ROOT, + expectedCacheDir, + ], + { + cwd: REPO_ROOT, + encoding: 'utf8', + env: { + ...env, + FORCE_COLOR: '0', + OPENCLAUDE_CONFIG_DIR: configDir, + }, + }, + ) + + if (result.error) { + throw result.error + } + + expect(result.stderr).toBe('') + expect(result.status).toBe(0) + } finally { + rmSync(configDir, { recursive: true, force: true }) + } + }) + + test('second build of unchanged fixture uses the cache', async () => withWritableConfigHome(async () => { + // First build (cold) + const result1 = await buildRepoMap({ + root: FIXTURE_ROOT, + maxTokens: 2048, + files: FIXTURE_FILES, + }) + expect(result1.cacheHit).toBe(false) + + // Second build (warm) + const result2 = await buildRepoMap({ + root: FIXTURE_ROOT, + maxTokens: 2048, + files: FIXTURE_FILES, + }) + expect(result2.cacheHit).toBe(true) + + // Output should be identical + expect(result2.map).toBe(result1.map) + })) + + test('modifying a file invalidates the rendered cache without clearing cache data', async () => withWritableConfigHome(async () => { + // Create a temp copy of the fixture + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-test-')) + try { + for (const f of FIXTURE_FILES) { + cpSync(join(FIXTURE_ROOT, f), join(tempDir, f)) + } + + // First build + const result1 = await buildRepoMap({ + root: tempDir, + maxTokens: 2048, + files: FIXTURE_FILES, + }) + expect(result1.cacheHit).toBe(false) + + // Touch one file to change its mtime + const targetFile = join(tempDir, 'fileE.ts') + const now = new Date() + utimesSync(targetFile, now, now) + + const cacheBeforeSecondBuild = loadCache(tempDir) + expect(getCachedTags( + cacheBeforeSecondBuild, + 'fileA.ts', + tempDir, + statFile(tempDir, 'fileA.ts') ?? undefined, + )).not.toBeNull() + expect(getCachedTags( + cacheBeforeSecondBuild, + 'fileE.ts', + tempDir, + statFile(tempDir, 'fileE.ts') ?? undefined, + )).toBeNull() + + // Second build — rendered cache should be invalidated because the map + // hash includes file stats, but unchanged files can still reuse tag cache. + const result2 = await buildRepoMap({ + root: tempDir, + maxTokens: 2048, + files: FIXTURE_FILES, + }) + // The per-file cache for fileE should miss (mtime changed), + // but other files should still hit the per-file cache + expect(result2.cacheHit).toBe(false) + + // Output should still be valid + expect(result2.map.length).toBeGreaterThan(0) + expect(result2.fileCount).toBe(result1.fileCount) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + })) +}) + +describe('gitFiles', () => { + test('ignores inherited git environment overrides', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-git-env-')) + const previousGitDir = process.env.GIT_DIR + const previousGitWorkTree = process.env.GIT_WORK_TREE + try { + writeFileSync( + join(tempDir, 'hello.ts'), + 'export function hello(): string { return "world" }\n', + ) + + process.env.GIT_DIR = join(process.cwd(), '.git') + process.env.GIT_WORK_TREE = process.cwd() + + expect(await getRepoFiles(tempDir)).toEqual(['hello.ts']) + } finally { + if (previousGitDir === undefined) delete process.env.GIT_DIR + else process.env.GIT_DIR = previousGitDir + if (previousGitWorkTree === undefined) delete process.env.GIT_WORK_TREE + else process.env.GIT_WORK_TREE = previousGitWorkTree + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('preserves leading whitespace in git-tracked file paths', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-git-paths-')) + const env = { ...process.env } + delete env.GIT_DIR + delete env.GIT_WORK_TREE + delete env.GIT_INDEX_FILE + + try { + writeFileSync( + join(tempDir, ' leading.ts'), + 'export function LeadingSpaceFile(): string { return "ok" }\n', + ) + writeFileSync( + join(tempDir, 'normal.ts'), + 'export function NormalFile(): string { return "ok" }\n', + ) + + execFileSync('git', ['init'], { cwd: tempDir, env, stdio: 'ignore' }) + execFileSync('git', ['add', '.'], { cwd: tempDir, env, stdio: 'ignore' }) + + const files = await getRepoFiles(tempDir) + expect(files).toContain(' leading.ts') + + const result = await buildRepoMap({ + root: tempDir, + maxTokens: 1024, + }) + + expect(result.map).toContain(' leading.ts:') + expect(result.map).toContain('LeadingSpaceFile') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('falls back gracefully when not in a git repo', async () => { + // Create a temp directory with source files but NO .git + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-nogit-')) + try { + writeFileSync( + join(tempDir, 'hello.ts'), + 'export function hello(): string { return "world" }\n', + ) + writeFileSync( + join(tempDir, 'utils.ts'), + 'export function add(a: number, b: number): number { return a + b }\n', + ) + + const result = await buildRepoMap({ + root: tempDir, + maxTokens: 1024, + }) + + // Should succeed without throwing + expect(result.map.length).toBeGreaterThan(0) + expect(result.totalFileCount).toBeGreaterThan(0) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) +}) + +describe('error handling', () => { + test('no crash on malformed source file', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-malformed-')) + try { + // Valid file + writeFileSync( + join(tempDir, 'good.ts'), + 'export function good(): number { return 1 }\n', + ) + // Malformed file — severe syntax errors + writeFileSync( + join(tempDir, 'bad.ts'), + '}{}{}{export classclass [[[ function ,,, @@@ ###\n', + ) + + const result = await buildRepoMap({ + root: tempDir, + maxTokens: 1024, + files: ['good.ts', 'bad.ts'], + }) + + // Should complete successfully + expect(result.map.length).toBeGreaterThan(0) + // The good file should be in the output + expect(result.map).toContain('good.ts') + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) +}) + +describe('rendered cache invalidation', () => { + test('reflects file edits without manual invalidation', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-stale-')) + try { + writeFileSync( + join(tempDir, 'main.ts'), + 'export function oldName(): void {}\n', + ) + + const first = await buildRepoMap({ root: tempDir, maxTokens: 1024 }) + expect(first.cacheHit).toBe(false) + expect(first.map).toContain('oldName') + + // Bump mtime forward so the change is visible on filesystems with + // coarse timestamp resolution. + const future = new Date(Date.now() + 2000) + writeFileSync( + join(tempDir, 'main.ts'), + 'export function newName(): void {}\n', + ) + utimesSync(join(tempDir, 'main.ts'), future, future) + + const second = await buildRepoMap({ root: tempDir, maxTokens: 1024 }) + expect(second.map).toContain('newName') + expect(second.map).not.toContain('oldName') + } finally { + invalidateCache(tempDir) + rmSync(tempDir, { recursive: true, force: true }) + } + }) +}) diff --git a/src/context/repoMap/symbolExtractor.ts b/src/context/repoMap/symbolExtractor.ts new file mode 100644 index 000000000..1435d3572 --- /dev/null +++ b/src/context/repoMap/symbolExtractor.ts @@ -0,0 +1,145 @@ +import { readFileSync } from 'fs' +import { join } from 'path' +import { getLanguageForFile } from './gitFiles.js' +import { createParser, loadLanguage, loadQuery } from './parser.js' +import type { FileTags, Tag } from './types.js' + +type TreeSitterQuery = { + matches(rootNode: unknown): Array<{ + pattern: number + captures: Array<{ + name: string + node: { + text: string + startPosition: { row: number; column: number } + endPosition: { row: number; column: number } + } + }> + }> + delete?: () => void +} + +type TreeSitterTree = { + rootNode: unknown + delete?: () => void +} + +const queryCache = new Map() +const queryLoadPromises = new Map>() + +async function getQuery(language: NonNullable>) { + const cached = queryCache.get(language) + if (cached) return cached + + const pending = queryLoadPromises.get(language) + if (pending) return await pending + + const promise = (async () => { + const querySource = loadQuery(language) + if (!querySource) return null + + const lang = await loadLanguage(language) + if (!lang) return null + + const { Query } = await import('web-tree-sitter') + const query = new Query(lang, querySource) as TreeSitterQuery + queryCache.set(language, query) + return query + })() + + queryLoadPromises.set(language, promise) + + try { + return await promise + } finally { + if (queryLoadPromises.get(language) === promise) { + queryLoadPromises.delete(language) + } + } +} + +export function clearSymbolExtractorCaches(): void { + for (const query of queryCache.values()) { + query.delete?.() + } + queryCache.clear() + queryLoadPromises.clear() +} + +/** + * Extract definition and reference tags from a single source file. + * Returns null if the file can't be parsed (unsupported language, parse error, etc). + */ +export async function extractTags( + filePath: string, + root: string, +): Promise { + const language = getLanguageForFile(filePath) + if (!language) return null + + const absolutePath = join(root, filePath) + let source: string + try { + source = readFileSync(absolutePath, 'utf-8') + } catch { + return null + } + + const lines = source.split('\n') + + const parser = await createParser(language) + if (!parser) return null + + let tree: TreeSitterTree | null = null + + try { + tree = parser.parse(source) as TreeSitterTree + const query = await getQuery(language) + if (!query) return null + + const matches = query.matches(tree.rootNode) + const tags: Tag[] = [] + const seen = new Set() // dedup by kind+subkind+name+line + + for (const match of matches) { + let name: string | null = null + let kind: 'def' | 'ref' | null = null + let subKind: string | undefined + let lineRow = 0 + + for (const capture of match.captures) { + const captureName = capture.name + + // Name captures: name.definition.X or name.reference.X + if (captureName.startsWith('name.definition.')) { + name = capture.node.text + kind = 'def' + subKind = captureName.slice('name.definition.'.length) + lineRow = capture.node.startPosition.row + } else if (captureName.startsWith('name.reference.')) { + name = capture.node.text + kind = 'ref' + subKind = captureName.slice('name.reference.'.length) + lineRow = capture.node.startPosition.row + } + } + + if (name && kind) { + const key = `${kind}:${subKind ?? ''}:${name}:${lineRow}` + if (!seen.has(key)) { + seen.add(key) + const line = lineRow + 1 // convert 0-based to 1-based + const signature = lines[lineRow]?.trimEnd() ?? '' + tags.push({ kind, name, line, signature, subKind }) + } + } + } + + return { path: filePath, tags } + } catch { + return null + } finally { + tree?.delete?.() + parser.delete() + } +} diff --git a/src/context/repoMap/tokenize.ts b/src/context/repoMap/tokenize.ts new file mode 100644 index 000000000..274a6efb6 --- /dev/null +++ b/src/context/repoMap/tokenize.ts @@ -0,0 +1,15 @@ +import { getEncoding, type Tiktoken } from 'js-tiktoken' + +let encoder: Tiktoken | null = null + +function getEncoder() { + if (!encoder) { + encoder = getEncoding('cl100k_base') + } + return encoder +} + +/** Count the number of tokens in a string using cl100k_base encoding. */ +export function countTokens(text: string): number { + return getEncoder().encode(text).length +} diff --git a/src/context/repoMap/types.ts b/src/context/repoMap/types.ts new file mode 100644 index 000000000..dd776fbc3 --- /dev/null +++ b/src/context/repoMap/types.ts @@ -0,0 +1,82 @@ +export interface Tag { + /** 'def' for definitions, 'ref' for references */ + kind: 'def' | 'ref' + /** Symbol name (e.g. function name, class name) */ + name: string + /** 1-based line number in the source file */ + line: number + /** The full line of source code at this position (used as signature for defs) */ + signature: string + /** Sub-kind from the query (e.g. 'function', 'class', 'method', 'type') */ + subKind?: string +} + +export interface FileTags { + /** Relative path from the repo root */ + path: string + /** All tags extracted from this file */ + tags: Tag[] +} + +export interface RepoMapOptions { + /** Root directory of the repo (defaults to cwd) */ + root?: string + /** Maximum token budget for the rendered map */ + maxTokens?: number + /** Files to boost in PageRank (relative paths) */ + focusFiles?: string[] + /** Symbol names whose defining files should be boosted in PageRank */ + focusSymbols?: string[] + /** Override the list of files to process (relative paths) */ + files?: string[] + /** Optional cancellation hook checked between expensive build steps */ + shouldContinue?: () => void +} + +export interface RepoMapResult { + /** The rendered repo map string */ + map: string + /** Whether the result came from cache */ + cacheHit: boolean + /** Time in milliseconds to build the map */ + buildTimeMs: number + /** Number of files included in the rendered map */ + fileCount: number + /** Total number of files processed */ + totalFileCount: number + /** Actual token count of the rendered map */ + tokenCount: number +} + +export interface CacheEntry { + tags: Tag[] + mtimeMs: number + size: number +} + +export interface FileStatFingerprint { + mtimeMs: number + size: number +} + +export interface RenderedCacheEntry { + map: string + fileCount: number + tokenCount: number + createdAt: number +} + +export interface CacheData { + version: number + entries: Record + renderedEntries: Record +} + +export interface CacheStats { + cacheDir: string + cacheFile: string | null + entryCount: number + exists: boolean +} + +export type SupportedLanguage = 'typescript' | 'tsx' | 'javascript' | 'python' diff --git a/src/tools.ts b/src/tools.ts index 0d93e368c..bdbec480b 100644 --- a/src/tools.ts +++ b/src/tools.ts @@ -48,6 +48,7 @@ import { TodoWriteTool } from './tools/TodoWriteTool/TodoWriteTool.js' import { ExitPlanModeV2Tool } from './tools/ExitPlanModeTool/ExitPlanModeV2Tool.js' import { TestingPermissionTool } from './tools/testing/TestingPermissionTool.js' import { GrepTool } from './tools/GrepTool/GrepTool.js' +import { RepoMapTool } from './tools/RepoMapTool/RepoMapTool.js' // Lazy require to break circular dependency: tools.ts -> TeamCreateTool/TeamDeleteTool -> ... -> tools.ts /* eslint-disable @typescript-eslint/no-require-imports */ const getTeamCreateTool = () => @@ -188,6 +189,7 @@ export function getAllBaseTools(): Tools { // trick as ripgrep). When available, find/grep in Claude's shell are aliased // to these fast tools, so the dedicated Glob/Grep tools are unnecessary. ...(hasEmbeddedSearchTools() ? [] : [GlobTool, GrepTool]), + RepoMapTool, ExitPlanModeV2Tool, FileReadTool, FileEditTool, diff --git a/src/tools/RepoMapTool/RepoMapTool.test.ts b/src/tools/RepoMapTool/RepoMapTool.test.ts new file mode 100644 index 000000000..7aa1781d0 --- /dev/null +++ b/src/tools/RepoMapTool/RepoMapTool.test.ts @@ -0,0 +1,299 @@ +import { beforeAll, describe, expect, test } from 'bun:test' +import { execFileSync } from 'child_process' +import { cpSync, mkdtempSync, rmSync } from 'fs' +import { tmpdir } from 'os' +import { join } from 'path' +import { initParser } from '../../context/repoMap/parser.js' +import { invalidateCache } from '../../context/repoMap/index.js' +import { getRepoFiles } from '../../context/repoMap/gitFiles.js' +import { getCwd, pwd, runWithCwdOverride } from '../../utils/cwd.js' +import { isAbortError } from '../../utils/errors.js' +import { runWithSdkContext, setCwdState } from '../../bootstrap/state.js' +import type { SessionId } from '../../types/ids.js' +import { RepoMapTool } from './RepoMapTool.js' +import { getToolUseSummary } from './UI.js' + +const FIXTURE_ROOT = join( + import.meta.dir, + '..', + '..', + 'context', + 'repoMap', + '__fixtures__', + 'mini-repo', +) +const FIXTURE_FILES = [ + 'fileA.ts', + 'fileB.ts', + 'fileC.ts', + 'fileD.ts', + 'fileE.ts', +] + +beforeAll(async () => { + await initParser() +}) + +function populateFixtureRepo(tempDir: string): void { + for (const f of FIXTURE_FILES) { + cpSync(join(FIXTURE_ROOT, f), join(tempDir, f)) + } + const env = { ...process.env } + delete env.GIT_DIR + delete env.GIT_WORK_TREE + delete env.GIT_INDEX_FILE + execFileSync('git', ['init'], { cwd: tempDir, env, stdio: 'ignore' }) + execFileSync('git', ['add', '.'], { cwd: tempDir, env, stdio: 'ignore' }) +} + +async function callRepoMapToolInCwd( + cwd: string, + input: Parameters[0], +) { + const previousCwd = getCwd() + return await runWithSdkContext({ + sessionId: 'repomap-tool-test' as SessionId, + sessionProjectDir: null, + cwd, + originalCwd: cwd, + }, () => runWithCwdOverride(cwd, async () => { + setCwdState(cwd) + expect(pwd()).toBe(cwd) + try { + return await RepoMapTool.call( + input, + { abortController: new AbortController() } as Parameters[1], + ) + } finally { + setCwdState(previousCwd) + } + })) +} + +describe('RepoMapTool schema', () => { + test('validates a minimal input {}', () => { + const schema = RepoMapTool.inputSchema + const result = schema.safeParse({}) + expect(result.success).toBe(true) + }) + + test('rejects max_tokens below 256', () => { + const schema = RepoMapTool.inputSchema + const result = schema.safeParse({ max_tokens: 100 }) + expect(result.success).toBe(false) + }) + + test('rejects max_tokens above 16384', () => { + const schema = RepoMapTool.inputSchema + const result = schema.safeParse({ max_tokens: 20000 }) + expect(result.success).toBe(false) + }) + + test('accepts focus_files as string[]', () => { + const schema = RepoMapTool.inputSchema + const result = schema.safeParse({ + focus_files: ['src/tools/', 'src/context.ts'], + }) + expect(result.success).toBe(true) + }) +}) + +describe('RepoMapTool call', () => { + test('call returns the declared output shape', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-tool-')) + try { + populateFixtureRepo(tempDir) + expect(await getRepoFiles(tempDir)).toEqual(FIXTURE_FILES) + + const result = await callRepoMapToolInCwd(tempDir, { max_tokens: 256 }) + + expect(typeof result.data.rendered).toBe('string') + expect(typeof result.data.token_count).toBe('number') + expect(typeof result.data.file_count).toBe('number') + expect(typeof result.data.total_file_count).toBe('number') + expect(typeof result.data.cache_hit).toBe('boolean') + expect(typeof result.data.build_time_ms).toBe('number') + expect(result.data.total_file_count).toBe(FIXTURE_FILES.length) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('returns a rendered map for a directory', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-tool-')) + try { + populateFixtureRepo(tempDir) + + const { buildRepoMap } = await import( + '../../context/repoMap/index.js' + ) + const result = await buildRepoMap({ + root: tempDir, + maxTokens: 1024, + }) + + expect(result.map.length).toBeGreaterThan(0) + expect(result.fileCount).toBeGreaterThan(0) + expect(result.totalFileCount).toBe(5) + expect(result.tokenCount).toBeGreaterThan(0) + expect(result.tokenCount).toBeLessThanOrEqual(1024) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('respects max_tokens parameter', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-tool-')) + try { + populateFixtureRepo(tempDir) + + const { buildRepoMap } = await import( + '../../context/repoMap/index.js' + ) + + const small = await buildRepoMap({ root: tempDir, maxTokens: 256 }) + const large = await buildRepoMap({ root: tempDir, maxTokens: 4096 }) + + expect(small.tokenCount).toBeLessThanOrEqual(256) + // Large budget should include more or equal content + expect(large.map.length).toBeGreaterThanOrEqual(small.map.length) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('focus_files boosts specified files in the ranking', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-tool-')) + try { + populateFixtureRepo(tempDir) + + const { buildRepoMap } = await import( + '../../context/repoMap/index.js' + ) + + // Without focus, fileE is ranked last (isolated) + const noFocus = await buildRepoMap({ root: tempDir, maxTokens: 2048 }) + const lines = noFocus.map.split('\n') + const fileEPos = lines.findIndex(l => l === 'fileE.ts:') + + // With focus on fileE + invalidateCache(tempDir) + const withFocus = await buildRepoMap({ + root: tempDir, + maxTokens: 2048, + focusFiles: ['fileE.ts'], + }) + const focusLines = withFocus.map.split('\n') + const fileEFocusPos = focusLines.findIndex(l => l === 'fileE.ts:') + + // fileE should rank higher (earlier position) with focus + expect(fileEFocusPos).toBeLessThan(fileEPos) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('focus_symbols resolves matching files through the public tool call', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-tool-')) + try { + populateFixtureRepo(tempDir) + + const result = await callRepoMapToolInCwd(tempDir, { + max_tokens: 2048, + focus_symbols: ['ConsoleLogger'], + }) + + expect(result.data.rendered).toContain('fileE.ts:') + expect(result.data.rendered).toContain( + 'export class ConsoleLogger implements Logger', + ) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) + + test('does not build a repository map when already aborted', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'repomap-tool-abort-')) + try { + populateFixtureRepo(tempDir) + const abortController = new AbortController() + abortController.abort('test-abort') + + let error: unknown + try { + await runWithSdkContext({ + sessionId: 'repomap-tool-abort-test' as SessionId, + sessionProjectDir: null, + cwd: tempDir, + originalCwd: tempDir, + }, () => runWithCwdOverride(tempDir, () => + RepoMapTool.call( + { max_tokens: 256 }, + { abortController } as Parameters[1], + ), + )) + } catch (err) { + error = err + } + + expect(isAbortError(error)).toBe(true) + } finally { + rmSync(tempDir, { recursive: true, force: true }) + invalidateCache(tempDir) + } + }) +}) + +describe('RepoMapTool properties', () => { + test('is marked read-only and concurrency-safe', () => { + expect(RepoMapTool.isReadOnly()).toBe(true) + expect(RepoMapTool.isConcurrencySafe()).toBe(true) + }) + + test('exposes a path hook for read permission grants', () => { + expect(typeof RepoMapTool.getPath).toBe('function') + expect(RepoMapTool.getPath?.()).toBe(pwd()) + }) +}) + +describe('RepoMapTool result mapping', () => { + test('maps output to a tool result block without a duplicate output type', () => { + const block = RepoMapTool.mapToolResultToToolResultBlockParam( + { + rendered: 'file.ts:\n export const value = 1', + token_count: 10, + file_count: 1, + total_file_count: 1, + cache_hit: false, + build_time_ms: 12, + }, + 'toolu_123', + ) + + expect(block.tool_use_id).toBe('toolu_123') + expect(block.type).toBe('tool_result') + expect(block.content).toContain('Repository map: 1 files ranked') + expect(block.content).toContain('file.ts:') + }) +}) + +describe('RepoMapTool UI', () => { + test('getToolUseSummary returns descriptive string including focus', () => { + expect(getToolUseSummary(undefined)).toBe('Repository map') + expect(getToolUseSummary({})).toBe('Repository map') + expect(getToolUseSummary({ focus_files: ['src/tools/'] })).toContain( + 'focus:', + ) + expect(getToolUseSummary({ focus_files: ['src/tools/'] })).toContain( + 'src/tools/', + ) + expect( + getToolUseSummary({ focus_symbols: ['buildTool'] }), + ).toContain('buildTool') + }) +}) diff --git a/src/tools/RepoMapTool/RepoMapTool.ts b/src/tools/RepoMapTool/RepoMapTool.ts new file mode 100644 index 000000000..b421b26b5 --- /dev/null +++ b/src/tools/RepoMapTool/RepoMapTool.ts @@ -0,0 +1,160 @@ +import { z } from 'zod/v4' +import { buildTool, type ToolDef } from '../../Tool.js' +import { pwd } from '../../utils/cwd.js' +import { lazySchema } from '../../utils/lazySchema.js' +import { checkReadPermissionForTool } from '../../utils/permissions/filesystem.js' +import type { PermissionDecision } from '../../utils/permissions/PermissionResult.js' +import { AbortError } from '../../utils/errors.js' +import { buildRepoMap } from '../../context/repoMap/index.js' +import { REPO_MAP_TOOL_NAME, getDescription } from './prompt.js' +import { + getToolUseSummary, + renderToolResultMessage, + renderToolUseErrorMessage, + renderToolUseMessage, +} from './UI.js' + +const inputSchema = lazySchema(() => + z.strictObject({ + max_tokens: z + .number() + .int() + .min(256) + .max(16384) + .optional() + .describe( + 'Maximum token budget for the rendered map. Higher values include more files. Default: 1024.', + ), + focus_files: z + .array(z.string()) + .optional() + .describe( + 'Relative file or directory paths to boost in the ranking (e.g. ["src/tools/", "src/context.ts"]).', + ), + focus_symbols: z + .array(z.string()) + .optional() + .describe( + 'Symbol names to boost — files defining these symbols rank higher (e.g. ["buildTool", "ToolUseContext"]).', + ), + }), +) +type InputSchema = ReturnType + +const outputSchema = lazySchema(() => + z.object({ + rendered: z.string(), + token_count: z.number(), + file_count: z.number(), + total_file_count: z.number(), + cache_hit: z.boolean(), + build_time_ms: z.number(), + }), +) +type OutputSchema = ReturnType + +export type Output = z.infer + +export const RepoMapTool = buildTool({ + name: REPO_MAP_TOOL_NAME, + searchHint: 'structural map of repository files and symbols', + maxResultSizeChars: 50_000, + async description() { + return getDescription() + }, + userFacingName() { + return 'Repository map' + }, + getToolUseSummary, + getActivityDescription(input) { + if (input?.focus_files?.length) { + return `Building repository map (focus: ${input.focus_files.join(', ')})` + } + return 'Building repository map' + }, + get inputSchema(): InputSchema { + return inputSchema() + }, + get outputSchema(): OutputSchema { + return outputSchema() + }, + isConcurrencySafe() { + return true + }, + isReadOnly() { + return true + }, + isSearchOrReadCommand() { + return { isSearch: false, isRead: true } + }, + getPath() { + return pwd() + }, + toAutoClassifierInput(input) { + const parts: string[] = ['repomap'] + if (input.focus_files?.length) parts.push(`focus: ${input.focus_files.join(',')}`) + return parts.join(' ') + }, + async checkPermissions(input, context): Promise { + const appState = context.getAppState() + return checkReadPermissionForTool( + RepoMapTool, + input, + appState.toolPermissionContext, + ) + }, + async prompt() { + return getDescription() + }, + renderToolUseMessage, + renderToolUseErrorMessage, + renderToolResultMessage, + extractSearchText({ rendered }) { + return rendered + }, + mapToolResultToToolResultBlockParam(output, toolUseID) { + const summary = [ + `Repository map: ${output.file_count} files ranked (${output.total_file_count} total), ${output.token_count} tokens`, + output.cache_hit ? '(cached)' : `(built in ${output.build_time_ms}ms)`, + ].join(' ') + + return { + tool_use_id: toolUseID, + type: 'tool_result', + content: `${summary}\n\n${output.rendered}`, + } + }, + async call( + { max_tokens = 1024, focus_files, focus_symbols }, + { abortController }, + ) { + throwIfAborted(abortController.signal) + const root = pwd() + + throwIfAborted(abortController.signal) + const result = await buildRepoMap({ + root, + maxTokens: max_tokens, + focusFiles: focus_files, + focusSymbols: focus_symbols, + shouldContinue: () => throwIfAborted(abortController.signal), + }) + + const output: Output = { + rendered: result.map, + token_count: result.tokenCount, + file_count: result.fileCount, + total_file_count: result.totalFileCount, + cache_hit: result.cacheHit, + build_time_ms: result.buildTimeMs, + } + + return { data: output } + }, +} satisfies ToolDef) + +function throwIfAborted(signal: AbortSignal): void { + if (signal.aborted) { + throw new AbortError() + } +} diff --git a/src/tools/RepoMapTool/UI.tsx b/src/tools/RepoMapTool/UI.tsx new file mode 100644 index 000000000..6fd949fc8 --- /dev/null +++ b/src/tools/RepoMapTool/UI.tsx @@ -0,0 +1,88 @@ +import type { ToolResultBlockParam } from '@anthropic-ai/sdk/resources/index.mjs' +import React from 'react' +import { FallbackToolUseErrorMessage } from '../../components/FallbackToolUseErrorMessage.js' +import { MessageResponse } from '../../components/MessageResponse.js' +import { TOOL_SUMMARY_MAX_LENGTH } from '../../constants/toolLimits.js' +import { Text } from '../../ink.js' +import type { ToolProgressData } from '../../Tool.js' +import type { ProgressMessage } from '../../types/message.js' +import { truncate } from '../../utils/format.js' +import type { Output } from './RepoMapTool.js' + +export function getToolUseSummary( + input: + | Partial<{ + max_tokens?: number + focus_files?: string[] + focus_symbols?: string[] + }> + | undefined, +): string | null { + if (!input) return 'Repository map' + const parts: string[] = [] + if (input.focus_files?.length) { + parts.push(input.focus_files.join(', ')) + } + if (input.focus_symbols?.length) { + parts.push(input.focus_symbols.join(', ')) + } + if (parts.length > 0) { + return truncate(`Repository map (focus: ${parts.join('; ')})`, TOOL_SUMMARY_MAX_LENGTH) + } + return 'Repository map' +} + +export function renderToolUseMessage( + input: Partial<{ + max_tokens?: number + focus_files?: string[] + focus_symbols?: string[] + }>, +): React.ReactNode { + const parts: string[] = [] + if (input.max_tokens) { + parts.push(`max_tokens: ${input.max_tokens}`) + } + if (input.focus_files?.length) { + parts.push(`focus: ${input.focus_files.join(', ')}`) + } + if (input.focus_symbols?.length) { + parts.push(`symbols: ${input.focus_symbols.join(', ')}`) + } + return parts.length > 0 ? parts.join(', ') : null +} + +export function renderToolResultMessage( + output: Output, + _progressMessages: ProgressMessage[], + { verbose }: { verbose: boolean }, +): React.ReactNode { + const summary = `${output.file_count} files ranked, ${output.token_count} tokens${output.cache_hit ? ' (cached)' : `, ${output.build_time_ms}ms`}` + + if (verbose) { + return ( + + + Built repository map: {summary} + {'\n'} + ({output.total_file_count} total files considered) + + + ) + } + + return ( + + + Built repository map: {summary} + + + ) +} + +export function renderToolUseErrorMessage( + result: ToolResultBlockParam['content'], + { verbose }: { verbose: boolean }, +): React.ReactNode { + return +} diff --git a/src/tools/RepoMapTool/prompt.ts b/src/tools/RepoMapTool/prompt.ts new file mode 100644 index 000000000..b7be48e86 --- /dev/null +++ b/src/tools/RepoMapTool/prompt.ts @@ -0,0 +1,31 @@ +export const REPO_MAP_TOOL_NAME = 'RepoMap' + +export function getDescription(): string { + return `Build a structural map of the repository showing ranked files and their key signatures (functions, classes, types, interfaces). + +## When to use +- At the start of a session on an unfamiliar repository to understand the codebase architecture +- Before cross-file refactors to identify which files are structurally connected +- When searching for where a concept or feature lives across the codebase +- When the user asks "how is this repo organized" or "what are the important files" + +## When NOT to use +- To read the contents of a specific file — use Read instead +- To search for exact text or patterns — use Grep instead +- To find files by name or glob pattern — use Glob instead +- When you already know which files to examine + +## How it works +The tool parses every supported source file (TypeScript, JavaScript, Python) using tree-sitter, extracts symbol definitions and references, builds a cross-file reference graph weighted by symbol importance (IDF), and ranks files using PageRank. The output is a token-budgeted summary showing the highest-ranked files with their key signatures (function/class/type declarations). + +## Parameters +- **max_tokens**: Controls how many files fit in the output. Use 1024 for a quick overview, 4096+ for comprehensive maps. Default: 1024. +- **focus_files**: Pass relative paths (e.g. \`["src/tools/"]\`) to boost specific files and their neighbors in the ranking. Use when the user mentions specific directories or files. +- **focus_symbols**: Pass symbol names (e.g. \`["buildTool", "ToolUseContext"]\`) to boost files that define those symbols. Use when the user asks about specific functions or types. + +## Important notes +- The map shows **signatures only**, not function bodies. Use Read to see implementations. +- Results are **auto-cached** on disk — repeat calls with the same parameters return instantly. +- Files are ranked by structural importance: files imported by many others rank highest. +` +} diff --git a/src/utils/cwd.test.ts b/src/utils/cwd.test.ts new file mode 100644 index 000000000..67b4b0c68 --- /dev/null +++ b/src/utils/cwd.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, test } from 'bun:test' +import { getCwd, runWithCwdOverride } from './cwd.js' + +function delay(ms: number): Promise { + return new Promise(resolve => setTimeout(resolve, ms)) +} + +describe('runWithCwdOverride', () => { + test('restores the outer cwd override after a nested synchronous override', () => { + runWithCwdOverride('/tmp/outer-cwd', () => { + expect(getCwd()).toBe('/tmp/outer-cwd') + + const inner = runWithCwdOverride('/tmp/inner-cwd', () => getCwd()) + expect(inner).toBe('/tmp/inner-cwd') + expect(getCwd()).toBe('/tmp/outer-cwd') + }) + }) + + test('keeps overlapping async cwd override windows isolated', async () => { + const first = runWithCwdOverride('/tmp/first-cwd', async () => { + expect(getCwd()).toBe('/tmp/first-cwd') + await delay(20) + expect(getCwd()).toBe('/tmp/first-cwd') + return getCwd() + }) + + const second = runWithCwdOverride('/tmp/second-cwd', async () => { + expect(getCwd()).toBe('/tmp/second-cwd') + await delay(5) + expect(getCwd()).toBe('/tmp/second-cwd') + return getCwd() + }) + + await expect(Promise.all([first, second])).resolves.toEqual([ + '/tmp/first-cwd', + '/tmp/second-cwd', + ]) + }) +}) diff --git a/src/utils/cwd.ts b/src/utils/cwd.ts index c4d1600aa..19ee95687 100644 --- a/src/utils/cwd.ts +++ b/src/utils/cwd.ts @@ -1,7 +1,8 @@ -import { AsyncLocalStorage } from 'async_hooks' +import { AsyncLocalStorage, createHook, executionAsyncId } from 'async_hooks' import { getCwdState, getOriginalCwd } from '../bootstrap/state.js' const cwdOverrideStorage = new AsyncLocalStorage() +const activeScopedOverrides = new Set>() /** * Run a function with an overridden working directory for the current async context. @@ -10,22 +11,83 @@ const cwdOverrideStorage = new AsyncLocalStorage() * agents to each see their own working directory without affecting each other. */ export function runWithCwdOverride(cwd: string, fn: () => T): T { - return cwdOverrideStorage.run(cwd, fn) + const scopedOverrides = new Map() + const initialAsyncId = executionAsyncId() + scopedOverrides.set(initialAsyncId, cwd) + activeScopedOverrides.add(scopedOverrides) + + const hook = createHook({ + init(asyncId, _type, triggerAsyncId) { + const inherited = scopedOverrides.get(triggerAsyncId) + if (inherited !== undefined) { + scopedOverrides.set(asyncId, inherited) + } + }, + destroy(asyncId) { + scopedOverrides.delete(asyncId) + }, + }) + + const cleanup = () => { + hook.disable() + activeScopedOverrides.delete(scopedOverrides) + scopedOverrides.clear() + } + + hook.enable() + + try { + const result = cwdOverrideStorage.run(cwd, fn) + if (isPromiseLike(result)) { + scopedOverrides.delete(initialAsyncId) + return result.finally(cleanup) as T + } + cleanup() + return result + } catch (error) { + cleanup() + throw error + } } /** * Get the current working directory */ export function pwd(): string { - return cwdOverrideStorage.getStore() ?? getCwdState() + return getCwdOverride() ?? getCwdState() +} + +function getCwdOverride(): string | undefined { + const asyncId = executionAsyncId() + return cwdOverrideStorage.getStore() ?? getScopedCwdOverride(asyncId) +} + +function getScopedCwdOverride(asyncId: number): string | undefined { + const scopedOverrideStack = [...activeScopedOverrides].reverse() + for (const scopedOverrides of scopedOverrideStack) { + const cwd = scopedOverrides.get(asyncId) + if (cwd !== undefined) return cwd + } + return undefined +} + +function isPromiseLike(value: T): value is T & { finally(onfinally: () => void): unknown } { + return ( + typeof value === 'object' && + value !== null && + 'finally' in value && + typeof value.finally === 'function' + ) } /** * Get the current working directory or the original working directory if the current one is not available */ export function getCwd(): string { + const override = getCwdOverride() + if (override !== undefined) return override try { - return pwd() + return getCwdState() } catch { return getOriginalCwd() } diff --git a/src/utils/user.test.ts b/src/utils/user.test.ts index f012f13c4..709f53cb3 100644 --- a/src/utils/user.test.ts +++ b/src/utils/user.test.ts @@ -6,7 +6,6 @@ import { } from '../test/sharedMutationLock.js' import * as realAuth from './auth.js' import * as realConfig from './config.js' -import * as realCwd from './cwd.js' import * as realEnv from './env.js' import * as realEnvUtils from './envUtils.js' @@ -19,17 +18,15 @@ async function importFreshUserModule() { async function importActualUserTestDeps() { const nonce = `${Date.now()}-${Math.random()}` - const [authModule, configModule, cwdModule, execaModule] = await Promise.all([ + const [authModule, configModule, execaModule] = await Promise.all([ import(`./auth.js?ts=${nonce}`), import(`./config.js?ts=${nonce}`), - import(`./cwd.js?ts=${nonce}`), import('execa'), ]) return { authModule, configModule, - cwdModule, execaModule, } } @@ -44,8 +41,7 @@ async function installCommonMocks(options?: { // every other test file that imports state.js (e.g. SDK CON-1 tests). // The dynamic import (importFreshUserModule) will use the real state.js, // which is fine — these tests only assert email, not sessionId. - const { authModule, configModule, cwdModule, execaModule } = - await importActualUserTestDeps() + const { authModule, configModule, execaModule } = await importActualUserTestDeps() mock.module('./auth.js', () => ({ ...authModule, @@ -67,11 +63,6 @@ async function installCommonMocks(options?: { getOrCreateUserID: () => 'device-test', })) - mock.module('./cwd.js', () => ({ - ...cwdModule, - getCwd: () => 'C:\\repo', - })) - mock.module('./env.js', () => ({ ...realEnv, env: { platform: 'windows' }, @@ -108,7 +99,6 @@ afterEach(() => { mock.restore() mock.module('./auth.js', () => realAuth) mock.module('./config.js', () => realConfig) - mock.module('./cwd.js', () => realCwd) mock.module('./env.js', () => realEnv) mock.module('./envUtils.js', () => realEnvUtils) mock.module('execa', () => realExeca)