cleveragents-core/robot/helper_safety_profile_enforcement.py

"""Robot Framework helper for Safety Profile Enforcement smoke tests.

Provides a CLI-style interface for Robot to invoke safety profile
resolution and enforcement operations.  Exit code 0 = success, 1 = failure.

Usage:
    python robot/helper_safety_profile_enforcement.py resolve-precedence
    python robot/helper_safety_profile_enforcement.py resolve-default
    python robot/helper_safety_profile_enforcement.py enforce-unsafe-blocked
    python robot/helper_safety_profile_enforcement.py enforce-unsafe-allowed
    python robot/helper_safety_profile_enforcement.py enforce-category-blocked
    python robot/helper_safety_profile_enforcement.py enforce-category-allowed
    python robot/helper_safety_profile_enforcement.py enforce-checkpoint
    python robot/helper_safety_profile_enforcement.py enforce-sandbox-blocked
    python robot/helper_safety_profile_enforcement.py enforce-sandbox-allowed
    python robot/helper_safety_profile_enforcement.py enforce-human-approval
    python robot/helper_safety_profile_enforcement.py enforce-cost-limit
    python robot/helper_safety_profile_enforcement.py enforce-retry-limit
    python robot/helper_safety_profile_enforcement.py backward-compat
    python robot/helper_safety_profile_enforcement.py context-summary
"""

from __future__ import annotations

import sys
from collections.abc import Callable
from pathlib import Path

# Ensure src is importable.
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from cleveragents.domain.models.core.safety_profile import (  # noqa: E402
    DEFAULT_SAFETY_PROFILE,
    SafetyProfile,
    SafetyProfileProvenance,
    resolve_safety_profile,
)
from cleveragents.domain.models.core.tool import (  # noqa: E402
    Tool,
    ToolCapability,
    ToolSource,
)
from cleveragents.tool.context import ToolExecutionContext  # noqa: E402
from cleveragents.tool.lifecycle import (  # noqa: E402
    ToolCheckpointRequiredError,
    ToolCostLimitExceededError,
    ToolDescriptor,
    ToolHumanApprovalRequiredError,
    ToolResult,
    ToolRetryLimitExceededError,
    ToolRuntime,
    ToolSafetyViolationError,
    ToolSandboxRequiredError,
)

# ---------------------------------------------------------------------------
# Stub tool instance
# ---------------------------------------------------------------------------


class _StubInstance:
    """Minimal ToolInstance for integration smoke tests."""

    def __init__(self, name: str) -> None:
        self._name = name

    def discover(self) -> ToolDescriptor:
        return ToolDescriptor(name=self._name, description=f"stub {self._name}")

    def activate(self, ctx: ToolExecutionContext) -> None:
        pass

    def execute(self, params: dict, ctx: ToolExecutionContext) -> ToolResult:  # type: ignore[type-arg]
        return ToolResult(success=True, data={"ok": True})

    def deactivate(self, ctx: ToolExecutionContext) -> None:
        pass


def _make_tool(name: str, *, unsafe: bool = False, read_only: bool = False) -> Tool:
    return Tool(
        name=name,
        description=f"Test tool {name}",
        source=ToolSource.CUSTOM,
        code="pass",
        capability=ToolCapability(read_only=read_only, unsafe=unsafe),
    )


def _make_runtime() -> ToolRuntime:
    rt = ToolRuntime()
    for name, unsafe, ro in [
        ("test/safe", False, True),
        ("test/unsafe", True, False),
    ]:
        t = _make_tool(name, unsafe=unsafe, read_only=ro)
        rt.register_tool(t, _StubInstance(name))
    return rt


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------


def _cmd_resolve_precedence() -> int:
    """Test that plan-level profile wins over all others."""
    plan = SafetyProfile(allow_unsafe_tools=True, require_sandbox=False)
    action = SafetyProfile(allow_unsafe_tools=False, require_sandbox=True)
    resolved, prov = resolve_safety_profile(plan_profile=plan, action_profile=action)
    if prov != SafetyProfileProvenance.PLAN:
        print(f"resolve-precedence-fail: expected PLAN, got {prov}")
        return 1
    if resolved.allow_unsafe_tools is not True:
        print("resolve-precedence-fail: allow_unsafe_tools should be True")
        return 1
    print("resolve-precedence-ok")
    return 0


def _cmd_resolve_default() -> int:
    """Test that all-None returns DEFAULT_SAFETY_PROFILE."""
    resolved, prov = resolve_safety_profile()
    if prov != SafetyProfileProvenance.GLOBAL:
        print(f"resolve-default-fail: expected GLOBAL, got {prov}")
        return 1
    if resolved is not DEFAULT_SAFETY_PROFILE:
        print("resolve-default-fail: not DEFAULT_SAFETY_PROFILE instance")
        return 1
    print("resolve-default-ok")
    return 0


def _cmd_enforce_unsafe_blocked() -> int:
    """Test that unsafe tool is blocked by safety profile."""
    rt = _make_runtime()
    profile = SafetyProfile(
        allow_unsafe_tools=False, require_sandbox=False, require_checkpoints=False
    )
    ctx = ToolExecutionContext(plan_id="robot-001", safety_profile=profile)
    try:
        rt.execute("test/unsafe", {}, ctx)
        print("enforce-unsafe-blocked-fail: no error raised")
        return 1
    except ToolSafetyViolationError:
        print("enforce-unsafe-blocked-ok")
        return 0


def _cmd_enforce_unsafe_allowed() -> int:
    """Test that unsafe tool runs when profile allows it."""
    rt = _make_runtime()
    profile = SafetyProfile(
        allow_unsafe_tools=True, require_sandbox=False, require_checkpoints=False
    )
    ctx = ToolExecutionContext(plan_id="robot-002", safety_profile=profile)
    result = rt.execute("test/unsafe", {}, ctx)
    if not result.success:
        print("enforce-unsafe-allowed-fail: execution not successful")
        return 1
    print("enforce-unsafe-allowed-ok")
    return 0


def _cmd_enforce_category_blocked() -> int:
    """Test that tool in wrong skill category is blocked."""
    rt = _make_runtime()
    profile = SafetyProfile(
        allowed_skill_categories=["code"],
        require_sandbox=False,
        require_checkpoints=False,
    )
    ctx = ToolExecutionContext(
        plan_id="robot-003",
        safety_profile=profile,
        metadata={"tool_skill_category": "deploy"},
    )
    try:
        rt.execute("test/safe", {}, ctx)
        print("enforce-category-blocked-fail: no error raised")
        return 1
    except ToolSafetyViolationError:
        print("enforce-category-blocked-ok")
        return 0


def _cmd_enforce_category_allowed() -> int:
    """Test that tool in allowed skill category runs."""
    rt = _make_runtime()
    profile = SafetyProfile(
        allowed_skill_categories=["code"],
        require_sandbox=False,
        require_checkpoints=False,
    )
    ctx = ToolExecutionContext(
        plan_id="robot-004",
        safety_profile=profile,
        metadata={"tool_skill_category": "code"},
    )
    result = rt.execute("test/safe", {}, ctx)
    if not result.success:
        print("enforce-category-allowed-fail: execution not successful")
        return 1
    print("enforce-category-allowed-ok")
    return 0


def _cmd_enforce_checkpoint() -> int:
    """Test that safety profile require_checkpoints blocks non-checkpointable tools."""
    rt = _make_runtime()
    profile = SafetyProfile(require_checkpoints=True, require_sandbox=False)
    ctx = ToolExecutionContext(plan_id="robot-005", safety_profile=profile)
    try:
        rt.execute("test/safe", {}, ctx)
        print("enforce-checkpoint-fail: no error raised")
        return 1
    except ToolCheckpointRequiredError:
        print("enforce-checkpoint-ok")
        return 0


def _cmd_backward_compat() -> int:
    """Test that unsafe tools run when no safety profile on context."""
    rt = _make_runtime()
    ctx = ToolExecutionContext(plan_id="robot-006")
    result = rt.execute("test/unsafe", {}, ctx)
    if not result.success:
        print("backward-compat-fail: execution not successful")
        return 1
    print("backward-compat-ok")
    return 0


def _cmd_context_summary() -> int:
    """Test that context summary includes safety profile flag."""
    profile = SafetyProfile(require_sandbox=False, require_checkpoints=False)
    ctx = ToolExecutionContext(plan_id="robot-007", safety_profile=profile)
    summary = ctx.as_summary()
    if summary.get("has_safety_profile") is not True:
        print(f"context-summary-fail: {summary}")
        return 1
    ctx2 = ToolExecutionContext(plan_id="robot-008")
    summary2 = ctx2.as_summary()
    if summary2.get("has_safety_profile") is not False:
        print(f"context-summary-fail: {summary2}")
        return 1
    print("context-summary-ok")
    return 0


def _cmd_enforce_sandbox_blocked() -> int:
    """Test that writing tool is blocked when sandbox required but missing."""
    rt = ToolRuntime()
    tool = Tool(
        name="test/writer",
        description="writes",
        source=ToolSource.CUSTOM,
        code="pass",
        capability=ToolCapability(writes=True),
    )
    rt.register_tool(tool, _StubInstance("test/writer"))
    profile = SafetyProfile(require_sandbox=True, require_checkpoints=False)
    ctx = ToolExecutionContext(
        plan_id="robot-009", safety_profile=profile, sandbox_id=None
    )
    try:
        rt.execute("test/writer", {}, ctx)
        print("enforce-sandbox-blocked-fail: no error raised")
        return 1
    except ToolSandboxRequiredError:
        print("enforce-sandbox-blocked-ok")
        return 0


def _cmd_enforce_sandbox_allowed() -> int:
    """Test that writing tool runs when sandbox is provided."""
    rt = ToolRuntime()
    tool = Tool(
        name="test/writer",
        description="writes",
        source=ToolSource.CUSTOM,
        code="pass",
        capability=ToolCapability(writes=True),
    )
    rt.register_tool(tool, _StubInstance("test/writer"))
    profile = SafetyProfile(require_sandbox=True, require_checkpoints=False)
    ctx = ToolExecutionContext(
        plan_id="robot-010",
        safety_profile=profile,
        sandbox_id="sandbox-001",
    )
    result = rt.execute("test/writer", {}, ctx)
    if not result.success:
        print("enforce-sandbox-allowed-fail: not successful")
        return 1
    print("enforce-sandbox-allowed-ok")
    return 0


def _cmd_enforce_human_approval() -> int:
    """Test that tools are blocked when human approval is required."""
    rt = _make_runtime()
    profile = SafetyProfile(
        require_human_approval=True,
        require_sandbox=False,
        require_checkpoints=False,
    )
    ctx = ToolExecutionContext(plan_id="robot-011", safety_profile=profile)
    try:
        rt.execute("test/safe", {}, ctx)
        print("enforce-human-approval-fail: no error raised")
        return 1
    except ToolHumanApprovalRequiredError:
        print("enforce-human-approval-ok")
        return 0


def _cmd_enforce_cost_limit() -> int:
    """Test that tools are blocked when cost limit is exceeded."""
    rt = _make_runtime()
    profile = SafetyProfile(
        max_cost_per_plan=10.0,
        require_sandbox=False,
        require_checkpoints=False,
    )
    ctx = ToolExecutionContext(
        plan_id="robot-012",
        safety_profile=profile,
        accumulated_cost=10.0,
    )
    try:
        rt.execute("test/safe", {}, ctx)
        print("enforce-cost-limit-fail: no error raised")
        return 1
    except ToolCostLimitExceededError:
        print("enforce-cost-limit-ok")
        return 0


def _cmd_enforce_retry_limit() -> int:
    """Test that tools are blocked when retry limit is exceeded."""
    rt = _make_runtime()
    profile = SafetyProfile(
        max_retries_per_step=3,
        require_sandbox=False,
        require_checkpoints=False,
    )
    ctx = ToolExecutionContext(
        plan_id="robot-013",
        safety_profile=profile,
        step_retry_count=4,
    )
    try:
        rt.execute("test/safe", {}, ctx)
        print("enforce-retry-limit-fail: no error raised")
        return 1
    except ToolRetryLimitExceededError:
        print("enforce-retry-limit-ok")
        return 0


# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

_COMMANDS: dict[str, Callable[[], int]] = {
    "resolve-precedence": _cmd_resolve_precedence,
    "resolve-default": _cmd_resolve_default,
    "enforce-unsafe-blocked": _cmd_enforce_unsafe_blocked,
    "enforce-unsafe-allowed": _cmd_enforce_unsafe_allowed,
    "enforce-category-blocked": _cmd_enforce_category_blocked,
    "enforce-category-allowed": _cmd_enforce_category_allowed,
    "enforce-checkpoint": _cmd_enforce_checkpoint,
    "enforce-sandbox-blocked": _cmd_enforce_sandbox_blocked,
    "enforce-sandbox-allowed": _cmd_enforce_sandbox_allowed,
    "enforce-human-approval": _cmd_enforce_human_approval,
    "enforce-cost-limit": _cmd_enforce_cost_limit,
    "enforce-retry-limit": _cmd_enforce_retry_limit,
    "backward-compat": _cmd_backward_compat,
    "context-summary": _cmd_context_summary,
}


def main() -> int:
    """Entry point called by Robot Framework ``Run Process``."""
    if len(sys.argv) < 2:
        print("Usage: helper_safety_profile_enforcement.py <command>")
        return 1
    command = sys.argv[1]
    handler = _COMMANDS.get(command)
    if handler is None:
        print(f"Unknown command: {command}")
        return 1
    return handler()


if __name__ == "__main__":
    sys.exit(main())