vaultcompute 0.1.0a1__py3-none-any.whl

vaultcompute/__init__.py

@@ -0,0 +1,18 @@
+"""VaultCompute — privacy proxy for LLM tool calls."""
+
+from vaultcompute.config import describe_config
+from vaultcompute.core.capabilities import TableQueryCapability
+from vaultcompute.core.rehydrator import PLACEHOLDER_PROMPT, rehydrate
+from vaultcompute.core.tokenizer import describe_schema
+from vaultcompute.errors import ProtectionError
+from vaultcompute.session import VaultComputeSession
+
+__all__ = [
+    "PLACEHOLDER_PROMPT",
+    "VaultComputeSession",
+    "ProtectionError",
+    "TableQueryCapability",
+    "describe_config",
+    "describe_schema",
+    "rehydrate",
+]

vaultcompute/__main__.py

@@ -0,0 +1,3 @@
+from vaultcompute.cli import main
+
+raise SystemExit(main())

vaultcompute/audit.py

@@ -0,0 +1,253 @@
+"""`vaultcompute audit` — check a transcript against the vault.
+
+A privacy tool has an observability problem that most tools do not: **when it
+works, the screen looks exactly the same as when it is not installed.** In
+Mode C especially, the display hook puts real values back before you read them,
+so the frontend cannot tell you whether anything was ever hidden.
+
+The ground truth is the transcript — what the model actually received — and the
+answer is not a grep, because knowing whether a value leaked means knowing
+which values were supposed to be hidden. That is what the vault holds.
+
+So: for every record in the vault, look for its value in the transcript. A hit
+is a leak. It reports placeholders too, because "no leaks" is also what an
+empty config produces, and the two need telling apart.
+
+What it cannot tell you: a field you never declared was never tokenized, so
+there is no vault record and nothing to look for. This audit measures whether
+VaultCompute kept the promises your config made, not whether the config named
+everything it should have.
+
+One more case worth naming: a ``vault_compute`` result can be a string
+literal the model already typed into its own code — picking between two
+already-known names based on a hidden comparison, say. That text was never
+secret; it sat in the transcript, in plain sight, in the tool call's own
+arguments, before the token wrapping it even existed. Flagging its later
+reappearance as a leak was a real false positive. It is excluded now —
+reported separately, as "explained" — but only for records ``vault_compute``
+minted, and only when the matched text is provably a literal the model wrote
+itself; a value actually pulled out of ``resolve(...)`` (a computed sum, say)
+is still flagged exactly as before.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Iterator
+
+from vaultcompute.core.rehydrator import TOKEN_PATTERN
+from vaultcompute.ports.token_store import TokenStore
+
+#: A vault_compute tool call's `code` argument, as it sits in a
+#: transcript. Not scoped to that specific tool by name — a "code" field from
+#: something else would only ever add extra literals to the allow-list, never
+#: remove one, which is the safe direction for a heuristic like this.
+_COMPUTE_CODE = re.compile(r'"code"\s*:\s*"((?:\\.|[^"\\])*)"')
+#: A quoted string literal inside Python source.
+_STRING_LITERAL = re.compile(r"""["']((?:\\.|[^"'\\])*)["']""")
+
+
+def _literals_the_model_already_wrote(transcript: str) -> set[str]:
+    """String literals appearing in any vault_compute call's own code.
+
+    These were typed by the model, not extracted from the vault — the model
+    already had them. Matching one later isn't proof of a leak.
+    """
+    literals: set[str] = set()
+    for call in _COMPUTE_CODE.finditer(transcript):
+        try:
+            code = json.loads(f'"{call.group(1)}"')
+        except json.JSONDecodeError:
+            continue
+        literals.update(m.group(1) for m in _STRING_LITERAL.finditer(code))
+    return literals
+
+
+#: Values shorter than this are skipped. "Eng", "1", "true" occur in any
+#: transcript for reasons that have nothing to do with a leak, and reporting
+#: them would bury the real finding under noise.
+MIN_INTERESTING = 5
+
+
+@dataclass
+class Finding:
+    token: str
+    value: str
+    semantic_type: str | None
+
+
+@dataclass
+class Report:
+    records: int = 0
+    placeholders_seen: int = 0
+    unresolved: set[str] = field(default_factory=set)
+    leaks: list[Finding] = field(default_factory=list)
+    skipped_as_too_short: int = 0
+    #: A vault_compute record's value matched the transcript, but the match is
+    #: a literal the model wrote into its own code — not a leak, see
+    #: `_literals_the_model_already_wrote`.
+    explained: list[Finding] = field(default_factory=list)
+    compute_attempts: int = 0
+    compute_outcomes: dict[str, int] = field(default_factory=dict)
+
+    @property
+    def suspicious_compute(self) -> bool:
+        return self.compute_outcomes.get("blocked", 0) > 0
+
+    @property
+    def passed(self) -> bool:
+        """Suitable for CI: neither a value leak nor a blocked probing burst."""
+        return self.clean and not self.suspicious_compute
+
+    @property
+    def clean(self) -> bool:
+        return not self.leaks
+
+    def render(self) -> str:
+        lines = [
+            f"vault records            : {self.records}",
+            f"placeholders in transcript: {self.placeholders_seen}",
+        ]
+        if self.compute_attempts:
+            outcomes = ", ".join(
+                f"{name}={count}" for name, count in sorted(self.compute_outcomes.items())
+            )
+            lines.append(f"secret compute attempts  : {self.compute_attempts} ({outcomes})")
+        if self.suspicious_compute:
+            lines.append(
+                "SUSPICIOUS COMPUTE       : the lineage rate limit blocked "
+                f"{self.compute_outcomes['blocked']} attempt(s)"
+            )
+        if self.unresolved:
+            lines.append(
+                f"placeholders that do not resolve: {len(self.unresolved)} "
+                f"(expired, or from another session)"
+            )
+        if self.skipped_as_too_short:
+            lines.append(
+                f"values too short to check : {self.skipped_as_too_short} "
+                f"(under {MIN_INTERESTING} characters, would match anything)"
+            )
+        if self.explained:
+            lines.append(
+                f"matched but explained     : {len(self.explained)} "
+                f"(the model wrote this text itself, in compute code — not a leak)"
+            )
+        if self.leaks:
+            lines.append("")
+            lines.append(f"LEAKED — {len(self.leaks)} hidden value(s) found in the transcript:")
+            for f in self.leaks:
+                what = f" ({f.semantic_type})" if f.semantic_type else ""
+                lines.append(f"  {f.token}{what} -> {f.value[:60]}")
+        elif self.records == 0:
+            lines.append("")
+            lines.append(
+                "No vault records. Either nothing ran yet, or nothing was declared "
+                "— an empty config leaks nothing and protects nothing."
+            )
+        elif self.placeholders_seen == 0:
+            lines.append("")
+            lines.append(
+                "Vault has records but the transcript has no placeholders. Different "
+                "session, or the wrong transcript."
+            )
+        else:
+            lines.append("")
+            lines.append("No hidden value appears in the transcript.")
+        return "\n".join(lines)
+
+
+def _searchable(value: Any) -> Iterator[str]:
+    """Every scalar inside a value, as text — a table hides one per cell."""
+    if isinstance(value, bool) or value is None:
+        return
+    if isinstance(value, (int, float)):
+        yield str(value)
+    elif isinstance(value, str):
+        yield value
+    elif isinstance(value, dict):
+        for v in value.values():
+            yield from _searchable(v)
+    elif isinstance(value, list):
+        for v in value:
+            yield from _searchable(v)
+
+
+def audit(transcript: str, store: TokenStore, session_id: str) -> Report:
+    report = Report()
+    seen = set(TOKEN_PATTERN.findall(transcript))
+    report.placeholders_seen = len(seen)
+
+    records = store.find_by_session(session_id)
+    report.records = len(records)
+    attempts = store.find_compute_attempts(session_id)
+    report.compute_attempts = len(attempts)
+    report.compute_outcomes = dict(Counter(attempt.outcome for attempt in attempts))
+
+    for token in seen:
+        if store.get(token) is None:
+            report.unresolved.add(token)
+
+    literals = _literals_the_model_already_wrote(transcript)
+
+    for record in records:
+        for text in _searchable(record.value):
+            if len(text) < MIN_INTERESTING:
+                report.skipped_as_too_short += 1
+                continue
+            if text not in transcript:
+                continue
+            finding = Finding(token=record.token, value=text, semantic_type=record.semantic_type)
+            if record.lineage.op == "vault_compute" and text in literals:
+                # The model typed this itself; it was never pulled out of the
+                # vault. Keep checking the record's other cells rather than
+                # stopping here — an explained match on one cell must not hide
+                # a genuine leak on another.
+                report.explained.append(finding)
+                continue
+            report.leaks.append(finding)
+            break
+    return report
+
+
+def session_ids_in(path: Path) -> list[str]:
+    """The session ids a Claude Code transcript belongs to.
+
+    Saves the caller from knowing one: the file names its own session, and that
+    is the id the hooks minted under.
+    """
+    found: list[str] = []
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        try:
+            sid = json.loads(line).get("sessionId")
+        except (json.JSONDecodeError, AttributeError):
+            continue
+        if isinstance(sid, str) and sid not in found:
+            found.append(sid)
+    return found
+
+
+def read_transcript(path: Path) -> str:
+    """The transcript as one blob.
+
+    Claude Code writes JSONL; anything else is read as plain text, so this also
+    works on a log you captured yourself.
+    """
+    raw = path.read_text(encoding="utf-8", errors="replace")
+    if path.suffix != ".jsonl":
+        return raw
+    # Re-serialising each line normalises escaping, so a placeholder written as
+    # ⟦ in the file is found the same as one written literally.
+    out = []
+    for line in raw.splitlines():
+        if not line.strip():
+            continue
+        try:
+            out.append(json.dumps(json.loads(line), ensure_ascii=False))
+        except json.JSONDecodeError:
+            out.append(line)
+    return "\n".join(out)

vaultcompute/cli.py

@@ -0,0 +1,224 @@
+"""VaultCompute CLI entry point."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import sys
+from pathlib import Path
+
+from vaultcompute import audit as audit_mod
+from vaultcompute import hooks, mcp_server
+from vaultcompute.config import VaultComputeConfig, build_token_store, load_config
+from vaultcompute.core.policy import SessionBoundPolicy
+from vaultcompute.hosts import CLAUDE_CODE, HOSTS
+from vaultcompute.proxy import run_proxy
+
+USAGE = (
+    "vaultcompute [--config PATH] -- <downstream-mcp-command> [args...]\n"
+    "vaultcompute hook <event> [--host claude-code|codex] [--config PATH]\n"
+    "vaultcompute mcp-server [--config PATH]\n"
+    "vaultcompute audit <transcript> [--session ID] [--config PATH]\n\n"
+    "Wraps the given stdio MCP server, tokenizing declared JSON tool results\n"
+    "and exposing configured operation tools + vaultcompute/rehydrate. Strict\n"
+    "protocol handling is the default; arbitrary Python is opt-in.\n\n"
+    "`hook` reads one Claude Code or Codex hook event as JSON on stdin and writes the\n"
+    "hook response on stdout. `mcp-server` exposes the operation tools selected\n"
+    "by compute.mode. Both need a shared vault\n"
+    "(storage.backend: sqlite): they run as separate processes.\n\n"
+    "`audit` checks a transcript against the vault and reports whether any\n"
+    "hidden value reached the model — the question the screen cannot answer."
+)
+
+
+def _split_argv(argv: list[str]) -> tuple[list[str], list[str]]:
+    if "--" not in argv:
+        return argv, []
+    idx = argv.index("--")
+    return argv[:idx], argv[idx + 1 :]
+
+
+def main(argv: list[str] | None = None) -> int:
+    # Windows' default console/subprocess codepage cannot represent the token
+    # delimiters U+27E6/U+27E7, and nothing upstream sets PYTHONIOENCODING for
+    # a `uv tool install`-ed binary a host invokes as a bare command name —
+    # which is exactly how the plugin's hooks and MCP server run it. Every
+    # subcommand below reads or prints through text-mode stdio, so this forces
+    # UTF-8 once, here, rather than corrupting or crashing at the first
+    # placeholder touched. Harmless to the proxy's own JSON-RPC path, which
+    # reads and writes raw bytes via .buffer and never goes through this text
+    # layer at all.
+    #
+    # stdin needs the same treatment as stdout/stderr, and the failure mode is
+    # worse: reading a mis-decoded ⟦ doesn't raise, it silently produces a
+    # different character. `hook message-display` still parses as valid JSON
+    # and returns cleanly — TOKEN_PATTERN just never matches, so the hook
+    # answers "nothing to do" and the host displays the raw placeholder
+    # forever. Confirmed by running the exact same call with and without
+    # PYTHONIOENCODING set: identical event, one resolves the token, the other
+    # produces no output and no error.
+    for stream in (sys.stdin, sys.stdout, sys.stderr):
+        if hasattr(stream, "reconfigure") and (stream.encoding or "").lower() != "utf-8":
+            stream.reconfigure(encoding="utf-8")
+
+    argv = list(sys.argv[1:] if argv is None else argv)
+
+    if argv and argv[0] == "hook":
+        return run_hook(argv[1:])
+
+    if argv and argv[0] == "mcp-server":
+        return run_mcp_server(argv[1:])
+
+    if argv and argv[0] == "audit":
+        return run_audit(argv[1:])
+
+    own, downstream = _split_argv(argv)
+    parser = argparse.ArgumentParser(
+        prog="vaultcompute",
+        description="Privacy proxy for stdio MCP servers.",
+        usage=USAGE,
+    )
+    parser.add_argument("--config", type=Path, default=None, help="Path to vaultcompute.yaml")
+    args = parser.parse_args(own)
+
+    if not downstream:
+        print("error: no downstream command after `--`", file=sys.stderr)
+        print(USAGE, file=sys.stderr)
+        return 2
+
+    asyncio.run(run_proxy(downstream, config_path=args.config))
+    return 0
+
+
+def run_hook(argv: list[str]) -> int:
+    """Handle one host lifecycle event.
+
+    Failure is asymmetric on purpose. Printing nothing tells the host to keep
+    what it had: for ``PostToolUse`` that is the untokenized tool result on its
+    way to the model, so anything going wrong there has to block instead. For
+    ``MessageDisplay`` the worst case is the user reading a placeholder, which
+    is a nuisance and not a disclosure, so that one stays quiet.
+    """
+    parser = argparse.ArgumentParser(prog="vaultcompute hook", usage=USAGE)
+    parser.add_argument("event")
+    parser.add_argument("--host", choices=list(HOSTS), default=CLAUDE_CODE)
+    parser.add_argument("--config", type=Path, default=Path("vaultcompute.yaml"))
+    args = parser.parse_args(argv)
+
+    if args.event not in hooks.events_for(args.host):
+        parser.error(
+            f"event {args.event!r} is not available for {args.host}; "
+            f"choose from {', '.join(hooks.events_for(args.host))}"
+        )
+
+    protects = args.event in (hooks.PRE_TOOL_USE, hooks.POST_TOOL_USE)
+
+    def fail(message: str) -> int:
+        print(f"[vaultcompute] {message}", file=sys.stderr)
+        if protects:
+            response = hooks.failure_response(args.host, args.event, message)
+            if response is not None:
+                print(json.dumps(response))
+        return 0
+
+    try:
+        config = load_config(args.config) if args.config.exists() else VaultComputeConfig()
+    except Exception as exc:  # a broken config must not silently disable protection
+        return fail(f"could not load {args.config}: {type(exc).__name__}")
+
+    if config.storage.backend == "memory":
+        return fail(
+            "hooks need storage.backend: sqlite — each hook run is a separate process, "
+            "and a memory vault would not survive between tokenizing and displaying"
+        )
+
+    try:
+        event = hooks.read_event(sys.stdin.read())
+    except Exception as exc:
+        return fail(f"unreadable hook input: {type(exc).__name__}")
+
+    try:
+        store = build_token_store(config)
+    except Exception as exc:
+        return fail(f"could not open vault: {type(exc).__name__}")
+    try:
+        response = hooks.dispatch(
+            args.event,
+            event,
+            config=config,
+            store=store,
+            policy=SessionBoundPolicy(),
+            host=args.host,
+        )
+    except Exception as exc:
+        # Only the exception type, never its message: the same reasoning as the
+        # compute sandbox, since a message can carry a value it touched.
+        return fail(f"{args.event} failed: {type(exc).__name__}")
+    finally:
+        close = getattr(store, "close", None)
+        if close is not None:
+            close()
+
+    if response is not None:
+        print(json.dumps(response, ensure_ascii=False))
+    return 0
+
+
+def run_mcp_server(argv: list[str]) -> int:
+    """Serve vault_compute over stdio MCP."""
+    parser = argparse.ArgumentParser(prog="vaultcompute mcp-server", usage=USAGE)
+    parser.add_argument("--config", type=Path, default=Path("vaultcompute.yaml"))
+    args = parser.parse_args(argv)
+
+    try:
+        mcp_server.run(args.config)
+    except mcp_server.SharedVaultRequired as exc:
+        print(f"[vaultcompute] {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+def run_audit(argv: list[str]) -> int:
+    """Answer "is it actually working" from the transcript, not the screen.
+
+    Exits non-zero when a hidden value or blocked probing burst is found, so it
+    can gate a pipeline.
+    """
+    parser = argparse.ArgumentParser(prog="vaultcompute audit", usage=USAGE)
+    parser.add_argument("transcript", type=Path)
+    parser.add_argument("--session", default=None, help="defaults to the transcript's own")
+    parser.add_argument("--config", type=Path, default=Path("vaultcompute.yaml"))
+    args = parser.parse_args(argv)
+
+    if not args.transcript.exists():
+        print(f"[vaultcompute] no such transcript: {args.transcript}", file=sys.stderr)
+        return 2
+
+    config = load_config(args.config) if args.config.exists() else VaultComputeConfig()
+    store = build_token_store(config)
+    try:
+        text = audit_mod.read_transcript(args.transcript)
+        sessions = [args.session] if args.session else audit_mod.session_ids_in(args.transcript)
+        if not sessions:
+            print(
+                "[vaultcompute] the transcript names no session; pass --session",
+                file=sys.stderr,
+            )
+            return 2
+        clean = True
+        for session in sessions:
+            if len(sessions) > 1:
+                print(f"\n--- session {session} ---")
+            report = audit_mod.audit(text, store, session)
+            print(report.render())
+            clean = clean and report.passed
+    finally:
+        close = getattr(store, "close", None)
+        if close is not None:
+            close()
+    return 0 if clean else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())