agentengine-sdk-python 0.6.4__py3-none-any.whl → 0.7.0__py3-none-any.whl

ksadk/terminal_exec_policy.py

@@ -0,0 +1,229 @@
+"""Shared remote terminal exec command policy."""
+
+from __future__ import annotations
+
+import os
+import re
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+
+SHELL_METACHARS = set("|&;<>()$`\\\n\r")
+FORBIDDEN_LAUNCHERS = {
+    "bash",
+    "sh",
+    "zsh",
+    "fish",
+    "python",
+    "python3",
+    "node",
+    "npx",
+    "pnpm",
+    "npm",
+    "yarn",
+    "uv",
+    "uvx",
+    "hermes",
+}
+
+TERMINAL_EXEC_ALLOWLIST_ENV = "KSADK_TERMINAL_EXEC_SUBCOMMAND_ALLOWLIST"
+
+COMMON_REMOTE_EXEC_PREFIXES: tuple[tuple[str, ...], ...] = (
+    ("cat",),
+    ("date",),
+    ("df",),
+    ("du",),
+    ("find",),
+    ("git", "diff"),
+    ("git", "log"),
+    ("git", "rev-parse"),
+    ("git", "show"),
+    ("git", "status"),
+    ("head",),
+    ("id",),
+    ("ls",),
+    ("ps",),
+    ("pwd",),
+    ("stat",),
+    ("tail",),
+    ("uname",),
+    ("wc",),
+    ("whoami",),
+)
+
+
+@dataclass(frozen=True)
+class TerminalExecPolicy:
+    name: str
+    default_exact: tuple[tuple[str, ...], ...] = ()
+    default_prefixes: tuple[tuple[str, ...], ...] = ()
+    default_bounded_prefixes: tuple[tuple[tuple[str, ...], int, int], ...] = ()
+    env_names: tuple[str, ...] = (TERMINAL_EXEC_ALLOWLIST_ENV,)
+
+
+GENERIC_TERMINAL_EXEC_POLICY = TerminalExecPolicy(
+    name="terminal",
+    default_prefixes=COMMON_REMOTE_EXEC_PREFIXES,
+)
+
+HERMES_TERMINAL_EXEC_POLICY = TerminalExecPolicy(
+    name="Hermes",
+    default_exact=(
+        ("cron", "list"),
+        ("cron", "status"),
+        ("config", "check"),
+        ("config", "env-path"),
+        ("config", "path"),
+        ("config", "show"),
+        ("doctor",),
+        ("gateway", "status"),
+        ("insights",),
+        ("sessions", "list"),
+        ("skills", "audit"),
+        ("skills", "check"),
+        ("skills", "list"),
+        ("status",),
+        ("tools", "list"),
+        ("version",),
+    ),
+    default_bounded_prefixes=(
+        (("sessions", "export"), 3, 3),
+        (("sessions", "show"), 3, 3),
+    ),
+)
+
+OPENCLAW_TERMINAL_EXEC_POLICY = TerminalExecPolicy(
+    name="OpenClaw",
+    default_prefixes=(
+        *COMMON_REMOTE_EXEC_PREFIXES,
+        ("openclaw", "channels", "login"),
+    ),
+)
+
+_ALLOWLIST_ENTRY_SPLIT_RE = re.compile(r"[\n;,]+")
+_ALLOWLIST_TOKEN_SPLIT_RE = re.compile(r"\s+")
+
+
+def normalize_exec_argv(
+    argv: Iterable[str],
+    *,
+    policy_name: str = "terminal",
+    allow_shell_metachars: bool = False,
+    allow_forbidden_launchers: bool = False,
+) -> list[str]:
+    normalized = [str(item).strip() for item in argv]
+    if not normalized:
+        raise ValueError(f"{policy_name} exec requires argv")
+    for index, item in enumerate(normalized):
+        if not item:
+            raise ValueError(f"{policy_name} exec argv contains an empty argument")
+        if not allow_shell_metachars and any(char in SHELL_METACHARS for char in item):
+            raise ValueError(f"{policy_name} exec does not allow shell metacharacters: {item}")
+        if index == 0 and item.startswith("-"):
+            raise ValueError(f"{policy_name} exec command is invalid: {item}")
+        if not allow_forbidden_launchers and index == 0 and item in FORBIDDEN_LAUNCHERS:
+            raise ValueError(f"{policy_name} exec launcher is not allowed: {item}")
+    return normalized
+
+
+def validate_terminal_exec_argv(
+    argv: Iterable[str],
+    *,
+    policy: TerminalExecPolicy = GENERIC_TERMINAL_EXEC_POLICY,
+) -> list[str]:
+    if _env_allowlist_all(policy.env_names):
+        return normalize_exec_argv(
+            argv,
+            policy_name=policy.name,
+            allow_shell_metachars=True,
+            allow_forbidden_launchers=True,
+        )
+
+    normalized = normalize_exec_argv(argv, policy_name=policy.name)
+    allowed_exact = tuple(tuple(item) for item in policy.default_exact)
+    allowed_prefixes = (
+        tuple(tuple(item) for item in policy.default_prefixes)
+        + _env_allowlist_prefixes(policy.env_names, policy_name=policy.name)
+    )
+
+    if _matches_exact_or_prefix(
+        normalized,
+        allowed_exact=allowed_exact,
+        allowed_prefixes=allowed_prefixes,
+        allowed_bounded_prefixes=policy.default_bounded_prefixes,
+    ):
+        return normalized
+    suggested_prefix = normalized[0]
+    raise ValueError(
+        f"{policy.name} exec subcommand is not allowed: {' '.join(normalized)}. "
+        f"Set {TERMINAL_EXEC_ALLOWLIST_ENV}='{suggested_prefix}' to allow this prefix, "
+        f"or {TERMINAL_EXEC_ALLOWLIST_ENV}='*' to allow all remote exec commands."
+    )
+
+
+def _env_allowlist_prefixes(
+    env_names: Sequence[str],
+    *,
+    policy_name: str,
+) -> tuple[tuple[str, ...], ...]:
+    prefixes: list[tuple[str, ...]] = []
+    seen: set[tuple[str, ...]] = set()
+    for env_name in env_names:
+        for prefix in _parse_allowlist_prefixes(os.getenv(env_name), policy_name=policy_name):
+            if prefix not in seen:
+                prefixes.append(prefix)
+                seen.add(prefix)
+    return tuple(prefixes)
+
+
+def _env_allowlist_all(env_names: Sequence[str]) -> bool:
+    for env_name in env_names:
+        raw = os.getenv(env_name)
+        if not raw:
+            continue
+        for entry in _ALLOWLIST_ENTRY_SPLIT_RE.split(raw):
+            if entry.strip() == "*":
+                return True
+    return False
+
+
+def _parse_allowlist_prefixes(raw: str | None, *, policy_name: str) -> tuple[tuple[str, ...], ...]:
+    if not raw:
+        return ()
+    prefixes: list[tuple[str, ...]] = []
+    for entry in _ALLOWLIST_ENTRY_SPLIT_RE.split(raw):
+        entry = entry.strip()
+        if not entry:
+            continue
+        if entry == "*":
+            continue
+        tokens = tuple(token for token in _ALLOWLIST_TOKEN_SPLIT_RE.split(entry) if token)
+        normalize_exec_argv(tokens, policy_name=policy_name)
+        prefixes.append(tokens)
+    return tuple(prefixes)
+
+
+def _matches_exact_or_prefix(
+    argv: Sequence[str],
+    *,
+    allowed_exact: Sequence[Sequence[str]],
+    allowed_prefixes: Sequence[Sequence[str]],
+    allowed_bounded_prefixes: Sequence[tuple[Sequence[str], int, int]],
+) -> bool:
+    argv_tuple = tuple(argv)
+    for item in allowed_exact:
+        if argv_tuple == tuple(item):
+            return True
+    for prefix in allowed_prefixes:
+        prefix_tuple = tuple(prefix)
+        if prefix_tuple and argv_tuple[: len(prefix_tuple)] == prefix_tuple:
+            return True
+    for prefix, min_len, max_len in allowed_bounded_prefixes:
+        prefix_tuple = tuple(prefix)
+        if (
+            prefix_tuple
+            and argv_tuple[: len(prefix_tuple)] == prefix_tuple
+            and len(argv_tuple) >= min_len
+            and len(argv_tuple) <= max_len
+        ):
+            return True
+    return False

ksadk/tools/__init__.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from ksadk.tools.gateway import (
+    ToolGateway,
+    ToolPolicy,
+    approval_interrupt_info_from_result,
+    check_command_policy,
+    default_tool_gateway,
+)
+
+__all__ = [
+    "ToolGateway",
+    "ToolPolicy",
+    "approval_interrupt_info_from_result",
+    "check_command_policy",
+    "default_tool_gateway",
+]

ksadk/tools/gateway.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+import os
+import hashlib
+import json
+import shlex
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+from uuid import uuid4
+
+
+@dataclass(frozen=True)
+class ToolPolicy:
+    risk_level: str = "low"
+    side_effects: Sequence[str] = field(default_factory=tuple)
+    requires_approval: bool | None = None
+
+
+_READ_ONLY_GIT_COMMANDS = {"status", "diff", "log", "show", "branch", "ls-files"}
+_DANGEROUS_COMMAND_TOKENS = {
+    "sudo",
+    "kubectl",
+    "docker",
+}
+_DANGEROUS_GIT_COMMANDS = {"reset", "clean", "push", "checkout"}
+
+
+class ToolGateway:
+    def __init__(self, policies: Mapping[str, ToolPolicy] | None = None):
+        self._policies = dict(policies or {})
+
+    def invoke(
+        self,
+        tool_name: str,
+        func: Callable[..., Any],
+        *args: Any,
+        approval: Mapping[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        policy = self._policies.get(tool_name, ToolPolicy())
+        if self._is_rejected(approval):
+            return self._approval_rejected(tool_name, policy, approval)
+        if self._is_approved(approval) or not self._requires_approval(policy):
+            return func(*args, **kwargs)
+        return self._approval_required(tool_name, policy)
+
+    def _requires_approval(self, policy: ToolPolicy) -> bool:
+        return tool_policy_requires_approval(policy, approval_mode=self._approval_mode())
+
+    @staticmethod
+    def _approval_mode() -> str:
+        value = os.environ.get("KSADK_TOOL_APPROVAL_MODE", "").strip().lower()
+        return value or "off"
+
+    @staticmethod
+    def _is_approved(approval: Mapping[str, Any] | None) -> bool:
+        if not approval:
+            return False
+        return bool(approval.get("approved") or approval.get("approve"))
+
+    @staticmethod
+    def _is_rejected(approval: Mapping[str, Any] | None) -> bool:
+        if not approval:
+            return False
+        if "approved" in approval:
+            return not bool(approval.get("approved"))
+        if "approve" in approval:
+            return not bool(approval.get("approve"))
+        return False
+
+    @staticmethod
+    def _approval_required(tool_name: str, policy: ToolPolicy) -> dict[str, Any]:
+        return {
+            "ok": False,
+            "type": "approval_required",
+            "approval_required": True,
+            "approval_request": {
+                "id": f"appr_{uuid4().hex}",
+                "tool_name": tool_name,
+                "risk_level": policy.risk_level,
+                "side_effects": list(policy.side_effects),
+            },
+        }
+
+    @staticmethod
+    def _approval_rejected(
+        tool_name: str,
+        policy: ToolPolicy,
+        approval: Mapping[str, Any] | None,
+    ) -> dict[str, Any]:
+        return {
+            "ok": False,
+            "type": "approval_rejected",
+            "approval_required": False,
+            "tool_name": tool_name,
+            "risk_level": policy.risk_level,
+            "reason": str((approval or {}).get("reason") or ""),
+        }
+
+
+def default_tool_gateway(policies: Mapping[str, ToolPolicy] | None = None) -> ToolGateway:
+    return ToolGateway(policies)
+
+
+def tool_policy_requires_approval(
+    policy: ToolPolicy,
+    *,
+    approval_mode: str | None = None,
+) -> bool:
+    mode = (approval_mode or os.environ.get("KSADK_TOOL_APPROVAL_MODE", "")).strip().lower() or "off"
+    if policy.requires_approval is not None:
+        return policy.requires_approval and mode != "off"
+    if mode != "strict":
+        return False
+    return policy.risk_level.lower() in {"medium", "high", "critical"}
+
+
+def check_command_policy(command: str) -> dict[str, Any]:
+    text = str(command or "").strip()
+    if not text:
+        return {"ok": False, "decision": "reject", "error_type": "command_required", "error_message": "command is required"}
+    try:
+        tokens = shlex.split(text)
+    except ValueError as exc:
+        return {"ok": False, "decision": "reject", "error_type": "invalid_command", "error_message": str(exc)}
+    if not tokens:
+        return {"ok": False, "decision": "reject", "error_type": "command_required", "error_message": "command is required"}
+    if _contains_recursive_rm(tokens):
+        return _command_rejected("recursive rm is not allowed without explicit approval")
+    if any(token in _DANGEROUS_COMMAND_TOKENS for token in tokens):
+        return _command_rejected(f"{tokens[0]} command is not allowed by default policy")
+    if tokens[0] == "git" and len(tokens) > 1:
+        subcommand = tokens[1]
+        if subcommand in _READ_ONLY_GIT_COMMANDS:
+            return {"ok": True, "decision": "allow", "reason": "read_only_git"}
+        if subcommand in _DANGEROUS_GIT_COMMANDS:
+            return _command_rejected(f"git {subcommand} is not allowed without explicit approval")
+    if _references_metadata_endpoint(tokens):
+        return _command_rejected("metadata/private endpoint access is not allowed by default policy")
+    return {"ok": True, "decision": "allow", "reason": "default_allow"}
+
+
+def _command_rejected(message: str) -> dict[str, Any]:
+    return {"ok": False, "decision": "reject", "error_type": "command_rejected", "error_message": message}
+
+
+def _contains_recursive_rm(tokens: Sequence[str]) -> bool:
+    for index, token in enumerate(tokens[:-1]):
+        if token == "rm" and any(
+            flag.startswith("-") and ("r" in flag or "R" in flag)
+            for flag in tokens[index + 1 : index + 3]
+        ):
+            return True
+    return False
+
+
+def _references_metadata_endpoint(tokens: Sequence[str]) -> bool:
+    joined = " ".join(tokens)
+    return "169.254.169.254" in joined or "metadata.google.internal" in joined
+
+
+def _canonical_tool_args(tool_args: Any) -> str:
+    try:
+        return json.dumps(tool_args or {}, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+    except TypeError:
+        return json.dumps(str(tool_args), ensure_ascii=False, separators=(",", ":"))
+
+
+def build_tool_receipt_idempotency_key(
+    *,
+    session_id: str,
+    run_id: str,
+    checkpoint_id: str | None = None,
+    tool_call_id: str | None = None,
+    tool_name: str,
+    tool_args: Any = None,
+) -> str:
+    payload = {
+        "session_id": str(session_id or ""),
+        "run_id": str(run_id or ""),
+        "checkpoint_id": str(checkpoint_id or ""),
+        "tool_call_id": str(tool_call_id or ""),
+        "tool_name": str(tool_name or ""),
+        "tool_args": _canonical_tool_args(tool_args),
+    }
+    digest = hashlib.sha256(
+        json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    ).hexdigest()
+    return f"tool_receipt:{digest}"
+
+
+def approval_interrupt_info_from_result(
+    result: Any,
+    *,
+    fallback_tool_name: str = "tool",
+    tool_args: Any = None,
+    run_id: str | None = None,
+) -> dict[str, Any] | None:
+    if not isinstance(result, Mapping):
+        return None
+    if str(result.get("type") or "") != "approval_required":
+        return None
+    approval_request = result.get("approval_request")
+    if not isinstance(approval_request, Mapping):
+        return None
+
+    request_id = str(
+        approval_request.get("approval_request_id")
+        or approval_request.get("id")
+        or f"appr_{uuid4().hex}"
+    )
+    interrupt = {
+        "id": request_id,
+        "approval_request_id": request_id,
+        "tool_name": str(approval_request.get("tool_name") or fallback_tool_name or "tool"),
+        "arguments": (
+            approval_request.get("arguments")
+            or approval_request.get("tool_args")
+            or approval_request.get("args")
+            or tool_args
+            or {}
+        ),
+        "risk_level": str(approval_request.get("risk_level") or result.get("risk_level") or ""),
+        "side_effects": list(approval_request.get("side_effects") or result.get("side_effects") or []),
+        "server_label": str(approval_request.get("server_label") or result.get("server_label") or "ksadk"),
+    }
+    if run_id:
+        interrupt["run_id"] = str(run_id)
+    return interrupt

ksadk/tools/result_budget.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from ksadk.sessions.local_service import resolve_local_session_dir
+
+
+_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9_.-]+")
+
+
+@dataclass(frozen=True)
+class ToolResultBudget:
+    max_chars: int = 50_000
+    preview_chars: int = 8_000
+    persist_threshold_chars: int = 50_000
+    persist_dir: Path | str | None = None
+
+    @property
+    def resolved_persist_dir(self) -> Path:
+        if self.persist_dir is not None:
+            return Path(self.persist_dir).expanduser().resolve()
+        configured = os.environ.get("KSADK_TOOL_RESULT_DIR", "").strip()
+        if configured:
+            return Path(configured).expanduser().resolve()
+        return (resolve_local_session_dir() / "tool-results").resolve()
+
+
+def default_tool_result_budget(
+    *,
+    max_chars: int | None = None,
+    preview_chars: int | None = None,
+    persist_threshold_chars: int | None = None,
+) -> ToolResultBudget:
+    return ToolResultBudget(
+        max_chars=_int_env("KSADK_TOOL_RESULT_MAX_CHARS", max_chars or 50_000),
+        preview_chars=_int_env("KSADK_TOOL_RESULT_PREVIEW_CHARS", preview_chars or 8_000),
+        persist_threshold_chars=_int_env(
+            "KSADK_TOOL_RESULT_PERSIST_THRESHOLD_CHARS",
+            persist_threshold_chars or max_chars or 50_000,
+        ),
+    )
+
+
+def budget_tool_output(
+    *,
+    tool_name: str,
+    field_name: str,
+    value: Any,
+    metadata: dict[str, Any] | None = None,
+    budget: ToolResultBudget | None = None,
+) -> dict[str, Any]:
+    active_budget = budget or default_tool_result_budget()
+    text, mime_type, extension = _stringify_value(value)
+    original_chars = len(text)
+    should_persist = original_chars > active_budget.persist_threshold_chars or original_chars > active_budget.max_chars
+    preview_limit = max(0, min(active_budget.preview_chars, active_budget.max_chars))
+    preview = text[:preview_limit] if should_persist else text[: active_budget.max_chars]
+    result: dict[str, Any] = {
+        field_name: preview,
+        "truncated": original_chars > len(preview),
+        "original_chars": original_chars,
+        "preview_chars": len(preview),
+    }
+    if should_persist:
+        path = _persist_tool_output(
+            tool_name=tool_name,
+            field_name=field_name,
+            text=text,
+            extension=extension,
+            metadata=metadata or {},
+            budget=active_budget,
+        )
+        result["persisted"] = {"path": str(path), "mime_type": mime_type}
+    return result
+
+
+def budget_text_fields(
+    payload: dict[str, Any],
+    *,
+    tool_name: str,
+    fields: tuple[str, ...],
+    metadata: dict[str, Any] | None = None,
+    budget: ToolResultBudget | None = None,
+) -> dict[str, Any]:
+    updated = dict(payload)
+    for field in fields:
+        if field not in updated:
+            continue
+        budgeted = budget_tool_output(
+            tool_name=tool_name,
+            field_name=field,
+            value=updated[field],
+            metadata={**(metadata or {}), "field_name": field},
+            budget=budget,
+        )
+        updated[field] = budgeted[field]
+        for key in ("truncated", "original_chars", "preview_chars", "persisted"):
+            if key in budgeted:
+                if key == "persisted":
+                    updated.setdefault("persisted_outputs", {})[field] = budgeted[key]
+                else:
+                    updated[f"{field}_{key}"] = budgeted[key]
+    return updated
+
+
+def _persist_tool_output(
+    *,
+    tool_name: str,
+    field_name: str,
+    text: str,
+    extension: str,
+    metadata: dict[str, Any],
+    budget: ToolResultBudget,
+) -> Path:
+    persist_dir = budget.resolved_persist_dir
+    persist_dir.mkdir(parents=True, exist_ok=True)
+    tool_use_id = str(
+        metadata.get("tool_use_id")
+        or metadata.get("tool_call_id")
+        or metadata.get("id")
+        or uuid4().hex
+    )
+    stem = _safe_filename(f"{tool_use_id}.{field_name}") or _safe_filename(f"{tool_name}.{field_name}.{uuid4().hex}")
+    path = (persist_dir / f"{stem}.{extension}").resolve()
+    if persist_dir not in path.parents and path != persist_dir:
+        raise ValueError("tool result path must stay inside persist_dir")
+    path.write_text(text, encoding="utf-8")
+    return path
+
+
+def _stringify_value(value: Any) -> tuple[str, str, str]:
+    if isinstance(value, str):
+        return value, "text/plain", "txt"
+    return json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True), "application/json", "json"
+
+
+def _safe_filename(value: str) -> str:
+    return _SAFE_NAME_RE.sub("_", value.strip()).strip("._")
+
+
+def _int_env(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None:
+        return int(default)
+    try:
+        return max(0, int(raw))
+    except ValueError:
+        return int(default)