Skip to content

feat(wiki-engine): add WASM tree-sitter AST track for code knowledge graph - #304

Open
m0Nst3r873 wants to merge 9 commits into
Tencent:mainfrom
m0Nst3r873:feature/ast-code-knowledge
Open

feat(wiki-engine): add WASM tree-sitter AST track for code knowledge graph#304
m0Nst3r873 wants to merge 9 commits into
Tencent:mainfrom
m0Nst3r873:feature/ast-code-knowledge

Conversation

@m0Nst3r873

Copy link
Copy Markdown
Collaborator

What & Why

The code knowledge graph extractors (src/wiki-engine/code-knowledge/) were purely regex / line-based. That has two concrete problems:

  1. False-positive edges. Dependency edges were built by path-substring matching (file.includes(importPath)), so an import "./user" wrongly links to both user.ts and user-repo.ts (any path containing user).
  2. No call relationships. Only DEPENDS_ON (imports) existed; call sites were invisible.

This PR adds a real AST track using web-tree-sitter (pure-WASM, no native toolchain), ported from the team-wiki reference implementation, for TypeScript/JavaScript, Python, and Go. It runs alongside the existing regex heuristic track (which still covers Java/Rust/config), and AST results win on merge.

How it works

  • New ast/ module: WASM parser registry (async one-time init), tree-sitter queries, symbol/import/call-site walk, import & call resolvers, and fact/edge adapters. Emits precise file-to-file DEPENDS_ON / REFERENCES edges tagged source: "code-ast" with confidence weights.
  • Dual-track with graceful fallback: if the WASM runtime can't load, or TEAMAI_SKIP_AST=1 is set, extraction falls back to heuristic-only and records an AST_UNAVAILABLE gap.
  • code-graph: AST relation facts build precise edges directly, instead of being re-fuzzed through the path-substring matcher.
  • enrich: manifest edges now preserve real AST relation/source provenance (deterministic rank-based merge) instead of hardcoding DEPENDS_ON / code-heuristic. Combined with upstream's resolveImportToModule resolver.
  • Pinned versions: web-tree-sitter@0.25.10 + tree-sitter-wasms@0.1.13 — the only ABI-14-compatible pair (0.26 rejects these grammars). Both are pure-JS deps resolved from node_modules at runtime; no .wasm files are bundled into dist/.

Accuracy (measured)

On a fixture with a naming-collision decoy (user.ts real target, user-repo.ts decoy that nobody imports):

Mode Edges False positives Precision
Regex only (TEAMAI_SKIP_AST=1) 3 1 (→ user-repo.ts) 67%
AST enabled 2 0 100%

Test Plan

  • npx tsc --noEmit — clean
  • npx vitest run — 2074/2074 pass (includes new src/__tests__/ast-extract.test.ts, 11 cases covering TS/Python/Go extraction, merge precedence, gap recording, and enrich provenance)
  • npm run build — success
  • E2E: real teamai codebase --extract on a multi-language sample repo produces code-ast DEPENDS_ON + REFERENCES edges; TEAMAI_SKIP_AST=1 falls back to heuristic-only + writes AST_UNAVAILABLE gap
  • Packaged form: npm pack + install into a clean consumer — all .wasm files resolve from node_modules, AST track engages via the installed CLI (verifies npx/npm distribution works)

Docs

README and usage-guide updated in both EN and zh-CN; new TEAMAI_SKIP_AST env var documented.

Compatibility notes for downstream consumers

  • recall graph-boost already registers REFERENCES in RELATION_WEIGHT and uses maxBoost semantics, so the new edges participate correctly with no change and can't inflate scores via duplicate edges.
  • GraphEdgeSource already contains code-ast; no schema change.

🤖 Generated with Claude Code

m0Nst3r873 and others added 2 commits August 20, 2026 20:40
…graph

The code knowledge graph extractors were purely regex/line-based, which
produced false-positive dependency edges (path-substring matching) and had
no notion of call relationships. This adds a real AST track using
web-tree-sitter (pure-WASM, no native toolchain) for TypeScript/JavaScript,
Python, and Go, ported from the team-wiki reference implementation.

- New ast/ module: WASM parser registry (async one-time init), tree-sitter
  queries, symbol/import/call-site walk, import & call resolvers, and
  fact/edge adapters. Emits precise file-to-file DEPENDS_ON / REFERENCES
  edges tagged source:"code-ast" with confidence weights.
- Dual-track: runs alongside the regex heuristic track (which still covers
  Java/Rust/config); AST facts win on merge. Falls back to heuristic-only
  and records an AST_UNAVAILABLE gap when the runtime is unavailable or
  TEAMAI_SKIP_AST=1.
- code-graph: AST relation facts build precise edges instead of being
  re-fuzzed through the path-substring matcher.
- enrich: manifest edges now preserve real AST relation/source provenance
  (deterministic rank-based merge) instead of hardcoding DEPENDS_ON /
  code-heuristic.
- Pinned web-tree-sitter@0.25.10 + tree-sitter-wasms@0.1.13 (only ABI-14
  compatible pair; 0.26 rejects these grammars).
- Docs (README + usage-guide, EN/zh-CN) and unit tests added.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- walk.ts: free the parsed tree-sitter Tree via try/finally { tree.delete() }
  to release WASM off-heap memory; JS GC does not reclaim it, so large repos
  would otherwise grow the emscripten heap unbounded.
- call-resolver.ts: memoize import bindings per source file in resolveCallSites
  (was rebuilt for every call site — O(callSites × imports)).
- import-bindings.ts: detect Python exports via module-level class/function
  definitions instead of the non-existent "__export__" node type, so Python
  symbol-level call resolution works (was always exported=false).
- merge-edges.ts / import-resolver.ts: drop unused findConflictingEdges,
  EdgeConflict, and clearTsconfigCache (speculative dead code).
- walk.ts: simplify a no-op ternary on the call receiver.
- tests: add a Python module-level-export regression case.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@m0Nst3r873

Copy link
Copy Markdown
Collaborator Author

Code review follow-up (pushed in 6555763)

Ran a senior-engineer review pass over the diff. No P0. Addressed the following in this branch:

Correctness / robustness (fixed):

  • WASM memory leakwalkFile now frees the parsed tree-sitter Tree via try/finally { tree.delete() }. JS GC does not reclaim WASM off-heap memory, so a large repo would have grown the emscripten heap unbounded.
  • Quadratic call resolutionresolveCallSites now memoizes import bindings per source file (was rebuilt for every call site).
  • Python export detection — was gated on a non-existent __export__ tree-sitter node, so every Python symbol was exported=false and symbol-level call resolution never fired. Now detects module-level class/def. Added a regression test.

Cleanup (fixed):

  • Removed speculative dead code (findConflictingEdges, EdgeConflict, clearTsconfigCache) and simplified a no-op ternary.

Known limitations (deliberately deferred — the regex heuristic track still fills these gaps)

These are AST-resolver precision gaps; when the AST track doesn't resolve an edge it degrades to the heuristic track, so no coverage is lost — only precision is bounded:

  • Python multi-segment imports (from a.b.c import x) don't map dotted paths to nested files yet — only single-segment sibling imports resolve.
  • Go cross-package imports using full module paths (github.com/org/repo/pkg) resolve to an external/unresolved gap (no go.mod module-prefix stripping).
  • IMPLEMENTS edges are handled downstream but no query capture emits them yet.

Happy to fold any of these into this PR if preferred, or track them as follow-ups.

Full suite green after the fixes: tsc clean, 2075 tests pass, build + real-CLI E2E verified.

Two AST-resolver precision improvements from PR review follow-up:

- Python multi-segment imports: `from a.b.c import x` now maps the dotted
  module to a nested path (a/b/c.py or a/b/c/__init__.py) instead of only
  resolving single-segment sibling imports.
- TS/TSX IMPLEMENTS edges: a class's `implements` clause now produces
  IMPLEMENTS edges (source:"code-ast") to the interface's defining file,
  resolved via same-file interface symbols or imported bindings. A separate
  query pattern avoids the capture-map collision when a class implements
  multiple interfaces; names that resolve to neither (ambient/global types)
  are skipped rather than emitting a spurious edge.

Adds tests for both; Go cross-package module resolution remains a known
follow-up.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@m0Nst3r873

Copy link
Copy Markdown
Collaborator Author

Addressed two of the three known limitations (pushed in 3add345)

  • Python multi-segment importsfrom a.b.c import x now maps the dotted module to a nested path (a/b/c.py or a/b/c/__init__.py). Previously only single-segment sibling imports resolved.
  • TS/TSX IMPLEMENTS edges — a class's implements clause now emits IMPLEMENTS edges (source: "code-ast") to the interface's defining file, resolved via same-file interface symbols or imported bindings. Uses a separate query pattern so a class implementing multiple interfaces doesn't collide in the capture map; names that resolve to neither (ambient/global types) are skipped rather than emitting a spurious edge.

Both covered by new tests (unit + real-CLI E2E). Full suite: tsc clean, 2078 tests pass.

Still deferred: Go cross-package imports (github.com/org/repo/pkg) — needs go.mod module-prefix stripping plus directory-level package resolution; lower value / higher cost, and the heuristic track still covers it. Happy to take it in a follow-up if desired.

m0Nst3r873 and others added 5 commits August 21, 2026 10:41
Keep README and usage-guide (EN/zh-CN) in sync with the new TS implements
edge support.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ping AST edges

Batch imports (--from-repo-list and --from-org, which run importFromRepo
concurrently at concurrency>=2) were losing AST edges: the final per-repo
graph-index.json committed to the team repo contained only code-heuristic
edges, while single --from-repo and concurrency=1 produced the correct
code-ast edges.

Root cause: each concurrent importFromRepo writes into the shared
teamwikiRoot (per-repo graph copy, facts/interfaces caches, source-manifest,
router/index, and reconcileKnowledge's global graph). Those are
read-modify-write operations on shared files, so parallel repos clobbered
each other's artifacts.

Fix: a module-level promise-chain mutex. The clone + extractCodebase phase
(the expensive part, which writes only to the per-repo cache dir) stays
parallel; only the team-repo write phase is serialized. A repo acquires the
lock after extract and releases it in the existing finally, so it is
exception-safe and never deadlocks. Verified: with the fix, concurrency=3
batch import produces the same code-ast edge counts as concurrency=1.

Adds unit tests asserting the mutex's mutual-exclusion, re-acquire, and FIFO
contract.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
During import, teamwiki artifacts sit uncommitted in the team-repo
working tree until the final aggregate+push. An ambient teamai hook
(SessionStart -> reportUsageToTeam) runs `git reset --hard HEAD` on that
same repo, wiping those artifacts before they are committed.

Add a cross-process file lock (import-lock.ts): the lock file lives in
the team-repo's PARENT dir so it is never reset away. It carries pid +
host + startedAt, with a 2h TTL, same-host PID liveness probe, and
dead-lock self-cleanup. Reference-counted for reentrant single/batch use.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Wire the cross-process lock into the import paths and honor it in the
reset path:

- reportUsageToTeam: before resetToCleanMaster, check isImportInProgress
  and skip reset+pull when an import holds the lock.
- importFromRepo: acquire the lock around the whole artifact-write phase
  (steps 4-7), release in finally.
- importFromRepoList: hold one outer batch lock across the whole run so
  per-repo skipAutoPush artifacts survive until the final push.

Add deep-enrich-graph-preserve regression test proving deepEnrich never
mutates the per-repo graph-index.json.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@m0Nst3r873

Copy link
Copy Markdown
Collaborator Author

Follow-up: fix import artifacts being wiped by an ambient git reset --hard

While validating the AST track on a real multi-repo import, the per-repo graph-index.json kept ending up as the old regex-only graph (0 code-ast edges) even though the AST track produced the correct edges (43 edges / 10 code-ast).

Root cause (not the AST code — a pre-existing import/hook race):
import writes teamwiki artifacts into the team-repo working tree and leaves them uncommitted until the final aggregate+push. Meanwhile an ambient teamai hook (SessionStart → pull() → reportUsageToTeam()) runs git reset --hard HEAD on that same repo (utils/git.ts resetToCleanMaster). Any hook firing during the import window reverts the uncommitted artifacts to the last committed version. On the affected machine the team-repo reflog showed 127 reset: moving to HEAD across a single import run. This is invisible to in-process fs instrumentation because the reset happens in a separate teamai process.

Fix — cooperative cross-process lock:

  • src/utils/import-lock.ts: lock file placed in the team-repo's parent dir (so it is never reset away), carrying pid+host+startedAt, 2h TTL, same-host PID-liveness probe, dead-lock self-cleanup, reference-counted for reentrant single/batch use.
  • reportUsageToTeam: skip reset --hard/pull when isImportInProgress() (checked right after the existing isDedicatedRepoRoot guard).
  • importFromRepo / importFromRepoList: hold the lock across the whole artifact-write phase (single) and the whole batch (until final push), released in finally.

Validation: positive + negative control — firing the real SessionStart hook with the lock held → uncommitted AST artifact survives; without the lock → it is wiped (reproduces the original bug). 9 new unit tests for the lock; full suite green; type-check clean.

@jeff-r2026
jeff-r2026 self-requested a review August 24, 2026 03:55
The AST import resolver only handled relative and sibling-module imports (based on the importing file directory). Python code importing via the top-level package name rooted at the repo root (e.g. from hai_flow.conf import config) failed to resolve, leaving the import-binding table empty. Since cross-file call edges are built from those bindings, a large pure-Python repo produced zero code-ast edges despite hundreds of resolved calls. Add resolveAbsolutePackageImport: map a dotted specifier onto a repo-root-relative path and probe for a module file or package __init__.py. On a real 669-file Python repo this took the graph from 0 code-ast edges to 948.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

@jeff-r2026 jeff-r2026 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the AST track worth its complexity? Not in this form.

I want to separate the direction (good) from this implementation (not mergeable as-is). The concern isn't complexity for its own sake — it's that most of the complexity's payoff isn't actually realized on real inputs.

The PR's own premise splits into two problems

① False-positive import edges (file.includes(importPath) links ./user to both user.ts and user-repo.ts).
This is a real bug — but it does not need an AST. The root cause is substring path matching; the fix is normalizing the import path and resolving it precisely to a file. That's a string/regex-level change, an order of magnitude smaller than a WASM runtime. The "67% → 100% precision" win in the PR table is almost entirely attributable to this, not to tree-sitter.

② Call relationships (REFERENCES edges).
This is the only part that genuinely can't be done with regex. So the real question collapses to: do we need a call graph right now, and is there a downstream consumer of REFERENCES edges yet? If not, the AST track pre-pays 2294 lines + a WASM runtime + an ABI-lock (0.25.10 ↔ 0.1.13) for a feature with no consumer.

More importantly: the multi-language AST is broken on its main paths

I verified these empirically (installed the actual WASM grammars and ran them). These are "core feature silently doesn't work," not polish:

  • Python relative imports are silently dropped. ast/queries.ts:44 only captures module_name: (dotted_name). from .util import boot / from . import x / from ..pkg import y parse as relative_import, match zero patterns, and produce no edge and no gap. Relative imports are the dominant intra-package Python style — so for a typical Python repo the AST track resolves almost no internal edges and doesn't even flag it.
  • Every Go symbol is judged non-exported. ast/import-bindings.ts:111 slices the source before the declaration (source.slice(startIndex-20, startIndex)), so /^func\s+[A-Z]/ tests the trailing junk before func and always fails. call-resolver requires exported for cross-file resolution, so Go cross-file call edges to exported symbols never fire.
  • .jsx routed to the non-JSX grammar. ast/parser-registry.ts:69 maps .jsx → typescript (only .tsx → tsx). A .jsx file with any JSX element parses with hasError === true → corrupted tree → imports/symbols/calls extracted wrong or not at all. Should map to tsx.

The 100%-precision benchmark only covers a single TS naming-collision fixture — the one scenario that happens to dodge all three. The Python defaultBinding test also only passes because the fixture util.py has a single symbol, which masks a mis-resolution (call-resolver.ts:30 resolves the first named import to the first exported symbol regardless of name).

Lower-severity issues found in the same pass

  • call-resolver.ts:25 — named-only imports set localToFile with an empty-string key (last-writer-wins pollution).
  • import-bindings.ts:31 — inline type specifiers (import { type Foo, bar }) yield a bogus binding name "type Foo", losing the real Foo binding.
  • walk.ts:112 — capture name @call.member is reused for both the property identifier and the whole member-call, so byName resolves callNode to the property node (latent wrong-node).
  • index.ts:15 — the astAvailability memo isn't cleared by resetParserRegistryForTests, so the reset helper is misleadingly incomplete.

Recommendation

  1. Ship ① on its own first — replace file.includes() with precise path resolution. Zero new deps, captures the bulk of the claimed precision win.
  2. Gate ② on a real need — only add the call graph once something actually consumes REFERENCES.
  3. If we do want it now, it can't merge in this state — the Python relative-import, Go export, and .jsx bugs must be fixed and the benchmark must cover Python relative imports and Go cross-file calls, not just the one TS fixture.

Happy to prototype the AST-free ① as a comparison if that's useful.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants