626 lines
26 KiB
Python
626 lines
26 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import secrets
|
||
import sqlite3
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Optional
|
||
|
||
from hermes_constants import get_hermes_home
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_PERMISSION_RANK = {"user": 1, "pro": 2, "admin": 3}
|
||
_COOLDOWN_HOURS = 24
|
||
_DEFAULT_WARN_THRESHOLD = 80
|
||
_ALLOWED_EMAIL_DOMAINS = ("@bremen.com.tw", "@journeys.com.tw")
|
||
|
||
|
||
def _utcnow() -> datetime:
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def _iso(dt: Optional[datetime]) -> Optional[str]:
|
||
if dt is None:
|
||
return None
|
||
if dt.tzinfo is None:
|
||
dt = dt.replace(tzinfo=timezone.utc)
|
||
return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat()
|
||
|
||
|
||
def _parse_iso(value: Optional[str]) -> Optional[datetime]:
|
||
if not value:
|
||
return None
|
||
try:
|
||
dt = datetime.fromisoformat(value)
|
||
except ValueError:
|
||
return None
|
||
if dt.tzinfo is None:
|
||
dt = dt.replace(tzinfo=timezone.utc)
|
||
return dt.astimezone(timezone.utc)
|
||
|
||
|
||
def resolve_governance_db_path() -> Path:
|
||
hermes_home = get_hermes_home()
|
||
candidates = [
|
||
hermes_home.parent / "hermes-governance-admin" / "data" / "governance.db",
|
||
hermes_home / "governance.db",
|
||
]
|
||
for path in candidates:
|
||
if path.exists():
|
||
return path
|
||
return candidates[0]
|
||
|
||
|
||
def _connect(*, readonly: bool) -> sqlite3.Connection:
|
||
path = resolve_governance_db_path()
|
||
if not path.exists():
|
||
raise FileNotFoundError(str(path))
|
||
if readonly:
|
||
conn = sqlite3.connect(f"file:{path.as_posix()}?mode=ro", uri=True, timeout=0.5)
|
||
else:
|
||
conn = sqlite3.connect(str(path), timeout=0.5)
|
||
try:
|
||
conn.execute("PRAGMA journal_mode=WAL")
|
||
except Exception:
|
||
logger.debug("Failed to enable WAL on governance DB", exc_info=True)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA busy_timeout = 500")
|
||
return conn
|
||
|
||
|
||
def _platform_value(platform: Any) -> str:
|
||
return getattr(platform, "value", platform) or ""
|
||
|
||
|
||
def _normalize_email(email: str) -> str:
|
||
return (email or "").strip().lower()
|
||
|
||
|
||
def _is_allowed_email(email: str) -> bool:
|
||
return _normalize_email(email).endswith(_ALLOWED_EMAIL_DOMAINS)
|
||
|
||
|
||
def _bind_identity_enforcement_row(
|
||
conn: sqlite3.Connection,
|
||
*,
|
||
platform: str,
|
||
external_user_id: str,
|
||
external_name: Optional[str],
|
||
principal_id: str,
|
||
verified_email: str,
|
||
now_iso: str,
|
||
) -> None:
|
||
row = conn.execute(
|
||
"SELECT id FROM identity_enforcement WHERE platform = ? AND external_user_id = ? LIMIT 1",
|
||
(platform, external_user_id),
|
||
).fetchone()
|
||
if row is None:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO identity_enforcement(
|
||
platform, external_user_id, external_user_name, state, principal_id,
|
||
verified_email, pending_email, failure_cycles, last_failure_at,
|
||
last_failure_reason, resend_count, blocked_until, created_at, updated_at
|
||
) VALUES (?, ?, ?, 'verified', ?, ?, NULL, 0, NULL, NULL, 0, NULL, ?, ?)
|
||
""",
|
||
(platform, external_user_id, external_name, principal_id, verified_email, now_iso, now_iso),
|
||
)
|
||
return
|
||
conn.execute(
|
||
"""
|
||
UPDATE identity_enforcement
|
||
SET external_user_name = ?, state = 'verified', principal_id = ?, verified_email = ?,
|
||
pending_email = NULL, failure_cycles = 0, last_failure_at = NULL,
|
||
last_failure_reason = NULL, resend_count = 0, blocked_until = NULL, updated_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(external_name, principal_id, verified_email, now_iso, int(row["id"])),
|
||
)
|
||
|
||
|
||
def bind_verified_identity(
|
||
*,
|
||
platform: str,
|
||
external_user_id: str,
|
||
verified_email: str,
|
||
external_name: Optional[str] = None,
|
||
) -> dict[str, Any]:
|
||
platform = (platform or "").strip().lower()
|
||
external_user_id = str(external_user_id or "").strip()
|
||
verified_email = _normalize_email(verified_email)
|
||
external_name = (external_name or "").strip() or None
|
||
if not platform or not external_user_id or not verified_email:
|
||
return {"ok": False, "reason": "missing_identity_fields"}
|
||
if not _is_allowed_email(verified_email):
|
||
return {"ok": False, "reason": "email_domain_not_allowed"}
|
||
try:
|
||
conn = _connect(readonly=False)
|
||
except Exception as exc:
|
||
return {"ok": False, "reason": "db_unavailable", "error": str(exc)}
|
||
try:
|
||
now_iso = _iso(_utcnow()) or ""
|
||
principal = conn.execute(
|
||
"SELECT principal_id FROM principals WHERE primary_email = ? LIMIT 1",
|
||
(verified_email,),
|
||
).fetchone()
|
||
if principal is None:
|
||
principal = conn.execute(
|
||
"SELECT principal_id FROM external_identities WHERE verified_email = ? LIMIT 1",
|
||
(verified_email,),
|
||
).fetchone()
|
||
principal_id = str(principal["principal_id"]) if principal and principal["principal_id"] else secrets.token_hex(12)
|
||
if principal is None:
|
||
conn.execute(
|
||
"INSERT INTO principals(principal_id, primary_email, status, created_at, updated_at) VALUES (?, ?, 'verified', ?, ?)",
|
||
(principal_id, verified_email, now_iso, now_iso),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"UPDATE principals SET status = 'verified', updated_at = ? WHERE principal_id = ?",
|
||
(now_iso, principal_id),
|
||
)
|
||
|
||
existing_user = conn.execute(
|
||
"SELECT id, principal_id, verified_email FROM external_identities WHERE platform = ? AND external_user_id = ? LIMIT 1",
|
||
(platform, external_user_id),
|
||
).fetchone()
|
||
if existing_user is not None:
|
||
existing_principal_id = str(existing_user["principal_id"] or "")
|
||
existing_email = _normalize_email(str(existing_user["verified_email"] or ""))
|
||
if existing_principal_id != principal_id or existing_email != verified_email:
|
||
return {
|
||
"ok": False,
|
||
"reason": "identity_conflict",
|
||
"existing_principal_id": existing_principal_id,
|
||
"existing_email": existing_email,
|
||
}
|
||
conn.execute(
|
||
"""
|
||
UPDATE external_identities
|
||
SET external_name = ?, verified_email = ?, verified_at = ?, last_used_at = ?, updated_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(external_name, verified_email, now_iso, now_iso, now_iso, int(existing_user["id"])),
|
||
)
|
||
_bind_identity_enforcement_row(
|
||
conn,
|
||
platform=platform,
|
||
external_user_id=external_user_id,
|
||
external_name=external_name,
|
||
principal_id=principal_id,
|
||
verified_email=verified_email,
|
||
now_iso=now_iso,
|
||
)
|
||
conn.commit()
|
||
return {"ok": True, "reason": "updated", "principal_id": principal_id, "created": False}
|
||
|
||
existing_email = conn.execute(
|
||
"SELECT principal_id, external_user_id FROM external_identities WHERE platform = ? AND verified_email = ? LIMIT 1",
|
||
(platform, verified_email),
|
||
).fetchone()
|
||
if existing_email is not None and str(existing_email["external_user_id"] or "") != external_user_id:
|
||
return {
|
||
"ok": False,
|
||
"reason": "email_already_bound_on_platform",
|
||
"principal_id": str(existing_email["principal_id"] or ""),
|
||
"existing_external_user_id": str(existing_email["external_user_id"] or ""),
|
||
}
|
||
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO external_identities(
|
||
principal_id, platform, external_user_id, external_name,
|
||
verified_email, verified_at, last_used_at, created_at, updated_at
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
principal_id,
|
||
platform,
|
||
external_user_id,
|
||
external_name,
|
||
verified_email,
|
||
now_iso,
|
||
now_iso,
|
||
now_iso,
|
||
now_iso,
|
||
),
|
||
)
|
||
_bind_identity_enforcement_row(
|
||
conn,
|
||
platform=platform,
|
||
external_user_id=external_user_id,
|
||
external_name=external_name,
|
||
principal_id=principal_id,
|
||
verified_email=verified_email,
|
||
now_iso=now_iso,
|
||
)
|
||
conn.commit()
|
||
return {"ok": True, "reason": "created", "principal_id": principal_id, "created": True}
|
||
except Exception as exc:
|
||
logger.exception("bind_verified_identity failed")
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
return {"ok": False, "reason": "exception", "error": str(exc)}
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _combine_limit(values: list[int]) -> int:
|
||
if not values:
|
||
return -1
|
||
if any(v < 0 for v in values):
|
||
return -1
|
||
return max(values)
|
||
|
||
|
||
def _effective_permission_and_quota(conn: sqlite3.Connection, principal_id: str) -> tuple[str, dict[str, Any]]:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT p.permission_level, q.daily_token_limit, q.monthly_token_limit, q.warn_threshold_percent
|
||
FROM group_memberships gm
|
||
LEFT JOIN group_permission_policies p ON p.group_id = gm.group_id
|
||
LEFT JOIN group_quota_policies q ON q.group_id = gm.group_id
|
||
WHERE gm.principal_id = ?
|
||
""",
|
||
(principal_id,),
|
||
).fetchall()
|
||
permission_level = "user"
|
||
daily_values: list[int] = []
|
||
monthly_values: list[int] = []
|
||
warn_values: list[int] = []
|
||
for row in rows:
|
||
level = row["permission_level"] or "user"
|
||
if _PERMISSION_RANK.get(level, 0) > _PERMISSION_RANK.get(permission_level, 0):
|
||
permission_level = level
|
||
if row["daily_token_limit"] is not None:
|
||
daily_values.append(int(row["daily_token_limit"]))
|
||
if row["monthly_token_limit"] is not None:
|
||
monthly_values.append(int(row["monthly_token_limit"]))
|
||
if row["warn_threshold_percent"] is not None:
|
||
warn_values.append(int(row["warn_threshold_percent"]))
|
||
if permission_level == "admin":
|
||
return permission_level, {
|
||
"daily_token_limit": -1,
|
||
"monthly_token_limit": -1,
|
||
"warn_threshold_percent": max(warn_values) if warn_values else _DEFAULT_WARN_THRESHOLD,
|
||
"quota_source": "admin_exempt",
|
||
}
|
||
quota = {
|
||
"daily_token_limit": _combine_limit(daily_values),
|
||
"monthly_token_limit": _combine_limit(monthly_values),
|
||
"warn_threshold_percent": max(warn_values) if warn_values else _DEFAULT_WARN_THRESHOLD,
|
||
"quota_source": "group_merge" if daily_values or monthly_values else "default",
|
||
}
|
||
return permission_level, quota
|
||
|
||
|
||
def _lookup_principal_id(conn: sqlite3.Connection, platform: str, external_user_id: str) -> Optional[str]:
|
||
row = conn.execute(
|
||
"SELECT principal_id FROM external_identities WHERE platform = ? AND external_user_id = ? LIMIT 1",
|
||
(platform, external_user_id),
|
||
).fetchone()
|
||
if row and row["principal_id"]:
|
||
return str(row["principal_id"])
|
||
row = conn.execute(
|
||
"SELECT principal_id FROM identity_enforcement WHERE platform = ? AND external_user_id = ? AND principal_id IS NOT NULL LIMIT 1",
|
||
(platform, external_user_id),
|
||
).fetchone()
|
||
if row and row["principal_id"]:
|
||
return str(row["principal_id"])
|
||
return None
|
||
|
||
|
||
def _usage_totals(conn: sqlite3.Connection, principal_id: str, platform: str, event_date: str, event_month: str) -> tuple[int, int]:
|
||
daily_row = conn.execute(
|
||
"SELECT total_tokens FROM principal_usage_daily WHERE principal_id = ? AND platform = ? AND usage_date = ? LIMIT 1",
|
||
(principal_id, platform, event_date),
|
||
).fetchone()
|
||
monthly_row = conn.execute(
|
||
"SELECT total_tokens FROM principal_usage_monthly WHERE principal_id = ? AND platform = ? AND usage_month = ? LIMIT 1",
|
||
(principal_id, platform, event_month),
|
||
).fetchone()
|
||
return int((daily_row["total_tokens"] if daily_row else 0) or 0), int((monthly_row["total_tokens"] if monthly_row else 0) or 0)
|
||
|
||
|
||
def render_block_message(reason: str) -> str:
|
||
if reason == "principal_suspended":
|
||
return "你的帳號目前已被停權，暫時無法使用 Hermes。"
|
||
if reason == "daily_quota_exceeded":
|
||
return "你今天的使用額度已用完，請明天再試。"
|
||
if reason == "monthly_quota_exceeded":
|
||
return "你本月的使用額度已用完，請下個月再試。"
|
||
return "你目前無法使用 Hermes，請稍後再試或聯絡管理員。"
|
||
|
||
|
||
def enforce_runtime_access_for_identity(*, platform: str, external_user_id: str, estimated_tokens: int) -> dict[str, Any]:
|
||
try:
|
||
conn = _connect(readonly=True)
|
||
except Exception as exc:
|
||
return {
|
||
"allowed": True,
|
||
"reason": "fail_open_db_unavailable",
|
||
"fail_open": True,
|
||
"error": str(exc),
|
||
"quota": {"daily_token_limit": -1, "monthly_token_limit": -1, "warn_threshold_percent": _DEFAULT_WARN_THRESHOLD},
|
||
}
|
||
try:
|
||
principal_id = _lookup_principal_id(conn, platform, external_user_id)
|
||
if not principal_id:
|
||
return {
|
||
"allowed": True,
|
||
"reason": "fail_open_principal_missing",
|
||
"fail_open": True,
|
||
"quota": {"daily_token_limit": -1, "monthly_token_limit": -1, "warn_threshold_percent": _DEFAULT_WARN_THRESHOLD},
|
||
}
|
||
principal = conn.execute(
|
||
"SELECT principal_id, primary_email, status, suspend_reason FROM principals WHERE principal_id = ? LIMIT 1",
|
||
(principal_id,),
|
||
).fetchone()
|
||
if not principal:
|
||
return {
|
||
"allowed": True,
|
||
"reason": "fail_open_principal_row_missing",
|
||
"fail_open": True,
|
||
"principal_id": principal_id,
|
||
"quota": {"daily_token_limit": -1, "monthly_token_limit": -1, "warn_threshold_percent": _DEFAULT_WARN_THRESHOLD},
|
||
}
|
||
permission_level, quota = _effective_permission_and_quota(conn, principal_id)
|
||
if (principal["status"] or "").lower() == "suspended":
|
||
return {
|
||
"allowed": False,
|
||
"reason": "principal_suspended",
|
||
"principal_id": principal_id,
|
||
"permission_level": permission_level,
|
||
"quota": quota,
|
||
}
|
||
now = _utcnow()
|
||
event_date = now.date().isoformat()
|
||
event_month = now.strftime("%Y-%m")
|
||
daily_total, monthly_total = _usage_totals(conn, principal_id, platform, event_date, event_month)
|
||
projected_daily = daily_total + max(int(estimated_tokens or 0), 0)
|
||
projected_monthly = monthly_total + max(int(estimated_tokens or 0), 0)
|
||
if quota.get("daily_token_limit", -1) > 0 and projected_daily > quota["daily_token_limit"]:
|
||
return {
|
||
"allowed": False,
|
||
"reason": "daily_quota_exceeded",
|
||
"principal_id": principal_id,
|
||
"permission_level": permission_level,
|
||
"quota": quota,
|
||
"daily_total_tokens": daily_total,
|
||
"monthly_total_tokens": monthly_total,
|
||
}
|
||
if quota.get("monthly_token_limit", -1) > 0 and projected_monthly > quota["monthly_token_limit"]:
|
||
return {
|
||
"allowed": False,
|
||
"reason": "monthly_quota_exceeded",
|
||
"principal_id": principal_id,
|
||
"permission_level": permission_level,
|
||
"quota": quota,
|
||
"daily_total_tokens": daily_total,
|
||
"monthly_total_tokens": monthly_total,
|
||
}
|
||
return {
|
||
"allowed": True,
|
||
"reason": "ok",
|
||
"principal_id": principal_id,
|
||
"permission_level": permission_level,
|
||
"quota": quota,
|
||
"daily_total_tokens": daily_total,
|
||
"monthly_total_tokens": monthly_total,
|
||
}
|
||
except Exception as exc:
|
||
logger.exception("runtime governance enforcement failed; fail-open")
|
||
return {
|
||
"allowed": True,
|
||
"reason": "fail_open_exception",
|
||
"fail_open": True,
|
||
"error": str(exc),
|
||
"quota": {"daily_token_limit": -1, "monthly_token_limit": -1, "warn_threshold_percent": _DEFAULT_WARN_THRESHOLD},
|
||
}
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def enforce_runtime_access_for_source(source: Any, *, estimated_tokens: int) -> dict[str, Any]:
|
||
platform = _platform_value(getattr(source, "platform", None))
|
||
user_id = str(getattr(source, "user_id", "") or "")
|
||
if not platform or not user_id:
|
||
return {
|
||
"allowed": True,
|
||
"reason": "fail_open_missing_identity",
|
||
"fail_open": True,
|
||
"quota": {"daily_token_limit": -1, "monthly_token_limit": -1, "warn_threshold_percent": _DEFAULT_WARN_THRESHOLD},
|
||
}
|
||
return enforce_runtime_access_for_identity(platform=platform, external_user_id=user_id, estimated_tokens=estimated_tokens)
|
||
|
||
|
||
def _write_audit(conn: sqlite3.Connection, *, principal_id: str, platform: str, event_date: str, total_tokens: int, daily_total: int, monthly_total: int, alerts: list[str]) -> None:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO admin_audit_logs(actor, action, target_type, target_id, before_json, after_json, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
"runtime",
|
||
"runtime_usage_recorded",
|
||
"principal",
|
||
principal_id,
|
||
json.dumps({"platform": platform, "event_date": event_date}, ensure_ascii=False),
|
||
json.dumps({"total_tokens": total_tokens, "daily_total": daily_total, "monthly_total": monthly_total, "alerts": alerts}, ensure_ascii=False),
|
||
_iso(_utcnow()),
|
||
),
|
||
)
|
||
|
||
|
||
def record_runtime_usage_for_identity(
|
||
*,
|
||
platform: str,
|
||
external_user_id: str,
|
||
model: Optional[str],
|
||
input_tokens: int,
|
||
output_tokens: int,
|
||
session_id: Optional[str] = None,
|
||
session_key: Optional[str] = None,
|
||
message_id: Optional[str] = None,
|
||
question_count: int = 1,
|
||
event_time: Optional[datetime] = None,
|
||
) -> dict[str, Any]:
|
||
try:
|
||
conn = _connect(readonly=False)
|
||
except Exception as exc:
|
||
return {"allowed": True, "reason": "fail_open_db_unavailable", "fail_open": True, "error": str(exc)}
|
||
try:
|
||
principal_id = _lookup_principal_id(conn, platform, external_user_id)
|
||
if not principal_id:
|
||
return {"allowed": True, "reason": "fail_open_principal_missing", "fail_open": True}
|
||
principal = conn.execute(
|
||
"SELECT principal_id, status, last_notified_daily_80_at, last_notified_monthly_80_at FROM principals WHERE principal_id = ? LIMIT 1",
|
||
(principal_id,),
|
||
).fetchone()
|
||
if not principal:
|
||
return {"allowed": True, "reason": "fail_open_principal_row_missing", "fail_open": True, "principal_id": principal_id}
|
||
when = event_time or _utcnow()
|
||
if when.tzinfo is None:
|
||
when = when.replace(tzinfo=timezone.utc)
|
||
when = when.astimezone(timezone.utc)
|
||
event_date = when.date().isoformat()
|
||
event_month = when.strftime("%Y-%m")
|
||
input_tokens = int(input_tokens or 0)
|
||
output_tokens = int(output_tokens or 0)
|
||
total_tokens = input_tokens + output_tokens
|
||
permission_level, quota = _effective_permission_and_quota(conn, principal_id)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO usage_events(
|
||
principal_id, platform, session_id, session_key, message_id, model,
|
||
input_tokens, output_tokens, total_tokens, question_count,
|
||
event_date, event_month, created_at
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
principal_id,
|
||
platform,
|
||
session_id,
|
||
session_key,
|
||
message_id,
|
||
model,
|
||
input_tokens,
|
||
output_tokens,
|
||
total_tokens,
|
||
int(question_count or 1),
|
||
event_date,
|
||
event_month,
|
||
_iso(when),
|
||
),
|
||
)
|
||
daily_row = conn.execute(
|
||
"SELECT id, question_count, input_tokens, output_tokens, total_tokens FROM principal_usage_daily WHERE principal_id = ? AND platform = ? AND usage_date = ? LIMIT 1",
|
||
(principal_id, platform, event_date),
|
||
).fetchone()
|
||
if daily_row is None:
|
||
conn.execute(
|
||
"INSERT INTO principal_usage_daily(principal_id, platform, usage_date, question_count, input_tokens, output_tokens, total_tokens, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
|
||
(principal_id, platform, event_date, int(question_count or 1), input_tokens, output_tokens, total_tokens, _iso(when)),
|
||
)
|
||
daily_total = total_tokens
|
||
else:
|
||
daily_total = int(daily_row["total_tokens"] or 0) + total_tokens
|
||
conn.execute(
|
||
"UPDATE principal_usage_daily SET question_count = ?, input_tokens = ?, output_tokens = ?, total_tokens = ?, updated_at = ? WHERE id = ?",
|
||
(
|
||
int(daily_row["question_count"] or 0) + int(question_count or 1),
|
||
int(daily_row["input_tokens"] or 0) + input_tokens,
|
||
int(daily_row["output_tokens"] or 0) + output_tokens,
|
||
daily_total,
|
||
_iso(when),
|
||
int(daily_row["id"]),
|
||
),
|
||
)
|
||
monthly_row = conn.execute(
|
||
"SELECT id, question_count, input_tokens, output_tokens, total_tokens FROM principal_usage_monthly WHERE principal_id = ? AND platform = ? AND usage_month = ? LIMIT 1",
|
||
(principal_id, platform, event_month),
|
||
).fetchone()
|
||
if monthly_row is None:
|
||
conn.execute(
|
||
"INSERT INTO principal_usage_monthly(principal_id, platform, usage_month, question_count, input_tokens, output_tokens, total_tokens, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
|
||
(principal_id, platform, event_month, int(question_count or 1), input_tokens, output_tokens, total_tokens, _iso(when)),
|
||
)
|
||
monthly_total = total_tokens
|
||
else:
|
||
monthly_total = int(monthly_row["total_tokens"] or 0) + total_tokens
|
||
conn.execute(
|
||
"UPDATE principal_usage_monthly SET question_count = ?, input_tokens = ?, output_tokens = ?, total_tokens = ?, updated_at = ? WHERE id = ?",
|
||
(
|
||
int(monthly_row["question_count"] or 0) + int(question_count or 1),
|
||
int(monthly_row["input_tokens"] or 0) + input_tokens,
|
||
int(monthly_row["output_tokens"] or 0) + output_tokens,
|
||
monthly_total,
|
||
_iso(when),
|
||
int(monthly_row["id"]),
|
||
),
|
||
)
|
||
alerts: list[str] = []
|
||
threshold = max(int(quota.get("warn_threshold_percent", _DEFAULT_WARN_THRESHOLD)), 1) / 100.0
|
||
daily_limit = int(quota.get("daily_token_limit", -1) or -1)
|
||
monthly_limit = int(quota.get("monthly_token_limit", -1) or -1)
|
||
last_daily = _parse_iso(principal["last_notified_daily_80_at"])
|
||
last_monthly = _parse_iso(principal["last_notified_monthly_80_at"])
|
||
next_daily = principal["last_notified_daily_80_at"]
|
||
next_monthly = principal["last_notified_monthly_80_at"]
|
||
if daily_limit > 0 and daily_total >= int(daily_limit * threshold):
|
||
if last_daily is None or when - last_daily >= timedelta(hours=_COOLDOWN_HOURS):
|
||
alerts.append("daily_80")
|
||
next_daily = _iso(when)
|
||
if monthly_limit > 0 and monthly_total >= int(monthly_limit * threshold):
|
||
if last_monthly is None or when - last_monthly >= timedelta(hours=_COOLDOWN_HOURS):
|
||
alerts.append("monthly_80")
|
||
next_monthly = _iso(when)
|
||
if alerts:
|
||
conn.execute(
|
||
"UPDATE principals SET last_notified_daily_80_at = ?, last_notified_monthly_80_at = ?, updated_at = ? WHERE principal_id = ?",
|
||
(next_daily, next_monthly, _iso(when), principal_id),
|
||
)
|
||
_write_audit(
|
||
conn,
|
||
principal_id=principal_id,
|
||
platform=platform,
|
||
event_date=event_date,
|
||
total_tokens=total_tokens,
|
||
daily_total=daily_total,
|
||
monthly_total=monthly_total,
|
||
alerts=alerts,
|
||
)
|
||
conn.commit()
|
||
return {
|
||
"allowed": True,
|
||
"reason": "ok",
|
||
"principal_id": principal_id,
|
||
"permission_level": permission_level,
|
||
"threshold_alerts": alerts,
|
||
"daily_total_tokens": daily_total,
|
||
"monthly_total_tokens": monthly_total,
|
||
"quota": quota,
|
||
}
|
||
except Exception as exc:
|
||
logger.exception("runtime governance usage recording failed; ignored")
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
return {"allowed": True, "reason": "fail_open_exception", "fail_open": True, "error": str(exc)}
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def record_runtime_usage_for_source(source: Any, **kwargs: Any) -> dict[str, Any]:
|
||
platform = _platform_value(getattr(source, "platform", None))
|
||
user_id = str(getattr(source, "user_id", "") or "")
|
||
if not platform or not user_id:
|
||
return {"allowed": True, "reason": "fail_open_missing_identity", "fail_open": True}
|
||
return record_runtime_usage_for_identity(platform=platform, external_user_id=user_id, **kwargs)
|