anolisa-tokenless 0.7.13 → 0.8.0

package/adapters/tokenless/common/hooks/compress_schema_hook.py

@@ -1,9 +1,12 @@
 #!/usr/bin/env python3
 """Tokenless schema compression hook.
 
-Reads a BeforeModel JSON from stdin, extracts the tools array,
-invokes ``tokenless compress-schema --batch`` via subprocess, and
+Reads a BeforeModel JSON from stdin, extracts the tools array, forwards it
+to the unified ``tokenless compress`` Protocol v2 BeforeModel operation and
 writes a HookOutput JSON to stdout.
+The entry point returns the original array on no-savings, which this hook
+wraps exactly like a compressed one — the historical behavior of the
+``compress-schema`` flow it replaces.
 
 Hook point: **BeforeModel**
 
@@ -16,7 +19,6 @@ from __future__ import annotations
 import contextlib
 import json
 import os
-import subprocess
 import sys
 
 try:  # POSIX hosts (cosh / Cosh-NG) — the platforms these hooks target.
@@ -30,9 +32,11 @@ from hook_utils import (
     _TOKENLESS_FALLBACK,
     _TOKENLESS_LOCAL_LIB,
     _TOKENLESS_LOCAL_SHARE,
+    build_before_model_request,
     resolve_agent_id,
     resolve_binary,
     resolve_tool_call_id,
+    run_compress,
     secure_write_text,
     skip,
     warn,
@@ -42,6 +46,11 @@ from hook_utils import (
 
 _AGENT_ID = resolve_agent_id()
 
+# Below the extension manifests' 10 s host wrapper, so a pathological batch
+# is killed here (fail-open skip) instead of racing the host's kill of the
+# whole hook (the old subprocess timeout was 10 s against a 10 s wrapper).
+_COMPRESS_TIMEOUT = 8
+
 # One marker file holds the session keys that already emitted the "no tool
 # declarations" warning — one key per line, most recent last — so the warning
 # repeats at most once per session even though BeforeModel fires on every
@@ -103,14 +112,6 @@ def _marker_lock():
 # -- helpers -----------------------------------------------------------------
 
 
-def _is_json_array(data: str) -> bool:
-    try:
-        obj = json.loads(data)
-        return isinstance(obj, list)
-    except (json.JSONDecodeError, ValueError):
-        return False
-
-
 def _session_warn_key(session_id: str) -> str:
     """Normalize the session ID into the dedup key for the no-tools warning.
 
@@ -266,42 +267,33 @@ def main() -> None:
                 )
         skip()
 
-    tools_json = json.dumps(tools, separators=(",", ":"))
-
     # 4. Extract caller context
     session_id = input_data.get("session_id", "")
     tool_use_id = resolve_tool_call_id(_AGENT_ID, input_data)
 
-    # 5. Compress schemas via tokenless compress-schema --batch
-    cmd = [tokenless_bin, "compress-schema", "--batch", "--agent-id", _AGENT_ID]
-    if session_id:
-        cmd.extend(["--session-id", session_id])
-    if tool_use_id:
-        cmd.extend(["--tool-use-id", tool_use_id])
-
-    try:
-        proc = subprocess.run(
-            cmd,
-            input=tools_json,
-            capture_output=True,
-            text=True,
-            timeout=10,
-        )
-    except Exception:
+    # 5. Compress schemas via the unified entry point (one subprocess).
+    llm_request = input_data.get("llm_request", {})
+    visible_context = dict(llm_request) if isinstance(llm_request, dict) else {}
+    visible_context.pop("tools", None)
+    config = visible_context.get("config")
+    if isinstance(config, dict):
+        visible_context["config"] = dict(config)
+        visible_context["config"].pop("tools", None)
+    request = build_before_model_request(
+        tools,
+        visible_context,
+        _AGENT_ID,
+        session_id=session_id,
+        tool_use_id=tool_use_id,
+    )
+    response = run_compress(
+        tokenless_bin, request, _COMPRESS_TIMEOUT, "before_model"
+    )
+    if response is None:
         warn("Schema compression subprocess failed. Passing through unchanged.")
         skip()
-
-    if proc.returncode != 0:
-        detail = (proc.stderr or "").strip()[:200]
-        warn(
-            f"Schema compression failed with exit code {proc.returncode}: {detail}"
-            if detail
-            else f"Schema compression failed with exit code {proc.returncode}. Passing through unchanged."
-        )
-        skip()
-
-    compressed = proc.stdout.strip()
-    if not compressed or not _is_json_array(compressed):
+    compressed = response.get("tools")
+    if not isinstance(compressed, list):
         warn(
             "Schema compression returned invalid JSON. Passing through unchanged."
         )
@@ -313,7 +305,7 @@ def main() -> None:
             "hookEventName": "BeforeModel",
             "llm_request": {
                 "config": {
-                    "tools": json.loads(compressed),
+                    "tools": compressed,
                 },
             },
         },

package/adapters/tokenless/common/hooks/hook_utils.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
+import hashlib
 import json
 import os
 import re
 import shutil
 import subprocess
 import sys
+import time
 
 # -- Binary fallback paths ----------------------------------------------------
 #
@@ -35,16 +37,10 @@ def _user_path(*parts: str) -> str:
 
 
 _TOKENLESS_FALLBACK = "/usr/bin/tokenless"
-_TOKENLESS_LOCAL_SHARE = _user_path(
-    ".local", "share", "anolisa", "tokenless", "tokenless"
-)
-_TOKENLESS_LOCAL_LIB = _user_path(
-    ".local", "lib", "anolisa", "tokenless", "tokenless"
-)
+_TOKENLESS_LOCAL_SHARE = _user_path(".local", "share", "anolisa", "tokenless", "tokenless")
+_TOKENLESS_LOCAL_LIB = _user_path(".local", "lib", "anolisa", "tokenless", "tokenless")
 _RTK_FALLBACK = "/usr/libexec/anolisa/tokenless/rtk"
-_RTK_LOCAL_SHARE = _user_path(
-    ".local", "share", "anolisa", "tokenless", "rtk"
-)
+_RTK_LOCAL_SHARE = _user_path(".local", "share", "anolisa", "tokenless", "rtk")
 _RTK_LOCAL_LIB = _user_path(".local", "lib", "anolisa", "tokenless", "rtk")
 
 _TOKENLESS_HELPER_BINARIES = frozenset({"rtk"})
@@ -87,9 +83,7 @@ def _known_binary_paths(name: str, home: str | None = None) -> tuple[str, ...]:
     paths.append(os.path.join("/usr/local/bin", name))
     if name in _TOKENLESS_HELPER_BINARIES:
         # Anolisa CLI system mode.
-        paths.append(
-            os.path.join("/usr/local/libexec/anolisa/tokenless", name)
-        )
+        paths.append(os.path.join("/usr/local/libexec/anolisa/tokenless", name))
     paths.append(os.path.join("/usr/bin", name))
     if name in _TOKENLESS_HELPER_BINARIES:
         paths.extend(
@@ -103,16 +97,13 @@ def _known_binary_paths(name: str, home: str | None = None) -> tuple[str, ...]:
         if user_home:
             paths.extend(
                 [
-                    os.path.join(
-                        user_home, ".local", "share", "anolisa", "tokenless", name
-                    ),
-                    os.path.join(
-                        user_home, ".local", "lib", "anolisa", "tokenless", name
-                    ),
+                    os.path.join(user_home, ".local", "share", "anolisa", "tokenless", name),
+                    os.path.join(user_home, ".local", "lib", "anolisa", "tokenless", name),
                 ]
             )
     return tuple(paths)
 
+
 # -- Unified tool categorization ----------------------------------------------
 
 # Tool categories are loaded from tool_categories.json, which serves as the
@@ -126,16 +117,38 @@ _TOOL_CATEGORIES_PATH = os.path.join(os.path.dirname(__file__), "tool_categories
 # invalid. Matches the minimum safe classification from before the JSON was
 # introduced, ensuring content-retrieval tools are never accidentally compressed.
 _FALLBACK_SKIP_TOOLS = [
-    "Read", "read", "read_file", "read_many_files",
-    "Glob", "glob", "search_file", "list_directory", "list_dir",
-    "Grep", "grep", "grep_code", "grep_search", "search_files",
-    "Lsp", "lsp",
-    "NotebookRead", "notebook_read", "notebookread",
+    "Read",
+    "read",
+    "read_file",
+    "read_many_files",
+    "Glob",
+    "glob",
+    "search_file",
+    "list_directory",
+    "list_dir",
+    "Grep",
+    "grep",
+    "grep_code",
+    "grep_search",
+    "search_files",
+    "Lsp",
+    "lsp",
+    "NotebookRead",
+    "notebook_read",
+    "notebookread",
 ]
 _FALLBACK_SHELL_TOOLS = [
-    "Bash", "bash", "Shell", "shell", "exec", "terminal",
-    "run_shell_command", "run_in_terminal", "get_terminal_output",
-    "execute_command", "process",
+    "Bash",
+    "bash",
+    "Shell",
+    "shell",
+    "exec",
+    "terminal",
+    "run_shell_command",
+    "run_in_terminal",
+    "get_terminal_output",
+    "execute_command",
+    "process",
 ]
 
 
@@ -187,6 +200,29 @@ SKIP_TOOLS: set[str] = set(_tool_categories.get("layer_1_skip", {}).get("tools",
 # These tools produce text output that can be safely truncated if too long.
 SHELL_TOOLS: set[str] = set(_tool_categories.get("layer_2_shell", {}).get("tools", []))
 
+_TOKENLESS_RETRIEVE_COMMAND_RE = re.compile(
+    r"^[ \t]*(?:\"tokenless\"|'tokenless'|tokenless)[ \t]+retrieve[ \t]+"
+    r"(?:\"(?:[0-9a-f]{24}|<<tokenless:[0-9a-f]{24}>>)\"|"
+    r"'(?:[0-9a-f]{24}|<<tokenless:[0-9a-f]{24}>>)'|[0-9a-f]{24})[ \t]*$",
+    re.IGNORECASE,
+)
+
+
+def tokenless_retrieve_command_available() -> bool:
+    """Return whether a Marker command can invoke bare ``tokenless``."""
+    return shutil.which("tokenless") is not None
+
+
+def is_tokenless_retrieve_command(tool_name: str, arguments: object) -> bool:
+    """Recognize the exact local recovery command emitted by Tokenless markers."""
+    if tool_name not in SHELL_TOOLS or not isinstance(arguments, dict):
+        return False
+    command = arguments.get("command")
+    if not isinstance(command, str):
+        return False
+    return _TOKENLESS_RETRIEVE_COMMAND_RE.fullmatch(command) is not None
+
+
 # Layer 3: API tools (zero-truncation).
 # These tools return structured data or API responses that should not be truncated.
 # No explicit set needed; tools not in SKIP_TOOLS or SHELL_TOOLS are Layer 3.
@@ -196,7 +232,7 @@ SHELL_TOOLS: set[str] = set(_tool_categories.get("layer_2_shell", {}).get("tools
 # is missing or the file failed to load.
 
 # Layer 2 thresholds: moderate truncation for shell/exec output.
-# Restores old ResponseCompressor defaults for shell commands (git log, ls,
+# Restores the old JSON compression defaults for shell commands (git log, ls,
 # cat, etc.) where truncation is acceptable.
 # 64K strings: 95% of real shell output (git diff ~63K, git log ~34K) preserved.
 # 128 arrays: 95% of result sets (test results, audit reports) preserved.
@@ -211,11 +247,6 @@ _TRUNCATE_STRINGS_AT = _layer3_thr.get("truncate_strings_at", 1_048_576)
 _TRUNCATE_ARRAYS_AT = _layer3_thr.get("truncate_arrays_at", 65_536)
 _MAX_DEPTH = _layer3_thr.get("max_depth", 32)
 
-# Backward-compatible alias — direct reference (not a copy) so consumers see
-# the same set as SKIP_TOOLS. Used by compress_toon_hook.py for the standalone
-# TOON-only path where "content retrieval" is the more descriptive name.
-CONTENT_RETRIEVAL_TOOLS = SKIP_TOOLS
-
 
 def get_thresholds(tool_name: str) -> tuple[int, int, int]:
     """Return (truncate_strings_at, truncate_arrays_at, max_depth) for a tool.
@@ -231,9 +262,9 @@ def get_thresholds(tool_name: str) -> tuple[int, int, int]:
 
 # -- Shared environment error patterns ----------------------------------------
 #
-# Superset of patterns from both the codex and hook adapters. Uses regex
+# Superset of patterns from both the Codex diagnostics and shared hook adapters. Uses regex
 # matching (case-insensitive) so patterns like "/bin/sh:.*: not found" work
-# correctly. Both codex/scripts/compress-response and compress_response_hook
+# correctly. Both codex/scripts/response-diagnostics and compress_response_hook
 # import this list and the classify_env_error() function below.
 
 ENV_PATTERNS: list[tuple[list[str], str, str]] = [
@@ -315,7 +346,7 @@ def classify_env_error(tool_response) -> tuple[str | None, str | None]:
     Accepts either a parsed dict (with stderr/error/exit_code fields) or a
     plain string. Returns (category_tag, fix_hint) or (None, None).
 
-    Shared by codex/scripts/compress-response and compress_response_hook.
+    Shared by codex/scripts/response-diagnostics and compress_response_hook.
     """
     if isinstance(tool_response, dict):
         text = str(tool_response.get("stderr", "")) + str(tool_response.get("error", ""))
@@ -346,6 +377,9 @@ def classify_env_error(tool_response) -> tuple[str | None, str | None]:
 
 _CONTEXT_DIR = os.path.join(os.path.expanduser("~"), ".tokenless")
 _CONTEXT_FILE = os.path.join(_CONTEXT_DIR, ".rewrite-context")
+_OPTIMIZATION_STATE_DIR = os.path.join(_CONTEXT_DIR, "hook-state")
+_OPTIMIZATION_STATE_TTL_SECONDS = 24 * 60 * 60
+_OPTIMIZATION_STATE_MAX_FILES = 1024
 
 # -- Binary resolution (cached) -----------------------------------------------
 
@@ -412,9 +446,7 @@ def unwrap_string_json(raw: str) -> str | None:
             # characters (code points), not \uXXXX escape sequences, so
             # string-wrapped payloads are measured the same way as the
             # dict/list branch and the OpenClaw adapter.
-            return json.dumps(
-                inner_obj, separators=(",", ":"), ensure_ascii=False
-            )
+            return json.dumps(inner_obj, separators=(",", ":"), ensure_ascii=False)
         return None
     return raw
 
@@ -474,6 +506,77 @@ def write_context(agent_id: str, session_id: str, tool_use_id: str) -> None:
     secure_write_text(_CONTEXT_FILE, f"{agent_id}\n{session_id}\n{tool_use_id}\n")
 
 
+def _optimization_state_path(agent_id: str, session_id: str, tool_use_id: str) -> str:
+    identity = "\0".join((agent_id, session_id, tool_use_id)).encode()
+    digest = hashlib.sha256(identity).hexdigest()
+    return os.path.join(_OPTIMIZATION_STATE_DIR, digest)
+
+
+def _prune_optimization_states() -> None:
+    """Bound abandoned per-call state when a host omits PostToolUse."""
+    try:
+        entries = os.scandir(_OPTIMIZATION_STATE_DIR)
+    except FileNotFoundError:
+        return
+
+    cutoff = time.time() - _OPTIMIZATION_STATE_TTL_SECONDS
+    live = []
+    with entries:
+        for entry in entries:
+            try:
+                if not entry.is_file(follow_symlinks=False):
+                    continue
+                modified = entry.stat(follow_symlinks=False).st_mtime
+                if ".consuming." in entry.name and modified > cutoff:
+                    continue
+                if modified <= cutoff:
+                    os.unlink(entry.path)
+                else:
+                    live.append((modified, entry.path))
+            except FileNotFoundError:
+                continue
+
+    excess = len(live) - _OPTIMIZATION_STATE_MAX_FILES + 1
+    if excess > 0:
+        for _, path in sorted(live)[:excess]:
+            try:
+                os.unlink(path)
+            except FileNotFoundError:
+                pass
+
+
+def mark_rtk_optimized(agent_id: str, session_id: str, tool_use_id: str) -> None:
+    """Persist RTK ownership for one tool call before applying its rewrite."""
+    _prune_optimization_states()
+    secure_write_text(_optimization_state_path(agent_id, session_id, tool_use_id), "rtk\n")
+
+
+def consume_output_optimization(agent_id: str, session_id: str, tool_use_id: str) -> str:
+    """Consume one tool call's optimization state for its final result."""
+    if not tool_use_id:
+        return "none"
+    path = _optimization_state_path(agent_id, session_id, tool_use_id)
+    consuming_path = f"{path}.consuming.{os.getpid()}"
+    try:
+        os.rename(path, consuming_path)
+    except FileNotFoundError:
+        return "none"
+
+    try:
+        flags = os.O_RDONLY
+        if hasattr(os, "O_NOFOLLOW"):
+            flags |= os.O_NOFOLLOW
+        fd = os.open(consuming_path, flags)
+        with os.fdopen(fd) as state_file:
+            state = state_file.read()
+    finally:
+        try:
+            os.unlink(consuming_path)
+        except FileNotFoundError:
+            pass
+    return "rtk" if state == "rtk\n" else "none"
+
+
 def forward_stderr(proc: subprocess.CompletedProcess) -> None:
     """Forward subprocess stderr on failure (non-zero exit) via warn()."""
     if proc.returncode != 0 and proc.stderr:
@@ -496,6 +599,144 @@ def run(args: list[str], input_data: str, timeout: int = 3) -> subprocess.Comple
         return None
 
 
+def _attribution(
+    agent_id: str,
+    session_id: str = "",
+    tool_use_id: str = "",
+) -> dict:
+    """Build the shared Protocol v2 attribution object."""
+    value = {"agent_id": agent_id}
+    if session_id:
+        value["session_id"] = session_id
+    if tool_use_id:
+        value["tool_use_id"] = tool_use_id
+    return value
+
+
+def build_before_model_request(
+    tools: list,
+    visible_context: object,
+    agent_id: str,
+    session_id: str = "",
+    tool_use_id: str = "",
+) -> dict:
+    """Build a Protocol v2 BeforeModel transport request."""
+    return {
+        "protocol_version": 2,
+        "operation": "before_model",
+        "attribution": _attribution(agent_id, session_id, tool_use_id),
+        "input": {
+            "tools": tools,
+            "visible_context": visible_context,
+            "capabilities": {
+                "replace_tools": True,
+                # A local CLI is not marker-scoped Agent authorization.
+                "recovery": {"kind": "none"},
+            },
+        },
+    }
+
+
+def build_pre_tool_request(
+    arguments: dict,
+    agent_id: str,
+    tool_name: str,
+    command_field: str,
+    session_id: str = "",
+    tool_use_id: str = "",
+    replace_arguments: bool = True,
+    block_and_suggest: bool = False,
+) -> dict:
+    """Build a Protocol v2 PreTool transport request."""
+    return {
+        "protocol_version": 2,
+        "operation": "pre_tool",
+        "attribution": _attribution(agent_id, session_id, tool_use_id),
+        "input": {
+            "tool_name": tool_name,
+            "arguments": arguments,
+            "command_field": command_field,
+            "capabilities": {
+                "replace_arguments": replace_arguments,
+                "block_and_suggest": block_and_suggest,
+            },
+        },
+    }
+
+
+def build_post_tool_request(
+    content: str,
+    agent_id: str,
+    tool_name: str,
+    status: str,
+    content_origin: str,
+    output_optimization: str,
+    *,
+    result_kind: str,
+    recovery: dict[str, str],
+    session_id: str = "",
+    tool_use_id: str = "",
+    replace_output: bool = False,
+    replace_with_text: bool = False,
+) -> dict:
+    """Build a Protocol v2 PostTool transport request."""
+    return {
+        "protocol_version": 2,
+        "operation": "post_tool",
+        "attribution": _attribution(agent_id, session_id, tool_use_id),
+        "input": {
+            "result_kind": result_kind,
+            "tool_name": tool_name,
+            "content": content,
+            "status": status,
+            "content_origin": content_origin,
+            "output_optimization": output_optimization,
+            "capabilities": {
+                "replace_output": replace_output,
+                "recovery": recovery,
+                "replace_with_text": replace_with_text,
+            },
+        },
+    }
+
+
+def run_compress(
+    tokenless_bin: str, request: dict, timeout: int, expected_operation: str
+) -> dict | None:
+    """Run ``tokenless compress`` on one request; None on any failure.
+
+    The single Tokenless subprocess of a hook invocation (roadmap §5.6).
+    Fail-open: a dead binary, non-zero exit, or malformed stdout all
+    return None so the caller passes the original through.
+    """
+    proc = run(
+        [tokenless_bin, "compress"],
+        json.dumps(request, ensure_ascii=False),
+        timeout=timeout,
+    )
+    if proc is None or proc.returncode != 0 or not proc.stdout.strip():
+        if proc is not None and proc.returncode != 0:
+            warn(f"tokenless compress exited {proc.returncode}")
+        return None
+    response = try_parse_json(proc.stdout.strip())
+    if not isinstance(response, dict):
+        warn("tokenless compress returned malformed output")
+        return None
+    if response.get("protocol_version") != 2:
+        # Version-skewed binary: never trust a response whose contract
+        # this adapter does not speak.
+        warn("tokenless compress returned an unsupported protocol version")
+        return None
+    if response.get("operation") != expected_operation:
+        warn("tokenless compress returned a mismatched operation")
+        return None
+    result = response.get("result")
+    if not isinstance(result, dict):
+        warn("tokenless compress returned a malformed operation result")
+        return None
+    return result
+
+
 def detect_cosh_ng_runtime() -> tuple | None:
     """Detect if we are running under Cosh-NG and return its version.
 
@@ -547,7 +788,7 @@ def _agent_id_from_argv(argv: list[str] | None) -> str | None:
             following = args[index + 1] if index + 1 < len(args) else ""
             return following or None
         if arg.startswith("--agent-id="):
-            return arg[len("--agent-id="):] or None
+            return arg[len("--agent-id=") :] or None
     return None
 
 
@@ -575,11 +816,7 @@ def resolve_agent_id(default: str = "unknown", argv: list[str] | None = None) ->
     """
     if detect_cosh_ng_runtime() is not None:
         return "cosh-ng"
-    return (
-        _agent_id_from_argv(argv)
-        or os.environ.get("TOKENLESS_AGENT_ID")
-        or default
-    )
+    return _agent_id_from_argv(argv) or os.environ.get("TOKENLESS_AGENT_ID") or default
 
 
 def parse_version(version_str: str) -> tuple | None: