@ictechgy/context-guard 0.4.10 → 0.4.12

package/plugins/context-guard/bin/context-guard-tool-prune

@@ -38,6 +38,11 @@ MAX_TOP = 200
 MAX_DEFERRED_TOP = 1_000
 MAX_NAMESPACE_TOP = 200
 MAX_LABEL_CHARS = 160
+NO_FOLLOW_SUPPORTED = hasattr(os, "O_NOFOLLOW")
+DIR_FD_OPEN_SUPPORTED = bool(os.supports_dir_fd and os.open in os.supports_dir_fd)
+DIR_FD_MKDIR_SUPPORTED = bool(os.supports_dir_fd and os.mkdir in os.supports_dir_fd)
+DIR_FD_STAT_SUPPORTED = bool(os.supports_dir_fd and os.stat in os.supports_dir_fd)
+DIR_FD_UNLINK_SUPPORTED = bool(os.supports_dir_fd and os.unlink in os.supports_dir_fd)
 MAX_DESCRIPTION_CHARS = 360
 MAX_OMITTED_TOOLS = 30
 TOKEN_PROXY_CHARS_PER_TOKEN = 4
@@ -82,6 +87,8 @@ class Candidate:
     index: int
     score: float = 0.0
     rank: int = 0
+    schema_bytes: int = 0
+    parameter_terms: frozenset[str] | None = None
 
 
 def fail(message: str) -> NoReturn:
@@ -201,12 +208,26 @@ def sanitize_value(value: Any, *, sensitive_context: bool = False, sensitive_sch
 
 
 def read_limited_path(path: Path, max_bytes: int) -> str:
+    if not NO_FOLLOW_SUPPORTED:
+        fail("catalog reads require O_NOFOLLOW support")
+    reject_parent_traversal(path, label="catalog")
+    # Preserve clear diagnostics for stable symlink paths, then anchor the real
+    # read to an opened no-follow parent fd so parent/leaf swaps after this
+    # precheck still fail closed.
     reject_symlink_components(path)
-    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
+    parent_fd = open_private_directory_no_follow(path.parent, label="catalog directory", create=False)
+    flags = os.O_RDONLY | os.O_NOFOLLOW
+    if hasattr(os, "O_CLOEXEC"):
+        flags |= os.O_CLOEXEC
+    leaf = path.name
+    if leaf in {"", ".", ".."}:
+        os.close(parent_fd)
+        fail("catalog must name a regular file")
     try:
-        fd = os.open(str(path), flags)
+        fd = os.open(leaf, flags, dir_fd=parent_fd)
     except OSError as exc:
-        fail(f"catalog read failed: {exc}")
+        os.close(parent_fd)
+        fail(f"catalog read failed: {os_error_detail(exc)}")
     try:
         st = os.fstat(fd)
         if not stat.S_ISREG(st.st_mode):
@@ -216,6 +237,7 @@ def read_limited_path(path: Path, max_bytes: int) -> str:
         data = os.read(fd, max_bytes + 1)
     finally:
         os.close(fd)
+        os.close(parent_fd)
     if len(data) > max_bytes:
         fail(f"catalog exceeds --max-catalog-bytes: > {max_bytes}")
     return data.decode("utf-8", errors="replace")
@@ -256,7 +278,15 @@ def tool_schema_from_dict(raw: dict[str, Any], *, fallback_name: str | None = No
         schema["description"] = description
     if server and "server" not in schema:
         schema["server"] = server
-    return Candidate(name=name, server=cap_text(server, MAX_LABEL_CHARS) if server else None, description=description, schema=schema, index=index)
+    return Candidate(
+        name=name,
+        server=cap_text(server, MAX_LABEL_CHARS) if server else None,
+        description=description,
+        schema=schema,
+        index=index,
+        schema_bytes=byte_len_json(schema),
+        parameter_terms=frozenset(terms(" ".join(collect_parameter_text(schema)))),
+    )
 
 
 def normalize_catalog(raw: Any) -> list[Candidate]:
@@ -342,7 +372,11 @@ def score_candidate(candidate: Candidate, query_terms: set[str]) -> float:
         return 0.0
     name_terms = terms(candidate.name)
     desc_terms = terms(candidate.description)
-    parameter_terms = terms(" ".join(collect_parameter_text(candidate.schema)))
+    parameter_terms = (
+        set(candidate.parameter_terms)
+        if candidate.parameter_terms is not None
+        else terms(" ".join(collect_parameter_text(candidate.schema)))
+    )
     score = 0.0
     score += 4.0 * len(query_terms & name_terms)
     score += 1.5 * len(query_terms & desc_terms)
@@ -359,14 +393,38 @@ def rank_candidates(candidates: list[Candidate], query: str) -> list[Candidate]:
     query_terms = terms(query)
     scored: list[Candidate] = []
     for cand in candidates:
-        scored.append(Candidate(cand.name, cand.server, cand.description, cand.schema, cand.index, score_candidate(cand, query_terms), 0))
+        scored.append(Candidate(
+            cand.name,
+            cand.server,
+            cand.description,
+            cand.schema,
+            cand.index,
+            score_candidate(cand, query_terms),
+            0,
+            schema_bytes=cand.schema_bytes,
+            parameter_terms=cand.parameter_terms,
+        ))
     scored.sort(key=lambda item: (-item.score, item.index))
     ranked: list[Candidate] = []
     for rank, cand in enumerate(scored, start=1):
-        ranked.append(Candidate(cand.name, cand.server, cand.description, cand.schema, cand.index, cand.score, rank))
+        ranked.append(Candidate(
+            cand.name,
+            cand.server,
+            cand.description,
+            cand.schema,
+            cand.index,
+            cand.score,
+            rank,
+            schema_bytes=cand.schema_bytes,
+            parameter_terms=cand.parameter_terms,
+        ))
     return ranked
 
 
+def candidate_schema_bytes(cand: Candidate) -> int:
+    return cand.schema_bytes if cand.schema_bytes > 0 else byte_len_json(cand.schema)
+
+
 def normalized_link_target(parent: Path, raw_target: str) -> Path:
     target = Path(raw_target)
     if not target.is_absolute():
@@ -409,15 +467,142 @@ def reject_symlink_components(path: Path) -> None:
             fail(f"refusing path through non-directory component: {current}")
 
 
+def dir_fd_replace_supported() -> bool:
+    try:
+        import inspect
+
+        signature = inspect.signature(os.replace)
+    except (TypeError, ValueError):
+        return True
+    return "src_dir_fd" in signature.parameters and "dst_dir_fd" in signature.parameters
+
+
+DIR_FD_REPLACE_SUPPORTED = dir_fd_replace_supported()
+
+
+def reject_parent_traversal(path: Path, *, label: str) -> None:
+    if ".." in path.parts:
+        fail(f"{label} must not contain parent traversal")
+
+
+def os_error_detail(exc: OSError) -> str:
+    detail = exc.strerror or str(exc) or exc.__class__.__name__
+    if exc.errno is not None:
+        return f"{detail} (errno {exc.errno})"
+    return detail
+
+
+def no_follow_dir_flags() -> int:
+    if not NO_FOLLOW_SUPPORTED:
+        fail("private store IO requires O_NOFOLLOW support")
+    flags = os.O_RDONLY | os.O_NOFOLLOW
+    if hasattr(os, "O_CLOEXEC"):
+        flags |= os.O_CLOEXEC
+    if hasattr(os, "O_DIRECTORY"):
+        flags |= os.O_DIRECTORY
+    return flags
+
+
+def private_temp_file_flags() -> int:
+    if not NO_FOLLOW_SUPPORTED:
+        fail("private store IO requires O_NOFOLLOW support")
+    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW
+    if hasattr(os, "O_CLOEXEC"):
+        flags |= os.O_CLOEXEC
+    if hasattr(os, "O_NOCTTY"):
+        flags |= os.O_NOCTTY
+    return flags
+
+
+def open_private_directory_no_follow(path: Path, *, label: str, create: bool) -> int:
+    reject_parent_traversal(path, label=label)
+    path = normalize_allowed_first_absolute_symlink(path.expanduser())
+    if not DIR_FD_OPEN_SUPPORTED:
+        fail(f"{label} requires dir_fd open support")
+    if create and not DIR_FD_MKDIR_SUPPORTED:
+        fail(f"{label} requires dir_fd mkdir support")
+    flags = no_follow_dir_flags()
+    if path.is_absolute():
+        root_flags = os.O_RDONLY | (os.O_CLOEXEC if hasattr(os, "O_CLOEXEC") else 0)
+        current_fd = os.open(path.anchor or os.sep, root_flags)
+        parts = path.parts[1:]
+    else:
+        current_fd = os.open(".", flags)
+        parts = path.parts
+    try:
+        for part in parts:
+            if part in {"", "."}:
+                continue
+            if part == "..":
+                fail(f"{label} must not contain parent traversal")
+            try:
+                next_fd = os.open(part, flags, dir_fd=current_fd)
+            except FileNotFoundError:
+                if not create:
+                    raise
+                os.mkdir(part, 0o700, dir_fd=current_fd)
+                next_fd = os.open(part, flags, dir_fd=current_fd)
+            try:
+                if not stat.S_ISDIR(os.fstat(next_fd).st_mode):
+                    fail(f"{label} must not traverse non-directory components")
+            except Exception:
+                os.close(next_fd)
+                raise
+            os.close(current_fd)
+            current_fd = next_fd
+        owned_fd = current_fd
+        current_fd = -1
+        return owned_fd
+    except OSError as exc:
+        fail(f"could not inspect {label}: {os_error_detail(exc)}")
+    finally:
+        if current_fd >= 0:
+            os.close(current_fd)
+
+
+def precheck_private_leaf(parent_fd: int, leaf: str, *, label: str) -> None:
+    if not DIR_FD_STAT_SUPPORTED:
+        fail(f"{label} requires dir_fd stat support")
+    try:
+        st = os.stat(leaf, dir_fd=parent_fd, follow_symlinks=False)
+    except FileNotFoundError:
+        return
+    except OSError as exc:
+        fail(f"could not inspect {label}: {os_error_detail(exc)}")
+    if not stat.S_ISREG(st.st_mode):
+        fail(f"{label} must be missing or a regular file")
+
+
+def write_all_fd(fd: int, data: bytes) -> None:
+    view = memoryview(data)
+    offset = 0
+    while offset < len(view):
+        written = os.write(fd, view[offset:])
+        if written <= 0:
+            raise OSError("short write")
+        offset += written
+
+
+def fsync_best_effort(fd: int) -> None:
+    try:
+        os.fsync(fd)
+    except OSError:
+        pass
+
+
 def ensure_private_dir(path: Path) -> None:
-    path = normalize_allowed_first_absolute_symlink(path)
-    reject_symlink_components(path)
+    reject_parent_traversal(path, label="store directory")
     try:
-        path.mkdir(parents=True, exist_ok=True)
-        reject_symlink_components(path)
-        os.chmod(path, 0o700)
+        fd = open_private_directory_no_follow(path, label="store directory", create=True)
     except OSError as exc:
         fail(f"store directory unavailable: {exc}")
+    try:
+        try:
+            os.fchmod(fd, 0o700)
+        except OSError:
+            pass
+    finally:
+        os.close(fd)
 
 
 def write_private_json_atomic(path: Path, data: dict[str, Any], *, max_bytes: int, label: str) -> int:
@@ -425,43 +610,77 @@ def write_private_json_atomic(path: Path, data: dict[str, Any], *, max_bytes: in
     size = byte_len_text(text)
     if size > max_bytes:
         fail(f"{label} exceeds size cap: {size} > {max_bytes}")
-    ensure_private_dir(path.parent)
-    tmp = path.with_name(path.name + f".tmp-{os.getpid()}-{time.time_ns()}")
-    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
-    try:
-        fd = os.open(str(tmp), flags, 0o600)
-    except OSError as exc:
-        fail(f"{label} write failed: {exc}")
+    reject_parent_traversal(path, label=label)
+    if not DIR_FD_REPLACE_SUPPORTED:
+        fail(f"{label} write requires dir_fd replace support")
+    if not DIR_FD_UNLINK_SUPPORTED:
+        fail(f"{label} write requires dir_fd unlink support")
+    if not DIR_FD_STAT_SUPPORTED:
+        fail(f"{label} write requires dir_fd stat support")
+    parent_fd = open_private_directory_no_follow(path.parent, label="store directory", create=True)
+    fd = -1
+    temp_leaf: str | None = None
     try:
-        with os.fdopen(fd, "w", encoding="utf-8", newline="") as handle:
-            handle.write(text)
-            handle.flush()
-            try:
-                os.fsync(handle.fileno())
-            except OSError:
-                pass
-        os.replace(tmp, path)
         try:
-            os.chmod(path, 0o600)
+            os.fchmod(parent_fd, 0o700)
         except OSError:
             pass
+        leaf = path.name
+        if leaf in {"", ".", ".."}:
+            fail(f"{label} must name a regular file")
+        precheck_private_leaf(parent_fd, leaf, label=label)
+        for _attempt in range(20):
+            candidate = f".{leaf}.{os.getpid()}.{time.time_ns()}.tmp"
+            try:
+                fd = os.open(candidate, private_temp_file_flags(), 0o600, dir_fd=parent_fd)
+                temp_leaf = candidate
+                break
+            except FileExistsError:
+                continue
+        if fd < 0 or temp_leaf is None:
+            fail(f"{label} write failed: could not create temporary file")
+        if not stat.S_ISREG(os.fstat(fd).st_mode):
+            fail(f"{label} temporary file must be a regular file")
+        os.fchmod(fd, 0o600)
+        write_all_fd(fd, text.encode("utf-8"))
+        fsync_best_effort(fd)
+        os.close(fd)
+        fd = -1
+        fsync_best_effort(parent_fd)
+        os.replace(temp_leaf, leaf, src_dir_fd=parent_fd, dst_dir_fd=parent_fd)
+        temp_leaf = None
+        fsync_best_effort(parent_fd)
+    except OSError as exc:
+        fail(f"{label} write failed: {os_error_detail(exc)}")
     except Exception:
-        try:
-            tmp.unlink()
-        except OSError:
-            pass
         raise
+    finally:
+        if fd >= 0:
+            os.close(fd)
+        if temp_leaf is not None:
+            try:
+                os.unlink(temp_leaf, dir_fd=parent_fd)
+            except OSError:
+                pass
+        os.close(parent_fd)
     return size
 
 
 def read_private_text(path: Path, *, max_bytes: int, label: str) -> tuple[str, int]:
-    if path.is_symlink():
-        fail(f"{label} must not be a symlink")
-    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
+    reject_parent_traversal(path, label=label)
+    parent_fd = open_private_directory_no_follow(path.parent, label=f"{label} directory", create=False)
+    flags = os.O_RDONLY | os.O_NOFOLLOW
+    if hasattr(os, "O_CLOEXEC"):
+        flags |= os.O_CLOEXEC
+    leaf = path.name
+    if leaf in {"", ".", ".."}:
+        os.close(parent_fd)
+        fail(f"{label} must name a regular file")
     try:
-        fd = os.open(str(path), flags)
+        fd = os.open(leaf, flags, dir_fd=parent_fd)
     except OSError as exc:
-        fail(f"{label} read failed: {exc}")
+        os.close(parent_fd)
+        fail(f"{label} read failed: {os_error_detail(exc)}")
     try:
         st = os.fstat(fd)
         if not stat.S_ISREG(st.st_mode):
@@ -471,32 +690,16 @@ def read_private_text(path: Path, *, max_bytes: int, label: str) -> tuple[str, i
         data = os.read(fd, max_bytes + 1)
     finally:
         os.close(fd)
+        os.close(parent_fd)
     if len(data) > max_bytes:
         fail(f"{label} exceeds trusted size cap: > {max_bytes}")
     return data.decode("utf-8", errors="replace"), len(data)
 
 
 def read_private_json(path: Path, *, max_bytes: int, label: str) -> dict[str, Any]:
-    if path.is_symlink():
-        fail(f"{label} must not be a symlink")
-    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
-    try:
-        fd = os.open(str(path), flags)
-    except OSError as exc:
-        fail(f"{label} read failed: {exc}")
-    try:
-        st = os.fstat(fd)
-        if not stat.S_ISREG(st.st_mode):
-            fail(f"{label} must be a regular file")
-        if st.st_size > max_bytes:
-            fail(f"{label} exceeds trusted size cap: {st.st_size} > {max_bytes}")
-        data = os.read(fd, max_bytes + 1)
-    finally:
-        os.close(fd)
-    if len(data) > max_bytes:
-        fail(f"{label} exceeds trusted size cap: > {max_bytes}")
+    text, _size = read_private_text(path, max_bytes=max_bytes, label=label)
     try:
-        parsed = json.loads(data.decode("utf-8", errors="replace"))
+        parsed = json.loads(text)
     except json.JSONDecodeError as exc:
         fail(f"{label} is malformed JSON: {exc.msg}")
     if not isinstance(parsed, dict):
@@ -542,7 +745,7 @@ def build_payload(receipt_id: str, ranked: list[Candidate], query: str, redactio
                 "description": cand.description,
                 "score": cand.score,
                 "rank": cand.rank,
-                "schema_bytes": byte_len_json(cand.schema),
+                "schema_bytes": candidate_schema_bytes(cand),
                 "schema": cand.schema,
             }
             for cand in ranked
@@ -574,7 +777,7 @@ def retrieval_command(receipt_id: str, *, store_dir: str, tool_name: str | None
 
 
 def selected_tool_record(cand: Candidate, receipt_id: str, budget_left: int, *, store_dir: str) -> tuple[dict[str, Any], int]:
-    schema_size = byte_len_json(cand.schema)
+    schema_size = candidate_schema_bytes(cand)
     record: dict[str, Any] = {
         "name": cand.name,
         "server": cand.server,
@@ -600,7 +803,7 @@ def deferred_tool_record(cand: Candidate, receipt_id: str, *, store_dir: str) ->
         "score": cand.score,
         "rank": cand.rank,
         "description": cand.description,
-        "schema_bytes": byte_len_json(cand.schema),
+        "schema_bytes": candidate_schema_bytes(cand),
         "reason": "deferred_after_core_top",
         "retrieval": retrieval_command(receipt_id, store_dir=store_dir, tool_name=cand.name),
     }
@@ -843,8 +1046,15 @@ def defer_report(args: argparse.Namespace) -> str:
         store_dir=args.store_dir,
         namespace_top=namespace_top,
     )
-    all_schema_bytes = sum(byte_len_json(cand.schema) for cand in ranked)
+    all_schema_bytes = sum(candidate_schema_bytes(cand) for cand in ranked)
+    listed_deferred_schema_bytes = sum(candidate_schema_bytes(cand) for cand in deferred_candidates)
+    total_deferred_schema_bytes = sum(candidate_schema_bytes(cand) for cand in ranked[core_top:])
     tool_stub_report_bytes = byte_len_json(core_tools) + byte_len_json(deferred_tools)
+    all_schema_tokens = proxy_tokens(all_schema_bytes)
+    inline_core_schema_tokens = proxy_tokens(core_schema_bytes)
+    listed_deferred_schema_tokens = proxy_tokens(listed_deferred_schema_bytes)
+    total_deferred_schema_tokens = proxy_tokens(total_deferred_schema_bytes)
+    tool_stub_report_tokens = proxy_tokens(tool_stub_report_bytes)
     result = {
         "tool": TOOL_NAME,
         "schema_version": DEFER_SCHEMA_VERSION,
@@ -862,6 +1072,7 @@ def defer_report(args: argparse.Namespace) -> str:
         "deferred_tools_truncated_count": max(0, len(ranked) - core_top - len(deferred_tools)),
         "deferred_namespaces": deferred_namespaces,
         "deferred_namespaces_truncated_count": deferred_namespaces_truncated_count,
+        "deferred_schema_retrieval_required_before_use": True,
         "receipt": {
             **receipt,
             "bytes": receipt_size,
@@ -871,9 +1082,21 @@ def defer_report(args: argparse.Namespace) -> str:
             "method": "char4_proxy",
             "chars_per_token": TOKEN_PROXY_CHARS_PER_TOKEN,
             "all_schema_bytes": all_schema_bytes,
+            "inline_core_schema_bytes": core_schema_bytes,
+            "listed_deferred_schema_bytes": listed_deferred_schema_bytes,
+            "total_deferred_schema_bytes": total_deferred_schema_bytes,
             "tool_stub_report_bytes": tool_stub_report_bytes,
-            "all_schema_tokens_estimated": proxy_tokens(all_schema_bytes),
-            "tool_stub_report_tokens_estimated": proxy_tokens(tool_stub_report_bytes),
+            "all_schema_tokens_estimated": all_schema_tokens,
+            "inline_core_schema_tokens_estimated": inline_core_schema_tokens,
+            "listed_deferred_schema_tokens_estimated": listed_deferred_schema_tokens,
+            "total_deferred_schema_tokens_estimated": total_deferred_schema_tokens,
+            "tool_stub_report_tokens_estimated": tool_stub_report_tokens,
+            "gross_listed_deferred_schema_tokens_avoided": listed_deferred_schema_tokens,
+            "gross_total_deferred_schema_tokens_avoided": total_deferred_schema_tokens,
+            "net_initial_report_tokens_delta": tool_stub_report_tokens - all_schema_tokens,
+            "net_initial_report_tokens_delta_semantics": "tool_stub_report_tokens_estimated_minus_all_schema_tokens_estimated",
+            "estimated_initial_schema_tokens_avoided": max(0, all_schema_tokens - tool_stub_report_tokens),
+            "estimated_initial_schema_tokens_avoided_semantics": "max(0, all_schema_tokens_estimated - tool_stub_report_tokens_estimated)",
             "claim_boundary": "proxy_only_not_provider_billed_tokens",
         },
         "provider_patterns": [
@@ -899,11 +1122,13 @@ def defer_report(args: argparse.Namespace) -> str:
             "provider_tool_search_configured": False,
             "hosted_api_token_or_cost_savings_claim_allowed": False,
             "requires_provider_measured_matched_tasks_for_savings_claims": True,
+            "deferred_schema_retrieval_required_before_use": True,
         },
         "redaction": {"redacted_values": total_redactions},
         "caveats": [
             "Deferred loading is an application strategy report, not a native provider integration.",
             "Token proxy values are char/4 estimates over sanitized local JSON, not billed provider tokens.",
+            "Deferred schema token fields are initial-prompt proxy accounting; full schemas must be retrieved before deferred tool use.",
             "Use receipt get commands to retrieve full sanitized schemas before using deferred tools.",
         ],
     }