diff --git a/README.md b/README.md index 9c72a80a..200f72a7 100644 --- a/README.md +++ b/README.md @@ -215,6 +215,13 @@ teamai codebase --lint # health check The graph stores components, interfaces, configs, and cross-repo import edges. `teamai recall` uses it for graph-boosted re-ranking. When a recall hit comes from a codebase page, the result includes a `Sources:` line listing the relevant source file paths — giving agents a direct starting point for code changes instead of re-exploring the repo. +Edges come from two tracks that run together, with AST results taking precedence on overlap: + +- **AST track** (TypeScript/JavaScript, Python, Go): a WASM [tree-sitter](https://tree-sitter.github.io/) parser resolves `import`/`require`, call sites, and TS `implements` clauses to precise file-to-file `DEPENDS_ON` / `REFERENCES` / `IMPLEMENTS` edges (tagged `code-ast`, with confidence weights). +- **Heuristic track** (all languages, including Java/Rust): regex-based extraction (tagged `code-heuristic`), which also covers languages the AST track does not. + +The WASM parser is a pure-JavaScript dependency — no native toolchain is required. If it fails to load for any reason, extraction falls back to the heuristic track and records an `AST_UNAVAILABLE` gap. Set `TEAMAI_SKIP_AST=1` to force heuristic-only extraction. + ## Commands | Command | Description | diff --git a/README.zh-CN.md b/README.zh-CN.md index 40eb3c7b..b437dff4 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -215,6 +215,13 @@ teamai codebase --lint # 健康检查 图谱存储组件、接口、配置和跨仓库依赖边。`teamai recall` 利用图谱进行增强排名。 当召回命中 codebase 页面时，结果会附带一行 `Sources:`，列出相关源文件路径，供 agent 直接作为代码改动的入口，无需重新探索代码库。 +依赖边来自两条并行的提取轨道，重叠时以 AST 结果优先： + +- **AST 轨**（TypeScript/JavaScript、Python、Go）：使用 WASM 版 [tree-sitter](https://tree-sitter.github.io/) 解析器，将 `import`/`require`、调用点、以及 TS `implements` 子句解析为精确的文件到文件 `DEPENDS_ON` / `REFERENCES` / `IMPLEMENTS` 边（标记为 `code-ast`，带置信度权重）。 +- **启发式轨**（所有语言，含 Java/Rust）：基于正则的提取（标记为 `code-heuristic`），同时覆盖 AST 轨未支持的语言。 + +WASM 解析器是纯 JavaScript 依赖，无需任何原生编译工具链。若因任何原因加载失败，提取会降级到启发式轨并记录一条 `AST_UNAVAILABLE` gap。设置 `TEAMAI_SKIP_AST=1` 可强制仅使用启发式提取。 + ## 命令一览 | 命令 | 说明 | diff --git a/docs/usage-guide.md b/docs/usage-guide.md index 45dbff4c..225ea69b 100644 --- a/docs/usage-guide.md +++ b/docs/usage-guide.md @@ -774,6 +774,7 @@ Configurable environment variables: | `TEAMAI_SKILL_DOWNLOAD_HOSTS` | Allowlist of hosts for skill `download_url` (empty = allow all) | | `TEAMAI_ALLOW_SANDBOX_REPORT` | Set to `1` to force report/sync inside a CloudStudio sandbox (see note below) | | `TEAMAI_DISABLE_REMOTE_CMD` | Set to `1` to reject server-pushed `uninstall_teamai`, `install_hook_rule`, and `uninstall_hook_rule` commands (they are acked `failed`) | +| `TEAMAI_SKIP_AST` | Set to `1` to force heuristic-only code extraction, skipping the WASM tree-sitter AST track | > **Privacy:** The install path and machine id are only hashed locally to derive `local_agent_id` — they are never reported. @@ -815,6 +816,8 @@ teamai import --from-repo https://github.com/org/repo --skip-enrich The graph stores components, interfaces, configs, and cross-repo dependencies. `teamai recall` uses the graph for BM25 + graph-boosted ranking. +Dependency edges are extracted by two parallel tracks: a WASM tree-sitter **AST track** (TypeScript/JavaScript, Python, Go) that resolves imports, calls, and TS `implements` clauses to precise file-to-file edges (`code-ast`), and a regex **heuristic track** (all languages, `code-heuristic`) that also covers languages the AST track does not. AST results win on overlap. The AST parser needs no native toolchain; on load failure, extraction falls back to heuristics and records an `AST_UNAVAILABLE` gap. Set `TEAMAI_SKIP_AST=1` to force heuristic-only extraction. + ```bash # Graph health check teamai codebase --lint diff --git a/docs/usage-guide.zh-CN.md b/docs/usage-guide.zh-CN.md index 66e783d1..c6fe88f7 100644 --- a/docs/usage-guide.zh-CN.md +++ b/docs/usage-guide.zh-CN.md @@ -770,6 +770,7 @@ agent hook 规则： | `TEAMAI_SKILL_DOWNLOAD_HOSTS` | skill `download_url` host 白名单（空 = 全部放行） | | `TEAMAI_ALLOW_SANDBOX_REPORT` | 设为 `1` 可强制在 CloudStudio 沙箱内 report/sync（见下方说明） | | `TEAMAI_DISABLE_REMOTE_CMD` | 设为 `1` 可拒绝服务端下发的 `uninstall_teamai`、`install_hook_rule`、`uninstall_hook_rule` 命令（会 ack `failed`） | +| `TEAMAI_SKIP_AST` | 设为 `1` 时强制仅用启发式提取，跳过 WASM tree-sitter AST 轨 | > **隐私**：install path 和 machine id 仅在本地哈希以派生 `local_agent_id`，不会上报。 @@ -810,6 +811,8 @@ teamai import --from-repo https://github.com/org/repo --skip-enrich 图谱存储组件、接口、配置和跨仓库依赖关系。`teamai recall` 利用图谱进行 BM25 + graph-boost 增强排名。 +依赖边由两条并行轨道提取：WASM tree-sitter **AST 轨**（TypeScript/JavaScript、Python、Go），将 import、调用、以及 TS `implements` 子句解析为精确的文件到文件边（`code-ast`）；以及正则 **启发式轨**（所有语言，`code-heuristic`），同时覆盖 AST 轨未支持的语言。重叠时 AST 结果优先。AST 解析器无需原生编译工具链；加载失败时提取会降级到启发式并记录一条 `AST_UNAVAILABLE` gap。设置 `TEAMAI_SKIP_AST=1` 可强制仅用启发式提取。 + ```bash # 图谱健康检查 teamai codebase --lint diff --git a/package-lock.json b/package-lock.json index 39173850..fd7fabbf 100644 --- a/package-lock.json +++ b/package-lock.json @@ -18,6 +18,8 @@ "ora": "^8.1.0", "simple-git": "^3.36.0", "smol-toml": "^1.3.1", + "tree-sitter-wasms": "0.1.13", + "web-tree-sitter": "0.25.10", "yaml": "^2.6.0", "zod": "^3.24.0" }, @@ -4932,6 +4934,15 @@ "tree-kill": "cli.js" } }, + "node_modules/tree-sitter-wasms": { + "version": "0.1.13", + "resolved": "https://mirrors.tencent.com/npm/tree-sitter-wasms/-/tree-sitter-wasms-0.1.13.tgz", + "integrity": "sha512-wT+cR6DwaIz80/vho3AvSF0N4txuNx/5bcRKoXouOfClpxh/qqrF4URNLQXbbt8MaAxeksZcZd1j8gcGjc+QxQ==", + "license": "Unlicense", + "dependencies": { + "tree-sitter-wasms": "^0.1.11" + } + }, "node_modules/trim-newlines": { "version": "3.0.1", "resolved": "https://mirrors.tencent.com/npm/trim-newlines/-/trim-newlines-3.0.1.tgz", @@ -5643,6 +5654,20 @@ "integrity": "sha512-whLdWMYL2TwI08hn8/ZqAbrVemu0LNaNNJZX73O6qaIdCTfXutsLhMkjdENX0qhsQ9uIimo4/aQOmXkoon2nDQ==", "dev": true }, + "node_modules/web-tree-sitter": { + "version": "0.25.10", + "resolved": "https://mirrors.tencent.com/npm/web-tree-sitter/-/web-tree-sitter-0.25.10.tgz", + "integrity": "sha512-Y09sF44/13XvgVKgO2cNDw5rGk6s26MgoZPXLESvMXeefBf7i6/73eFurre0IsTW6E14Y0ArIzhUMmjoc7xyzA==", + "license": "MIT", + "peerDependencies": { + "@types/emscripten": "^1.40.0" + }, + "peerDependenciesMeta": { + "@types/emscripten": { + "optional": true + } + } + }, "node_modules/which": { "version": "2.0.2", "resolved": "https://mirrors.tencent.com/npm/which/-/which-2.0.2.tgz", diff --git a/package.json b/package.json index b7ce68d0..cfc257c5 100644 --- a/package.json +++ b/package.json @@ -55,6 +55,8 @@ "ora": "^8.1.0", "simple-git": "^3.36.0", "smol-toml": "^1.3.1", + "tree-sitter-wasms": "0.1.13", + "web-tree-sitter": "0.25.10", "yaml": "^2.6.0", "zod": "^3.24.0" }, diff --git a/src/__tests__/ast-extract.test.ts b/src/__tests__/ast-extract.test.ts new file mode 100644 index 00000000..25accb67 --- /dev/null +++ b/src/__tests__/ast-extract.test.ts @@ -0,0 +1,271 @@ +import { describe, it, expect, beforeEach } from 'vitest'; + +import type { CodeCollectedFile } from '../wiki-engine/code-knowledge/code-collector.js'; +import { + astAvailable, + extractStructuralGraphAsFacts, +} from '../wiki-engine/code-knowledge/ast/index.js'; +import { resetParserRegistryForTests } from '../wiki-engine/code-knowledge/ast/parser-registry.js'; +import { mergeCodeFacts } from '../wiki-engine/code-knowledge/ast/merge-edges.js'; +import type { CodeFact } from '../wiki-engine/code-knowledge/code-extractors.js'; +import { parseEdgeProvenance, edgeProvenanceRank, edgeReason } from '../enrich-with-ai.js'; + +/** + * Build an in-memory collected file for AST extraction tests. + * + * repoRoot is a virtual path; import resolution relies on the in-memory + * known-files set (relativePath), so no real files need to exist on disk. + */ +function makeFile(relativePath: string, content: string): CodeCollectedFile { + const language = relativePath.endsWith('.py') + ? 'python' + : relativePath.endsWith('.go') + ? 'go' + : 'typescript'; + return { + path: `/virtual/${relativePath}`, + relativePath, + language, + sha256: 'test', + content, + }; +} + +const REPO_ROOT = '/virtual'; + +describe('AST structural extraction (web-tree-sitter WASM)', () => { + beforeEach(() => { + resetParserRegistryForTests(); + }); + + it('reports AST as available in this environment', () => { + expect(astAvailable()).toBe(true); + }); + + it('resolves a TypeScript relative import into a precise code-ast DEPENDS_ON edge', async () => { + const files = [ + makeFile('src/a.ts', 'import { helper } from "./b";\nexport function run() {\n return helper();\n}\n'), + makeFile('src/b.ts', 'export function helper() {\n return 42;\n}\n'), + ]; + + const { facts, result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + const dependsOn = result.edges.find( + (e) => e.from === 'src/a.ts' && e.to === 'src/b.ts' && e.relation === 'DEPENDS_ON', + ); + expect(dependsOn).toBeDefined(); + expect(dependsOn?.source).toBe('code-ast'); + + // The adapted CodeFact carries the resolved target and a (code-ast) marker. + const astFact = facts.find((f) => f.file === 'src/a.ts' && f.name === 'src/b.ts'); + expect(astFact).toBeDefined(); + expect(astFact?.kind).toBe('relation'); + expect(astFact?.detail).toContain('(code-ast)'); + + expect(result.stats.imports).toBeGreaterThanOrEqual(1); + expect(result.stats.importsResolved).toBeGreaterThanOrEqual(1); + }); + + it('resolves a cross-file call into a code-ast REFERENCES edge', async () => { + const files = [ + makeFile('src/a.ts', 'import { helper } from "./b";\nexport function run() {\n return helper();\n}\n'), + makeFile('src/b.ts', 'export function helper() {\n return 42;\n}\n'), + ]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + const references = result.edges.find( + (e) => e.from === 'src/a.ts' && e.to === 'src/b.ts' && e.relation === 'REFERENCES', + ); + expect(references).toBeDefined(); + expect(references?.source).toBe('code-ast'); + }); + + it('extracts symbols and resolves a sibling module import from Python', async () => { + // "from util import boot" captures module_name "util", which resolves + // relative to main.py's directory → sibling pkg/util.py. + const files = [ + makeFile('pkg/main.py', 'from util import boot\n\nclass App:\n def start(self):\n return boot()\n'), + makeFile('pkg/util.py', 'def boot():\n return 1\n'), + ]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + // class App + def start + def boot = at least 3 symbols + expect(result.stats.symbols).toBeGreaterThanOrEqual(3); + const edge = result.edges.find((e) => e.from === 'pkg/main.py' && e.to === 'pkg/util.py'); + expect(edge).toBeDefined(); + expect(edge?.source).toBe('code-ast'); + }); + + it('treats module-level Python defs as exported so cross-file calls resolve to a symbol', async () => { + // Regression: tree-sitter-python has no "__export__" node, so export + // detection must fall back to module-level class/def. Without it every + // Python symbol is exported=false and symbol-level call resolution never + // fires (only the coarser file-level edge survives). + const files = [ + makeFile('pkg/main.py', 'from util import boot\n\ndef run():\n return boot()\n'), + makeFile('pkg/util.py', 'def boot():\n return 1\n'), + ]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + const bootSymbol = result.symbols.find((s) => s.name === 'boot' && s.file === 'pkg/util.py'); + expect(bootSymbol?.exported).toBe(true); + + // The call boot() in main.py resolves to boot's symbol id in util.py. + const call = result.callSites.find((c) => c.calleeText === 'boot' && c.fromFile === 'pkg/main.py'); + expect(call?.resolvedTargetId).toBe('pkg/util.py#Function:boot'); + expect(call?.resolvedTargetFile).toBe('pkg/util.py'); + }); + + it('extracts symbols from Go', async () => { + const goSrc = [ + 'package main', + '', + 'import "fmt"', + '', + 'type Server struct{}', + '', + 'func Run() {', + ' fmt.Println("hi")', + '}', + '', + ].join('\n'); + const files = [makeFile('main.go', goSrc)]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + // Go: func Run + type Server = 2 symbols + expect(result.stats.symbols).toBeGreaterThanOrEqual(2); + // "fmt" is an external package import → recorded as an EXTERNAL_IMPORT gap + expect(result.gaps.some((g) => g.kind === 'EXTERNAL_IMPORT')).toBe(true); + }); + + it('resolves a multi-segment Python import (from a.b.c import x) to a nested file', async () => { + const files = [ + makeFile('app/main.py', 'from pkg.sub.helper import boot\n\ndef run():\n return boot()\n'), + makeFile('app/pkg/sub/helper.py', 'def boot():\n return 1\n'), + ]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + const edge = result.edges.find( + (e) => e.from === 'app/main.py' && e.to === 'app/pkg/sub/helper.py', + ); + expect(edge).toBeDefined(); + expect(edge?.relation).toBe('DEPENDS_ON'); + expect(edge?.source).toBe('code-ast'); + }); + + it('produces an IMPLEMENTS edge for a TS class implementing an imported interface', async () => { + const files = [ + makeFile('src/svc.ts', 'import { IFoo } from "./iface";\n\nexport class Svc implements IFoo {\n run() {}\n}\n'), + makeFile('src/iface.ts', 'export interface IFoo {\n run(): void;\n}\n'), + ]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + const impl = result.edges.find( + (e) => e.relation === 'IMPLEMENTS' && e.from === 'src/svc.ts' && e.to === 'src/iface.ts', + ); + expect(impl).toBeDefined(); + expect(impl?.source).toBe('code-ast'); + }); + + it('produces an IMPLEMENTS edge for a same-file interface', async () => { + const files = [ + makeFile('src/only.ts', 'export interface IBar {\n go(): void;\n}\n\nexport class Impl implements IBar {\n go() {}\n}\n'), + ]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + const impl = result.edges.find((e) => e.relation === 'IMPLEMENTS' && e.from === 'src/only.ts'); + expect(impl).toBeDefined(); + expect(impl?.to).toBe('src/only.ts'); + }); + + it('records unresolved external imports as gaps', async () => { + const files = [ + makeFile('src/a.ts', 'import { thing } from "some-external-pkg";\nexport const x = thing;\n'), + ]; + + const { result } = await extractStructuralGraphAsFacts({ repoRoot: REPO_ROOT, files }); + + expect(result.gaps.some((g) => g.kind === 'EXTERNAL_IMPORT')).toBe(true); + }); + + it('mergeCodeFacts lets AST relation facts win over heuristic relation facts on the same line', () => { + const astFacts: CodeFact[] = [ + { + kind: 'relation', + name: 'src/b.ts', + file: 'src/a.ts', + lineStart: 1, + lineEnd: 1, + detail: 'DEPENDS_ON → src/b.ts (code-ast)', + confidence: 'EXTRACTED', + evidenceType: 'usage', + }, + ]; + const heuristicFacts: CodeFact[] = [ + { + kind: 'relation', + name: './b', + file: 'src/a.ts', + lineStart: 1, + lineEnd: 1, + detail: 'import ./b', + confidence: 'INFERRED', + evidenceType: 'usage', + }, + ]; + + const merged = mergeCodeFacts(astFacts, heuristicFacts); + const relations = merged.filter((f) => f.kind === 'relation' && f.file === 'src/a.ts' && f.lineStart === 1); + expect(relations).toHaveLength(1); + expect(relations[0]?.detail).toContain('(code-ast)'); + }); + + describe('enrich edge provenance', () => { + it('preserves AST relation and source from a code-ast fact detail', () => { + expect(parseEdgeProvenance('REFERENCES → src/repo.ts (code-ast)')).toEqual({ + relation: 'REFERENCES', + source: 'code-ast', + }); + expect(parseEdgeProvenance('DEPENDS_ON → src/b.ts (code-ast)')).toEqual({ + relation: 'DEPENDS_ON', + source: 'code-ast', + }); + expect(parseEdgeProvenance('IMPLEMENTS → src/iface.ts (code-ast)')).toEqual({ + relation: 'IMPLEMENTS', + source: 'code-ast', + }); + }); + + it('falls back to DEPENDS_ON / code-heuristic for a raw regex fact detail', () => { + expect(parseEdgeProvenance('import { x } from "./y";')).toEqual({ + relation: 'DEPENDS_ON', + source: 'code-heuristic', + }); + }); + + it('ranks code-ast DEPENDS_ON above REFERENCES/IMPLEMENTS and heuristic (deterministic merge)', () => { + const astDepends = { relation: 'DEPENDS_ON', source: 'code-ast' as const }; + const astRefs = { relation: 'REFERENCES', source: 'code-ast' as const }; + const astImpl = { relation: 'IMPLEMENTS', source: 'code-ast' as const }; + const heuristic = { relation: 'DEPENDS_ON', source: 'code-heuristic' as const }; + + expect(edgeProvenanceRank(astDepends)).toBeGreaterThan(edgeProvenanceRank(astRefs)); + expect(edgeProvenanceRank(astRefs)).toBeGreaterThan(edgeProvenanceRank(astImpl)); + expect(edgeProvenanceRank(astImpl)).toBeGreaterThan(edgeProvenanceRank(heuristic)); + expect(edgeProvenanceRank(heuristic)).toBe(0); + }); + + it('edgeReason matches the resolved relation', () => { + expect(edgeReason('a', 'b', 'REFERENCES')).toBe('a references b'); + expect(edgeReason('a', 'b', 'IMPLEMENTS')).toBe('a implements b'); + expect(edgeReason('a', 'b', 'DEPENDS_ON')).toBe('a imports from b'); + }); + }); +}); diff --git a/src/__tests__/ast-import-resolver.test.ts b/src/__tests__/ast-import-resolver.test.ts new file mode 100644 index 00000000..1e14b384 --- /dev/null +++ b/src/__tests__/ast-import-resolver.test.ts @@ -0,0 +1,107 @@ +import { describe, it, expect } from 'vitest'; + +import { resolveImportSpecifier } from '../wiki-engine/code-knowledge/ast/import-resolver.js'; + +/** + * Build a fileExists predicate backed by an in-memory Set. + * + * No real filesystem access occurs; the predicate simply checks whether the + * given repo-root-relative path is a member of `known`. + */ +function makeExists(known: Set): (relativePath: string) => Promise { + return async (relativePath: string) => known.has(relativePath); +} + +const REPO_ROOT = '/repo'; + +describe('resolveImportSpecifier — Python absolute package imports', () => { + it('resolves absolute package import to a module .py file', async () => { + const known = new Set(['hai_flow/conf.py']); + const result = await resolveImportSpecifier( + REPO_ROOT, + 'hai_flow/app.py', + 'hai_flow.conf', + makeExists(known), + ); + expect(result).toBeDefined(); + expect(result?.targetFile).toBe('hai_flow/conf.py'); + expect(result?.confidence).toBe('EXTRACTED'); + }); + + it('resolves absolute package import to a package __init__.py', async () => { + const known = new Set(['hai_flow/core/__init__.py']); + const result = await resolveImportSpecifier( + REPO_ROOT, + 'hai_flow/app.py', + 'hai_flow.core', + makeExists(known), + ); + expect(result).toBeDefined(); + expect(result?.targetFile).toBe('hai_flow/core/__init__.py'); + expect(result?.confidence).toBe('EXTRACTED'); + }); + + it('resolves absolute package import from a deeply nested file using repoRoot, not fromDir', async () => { + // fromFile is three directories deep; the target is rooted at repoRoot. + // This verifies that resolution is NOT relative to the importing file. + const known = new Set(['hai_flow/utils/string_util.py']); + const result = await resolveImportSpecifier( + REPO_ROOT, + 'hai_flow/api/v2/views.py', + 'hai_flow.utils.string_util', + makeExists(known), + ); + expect(result).toBeDefined(); + expect(result?.targetFile).toBe('hai_flow/utils/string_util.py'); + expect(result?.confidence).toBe('EXTRACTED'); + }); + + it('still resolves relative Python imports correctly (no regression)', async () => { + // Specifier must be "./helper" (not ".helper") so that path.join strips + // the leading dot and produces "hai_flow/api/helper", not the hidden-file + // path "hai_flow/api/.helper". This verifies resolveRelativeImport still + // fires and is not preempted by the new resolveAbsolutePackageImport branch. + const known = new Set(['hai_flow/api/helper.py']); + const result = await resolveImportSpecifier( + REPO_ROOT, + 'hai_flow/api/views.py', + './helper', + makeExists(known), + ); + expect(result).toBeDefined(); + expect(result?.targetFile).toBe('hai_flow/api/helper.py'); + }); + + it('returns undefined for an external package when no matching file exists in the repo', async () => { + const known = new Set(); + + const resultOs = await resolveImportSpecifier( + REPO_ROOT, + 'hai_flow/app.py', + 'os', + makeExists(known), + ); + expect(resultOs).toBeUndefined(); + + const resultApscheduler = await resolveImportSpecifier( + REPO_ROOT, + 'hai_flow/app.py', + 'apscheduler.schedulers.background', + makeExists(known), + ); + expect(resultApscheduler).toBeUndefined(); + }); + + it('does not apply Python dot-path mapping for TypeScript source files', async () => { + // pkg.mod looks like a Python absolute import but fromFile is .ts, so + // the absolute-package branch is skipped; no tsconfig present, so undefined. + const known = new Set(['pkg/mod.py']); + const result = await resolveImportSpecifier( + REPO_ROOT, + 'app.ts', + 'pkg.mod', + makeExists(known), + ); + expect(result).toBeUndefined(); + }); +}); diff --git a/src/__tests__/deep-enrich-graph-preserve.test.ts b/src/__tests__/deep-enrich-graph-preserve.test.ts new file mode 100644 index 00000000..9a5303f9 --- /dev/null +++ b/src/__tests__/deep-enrich-graph-preserve.test.ts @@ -0,0 +1,87 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { mkdtemp, mkdir, writeFile, readFile, rm } from 'node:fs/promises'; +import os from 'node:os'; +import path from 'node:path'; + +// Mock the AI client so deepEnrich's LLM calls fail instantly instead of +// waiting on real 600s timeouts. deepEnrich treats these as non-blocking skips. +vi.mock('../utils/ai-client.js', () => ({ + getAICliName: () => 'mock-cli', + callClaude: vi.fn(async () => { + throw new Error('mock: AI unavailable'); + }), + // Reject to trigger deepEnrich's sequential fallback, where each callClaude + // then rejects → component skipped. Mirrors the real 600s-timeout skip path. + callClaudeParallel: vi.fn(async () => { + throw new Error('mock: AI batch unavailable'); + }), +})); + +import { deepEnrich } from '../deep-enrich.js'; + +const AST_GRAPH = { + schemaVersion: 'team-wiki.graph-index.v1', + generatedAt: '2026-01-01T00:00:00Z', + nodes: [ + { slug: 'component/A', type: 'component', confidence: 'EXTRACTED', title: 'A' }, + { slug: 'component/B', type: 'component', confidence: 'EXTRACTED', title: 'B' }, + ], + edges: [ + { from: 'a.ts', to: 'b.ts', relation: 'DEPENDS_ON', source: 'code-ast', weight: 0.9 }, + { from: 'a.ts', to: 'c.ts', relation: 'REFERENCES', source: 'code-heuristic', weight: 0.8 }, + ], +}; + +const MANIFEST = { + schemaVersion: 'team-wiki.codebase-output-manifest.v2', + project: 'faketest', + generatedAt: '2026-01-01T00:00:00Z', + components: [ + { + slug: 'A', + docPath: 'evidence/code/faketest/A.md', + title: 'A', + category: 'component', + confidence: 'INFERRED', + responsibilities: [], + entrypoints: [], + }, + ], + edges: [], +}; + +function astCount(g: { edges: Array<{ source?: string }> }): number { + return g.edges.filter((e) => e.source === 'code-ast').length; +} + +describe('deepEnrich preserves per-repo graph-index.json (AST edges)', () => { + let root: string; + + beforeEach(async () => { + root = await mkdtemp(path.join(os.tmpdir(), 'de-graph-')); + }); + afterEach(async () => { + await rm(root, { recursive: true, force: true }); + }); + + it('does not mutate evidence/.indices/graph-index.json across a full deepEnrich run', async () => { + const wikiRoot = path.join(root, 'teamwiki'); + const evidenceDir = path.join(wikiRoot, 'evidence', 'code', 'faketest'); + const idxDir = path.join(evidenceDir, '.indices'); + await mkdir(idxDir, { recursive: true }); + const graphPath = path.join(idxDir, 'graph-index.json'); + await writeFile(graphPath, JSON.stringify(AST_GRAPH, null, 2)); + await writeFile(path.join(evidenceDir, '_manifest.json'), JSON.stringify(MANIFEST, null, 2)); + + const before = JSON.parse(await readFile(graphPath, 'utf8')); + expect(before.edges).toHaveLength(2); + expect(astCount(before)).toBe(1); + + await deepEnrich({ project: 'faketest', evidenceDir, wikiRoot }); + + const after = JSON.parse(await readFile(graphPath, 'utf8')); + // The per-repo graph must be untouched by enrichment. + expect(after.edges).toHaveLength(2); + expect(astCount(after)).toBe(1); + }); +}); diff --git a/src/__tests__/import-lock.test.ts b/src/__tests__/import-lock.test.ts new file mode 100644 index 00000000..a4f5416a --- /dev/null +++ b/src/__tests__/import-lock.test.ts @@ -0,0 +1,120 @@ +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import { mkdtemp, rm, writeFile, stat } from 'node:fs/promises'; +import path from 'node:path'; +import os from 'node:os'; +import { acquireImportLock, isImportInProgress } from '../utils/import-lock.js'; + +/** + * Mirrors the lockPathFor formula from the implementation. + * Lock is placed in the parent of the resolved teamRepoPath. + */ +function lockPathFor(teamRepoPath: string): string { + return path.join(path.dirname(path.resolve(teamRepoPath)), '.teamai-import.lock'); +} + +/** Returns true if a file exists at the given absolute path. */ +async function fileExists(filePath: string): Promise { + try { + await stat(filePath); + return true; + } catch { + return false; + } +} + +describe('import-lock', () => { + let root: string; + let teamRepo: string; + let lockPath: string; + + beforeEach(async () => { + root = await mkdtemp(path.join(os.tmpdir(), 'implock-')); + teamRepo = path.join(root, 'team-repo'); + lockPath = lockPathFor(teamRepo); + }); + + afterEach(async () => { + await rm(root, { recursive: true, force: true }); + }); + + it('acquire writes lock file, release deletes it', async () => { + const release = await acquireImportLock(teamRepo); + expect(await fileExists(lockPath)).toBe(true); + await release(); + expect(await fileExists(lockPath)).toBe(false); + }); + + it('isImportInProgress is true while locked and false after release', async () => { + const release = await acquireImportLock(teamRepo); + expect(await isImportInProgress(teamRepo)).toBe(true); + await release(); + expect(await isImportInProgress(teamRepo)).toBe(false); + }); + + it('isImportInProgress is false when no lock file exists', async () => { + expect(await isImportInProgress(teamRepo)).toBe(false); + }); + + it('first release keeps lock alive (ref count 2 → 1), second release removes it (1 → 0)', async () => { + const release1 = await acquireImportLock(teamRepo); + const release2 = await acquireImportLock(teamRepo); + + await release1(); + expect(await fileExists(lockPath)).toBe(true); + expect(await isImportInProgress(teamRepo)).toBe(true); + + await release2(); + expect(await fileExists(lockPath)).toBe(false); + expect(await isImportInProgress(teamRepo)).toBe(false); + }); + + it('release is idempotent: calling the same release twice does not throw', async () => { + const release = await acquireImportLock(teamRepo); + await release(); + await expect(release()).resolves.toBeUndefined(); + }); + + it('stale lock (startedAt 3 hours ago) is inactive and the lock file is removed', async () => { + const staleMeta = { + pid: process.pid, + host: os.hostname(), + startedAt: new Date(Date.now() - 3 * 3600 * 1000).toISOString(), + teamRepo, + }; + await writeFile(lockPath, JSON.stringify(staleMeta, null, 2), 'utf8'); + + expect(await isImportInProgress(teamRepo)).toBe(false); + expect(await fileExists(lockPath)).toBe(false); + }); + + it('lock for a dead process (pid 2147483646) is inactive and the lock file is removed', async () => { + const deadMeta = { + pid: 2147483646, + host: os.hostname(), + startedAt: new Date().toISOString(), + teamRepo, + }; + await writeFile(lockPath, JSON.stringify(deadMeta, null, 2), 'utf8'); + + expect(await isImportInProgress(teamRepo)).toBe(false); + expect(await fileExists(lockPath)).toBe(false); + }); + + it('lock from a different host within TTL is treated as active', async () => { + const remoteMeta = { + pid: process.pid, + host: 'some-other-host-xyz', + startedAt: new Date().toISOString(), + teamRepo, + }; + await writeFile(lockPath, JSON.stringify(remoteMeta, null, 2), 'utf8'); + + expect(await isImportInProgress(teamRepo)).toBe(true); + }); + + it('malformed lock file with fresh mtime is treated as active', async () => { + await writeFile(lockPath, 'not-json{', 'utf8'); + + expect(await isImportInProgress(teamRepo)).toBe(true); + }); +}); diff --git a/src/codebase-extract.ts b/src/codebase-extract.ts index b6359cc4..6340f9ba 100644 --- a/src/codebase-extract.ts +++ b/src/codebase-extract.ts @@ -21,6 +21,10 @@ import { buildIndexHubOverlay, mergeGraphs, saveGraphIndex, + extractStructuralGraphAsFacts, + astAvailable, + mergeCodeFacts, + formatAstStatsSummary, } from './wiki-engine/adapters/index.js'; import type { CodeFact, InterfaceInventory, CallChain } from './wiki-engine/adapters/index.js'; import { @@ -612,6 +616,49 @@ export async function extractCodebase(opts: ExtractCodebaseOptions): Promise 0 && astAvailable()) { + try { + const { facts: astFacts, result: astResult } = await extractStructuralGraphAsFacts({ + repoRoot: root, + files, + }); + facts = mergeCodeFacts(astFacts, facts); + let gapSeq = 0; + for (const gap of astResult.gaps) { + astGaps.push({ + id: `AST-${gap.kind}-${gapSeq++}`, + kind: gap.kind, + description: gap.message, + source: gap.sources.join(', '), + }); + } + if (!opts.json) { + console.log(chalk.dim(` [AST: ${formatAstStatsSummary(astResult.stats)}]`)); + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + astGaps.push({ + id: 'AST-UNAVAILABLE-0', + kind: 'AST_UNAVAILABLE', + description: `code-ast failed: ${message}; used code-heuristic only.`, + source: 'code-ast', + }); + } + } else if (files.length > 0) { + astGaps.push({ + id: 'AST-UNAVAILABLE-0', + kind: 'AST_UNAVAILABLE', + description: 'web-tree-sitter WASM runtime unavailable or TEAMAI_SKIP_AST=1; used code-heuristic only.', + source: 'code-ast', + }); + } + const graph: GraphIndex = buildCodeGraph(facts); // Call chain tracing (entry → orchestration → service → data) @@ -733,7 +780,7 @@ export async function extractCodebase(opts: ExtractCodebaseOptions): Promise → (code-ast)"); + // regex heuristic facts carry a raw source line, so they keep the defaults. const moduleImports = ctx.facts.filter(f => f.kind === 'relation' && f.file.startsWith(name + '/')); - const targetModules = new Set(); + const targetEdges = new Map(); for (const imp of moduleImports) { const resolved = resolveImportToModule(imp.file, imp.name); - if (resolved && resolved !== name) { - targetModules.add(resolved); + if (!resolved || resolved === name) { + continue; + } + const provenance = parseEdgeProvenance(imp.detail); + const existing = targetEdges.get(resolved); + // Deterministic winner regardless of fact order: highest provenance rank wins. + if (!existing || edgeProvenanceRank(provenance) > edgeProvenanceRank(existing)) { + targetEdges.set(resolved, provenance); } } - for (const target of targetModules) { + for (const [target, provenance] of targetEdges) { if (moduleResults.some(m => m.name === target)) { edges.push({ from: name, to: target, - relation: 'DEPENDS_ON', + relation: provenance.relation, confidence: 'EXTRACTED', - source: 'code-heuristic', - reason: `${name} imports from ${target}`, + source: provenance.source, + reason: edgeReason(name, target, provenance.relation), }); } } @@ -227,6 +235,61 @@ export async function enrichWithAI(ctx: EnrichContext): Promise → (code-ast)"`; anything else + * (regex heuristic facts, whose detail is a raw source line) falls back to a + * DEPENDS_ON / code-heuristic default. + */ +export function parseEdgeProvenance(detail: string): { relation: string; source: ManifestEdgeSource } { + if (detail.includes('(code-ast)')) { + const relation = detail.startsWith('REFERENCES') + ? 'REFERENCES' + : detail.startsWith('IMPLEMENTS') + ? 'IMPLEMENTS' + : 'DEPENDS_ON'; + return { relation, source: 'code-ast' }; + } + return { relation: 'DEPENDS_ON', source: 'code-heuristic' }; +} + +/** + * Rank an edge provenance so a deterministic winner emerges when several facts + * describe the same module pair (independent of fact ordering). + * + * AST provenance always outranks heuristic. Among AST edges the import + * dependency (DEPENDS_ON) is the canonical module-level relation and wins, + * keeping the edge consistent with the "imports from" reason. + */ +export function edgeProvenanceRank(provenance: { relation: string; source: ManifestEdgeSource }): number { + if (provenance.source !== 'code-ast') { + return 0; + } + switch (provenance.relation) { + case 'DEPENDS_ON': + return 3; + case 'REFERENCES': + return 2; + case 'IMPLEMENTS': + return 1; + default: + return 1; + } +} + +/** Build a human-readable edge reason that matches the resolved relation. */ +export function edgeReason(from: string, to: string, relation: string): string { + switch (relation) { + case 'REFERENCES': + return `${from} references ${to}`; + case 'IMPLEMENTS': + return `${from} implements ${to}`; + default: + return `${from} imports from ${to}`; + } +} + export async function writeManifest(manifest: CodebaseOutputManifestV2, outputDir: string): Promise { await mkdir(outputDir, { recursive: true }); const manifestPath = path.join(outputDir, '_manifest.json'); diff --git a/src/import-repo-list.ts b/src/import-repo-list.ts index 2cae7eb8..b2f20d3c 100644 --- a/src/import-repo-list.ts +++ b/src/import-repo-list.ts @@ -74,10 +74,23 @@ export async function importFromRepoList( // 1. 加载白名单 const repoListFile = await loadRepoList(listPath); + let releaseBatchLock: (() => Promise) | null = null; + if (!dryRun) { + try { + const { autoDetectInit } = await import('./config.js'); + const { localConfig: lc } = await autoDetectInit(); + const { acquireImportLock } = await import('./utils/import-lock.js'); + releaseBatchLock = await acquireImportLock(lc.repo.localPath); + } catch (e) { + log.debug(`[import-lock] batch lock acquire skipped: ${(e as Error).message}`); + } + } + const succeeded: number[] = []; const failed: Array<{ url: string; error: string }> = []; const skipped: Array<{ url: string; reason: string }> = []; + try { // 2. 分拣 org entry（暂不支持）与单仓 entry const singleEntries: ReturnType = []; for (const item of repoListFile.repos) { @@ -198,6 +211,12 @@ export async function importFromRepoList( } } + } finally { + if (releaseBatchLock) { + await releaseBatchLock(); + } + } + return { succeeded: succeeded.length, failed, diff --git a/src/import-repo.ts b/src/import-repo.ts index 4fa5e145..51e0bdc6 100644 --- a/src/import-repo.ts +++ b/src/import-repo.ts @@ -301,6 +301,9 @@ export async function importFromRepo(opts: ImportFromRepoOptions): Promise teamRepoDir = path.join(process.cwd(), '.teamai', 'team-repo'); } + const { acquireImportLock } = await import('./utils/import-lock.js'); + const releaseImportLock = await acquireImportLock(teamRepoDir); + try { // 4. Generate teamwiki/ knowledge graph artifacts + append AI narrative to overview.md const teamwikiRoot = output ? path.resolve(output, '..', 'teamwiki') @@ -482,4 +485,7 @@ export async function importFromRepo(opts: ImportFromRepoOptions): Promise log.debug(`[cache-index] touchCacheEntry failed: ${String(touchErr)}`); } } + } finally { + await releaseImportLock(); + } } diff --git a/src/team-push.ts b/src/team-push.ts index 2bd21055..9e321622 100644 --- a/src/team-push.ts +++ b/src/team-push.ts @@ -371,6 +371,11 @@ export async function reportUsageToTeam( log.debug(`Skipping report: ${repoPath} is not a dedicated team-repo root (safety guard)`); return; } + const { isImportInProgress } = await import('./utils/import-lock.js'); + if (await isImportInProgress(repoPath)) { + log.debug(`Skipping report: import in progress for ${repoPath} (would reset uncommitted artifacts)`); + return; + } const yamlPath = path.join(repoPath, 'teamai.yaml'); const workingContent = await readFileSafe(yamlPath); const committedContent = workingContent === null diff --git a/src/utils/import-lock.ts b/src/utils/import-lock.ts new file mode 100644 index 00000000..0671453f --- /dev/null +++ b/src/utils/import-lock.ts @@ -0,0 +1,174 @@ +import { writeFile, readFile, unlink, stat, mkdir } from 'node:fs/promises'; +import path from 'node:path'; +import os from 'node:os'; +import { log } from './logger.js'; + +/** Duration after which a lock file is considered stale and safe to ignore. */ +const LOCK_TTL_MS = 2 * 60 * 60 * 1000; // 2 hours; a lock older than this is stale + +/** Module-level reference counts for reentrant locking within the same process. */ +const refCounts = new Map(); + +/** Metadata written into the lock file. */ +interface ImportLockMeta { + pid: number; + host: string; + startedAt: string; // ISO + teamRepo: string; // resolved team repo path +} + +/** + * Returns the lock file path for a given team repo directory. + * Placed in the parent of the resolved team repo path so it is never touched + * by git operations (reset, clean, checkout) running inside the repo itself. + */ +function lockPathFor(teamRepoPath: string): string { + return path.join(path.dirname(path.resolve(teamRepoPath)), '.teamai-import.lock'); +} + +/** + * Acquire the import lock for the given team repo directory. + * Returns an idempotent async release function. + * + * Re-entrant within the same process: a reference count is maintained so that + * nested acquire/release pairs do not prematurely delete the lock file written + * by the outermost caller. The lock file is only created on the first acquire + * (ref count 0 → 1) and only deleted on the last release (ref count 1 → 0). + * + * Lock acquisition is best-effort: if writing the lock file fails, the error is + * logged at debug level and execution continues — the lock is a protective hint, + * not a hard barrier that should block an import from running. + * + * @param teamRepoPath - Path to the team repo directory (will be resolved). + * @returns An async release function. Calling it multiple times is safe (idempotent). + */ +export async function acquireImportLock(teamRepoPath: string): Promise<() => Promise> { + const lockPath = lockPathFor(teamRepoPath); + + const count = refCounts.get(lockPath) ?? 0; + refCounts.set(lockPath, count + 1); + + if (count === 0) { + // First acquire: write the lock file. + try { + await mkdir(path.dirname(lockPath), { recursive: true }); + const meta: ImportLockMeta = { + pid: process.pid, + host: os.hostname(), + startedAt: new Date().toISOString(), + teamRepo: path.resolve(teamRepoPath), + }; + await writeFile(lockPath, JSON.stringify(meta, null, 2), 'utf8'); + } catch (e) { + log.debug(`[import-lock] Failed to write lock file ${lockPath}: ${(e as Error).message}`); + } + } + + let released = false; + + return async (): Promise => { + if (released) return; + released = true; + + const cur = refCounts.get(lockPath) ?? 0; + const next = Math.max(0, cur - 1); + + if (next === 0) { + refCounts.delete(lockPath); + try { + await unlink(lockPath); + } catch (e) { + log.debug(`[import-lock] Failed to remove lock file ${lockPath}: ${(e as Error).message}`); + } + } else { + refCounts.set(lockPath, next); + } + }; +} + +/** + * Check whether an import is currently in progress for the given team repo. + * + * Intended to be called by reportUsageToTeam before executing git reset --hard. + * Returns true when a live import lock is detected, in which case the caller + * should skip the reset and pull to avoid overwriting uncommitted import artifacts. + * + * Detection strategy: + * - Same host: use process.kill(pid, 0) to probe whether the locking process is alive. + * - Different host: cannot probe remotely; treat any non-stale lock as active. + * - Stale lock (older than LOCK_TTL_MS): treat as abandoned and clean up. + * - Malformed lock: fall back to file mtime; a fresh but unreadable file is treated + * as active to err on the side of protecting uncommitted data. + * + * @param teamRepoPath - Path to the team repo directory (will be resolved). + * @returns True if an active import lock is detected, false otherwise. + */ +export async function isImportInProgress(teamRepoPath: string): Promise { + const lockPath = lockPathFor(teamRepoPath); + + let raw: string; + try { + raw = await readFile(lockPath, 'utf8'); + } catch { + // File absent (ENOENT) or unreadable — no lock. + return false; + } + + let meta: ImportLockMeta; + try { + meta = JSON.parse(raw) as ImportLockMeta; + + // Check TTL. + const parsed = Date.parse(meta.startedAt); + const age = Number.isNaN(parsed) ? Infinity : Date.now() - parsed; + + if (age > LOCK_TTL_MS) { + // Stale lock — clean up and report no active import. + try { + await unlink(lockPath); + } catch (e) { + log.debug(`[import-lock] Failed to remove stale lock ${lockPath}: ${(e as Error).message}`); + } + return false; + } + + if (meta.host === os.hostname()) { + // Same host: probe the locking process via signal 0. + try { + process.kill(meta.pid, 0); + // No exception thrown — process is alive — import is in progress. + return true; + } catch { + // Process is dead (ESRCH) — stale lock. + try { + await unlink(lockPath); + } catch (e) { + log.debug( + `[import-lock] Failed to remove dead-process lock ${lockPath}: ${(e as Error).message}`, + ); + } + return false; + } + } + + // Different host — cannot probe; treat as active within TTL. + return true; + } catch { + // JSON parse failed — fall back to mtime heuristic. + try { + const { mtimeMs } = await stat(lockPath); + if (Date.now() - mtimeMs <= LOCK_TTL_MS) { + // Fresh but unreadable — treat as active to protect uncommitted data. + return true; + } + try { + await unlink(lockPath); + } catch (e) { + log.debug(`[import-lock] Failed to remove malformed lock ${lockPath}: ${(e as Error).message}`); + } + return false; + } catch { + return false; + } + } +} diff --git a/src/wiki-engine/adapters/index.ts b/src/wiki-engine/adapters/index.ts index 90fe2c29..630d34c9 100644 --- a/src/wiki-engine/adapters/index.ts +++ b/src/wiki-engine/adapters/index.ts @@ -10,6 +10,13 @@ export { extractCodeFacts } from '../code-knowledge/code-extractors.js'; export type { CodeFact, CodeFactKind, CodeEvidenceType } from '../code-knowledge/code-extractors.js'; export { buildCodeGraph } from '../code-knowledge/code-graph.js'; +export { + extractStructuralGraphAsFacts, + astAvailable, + mergeCodeFacts, + formatAstStatsSummary, +} from '../code-knowledge/ast/index.js'; +export type { StructuralGraphResult, AstExtractionGap } from '../code-knowledge/ast/index.js'; export { detectCodeIncrementalChanges } from '../code-knowledge/code-incremental.js'; diff --git a/src/wiki-engine/code-knowledge/ast/adapt-code-facts.ts b/src/wiki-engine/code-knowledge/ast/adapt-code-facts.ts new file mode 100644 index 00000000..43268ff0 --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/adapt-code-facts.ts @@ -0,0 +1,49 @@ +import { type CodeFact, mapKindToEvidenceType } from "../code-extractors.js"; +import type { AstImport, StructuralEdge } from "./types.js"; + +type UnresolvedGap = { kind: string; message: string; sources: string[] }; + +export function structuralEdgesToCodeFacts(edges: StructuralEdge[]): CodeFact[] { + const facts: CodeFact[] = []; + for (const edge of edges) { + const evidence = edge.evidence[0]; + const lineStart = evidence?.lineStart ?? 1; + const lineEnd = evidence?.lineEnd ?? lineStart; + const name = edge.to; + + facts.push({ + kind: "relation", + name, + file: edge.from, + lineStart, + lineEnd, + detail: `${edge.relation} → ${edge.to} (${edge.source})`, + confidence: edge.confidence, + evidenceType: mapKindToEvidenceType("relation") + }); + } + return facts; +} + +export function unresolvedImportsToGaps(imports: AstImport[], resolvedKeys: Set): UnresolvedGap[] { + const gaps: UnresolvedGap[] = []; + for (const imp of imports) { + const key = `${imp.fromFile}:${imp.line}`; + if (resolvedKeys.has(key)) continue; + if (!imp.specifier.startsWith(".") && !imp.specifier.startsWith("/") && !imp.specifier.startsWith("@")) { + gaps.push({ + kind: "EXTERNAL_IMPORT", + message: `External package import not resolved: ${imp.specifier}`, + sources: [`${imp.fromFile}:${imp.line}`] + }); + } else { + gaps.push({ + kind: "UNRESOLVED_IMPORT", + message: `Could not resolve import "${imp.specifier}" from ${imp.fromFile}`, + sources: [`${imp.fromFile}:${imp.line}`] + }); + } + } + return gaps; +} + diff --git a/src/wiki-engine/code-knowledge/ast/call-resolver.ts b/src/wiki-engine/code-knowledge/ast/call-resolver.ts new file mode 100644 index 00000000..805554cb --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/call-resolver.ts @@ -0,0 +1,140 @@ +import type { AstCallSite, AstImport, AstSymbol } from "./types.js"; +import type { ResolvedImport } from "./import-resolver.js"; + +export interface ImportBindingMap { + /** Local name → exported symbol id in target file */ + localToSymbolId: Map; + /** Local name → resolved target file */ + localToFile: Map; +} + +export function buildImportBindingsForFile( + fromFile: string, + imports: AstImport[], + resolved: Map, + symbolsByFile: Map +): ImportBindingMap { + const localToSymbolId = new Map(); + const localToFile = new Map(); + + for (const imp of imports.filter((i) => i.fromFile === fromFile)) { + const key = `${imp.fromFile}:${imp.line}`; + const target = resolved.get(key); + if (!target) continue; + + localToFile.set(imp.defaultBinding ?? imp.namespaceBinding ?? "", target.targetFile); + + const targetSymbols = symbolsByFile.get(target.targetFile) ?? []; + + if (imp.defaultBinding) { + const def = targetSymbols.find((s) => s.exported && (s.kind === "function" || s.kind === "class")); + if (def) localToSymbolId.set(imp.defaultBinding, def.id); + localToFile.set(imp.defaultBinding, target.targetFile); + } + + if (imp.namespaceBinding) { + localToFile.set(imp.namespaceBinding, target.targetFile); + } + + for (const name of imp.namedBindings ?? []) { + const local = name; + const exported = targetSymbols.find((s) => s.name === name && s.exported); + if (exported) localToSymbolId.set(local, exported.id); + localToFile.set(local, target.targetFile); + } + } + + return { localToSymbolId, localToFile }; +} + +export function resolveCallSites( + callSites: AstCallSite[], + imports: AstImport[], + resolved: Map, + symbolsByFile: Map +): AstCallSite[] { + const bindingsByFile = new Map(); + return callSites.map((site) => { + let bindings = bindingsByFile.get(site.fromFile); + if (!bindings) { + bindings = buildImportBindingsForFile(site.fromFile, imports, resolved, symbolsByFile); + bindingsByFile.set(site.fromFile, bindings); + } + return resolveOneCall(site, symbolsByFile, bindings); + }); +} + +function resolveOneCall( + site: AstCallSite, + symbolsByFile: Map, + bindings: ImportBindingMap +): AstCallSite { + const callee = site.calleeText; + + if (!callee.includes(".")) { + const localSymbols = symbolsByFile.get(site.fromFile) ?? []; + const sameFile = localSymbols.find((s) => s.name === callee && (s.kind === "function" || s.kind === "class")); + if (sameFile) { + return { + ...site, + resolvedTargetId: sameFile.id, + resolvedTargetFile: site.fromFile, + confidence: "EXTRACTED" + }; + } + + const importedId = bindings.localToSymbolId.get(callee); + const importedFile = bindings.localToFile.get(callee); + if (importedId) { + return { + ...site, + resolvedTargetId: importedId, + resolvedTargetFile: importedFile, + confidence: "EXTRACTED" + }; + } + if (importedFile) { + return { ...site, resolvedTargetFile: importedFile, confidence: "INFERRED" }; + } + + return site; + } + + const [recv, member] = callee.split(".", 2); + if (!recv || !member) return site; + + const importedFile = bindings.localToFile.get(recv); + if (importedFile) { + const targetSymbols = symbolsByFile.get(importedFile) ?? []; + const sym = targetSymbols.find((s) => s.name === member); + if (sym) { + return { + ...site, + resolvedTargetId: sym.id, + resolvedTargetFile: importedFile, + receiver: recv, + confidence: "EXTRACTED" + }; + } + return { ...site, resolvedTargetFile: importedFile, receiver: recv, confidence: "INFERRED" }; + } + + const localSymbols = symbolsByFile.get(site.fromFile) ?? []; + const localClass = localSymbols.find((s) => s.name === recv && s.kind === "class"); + if (localClass) { + return { ...site, resolvedTargetFile: site.fromFile, confidence: "INFERRED" }; + } + + return site; +} + +export function callResolutionWeight(confidence: AstCallSite["confidence"]): number { + switch (confidence) { + case "EXTRACTED": + return 0.85; + case "INFERRED": + return 0.75; + default: + return 0.5; + } +} diff --git a/src/wiki-engine/code-knowledge/ast/import-bindings.ts b/src/wiki-engine/code-knowledge/ast/import-bindings.ts new file mode 100644 index 00000000..d54839b2 --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/import-bindings.ts @@ -0,0 +1,142 @@ +import type { Node } from "web-tree-sitter"; + +import type { AstImport } from "./types.js"; +import type { GrammarVariant } from "./parser-registry.js"; + +type ImportBindingResult = Pick; + +/** Parse import statement text into binding metadata (language-aware). */ +export function parseImportBindings( + importText: string, + variant: GrammarVariant +): ImportBindingResult { + if (variant === "python") { + return parsePythonImportBindings(importText); + } + if (variant === "go") { + return {}; + } + return parseTsImportBindings(importText); +} + +function parseTsImportBindings(importText: string): ImportBindingResult { + const namedBindings: string[] = []; + let defaultBinding: string | undefined; + let namespaceBinding: string | undefined; + + const brace = /\{([^}]+)\}/u.exec(importText); + if (brace) { + for (const part of brace[1].split(",")) { + const name = part.trim().split(/\s+as\s+/u)[0]?.trim(); + if (name && name !== "type") namedBindings.push(name); + } + } + + const ns = /\*\s+as\s+([A-Za-z_$][\w$]*)/u.exec(importText); + if (ns) namespaceBinding = ns[1]; + + const dm1 = /^import\s+([A-Za-z_$][\w$]*)\s*,/u.exec(importText); + const dm2 = /^import\s+([A-Za-z_$][\w$]*)\s+from/u.exec(importText); + const defaultMatch = dm1 ?? dm2; + if (defaultMatch && !importText.includes("{")) { + defaultBinding = defaultMatch[1]; + } + + return { namedBindings: namedBindings.length > 0 ? namedBindings : undefined, defaultBinding, namespaceBinding }; +} + +function parsePythonImportBindings(importText: string): ImportBindingResult { + const namedBindings: string[] = []; + let defaultBinding: string | undefined; + let namespaceBinding: string | undefined; + + const fromImport = /^from\s+[\w.]+\s+import\s+(.+)$/u.exec(importText.trim()); + if (fromImport) { + const tail = fromImport[1]; + if (tail.startsWith("(") && tail.endsWith(")")) { + for (const part of tail.slice(1, -1).split(",")) { + const token = part.trim().split(/\s+as\s+/u)[0]?.trim(); + if (token && token !== "*") namedBindings.push(token); + } + } else if (tail === "*") { + namespaceBinding = "*"; + } else { + for (const part of tail.split(",")) { + const token = part.trim().split(/\s+as\s+/u)[0]?.trim(); + if (!token) continue; + if (!defaultBinding) defaultBinding = token; + else namedBindings.push(token); + } + } + } + + const plain = /^import\s+([\w.]+)(?:\s+as\s+(\w+))?/u.exec(importText.trim()); + if (plain) { + if (plain[2]) namespaceBinding = plain[2]; + else defaultBinding = plain[1].split(".").pop(); + } + + return { namedBindings: namedBindings.length > 0 ? namedBindings : undefined, defaultBinding, namespaceBinding }; +} + +/** Normalize import specifier from tree-sitter capture text. */ +export function normalizeImportSpecifier(specText: string, variant: GrammarVariant): string { + if (variant === "go") { + return specText.replace(/^["`]|["`]$/gu, ""); + } + return specText.replace(/^['"]|['"]$/gu, ""); +} + +export function isTypeOnlyImport(importText: string, variant: GrammarVariant): boolean { + if (variant === "typescript" || variant === "tsx") { + return /^\s*import\s+type\b/u.test(importText); + } + return false; +} + +export function isExportedSymbol( + variant: GrammarVariant, + startIndex: number, + source: string, + lineStart: number, + exportLineStarts: Set +): boolean { + if (variant === "typescript" || variant === "tsx") { + return exportLineStarts.has(lineStart) || isTsExportedDeclaration(startIndex, source); + } + if (variant === "python") { + return exportLineStarts.has(lineStart); + } + if (variant === "go") { + const decl = source.slice(Math.max(0, startIndex - 20), startIndex); + return /^func\s+[A-Z]/u.test(decl.trimStart()) || /^type\s+[A-Z]/u.test(decl.trimStart()); + } + return false; +} + +function isTsExportedDeclaration(startIndex: number, source: string): boolean { + const prefix = source.slice(Math.max(0, startIndex - 80), startIndex); + return /\bexport\s+(?:default\s+)?$/u.test(prefix.trimEnd()); +} + +export function collectExportLineStarts( + variant: GrammarVariant, + root: Node +): Set { + const lines = new Set(); + if (variant === "typescript" || variant === "tsx") { + for (const node of root.descendantsOfType("export_statement")) { + if (node) lines.add(node.startPosition.row + 1); + } + } + if (variant === "python") { + // Python has no export keyword: treat module-level class/function + // definitions as importable (tree-sitter-python has no "__export__" node). + for (const node of root.descendantsOfType(["class_definition", "function_definition"])) { + if (node && node.parent?.type === "module") { + lines.add(node.startPosition.row + 1); + } + } + } + return lines; +} diff --git a/src/wiki-engine/code-knowledge/ast/import-resolver.ts b/src/wiki-engine/code-knowledge/ast/import-resolver.ts new file mode 100644 index 00000000..9c938063 --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/import-resolver.ts @@ -0,0 +1,273 @@ +import { readFile, stat } from "node:fs/promises"; +import path from "node:path"; + +import { toPosix } from "../../core/wiki-protocol.js"; + +const EXTENSIONS_TS = [".ts", ".tsx", ".js", ".jsx", ".mjs", ".cjs"]; +const INDEX_SUFFIXES_TS = ["/index.ts", "/index.tsx", "/index.js", "/index.jsx"]; +const EXTENSIONS_PY = [".py"]; +const INDEX_SUFFIXES_PY = ["/__init__.py"]; +const EXTENSIONS_GO = [".go"]; +/** Matches a dotted module identifier (e.g. `pkg.sub.mod`), used by Python/Go import resolution. */ +const MODULE_IDENTIFIER_RE = /^[A-Za-z_][\w.]*$/u; + +export interface TsconfigPathsConfig { + baseUrl: string; + paths: Record; +} + +export interface ResolvedImport { + targetFile: string; + confidence: "EXTRACTED" | "AMBIGUOUS"; + candidates?: string[]; +} + +const tsconfigCache = new Map(); + +/** Find nearest tsconfig.json upward from a file directory. */ +export async function loadTsconfigPaths( + repoRoot: string, + fromRelativeFile: string +): Promise { + const absDir = path.dirname(path.resolve(repoRoot, fromRelativeFile)); + let dir = absDir; + const root = path.resolve(repoRoot); + + while (dir === root || dir.startsWith(root + path.sep)) { + if (tsconfigCache.has(dir)) { + return tsconfigCache.get(dir) ?? null; + } + + const configPath = path.join(dir, "tsconfig.json"); + try { + const raw = JSON.parse(await readFile(configPath, "utf8")) as { + compilerOptions?: { baseUrl?: string; paths?: Record }; + }; + const configDir = path.dirname(configPath); + const baseUrl = raw.compilerOptions?.baseUrl ?? "."; + const pathsMap = raw.compilerOptions?.paths ?? {}; + const resolved: TsconfigPathsConfig = { + baseUrl: toPosix(path.resolve(configDir, baseUrl)), + paths: Object.fromEntries( + Object.entries(pathsMap).map(([key, values]) => [ + key, + values.map((v) => { + const withoutStar = v.replace(/\*$/u, ""); + return toPosix(path.resolve(configDir, baseUrl, withoutStar)); + }) + ]) + ) + }; + tsconfigCache.set(dir, resolved); + return resolved; + } catch { + // try parent + } + + if (dir === root) break; + const parent = path.dirname(dir); + if (parent === dir) break; + dir = parent; + } + + tsconfigCache.set(absDir, null); + return null; +} + + +export async function resolveImportSpecifier( + repoRoot: string, + fromFile: string, + specifier: string, + fileExists: (relativePath: string) => Promise +): Promise { + if (specifier.startsWith(".") || specifier.startsWith("/")) { + return resolveRelativeImport(fromFile, specifier, fileExists); + } + + if (fromFile.endsWith(".py") || fromFile.endsWith(".go")) { + const sibling = await resolveSiblingModuleImport(fromFile, specifier, fileExists); + if (sibling) return sibling; + if (fromFile.endsWith(".py")) { + const absolutePkg = await resolveAbsolutePackageImport(fromFile, specifier, fileExists); + if (absolutePkg) return absolutePkg; + } + } + + const tsconfig = await loadTsconfigPaths(repoRoot, fromFile); + if (tsconfig) { + const mapped = await resolvePathsMapping(repoRoot, tsconfig, specifier, fromFile, fileExists); + if (mapped) return mapped; + } + + return undefined; +} + +async function resolveRelativeImport( + fromFile: string, + specifier: string, + fileExists: (relativePath: string) => Promise +): Promise { + const fromDir = path.dirname(fromFile); + const base = toPosix(path.normalize(path.join(fromDir, specifier))); + const candidates = await expandModulePaths(base, fromFile, fileExists); + return pickCandidate(candidates); +} + +/** Resolve `from b import` (Python) or local package dir imports (Go) via sibling path. */ +async function resolveSiblingModuleImport( + fromFile: string, + specifier: string, + fileExists: (relativePath: string) => Promise +): Promise { + if (!MODULE_IDENTIFIER_RE.test(specifier)) { + return undefined; + } + const fromDir = path.dirname(fromFile); + // Python dotted modules (a.b.c) map to nested paths (a/b/c); other langs keep the specifier as-is. + const modulePath = fromFile.endsWith(".py") ? specifier.replace(/\./gu, "/") : specifier; + const base = toPosix(path.normalize(path.join(fromDir, modulePath))); + const candidates = await expandModulePaths(base, fromFile, fileExists); + return pickCandidate(candidates); +} + +/** + * Resolve a Python absolute package import from the repo root. + * + * Python code commonly imports via the top-level package name rooted at the + * repository root (e.g. `from hai_flow.conf import config`), not relative to the + * importing file's directory. This maps the dotted specifier directly onto a + * repo-root-relative path (a.b.c -> a/b/c) and probes for a module file or + * package `__init__.py`. Only applies to Python source files. + * + * @param fromFile Repo-root-relative path of the importing file. + * @param specifier The (non-relative) import specifier, e.g. "hai_flow.conf". + * @param fileExists Predicate checking a repo-root-relative path exists. + * @returns Resolved import, or undefined when no repo-root path matches. + */ +async function resolveAbsolutePackageImport( + fromFile: string, + specifier: string, + fileExists: (relativePath: string) => Promise +): Promise { + if (!fromFile.endsWith(".py")) { + return undefined; + } + if (!MODULE_IDENTIFIER_RE.test(specifier)) { + return undefined; + } + const modulePath = specifier.replace(/\./gu, "/"); + const candidates = await expandModulePaths(modulePath, fromFile, fileExists); + return pickCandidate(candidates); +} + +async function resolvePathsMapping( + repoRoot: string, + config: TsconfigPathsConfig, + specifier: string, + fromFile: string, + fileExists: (relativePath: string) => Promise +): Promise { + const root = path.resolve(repoRoot); + const entries = Object.entries(config.paths).sort((a, b) => b[0].length - a[0].length); + + for (const [pattern, targets] of entries) { + if (pattern.endsWith("/*")) { + const prefix = pattern.slice(0, -2); + if (!specifier.startsWith(`${prefix}/`)) continue; + const rest = specifier.slice(prefix.length + 1); + for (const targetBase of targets) { + const abs = toPosix(path.join(targetBase, rest)); + const rel = toPosix(path.relative(root, abs)); + const candidates = await expandModulePaths(rel, fromFile, fileExists); + const picked = pickCandidate(candidates); + if (picked) return picked; + } + } else if (pattern === specifier) { + for (const target of targets) { + const rel = toPosix(path.relative(root, target)); + const candidates = await expandModulePaths(rel, fromFile, fileExists); + const picked = pickCandidate(candidates); + if (picked) return picked; + } + } + } + + return undefined; +} + +function moduleExtensions(fromFile: string): { extensions: string[]; indexSuffixes: string[] } { + if (fromFile.endsWith(".py")) { + return { extensions: EXTENSIONS_PY, indexSuffixes: INDEX_SUFFIXES_PY }; + } + if (fromFile.endsWith(".go")) { + return { extensions: EXTENSIONS_GO, indexSuffixes: [] }; + } + return { extensions: EXTENSIONS_TS, indexSuffixes: INDEX_SUFFIXES_TS }; +} + +/** Go: import "./pkg" often maps to pkg/pkg.go when the directory name matches the last segment. */ +async function expandGoPackageDir( + normalized: string, + fileExists: (relativePath: string) => Promise +): Promise { + const found: string[] = []; + const seg = normalized.split("/").pop(); + if (seg) { + const nested = `${normalized}/${seg}.go`; + if (await fileExists(nested)) found.push(nested); + } + return found; +} + +async function expandModulePaths( + baseRelative: string, + fromFile: string, + fileExists: (relativePath: string) => Promise +): Promise { + const found: string[] = []; + const normalized = baseRelative.replace(/^\.\//u, ""); + const { extensions, indexSuffixes } = moduleExtensions(fromFile); + + for (const ext of extensions) { + const candidate = `${normalized}${ext}`; + if (await fileExists(candidate)) found.push(candidate); + } + + for (const indexSuffix of indexSuffixes) { + const candidate = `${normalized}${indexSuffix}`; + if (await fileExists(candidate)) found.push(candidate); + } + + if (fromFile.endsWith(".go")) { + found.push(...(await expandGoPackageDir(normalized, fileExists))); + } + + if (await fileExists(normalized)) found.push(normalized); + + return [...new Set(found)]; +} + +function pickCandidate(candidates: string[]): ResolvedImport | undefined { + if (candidates.length === 0) return undefined; + if (candidates.length === 1) { + return { targetFile: candidates[0]!, confidence: "EXTRACTED" }; + } + return { targetFile: candidates[0]!, confidence: "AMBIGUOUS", candidates }; +} + +export async function buildFileExistenceChecker( + repoRoot: string, + knownFiles: Set +): Promise<(relativePath: string) => Promise> { + return async (relativePath: string) => { + if (knownFiles.has(relativePath)) return true; + try { + const abs = path.resolve(repoRoot, relativePath); + const s = await stat(abs); + return s.isFile(); + } catch { + return false; + } + }; +} diff --git a/src/wiki-engine/code-knowledge/ast/index.ts b/src/wiki-engine/code-knowledge/ast/index.ts new file mode 100644 index 00000000..c2bf0120 --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/index.ts @@ -0,0 +1,217 @@ +import { createRequire } from "node:module"; + +import type { CodeCollectedFile } from "../code-collector.js"; +import { type CodeFact } from "../code-extractors.js"; +import { structuralEdgesToCodeFacts, unresolvedImportsToGaps } from "./adapt-code-facts.js"; +import { buildImportBindingsForFile, callResolutionWeight, resolveCallSites } from "./call-resolver.js"; +import { buildFileExistenceChecker, resolveImportSpecifier } from "./import-resolver.js"; +import { ensureAstReady } from "./parser-registry.js"; +import type { AstExtractionGap, AstImplementsSite, StructuralEdge, StructuralGraphResult } from "./types.js"; +import { isAstParseableFile, walkFile } from "./walk.js"; + +const require = createRequire(import.meta.url); + +let astAvailability: boolean | undefined; + +/** Whether the WASM tree-sitter runtime and grammars can be loaded (honors TEAMAI_SKIP_AST). */ +export function astAvailable(): boolean { + if (process.env.TEAMAI_SKIP_AST === "1") { + return false; + } + if (astAvailability !== undefined) { + return astAvailability; + } + try { + require.resolve("web-tree-sitter"); + require.resolve("web-tree-sitter/tree-sitter.wasm"); + require.resolve("tree-sitter-wasms/out/tree-sitter-typescript.wasm"); + astAvailability = true; + } catch { + astAvailability = false; + } + return astAvailability; +} + +export interface ExtractStructuralGraphOptions { + repoRoot: string; + files: CodeCollectedFile[]; +} + +export async function extractStructuralGraph( + options: ExtractStructuralGraphOptions +): Promise { + await ensureAstReady(); + + const { repoRoot, files } = options; + const symbols: StructuralGraphResult["symbols"] = []; + const imports: StructuralGraphResult["imports"] = []; + const callSites: StructuralGraphResult["callSites"] = []; + const implementsSites: AstImplementsSite[] = []; + const gaps: AstExtractionGap[] = []; + const edges: StructuralEdge[] = []; + + const knownFiles = new Set(files.map((f) => f.relativePath)); + const fileExists = await buildFileExistenceChecker(repoRoot, knownFiles); + + let filesParsed = 0; + let filesSkipped = 0; + + for (const file of files) { + if (!isAstParseableFile(file.relativePath)) { + filesSkipped++; + continue; + } + const walked = walkFile(file); + if (walked.parseErrors.length > 0) { + for (const err of walked.parseErrors) { + gaps.push({ kind: "PARSE_SKIP", message: err, sources: [file.relativePath] }); + } + if (walked.symbols.length === 0 && walked.imports.length === 0) { + filesSkipped++; + continue; + } + } + filesParsed++; + symbols.push(...walked.symbols); + imports.push(...walked.imports); + callSites.push(...walked.callSites); + implementsSites.push(...walked.implementsSites); + } + + const symbolsByFile = new Map(); + for (const sym of symbols) { + const list = symbolsByFile.get(sym.file) ?? []; + list.push(sym); + symbolsByFile.set(sym.file, list); + } + + const resolvedImports = new Map>>(); + const resolvedKeys = new Set(); + + for (const imp of imports) { + const key = `${imp.fromFile}:${imp.line}`; + if (imp.isTypeOnly) continue; + + const resolved = await resolveImportSpecifier(repoRoot, imp.fromFile, imp.specifier, fileExists); + resolvedImports.set(key, resolved); + if (resolved) { + resolvedKeys.add(key); + edges.push({ + from: imp.fromFile, + to: resolved.targetFile, + relation: "DEPENDS_ON", + source: "code-ast", + weight: resolved.confidence === "EXTRACTED" ? 0.9 : 0.5, + confidence: resolved.confidence, + evidence: [ + { + ref: imp.fromFile, + lineStart: imp.line, + lineEnd: imp.line, + note: `resolved import ${imp.specifier}` + } + ] + }); + } + } + + gaps.push(...unresolvedImportsToGaps(imports.filter((i) => !i.isTypeOnly), resolvedKeys)); + + const resolvedCalls = resolveCallSites(callSites, imports, resolvedImports, symbolsByFile); + + for (const call of resolvedCalls) { + if (!call.resolvedTargetFile || call.resolvedTargetFile === call.fromFile) { + continue; + } + + edges.push({ + from: call.fromFile, + to: call.resolvedTargetFile, + relation: "REFERENCES", + source: "code-ast", + weight: callResolutionWeight(call.confidence), + confidence: call.confidence, + evidence: [ + { + ref: call.fromFile, + lineStart: call.line, + lineEnd: call.line, + note: `call ${call.calleeText}` + } + ] + }); + } + + // IMPLEMENTS edges: resolve each implemented interface name to its defining + // file via (a) same-file interface symbols or (b) imported bindings. Names + // that resolve to neither (e.g. ambient/global types) are skipped. + for (const site of implementsSites) { + const bindings = buildImportBindingsForFile(site.fromFile, imports, resolvedImports, symbolsByFile); + const localInterfaces = symbolsByFile.get(site.fromFile) ?? []; + for (const ifaceName of site.ifaceNames) { + let targetFile: string | undefined; + const sameFile = localInterfaces.find((s) => s.name === ifaceName && s.kind === "interface"); + if (sameFile) { + targetFile = site.fromFile; + } else { + targetFile = bindings.localToFile.get(ifaceName); + } + if (!targetFile) { + continue; + } + edges.push({ + from: site.fromFile, + to: targetFile, + relation: "IMPLEMENTS", + source: "code-ast", + weight: 0.9, + confidence: "EXTRACTED", + evidence: [ + { + ref: site.fromFile, + lineStart: site.line, + lineEnd: site.line, + note: `${site.className} implements ${ifaceName}` + } + ] + }); + } + } + + const stats = { + symbols: symbols.length, + imports: imports.length, + importsResolved: resolvedKeys.size, + calls: callSites.length, + callsResolved: resolvedCalls.filter((c) => c.resolvedTargetFile).length, + edges: edges.length, + filesParsed, + filesSkipped + }; + + return { + symbols, + imports, + callSites: resolvedCalls, + edges, + gaps, + stats + }; +} + +export async function extractStructuralGraphAsFacts( + options: ExtractStructuralGraphOptions +): Promise<{ facts: CodeFact[]; result: StructuralGraphResult }> { + const result = await extractStructuralGraph(options); + const facts = structuralEdgesToCodeFacts(result.edges); + return { facts, result }; +} + +export function formatAstStatsSummary(stats: StructuralGraphResult["stats"]): string { + const imports = `${stats.imports} imports (${stats.importsResolved} resolved)`; + const calls = `${stats.calls} calls (${stats.callsResolved} resolved)`; + return `ast: ${stats.symbols} symbols, ${imports}, ${calls}, ${stats.edges} edges`; +} + +export * from "./types.js"; +export { mergeCodeFacts } from "./merge-edges.js"; diff --git a/src/wiki-engine/code-knowledge/ast/merge-edges.ts b/src/wiki-engine/code-knowledge/ast/merge-edges.ts new file mode 100644 index 00000000..4adf30d2 --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/merge-edges.ts @@ -0,0 +1,35 @@ +import type { CodeFact } from "../code-extractors.js"; + +function factKey(fact: CodeFact): string { + return `${fact.kind}:${fact.name}:${fact.file}:${fact.lineStart}`; +} + +/** + * Merge AST-derived facts with heuristic facts. AST wins on duplicate keys; + * heuristic relation facts at the same file:line as an AST relation are dropped. + */ +export function mergeCodeFacts(astFacts: CodeFact[], heuristicFacts: CodeFact[]): CodeFact[] { + const astRelationLines = new Set( + astFacts.filter((f) => f.kind === "relation").map((f) => `${f.file}:${f.lineStart}`) + ); + + const astKeys = new Set(astFacts.map(factKey)); + const filteredHeuristic = heuristicFacts.filter((fact) => { + if (fact.kind === "relation" && astRelationLines.has(`${fact.file}:${fact.lineStart}`)) { + return false; + } + return !astKeys.has(factKey(fact)); + }); + + const merged = [...astFacts, ...filteredHeuristic]; + const seen = new Set(); + const result: CodeFact[] = []; + for (const fact of merged) { + const key = factKey(fact); + if (seen.has(key)) continue; + seen.add(key); + result.push(fact); + } + return result; +} + diff --git a/src/wiki-engine/code-knowledge/ast/parser-registry.ts b/src/wiki-engine/code-knowledge/ast/parser-registry.ts new file mode 100644 index 00000000..ea635d5c --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/parser-registry.ts @@ -0,0 +1,106 @@ +import { createRequire } from "node:module"; + +import { Language, Parser, Query } from "web-tree-sitter"; + +import { GO_AST_QUERY_SOURCE, PYTHON_AST_QUERY_SOURCE, TS_AST_QUERY_SOURCE } from "./queries.js"; + +export type GrammarVariant = "typescript" | "tsx" | "python" | "go"; + +const require = createRequire(import.meta.url); + +const GRAMMAR_WASM: Record = { + typescript: "tree-sitter-wasms/out/tree-sitter-typescript.wasm", + tsx: "tree-sitter-wasms/out/tree-sitter-tsx.wasm", + python: "tree-sitter-wasms/out/tree-sitter-python.wasm", + go: "tree-sitter-wasms/out/tree-sitter-go.wasm" +}; + +const QUERY_SOURCE: Record = { + typescript: TS_AST_QUERY_SOURCE, + tsx: TS_AST_QUERY_SOURCE, + python: PYTHON_AST_QUERY_SOURCE, + go: GO_AST_QUERY_SOURCE +}; + +let parserInstance: Parser | undefined; +let initPromise: Promise | undefined; +const languages = new Map(); +const queries = new Map(); + +/** + * Initialize the WASM runtime, parser, grammars and queries once. + * + * Idempotent: concurrent and repeat callers await the same in-flight promise. + * Must be awaited before any synchronous getParser/getLanguage/getQuery call. + */ +export async function ensureAstReady(): Promise { + if (!initPromise) { + initPromise = initAst(); + } + return initPromise; +} + +async function initAst(): Promise { + await Parser.init({ + locateFile: () => require.resolve("web-tree-sitter/tree-sitter.wasm") + }); + parserInstance = new Parser(); + for (const variant of Object.keys(GRAMMAR_WASM) as GrammarVariant[]) { + const language = await Language.load(require.resolve(GRAMMAR_WASM[variant])); + languages.set(variant, language); + queries.set(variant, new Query(language, QUERY_SOURCE[variant])); + } +} + +/** Lazy singleton parser. ensureAstReady() must have resolved first. */ +export function getParser(): Parser { + if (!parserInstance) { + throw new Error("AST parser not initialized; call ensureAstReady() first"); + } + return parserInstance; +} + +export function grammarForExtension(ext: string): GrammarVariant | undefined { + switch (ext.toLowerCase()) { + case ".ts": + case ".mts": + case ".cts": + case ".js": + case ".jsx": + case ".mjs": + case ".cjs": + return "typescript"; + case ".tsx": + return "tsx"; + case ".py": + case ".pyi": + return "python"; + case ".go": + return "go"; + default: + return undefined; + } +} + +export function getLanguage(variant: GrammarVariant): Language { + const language = languages.get(variant); + if (!language) { + throw new Error(`No tree-sitter grammar loaded for variant: ${variant}`); + } + return language; +} + +export function getQuery(variant: GrammarVariant): Query { + const query = queries.get(variant); + if (!query) { + throw new Error(`No tree-sitter query registered for variant: ${variant}`); + } + return query; +} + +export function resetParserRegistryForTests(): void { + parserInstance = undefined; + initPromise = undefined; + languages.clear(); + queries.clear(); +} diff --git a/src/wiki-engine/code-knowledge/ast/queries.ts b/src/wiki-engine/code-knowledge/ast/queries.ts new file mode 100644 index 00000000..fe82aa88 --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/queries.ts @@ -0,0 +1,103 @@ +/** Tree-sitter query sources per grammar variant. */ + +export const TS_AST_QUERY_SOURCE = ` +(import_statement + source: (string (string_fragment) @import.spec) +) @import.stmt + +(export_statement) @export.stmt + +(class_declaration + name: (type_identifier) @symbol.name +) @symbol.class + +(function_declaration + name: (identifier) @symbol.name +) @symbol.function + +(interface_declaration + name: (type_identifier) @symbol.name +) @symbol.interface + +(call_expression + function: (identifier) @call.callee +) @call.stmt + +(call_expression + function: (member_expression + object: (identifier) @call.receiver + property: (property_identifier) @call.member + ) +) @call.member + +(class_declaration + name: (type_identifier) @impl.class + (class_heritage + (implements_clause + (type_identifier) @impl.iface + ) + ) +) @impl.stmt +`; + +export const PYTHON_AST_QUERY_SOURCE = ` +(import_from_statement + module_name: (dotted_name) @import.spec +) @import.stmt + +(import_statement + name: (dotted_name) @import.spec +) @import.stmt + +(class_definition + name: (identifier) @symbol.name +) @symbol.class + +(function_definition + name: (identifier) @symbol.name +) @symbol.function + +(call + function: (identifier) @call.callee +) @call.stmt + +(call + function: (attribute + object: (identifier) @call.receiver + attribute: (identifier) @call.member + ) +) @call.member +`; + +export const GO_AST_QUERY_SOURCE = ` +(import_declaration + (import_spec + path: (interpreted_string_literal) @import.spec + ) +) @import.stmt + +(function_declaration + name: (identifier) @symbol.name +) @symbol.function + +(method_declaration + name: (field_identifier) @symbol.name +) @symbol.function + +(type_declaration + (type_spec + name: (type_identifier) @symbol.name + ) +) @symbol.class + +(call_expression + function: (identifier) @call.callee +) @call.stmt + +(call_expression + function: (selector_expression + operand: (identifier) @call.receiver + field: (field_identifier) @call.member + ) +) @call.member +`; diff --git a/src/wiki-engine/code-knowledge/ast/types.ts b/src/wiki-engine/code-knowledge/ast/types.ts new file mode 100644 index 00000000..070e846f --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/types.ts @@ -0,0 +1,79 @@ +import type { ManifestConfidence } from "../../manifest-schema.js"; +import type { WikiEvidence } from "../../core/wiki-protocol.js"; + +export type AstSymbolKind = "function" | "class" | "interface" | "method" | "variable"; + +export interface AstSymbol { + id: string; + kind: AstSymbolKind; + name: string; + file: string; + lineStart: number; + lineEnd: number; + exported: boolean; +} + +export interface AstImport { + fromFile: string; + specifier: string; + line: number; + isTypeOnly: boolean; + namedBindings?: string[]; + defaultBinding?: string; + namespaceBinding?: string; +} + +export interface AstCallSite { + fromFile: string; + line: number; + calleeText: string; + receiver?: string; + resolvedTargetId?: string; + resolvedTargetFile?: string; + confidence: ManifestConfidence; +} + +export interface AstImplementsSite { + fromFile: string; + className: string; + ifaceNames: string[]; + line: number; +} + +export type StructuralRelation = "DEPENDS_ON" | "REFERENCES" | "IMPLEMENTS"; + +export interface StructuralEdge { + from: string; + to: string; + relation: StructuralRelation; + source: "code-ast"; + weight: number; + evidence: WikiEvidence[]; + confidence: ManifestConfidence; +} + +export interface AstExtractionGap { + kind: string; + message: string; + sources: string[]; +} + +export interface AstExtractionStats { + symbols: number; + imports: number; + importsResolved: number; + calls: number; + callsResolved: number; + edges: number; + filesParsed: number; + filesSkipped: number; +} + +export interface StructuralGraphResult { + symbols: AstSymbol[]; + imports: AstImport[]; + callSites: AstCallSite[]; + edges: StructuralEdge[]; + gaps: AstExtractionGap[]; + stats: AstExtractionStats; +} diff --git a/src/wiki-engine/code-knowledge/ast/walk.ts b/src/wiki-engine/code-knowledge/ast/walk.ts new file mode 100644 index 00000000..c71b5ae5 --- /dev/null +++ b/src/wiki-engine/code-knowledge/ast/walk.ts @@ -0,0 +1,154 @@ +import path from "node:path"; + +import type { CodeCollectedFile } from "../code-collector.js"; +import { + collectExportLineStarts, + isExportedSymbol, + isTypeOnlyImport, + normalizeImportSpecifier, + parseImportBindings +} from "./import-bindings.js"; +import { grammarForExtension, getLanguage, getParser, getQuery } from "./parser-registry.js"; +import type { AstCallSite, AstImplementsSite, AstImport, AstSymbol, AstSymbolKind } from "./types.js"; + +export interface FileWalkResult { + symbols: AstSymbol[]; + imports: AstImport[]; + callSites: AstCallSite[]; + implementsSites: AstImplementsSite[]; + parseErrors: string[]; +} + +const MAX_FILE_BYTES = 512 * 1024; + +export function isAstParseableFile(relativePath: string): boolean { + return grammarForExtension(path.extname(relativePath)) !== undefined; +} + +export function walkFile(file: CodeCollectedFile): FileWalkResult { + const symbols: AstSymbol[] = []; + const imports: AstImport[] = []; + const callSites: AstCallSite[] = []; + const implementsSites: AstImplementsSite[] = []; + const parseErrors: string[] = []; + + if (!isAstParseableFile(file.relativePath)) { + return { symbols, imports, callSites, implementsSites, parseErrors }; + } + + if (Buffer.byteLength(file.content, "utf8") > MAX_FILE_BYTES) { + parseErrors.push(`skipped large file: ${file.relativePath}`); + return { symbols, imports, callSites, implementsSites, parseErrors }; + } + + const variant = grammarForExtension(path.extname(file.relativePath))!; + const language = getLanguage(variant); + const parser = getParser(); + parser.setLanguage(language); + + let tree; + try { + tree = parser.parse(file.content); + } catch (error) { + parseErrors.push(`parse failed: ${file.relativePath}: ${error instanceof Error ? error.message : String(error)}`); + return { symbols, imports, callSites, implementsSites, parseErrors }; + } + + if (!tree) { + parseErrors.push(`parse returned null: ${file.relativePath}`); + return { symbols, imports, callSites, implementsSites, parseErrors }; + } + + try { + const query = getQuery(variant); + const exportLineStarts = collectExportLineStarts(variant, tree.rootNode); + + for (const match of query.matches(tree.rootNode)) { + const byName = new Map(match.captures.map((c) => [c.name, c.node])); + + if (byName.has("import.stmt")) { + const stmt = byName.get("import.stmt")!; + const specNode = byName.get("import.spec"); + if (!specNode) continue; + const specifier = normalizeImportSpecifier(specNode.text, variant); + const line = stmt.startPosition.row + 1; + const isTypeOnly = isTypeOnlyImport(stmt.text, variant); + imports.push({ + fromFile: file.relativePath, + specifier, + line, + isTypeOnly, + ...parseImportBindings(stmt.text, variant) + }); + continue; + } + + const symbolName = byName.get("symbol.name")?.text; + if (symbolName) { + const decl = + byName.get("symbol.class") ?? byName.get("symbol.function") ?? byName.get("symbol.interface"); + if (!decl) continue; + const kind: AstSymbolKind = byName.has("symbol.class") + ? "class" + : byName.has("symbol.interface") + ? "interface" + : "function"; + const lineStart = decl.startPosition.row + 1; + const lineEnd = decl.endPosition.row + 1; + const exported = isExportedSymbol(variant, decl.startIndex, file.content, lineStart, exportLineStarts); + symbols.push({ + id: symbolId(file.relativePath, kind, symbolName), + kind, + name: symbolName, + file: file.relativePath, + lineStart, + lineEnd, + exported + }); + continue; + } + + if (byName.has("call.stmt") || byName.has("call.member")) { + const callNode = byName.get("call.stmt") ?? byName.get("call.member")!; + const line = callNode.startPosition.row + 1; + const callee = byName.get("call.callee")?.text; + const receiver = byName.get("call.receiver")?.text; + const member = byName.get("call.member")?.text; + const calleeText = callee ?? (receiver && member ? `${receiver}.${member}` : callNode.text); + callSites.push({ + fromFile: file.relativePath, + line, + calleeText, + receiver, + confidence: "INFERRED" + }); + continue; + } + + if (byName.has("impl.stmt")) { + const classNode = byName.get("impl.class"); + const ifaceNames = match.captures + .filter((c) => c.name === "impl.iface") + .map((c) => c.node.text); + if (classNode && ifaceNames.length > 0) { + implementsSites.push({ + fromFile: file.relativePath, + className: classNode.text, + ifaceNames, + line: classNode.startPosition.row + 1 + }); + } + continue; + } + } + } finally { + tree.delete(); + } + + return { symbols, imports, callSites, implementsSites, parseErrors }; +} + +function symbolId(file: string, kind: AstSymbolKind, name: string): string { + const kindLabel = kind.charAt(0).toUpperCase() + kind.slice(1); + return `${file}#${kindLabel}:${name}`; +} diff --git a/src/wiki-engine/code-knowledge/code-graph.ts b/src/wiki-engine/code-knowledge/code-graph.ts index 4983b099..f91b7585 100644 --- a/src/wiki-engine/code-knowledge/code-graph.ts +++ b/src/wiki-engine/code-knowledge/code-graph.ts @@ -5,6 +5,7 @@ import { type GraphIndex, type GraphNode, type GraphEdge, + type RelationType, createGraphIndex, } from "../core/graph-index.schema.js"; @@ -28,6 +29,13 @@ export function buildCodeGraph(facts: CodeFact[]): GraphIndex { const edges: GraphEdge[] = facts .filter((fact) => fact.kind === "relation") .flatMap((fact) => { + // AST-derived relation facts carry a resolved target file in fact.name and a + // "(code-ast)" marker in detail — trust them directly instead of fuzzy matching. + const astEdge = parseAstRelationFact(fact); + if (astEdge) { + return [astEdge]; + } + // Heuristic relation facts: fuzzy path-substring match against known node files. const targets = [...nodeFiles].filter((file) => relationMayTarget(fact.name, file)); return targets.map((file) => ({ from: fact.file, @@ -41,6 +49,31 @@ export function buildCodeGraph(facts: CodeFact[]): GraphIndex { return createGraphIndex(nodes, edges); } +/** + * Parse an AST-derived relation fact into a precise graph edge. + * + * AST facts encode their edge as `detail = " → (code-ast)"` + * with `fact.name` already holding the resolved target file. Returns undefined + * for heuristic (regex) relation facts, which are resolved by fuzzy matching. + */ +function parseAstRelationFact(fact: CodeFact): GraphEdge | undefined { + if (!fact.detail.includes("(code-ast)")) { + return undefined; + } + const relation: RelationType = fact.detail.startsWith("REFERENCES") + ? "REFERENCES" + : fact.detail.startsWith("IMPLEMENTS") + ? "IMPLEMENTS" + : "DEPENDS_ON"; + return { + from: fact.file, + to: fact.name, + relation, + weight: fact.confidence === "EXTRACTED" ? 0.9 : 0.5, + source: "code-ast", + }; +} + function relationMayTarget(importTarget: string, file: string): boolean { const normalized = importTarget.replace(/^\.\//u, "").replace(/\.\.\//g, "").replace(/\.(ts|tsx|js|jsx)$/u, ""); if (normalized.length < 3) return false; // Skip very short matches to reduce false positives