mirror of
https://github.com/val1813/kwcode.git
synced 2026-09-03 06:34:30 +08:00
a502c9db63
三个模块按顺序实施，424测试全绿： 模块A - 多语言支持： - ast-grep预定义模板查询（LLM只填参数不写pattern） - TreeSitterParser扩展支持JS/TS/Go/Rust/Java（可选依赖） - 多语言Verifier（jest/go test/cargo test/mvn test） - 4个新SKILL.md（golang/typescript/rust/java） 模块B - Server + TUI： - FastAPI + SSE事件流，端口7355 - kwcode serve命令 + kwcode --tui选项 - Textual TUI（文件树+事件流+输入框） - pipeline_factory共享构建 模块C - VSCode插件： - 薄客户端连接kwcode server - SSE事件流 + Webview面板 - 文件保存触发RIG刷新 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
529 lines
21 KiB
Python
529 lines
21 KiB
Python
"""
|
||
Code graph builder: full + incremental build, persisted to SQLite.
|
||
LOC-RED-2: Graph data persisted to SQLite, survives restart.
|
||
LOC-RED-4: Incremental update for modified files only.
|
||
"""
|
||
|
||
import logging
|
||
import os
|
||
import sqlite3
|
||
import subprocess
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
from kaiwu.ast_engine.parser import TreeSitterParser
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
DB_PATH = Path.home() / ".kwcode" / "graph.db"
|
||
|
||
_BASE_EXTENSIONS = {".py"} # Always supported (tree-sitter-python)
|
||
|
||
# Dynamically add extensions for available tree-sitter bindings
|
||
def _get_supported_extensions() -> set:
|
||
"""Return supported extensions based on installed tree-sitter bindings."""
|
||
exts = set(_BASE_EXTENSIONS)
|
||
try:
|
||
from kaiwu.ast_engine.parser import TreeSitterParser
|
||
_p = TreeSitterParser()
|
||
lang_ext_map = {
|
||
"javascript": {".js", ".mjs"},
|
||
"typescript": {".ts", ".tsx"},
|
||
"go": {".go"},
|
||
"rust": {".rs"},
|
||
"java": {".java"},
|
||
}
|
||
for lang in _p.supported_languages():
|
||
if lang in lang_ext_map:
|
||
exts.update(lang_ext_map[lang])
|
||
except Exception:
|
||
pass
|
||
return exts
|
||
|
||
SUPPORTED_EXTENSIONS = _get_supported_extensions()
|
||
|
||
SKIP_DIRS = {
|
||
".git", "__pycache__", "node_modules", ".venv", "venv",
|
||
"env", "dist", "build", ".tox", "htmlcov", ".pytest_cache",
|
||
".eggs", ".mypy_cache",
|
||
}
|
||
|
||
SKIP_FILE_PATTERNS = {"test_", "_test.", "conftest."}
|
||
|
||
|
||
class GraphBuilder:
|
||
"""Builds and maintains a code call graph in SQLite."""
|
||
|
||
def __init__(self, project_root: str):
|
||
self.project_root = str(Path(project_root).resolve())
|
||
self.db_path = DB_PATH
|
||
self._parser = TreeSitterParser()
|
||
self._init_db()
|
||
|
||
def _get_conn(self) -> sqlite3.Connection:
|
||
conn = sqlite3.connect(str(self.db_path), timeout=10.0)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA journal_mode=WAL")
|
||
conn.execute("PRAGMA synchronous=NORMAL")
|
||
return conn
|
||
|
||
def _init_db(self):
|
||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||
with self._get_conn() as conn:
|
||
conn.executescript("""
|
||
CREATE TABLE IF NOT EXISTS nodes (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT NOT NULL,
|
||
qualified TEXT,
|
||
file_path TEXT NOT NULL,
|
||
start_line INTEGER,
|
||
end_line INTEGER,
|
||
node_type TEXT DEFAULT 'function',
|
||
docstring TEXT DEFAULT '',
|
||
search_text TEXT,
|
||
task_count INTEGER DEFAULT 0,
|
||
success_count INTEGER DEFAULT 0,
|
||
last_modified TEXT,
|
||
project_root TEXT NOT NULL,
|
||
UNIQUE(qualified, file_path, project_root)
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS edges (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
from_id INTEGER NOT NULL REFERENCES nodes(id),
|
||
to_id INTEGER NOT NULL REFERENCES nodes(id),
|
||
edge_type TEXT NOT NULL,
|
||
project_root TEXT NOT NULL,
|
||
UNIQUE(from_id, to_id, edge_type)
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS graph_meta (
|
||
project_root TEXT PRIMARY KEY,
|
||
last_built TEXT,
|
||
last_commit TEXT,
|
||
file_count INTEGER,
|
||
node_count INTEGER,
|
||
edge_count INTEGER,
|
||
build_time_ms INTEGER
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_nodes_project
|
||
ON nodes(project_root);
|
||
CREATE INDEX IF NOT EXISTS idx_nodes_file
|
||
ON nodes(file_path, project_root);
|
||
CREATE INDEX IF NOT EXISTS idx_nodes_name
|
||
ON nodes(name, project_root);
|
||
CREATE INDEX IF NOT EXISTS idx_edges_from ON edges(from_id);
|
||
CREATE INDEX IF NOT EXISTS idx_edges_to ON edges(to_id);
|
||
""")
|
||
|
||
def needs_rebuild(self) -> bool:
|
||
current_commit = self._get_current_commit()
|
||
if not current_commit:
|
||
# Not a git repo — check if we ever built
|
||
with self._get_conn() as conn:
|
||
row = conn.execute(
|
||
"SELECT last_built FROM graph_meta WHERE project_root=?",
|
||
(self.project_root,)
|
||
).fetchone()
|
||
return row is None
|
||
|
||
with self._get_conn() as conn:
|
||
row = conn.execute(
|
||
"SELECT last_commit FROM graph_meta WHERE project_root=?",
|
||
(self.project_root,)
|
||
).fetchone()
|
||
|
||
if not row:
|
||
return True
|
||
return row["last_commit"] != current_commit
|
||
|
||
def get_last_commit(self) -> Optional[str]:
|
||
with self._get_conn() as conn:
|
||
row = conn.execute(
|
||
"SELECT last_commit FROM graph_meta WHERE project_root=?",
|
||
(self.project_root,)
|
||
).fetchone()
|
||
return row["last_commit"] if row else None
|
||
|
||
def build_full(self) -> dict:
|
||
t0 = time.perf_counter()
|
||
logger.info("[graph] full build: %s", self.project_root)
|
||
|
||
with self._get_conn() as conn:
|
||
conn.execute("DELETE FROM edges WHERE project_root=?", (self.project_root,))
|
||
conn.execute("DELETE FROM nodes WHERE project_root=?", (self.project_root,))
|
||
|
||
source_files = self._collect_source_files()
|
||
logger.info("[graph] found %d source files", len(source_files))
|
||
|
||
node_count = 0
|
||
edge_count = 0
|
||
for fpath in source_files:
|
||
try:
|
||
n, e = self._parse_file(fpath)
|
||
node_count += n
|
||
edge_count += e
|
||
except Exception as ex:
|
||
logger.warning("[graph] parse failed %s: %s", fpath, ex)
|
||
|
||
current_commit = self._get_current_commit()
|
||
elapsed_ms = int((time.perf_counter() - t0) * 1000)
|
||
|
||
with self._get_conn() as conn:
|
||
conn.execute("""
|
||
INSERT OR REPLACE INTO graph_meta
|
||
(project_root, last_built, last_commit,
|
||
file_count, node_count, edge_count, build_time_ms)
|
||
VALUES (?, datetime('now'), ?, ?, ?, ?, ?)
|
||
""", (self.project_root, current_commit or "",
|
||
len(source_files), node_count, edge_count, elapsed_ms))
|
||
|
||
logger.info("[graph] full build done: %d nodes %d edges %dms",
|
||
node_count, edge_count, elapsed_ms)
|
||
return {
|
||
"node_count": node_count,
|
||
"edge_count": edge_count,
|
||
"file_count": len(source_files),
|
||
"elapsed_ms": elapsed_ms,
|
||
}
|
||
|
||
def update_files(self, file_paths: list[str]) -> dict:
|
||
t0 = time.perf_counter()
|
||
node_count = 0
|
||
edge_count = 0
|
||
|
||
for file_path in file_paths:
|
||
try:
|
||
rel_path = os.path.relpath(file_path, self.project_root).replace("\\", "/")
|
||
except ValueError:
|
||
continue
|
||
|
||
# Delete old nodes/edges for this file
|
||
with self._get_conn() as conn:
|
||
old_ids = [row[0] for row in conn.execute(
|
||
"SELECT id FROM nodes WHERE file_path=? AND project_root=?",
|
||
(rel_path, self.project_root)
|
||
).fetchall()]
|
||
if old_ids:
|
||
ph = ",".join("?" * len(old_ids))
|
||
conn.execute(
|
||
f"DELETE FROM edges WHERE from_id IN ({ph}) OR to_id IN ({ph})",
|
||
old_ids + old_ids
|
||
)
|
||
conn.execute(f"DELETE FROM nodes WHERE id IN ({ph})", old_ids)
|
||
|
||
# Re-parse
|
||
if os.path.exists(file_path):
|
||
try:
|
||
n, e = self._parse_file(file_path)
|
||
node_count += n
|
||
edge_count += e
|
||
except Exception as ex:
|
||
logger.warning("[graph] incremental update failed %s: %s", file_path, ex)
|
||
|
||
elapsed_ms = int((time.perf_counter() - t0) * 1000)
|
||
logger.info("[graph] incremental update: %d files %d nodes %d edges %dms",
|
||
len(file_paths), node_count, edge_count, elapsed_ms)
|
||
return {
|
||
"files": len(file_paths),
|
||
"node_count": node_count,
|
||
"edge_count": edge_count,
|
||
"elapsed_ms": elapsed_ms,
|
||
}
|
||
|
||
def _collect_source_files(self) -> list[str]:
|
||
files = []
|
||
for dirpath, dirnames, filenames in os.walk(self.project_root):
|
||
dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")]
|
||
for fname in sorted(filenames):
|
||
ext = os.path.splitext(fname)[1].lower()
|
||
if ext not in SUPPORTED_EXTENSIONS:
|
||
continue
|
||
if any(fname.startswith(p) or p in fname for p in SKIP_FILE_PATTERNS):
|
||
continue
|
||
files.append(os.path.join(dirpath, fname))
|
||
return files
|
||
|
||
def _parse_file(self, file_path: str) -> tuple[int, int]:
|
||
rel_path = os.path.relpath(file_path, self.project_root).replace("\\", "/")
|
||
|
||
tree = self._parser.parse_file(file_path)
|
||
if tree is None:
|
||
return 0, 0
|
||
|
||
try:
|
||
with open(file_path, "r", encoding="utf-8", errors="replace") as f:
|
||
source = f.read().encode("utf-8")
|
||
except Exception:
|
||
return 0, 0
|
||
|
||
functions = self._parser.extract_functions(tree, source)
|
||
calls = self._parser.extract_calls(tree, source)
|
||
|
||
node_count = 0
|
||
edge_count = 0
|
||
|
||
with self._get_conn() as conn:
|
||
# Insert nodes
|
||
for func in functions:
|
||
name = func["name"]
|
||
# qualified = name (already includes Class.method from parser)
|
||
# Rich search_text: name + path components + short name
|
||
path_parts = rel_path.replace("/", " ").replace("\\", " ").replace(".", " ").replace("_", " ")
|
||
short_name = name.split(".")[-1] if "." in name else name
|
||
search_text = f"{name} {short_name} {rel_path} {path_parts}"
|
||
try:
|
||
conn.execute("""
|
||
INSERT OR IGNORE INTO nodes
|
||
(name, qualified, file_path, start_line, end_line,
|
||
node_type, search_text, project_root, last_modified)
|
||
VALUES (?, ?, ?, ?, ?, 'function', ?, ?, datetime('now'))
|
||
""", (name, name, rel_path,
|
||
func.get("start_line"), func.get("end_line"),
|
||
search_text, self.project_root))
|
||
node_count += 1
|
||
except sqlite3.IntegrityError:
|
||
pass
|
||
|
||
# Insert call edges
|
||
for call in calls:
|
||
caller = call.get("in_function")
|
||
callee = call.get("name")
|
||
if not caller or not callee:
|
||
continue
|
||
|
||
caller_row = conn.execute(
|
||
"SELECT id FROM nodes WHERE name=? AND file_path=? AND project_root=? LIMIT 1",
|
||
(caller, rel_path, self.project_root)
|
||
).fetchone()
|
||
|
||
# Callee could be in any file
|
||
callee_row = conn.execute(
|
||
"SELECT id FROM nodes WHERE name=? AND project_root=? LIMIT 1",
|
||
(callee, self.project_root)
|
||
).fetchone()
|
||
if not callee_row:
|
||
# Try short name match (e.g. "bar" -> "Foo.bar")
|
||
callee_row = conn.execute(
|
||
"SELECT id FROM nodes WHERE name LIKE ? AND project_root=? LIMIT 1",
|
||
(f"%.{callee}", self.project_root)
|
||
).fetchone()
|
||
|
||
if caller_row and callee_row:
|
||
try:
|
||
conn.execute("""
|
||
INSERT OR IGNORE INTO edges
|
||
(from_id, to_id, edge_type, project_root)
|
||
VALUES (?, ?, 'CALLS', ?)
|
||
""", (caller_row["id"], callee_row["id"], self.project_root))
|
||
edge_count += 1
|
||
except sqlite3.IntegrityError:
|
||
pass
|
||
|
||
return node_count, edge_count
|
||
|
||
def _get_current_commit(self) -> str:
|
||
try:
|
||
result = subprocess.run(
|
||
["git", "rev-parse", "HEAD"],
|
||
cwd=self.project_root,
|
||
capture_output=True, text=True, timeout=3
|
||
)
|
||
return result.stdout.strip() if result.returncode == 0 else ""
|
||
except Exception:
|
||
return ""
|
||
|
||
def export_rig(self) -> dict:
|
||
"""Scan project and output .kaiwu/rig.json with structure map."""
|
||
import json
|
||
import re
|
||
|
||
rig = {
|
||
"files": {},
|
||
"api_routes": {},
|
||
"test_coverage": {},
|
||
"frontend_api_calls": {},
|
||
"language_stats": {},
|
||
}
|
||
|
||
# --- Language stats ---
|
||
try:
|
||
from kaiwu.ast_engine.language_detector import detect_project_languages
|
||
rig["language_stats"] = detect_project_languages(self.project_root)
|
||
except Exception:
|
||
pass
|
||
|
||
# --- Collect all Python files (including test files) ---
|
||
all_py_files = []
|
||
for dirpath, dirnames, filenames in os.walk(self.project_root):
|
||
dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")]
|
||
for fname in sorted(filenames):
|
||
if fname.endswith(".py"):
|
||
all_py_files.append(os.path.join(dirpath, fname))
|
||
|
||
# --- Parse each Python file for exports/imports/routes ---
|
||
re_import = re.compile(r"^\s*import\s+([\w.]+)", re.MULTILINE)
|
||
re_from_import = re.compile(r"^\s*from\s+([\w.]+)\s+import", re.MULTILINE)
|
||
re_top_def = re.compile(r"^(?:def|class)\s+(\w+)", re.MULTILINE)
|
||
re_route = re.compile(
|
||
r"@(?:app|router)\.(route|get|post|put|delete|patch)\(\s*['\"]([^'\"]+)['\"]",
|
||
re.MULTILINE,
|
||
)
|
||
|
||
for fpath in all_py_files:
|
||
rel_path = os.path.relpath(fpath, self.project_root).replace("\\", "/")
|
||
try:
|
||
with open(fpath, "r", encoding="utf-8", errors="replace") as f:
|
||
source = f.read()
|
||
except Exception:
|
||
continue
|
||
|
||
# Exports: top-level def/class names
|
||
exports = re_top_def.findall(source)
|
||
|
||
# Imports: module names
|
||
imports = list(set(re_import.findall(source) + re_from_import.findall(source)))
|
||
|
||
if exports or imports:
|
||
rig["files"][rel_path] = {
|
||
"exports": exports,
|
||
"imports": imports,
|
||
}
|
||
|
||
# API routes
|
||
for match in re_route.finditer(source):
|
||
method_raw, path = match.group(1), match.group(2)
|
||
method = method_raw.upper() if method_raw != "route" else "GET"
|
||
# Find the function defined right after the decorator
|
||
after = source[match.end():]
|
||
func_match = re.search(r"^\s*def\s+(\w+)", after, re.MULTILINE)
|
||
func_name = func_match.group(1) if func_match else "unknown"
|
||
route_key = f"{method} {path}"
|
||
rig["api_routes"][route_key] = f"{rel_path}:{func_name}"
|
||
|
||
# --- Test coverage: match test_foo.py -> foo.py ---
|
||
source_files_by_name = {}
|
||
for rel_path in rig["files"]:
|
||
basename = os.path.basename(rel_path)
|
||
source_files_by_name[basename] = rel_path
|
||
|
||
for fpath in all_py_files:
|
||
rel_path = os.path.relpath(fpath, self.project_root).replace("\\", "/")
|
||
basename = os.path.basename(rel_path)
|
||
# Match test_foo.py -> foo.py
|
||
if basename.startswith("test_"):
|
||
target_name = basename[5:] # strip "test_"
|
||
if target_name in source_files_by_name:
|
||
target_rel = source_files_by_name[target_name]
|
||
rig["test_coverage"].setdefault(target_rel, []).append(rel_path)
|
||
|
||
# --- Frontend API calls: scan .js/.ts files ---
|
||
re_axios_method = re.compile(
|
||
r"""axios\.(get|post|put|delete|patch)\(\s*[`'"]([^`'"]+)[`'"]""",
|
||
re.MULTILINE,
|
||
)
|
||
re_fetch_method = re.compile(
|
||
r"""fetch\(\s*[`'"]([^`'"]+)[`'"].*?method:\s*['\"](\w+)['\"]""",
|
||
re.DOTALL,
|
||
)
|
||
|
||
for dirpath, dirnames, filenames in os.walk(self.project_root):
|
||
dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")]
|
||
for fname in sorted(filenames):
|
||
if not fname.endswith((".js", ".ts", ".jsx", ".tsx")):
|
||
continue
|
||
fpath = os.path.join(dirpath, fname)
|
||
rel_path = os.path.relpath(fpath, self.project_root).replace("\\", "/")
|
||
try:
|
||
with open(fpath, "r", encoding="utf-8", errors="replace") as f:
|
||
content = f.read()
|
||
except Exception:
|
||
continue
|
||
|
||
# axios.post("/path", ...) style
|
||
for match in re_axios_method.finditer(content):
|
||
method = match.group(1).upper()
|
||
path = match.group(2)
|
||
# Find last function/const def before this call
|
||
before = content[:match.start()]
|
||
func_matches = re.findall(
|
||
r"(?:function\s+(\w+)|(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?\(?)",
|
||
before,
|
||
)
|
||
func_name = "anonymous"
|
||
if func_matches:
|
||
last = func_matches[-1]
|
||
func_name = last[0] or last[1]
|
||
route_key = f"{method} {path}"
|
||
rig["frontend_api_calls"][route_key] = f"{rel_path}:{func_name}"
|
||
|
||
# fetch("/path", {method: "POST"}) style
|
||
for match in re_fetch_method.finditer(content):
|
||
path = match.group(1)
|
||
method = match.group(2).upper()
|
||
before = content[:match.start()]
|
||
func_matches = re.findall(
|
||
r"(?:function\s+(\w+)|(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?\(?)",
|
||
before,
|
||
)
|
||
func_name = "anonymous"
|
||
if func_matches:
|
||
last = func_matches[-1]
|
||
func_name = last[0] or last[1]
|
||
route_key = f"{method} {path}"
|
||
rig["frontend_api_calls"][route_key] = f"{rel_path}:{func_name}"
|
||
|
||
# --- Write output ---
|
||
out_dir = os.path.join(self.project_root, ".kaiwu")
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
out_path = os.path.join(out_dir, "rig.json")
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
json.dump(rig, f, indent=2, ensure_ascii=False)
|
||
|
||
# --- Write summary (lightweight, <5KB, for Gate/Locator prompt injection) ---
|
||
summary = self._build_rig_summary(rig)
|
||
summary_path = os.path.join(out_dir, "rig_summary.json")
|
||
with open(summary_path, "w", encoding="utf-8") as f:
|
||
json.dump(summary, f, indent=2, ensure_ascii=False)
|
||
|
||
logger.info("[graph] exported rig.json: %d files, %d routes, %d test mappings",
|
||
len(rig["files"]), len(rig["api_routes"]), len(rig["test_coverage"]))
|
||
return rig
|
||
|
||
@staticmethod
|
||
def _build_rig_summary(rig: dict) -> dict:
|
||
"""
|
||
生成 Gate/Locator 可注入的精简骨架。
|
||
只保留：api_routes、test_coverage、frontend_api_calls、文件路径列表（截断）。
|
||
不含详细 exports/imports，目标 <5KB。
|
||
"""
|
||
import json
|
||
|
||
summary = {
|
||
"api_routes": rig.get("api_routes", {}),
|
||
"test_coverage": rig.get("test_coverage", {}),
|
||
"frontend_api_calls": rig.get("frontend_api_calls", {}),
|
||
}
|
||
|
||
# File list: only include paths, cap to keep total <5KB
|
||
all_files = sorted(rig.get("files", {}).keys())
|
||
# Budget: 4096 bytes total, subtract non-file content
|
||
non_file_size = len(json.dumps(summary, ensure_ascii=False))
|
||
file_budget = 4096 - non_file_size - 50 # 50 bytes overhead for "files" key + brackets
|
||
# Add files until budget exhausted
|
||
included = []
|
||
running = 0
|
||
for f in all_files:
|
||
entry_size = len(f) + 5 # quotes + comma + space
|
||
if running + entry_size > file_budget:
|
||
break
|
||
included.append(f)
|
||
running += entry_size
|
||
summary["files"] = included
|
||
if len(included) < len(all_files):
|
||
summary["_files_truncated"] = len(all_files)
|
||
|
||
return summary
|