agentengine-sdk-python 0.6.2__py3-none-any.whl → 0.6.9__py3-none-any.whl

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)