witan-code 0.2.0__py3-none-any.whl

witan_code/config.py

@@ -0,0 +1,65 @@
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+# Bundled query files, resolved relative to this file.
+_QUERIES_DIR = Path(__file__).parent / "queries"
+_SCHEMA_FILE = Path(__file__).parent / "schema" / "code-schema.pg"
+_BRIDGE_SCHEMA_FILE = Path(__file__).parent / "schema" / "bridge-schema.pg"
+_DEFAULT_CODE_DIR = Path.home() / ".local" / "share" / "witan" / "code"
+
+# Filename of the single shared cross-repo bridge store, a sibling of the
+# per-repo `<slug>.omni` stores in code_dir. Not routed through sanitize_slug
+# (whose .strip("_") would eat the leading underscore); no real repo slug
+# resolves to this name.
+BRIDGE_STORE_NAME = "_bridge.omni"
+
+
+@dataclass(frozen=True)
+class Config:
+    code_dir: Path
+    """Directory holding per-repo code stores (one ``<slug>.omni`` each)."""
+
+    author: str
+    """Attribution string (carried for parity with Layer 1; unused on inserts)."""
+
+    queries_dir: Path
+    """Directory containing code_read.gq, code_mutations.gq, delete.gq."""
+
+    schema_file: Path
+    """Path to code-schema.pg, used to lazily init a per-repo store."""
+
+    bridge_schema_file: Path
+    """Path to bridge-schema.pg, used to lazily init the shared bridge store."""
+
+
+def load() -> Config:
+    """Load config from environment. All variables are optional with defaults."""
+    return Config(
+        code_dir=Path(os.environ.get("WITAN_CODE_DIR", str(_DEFAULT_CODE_DIR))),
+        author=os.environ.get(
+            "WITAN_AUTHOR",
+            os.environ.get("USER", "unknown"),
+        ),
+        queries_dir=_QUERIES_DIR,
+        schema_file=_SCHEMA_FILE,
+        bridge_schema_file=_BRIDGE_SCHEMA_FILE,
+    )
+
+
+def sanitize_slug(slug: str) -> str:
+    """Make a repo slug safe for use as a filename component."""
+    return re.sub(r"[/:]+", "_", slug).strip("_")
+
+
+def store_path(slug: str, code_dir: Path | None = None) -> Path:
+    """Resolve the per-repo store path for ``slug``."""
+    base = code_dir or load().code_dir
+    return base / f"{sanitize_slug(slug)}.omni"
+
+
+def bridge_store_path(code_dir: Path | None = None) -> Path:
+    """Resolve the shared cross-repo bridge store path."""
+    base = code_dir or load().code_dir
+    return base / BRIDGE_STORE_NAME

witan_code/context.py

@@ -0,0 +1,98 @@
+"""Self-contained UserPromptSubmit status block for witan-code.
+
+Mirrors the shape of ``witan``'s own ``inject-context`` hook, but is a
+deliberately independent implementation (no cross-package import — see
+``graph.py``'s docstring) so witan-code stays fully usable, and self-announcing,
+when installed standalone without ``witan``.
+
+Tells the agent whether the current repo has a code graph ready to query (and
+its rough size/freshness) — without this, an agent has no signal that
+``code_*`` tools exist or are populated, short of trying one and seeing what
+comes back.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from pathlib import Path
+
+from . import config as cfg_module
+from . import repo as repo_module
+from .cli import _code_store_stats, _dir_stats
+
+# Matches the lock directory hooks.session_init() creates around a background
+# SessionStart index, so this hook can report "indexing in progress" instead
+# of a misleadingly empty/stale store. Keyed on a hash of the project
+# directory (not the raw sanitized path) so two distinct paths can't collide
+# on the same lock file (e.g. "/tmp/a/b" and "/tmp/a_b" both sanitizing to
+# "_tmp_a_b") and so a deep/long checkout path can't blow past a filesystem's
+# filename length limit and silently fail the `mkdir`.
+_LOCK_PREFIX = "codegraph-init-"
+
+
+def _lock_digest(project_dir: Path) -> str:
+    return hashlib.sha256(str(project_dir).encode()).hexdigest()[:16]
+
+
+def _project_dir() -> Path:
+    try:
+        cwd = os.getcwd()
+    except OSError:  # e.g. the cwd was deleted out from under this process
+        cwd = "/"
+    return Path(os.environ.get("CLAUDE_PROJECT_DIR", cwd))
+
+
+def _lock_path(project_dir: Path) -> Path:
+    tmp = Path(os.environ.get("TMPDIR", "/tmp"))
+    return tmp / f"{_LOCK_PREFIX}{_lock_digest(project_dir)}.lock"
+
+
+def indexing_in_progress() -> bool:
+    """Whether hooks.session_init()'s background index is still running."""
+    return _lock_path(_project_dir()).is_dir()
+
+
+def inject_context() -> str:
+    """A short markdown status block, or "" when there's nothing worth saying.
+
+    Silent when the repo has neither a store nor an index in flight (nothing
+    to report), so this hook adds no noise for repos that don't use witan-code.
+    """
+    cfg = cfg_module.load()
+    slug = repo_module.detect()
+    if slug is None:
+        return ""
+
+    store = cfg_module.store_path(slug, cfg.code_dir)
+    in_progress = indexing_in_progress()
+
+    if not store.exists():
+        if not in_progress:
+            return ""
+        return (
+            "## Code Graph\n\n"
+            f"Indexing `{slug}` for the first time in the background — "
+            "`code_*` tools may return partial or empty results until it "
+            "finishes.\n"
+        )
+
+    repo_uri, file_count = _code_store_stats(store)
+    try:
+        _, last_indexed = _dir_stats(store)
+        freshness = f", last updated {last_indexed}"
+    except OSError:  # e.g. a file vanished mid-walk — degrade, don't blank the block
+        freshness = ""
+    lines = [
+        "## Code Graph",
+        "",
+        f"`{repo_uri}` is indexed: {file_count} files{freshness}.",
+    ]
+    if in_progress:
+        lines.append("A background reindex is currently running.")
+    lines.append(
+        "Prefer `code_search_symbol` / `code_find_definition` / "
+        "`code_find_references` / `code_callers` / `code_impact` over grep "
+        "for symbol lookups, call graphs, and change-impact analysis."
+    )
+    return "\n".join(lines) + "\n"

witan_code/edges.py

@@ -0,0 +1,171 @@
+"""Typed cross-repo edge precision tiers (docs/EDGE_PRECISION_TIERS.md).
+
+Kythe-inspired ``ref/call/direct`` vs ``ref/call`` pattern: replace the
+previous binary consumer/provider model with typed edges that carry the
+precision of *how* a cross-repo link was established, so callers filter by a
+minimum trust floor instead of a confidence threshold they have to pick
+themselves.
+
+Tiers, in decreasing trust order:
+  * ``precise``   — Stage 2: canonical symbol string join (``stitch``).
+  * ``heuristic`` — Stage 3: the pre-existing ``(kind, key_norm)`` binding
+    grouping, confidence-scored (``bridge_extractors.adjust_confidence``,
+    ``visualize.cross_repo_edges``).
+  * ``fuzzy``     — future: embedding/BM25 route similarity (tracked
+    separately: "Research: REST-route cross-service matching via embedding
+    similarity"). No fuzzy edges exist yet, so this tier is currently
+    identical to ``heuristic``.
+
+``min_precision`` is a FLOOR: ``"precise"`` returns only the precise tier;
+``"heuristic"`` (the default, preserving pre-existing behavior) returns
+precise + heuristic; ``"fuzzy"`` returns everything. A heuristic edge is
+suppressed whenever a precise edge already covers the same
+``(consumer_repo, provider_repo, kind, key_norm)`` triple, so the same
+logical link never appears twice at two precision tiers.
+"""
+
+from dataclasses import dataclass
+
+from . import stitch
+from .visualize import cross_repo_edges as _confidence_filter
+
+PRECISION_TIERS = ("precise", "heuristic", "fuzzy")
+_TIER_ORDER = {tier: i for i, tier in enumerate(PRECISION_TIERS)}
+
+
+@dataclass(frozen=True)
+class TypedEdge:
+    precision: str  # "precise" | "heuristic" | "fuzzy"
+    consumer_repo: str
+    provider_repo: str
+    kind: str
+    key_norm: str
+    canonical_symbol: str | None  # the consumer's symbol; precise tier only
+    confidence: float
+    evidence: tuple[dict, ...]
+
+    def as_dict(self) -> dict:
+        return {
+            "precision": self.precision,
+            "consumer_repo": self.consumer_repo,
+            "provider_repo": self.provider_repo,
+            "kind": self.kind,
+            "key_norm": self.key_norm,
+            "canonical_symbol": self.canonical_symbol,
+            "confidence": self.confidence,
+            "evidence": list(self.evidence),
+        }
+
+
+def _pair(edge) -> tuple:
+    return (edge.consumer_repo, edge.provider_repo, edge.kind, edge.key_norm)
+
+
+def _precise_edges(repo_symbol_rows: list[dict]) -> list[TypedEdge]:
+    precise, _ = stitch.resolve(repo_symbol_rows)
+    return [
+        TypedEdge(
+            precision="precise",
+            consumer_repo=e.consumer_repo,
+            provider_repo=e.provider_repo,
+            kind=e.kind,
+            key_norm=e.key_norm,
+            canonical_symbol=e.consumer_symbol,
+            confidence=1.0,
+            evidence=e.evidence,
+        )
+        for e in precise
+    ]
+
+
+def _heuristic_edges(
+    binding_rows: list[dict], *, min_confidence: float
+) -> list[TypedEdge]:
+    """Group raw ``InterfaceBinding`` rows into (consumer, provider) edges.
+
+    Mirrors ``visualize.build_graph``'s grouping/self-providing/service-anchor
+    rules but keeps per-occurrence evidence instead of collapsing into a
+    render-oriented ``DepGraph``.
+    """
+    filtered = _confidence_filter(binding_rows, min_confidence=min_confidence)
+
+    groups: dict[tuple[str, str], dict] = {}
+    for b in filtered:
+        if b["kind"] == "service":
+            continue  # service anchors aren't repo-to-repo edges (see visualize.build_graph)
+        group = groups.setdefault((b["kind"], b["key_norm"]), {"p": {}, "c": {}})
+        bucket = group["p"] if b["role"] == "provider" else group["c"]
+        repo_bucket = bucket.setdefault(b["repo"], {"confidence": None, "evidence": []})
+        if b["role"] != "provider":
+            conf = float(b["confidence"]) if b.get("confidence") is not None else 1.0
+            if repo_bucket["confidence"] is None or conf > repo_bucket["confidence"]:
+                repo_bucket["confidence"] = conf
+        if b.get("file"):
+            repo_bucket["evidence"].append(
+                {"repo": b["repo"], "file": b["file"], "line": b.get("line")}
+            )
+
+    out: list[TypedEdge] = []
+    for (kind, key_norm), group in groups.items():
+        for cons_repo, cons in group["c"].items():
+            if cons_repo in group["p"]:
+                continue  # self-providing: the repo serves its own route
+            for prov_repo in group["p"]:
+                if cons_repo == prov_repo:
+                    continue
+                out.append(
+                    TypedEdge(
+                        precision="heuristic",
+                        consumer_repo=cons_repo,
+                        provider_repo=prov_repo,
+                        kind=kind,
+                        key_norm=key_norm,
+                        canonical_symbol=None,
+                        confidence=(
+                            cons["confidence"]
+                            if cons["confidence"] is not None
+                            else 1.0
+                        ),
+                        evidence=tuple(cons["evidence"]),
+                    )
+                )
+    return out
+
+
+def cross_repo_edges(
+    repo_symbol_rows: list[dict],
+    binding_rows: list[dict],
+    *,
+    min_precision: str = "heuristic",
+    min_confidence: float = 0.5,
+) -> list[TypedEdge]:
+    """Merge Stage 2 precise + Stage 3 heuristic edges into one typed, filterable list.
+
+    ``repo_symbol_rows`` is a full ``RepoSymbol`` dump (``all_repo_symbols``);
+    ``binding_rows`` is a full ``InterfaceBinding`` dump (``all_bindings``).
+    See the module docstring for tier semantics.
+    """
+    if min_precision not in _TIER_ORDER:
+        raise ValueError(f"min_precision must be one of {PRECISION_TIERS!r}")
+
+    precise = _precise_edges(repo_symbol_rows)
+    out = list(precise)
+    if _TIER_ORDER[min_precision] >= _TIER_ORDER["heuristic"]:
+        covered = {_pair(e) for e in precise}
+        heuristic = _heuristic_edges(binding_rows, min_confidence=min_confidence)
+        out.extend(e for e in heuristic if _pair(e) not in covered)
+    return out
+
+
+def precise_pairs(repo_symbol_rows: list[dict]) -> frozenset[tuple[str, str, str, str]]:
+    """``(consumer_repo, provider_repo, kind, key_norm)`` covered by a Stage-2 join.
+
+    A cheap membership-test helper for tools that need to know whether a
+    *specific* binding participates in a precise edge without building the
+    full merged/typed edge list — e.g. filtering ``code_interface_*`` rows by
+    ``min_precision="precise"``.
+    """
+    precise, _ = stitch.resolve(repo_symbol_rows)
+    return frozenset(
+        (e.consumer_repo, e.provider_repo, e.kind, e.key_norm) for e in precise
+    )

witan_code/elicit.py

@@ -0,0 +1,92 @@
+"""Additive MCP elicitation helpers for the witan-code tools.
+
+FastMCP 3.4.3 supports ``ctx.elicit``, but the connected client may not, and
+several ``code_*`` tools also run under headless automation (background
+indexers, other MCP clients without elicitation support). These helpers keep
+elicitation strictly *additive*: when the client can't elicit — or ``ctx`` is
+absent — they return the caller's ``default`` so behavior is exactly what it
+was before elicitation existed. Only an *explicit* user decline changes the
+outcome.
+
+Call sites in ``server.py``:
+
+- ``code_symbols_in_file`` offers to elicit a repo URI when none can be
+  detected (``text``).
+- ``code_find_references``/``code_callers``/``code_impact``/
+  ``code_symbols_in_file`` and the bridge-backed interface/cross-repo tools
+  offer to index now when the relevant store is missing (``confirm``, via the
+  ``_confirm_and_reindex``/``_confirm_and_reindex_bridge`` helpers in
+  ``server.py``).
+- ``code_find_definition`` offers to narrow a multi-repo name match down to
+  one repo (``choose_repo``).
+
+Never call these from a tool that must stay non-interactive under automation
+— ``code_reindex`` never takes a ``Context`` (see its docstring), and
+``code_search_symbol`` returns many results by design, not an error case.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from fastmcp.server.elicitation import AcceptedElicitation
+
+if TYPE_CHECKING:
+    from fastmcp import Context
+
+
+async def confirm(
+    ctx: Context | None, message: str, *, default_when_unsupported: bool
+) -> bool:
+    """Ask a yes/no question.
+
+    ``accept`` → the chosen bool; ``decline``/``cancel`` → ``False``; and when
+    elicitation is unsupported or errors (headless client, no ``ctx``) →
+    ``default_when_unsupported`` — pick that so the non-interactive path keeps
+    today's behavior (here, always ``False``: never index without being asked).
+    """
+    if ctx is None:
+        return default_when_unsupported
+    try:
+        result = await ctx.elicit(message, response_type=bool)
+    except Exception:  # noqa: BLE001 — any elicit failure means "can't ask"
+        return default_when_unsupported
+    if isinstance(result, AcceptedElicitation):
+        return bool(result.data)
+    return False
+
+
+async def text(ctx: Context | None, message: str, *, default: str) -> str:
+    """Ask for a line of text. A non-empty accepted value is returned; a
+    decline/cancel, an empty value, an unsupported client, or no ``ctx`` all
+    fall back to ``default``."""
+    if ctx is None:
+        return default
+    try:
+        result = await ctx.elicit(message, response_type=str)
+    except Exception:  # noqa: BLE001
+        return default
+    if isinstance(result, AcceptedElicitation) and (result.data or "").strip():
+        return result.data.strip()
+    return default
+
+
+async def choose_repo(
+    ctx: Context | None, message: str, repos: list[str]
+) -> str | None:
+    """Ask the caller to narrow a multi-repo match down to one of ``repos``.
+
+    Reuses ``text`` for the primitive. The answer is matched against ``repos``
+    case-insensitively and whitespace-stripped; an exact match returns that
+    repo, and anything else — empty, declined, unsupported, or a value that
+    doesn't match any candidate — returns ``None``, meaning "keep every
+    match" (today's behavior, unchanged).
+    """
+    answer = await text(ctx, message, default="")
+    if not answer:
+        return None
+    normalized = answer.strip().casefold()
+    for repo in repos:
+        if repo.casefold() == normalized:
+            return repo
+    return None

witan_code/extensions/pi/codegraph.ts

@@ -0,0 +1,121 @@
+/**
+ * witan-code — Pi extension
+ *
+ * Pi equivalent of the four Claude Code code-graph hooks:
+ *  - session_start     : seed/refresh the whole repo's Layer-2 code graph in
+ *                        the background (first session builds it, later
+ *                        sessions re-hash and skip unchanged files).
+ *  - edit / write      : incrementally re-index the edited file after the
+ *                        tool runs.
+ *  - before_agent_start: report whether the code graph is indexed (file
+ *                        count, last-updated) or still being built, with a
+ *                        nudge toward code_* tools over grep.
+ *  - session_shutdown  : opportunistically compact the current repo's store
+ *                        and the shared cross-repo bridge store (throttled;
+ *                        see witan_code.maintenance) — no session-id
+ *                        dependency, unlike workflow-session-checkpoint, so
+ *                        this one *is* mirrored under Pi.
+ *
+ * Best-effort and non-blocking: a missing CLI, non-git dir, or parse failure
+ * never disrupts the session. Requires `witan-code` on PATH
+ * (`witan-code setup --agent pi`, or `uv tool install --editable
+ * mcp/servers/witan-code`); otherwise it silently no-ops.
+ *
+ * Install: `witan-code setup --agent pi`, or symlink into
+ * ~/.pi/agent/extensions/ (see configs/pi/README.md).
+ */
+
+import { execSync, spawn, spawnSync } from "node:child_process";
+import { existsSync } from "node:fs";
+import { isAbsolute, resolve } from "node:path";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+const SRC_EXT = /\.(py|pyi|ts|tsx|js|jsx|mjs|cjs)$/;
+const EDIT_TOOLS = new Set(["edit", "write"]);
+
+function inGitRepo(cwd: string): boolean {
+	try {
+		execSync("git rev-parse --is-inside-work-tree", {
+			cwd,
+			stdio: ["ignore", "ignore", "ignore"],
+		});
+		return true;
+	} catch {
+		return false;
+	}
+}
+
+/** Run a witan-code subcommand detached in the background; ignore all failures. */
+function runInBackground(args: string[], cwd?: string): void {
+	try {
+		const child = spawn("witan-code", args, {
+			detached: true,
+			stdio: "ignore",
+			...(cwd ? { cwd } : {}),
+		});
+		child.on("error", () => {}); // CLI not installed, etc.
+		child.unref();
+	} catch {
+		/* ignore */
+	}
+}
+
+function editedPath(event: any, cwd: string): string | null {
+	const input = event?.input ?? {};
+	const raw = input.path ?? input.file_path ?? input.filename;
+	if (typeof raw !== "string" || raw.length === 0) return null;
+	const abs = isAbsolute(raw) ? raw : resolve(cwd, raw);
+	return SRC_EXT.test(abs) ? abs : null;
+}
+
+export default function codegraphExtension(pi: ExtensionAPI): void {
+	// Paths captured at tool_call, consumed at tool_result (FIFO) — covers Pi
+	// builds where tool_result carries no input.
+	const pending: string[] = [];
+
+	// Seed / refresh the whole repo once per session (incremental).
+	pi.on("session_start", async (_event, ctx) => {
+		if (inGitRepo(ctx.cwd)) runInBackground(["index", ctx.cwd]);
+	});
+
+	// Push one entry per edit tool_call (path or "") so the FIFO stays paired
+	// 1:1 with tool_result, which always shifts exactly one.
+	pi.on("tool_call", (event: any, ctx) => {
+		if (!EDIT_TOOLS.has(event?.toolName)) return;
+		pending.push(editedPath(event, ctx.cwd) ?? "");
+	});
+
+	// After a successful edit/write, re-index just that file.
+	pi.on("tool_result", (event: any, ctx) => {
+		if (!EDIT_TOOLS.has(event?.toolName)) return;
+		const queued = pending.shift() ?? "";
+		if (event?.isError) return;
+		const p = editedPath(event, ctx.cwd) ?? queued;
+		if (p && existsSync(p)) runInBackground(["index", p]);
+	});
+
+	// Report code-graph readiness before each turn (mirrors the Claude
+	// `witan-code inject-context` UserPromptSubmit hook).
+	pi.on("before_agent_start", async (event: any, ctx: any) => {
+		try {
+			const r = spawnSync("witan-code", ["inject-context"], {
+				encoding: "utf8",
+				timeout: 5000,
+				cwd: ctx?.cwd,
+			});
+			const text = (r.stdout ?? "").trim();
+			if (r.status !== 0 || !text) return;
+			return { systemPrompt: `${event.systemPrompt ?? ""}\n\n${text}` };
+		} catch {
+			return;
+		}
+	});
+
+	// Opportunistically compact the store(s) on session end (mirrors the
+	// Claude `witan-code checkpoint` Stop hook). Detached and non-blocking,
+	// like session_start's index — session_shutdown fires before teardown,
+	// not after, so this must not wait on the child process.
+	pi.on("session_shutdown", async (_event, ctx: any) => {
+		if (inGitRepo(ctx?.cwd)) runInBackground(["checkpoint"], ctx.cwd);
+	});
+}