116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
"""
|
|
Phase 2 #2A: Four-Level Tool Risk Approval System
|
|
Implements: read(0) / write_local(1) / exec(2) / external(3)
|
|
Appended to chat_service.py
|
|
"""
|
|
|
|
# Risk level constants
|
|
TOOL_RISK_LEVELS = {
|
|
"read": 0,
|
|
"write_local": 1,
|
|
"exec": 2,
|
|
"external": 3,
|
|
}
|
|
|
|
# Risk level defaults by tool type
|
|
DEFAULT_RISK_LEVELS = {
|
|
"skill_web_search": "read",
|
|
"skill_code_interpreter": "exec",
|
|
"skill_file_read": "read",
|
|
"skill_file_write": "write_local",
|
|
"mcp_default": "external",
|
|
"skill_default": "write_local",
|
|
}
|
|
|
|
|
|
def assess_tool_risk(func_name: str, mcp_info: dict = None, skill_info: dict = None) -> dict:
|
|
"""Determine risk level of a tool call."""
|
|
risk = "write_local" # default
|
|
|
|
if func_name.startswith("mcp__"):
|
|
# MCP tools default to external (highest risk)
|
|
risk = "external"
|
|
if mcp_info and mcp_info.get("tool_name", ""):
|
|
tn = mcp_info.get("tool_name", "").lower()
|
|
if any(k in tn for k in ["search", "find", "get", "list", "read"]):
|
|
risk = "read"
|
|
elif any(k in tn for k in ["create", "update", "write", "save"]):
|
|
risk = "write_local"
|
|
elif func_name.startswith("install_app"):
|
|
risk = "write_local"
|
|
elif skill_info:
|
|
cap = skill_info.get("capability", "").lower()
|
|
if any(k in cap for k in ["search", "read", "query", "analyze"]):
|
|
risk = "read"
|
|
elif any(k in cap for k in ["generate", "create", "write"]):
|
|
risk = "write_local"
|
|
elif any(k in cap for k in ["execute", "run", "code"]):
|
|
risk = "exec"
|
|
|
|
level = TOOL_RISK_LEVELS.get(risk, 1)
|
|
return {
|
|
"risk_level": risk,
|
|
"numeric_level": level,
|
|
"requires_approval": level >= 2, # exec + external require approval
|
|
"auto_allow": level == 0, # read is auto-allowed
|
|
}
|
|
|
|
|
|
async def create_approval_request(
|
|
user_id: str, tool_name: str, risk_level: str,
|
|
func_args: dict, reason: str = ""
|
|
) -> dict:
|
|
"""Create an approval request in DB."""
|
|
try:
|
|
from app.database import SessionLocal
|
|
from sqlalchemy import text
|
|
db = SessionLocal()
|
|
import uuid
|
|
apr_id = str(uuid.uuid4())
|
|
db.execute(text(
|
|
"INSERT INTO pending_approvals (id, user_id, tool_name, risk_level, "
|
|
"func_args, reason, status, created_at) "
|
|
"VALUES (:id, :uid, :tn, :rl, :fa, :r, 'pending', NOW())"
|
|
), {"id": apr_id, "uid": user_id, "tn": tool_name,
|
|
"rl": risk_level, "fa": str(func_args)[:1000], "r": reason})
|
|
db.commit()
|
|
db.close()
|
|
return {"approval_id": apr_id, "status": "pending"}
|
|
except Exception:
|
|
return {"approval_id": None, "status": "auto_blocked"}
|
|
|
|
|
|
async def check_approval_status(approval_id: str) -> dict:
|
|
"""Check if an approval request has been resolved."""
|
|
try:
|
|
from app.database import SessionLocal
|
|
from sqlalchemy import text
|
|
db = SessionLocal()
|
|
row = db.execute(text(
|
|
"SELECT status, reviewer_id FROM pending_approvals WHERE id = :id"
|
|
), {"id": approval_id}).fetchone()
|
|
db.close()
|
|
if row:
|
|
return {"status": row[0], "reviewer_id": row[1]}
|
|
return {"status": "not_found"}
|
|
except Exception:
|
|
return {"status": "error"}
|
|
|
|
|
|
async def approve_or_reject(approval_id: str, action: str, reviewer_id: str) -> dict:
|
|
"""Approve or reject a pending approval."""
|
|
try:
|
|
from app.database import SessionLocal
|
|
from sqlalchemy import text
|
|
db = SessionLocal()
|
|
new_status = "approved" if action == "approve" else "rejected"
|
|
db.execute(text(
|
|
"UPDATE pending_approvals SET status = :st, reviewer_id = :rid, "
|
|
"reviewed_at = NOW() WHERE id = :id"
|
|
), {"st": new_status, "rid": reviewer_id, "id": approval_id})
|
|
db.commit()
|
|
db.close()
|
|
return {"approval_id": approval_id, "status": new_status}
|
|
except Exception as e:
|
|
return {"status": "error", "error": str(e)}
|