Agent/cyber_agent/tools/approval.py

"""Persistent contracts for human approval of tool-call batches."""

from __future__ import annotations

import json
import uuid
from datetime import datetime
from hashlib import sha256
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator


def canonical_arguments(arguments: dict[str, Any]) -> str:
    return json.dumps(
        arguments,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    )


def tool_argument_hash(
    *,
    tool_call_id: str,
    tool_name: str,
    arguments: dict[str, Any],
) -> str:
    payload = f"{tool_call_id}\n{tool_name}\n{canonical_arguments(arguments)}"
    return sha256(payload.encode("utf-8")).hexdigest()


class _StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class ToolApprovalCallV1(_StrictModel):
    tool_call_id: str = Field(min_length=1)
    tool_name: str = Field(min_length=1)
    original_arguments: dict[str, Any]
    effective_arguments: dict[str, Any]
    argument_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
    editable_params: list[str]
    requires_confirmation: bool
    decision: Literal["pending", "approved", "rejected", "auto_approved"]
    execution_status: Literal[
        "not_started",
        "executing",
        "executed",
        "rejected",
        "execution_unknown",
    ] = "not_started"
    result_message_id: str | None = None
    decided_at: str | None = None

    @model_validator(mode="after")
    def validate_hash(self) -> "ToolApprovalCallV1":
        expected = tool_argument_hash(
            tool_call_id=self.tool_call_id,
            tool_name=self.tool_name,
            arguments=self.effective_arguments,
        )
        if self.argument_hash != expected:
            raise ValueError("tool approval argument hash mismatch")
        return self


class ToolApprovalBatchV1(_StrictModel):
    schema_version: Literal[1] = 1
    batch_id: str = Field(min_length=1)
    trace_id: str = Field(min_length=1)
    assistant_message_id: str = Field(min_length=1)
    assistant_sequence: int = Field(gt=0)
    status: Literal[
        "pending",
        "decided",
        "executing",
        "recoverable_candidate_command",
        "completed",
        "cancelled",
        "execution_unknown",
    ] = "pending"
    calls: list[ToolApprovalCallV1] = Field(min_length=1)
    created_at: str
    updated_at: str

    @classmethod
    def create(
        cls,
        *,
        trace_id: str,
        assistant_message_id: str,
        assistant_sequence: int,
        calls: list[ToolApprovalCallV1],
    ) -> "ToolApprovalBatchV1":
        now = datetime.now().isoformat()
        return cls(
            batch_id=f"approval_{uuid.uuid4().hex}",
            trace_id=trace_id,
            assistant_message_id=assistant_message_id,
            assistant_sequence=assistant_sequence,
            calls=calls,
            created_at=now,
            updated_at=now,
        )

    @property
    def pending_calls(self) -> list[ToolApprovalCallV1]:
        return [call for call in self.calls if call.decision == "pending"]


def approval_grant(batch: ToolApprovalBatchV1, call: ToolApprovalCallV1) -> dict[str, Any]:
    """Build the internal, hash-bound grant consumed by ToolRegistry.execute."""
    return {
        "batch_id": batch.batch_id,
        "trace_id": batch.trace_id,
        "tool_call_id": call.tool_call_id,
        "tool_name": call.tool_name,
        "argument_hash": call.argument_hash,
        "decision": call.decision,
    }