mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-09 05:50:44 +08:00
5168 lines
197 KiB
Python
5168 lines
197 KiB
Python
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
import sqlite3
|
||
import sys
|
||
import hashlib
|
||
import uuid
|
||
from collections.abc import Iterator
|
||
from contextlib import contextmanager
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Optional
|
||
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
from svc.persistence.pg_adapter import PgConnShim, connect_postgres
|
||
from svc.persistence.pg_compat import scrub_nul_bytes_from_jsonable, scrub_nul_bytes_from_text
|
||
|
||
if sys.platform == "win32":
|
||
import ctypes
|
||
from ctypes import wintypes
|
||
|
||
|
||
def utc_now_iso() -> str:
|
||
return datetime.now(timezone.utc).isoformat()
|
||
|
||
|
||
_LOG = logging.getLogger(__name__)
|
||
|
||
|
||
def _chat_message_persist_log_enabled() -> bool:
|
||
v = str(os.getenv("AIA_LOG_CHAT_MESSAGE_PERSIST") or "").strip().lower()
|
||
return v in ("1", "true", "yes", "on")
|
||
|
||
|
||
# 内置模型配置（不可删除；rule 不在设置界面展示）
|
||
LLM_BUILTIN_OLLAMA_PROFILE_ID = "00000001-0000-4000-8000-000000000001"
|
||
LLM_BUILTIN_RULE_PROFILE_ID = "00000001-0000-4000-8000-000000000002"
|
||
_BUILTIN_OLLAMA_MODEL_SEED = "qwen2.5:7b"
|
||
_BUILTIN_OLLAMA_BASE_SEED = "http://127.0.0.1:11434/v1"
|
||
|
||
# 与 src.agents.specialists.AGENT_PROFILE_BINDINGS_KEY 保持一致（避免 sqlite_store 导入 agents 环依赖）
|
||
_MODEL_BINDINGS_KEY = "agent_profile_bindings"
|
||
|
||
|
||
def is_administrator_model_pool(username: str | None) -> bool:
|
||
"""控制台「全局模型池」仅 administrator 用户名使用；其余用户各自一套 profile + active + bindings。"""
|
||
return str(username or "").strip().lower() == "administrator"
|
||
|
||
|
||
def active_llm_profile_setting_key(user_id: str, username: str | None) -> str:
|
||
return "active_llm_profile_id" if is_administrator_model_pool(username) else f"active_llm_profile_id:{user_id}"
|
||
|
||
|
||
def agent_profile_bindings_setting_key(user_id: str, username: str | None) -> str:
|
||
return _MODEL_BINDINGS_KEY if is_administrator_model_pool(username) else f"{_MODEL_BINDINGS_KEY}:{user_id}"
|
||
|
||
|
||
def _tool_row_assistant_message_id(tool_calls_text: str | None) -> int | None:
|
||
if not tool_calls_text:
|
||
return None
|
||
try:
|
||
meta = json.loads(tool_calls_text)
|
||
if isinstance(meta, dict):
|
||
aid = meta.get("assistant_message_id")
|
||
if aid is not None:
|
||
return int(aid)
|
||
except (json.JSONDecodeError, TypeError, ValueError):
|
||
pass
|
||
return None
|
||
|
||
|
||
def _trim_messages_start_index(rows: list[Any], keep_last: int) -> int | None:
|
||
"""计算应保留的起始索引，避免 tool 消息与对应 assistant 消息断链。"""
|
||
n = len(rows)
|
||
if n <= keep_last:
|
||
return None
|
||
start = n - keep_last
|
||
id_list = [int(r["id"]) for r in rows]
|
||
id_to_index = {mid: i for i, mid in enumerate(id_list)}
|
||
while 0 <= start < n:
|
||
r = rows[start]
|
||
if str(r["role"]) != "tool":
|
||
break
|
||
aid = _tool_row_assistant_message_id(r["tool_calls"])
|
||
if aid is None:
|
||
break
|
||
j = id_to_index.get(aid)
|
||
if j is None or j >= start:
|
||
break
|
||
start = j
|
||
return start
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ChatSession:
|
||
id: str
|
||
title: str
|
||
created_at: str
|
||
last_message_at: str | None = None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ChatMessage:
|
||
id: int
|
||
session_id: str
|
||
role: str
|
||
content: str
|
||
tool_calls: Optional[str]
|
||
timestamp: str
|
||
attachments: Optional[str] = None
|
||
turn_uuid: Optional[str] = None
|
||
event_type: Optional[str] = None
|
||
event_payload: Optional[str] = None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SessionMessagesMeta:
|
||
session_id: str
|
||
message_count: int
|
||
last_message_id: int | None
|
||
last_message_at: str | None
|
||
|
||
@property
|
||
def revision_key(self) -> tuple[str, int, int | None, str | None]:
|
||
return (self.session_id, self.message_count, self.last_message_id, self.last_message_at)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SessionsListMeta:
|
||
session_count: int
|
||
latest_activity_at: str | None
|
||
|
||
@property
|
||
def revision_key(self) -> tuple[int, str | None]:
|
||
return (self.session_count, self.latest_activity_at)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class OclawTask:
|
||
id: str
|
||
tenant_id: str
|
||
session_id: str
|
||
task_type: str
|
||
status: str
|
||
payload: str
|
||
result: str
|
||
attempt_count: int
|
||
claimed_by: str | None
|
||
lease_expires_at: str | None
|
||
last_error: str
|
||
created_at: str
|
||
updated_at: str
|
||
finished_at: str | None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class OclawRun:
|
||
run_id: str
|
||
tenant_id: str
|
||
session_id: str
|
||
status: str
|
||
payload: str
|
||
created_at: str
|
||
updated_at: str
|
||
|
||
|
||
class _CryptoError(RuntimeError):
|
||
pass
|
||
|
||
|
||
def _fernet() -> Any | None:
|
||
"""Best-effort Fernet cipher for non-Windows secret storage.
|
||
|
||
Enable by setting AIA_ASSISTANT_MASTER_KEY (any non-empty string).
|
||
"""
|
||
mk = (os.getenv("AIA_ASSISTANT_MASTER_KEY") or "").strip()
|
||
if not mk:
|
||
return None
|
||
try:
|
||
from cryptography.fernet import Fernet # type: ignore
|
||
except Exception:
|
||
return None
|
||
# Derive a stable 32-byte key from master key string.
|
||
digest = hashlib.sha256(mk.encode("utf-8")).digest()
|
||
key = base64.urlsafe_b64encode(digest)
|
||
try:
|
||
return Fernet(key)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _dpapi_encrypt(plain: bytes) -> bytes:
|
||
if sys.platform != "win32":
|
||
raise _CryptoError("DPAPI 仅支持 Windows")
|
||
|
||
class DATA_BLOB(ctypes.Structure):
|
||
_fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_byte))]
|
||
|
||
crypt32 = ctypes.WinDLL("crypt32", use_last_error=True)
|
||
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
||
|
||
in_blob = DATA_BLOB()
|
||
in_blob.cbData = len(plain)
|
||
buf = ctypes.create_string_buffer(plain, len(plain))
|
||
in_blob.pbData = ctypes.cast(buf, ctypes.POINTER(ctypes.c_byte))
|
||
|
||
out_blob = DATA_BLOB()
|
||
res = crypt32.CryptProtectData(
|
||
ctypes.byref(in_blob),
|
||
None,
|
||
None,
|
||
None,
|
||
None,
|
||
0x1,
|
||
ctypes.byref(out_blob),
|
||
)
|
||
if not res:
|
||
raise _CryptoError(f"CryptProtectData 失败: {ctypes.get_last_error()}")
|
||
|
||
try:
|
||
return ctypes.string_at(out_blob.pbData, out_blob.cbData)
|
||
finally:
|
||
kernel32.LocalFree(out_blob.pbData)
|
||
|
||
|
||
def _dpapi_decrypt(cipher: bytes) -> bytes:
|
||
if sys.platform != "win32":
|
||
raise _CryptoError("DPAPI 仅支持 Windows")
|
||
|
||
class DATA_BLOB(ctypes.Structure):
|
||
_fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_byte))]
|
||
|
||
crypt32 = ctypes.WinDLL("crypt32", use_last_error=True)
|
||
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
||
|
||
in_blob = DATA_BLOB()
|
||
in_blob.cbData = len(cipher)
|
||
buf = ctypes.create_string_buffer(cipher, len(cipher))
|
||
in_blob.pbData = ctypes.cast(buf, ctypes.POINTER(ctypes.c_byte))
|
||
|
||
out_blob = DATA_BLOB()
|
||
res = crypt32.CryptUnprotectData(
|
||
ctypes.byref(in_blob),
|
||
None,
|
||
None,
|
||
None,
|
||
None,
|
||
0x1,
|
||
ctypes.byref(out_blob),
|
||
)
|
||
if not res:
|
||
raise _CryptoError(f"CryptUnprotectData 失败: {ctypes.get_last_error()}")
|
||
|
||
try:
|
||
return ctypes.string_at(out_blob.pbData, out_blob.cbData)
|
||
finally:
|
||
kernel32.LocalFree(out_blob.pbData)
|
||
|
||
|
||
def _encode_secret(plain_text: str) -> str:
|
||
plain_bytes = plain_text.encode("utf-8")
|
||
try:
|
||
enc = _dpapi_encrypt(plain_bytes)
|
||
return "dpapi:" + base64.b64encode(enc).decode("ascii")
|
||
except Exception:
|
||
f = _fernet()
|
||
if f is not None:
|
||
try:
|
||
token = f.encrypt(plain_bytes)
|
||
return "fernet:" + token.decode("ascii")
|
||
except Exception:
|
||
pass
|
||
return "b64:" + base64.b64encode(plain_bytes).decode("ascii")
|
||
|
||
|
||
def _decode_secret(secret_text: str) -> str:
|
||
if secret_text.startswith("dpapi:"):
|
||
cipher = base64.b64decode(secret_text[len("dpapi:") :].encode("ascii"))
|
||
plain = _dpapi_decrypt(cipher)
|
||
return plain.decode("utf-8")
|
||
if secret_text.startswith("fernet:"):
|
||
token = secret_text[len("fernet:") :].strip().encode("ascii")
|
||
f = _fernet()
|
||
if f is None:
|
||
raise _CryptoError("missing AIA_ASSISTANT_MASTER_KEY or cryptography for fernet secret")
|
||
plain = f.decrypt(token)
|
||
return plain.decode("utf-8")
|
||
if secret_text.startswith("b64:"):
|
||
plain = base64.b64decode(secret_text[len("b64:") :].encode("ascii"))
|
||
return plain.decode("utf-8")
|
||
raise _CryptoError("未知的密钥编码格式")
|
||
|
||
|
||
class SqliteStore:
|
||
@staticmethod
|
||
def _cap_json_for_log(obj: Any, *, max_chars: int, keep_keys: tuple[str, ...] = ("ok", "error_code", "error")) -> Any:
|
||
cap = max(2000, int(max_chars or 0))
|
||
try:
|
||
blob = json.dumps(obj, ensure_ascii=False, default=str)
|
||
except Exception:
|
||
blob = json.dumps({"_log_cap": True, "repr": repr(obj)}, ensure_ascii=False)
|
||
if len(blob) <= cap:
|
||
return obj
|
||
slim: dict[str, Any] = {"_log_cap": True, "bytes": len(blob)}
|
||
if isinstance(obj, dict):
|
||
for k in keep_keys:
|
||
if k in obj:
|
||
slim[k] = obj.get(k)
|
||
slim["preview"] = blob[: min(cap, 4000)]
|
||
return slim
|
||
def __init__(self, db_path: str | Path | None = None, *, postgres_url: str | None = None) -> None:
|
||
if postgres_url:
|
||
self._postgres_url = str(postgres_url).strip()
|
||
self._use_pg = True
|
||
self.db_path = self._postgres_url
|
||
else:
|
||
if db_path is None:
|
||
raise TypeError("SqliteStore requires db_path unless postgres_url= is set")
|
||
self.db_path = str(db_path)
|
||
self._postgres_url = ""
|
||
self._use_pg = False
|
||
try:
|
||
self._init_db()
|
||
except Exception as exc:
|
||
if self._use_pg:
|
||
raise RuntimeError(
|
||
"Assistant PostgreSQL initialization failed (connection, permissions, or missing tables). "
|
||
"Apply schema with `alembic upgrade head` (or `svc/persistence/ddl/postgresql_bootstrap.sql`), "
|
||
"verify AIA_ASSISTANT_DATABASE_URL, and inspect the chained exception."
|
||
) from exc
|
||
raise
|
||
|
||
@contextmanager
|
||
def _connect(self) -> Iterator[Any]:
|
||
if self._use_pg:
|
||
raw = connect_postgres(self._postgres_url)
|
||
shim = PgConnShim(raw)
|
||
try:
|
||
yield shim
|
||
raw.commit()
|
||
except BaseException:
|
||
raw.rollback()
|
||
raise
|
||
finally:
|
||
raw.close()
|
||
return
|
||
conn = sqlite3.connect(self.db_path, timeout=30.0)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA foreign_keys = ON;")
|
||
conn.execute("PRAGMA journal_mode = WAL;")
|
||
conn.execute("PRAGMA synchronous = NORMAL;")
|
||
conn.execute("PRAGMA busy_timeout = 30000;")
|
||
try:
|
||
with conn:
|
||
yield conn
|
||
finally:
|
||
conn.close()
|
||
|
||
def _shared_sa_engine(self):
|
||
"""SQLAlchemy engine for the same database as :meth:`_connect` (per SQLite file or PostgreSQL URL)."""
|
||
from svc.persistence.db.engine import engine_for_sqlite_file, get_assistant_engine
|
||
|
||
if self._use_pg:
|
||
return get_assistant_engine()
|
||
return engine_for_sqlite_file(str(self.db_path))
|
||
|
||
def _app_settings_repo(self):
|
||
"""Lazily construct SA repository for ``app_setting`` (same DB URL as raw :meth:`_connect`)."""
|
||
from svc.persistence.sa_repos.app_settings import AppSettingsSaRepository
|
||
|
||
r = self.__dict__.get("_app_settings_sa")
|
||
if r is None:
|
||
r = AppSettingsSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_app_settings_sa"] = r
|
||
return r
|
||
|
||
def _auth_sessions_repo(self):
|
||
"""Lazily construct SA repository for ``auth_session`` (same DB URL as raw :meth:`_connect`)."""
|
||
from svc.persistence.sa_repos.auth_sessions import AuthSessionsSaRepository
|
||
|
||
r = self.__dict__.get("_auth_sessions_sa")
|
||
if r is None:
|
||
r = AuthSessionsSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_auth_sessions_sa"] = r
|
||
return r
|
||
|
||
def _chat_sessions_repo(self):
|
||
"""Lazily construct SA repository for ``chat_session`` / ``ui_session_owner`` list paths."""
|
||
from svc.persistence.sa_repos.chat_sessions import ChatSessionsSaRepository
|
||
|
||
r = self.__dict__.get("_chat_sessions_sa")
|
||
if r is None:
|
||
r = ChatSessionsSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_chat_sessions_sa"] = r
|
||
return r
|
||
|
||
def _ui_session_owner_repo(self):
|
||
from svc.persistence.sa_repos.ui_session_owner import UiSessionOwnerSaRepository
|
||
|
||
r = self.__dict__.get("_ui_session_owner_sa")
|
||
if r is None:
|
||
r = UiSessionOwnerSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_ui_session_owner_sa"] = r
|
||
return r
|
||
|
||
def _tenant_repo(self):
|
||
from svc.persistence.sa_repos.tenant_bind_code import TenantSaRepository
|
||
|
||
r = self.__dict__.get("_tenant_sa")
|
||
if r is None:
|
||
r = TenantSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_tenant_sa"] = r
|
||
return r
|
||
|
||
def _bind_code_repo(self):
|
||
from svc.persistence.sa_repos.tenant_bind_code import BindCodeSaRepository
|
||
|
||
r = self.__dict__.get("_bind_code_sa")
|
||
if r is None:
|
||
r = BindCodeSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_bind_code_sa"] = r
|
||
return r
|
||
|
||
def _app_users_repo(self):
|
||
from svc.persistence.sa_repos.app_users import AppUsersSaRepository
|
||
|
||
r = self.__dict__.get("_app_users_sa")
|
||
if r is None:
|
||
r = AppUsersSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_app_users_sa"] = r
|
||
return r
|
||
|
||
def _chat_messages_repo(self):
|
||
"""Lazily construct SA repository for ``chat_message`` hot paths."""
|
||
from svc.persistence.sa_repos.chat_messages import ChatMessagesSaRepository
|
||
|
||
r = self.__dict__.get("_chat_messages_sa")
|
||
if r is None:
|
||
r = ChatMessagesSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_chat_messages_sa"] = r
|
||
return r
|
||
|
||
def _admin_user_stats_repo(self):
|
||
from svc.persistence.sa_repos.admin_user_stats import AdminUserStatsSaRepository
|
||
|
||
r = self.__dict__.get("_admin_user_stats_sa")
|
||
if r is None:
|
||
r = AdminUserStatsSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_admin_user_stats_sa"] = r
|
||
return r
|
||
|
||
def _session_tool_health_repo(self):
|
||
from svc.persistence.sa_repos.session_tool_health import SessionToolHealthSaRepository
|
||
|
||
r = self.__dict__.get("_session_tool_health_sa")
|
||
if r is None:
|
||
r = SessionToolHealthSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_session_tool_health_sa"] = r
|
||
return r
|
||
|
||
def _tool_log_queries_repo(self):
|
||
from svc.persistence.sa_repos.tool_log_queries import ToolLogQueriesSaRepository
|
||
|
||
r = self.__dict__.get("_tool_log_queries_sa")
|
||
if r is None:
|
||
r = ToolLogQueriesSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_tool_log_queries_sa"] = r
|
||
return r
|
||
|
||
def _trace_events_repo(self):
|
||
from svc.persistence.sa_repos.trace_events import TraceEventsSaRepository
|
||
|
||
r = self.__dict__.get("_trace_events_sa")
|
||
if r is None:
|
||
r = TraceEventsSaRepository(self._shared_sa_engine())
|
||
self.__dict__["_trace_events_sa"] = r
|
||
return r
|
||
|
||
def _init_db(self) -> None:
|
||
if self._use_pg:
|
||
self._init_db_postgresql()
|
||
return
|
||
Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS chat_session (
|
||
id TEXT PRIMARY KEY,
|
||
title TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
last_message_at TEXT
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_chat_session_activity ON chat_session(COALESCE(last_message_at, created_at) DESC, created_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS tenant (
|
||
id TEXT PRIMARY KEY,
|
||
name TEXT NOT NULL,
|
||
created_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS app_user (
|
||
id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
username TEXT,
|
||
display_name TEXT NOT NULL,
|
||
role TEXT NOT NULL,
|
||
password_hash TEXT,
|
||
is_active INTEGER NOT NULL DEFAULT 1,
|
||
created_at TEXT NOT NULL,
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
user_cols = {row[1] for row in conn.execute("PRAGMA table_info(app_user)").fetchall()}
|
||
if "username" not in user_cols:
|
||
conn.execute("ALTER TABLE app_user ADD COLUMN username TEXT")
|
||
if "password_hash" not in user_cols:
|
||
conn.execute("ALTER TABLE app_user ADD COLUMN password_hash TEXT")
|
||
if "is_active" not in user_cols:
|
||
conn.execute("ALTER TABLE app_user ADD COLUMN is_active INTEGER NOT NULL DEFAULT 1")
|
||
if "avatar_attachment_id" not in user_cols:
|
||
conn.execute("ALTER TABLE app_user ADD COLUMN avatar_attachment_id TEXT")
|
||
conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_app_user_tenant_username ON app_user(tenant_id, username)")
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS channel_identity (
|
||
tenant_id TEXT NOT NULL,
|
||
channel TEXT NOT NULL,
|
||
external_user_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
PRIMARY KEY (tenant_id, channel, external_user_id),
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS channel_identity_v2 (
|
||
tenant_id TEXT NOT NULL,
|
||
channel TEXT NOT NULL,
|
||
account_id TEXT NOT NULL,
|
||
external_user_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
PRIMARY KEY (tenant_id, channel, account_id, external_user_id),
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS bind_code (
|
||
code TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
role TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
used_at TEXT,
|
||
used_by_external_user_id TEXT,
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS channel_session (
|
||
tenant_id TEXT NOT NULL,
|
||
channel TEXT NOT NULL,
|
||
external_chat_id TEXT NOT NULL,
|
||
external_user_id TEXT NOT NULL,
|
||
session_id TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
PRIMARY KEY (tenant_id, channel, external_chat_id, external_user_id),
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(session_id) REFERENCES chat_session(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS channel_session_v2 (
|
||
tenant_id TEXT NOT NULL,
|
||
channel TEXT NOT NULL,
|
||
account_id TEXT NOT NULL,
|
||
external_chat_id TEXT NOT NULL,
|
||
external_user_id TEXT NOT NULL,
|
||
session_id TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
PRIMARY KEY (tenant_id, channel, account_id, external_chat_id, external_user_id),
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(session_id) REFERENCES chat_session(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS user_channel_account (
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
channel TEXT NOT NULL,
|
||
account_id TEXT NOT NULL,
|
||
name TEXT NOT NULL DEFAULT '',
|
||
config TEXT NOT NULL DEFAULT '{}',
|
||
is_active INTEGER NOT NULL DEFAULT 1,
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
PRIMARY KEY (tenant_id, user_id, channel, account_id),
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_user_channel_account_channel_account ON user_channel_account(channel, account_id, is_active)"
|
||
)
|
||
uca_cols = {row[1] for row in conn.execute("PRAGMA table_info(user_channel_account)").fetchall()}
|
||
if "name" not in uca_cols:
|
||
conn.execute("ALTER TABLE user_channel_account ADD COLUMN name TEXT NOT NULL DEFAULT ''")
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS todo_item (
|
||
id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
owner_user_id TEXT NOT NULL,
|
||
assignee_user_id TEXT,
|
||
title TEXT NOT NULL,
|
||
due_at TEXT,
|
||
status TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(owner_user_id) REFERENCES app_user(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(assignee_user_id) REFERENCES app_user(id) ON DELETE SET NULL
|
||
);
|
||
"""
|
||
)
|
||
sess_cols = {row[1] for row in conn.execute("PRAGMA table_info(chat_session)").fetchall()}
|
||
if "last_message_at" not in sess_cols:
|
||
conn.execute("ALTER TABLE chat_session ADD COLUMN last_message_at TEXT")
|
||
conn.execute(
|
||
"""
|
||
UPDATE chat_session SET last_message_at = (
|
||
SELECT MAX(timestamp) FROM chat_message
|
||
WHERE chat_message.session_id = chat_session.id
|
||
)
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS chat_message (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
session_id TEXT NOT NULL,
|
||
role TEXT NOT NULL,
|
||
content TEXT NOT NULL,
|
||
tool_calls TEXT,
|
||
attachments TEXT,
|
||
turn_uuid TEXT,
|
||
event_type TEXT,
|
||
event_payload TEXT,
|
||
timestamp TEXT NOT NULL,
|
||
FOREIGN KEY(session_id) REFERENCES chat_session(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
msg_cols = {row[1] for row in conn.execute("PRAGMA table_info(chat_message)").fetchall()}
|
||
if "turn_uuid" not in msg_cols:
|
||
conn.execute("ALTER TABLE chat_message ADD COLUMN turn_uuid TEXT")
|
||
if "event_type" not in msg_cols:
|
||
conn.execute("ALTER TABLE chat_message ADD COLUMN event_type TEXT")
|
||
if "event_payload" not in msg_cols:
|
||
conn.execute("ALTER TABLE chat_message ADD COLUMN event_payload TEXT")
|
||
conn.execute("CREATE INDEX IF NOT EXISTS idx_chat_message_session_turn_uuid ON chat_message(session_id, turn_uuid)")
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS ui_session_owner (
|
||
session_id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
FOREIGN KEY(session_id) REFERENCES chat_session(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_ui_session_owner_tenant_user ON ui_session_owner(tenant_id, user_id, created_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_ui_session_owner_tenant_user_session ON ui_session_owner(tenant_id, user_id, session_id)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_ui_session_owner_tenant_session ON ui_session_owner(tenant_id, session_id)"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS user_workspace_path_allowlist (
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
extra_roots TEXT NOT NULL DEFAULT '',
|
||
allow_any_path INTEGER NOT NULL DEFAULT 0,
|
||
allow_high_risk_public_tools INTEGER NOT NULL DEFAULT 0,
|
||
updated_at TEXT NOT NULL,
|
||
PRIMARY KEY (tenant_id, user_id),
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
uw_cols = {row[1] for row in conn.execute("PRAGMA table_info(user_workspace_path_allowlist)").fetchall()}
|
||
if "allow_high_risk_public_tools" not in uw_cols:
|
||
conn.execute(
|
||
"ALTER TABLE user_workspace_path_allowlist ADD COLUMN allow_high_risk_public_tools INTEGER NOT NULL DEFAULT 0"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS auth_session (
|
||
session_token_hash TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
role TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
expires_at TEXT NOT NULL,
|
||
last_seen_at TEXT NOT NULL,
|
||
revoked_at TEXT,
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute("CREATE INDEX IF NOT EXISTS idx_auth_session_user_expires ON auth_session(user_id, expires_at)")
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS role_permission (
|
||
role TEXT NOT NULL,
|
||
permission TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
PRIMARY KEY (role, permission)
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS user_permission (
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
permission TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
PRIMARY KEY (tenant_id, user_id, permission),
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS admin_audit_log (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
actor_tenant_id TEXT NOT NULL,
|
||
actor_user_id TEXT NOT NULL,
|
||
action TEXT NOT NULL,
|
||
target_type TEXT NOT NULL,
|
||
target_id TEXT NOT NULL,
|
||
status TEXT NOT NULL,
|
||
detail TEXT NOT NULL DEFAULT '{}',
|
||
timestamp TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute("CREATE INDEX IF NOT EXISTS idx_admin_audit_actor_ts ON admin_audit_log(actor_user_id, timestamp DESC)")
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS attachment_acl (
|
||
attachment_id TEXT NOT NULL,
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
session_id TEXT NOT NULL,
|
||
source TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
PRIMARY KEY (attachment_id, tenant_id, user_id, session_id, source),
|
||
FOREIGN KEY(session_id) REFERENCES chat_session(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(tenant_id) REFERENCES tenant(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_attachment_acl_tenant_attachment ON attachment_acl(tenant_id, attachment_id, created_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_attachment_acl_user_attachment ON attachment_acl(tenant_id, user_id, attachment_id, created_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_attachment_acl_session_attachment ON attachment_acl(session_id, attachment_id, created_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS tool_log (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
session_id TEXT NOT NULL,
|
||
tool_name TEXT NOT NULL,
|
||
specialist TEXT NOT NULL DEFAULT '',
|
||
args TEXT NOT NULL,
|
||
result TEXT NOT NULL,
|
||
timestamp TEXT NOT NULL,
|
||
duration_ms INTEGER,
|
||
FOREIGN KEY(session_id) REFERENCES chat_session(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
cols = {row[1] for row in conn.execute("PRAGMA table_info(tool_log)").fetchall()}
|
||
if "duration_ms" not in cols:
|
||
conn.execute("ALTER TABLE tool_log ADD COLUMN duration_ms INTEGER")
|
||
if "specialist" not in cols:
|
||
conn.execute("ALTER TABLE tool_log ADD COLUMN specialist TEXT NOT NULL DEFAULT ''")
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS tool_plugin (
|
||
plugin_name TEXT NOT NULL,
|
||
plugin_version TEXT NOT NULL,
|
||
entry_point TEXT NOT NULL,
|
||
enabled INTEGER NOT NULL DEFAULT 1,
|
||
updated_at TEXT NOT NULL,
|
||
PRIMARY KEY (plugin_name, entry_point)
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS mcp_server_registry (
|
||
server_id TEXT PRIMARY KEY,
|
||
source_type TEXT NOT NULL,
|
||
source_ref TEXT NOT NULL,
|
||
version TEXT NOT NULL DEFAULT '',
|
||
entry_command TEXT NOT NULL DEFAULT '',
|
||
entry_args TEXT NOT NULL DEFAULT '[]',
|
||
env_schema TEXT NOT NULL DEFAULT '{}',
|
||
required_permissions TEXT NOT NULL DEFAULT '[]',
|
||
risk_level TEXT NOT NULL DEFAULT 'high',
|
||
timeout_s REAL NOT NULL DEFAULT 30,
|
||
enabled INTEGER NOT NULL DEFAULT 0,
|
||
updated_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS mcp_server_installation (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
server_id TEXT NOT NULL,
|
||
status TEXT NOT NULL,
|
||
error_code TEXT NOT NULL DEFAULT '',
|
||
detail TEXT NOT NULL DEFAULT '{}',
|
||
install_command TEXT NOT NULL DEFAULT '',
|
||
timestamp TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS mcp_server_health (
|
||
server_id TEXT PRIMARY KEY,
|
||
status TEXT NOT NULL,
|
||
detail TEXT NOT NULL DEFAULT '{}',
|
||
checked_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS mcp_server_tool (
|
||
server_id TEXT NOT NULL,
|
||
tool_name TEXT NOT NULL,
|
||
description TEXT NOT NULL DEFAULT '',
|
||
parameters TEXT NOT NULL DEFAULT '{}',
|
||
updated_at TEXT NOT NULL,
|
||
PRIMARY KEY (server_id, tool_name)
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS app_setting (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT NOT NULL,
|
||
is_secret INTEGER NOT NULL DEFAULT 0,
|
||
updated_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS knowledge_chunk (
|
||
chunk_id TEXT PRIMARY KEY,
|
||
source TEXT NOT NULL,
|
||
content TEXT NOT NULL,
|
||
metadata TEXT NOT NULL DEFAULT '{}',
|
||
updated_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS knowledge_embedding (
|
||
chunk_id TEXT NOT NULL,
|
||
model TEXT NOT NULL,
|
||
dim INTEGER NOT NULL,
|
||
vector_json TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
PRIMARY KEY (chunk_id, model),
|
||
FOREIGN KEY(chunk_id) REFERENCES knowledge_chunk(chunk_id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS memory_item (
|
||
memory_id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
session_id TEXT NOT NULL,
|
||
memory_type TEXT NOT NULL,
|
||
content TEXT NOT NULL,
|
||
confidence REAL NOT NULL DEFAULT 0,
|
||
source TEXT NOT NULL DEFAULT 'memory',
|
||
metadata TEXT NOT NULL DEFAULT '{}',
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
expires_at TEXT
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS memory_vector (
|
||
memory_id TEXT NOT NULL,
|
||
model TEXT NOT NULL,
|
||
dim INTEGER NOT NULL,
|
||
vector_json TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
PRIMARY KEY (memory_id, model),
|
||
FOREIGN KEY(memory_id) REFERENCES memory_item(memory_id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS memory_hit_log (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
tenant_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
session_id TEXT,
|
||
memory_id TEXT,
|
||
query_text TEXT NOT NULL,
|
||
score REAL NOT NULL DEFAULT 0,
|
||
source TEXT NOT NULL DEFAULT 'memory',
|
||
timestamp TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_memory_item_tenant_user_updated ON memory_item(tenant_id, user_id, updated_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_memory_hit_log_tenant_user_ts ON memory_hit_log(tenant_id, user_id, timestamp DESC)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_memory_item_session_updated ON memory_item(session_id, updated_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS oclaw_task (
|
||
id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
session_id TEXT NOT NULL,
|
||
task_type TEXT NOT NULL DEFAULT 'async_turn',
|
||
status TEXT NOT NULL,
|
||
payload TEXT NOT NULL DEFAULT '{}',
|
||
result TEXT NOT NULL DEFAULT '{}',
|
||
attempt_count INTEGER NOT NULL DEFAULT 0,
|
||
claimed_by TEXT,
|
||
lease_expires_at TEXT,
|
||
last_error TEXT NOT NULL DEFAULT '',
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
finished_at TEXT
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS oclaw_run (
|
||
run_id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
session_id TEXT NOT NULL,
|
||
status TEXT NOT NULL,
|
||
payload TEXT NOT NULL DEFAULT '{}',
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS oclaw_attempt (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
run_id TEXT NOT NULL,
|
||
tenant_id TEXT NOT NULL,
|
||
session_id TEXT NOT NULL,
|
||
attempt_no INTEGER NOT NULL,
|
||
status TEXT NOT NULL,
|
||
reason TEXT NOT NULL DEFAULT '',
|
||
payload TEXT NOT NULL DEFAULT '{}',
|
||
created_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_oclaw_task_status_updated ON oclaw_task(status, updated_at)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_oclaw_task_tenant_session ON oclaw_task(tenant_id, session_id, created_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_oclaw_run_tenant_session ON oclaw_run(tenant_id, session_id, updated_at DESC)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_oclaw_attempt_run_no ON oclaw_attempt(run_id, attempt_no)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_knowledge_chunk_source_updated ON knowledge_chunk(source, updated_at)"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS agent_audit_log (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
session_id TEXT NOT NULL,
|
||
specialist TEXT NOT NULL,
|
||
task_kind TEXT NOT NULL,
|
||
action TEXT NOT NULL,
|
||
payload TEXT NOT NULL,
|
||
status TEXT NOT NULL,
|
||
reason TEXT NOT NULL,
|
||
duration_ms INTEGER NOT NULL DEFAULT 0,
|
||
timestamp TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS agent_eval_log (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
session_id TEXT NOT NULL,
|
||
specialist TEXT NOT NULL,
|
||
task_kind TEXT NOT NULL,
|
||
success INTEGER NOT NULL,
|
||
latency_ms INTEGER NOT NULL,
|
||
cost_hint REAL NOT NULL DEFAULT 0,
|
||
notes TEXT NOT NULL DEFAULT '',
|
||
timestamp TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS trace_event (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
session_id TEXT NOT NULL,
|
||
trace_id TEXT NOT NULL,
|
||
span_id TEXT NOT NULL,
|
||
parent_span_id TEXT,
|
||
event_type TEXT NOT NULL,
|
||
payload TEXT NOT NULL DEFAULT '{}',
|
||
timestamp TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_trace_event_session_id_id ON trace_event(session_id, id)"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS llm_profile (
|
||
id TEXT PRIMARY KEY,
|
||
name TEXT NOT NULL,
|
||
mode TEXT NOT NULL,
|
||
model TEXT,
|
||
base_url TEXT,
|
||
api_key TEXT,
|
||
updated_at TEXT NOT NULL
|
||
);
|
||
"""
|
||
)
|
||
prof_cols = {row[1] for row in conn.execute("PRAGMA table_info(llm_profile)").fetchall()}
|
||
if "is_builtin" not in prof_cols:
|
||
conn.execute("ALTER TABLE llm_profile ADD COLUMN is_builtin INTEGER NOT NULL DEFAULT 0")
|
||
if "hide_in_ui" not in prof_cols:
|
||
conn.execute("ALTER TABLE llm_profile ADD COLUMN hide_in_ui INTEGER NOT NULL DEFAULT 0")
|
||
if "owner_user_id" not in prof_cols:
|
||
conn.execute("ALTER TABLE llm_profile ADD COLUMN owner_user_id TEXT")
|
||
if "thinking_mode_enabled" not in prof_cols:
|
||
conn.execute("ALTER TABLE llm_profile ADD COLUMN thinking_mode_enabled INTEGER NOT NULL DEFAULT 0")
|
||
if "reasoning_effort" not in prof_cols:
|
||
conn.execute("ALTER TABLE llm_profile ADD COLUMN reasoning_effort TEXT")
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS llm_profile_user_grant (
|
||
id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
profile_id TEXT NOT NULL,
|
||
user_id TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
created_by_user_id TEXT,
|
||
UNIQUE(tenant_id, profile_id, user_id),
|
||
FOREIGN KEY(profile_id) REFERENCES llm_profile(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(user_id) REFERENCES app_user(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_llm_profile_grant_user ON llm_profile_user_grant(tenant_id, user_id)"
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_llm_profile_grant_profile ON llm_profile_user_grant(tenant_id, profile_id)"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS llm_profile_tenant_grant (
|
||
id TEXT PRIMARY KEY,
|
||
tenant_id TEXT NOT NULL,
|
||
profile_id TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
created_by_user_id TEXT,
|
||
UNIQUE(tenant_id, profile_id),
|
||
FOREIGN KEY(profile_id) REFERENCES llm_profile(id) ON DELETE CASCADE
|
||
);
|
||
"""
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_llm_profile_tenant_grant ON llm_profile_tenant_grant(tenant_id, profile_id)"
|
||
)
|
||
self._seed_builtin_llm_profiles(conn)
|
||
self._seed_default_permissions(conn)
|
||
conn.execute(
|
||
"""
|
||
UPDATE llm_profile SET model = ?, updated_at = ?
|
||
WHERE id = ? AND (model IS NULL OR model = 'llama3.2')
|
||
""",
|
||
(_BUILTIN_OLLAMA_MODEL_SEED, utc_now_iso(), LLM_BUILTIN_OLLAMA_PROFILE_ID),
|
||
)
|
||
conn.execute(
|
||
"CREATE INDEX IF NOT EXISTS idx_chat_message_session_id_id ON chat_message(session_id, id)"
|
||
)
|
||
# 若历史上曾在未开启 foreign_keys 的连接里删除过 ``chat_session``，或表创建时未带外键，
|
||
# 会留下指向已不存在会话的行。每次初始化库时做一次幂等清理。
|
||
self._prune_rows_for_missing_chat_session(conn)
|
||
|
||
self._prune_orphan_chat_message_and_tool_log_outside_db_transaction()
|
||
|
||
def _init_db_postgresql(self) -> None:
|
||
"""PostgreSQL: tables from Alembic migration; run seeds and housekeeping."""
|
||
with self._connect() as conn:
|
||
self._seed_builtin_llm_profiles(conn)
|
||
self._seed_default_permissions(conn)
|
||
conn.execute(
|
||
"""
|
||
UPDATE llm_profile SET model = ?, updated_at = ?
|
||
WHERE id = ? AND (model IS NULL OR model = 'llama3.2')
|
||
""",
|
||
(_BUILTIN_OLLAMA_MODEL_SEED, utc_now_iso(), LLM_BUILTIN_OLLAMA_PROFILE_ID),
|
||
)
|
||
self._prune_rows_for_missing_chat_session(conn)
|
||
|
||
# PostgreSQL: ``chat_message`` / ``tool_log`` reference ``chat_session`` with FK.
|
||
# Do **not** run :meth:`_prune_orphan_chat_message_and_tool_log_outside_db_transaction` on
|
||
# every process start: the legacy ``NOT IN (SELECT id FROM chat_session)`` housekeeping
|
||
# can race concurrent inserts or interact badly with NULL / visibility, making admin chat
|
||
# ``loadMessages`` return empty right after a turn (bubbles flash then vanish).
|
||
|
||
def _prune_orphan_chat_message_and_tool_log_outside_db_transaction(self) -> None:
|
||
"""Remove orphan ``chat_message`` / ``tool_log`` rows (must not run inside raw ``_connect`` on SQLite).
|
||
|
||
SQLAlchemy uses its own pooled connection; running these deletes while a raw sqlite3 write
|
||
transaction is open would deadlock with ``database is locked``.
|
||
"""
|
||
self._chat_messages_repo().delete_messages_where_session_missing()
|
||
self._tool_log_queries_repo().delete_tool_logs_where_session_missing()
|
||
|
||
def _prune_rows_for_missing_chat_session(self, conn: Any) -> None:
|
||
sid_alive = "(SELECT id FROM chat_session)"
|
||
conn.execute(
|
||
"DELETE FROM oclaw_attempt WHERE run_id NOT IN (SELECT run_id FROM oclaw_run) "
|
||
f"OR run_id IN (SELECT run_id FROM oclaw_run WHERE session_id NOT IN {sid_alive})"
|
||
)
|
||
conn.execute(f"DELETE FROM oclaw_run WHERE session_id NOT IN {sid_alive}")
|
||
conn.execute(f"DELETE FROM oclaw_task WHERE session_id NOT IN {sid_alive}")
|
||
conn.execute(f"DELETE FROM attachment_acl WHERE session_id NOT IN {sid_alive}")
|
||
|
||
def _seed_builtin_llm_profiles(self, conn: Any) -> None:
|
||
ts = utc_now_iso()
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO llm_profile
|
||
(id, name, mode, model, base_url, api_key, updated_at, is_builtin, hide_in_ui, owner_user_id)
|
||
VALUES (?, ?, 'ollama', ?, ?, NULL, ?, 1, 0, NULL)
|
||
""",
|
||
(
|
||
LLM_BUILTIN_OLLAMA_PROFILE_ID,
|
||
"本地 Ollama（默认）",
|
||
_BUILTIN_OLLAMA_MODEL_SEED,
|
||
_BUILTIN_OLLAMA_BASE_SEED,
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def _seed_default_permissions(self, conn: Any) -> None:
|
||
ts = utc_now_iso()
|
||
defaults = {
|
||
"owner": {
|
||
"admin:read",
|
||
"admin:user:read",
|
||
"admin:user:write",
|
||
"admin:user:delete",
|
||
"admin:tenant:read",
|
||
"admin:tenant:write",
|
||
"admin:workspace_paths:read",
|
||
"admin:workspace_paths:write",
|
||
"admin:memory:write",
|
||
"admin:runtime:write",
|
||
},
|
||
"admin": {
|
||
"admin:read",
|
||
"admin:user:read",
|
||
"admin:user:write",
|
||
"admin:user:delete",
|
||
"admin:tenant:read",
|
||
"admin:tenant:write",
|
||
"admin:workspace_paths:read",
|
||
"admin:workspace_paths:write",
|
||
"admin:memory:write",
|
||
"admin:runtime:write",
|
||
},
|
||
"member": {
|
||
"admin:read",
|
||
"admin:tenant:read",
|
||
"admin:workspace_paths:read",
|
||
"admin:workspace_paths:write",
|
||
},
|
||
"guest": {"admin:read"},
|
||
}
|
||
for role, permissions in defaults.items():
|
||
for perm in permissions:
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO role_permission(role, permission, created_at)
|
||
VALUES (?, ?, ?)
|
||
""",
|
||
(role, perm, ts),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO llm_profile
|
||
(id, name, mode, model, base_url, api_key, updated_at, is_builtin, hide_in_ui, owner_user_id)
|
||
VALUES (?, ?, 'rule', NULL, NULL, NULL, ?, 1, 1, NULL)
|
||
""",
|
||
(LLM_BUILTIN_RULE_PROFILE_ID, "内置规则兜底", ts),
|
||
)
|
||
|
||
def create_session(self, title: str) -> ChatSession:
|
||
session_id = uuid.uuid4().hex
|
||
created_at = utc_now_iso()
|
||
self._chat_sessions_repo().insert_chat_session(
|
||
session_id=session_id, title=title, created_at=created_at
|
||
)
|
||
return ChatSession(id=session_id, title=title, created_at=created_at, last_message_at=None)
|
||
|
||
def create_session_for_user(self, *, title: str, tenant_id: str, user_id: str) -> ChatSession:
|
||
s = self.create_session(title)
|
||
try:
|
||
self._ui_session_owner_repo().upsert_replace(
|
||
session_id=str(s.id),
|
||
tenant_id=str(tenant_id),
|
||
user_id=str(user_id),
|
||
created_at=utc_now_iso(),
|
||
)
|
||
except IntegrityError:
|
||
# ``chat_session`` is already committed; without ``ui_session_owner`` the session is invisible
|
||
# to ``list_sessions_for_user`` / ``get_session_for_user`` (INNER JOIN). PostgreSQL enforces
|
||
# FK from ``ui_session_owner`` to ``tenant`` / ``app_user``; a failed owner row leaves a
|
||
# "ghost" session that looks like PG-specific data loss vs SQLite (weaker FK history).
|
||
try:
|
||
self.delete_session(str(s.id))
|
||
except Exception:
|
||
pass
|
||
raise
|
||
return s
|
||
|
||
def ensure_ui_session_owner(self, *, session_id: str, tenant_id: str, user_id: str) -> None:
|
||
self._ui_session_owner_repo().insert_ignore(
|
||
session_id=str(session_id),
|
||
tenant_id=str(tenant_id),
|
||
user_id=str(user_id),
|
||
created_at=utc_now_iso(),
|
||
)
|
||
|
||
def get_session(self, session_id: str) -> Optional[ChatSession]:
|
||
return self._chat_sessions_repo().fetch_chat_session_by_id(session_id=session_id)
|
||
|
||
def get_ui_session_owner(self, *, session_id: str) -> dict[str, Any] | None:
|
||
sid = str(session_id or "").strip()
|
||
if not sid:
|
||
return None
|
||
row = self._ui_session_owner_repo().fetch_by_session_id(session_id=sid)
|
||
if not row:
|
||
return None
|
||
return {
|
||
"tenant_id": str(row["tenant_id"] or ""),
|
||
"user_id": str(row["user_id"] or ""),
|
||
"created_at": str(row["created_at"] or ""),
|
||
}
|
||
|
||
def backfill_orphan_chat_sessions_for_user(self, *, tenant_id: str, user_id: str) -> int:
|
||
"""将**当前库中所有**尚无 ``ui_session_owner`` 的 ``chat_session`` 归属到指定用户。
|
||
|
||
**危险**：多用户环境下会把他人历史会话一并划给该用户，造成会话串用。
|
||
已从 HTTP 列表接口移除自动调用；仅保留供单租户数据修复时在 Python 控制台等场景**显式**调用。
|
||
"""
|
||
ts = utc_now_iso()
|
||
return self._ui_session_owner_repo().backfill_orphan_sessions_for_user(
|
||
tenant_id=str(tenant_id), user_id=str(user_id), default_created_at=ts
|
||
)
|
||
|
||
def get_session_for_user(self, *, session_id: str, tenant_id: str, user_id: str) -> Optional[ChatSession]:
|
||
return self._chat_sessions_repo().fetch_chat_session_for_user(
|
||
session_id=session_id, tenant_id=tenant_id, user_id=user_id
|
||
)
|
||
|
||
def list_sessions(
|
||
self,
|
||
limit: int | None = None,
|
||
offset: int = 0,
|
||
) -> list[ChatSession]:
|
||
return self._chat_sessions_repo().list_chat_sessions_global(limit=limit, offset=int(offset))
|
||
|
||
def list_sessions_for_user(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
limit: int | None = None,
|
||
offset: int = 0,
|
||
) -> list[ChatSession]:
|
||
return self._chat_sessions_repo().list_chat_sessions_for_user(
|
||
tenant_id=tenant_id,
|
||
user_id=user_id,
|
||
limit=limit,
|
||
offset=int(offset),
|
||
)
|
||
|
||
def count_sessions(self) -> int:
|
||
return self._chat_sessions_repo().count_chat_sessions_global()
|
||
|
||
def get_sessions_list_meta(self) -> SessionsListMeta:
|
||
return self._chat_sessions_repo().sessions_list_meta_global()
|
||
|
||
def get_sessions_list_meta_for_user(self, *, tenant_id: str, user_id: str) -> SessionsListMeta:
|
||
return self._chat_sessions_repo().sessions_list_meta_for_user(
|
||
tenant_id=tenant_id, user_id=user_id
|
||
)
|
||
|
||
def list_sessions_for_tenant(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
limit: int | None = None,
|
||
offset: int = 0,
|
||
) -> list[ChatSession]:
|
||
"""All chat sessions that belong to ``tenant_id`` via ``ui_session_owner`` (any user)."""
|
||
return self._chat_sessions_repo().list_chat_sessions_for_tenant(
|
||
tenant_id=tenant_id, limit=limit, offset=int(offset)
|
||
)
|
||
|
||
def get_sessions_list_meta_for_tenant(self, *, tenant_id: str) -> SessionsListMeta:
|
||
return self._chat_sessions_repo().sessions_list_meta_for_tenant(tenant_id=tenant_id)
|
||
|
||
def get_session_in_tenant(self, *, session_id: str, tenant_id: str) -> Optional[ChatSession]:
|
||
"""Session exists and is linked to this tenant (``administrator`` global browse)."""
|
||
return self._chat_sessions_repo().fetch_chat_session_in_tenant(
|
||
session_id=session_id, tenant_id=tenant_id
|
||
)
|
||
|
||
@staticmethod
|
||
def _attachments_contain_attachment_id(raw_attachments: Any, *, attachment_id: str) -> bool:
|
||
aid = str(attachment_id or "").strip()
|
||
if not aid:
|
||
return False
|
||
obj = raw_attachments
|
||
if isinstance(raw_attachments, str):
|
||
s = str(raw_attachments or "").strip()
|
||
if not s:
|
||
return False
|
||
try:
|
||
obj = json.loads(s)
|
||
except Exception:
|
||
return False
|
||
if isinstance(obj, dict):
|
||
items = [obj]
|
||
elif isinstance(obj, list):
|
||
items = obj
|
||
else:
|
||
return False
|
||
for it in items:
|
||
if not isinstance(it, dict):
|
||
continue
|
||
if str(it.get("attachment_id") or "").strip() == aid:
|
||
return True
|
||
return False
|
||
|
||
def attachment_referenced_by_user(self, *, tenant_id: str, user_id: str, attachment_id: str, scan_limit: int = 2000) -> bool:
|
||
aid = str(attachment_id or "").strip()
|
||
tid = str(tenant_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
if not aid or not tid or not uid:
|
||
return False
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT m.attachments
|
||
FROM chat_message m
|
||
INNER JOIN ui_session_owner o ON o.session_id = m.session_id
|
||
WHERE o.tenant_id = ? AND o.user_id = ? AND m.attachments IS NOT NULL AND m.attachments <> ''
|
||
ORDER BY m.id DESC
|
||
LIMIT ?
|
||
""",
|
||
(tid, uid, int(max(1, scan_limit))),
|
||
).fetchall()
|
||
for r in rows:
|
||
if self._attachments_contain_attachment_id(r["attachments"], attachment_id=aid):
|
||
return True
|
||
return False
|
||
|
||
def attachment_referenced_in_tenant(self, *, tenant_id: str, attachment_id: str, scan_limit: int = 4000) -> bool:
|
||
aid = str(attachment_id or "").strip()
|
||
tid = str(tenant_id or "").strip()
|
||
if not aid or not tid:
|
||
return False
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT m.attachments
|
||
FROM chat_message m
|
||
INNER JOIN ui_session_owner o ON o.session_id = m.session_id
|
||
WHERE o.tenant_id = ? AND m.attachments IS NOT NULL AND m.attachments <> ''
|
||
ORDER BY m.id DESC
|
||
LIMIT ?
|
||
""",
|
||
(tid, int(max(1, scan_limit))),
|
||
).fetchall()
|
||
for r in rows:
|
||
if self._attachments_contain_attachment_id(r["attachments"], attachment_id=aid):
|
||
return True
|
||
return False
|
||
|
||
def link_attachment_acl(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
session_id: str,
|
||
attachment_id: str,
|
||
source: str,
|
||
) -> None:
|
||
tid = str(tenant_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
sid = str(session_id or "").strip()
|
||
aid = str(attachment_id or "").strip().lower()
|
||
src = str(source or "").strip() or "unknown"
|
||
if not tid or not uid or not sid or not aid:
|
||
return
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO attachment_acl
|
||
(attachment_id, tenant_id, user_id, session_id, source, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(aid, tid, uid, sid, src, ts),
|
||
)
|
||
|
||
@staticmethod
|
||
def _attachment_ids_for_acl_from_payload_items(items: list[dict[str, Any]]) -> list[str]:
|
||
"""Collect stable attachment ids from message attachment JSON (for ACL rows)."""
|
||
out: list[str] = []
|
||
seen: set[str] = set()
|
||
ref_types = {"image_ref", "video_ref", "text_ref", "binary_ref"}
|
||
relay_re = re.compile(r"^relay://attachments/[^/]+/([a-f0-9]{8,64})$", re.IGNORECASE)
|
||
for a in items:
|
||
if not isinstance(a, dict):
|
||
continue
|
||
typ = str(a.get("type") or "").strip().lower()
|
||
aid = str(a.get("attachment_id") or a.get("attachmentId") or "").strip().lower()
|
||
if typ in ref_types and aid:
|
||
if aid not in seen:
|
||
seen.add(aid)
|
||
out.append(aid)
|
||
continue
|
||
if typ != "relay_pointer":
|
||
continue
|
||
if not aid:
|
||
uri = str(a.get("pointer_uri") or "").strip()
|
||
m = relay_re.match(uri)
|
||
if m:
|
||
aid = str(m.group(1) or "").strip().lower()
|
||
if aid and aid not in seen:
|
||
seen.add(aid)
|
||
out.append(aid)
|
||
return out
|
||
|
||
def sync_attachment_acl_from_chat_message_attachments(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
role: str,
|
||
event_type: str | None,
|
||
attachments: Any,
|
||
) -> None:
|
||
"""Best-effort: link chat_message attachments to session owner so strict ACL downloads work.
|
||
|
||
Tool results already link in runtime; assistant rows (image specialist, inbound, etc.)
|
||
historically did not, which breaks ``AIA_ATTACHMENT_ACL_STRICT=1``.
|
||
"""
|
||
sid = str(session_id or "").strip()
|
||
if not sid or attachments is None:
|
||
return
|
||
owner = self.get_ui_session_owner(session_id=sid) or {}
|
||
tid = str(owner.get("tenant_id") or "").strip()
|
||
uid = str(owner.get("user_id") or "").strip()
|
||
if not tid or not uid:
|
||
return
|
||
items: list[dict[str, Any]] = []
|
||
if isinstance(attachments, list):
|
||
items = [x for x in attachments if isinstance(x, dict)]
|
||
elif isinstance(attachments, dict):
|
||
items = [attachments]
|
||
else:
|
||
return
|
||
ids = self._attachment_ids_for_acl_from_payload_items(items)
|
||
if not ids:
|
||
return
|
||
r = str(role or "").strip() or "-"
|
||
ev = str(event_type or "").strip() or "-"
|
||
src = f"chat_message:{r}:{ev}"[:240]
|
||
for aid in ids:
|
||
self.link_attachment_acl(
|
||
tenant_id=tid,
|
||
user_id=uid,
|
||
session_id=sid,
|
||
attachment_id=aid,
|
||
source=src,
|
||
)
|
||
|
||
def attachment_acl_allows_user(self, *, tenant_id: str, user_id: str, attachment_id: str) -> bool:
|
||
tid = str(tenant_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
aid = str(attachment_id or "").strip().lower()
|
||
if not tid or not uid or not aid:
|
||
return False
|
||
with self._connect() as conn:
|
||
r = conn.execute(
|
||
"""
|
||
SELECT 1
|
||
FROM attachment_acl
|
||
WHERE tenant_id = ? AND user_id = ? AND attachment_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(tid, uid, aid),
|
||
).fetchone()
|
||
return bool(r)
|
||
|
||
def attachment_acl_allows_tenant(self, *, tenant_id: str, attachment_id: str) -> bool:
|
||
tid = str(tenant_id or "").strip()
|
||
aid = str(attachment_id or "").strip().lower()
|
||
if not tid or not aid:
|
||
return False
|
||
with self._connect() as conn:
|
||
r = conn.execute(
|
||
"""
|
||
SELECT 1
|
||
FROM attachment_acl
|
||
WHERE tenant_id = ? AND attachment_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(tid, aid),
|
||
).fetchone()
|
||
return bool(r)
|
||
|
||
def backfill_attachment_acl_from_messages(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
limit_messages: int = 50_000,
|
||
) -> dict[str, Any]:
|
||
"""Best-effort backfill: scan chat_message.attachments and populate attachment_acl.
|
||
|
||
This is intended for one-off migration / operator maintenance.
|
||
"""
|
||
tid = str(tenant_id or "").strip()
|
||
lim = max(1, int(limit_messages))
|
||
if not tid:
|
||
return {"ok": False, "error": "tenant_id_required"}
|
||
inserted = 0
|
||
scanned_msgs = 0
|
||
scanned_atts = 0
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT m.session_id, m.attachments, o.user_id
|
||
FROM chat_message m
|
||
INNER JOIN ui_session_owner o ON o.session_id = m.session_id
|
||
WHERE o.tenant_id = ? AND m.attachments IS NOT NULL AND m.attachments <> ''
|
||
ORDER BY m.id DESC
|
||
LIMIT ?
|
||
""",
|
||
(tid, lim),
|
||
).fetchall()
|
||
ts = utc_now_iso()
|
||
for r in rows:
|
||
scanned_msgs += 1
|
||
sid = str(r["session_id"] or "").strip()
|
||
uid = str(r["user_id"] or "").strip()
|
||
if not sid or not uid:
|
||
continue
|
||
try:
|
||
obj = json.loads(str(r["attachments"] or ""))
|
||
except Exception:
|
||
continue
|
||
items = obj if isinstance(obj, list) else ([obj] if isinstance(obj, dict) else [])
|
||
for a in items:
|
||
if not isinstance(a, dict):
|
||
continue
|
||
scanned_atts += 1
|
||
aid = str(a.get("attachment_id") or "").strip().lower()
|
||
if not aid:
|
||
continue
|
||
src = "backfill:chat_message"
|
||
cur = conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO attachment_acl
|
||
(attachment_id, tenant_id, user_id, session_id, source, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(aid, tid, uid, sid, src, ts),
|
||
)
|
||
inserted += int(cur.rowcount or 0)
|
||
return {
|
||
"ok": True,
|
||
"tenant_id": tid,
|
||
"scanned_messages": int(scanned_msgs),
|
||
"scanned_attachments": int(scanned_atts),
|
||
"inserted": int(inserted),
|
||
}
|
||
|
||
def list_admin_sessions(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
user_id: str | None = None,
|
||
q: str | None = None,
|
||
active_only: bool = False,
|
||
active_window_minutes: int = 30,
|
||
limit: int = 100,
|
||
offset: int = 0,
|
||
) -> tuple[int, list[dict[str, Any]]]:
|
||
tid = str(tenant_id or "").strip()
|
||
if not tid:
|
||
return 0, []
|
||
uid = str(user_id or "").strip()
|
||
q_text = str(q or "").strip().lower()
|
||
win = max(1, int(active_window_minutes))
|
||
cutoff = (datetime.now(timezone.utc) - timedelta(minutes=win)).isoformat()
|
||
lim = max(1, min(int(limit), 500))
|
||
off = max(0, int(offset))
|
||
|
||
total, rows = self._chat_sessions_repo().list_admin_sessions(
|
||
tenant_id=tid,
|
||
user_id=uid or None,
|
||
search_lower=q_text or None,
|
||
active_only=active_only,
|
||
active_cutoff_iso=cutoff,
|
||
limit=lim,
|
||
offset=off,
|
||
)
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
last_at = str(r["last_message_at"] or r["created_at"] or "")
|
||
out.append(
|
||
{
|
||
**r,
|
||
"is_active_30m": bool(last_at and last_at >= cutoff),
|
||
}
|
||
)
|
||
return total, out
|
||
|
||
def list_admin_user_stats(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
q: str | None = None,
|
||
active_window_minutes: int = 30,
|
||
limit: int = 100,
|
||
offset: int = 0,
|
||
) -> tuple[int, list[dict[str, Any]], dict[str, Any]]:
|
||
tid = str(tenant_id or "").strip()
|
||
if not tid:
|
||
return 0, [], {
|
||
"total_tokens_est": 0,
|
||
"active_sessions_30m": 0,
|
||
"active_logins_30m": 0,
|
||
"users_count": 0,
|
||
}
|
||
q_text = str(q or "").strip().lower()
|
||
win = max(1, int(active_window_minutes))
|
||
cutoff = (datetime.now(timezone.utc) - timedelta(minutes=win)).isoformat()
|
||
lim = max(1, min(int(limit), 500))
|
||
off = max(0, int(offset))
|
||
|
||
pack = self._admin_user_stats_repo().fetch(
|
||
tenant_id=tid,
|
||
search_lower=q_text or None,
|
||
cutoff_iso=cutoff,
|
||
limit=lim,
|
||
offset=off,
|
||
)
|
||
total_row_c = int(pack["total_users"] or 0)
|
||
rows = pack["user_rows"]
|
||
total_active_sessions = pack["total_active_sessions_30m"]
|
||
total_active_logins = pack["total_active_logins_30m"]
|
||
|
||
token_by_user: dict[str, int] = {}
|
||
sessions_count_by_user: dict[str, int] = {}
|
||
active_sessions_by_user: dict[str, int] = {}
|
||
last_message_at_by_user: dict[str, str] = {}
|
||
active_logins_by_user: dict[str, int] = {}
|
||
last_seen_at_by_user: dict[str, str] = {}
|
||
|
||
for tr in pack["trace_rows"]:
|
||
uid = str(tr.get("user_id") or "")
|
||
if not uid:
|
||
continue
|
||
try:
|
||
payload = json.loads(tr.get("payload") or "{}")
|
||
except Exception:
|
||
payload = {}
|
||
p = int(payload.get("prompt_tokens_est") or 0)
|
||
r2 = int(payload.get("response_tokens_est") or 0)
|
||
token_by_user[uid] = int(token_by_user.get(uid, 0)) + max(0, p) + max(0, r2)
|
||
|
||
for sr in pack["sessions_count_rows"]:
|
||
sessions_count_by_user[str(sr.get("user_id") or "")] = int(sr.get("c") or 0)
|
||
|
||
for rr in pack["active_sess_rows"]:
|
||
uid = str(rr.get("user_id") or "")
|
||
active_sessions_by_user[uid] = int(rr.get("active_30m") or 0)
|
||
last_message_at_by_user[uid] = str(rr.get("last_message_at") or "")
|
||
|
||
for lr in pack["login_rows"]:
|
||
uid = str(lr.get("user_id") or "")
|
||
active_logins_by_user[uid] = int(lr.get("c") or 0)
|
||
last_seen_at_by_user[uid] = str(lr.get("last_seen_at") or "")
|
||
|
||
users: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
uid = str(r["user_id"] or "")
|
||
users.append(
|
||
{
|
||
"user_id": uid,
|
||
"username": str(r["username"] or ""),
|
||
"display_name": str(r["display_name"] or ""),
|
||
"role": str(r["role"] or ""),
|
||
"is_active": bool(int(r["is_active"] or 0)),
|
||
"total_tokens_est": int(token_by_user.get(uid, 0)),
|
||
"sessions_count": int(sessions_count_by_user.get(uid, 0)),
|
||
"active_sessions_30m": int(active_sessions_by_user.get(uid, 0)),
|
||
"active_login_30m": int(active_logins_by_user.get(uid, 0)),
|
||
"last_message_at": str(last_message_at_by_user.get(uid, "")),
|
||
"last_seen_at": str(last_seen_at_by_user.get(uid, "")),
|
||
}
|
||
)
|
||
|
||
totals = {
|
||
"total_tokens_est": int(sum(int(x.get("total_tokens_est") or 0) for x in users)),
|
||
"active_sessions_30m": int(total_active_sessions or 0),
|
||
"active_logins_30m": int(total_active_logins or 0),
|
||
"users_count": int(total_row_c or 0),
|
||
}
|
||
return int(total_row_c or 0), users, totals
|
||
|
||
def delete_session_in_tenant(self, *, session_id: str, tenant_id: str) -> bool:
|
||
"""Delete session if it belongs to tenant (used by administrator account)."""
|
||
return self._chat_sessions_repo().try_delete_chat_session_for_tenant(
|
||
session_id=session_id, tenant_id=tenant_id
|
||
)
|
||
|
||
def delete_session(self, session_id: str) -> None:
|
||
self._chat_sessions_repo().delete_chat_session_by_id(session_id=session_id)
|
||
|
||
def delete_session_for_user(self, *, session_id: str, tenant_id: str, user_id: str) -> bool:
|
||
return self._chat_sessions_repo().try_delete_chat_session_for_user(
|
||
session_id=session_id, tenant_id=tenant_id, user_id=user_id
|
||
)
|
||
|
||
def delete_message(self, *, session_id: str, message_id: int) -> bool:
|
||
return self._chat_messages_repo().delete_message_and_refresh_session(
|
||
session_id=session_id, message_id=message_id
|
||
)
|
||
|
||
def add_message(
|
||
self,
|
||
session_id: str,
|
||
role: str,
|
||
content: str,
|
||
tool_calls: Any | None = None,
|
||
attachments: Any | None = None,
|
||
turn_uuid: str | None = None,
|
||
event_type: str | None = None,
|
||
event_payload: Any | None = None,
|
||
timestamp: str | None = None,
|
||
) -> ChatMessage:
|
||
ts = timestamp or utc_now_iso()
|
||
tool_calls_text = None
|
||
if tool_calls is not None:
|
||
if isinstance(tool_calls, str):
|
||
tool_calls_text = tool_calls
|
||
else:
|
||
tool_calls_text = json.dumps(
|
||
scrub_nul_bytes_from_jsonable(tool_calls), ensure_ascii=False
|
||
)
|
||
attachments_text = None
|
||
if attachments is not None:
|
||
attachments_text = json.dumps(
|
||
scrub_nul_bytes_from_jsonable(attachments), ensure_ascii=False
|
||
)
|
||
event_payload_text = None
|
||
if event_payload is not None:
|
||
if isinstance(event_payload, str):
|
||
event_payload_text = event_payload
|
||
else:
|
||
event_payload_text = json.dumps(
|
||
scrub_nul_bytes_from_jsonable(event_payload), ensure_ascii=False, default=str
|
||
)
|
||
turn_uuid_text = str(turn_uuid or "").strip() or None
|
||
event_type_text = str(event_type or "").strip() or None
|
||
# PostgreSQL TEXT rejects NUL; SQLite accepts it — strip before SA insert.
|
||
content_clean = str(scrub_nul_bytes_from_text(str(content)) or "")
|
||
tool_calls_text = scrub_nul_bytes_from_text(tool_calls_text)
|
||
attachments_text = scrub_nul_bytes_from_text(attachments_text)
|
||
event_payload_text = scrub_nul_bytes_from_text(event_payload_text)
|
||
turn_uuid_text = scrub_nul_bytes_from_text(turn_uuid_text)
|
||
event_type_text = scrub_nul_bytes_from_text(event_type_text)
|
||
if turn_uuid_text is not None and not str(turn_uuid_text).strip():
|
||
turn_uuid_text = None
|
||
if event_type_text is not None and not str(event_type_text).strip():
|
||
event_type_text = None
|
||
msg_id = self._chat_messages_repo().insert_message_and_touch_session(
|
||
session_id=str(session_id if session_id is not None else ""),
|
||
role=str(role),
|
||
content=content_clean,
|
||
tool_calls=tool_calls_text,
|
||
attachments=attachments_text,
|
||
turn_uuid=turn_uuid_text,
|
||
event_type=event_type_text,
|
||
event_payload=event_payload_text,
|
||
timestamp=str(ts),
|
||
)
|
||
if _chat_message_persist_log_enabled():
|
||
_LOG.warning(
|
||
"chat_message_persisted backend=%s session_id=%s message_id=%s role=%s event_type=%s "
|
||
"content_len=%d has_tool_calls=%s",
|
||
"postgresql" if self._use_pg else "sqlite",
|
||
str(session_id or "").strip(),
|
||
int(msg_id),
|
||
str(role),
|
||
str(event_type_text or ""),
|
||
len(str(content_clean or "")),
|
||
bool(tool_calls_text),
|
||
)
|
||
try:
|
||
self.sync_attachment_acl_from_chat_message_attachments(
|
||
session_id=str(session_id or "").strip(),
|
||
role=str(role or ""),
|
||
event_type=str(event_type or "").strip() or None,
|
||
attachments=attachments,
|
||
)
|
||
except Exception:
|
||
pass
|
||
return ChatMessage(
|
||
id=msg_id,
|
||
session_id=session_id,
|
||
role=role,
|
||
content=content_clean,
|
||
tool_calls=tool_calls_text,
|
||
attachments=attachments_text,
|
||
turn_uuid=turn_uuid_text,
|
||
event_type=event_type_text,
|
||
event_payload=event_payload_text,
|
||
timestamp=ts,
|
||
)
|
||
|
||
def update_message_content(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
message_id: int,
|
||
content: str,
|
||
event_payload: Any | None = None,
|
||
) -> bool:
|
||
sid = str(session_id or "").strip()
|
||
if not sid:
|
||
return False
|
||
mid = int(message_id or 0)
|
||
if mid <= 0:
|
||
return False
|
||
event_payload_text = None
|
||
if event_payload is not None:
|
||
if isinstance(event_payload, str):
|
||
event_payload_text = event_payload
|
||
else:
|
||
event_payload_text = json.dumps(
|
||
scrub_nul_bytes_from_jsonable(event_payload), ensure_ascii=False, default=str
|
||
)
|
||
event_payload_text = scrub_nul_bytes_from_text(event_payload_text)
|
||
content_u = str(scrub_nul_bytes_from_text(str(content or "")) or "")
|
||
return self._chat_messages_repo().update_message_content(
|
||
session_id=sid,
|
||
message_id=mid,
|
||
content=content_u,
|
||
event_payload_text=event_payload_text,
|
||
)
|
||
|
||
def get_messages(self, session_id: str, limit: int = 200) -> list[ChatMessage]:
|
||
"""返回最近 ``limit`` 条消息，顺序为时间正序（窗口内最早的一条在前）。"""
|
||
if limit <= 0:
|
||
return []
|
||
sid = str(session_id or "").strip()
|
||
if not sid:
|
||
return []
|
||
lim = max(1, min(int(limit), 2000))
|
||
return self._chat_messages_repo().get_messages_recent_asc(session_id=sid, limit=lim)
|
||
|
||
def get_messages_after_id(self, *, session_id: str, after_id: int, limit: int = 200) -> list[ChatMessage]:
|
||
"""Return messages with id > after_id in ASC order (bounded by limit)."""
|
||
sid = str(session_id or "").strip()
|
||
if not sid:
|
||
return []
|
||
aid = int(after_id or 0)
|
||
lim = max(1, min(int(limit), 2000))
|
||
return self._chat_messages_repo().get_messages_after_id(
|
||
session_id=sid, after_id=aid, limit=lim
|
||
)
|
||
|
||
def add_tool_log(
|
||
self,
|
||
session_id: str,
|
||
tool_name: str,
|
||
args: dict[str, Any],
|
||
result: Any,
|
||
specialist: str | None = None,
|
||
timestamp: str | None = None,
|
||
duration_ms: int | None = None,
|
||
) -> None:
|
||
ts = timestamp or utc_now_iso()
|
||
raw_cap = str(self.get_setting("AIA_TOOL_LOG_MAX_CHARS") or "").strip()
|
||
if not raw_cap:
|
||
raw_cap = str(os.getenv("AIA_TOOL_LOG_MAX_CHARS") or "").strip()
|
||
cap = 200_000
|
||
if raw_cap.isdigit():
|
||
cap = max(20_000, min(int(raw_cap), 2_000_000))
|
||
args_capped = self._cap_json_for_log(args, max_chars=cap, keep_keys=())
|
||
result_capped = self._cap_json_for_log(result, max_chars=cap, keep_keys=("ok", "error_code", "error"))
|
||
self._tool_log_queries_repo().insert_tool_log(
|
||
session_id=str(session_id),
|
||
tool_name=str(tool_name),
|
||
specialist=str(specialist or ""),
|
||
args=json.dumps(args_capped, ensure_ascii=False, default=str),
|
||
result=json.dumps(result_capped, ensure_ascii=False, default=str),
|
||
timestamp=str(ts),
|
||
duration_ms=duration_ms,
|
||
)
|
||
|
||
def get_tool_logs(self, session_id: str, limit: int = 200) -> list[dict[str, Any]]:
|
||
rows = self._tool_log_queries_repo().list_tool_logs_asc(
|
||
session_id=str(session_id), limit=max(1, int(limit))
|
||
)
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
out.append(
|
||
{
|
||
"tool_name": r["tool_name"],
|
||
"specialist": str(r["specialist"] or ""),
|
||
"args": json.loads(r["args"]),
|
||
"result": json.loads(r["result"]),
|
||
"timestamp": r["timestamp"],
|
||
"duration_ms": r["duration_ms"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def list_session_tool_health(self, *, session_id: str | None = None, limit: int = 80) -> list[dict[str, Any]]:
|
||
lim = max(1, int(limit))
|
||
sid = str(session_id or "").strip() or None
|
||
raw = self._session_tool_health_repo().list_session_tool_health(session_id=sid, limit=lim)
|
||
out: list[dict[str, Any]] = []
|
||
for r in raw:
|
||
tool_count = int(r["tool_count"] or 0)
|
||
assistant_count = int(r["assistant_count"] or 0)
|
||
unhealthy = assistant_count > 0 and tool_count == 0
|
||
out.append(
|
||
{
|
||
"session_id": str(r["session_id"] or ""),
|
||
"title": str(r["title"] or ""),
|
||
"last_message_at": str(r["last_message_at"] or ""),
|
||
"user_count": int(r["user_count"] or 0),
|
||
"assistant_count": assistant_count,
|
||
"tool_count": tool_count,
|
||
"mcp_tool_count": int(r["mcp_tool_count"] or 0),
|
||
"last_tool_at": str(r["last_tool_at"] or ""),
|
||
"status": "warn_no_tool_calls" if unhealthy else "ok",
|
||
}
|
||
)
|
||
return out
|
||
|
||
def move_tool_logs_to_session(self, *, from_session_id: str, to_session_id: str) -> int:
|
||
src = str(from_session_id or "").strip()
|
||
dst = str(to_session_id or "").strip()
|
||
if not src or not dst or src == dst:
|
||
return 0
|
||
return self._tool_log_queries_repo().move_tool_logs_between_sessions(
|
||
from_session_id=src, to_session_id=dst
|
||
)
|
||
|
||
def list_mcp_tool_usage_summary(self, *, limit: int = 200) -> list[dict[str, Any]]:
|
||
rows = self._tool_log_queries_repo().list_mcp_tool_usage_summary(limit=max(1, int(limit)))
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
tool_name = str(r["tool_name"] or "")
|
||
parts = tool_name.split("__", 2)
|
||
server_id = parts[1] if len(parts) >= 3 else ""
|
||
tool_short = parts[2] if len(parts) >= 3 else tool_name
|
||
out.append(
|
||
{
|
||
"tool_name": tool_name,
|
||
"server_id": server_id,
|
||
"mcp_tool_name": tool_short,
|
||
"specialist": str(r["specialist"] or ""),
|
||
"count": int(r["n"] or 0),
|
||
"last_ts": str(r["last_ts"] or ""),
|
||
}
|
||
)
|
||
return out
|
||
|
||
def list_mcp_tool_aggregate_usage(self) -> dict[str, dict[str, Any]]:
|
||
"""Cross-session counts and last call time per MCP tool name (``mcp__*``)."""
|
||
rows = self._tool_log_queries_repo().list_mcp_tool_aggregate_usage()
|
||
out: dict[str, dict[str, Any]] = {}
|
||
for r in rows:
|
||
tn = str(r["tool_name"] or "")
|
||
if not tn:
|
||
continue
|
||
out[tn] = {"count": int(r["n"] or 0), "last_ts": str(r["last_ts"] or "")}
|
||
return out
|
||
|
||
def list_mcp_tool_call_logs(self, *, server_id: str | None = None, limit: int = 200) -> list[dict[str, Any]]:
|
||
rows = self._tool_log_queries_repo().list_mcp_tool_call_logs(
|
||
server_id=server_id, limit=max(1, int(limit))
|
||
)
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
args = json.loads(r["args"] or "{}")
|
||
except Exception:
|
||
args = {}
|
||
try:
|
||
result = json.loads(r["result"] or "{}")
|
||
except Exception:
|
||
result = {}
|
||
tool_name = str(r["tool_name"] or "")
|
||
parts = tool_name.split("__", 2)
|
||
row_server_id = parts[1] if len(parts) >= 3 else ""
|
||
row_tool_name = parts[2] if len(parts) >= 3 else tool_name
|
||
out.append(
|
||
{
|
||
"session_id": str(r["session_id"] or ""),
|
||
"tool_name": tool_name,
|
||
"server_id": row_server_id,
|
||
"mcp_tool_name": row_tool_name,
|
||
"specialist": str(r["specialist"] or ""),
|
||
"args": args,
|
||
"result": result,
|
||
"timestamp": str(r["timestamp"] or ""),
|
||
"duration_ms": int(r["duration_ms"] or 0),
|
||
}
|
||
)
|
||
return out
|
||
|
||
def delete_mcp_server(self, *, server_id: str) -> dict[str, int]:
|
||
sid = str(server_id or "").strip()
|
||
if not sid:
|
||
return {"registry": 0, "tools": 0, "health": 0, "install_logs": 0}
|
||
with self._connect() as conn:
|
||
cur_tools = conn.execute("DELETE FROM mcp_server_tool WHERE server_id = ?", (sid,))
|
||
cur_health = conn.execute("DELETE FROM mcp_server_health WHERE server_id = ?", (sid,))
|
||
cur_logs = conn.execute("DELETE FROM mcp_server_installation WHERE server_id = ?", (sid,))
|
||
cur_registry = conn.execute("DELETE FROM mcp_server_registry WHERE server_id = ?", (sid,))
|
||
return {
|
||
"registry": int(cur_registry.rowcount or 0),
|
||
"tools": int(cur_tools.rowcount or 0),
|
||
"health": int(cur_health.rowcount or 0),
|
||
"install_logs": int(cur_logs.rowcount or 0),
|
||
}
|
||
|
||
def count_messages(self, session_id: str) -> int:
|
||
return self._chat_messages_repo().count_messages(session_id=session_id)
|
||
|
||
def get_session_messages_meta(self, session_id: str) -> SessionMessagesMeta:
|
||
return self._chat_messages_repo().session_messages_meta(session_id=session_id)
|
||
|
||
def get_last_message_id(self, session_id: str) -> int | None:
|
||
return self._chat_messages_repo().last_message_id(session_id=session_id)
|
||
|
||
def ensure_default_session(self) -> ChatSession:
|
||
sessions = self.list_sessions(limit=1, offset=0)
|
||
if sessions:
|
||
return sessions[0]
|
||
lang = (os.getenv("AIA_ASSISTANT_LANG") or "").strip().lower()
|
||
title = "New Chat" if lang.startswith("en") else "新会话"
|
||
return self.create_session(title)
|
||
|
||
def rename_session(self, session_id: str, title: str) -> None:
|
||
self._chat_sessions_repo().rename_chat_session(session_id=session_id, title=title)
|
||
|
||
def trim_messages(self, session_id: str, keep_last: int) -> None:
|
||
if keep_last <= 0:
|
||
self.delete_session(session_id)
|
||
return
|
||
self._chat_messages_repo().trim_messages_keep_last(session_id=session_id, keep_last=keep_last)
|
||
|
||
def fork_session(self, source_session_id: str, up_to_message_id: int, title: str) -> ChatSession:
|
||
"""将 `id <= up_to_message_id` 的消息复制到新会话，并重映射 tool 的 assistant_message_id。"""
|
||
self._chat_messages_repo().fork_assert_anchor(
|
||
source_session_id=source_session_id,
|
||
up_to_message_id=up_to_message_id,
|
||
)
|
||
new_sess = self.create_session(title)
|
||
self._chat_messages_repo().fork_copy_messages_to_session(
|
||
source_session_id=source_session_id,
|
||
up_to_message_id=up_to_message_id,
|
||
new_session_id=new_sess.id,
|
||
)
|
||
return self.get_session(new_sess.id) or new_sess
|
||
|
||
def set_setting(self, key: str, value: str) -> None:
|
||
ts = utc_now_iso()
|
||
self._app_settings_repo().upsert_plain(key=key, value=value, updated_at=ts)
|
||
|
||
def get_setting(self, key: str) -> Optional[str]:
|
||
row = self._app_settings_repo().fetch_row(key=key)
|
||
if row is None:
|
||
return None
|
||
val, is_secret = row
|
||
if is_secret != 0:
|
||
return None
|
||
return str(val)
|
||
|
||
def set_secret(self, key: str, plain_text: str) -> None:
|
||
ts = utc_now_iso()
|
||
enc = _encode_secret(plain_text)
|
||
self._app_settings_repo().upsert_secret(key=key, encoded_value=enc, updated_at=ts)
|
||
|
||
def get_secret(self, key: str) -> Optional[str]:
|
||
row = self._app_settings_repo().fetch_row(key=key)
|
||
if row is None:
|
||
return None
|
||
val, is_secret = row
|
||
if is_secret != 1:
|
||
return None
|
||
try:
|
||
return _decode_secret(str(val))
|
||
except Exception:
|
||
return None
|
||
|
||
def delete_setting(self, key: str) -> None:
|
||
self._app_settings_repo().delete_key(key=key)
|
||
|
||
def migrate_secrets_to_fernet(self) -> dict[str, int]:
|
||
"""Migrate legacy b64 secrets to a stronger scheme for both app settings and llm profiles.
|
||
|
||
- On Windows: migrates to DPAPI when available.
|
||
- On non-Windows: migrates to Fernet (requires AIA_ASSISTANT_MASTER_KEY + cryptography).
|
||
|
||
This is safe to run multiple times.
|
||
"""
|
||
if sys.platform != "win32" and _fernet() is None:
|
||
raise _CryptoError("fernet is not available; set AIA_ASSISTANT_MASTER_KEY and install cryptography")
|
||
|
||
ts = utc_now_iso()
|
||
migrated_app_settings = self._app_settings_repo().migrate_b64_secrets(
|
||
ts=ts,
|
||
decode_secret=_decode_secret,
|
||
encode_secret=_encode_secret,
|
||
predicate_new_encoding=lambda enc: enc.startswith("fernet:") or enc.startswith("dpapi:"),
|
||
)
|
||
migrated_llm_profiles = 0
|
||
with self._connect() as conn:
|
||
profs = conn.execute(
|
||
"SELECT id, api_key FROM llm_profile WHERE api_key IS NOT NULL AND api_key LIKE 'b64:%'"
|
||
).fetchall()
|
||
for r in profs:
|
||
pid = str(r["id"] or "")
|
||
v = str(r["api_key"] or "")
|
||
if not pid or not v:
|
||
continue
|
||
try:
|
||
plain = _decode_secret(v)
|
||
except Exception:
|
||
continue
|
||
enc = _encode_secret(plain)
|
||
if enc != v and (enc.startswith("fernet:") or enc.startswith("dpapi:")):
|
||
conn.execute(
|
||
"UPDATE llm_profile SET api_key = ?, updated_at = ? WHERE id = ?",
|
||
(enc, ts, pid),
|
||
)
|
||
migrated_llm_profiles += 1
|
||
|
||
return {
|
||
"migrated_app_settings": int(migrated_app_settings),
|
||
"migrated_llm_profiles": int(migrated_llm_profiles),
|
||
}
|
||
|
||
def legacy_secret_stats(self) -> dict[str, Any]:
|
||
"""Return counts of legacy b64 secrets for UI warning."""
|
||
legacy_b64_app = self._app_settings_repo().count_legacy_b64_secrets()
|
||
with self._connect() as conn:
|
||
row2 = conn.execute(
|
||
"SELECT COUNT(1) AS n FROM llm_profile WHERE api_key IS NOT NULL AND api_key LIKE 'b64:%'"
|
||
).fetchone()
|
||
return {
|
||
"legacy_b64_app_settings": int(legacy_b64_app),
|
||
"legacy_b64_llm_profiles": int((row2["n"] if row2 else 0) or 0),
|
||
}
|
||
|
||
def create_llm_profile(
|
||
self,
|
||
name: str,
|
||
mode: str = "openai",
|
||
model: str | None = None,
|
||
base_url: str | None = None,
|
||
*,
|
||
owner_user_id: str | None = None,
|
||
) -> str:
|
||
profile_id = uuid.uuid4().hex
|
||
ts = utc_now_iso()
|
||
own = str(owner_user_id).strip() if owner_user_id else None
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO llm_profile
|
||
(id, name, mode, model, base_url, api_key, updated_at, is_builtin, hide_in_ui, owner_user_id, thinking_mode_enabled, reasoning_effort)
|
||
VALUES (?, ?, ?, ?, ?, NULL, ?, 0, 0, ?, 0, '')
|
||
""",
|
||
(profile_id, name, mode, model, base_url, ts, own),
|
||
)
|
||
return profile_id
|
||
|
||
def list_llm_profiles(
|
||
self,
|
||
*,
|
||
visible_only: bool = False,
|
||
viewer_user_id: str | None = None,
|
||
viewer_username: str | None = None,
|
||
viewer_tenant_id: str | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
"""无 viewer：全局池。administrator：全部 profile。其他用户：内置 + 本人 + 用户/租户级 grant。"""
|
||
clauses: list[str] = []
|
||
params: list[Any] = []
|
||
if visible_only:
|
||
clauses.append("COALESCE(hide_in_ui, 0) = 0")
|
||
if viewer_user_id is None and viewer_username is None:
|
||
clauses.append("(COALESCE(TRIM(owner_user_id), '') = '')")
|
||
elif is_administrator_model_pool(viewer_username):
|
||
pass
|
||
elif viewer_user_id:
|
||
tid = str(viewer_tenant_id or "").strip()
|
||
if tid:
|
||
clauses.append(
|
||
"(COALESCE(is_builtin, 0) = 1 OR TRIM(COALESCE(owner_user_id, '')) = ? "
|
||
"OR id IN (SELECT profile_id FROM llm_profile_user_grant "
|
||
"WHERE tenant_id = ? AND user_id = ?) "
|
||
"OR id IN (SELECT profile_id FROM llm_profile_tenant_grant WHERE tenant_id = ?))"
|
||
)
|
||
params.extend([str(viewer_user_id), tid, str(viewer_user_id), tid])
|
||
else:
|
||
clauses.append("(COALESCE(is_builtin, 0) = 1 OR owner_user_id = ?)")
|
||
params.append(str(viewer_user_id))
|
||
else:
|
||
clauses.append("(COALESCE(TRIM(owner_user_id), '') = '')")
|
||
where_sql = (" WHERE " + " AND ".join(clauses)) if clauses else ""
|
||
sql = f"""
|
||
SELECT id, name, mode, model, base_url, api_key, updated_at,
|
||
COALESCE(is_builtin, 0) AS is_builtin,
|
||
COALESCE(hide_in_ui, 0) AS hide_in_ui,
|
||
owner_user_id,
|
||
COALESCE(thinking_mode_enabled, 0) AS thinking_mode_enabled,
|
||
COALESCE(reasoning_effort, '') AS reasoning_effort
|
||
FROM llm_profile
|
||
{where_sql}
|
||
ORDER BY COALESCE(is_builtin, 0) DESC, COALESCE(hide_in_ui, 0) ASC, updated_at DESC
|
||
"""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(sql, params).fetchall()
|
||
uid = str(viewer_user_id or "").strip()
|
||
is_admin = is_administrator_model_pool(viewer_username)
|
||
no_viewer = viewer_user_id is None and viewer_username is None
|
||
tid = str(viewer_tenant_id or "").strip()
|
||
user_grant_ids: set[str] = set()
|
||
tenant_grant_ids: set[str] = set()
|
||
if uid and tid and not is_admin and not no_viewer:
|
||
with self._connect() as conn:
|
||
for gr in conn.execute(
|
||
"SELECT profile_id FROM llm_profile_user_grant WHERE tenant_id = ? AND user_id = ?",
|
||
(tid, uid),
|
||
).fetchall():
|
||
user_grant_ids.add(str(gr["profile_id"]))
|
||
for gr in conn.execute(
|
||
"SELECT profile_id FROM llm_profile_tenant_grant WHERE tenant_id = ?",
|
||
(tid,),
|
||
).fetchall():
|
||
tenant_grant_ids.add(str(gr["profile_id"]))
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
own = str(r["owner_user_id"] or "").strip() if r["owner_user_id"] is not None else ""
|
||
is_builtin = bool(int(r["is_builtin"] or 0))
|
||
pid = str(r["id"])
|
||
if no_viewer:
|
||
mutable = True
|
||
if is_builtin:
|
||
vis = "builtin"
|
||
elif not own:
|
||
vis = "global"
|
||
else:
|
||
vis = "owned"
|
||
elif is_admin:
|
||
mutable = True
|
||
if is_builtin:
|
||
vis = "builtin"
|
||
elif not own:
|
||
vis = "global"
|
||
elif uid and own == uid:
|
||
vis = "owned"
|
||
else:
|
||
vis = "other_user"
|
||
elif is_builtin:
|
||
mutable = True
|
||
vis = "builtin"
|
||
elif uid and own == uid:
|
||
mutable = True
|
||
vis = "owned"
|
||
else:
|
||
mutable = False
|
||
if pid in user_grant_ids:
|
||
vis = "grant_user"
|
||
else:
|
||
vis = "grant_tenant"
|
||
out.append(
|
||
{
|
||
"id": r["id"],
|
||
"name": r["name"],
|
||
"mode": r["mode"],
|
||
"model": r["model"],
|
||
"base_url": r["base_url"],
|
||
"has_key": r["api_key"] is not None and str(r["api_key"]) != "",
|
||
"updated_at": r["updated_at"],
|
||
"is_builtin": is_builtin,
|
||
"hide_in_ui": bool(int(r["hide_in_ui"] or 0)),
|
||
"owner_user_id": own,
|
||
"thinking_mode_enabled": bool(int(r["thinking_mode_enabled"] or 0)),
|
||
"reasoning_effort": str(r["reasoning_effort"] or "").strip().lower(),
|
||
"mutable": mutable,
|
||
"visibility_reason": vis,
|
||
}
|
||
)
|
||
return out
|
||
|
||
def get_llm_profile(self, profile_id: str) -> Optional[dict[str, Any]]:
|
||
with self._connect() as conn:
|
||
r = conn.execute(
|
||
"""
|
||
SELECT id, name, mode, model, base_url, api_key, updated_at,
|
||
COALESCE(is_builtin, 0) AS is_builtin,
|
||
COALESCE(hide_in_ui, 0) AS hide_in_ui,
|
||
owner_user_id,
|
||
COALESCE(thinking_mode_enabled, 0) AS thinking_mode_enabled,
|
||
COALESCE(reasoning_effort, '') AS reasoning_effort
|
||
FROM llm_profile
|
||
WHERE id = ?
|
||
""",
|
||
(profile_id,),
|
||
).fetchone()
|
||
if not r:
|
||
return None
|
||
return {
|
||
"id": r["id"],
|
||
"name": r["name"],
|
||
"mode": r["mode"],
|
||
"model": r["model"],
|
||
"base_url": r["base_url"],
|
||
"has_key": r["api_key"] is not None and str(r["api_key"]) != "",
|
||
"updated_at": r["updated_at"],
|
||
"is_builtin": bool(int(r["is_builtin"] or 0)),
|
||
"hide_in_ui": bool(int(r["hide_in_ui"] or 0)),
|
||
"owner_user_id": str(r["owner_user_id"] or "").strip() if r["owner_user_id"] is not None else "",
|
||
"thinking_mode_enabled": bool(int(r["thinking_mode_enabled"] or 0)),
|
||
"reasoning_effort": str(r["reasoning_effort"] or "").strip().lower(),
|
||
}
|
||
|
||
def update_llm_profile(
|
||
self,
|
||
profile_id: str,
|
||
name: str,
|
||
mode: str,
|
||
model: str | None,
|
||
base_url: str | None,
|
||
*,
|
||
thinking_mode_enabled: bool | None = None,
|
||
reasoning_effort: str | None = None,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
think_val = None if thinking_mode_enabled is None else (1 if bool(thinking_mode_enabled) else 0)
|
||
eff = None if reasoning_effort is None else str(reasoning_effort or "").strip().lower()
|
||
if eff is not None and eff not in ("", "low", "medium", "high"):
|
||
eff = ""
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
UPDATE llm_profile
|
||
SET name = ?, mode = ?, model = ?, base_url = ?,
|
||
thinking_mode_enabled = COALESCE(?, thinking_mode_enabled),
|
||
reasoning_effort = COALESCE(?, reasoning_effort),
|
||
updated_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(name, mode, model, base_url, think_val, eff, ts, profile_id),
|
||
)
|
||
|
||
def delete_llm_profile(self, profile_id: str) -> None:
|
||
prof = self.get_llm_profile(profile_id)
|
||
if prof and prof.get("is_builtin"):
|
||
raise ValueError("cannot_delete_builtin_llm_profile")
|
||
with self._connect() as conn:
|
||
conn.execute("DELETE FROM llm_profile_user_grant WHERE profile_id = ?", (profile_id,))
|
||
conn.execute("DELETE FROM llm_profile_tenant_grant WHERE profile_id = ?", (profile_id,))
|
||
conn.execute("DELETE FROM llm_profile WHERE id = ?", (profile_id,))
|
||
|
||
def grant_llm_profile_to_user(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
profile_id: str,
|
||
user_id: str,
|
||
created_by_user_id: str | None = None,
|
||
) -> str:
|
||
tid = str(tenant_id or "").strip()
|
||
pid = str(profile_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
if not tid or not pid or not uid:
|
||
raise ValueError("llm_grant_params_required")
|
||
if not self.get_llm_profile(pid):
|
||
raise ValueError("profile_not_found")
|
||
if not self.get_user_by_id(tenant_id=tid, user_id=uid):
|
||
raise ValueError("user_not_found")
|
||
gid = uuid.uuid4().hex
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO llm_profile_user_grant
|
||
(id, tenant_id, profile_id, user_id, created_at, created_by_user_id)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(gid, tid, pid, uid, ts, str(created_by_user_id).strip() if created_by_user_id else None),
|
||
)
|
||
row = conn.execute(
|
||
"""
|
||
SELECT id FROM llm_profile_user_grant
|
||
WHERE tenant_id = ? AND profile_id = ? AND user_id = ?
|
||
""",
|
||
(tid, pid, uid),
|
||
).fetchone()
|
||
if not row:
|
||
raise ValueError("llm_grant_failed")
|
||
return str(row["id"])
|
||
|
||
def revoke_llm_profile_grant(self, *, tenant_id: str, profile_id: str, user_id: str) -> int:
|
||
tid = str(tenant_id or "").strip()
|
||
pid = str(profile_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
if not tid or not pid or not uid:
|
||
return 0
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"DELETE FROM llm_profile_user_grant WHERE tenant_id = ? AND profile_id = ? AND user_id = ?",
|
||
(tid, pid, uid),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
|
||
def list_llm_profile_grants_for_profile(self, tenant_id: str, profile_id: str) -> list[dict[str, Any]]:
|
||
tid = str(tenant_id or "").strip()
|
||
pid = str(profile_id or "").strip()
|
||
if not tid or not pid:
|
||
return []
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT g.id AS grant_id, g.user_id, g.created_at, g.created_by_user_id,
|
||
u.username AS user_username, u.display_name AS user_display_name
|
||
FROM llm_profile_user_grant g
|
||
LEFT JOIN app_user u ON u.id = g.user_id AND u.tenant_id = g.tenant_id
|
||
WHERE g.tenant_id = ? AND g.profile_id = ?
|
||
ORDER BY g.created_at DESC
|
||
""",
|
||
(tid, pid),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
out.append(
|
||
{
|
||
"grant_id": r["grant_id"],
|
||
"user_id": r["user_id"],
|
||
"created_at": r["created_at"],
|
||
"created_by_user_id": r["created_by_user_id"] or "",
|
||
"username": str(r["user_username"] or "").strip(),
|
||
"display_name": str(r["user_display_name"] or "").strip(),
|
||
}
|
||
)
|
||
return out
|
||
|
||
def user_has_llm_profile_grant(self, tenant_id: str, user_id: str, profile_id: str) -> bool:
|
||
tid = str(tenant_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
pid = str(profile_id or "").strip()
|
||
if not tid or not uid or not pid:
|
||
return False
|
||
with self._connect() as conn:
|
||
n = conn.execute(
|
||
"""
|
||
SELECT 1 FROM llm_profile_user_grant
|
||
WHERE tenant_id = ? AND user_id = ? AND profile_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(tid, uid, pid),
|
||
).fetchone()
|
||
return n is not None
|
||
|
||
def grant_llm_profile_to_tenant(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
profile_id: str,
|
||
created_by_user_id: str | None = None,
|
||
) -> str:
|
||
tid = str(tenant_id or "").strip()
|
||
pid = str(profile_id or "").strip()
|
||
if not tid or not pid:
|
||
raise ValueError("llm_grant_params_required")
|
||
if not self.get_llm_profile(pid):
|
||
raise ValueError("profile_not_found")
|
||
gid = uuid.uuid4().hex
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO llm_profile_tenant_grant
|
||
(id, tenant_id, profile_id, created_at, created_by_user_id)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
""",
|
||
(gid, tid, pid, ts, str(created_by_user_id).strip() if created_by_user_id else None),
|
||
)
|
||
row = conn.execute(
|
||
"SELECT id FROM llm_profile_tenant_grant WHERE tenant_id = ? AND profile_id = ?",
|
||
(tid, pid),
|
||
).fetchone()
|
||
if not row:
|
||
raise ValueError("llm_grant_failed")
|
||
return str(row["id"])
|
||
|
||
def revoke_llm_profile_tenant_grant(self, *, tenant_id: str, profile_id: str) -> int:
|
||
tid = str(tenant_id or "").strip()
|
||
pid = str(profile_id or "").strip()
|
||
if not tid or not pid:
|
||
return 0
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"DELETE FROM llm_profile_tenant_grant WHERE tenant_id = ? AND profile_id = ?",
|
||
(tid, pid),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
|
||
def tenant_has_llm_profile_grant(self, tenant_id: str, profile_id: str) -> bool:
|
||
tid = str(tenant_id or "").strip()
|
||
pid = str(profile_id or "").strip()
|
||
if not tid or not pid:
|
||
return False
|
||
with self._connect() as conn:
|
||
n = conn.execute(
|
||
"""
|
||
SELECT 1 FROM llm_profile_tenant_grant
|
||
WHERE tenant_id = ? AND profile_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(tid, pid),
|
||
).fetchone()
|
||
return n is not None
|
||
|
||
def ensure_personal_llm_clone_from_global(self, user_id: str, username: str | None) -> None:
|
||
"""非 administrator：首次进入模型页/对话前，把全局池中的自定义 profile 复制一份并挂上独立 active/bindings。"""
|
||
if is_administrator_model_pool(username):
|
||
return
|
||
uid = str(user_id or "").strip()
|
||
if not uid:
|
||
return
|
||
from runtime.agents.specialists import dump_agent_profile_bindings, parse_agent_profile_bindings
|
||
|
||
seeded_key = f"llm_personal_pool_seeded:{uid}"
|
||
act_key = active_llm_profile_setting_key(uid, username)
|
||
bind_key = agent_profile_bindings_setting_key(uid, username)
|
||
ts = utc_now_iso()
|
||
|
||
def _get_setting_conn(conn: sqlite3.Connection, key: str) -> str | None:
|
||
row = conn.execute(
|
||
"SELECT value, is_secret FROM app_setting WHERE key = ?",
|
||
(key,),
|
||
).fetchone()
|
||
if not row or int(row["is_secret"] or 0) != 0:
|
||
return None
|
||
return str(row["value"])
|
||
|
||
def _upsert_setting_conn(conn: sqlite3.Connection, key: str, value: str) -> None:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO app_setting (key, value, is_secret, updated_at)
|
||
VALUES (?, ?, 0, ?)
|
||
ON CONFLICT(key) DO UPDATE SET value = excluded.value, is_secret = 0, updated_at = excluded.updated_at
|
||
""",
|
||
(key, value, ts),
|
||
)
|
||
|
||
with self._connect() as conn:
|
||
if str(_get_setting_conn(conn, seeded_key) or "").strip() == "1":
|
||
return
|
||
own_n = conn.execute(
|
||
"""
|
||
SELECT COUNT(*) AS c FROM llm_profile
|
||
WHERE owner_user_id = ? AND COALESCE(is_builtin, 0) = 0
|
||
""",
|
||
(uid,),
|
||
).fetchone()["c"]
|
||
if int(own_n or 0) > 0:
|
||
_upsert_setting_conn(conn, seeded_key, "1")
|
||
return
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT id, name, mode, model, base_url, api_key, hide_in_ui
|
||
FROM llm_profile
|
||
WHERE (owner_user_id IS NULL OR TRIM(owner_user_id) = '')
|
||
AND COALESCE(is_builtin, 0) = 0
|
||
""",
|
||
).fetchall()
|
||
id_map: dict[str, str] = {}
|
||
for r in rows:
|
||
old_id = str(r["id"])
|
||
new_id = uuid.uuid4().hex
|
||
id_map[old_id] = new_id
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO llm_profile
|
||
(id, name, mode, model, base_url, api_key, updated_at, is_builtin, hide_in_ui, owner_user_id)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, 0, ?, ?)
|
||
""",
|
||
(
|
||
new_id,
|
||
r["name"],
|
||
r["mode"],
|
||
r["model"],
|
||
r["base_url"],
|
||
r["api_key"],
|
||
ts,
|
||
int(r["hide_in_ui"] or 0),
|
||
uid,
|
||
),
|
||
)
|
||
global_active = str(_get_setting_conn(conn, "active_llm_profile_id") or "").strip()
|
||
if global_active:
|
||
if global_active in id_map:
|
||
_upsert_setting_conn(conn, act_key, id_map[global_active])
|
||
else:
|
||
_upsert_setting_conn(conn, act_key, global_active)
|
||
else:
|
||
_upsert_setting_conn(conn, act_key, LLM_BUILTIN_OLLAMA_PROFILE_ID)
|
||
raw_g = _get_setting_conn(conn, _MODEL_BINDINGS_KEY)
|
||
parsed = parse_agent_profile_bindings(raw_g)
|
||
for rid in list(parsed.keys()):
|
||
v = str(parsed.get(rid) or "").strip()
|
||
if v in id_map:
|
||
parsed[rid] = id_map[v]
|
||
_upsert_setting_conn(conn, bind_key, dump_agent_profile_bindings(parsed))
|
||
_upsert_setting_conn(conn, seeded_key, "1")
|
||
|
||
def set_llm_profile_secret(self, profile_id: str, plain_text: str) -> None:
|
||
ts = utc_now_iso()
|
||
enc = _encode_secret(plain_text)
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
UPDATE llm_profile
|
||
SET api_key = ?, updated_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(enc, ts, profile_id),
|
||
)
|
||
|
||
def clear_llm_profile_secret(self, profile_id: str) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
UPDATE llm_profile
|
||
SET api_key = NULL, updated_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(ts, profile_id),
|
||
)
|
||
|
||
def get_llm_profile_secret(self, profile_id: str) -> Optional[str]:
|
||
with self._connect() as conn:
|
||
r = conn.execute(
|
||
"SELECT api_key FROM llm_profile WHERE id = ?",
|
||
(profile_id,),
|
||
).fetchone()
|
||
if not r:
|
||
return None
|
||
val = r["api_key"]
|
||
if val is None or str(val) == "":
|
||
return None
|
||
try:
|
||
return _decode_secret(str(val))
|
||
except Exception:
|
||
return None
|
||
|
||
def upsert_knowledge_chunk(
|
||
self,
|
||
*,
|
||
chunk_id: str,
|
||
source: str,
|
||
content: str,
|
||
metadata: dict[str, Any] | None = None,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO knowledge_chunk (chunk_id, source, content, metadata, updated_at)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
ON CONFLICT(chunk_id) DO UPDATE SET
|
||
source = excluded.source,
|
||
content = excluded.content,
|
||
metadata = excluded.metadata,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(
|
||
chunk_id,
|
||
source,
|
||
content,
|
||
json.dumps(metadata or {}, ensure_ascii=False, default=str),
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def search_knowledge(self, *, query: str, limit: int = 3) -> list[dict[str, Any]]:
|
||
token = (query or "").strip()
|
||
if not token:
|
||
return []
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT source, content, metadata, updated_at
|
||
FROM knowledge_chunk
|
||
WHERE content LIKE ?
|
||
ORDER BY updated_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(f"%{token[:64]}%", max(1, int(limit))),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
meta = json.loads(r["metadata"])
|
||
except (json.JSONDecodeError, TypeError):
|
||
meta = {}
|
||
out.append(
|
||
{
|
||
"source": r["source"],
|
||
"content": r["content"],
|
||
"metadata": meta,
|
||
"updated_at": r["updated_at"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def get_knowledge_chunks(self, *, chunk_ids: list[str]) -> list[dict[str, Any]]:
|
||
ids = [str(x) for x in (chunk_ids or []) if str(x)]
|
||
if not ids:
|
||
return []
|
||
# SQLite parameter limit is high enough for our small uses; still cap.
|
||
ids = ids[:500]
|
||
qmarks = ",".join(["?"] * len(ids))
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT chunk_id, source, content, metadata, updated_at
|
||
FROM knowledge_chunk
|
||
WHERE chunk_id IN ({qmarks})
|
||
""",
|
||
tuple(ids),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
meta = json.loads(r["metadata"])
|
||
except Exception:
|
||
meta = {}
|
||
out.append(
|
||
{
|
||
"chunk_id": r["chunk_id"],
|
||
"source": r["source"],
|
||
"content": r["content"],
|
||
"metadata": meta,
|
||
"updated_at": r["updated_at"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def upsert_knowledge_embedding(
|
||
self,
|
||
*,
|
||
chunk_id: str,
|
||
model: str,
|
||
vector: list[float],
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
vec = [float(x) for x in (vector or [])]
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO knowledge_embedding (chunk_id, model, dim, vector_json, updated_at)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
ON CONFLICT(chunk_id, model) DO UPDATE SET
|
||
dim = excluded.dim,
|
||
vector_json = excluded.vector_json,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(str(chunk_id), str(model), int(len(vec)), json.dumps(vec), ts),
|
||
)
|
||
|
||
def ensure_memory_tables(self) -> None:
|
||
# tables are initialized in _init_db; this method is used by callers
|
||
# that need an explicit bootstrap hook.
|
||
with self._connect() as conn:
|
||
conn.execute("SELECT 1 FROM memory_item LIMIT 1")
|
||
|
||
def upsert_memory_item(
|
||
self,
|
||
*,
|
||
memory_id: str,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
session_id: str,
|
||
memory_type: str,
|
||
content: str,
|
||
confidence: float,
|
||
source: str,
|
||
metadata: dict[str, Any] | None = None,
|
||
created_at: str | None = None,
|
||
updated_at: str | None = None,
|
||
expires_at: str | None = None,
|
||
) -> None:
|
||
self.ensure_memory_tables()
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO memory_item
|
||
(memory_id, tenant_id, user_id, session_id, memory_type, content, confidence, source, metadata, created_at, updated_at, expires_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(memory_id) DO UPDATE SET
|
||
tenant_id = excluded.tenant_id,
|
||
user_id = excluded.user_id,
|
||
session_id = excluded.session_id,
|
||
memory_type = excluded.memory_type,
|
||
content = excluded.content,
|
||
confidence = excluded.confidence,
|
||
source = excluded.source,
|
||
metadata = excluded.metadata,
|
||
updated_at = excluded.updated_at,
|
||
expires_at = excluded.expires_at
|
||
""",
|
||
(
|
||
str(memory_id),
|
||
str(tenant_id),
|
||
str(user_id),
|
||
str(session_id),
|
||
str(memory_type),
|
||
str(content or ""),
|
||
float(confidence),
|
||
str(source or "memory"),
|
||
json.dumps(metadata or {}, ensure_ascii=False, default=str),
|
||
str(created_at or ts),
|
||
str(updated_at or ts),
|
||
str(expires_at) if expires_at else None,
|
||
),
|
||
)
|
||
|
||
def upsert_memory_vector(
|
||
self,
|
||
*,
|
||
memory_id: str,
|
||
model: str,
|
||
vector: list[float],
|
||
updated_at: str | None = None,
|
||
) -> None:
|
||
self.ensure_memory_tables()
|
||
ts = updated_at or utc_now_iso()
|
||
vec = [float(x) for x in (vector or [])]
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO memory_vector (memory_id, model, dim, vector_json, updated_at)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
ON CONFLICT(memory_id, model) DO UPDATE SET
|
||
dim = excluded.dim,
|
||
vector_json = excluded.vector_json,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(str(memory_id), str(model), len(vec), json.dumps(vec), ts),
|
||
)
|
||
|
||
@staticmethod
|
||
def _cosine_similarity(a: list[float], b: list[float]) -> float:
|
||
n = min(len(a), len(b))
|
||
if n <= 0:
|
||
return 0.0
|
||
dot = 0.0
|
||
na = 0.0
|
||
nb = 0.0
|
||
for i in range(n):
|
||
x = float(a[i])
|
||
y = float(b[i])
|
||
dot += x * y
|
||
na += x * x
|
||
nb += y * y
|
||
if na <= 1e-9 or nb <= 1e-9:
|
||
return 0.0
|
||
return float(dot / ((na**0.5) * (nb**0.5)))
|
||
|
||
def search_memory_vectors(
|
||
self,
|
||
*,
|
||
query_vector: list[float],
|
||
model: str,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
limit: int = 5,
|
||
) -> list[dict[str, Any]]:
|
||
self.ensure_memory_tables()
|
||
lim = max(1, min(int(limit), 100))
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT mi.memory_id, mi.tenant_id, mi.user_id, mi.session_id, mi.memory_type,
|
||
mi.content, mi.confidence, mi.source, mi.metadata, mi.created_at, mi.updated_at, mi.expires_at,
|
||
mv.vector_json
|
||
FROM memory_item mi
|
||
JOIN memory_vector mv ON mv.memory_id = mi.memory_id
|
||
WHERE mv.model = ? AND mi.tenant_id = ? AND mi.user_id = ?
|
||
AND (mi.expires_at IS NULL OR mi.expires_at > ?)
|
||
ORDER BY mi.updated_at DESC
|
||
LIMIT 5000
|
||
""",
|
||
(str(model), str(tenant_id), str(user_id), utc_now_iso()),
|
||
).fetchall()
|
||
scored: list[tuple[float, dict[str, Any]]] = []
|
||
q = [float(x) for x in (query_vector or [])]
|
||
for r in rows:
|
||
try:
|
||
vec = json.loads(r["vector_json"])
|
||
except Exception:
|
||
continue
|
||
score = self._cosine_similarity(q, [float(x) for x in (vec or [])])
|
||
try:
|
||
meta = json.loads(r["metadata"])
|
||
except Exception:
|
||
meta = {}
|
||
scored.append(
|
||
(
|
||
score,
|
||
{
|
||
"memory_id": r["memory_id"],
|
||
"tenant_id": r["tenant_id"],
|
||
"user_id": r["user_id"],
|
||
"session_id": r["session_id"],
|
||
"memory_type": r["memory_type"],
|
||
"content": r["content"],
|
||
"confidence": float(r["confidence"] or 0.0),
|
||
"source": r["source"],
|
||
"metadata": meta,
|
||
"created_at": r["created_at"],
|
||
"updated_at": r["updated_at"],
|
||
"expires_at": r["expires_at"],
|
||
"score": float(score),
|
||
},
|
||
)
|
||
)
|
||
scored.sort(key=lambda x: x[0], reverse=True)
|
||
return [x[1] for x in scored[:lim]]
|
||
|
||
def list_memory_items(
|
||
self,
|
||
*,
|
||
tenant_id: str | None = None,
|
||
user_id: str | None = None,
|
||
limit: int = 100,
|
||
offset: int = 0,
|
||
) -> list[dict[str, Any]]:
|
||
self.ensure_memory_tables()
|
||
where: list[str] = []
|
||
params: list[Any] = []
|
||
if tenant_id:
|
||
where.append("tenant_id = ?")
|
||
params.append(str(tenant_id))
|
||
if user_id:
|
||
where.append("user_id = ?")
|
||
params.append(str(user_id))
|
||
wsql = ("WHERE " + " AND ".join(where)) if where else ""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT memory_id, tenant_id, user_id, session_id, memory_type, content, confidence, source, metadata, created_at, updated_at, expires_at
|
||
FROM memory_item
|
||
{wsql}
|
||
ORDER BY updated_at DESC
|
||
LIMIT ? OFFSET ?
|
||
""",
|
||
(*params, max(1, min(int(limit), 500)), max(0, int(offset))),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
meta = json.loads(r["metadata"])
|
||
except Exception:
|
||
meta = {}
|
||
out.append(
|
||
{
|
||
"memory_id": r["memory_id"],
|
||
"tenant_id": r["tenant_id"],
|
||
"user_id": r["user_id"],
|
||
"session_id": r["session_id"],
|
||
"memory_type": r["memory_type"],
|
||
"content": r["content"],
|
||
"confidence": float(r["confidence"] or 0.0),
|
||
"source": r["source"],
|
||
"metadata": meta,
|
||
"created_at": r["created_at"],
|
||
"updated_at": r["updated_at"],
|
||
"expires_at": r["expires_at"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def delete_memory_item(self, *, memory_id: str) -> int:
|
||
self.ensure_memory_tables()
|
||
with self._connect() as conn:
|
||
cur = conn.execute("DELETE FROM memory_item WHERE memory_id = ?", (str(memory_id),))
|
||
return int(cur.rowcount or 0)
|
||
|
||
def add_memory_hit_log(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
session_id: str | None,
|
||
memory_id: str | None,
|
||
query_text: str,
|
||
score: float,
|
||
source: str,
|
||
timestamp: str | None = None,
|
||
) -> None:
|
||
self.ensure_memory_tables()
|
||
ts = timestamp or utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO memory_hit_log (tenant_id, user_id, session_id, memory_id, query_text, score, source, timestamp)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
str(tenant_id),
|
||
str(user_id),
|
||
str(session_id) if session_id else None,
|
||
str(memory_id) if memory_id else None,
|
||
str(query_text or ""),
|
||
float(score),
|
||
str(source or "memory"),
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def list_memory_hit_logs(
|
||
self,
|
||
*,
|
||
tenant_id: str | None = None,
|
||
user_id: str | None = None,
|
||
limit: int = 100,
|
||
) -> list[dict[str, Any]]:
|
||
self.ensure_memory_tables()
|
||
where: list[str] = []
|
||
params: list[Any] = []
|
||
if tenant_id:
|
||
where.append("tenant_id = ?")
|
||
params.append(str(tenant_id))
|
||
if user_id:
|
||
where.append("user_id = ?")
|
||
params.append(str(user_id))
|
||
wsql = ("WHERE " + " AND ".join(where)) if where else ""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT id, tenant_id, user_id, session_id, memory_id, query_text, score, source, timestamp
|
||
FROM memory_hit_log
|
||
{wsql}
|
||
ORDER BY id DESC
|
||
LIMIT ?
|
||
""",
|
||
(*params, max(1, min(int(limit), 500))),
|
||
).fetchall()
|
||
return [
|
||
{
|
||
"id": int(r["id"]),
|
||
"tenant_id": r["tenant_id"],
|
||
"user_id": r["user_id"],
|
||
"session_id": r["session_id"],
|
||
"memory_id": r["memory_id"],
|
||
"query_text": r["query_text"],
|
||
"score": float(r["score"] or 0.0),
|
||
"source": r["source"],
|
||
"timestamp": r["timestamp"],
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def clear_low_confidence_memory(self, *, max_confidence: float) -> int:
|
||
self.ensure_memory_tables()
|
||
with self._connect() as conn:
|
||
cur = conn.execute("DELETE FROM memory_item WHERE confidence <= ?", (float(max_confidence),))
|
||
return int(cur.rowcount or 0)
|
||
|
||
def list_knowledge_embeddings(self, *, model: str, limit: int = 5000) -> list[dict[str, Any]]:
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT chunk_id, model, dim, vector_json, updated_at
|
||
FROM knowledge_embedding
|
||
WHERE model = ?
|
||
ORDER BY updated_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(str(model), max(1, int(limit))),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
vec = json.loads(r["vector_json"])
|
||
except Exception:
|
||
vec = []
|
||
out.append(
|
||
{
|
||
"chunk_id": r["chunk_id"],
|
||
"model": r["model"],
|
||
"dim": int(r["dim"] or 0),
|
||
"vector": vec,
|
||
"updated_at": r["updated_at"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def add_agent_audit_log(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
specialist: str,
|
||
task_kind: str,
|
||
action: str,
|
||
payload: dict[str, Any],
|
||
status: str,
|
||
reason: str,
|
||
duration_ms: int = 0,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO agent_audit_log
|
||
(session_id, specialist, task_kind, action, payload, status, reason, duration_ms, timestamp)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
session_id,
|
||
specialist,
|
||
task_kind,
|
||
action,
|
||
json.dumps(payload, ensure_ascii=False, default=str),
|
||
status,
|
||
reason,
|
||
int(duration_ms),
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def add_agent_eval_log(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
specialist: str,
|
||
task_kind: str,
|
||
success: bool,
|
||
latency_ms: int,
|
||
cost_hint: float = 0.0,
|
||
notes: str = "",
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO agent_eval_log
|
||
(session_id, specialist, task_kind, success, latency_ms, cost_hint, notes, timestamp)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
session_id,
|
||
specialist,
|
||
task_kind,
|
||
1 if success else 0,
|
||
int(latency_ms),
|
||
float(cost_hint),
|
||
notes,
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def list_agent_eval_logs(self, *, limit: int = 200) -> list[dict[str, Any]]:
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT session_id, specialist, task_kind, success, latency_ms, cost_hint, notes, timestamp
|
||
FROM agent_eval_log
|
||
ORDER BY id DESC
|
||
LIMIT ?
|
||
""",
|
||
(max(1, int(limit)),),
|
||
).fetchall()
|
||
return [
|
||
{
|
||
"session_id": r["session_id"],
|
||
"specialist": r["specialist"],
|
||
"task_kind": r["task_kind"],
|
||
"success": bool(int(r["success"] or 0)),
|
||
"latency_ms": int(r["latency_ms"] or 0),
|
||
"cost_hint": float(r["cost_hint"] or 0.0),
|
||
"notes": r["notes"],
|
||
"timestamp": r["timestamp"],
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def list_agent_audit_logs(
|
||
self,
|
||
*,
|
||
limit: int = 200,
|
||
session_id: str | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
with self._connect() as conn:
|
||
if session_id:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT session_id, specialist, task_kind, action, payload, status, reason, duration_ms, timestamp
|
||
FROM agent_audit_log
|
||
WHERE session_id = ?
|
||
ORDER BY id DESC
|
||
LIMIT ?
|
||
""",
|
||
(str(session_id), max(1, int(limit))),
|
||
).fetchall()
|
||
else:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT session_id, specialist, task_kind, action, payload, status, reason, duration_ms, timestamp
|
||
FROM agent_audit_log
|
||
ORDER BY id DESC
|
||
LIMIT ?
|
||
""",
|
||
(max(1, int(limit)),),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
payload = json.loads(r["payload"])
|
||
except (json.JSONDecodeError, TypeError):
|
||
payload = {}
|
||
out.append(
|
||
{
|
||
"session_id": r["session_id"],
|
||
"specialist": r["specialist"],
|
||
"task_kind": r["task_kind"],
|
||
"action": r["action"],
|
||
"payload": payload,
|
||
"status": r["status"],
|
||
"reason": r["reason"],
|
||
"duration_ms": int(r["duration_ms"] or 0),
|
||
"timestamp": r["timestamp"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def upsert_tool_plugin(
|
||
self,
|
||
*,
|
||
plugin_name: str,
|
||
plugin_version: str,
|
||
entry_point: str,
|
||
enabled: bool = True,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO tool_plugin (plugin_name, plugin_version, entry_point, enabled, updated_at)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
ON CONFLICT(plugin_name, entry_point) DO UPDATE SET
|
||
plugin_version = excluded.plugin_version,
|
||
enabled = excluded.enabled,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(str(plugin_name), str(plugin_version), str(entry_point), 1 if enabled else 0, ts),
|
||
)
|
||
|
||
def list_tool_plugins(self) -> list[dict[str, Any]]:
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT plugin_name, plugin_version, entry_point, enabled, updated_at
|
||
FROM tool_plugin
|
||
ORDER BY updated_at DESC
|
||
"""
|
||
).fetchall()
|
||
return [
|
||
{
|
||
"plugin_name": r["plugin_name"],
|
||
"plugin_version": r["plugin_version"],
|
||
"entry_point": r["entry_point"],
|
||
"enabled": bool(int(r["enabled"] or 0)),
|
||
"updated_at": r["updated_at"],
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def upsert_mcp_server(
|
||
self,
|
||
*,
|
||
server_id: str,
|
||
source_type: str,
|
||
source_ref: str,
|
||
version: str = "",
|
||
entry_command: str = "",
|
||
entry_args: list[str] | None = None,
|
||
env_schema: dict[str, Any] | None = None,
|
||
required_permissions: list[str] | None = None,
|
||
risk_level: str = "high",
|
||
timeout_s: float = 30.0,
|
||
enabled: bool = False,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO mcp_server_registry
|
||
(server_id, source_type, source_ref, version, entry_command, entry_args, env_schema,
|
||
required_permissions, risk_level, timeout_s, enabled, updated_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(server_id) DO UPDATE SET
|
||
source_type = excluded.source_type,
|
||
source_ref = excluded.source_ref,
|
||
version = excluded.version,
|
||
entry_command = excluded.entry_command,
|
||
entry_args = excluded.entry_args,
|
||
env_schema = excluded.env_schema,
|
||
required_permissions = excluded.required_permissions,
|
||
risk_level = excluded.risk_level,
|
||
timeout_s = excluded.timeout_s,
|
||
enabled = excluded.enabled,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(
|
||
str(server_id),
|
||
str(source_type),
|
||
str(source_ref),
|
||
str(version or ""),
|
||
str(entry_command or ""),
|
||
json.dumps(entry_args or [], ensure_ascii=False),
|
||
json.dumps(env_schema or {}, ensure_ascii=False),
|
||
json.dumps(required_permissions or [], ensure_ascii=False),
|
||
str(risk_level or "high"),
|
||
float(timeout_s or 30.0),
|
||
1 if enabled else 0,
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def set_mcp_server_enabled(self, *, server_id: str, enabled: bool) -> int:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"UPDATE mcp_server_registry SET enabled = ?, updated_at = ? WHERE server_id = ?",
|
||
(1 if enabled else 0, ts, str(server_id)),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
|
||
def list_mcp_servers(self, *, enabled_only: bool = False) -> list[dict[str, Any]]:
|
||
sql = """
|
||
SELECT server_id, source_type, source_ref, version, entry_command, entry_args, env_schema,
|
||
required_permissions, risk_level, timeout_s, enabled, updated_at
|
||
FROM mcp_server_registry
|
||
"""
|
||
params: tuple[Any, ...] = ()
|
||
if enabled_only:
|
||
sql += " WHERE enabled = 1"
|
||
sql += " ORDER BY updated_at DESC"
|
||
with self._connect() as conn:
|
||
rows = conn.execute(sql, params).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
entry_args = json.loads(r["entry_args"] or "[]")
|
||
except Exception:
|
||
entry_args = []
|
||
try:
|
||
env_schema = json.loads(r["env_schema"] or "{}")
|
||
except Exception:
|
||
env_schema = {}
|
||
try:
|
||
required_permissions = json.loads(r["required_permissions"] or "[]")
|
||
except Exception:
|
||
required_permissions = []
|
||
out.append(
|
||
{
|
||
"server_id": r["server_id"],
|
||
"source_type": r["source_type"],
|
||
"source_ref": r["source_ref"],
|
||
"version": r["version"],
|
||
"entry_command": r["entry_command"],
|
||
"entry_args": entry_args if isinstance(entry_args, list) else [],
|
||
"env_schema": env_schema if isinstance(env_schema, dict) else {},
|
||
"required_permissions": required_permissions if isinstance(required_permissions, list) else [],
|
||
"risk_level": r["risk_level"],
|
||
"timeout_s": float(r["timeout_s"] or 30.0),
|
||
"enabled": bool(int(r["enabled"] or 0)),
|
||
"updated_at": r["updated_at"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def add_mcp_installation_log(
|
||
self,
|
||
*,
|
||
server_id: str,
|
||
status: str,
|
||
error_code: str = "",
|
||
detail: dict[str, Any] | None = None,
|
||
install_command: str = "",
|
||
) -> None:
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO mcp_server_installation
|
||
(server_id, status, error_code, detail, install_command, timestamp)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
str(server_id),
|
||
str(status),
|
||
str(error_code or ""),
|
||
json.dumps(detail or {}, ensure_ascii=False),
|
||
str(install_command or ""),
|
||
utc_now_iso(),
|
||
),
|
||
)
|
||
|
||
def list_mcp_installation_logs(self, *, server_id: str | None = None, limit: int = 100) -> list[dict[str, Any]]:
|
||
where = ""
|
||
params: list[Any] = []
|
||
if server_id:
|
||
where = "WHERE server_id = ?"
|
||
params.append(str(server_id))
|
||
params.append(max(1, int(limit)))
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT server_id, status, error_code, detail, install_command, timestamp
|
||
FROM mcp_server_installation
|
||
{where}
|
||
ORDER BY id DESC
|
||
LIMIT ?
|
||
""",
|
||
tuple(params),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
detail = json.loads(r["detail"] or "{}")
|
||
except Exception:
|
||
detail = {}
|
||
out.append(
|
||
{
|
||
"server_id": r["server_id"],
|
||
"status": r["status"],
|
||
"error_code": r["error_code"],
|
||
"detail": detail,
|
||
"install_command": r["install_command"],
|
||
"timestamp": r["timestamp"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def list_mcp_install_failure_summary(self, *, limit: int = 20) -> list[dict[str, Any]]:
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT server_id, error_code, COUNT(1) AS n, MAX(timestamp) AS last_ts
|
||
FROM mcp_server_installation
|
||
WHERE status = 'error'
|
||
GROUP BY server_id, error_code
|
||
ORDER BY n DESC, last_ts DESC
|
||
LIMIT ?
|
||
""",
|
||
(max(1, int(limit)),),
|
||
).fetchall()
|
||
return [
|
||
{
|
||
"server_id": str(r["server_id"] or ""),
|
||
"error_code": str(r["error_code"] or ""),
|
||
"count": int(r["n"] or 0),
|
||
"last_ts": str(r["last_ts"] or ""),
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def set_mcp_server_health(self, *, server_id: str, status: str, detail: dict[str, Any] | None = None) -> None:
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO mcp_server_health (server_id, status, detail, checked_at)
|
||
VALUES (?, ?, ?, ?)
|
||
ON CONFLICT(server_id) DO UPDATE SET
|
||
status = excluded.status,
|
||
detail = excluded.detail,
|
||
checked_at = excluded.checked_at
|
||
""",
|
||
(str(server_id), str(status), json.dumps(detail or {}, ensure_ascii=False), utc_now_iso()),
|
||
)
|
||
|
||
def list_mcp_server_health(self) -> list[dict[str, Any]]:
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"SELECT server_id, status, detail, checked_at FROM mcp_server_health ORDER BY checked_at DESC"
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
detail = json.loads(r["detail"] or "{}")
|
||
except Exception:
|
||
detail = {}
|
||
out.append(
|
||
{"server_id": r["server_id"], "status": r["status"], "detail": detail, "checked_at": r["checked_at"]}
|
||
)
|
||
return out
|
||
|
||
def replace_mcp_server_tools(self, *, server_id: str, tools: list[dict[str, Any]]) -> None:
|
||
ts = utc_now_iso()
|
||
sid = str(server_id or "").strip()
|
||
if not sid:
|
||
return
|
||
with self._connect() as conn:
|
||
conn.execute("DELETE FROM mcp_server_tool WHERE server_id = ?", (sid,))
|
||
for t in tools:
|
||
tool_name = str((t or {}).get("tool_name") or "").strip()
|
||
if not tool_name:
|
||
continue
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO mcp_server_tool (server_id, tool_name, description, parameters, updated_at)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
sid,
|
||
tool_name,
|
||
str((t or {}).get("description") or ""),
|
||
json.dumps((t or {}).get("parameters") or {}, ensure_ascii=False),
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def list_mcp_server_tools(self, *, server_id: str) -> list[dict[str, Any]]:
|
||
sid = str(server_id or "").strip()
|
||
if not sid:
|
||
return []
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT tool_name, description, parameters, updated_at
|
||
FROM mcp_server_tool
|
||
WHERE server_id = ?
|
||
ORDER BY tool_name ASC
|
||
""",
|
||
(sid,),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
parameters = json.loads(r["parameters"] or "{}")
|
||
except Exception:
|
||
parameters = {}
|
||
out.append(
|
||
{
|
||
"tool_name": r["tool_name"],
|
||
"description": r["description"],
|
||
"parameters": parameters if isinstance(parameters, dict) else {},
|
||
"updated_at": r["updated_at"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def add_trace_event(
|
||
self,
|
||
*,
|
||
session_id: str,
|
||
trace_id: str,
|
||
span_id: str,
|
||
parent_span_id: str | None,
|
||
event_type: str,
|
||
payload: dict[str, Any],
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
self._trace_events_repo().insert_one(
|
||
session_id=str(session_id),
|
||
trace_id=str(trace_id),
|
||
span_id=str(span_id),
|
||
parent_span_id=str(parent_span_id) if parent_span_id else None,
|
||
event_type=str(event_type),
|
||
payload=json.dumps(payload or {}, ensure_ascii=False, default=str),
|
||
timestamp=ts,
|
||
)
|
||
|
||
def add_trace_events_batch(self, events: list[dict[str, Any]]) -> None:
|
||
ts = utc_now_iso()
|
||
batch: list[dict[str, Any]] = []
|
||
for e in events or []:
|
||
try:
|
||
session_id = str(e.get("session_id") or "").strip()
|
||
trace_id = str(e.get("trace_id") or "").strip()
|
||
span_id = str(e.get("span_id") or "").strip()
|
||
if (not session_id) or (not trace_id) or (not span_id):
|
||
continue
|
||
parent_span_id = str(e.get("parent_span_id") or "").strip() or None
|
||
event_type = str(e.get("event_type") or "").strip()
|
||
payload = e.get("payload") if isinstance(e.get("payload"), dict) else {}
|
||
batch.append(
|
||
{
|
||
"session_id": session_id,
|
||
"trace_id": trace_id,
|
||
"span_id": span_id,
|
||
"parent_span_id": parent_span_id,
|
||
"event_type": event_type,
|
||
"payload": json.dumps(payload or {}, ensure_ascii=False, default=str),
|
||
"timestamp": ts,
|
||
}
|
||
)
|
||
except Exception:
|
||
continue
|
||
self._trace_events_repo().insert_many(batch)
|
||
|
||
def list_trace_events(self, *, session_id: str, limit: int = 300) -> list[dict[str, Any]]:
|
||
rows = self._trace_events_repo().list_trace_events_desc(
|
||
session_id=str(session_id), limit=max(1, int(limit))
|
||
)
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
payload = json.loads(r["payload"])
|
||
except Exception:
|
||
payload = {}
|
||
out.append(
|
||
{
|
||
"trace_id": r["trace_id"],
|
||
"span_id": r["span_id"],
|
||
"parent_span_id": r["parent_span_id"],
|
||
"event_type": r["event_type"],
|
||
"payload": payload,
|
||
"timestamp": r["timestamp"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def list_trace_events_for_trace(
|
||
self, *, session_id: str, trace_id: str, limit: int = 500
|
||
) -> list[dict[str, Any]]:
|
||
rows = self._trace_events_repo().list_trace_events_for_trace_asc(
|
||
session_id=session_id, trace_id=trace_id, limit=max(1, int(limit))
|
||
)
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
payload = json.loads(r["payload"])
|
||
except Exception:
|
||
payload = {}
|
||
out.append(
|
||
{
|
||
"trace_id": r["trace_id"],
|
||
"span_id": r["span_id"],
|
||
"parent_span_id": r["parent_span_id"],
|
||
"event_type": r["event_type"],
|
||
"payload": payload,
|
||
"timestamp": r["timestamp"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def get_turn_time_window(self, *, session_id: str, trace_id: str) -> tuple[str | None, str | None]:
|
||
"""Best-effort (start_ts, end_ts) for one trace_id in a session."""
|
||
sid = str(session_id or "").strip()
|
||
tid = str(trace_id or "").strip()
|
||
if not sid or not tid:
|
||
return None, None
|
||
rows = self._trace_events_repo().list_event_type_timestamp_for_trace(
|
||
session_id=sid, trace_id=tid
|
||
)
|
||
if not rows:
|
||
return None, None
|
||
start = None
|
||
end = None
|
||
for r in rows:
|
||
et = str(r["event_type"] or "")
|
||
ts = str(r["timestamp"] or "")
|
||
if not start:
|
||
start = ts
|
||
end = ts
|
||
if et == "turn_started":
|
||
start = ts
|
||
if et == "turn_finished":
|
||
end = ts
|
||
return start, end
|
||
|
||
def list_messages_in_time_window(
|
||
self, *, session_id: str, start_ts: str | None, end_ts: str | None, limit: int = 500
|
||
) -> list[dict[str, Any]]:
|
||
sid = str(session_id or "").strip()
|
||
if not sid:
|
||
return []
|
||
start = str(start_ts or "").strip()
|
||
end = str(end_ts or "").strip()
|
||
if not start or not end:
|
||
return []
|
||
lim = max(1, min(int(limit), 2000))
|
||
return self._chat_messages_repo().list_messages_in_time_window(
|
||
session_id=sid, start_ts=start, end_ts=end, limit=lim
|
||
)
|
||
|
||
# ----------------------------
|
||
# Tenant / User / Bind Codes
|
||
# ----------------------------
|
||
def create_tenant(self, name: str) -> dict[str, Any]:
|
||
tid = str(uuid.uuid4())
|
||
ts = utc_now_iso()
|
||
nm = str(name or "").strip() or "Team"
|
||
self._tenant_repo().insert_tenant(tenant_id=tid, name=nm, created_at=ts)
|
||
return {"id": tid, "name": name, "created_at": ts}
|
||
|
||
def delete_tenant(self, *, tenant_id: str) -> int:
|
||
tid = str(tenant_id or "").strip()
|
||
if not tid:
|
||
return 0
|
||
return self._tenant_repo().delete_tenant(tenant_id=tid)
|
||
|
||
def list_tenants(self, *, limit: int = 200) -> list[dict[str, Any]]:
|
||
return self._tenant_repo().list_tenants(limit=max(1, int(limit)))
|
||
|
||
def create_user(self, *, tenant_id: str, display_name: str, role: str) -> dict[str, Any]:
|
||
uid = str(uuid.uuid4())
|
||
ts = utc_now_iso()
|
||
username = (str(display_name or "").strip() or "user").lower().replace(" ", "_")
|
||
repo = self._app_users_repo()
|
||
suffix = repo.count_by_tenant_username(tenant_id=str(tenant_id), username=username)
|
||
if suffix > 0:
|
||
username = f"{username}_{suffix+1}"
|
||
disp = str(display_name or "").strip() or "User"
|
||
repo.insert_user(
|
||
user_id=uid,
|
||
tenant_id=str(tenant_id),
|
||
username=username,
|
||
display_name=disp,
|
||
role=str(role or "member"),
|
||
password_hash="",
|
||
is_active=1,
|
||
created_at=ts,
|
||
avatar_attachment_id=None,
|
||
)
|
||
return {
|
||
"id": uid,
|
||
"tenant_id": tenant_id,
|
||
"username": username,
|
||
"display_name": display_name,
|
||
"role": role,
|
||
"is_active": True,
|
||
"created_at": ts,
|
||
}
|
||
|
||
def list_users(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
limit: int = 500,
|
||
offset: int = 0,
|
||
q: str | None = None,
|
||
include_inactive: bool = True,
|
||
) -> list[dict[str, Any]]:
|
||
return self._app_users_repo().list_users_for_tenant(
|
||
tenant_id=str(tenant_id),
|
||
limit=limit,
|
||
offset=offset,
|
||
q=q,
|
||
include_inactive=include_inactive,
|
||
)
|
||
|
||
def get_user_by_username(self, *, tenant_id: str, username: str) -> dict[str, Any] | None:
|
||
return self._app_users_repo().fetch_by_tenant_and_username(
|
||
tenant_id=str(tenant_id), username=str(username)
|
||
)
|
||
|
||
def get_user_by_username_global(self, *, username: str) -> dict[str, Any] | None:
|
||
return self._app_users_repo().fetch_first_by_username_global(username=str(username))
|
||
|
||
def get_user_by_id(self, *, tenant_id: str, user_id: str) -> dict[str, Any] | None:
|
||
return self._app_users_repo().fetch_by_tenant_and_id(
|
||
tenant_id=str(tenant_id), user_id=str(user_id)
|
||
)
|
||
|
||
def create_user_account(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
username: str,
|
||
display_name: str,
|
||
role: str,
|
||
password_hash: str,
|
||
is_active: bool = True,
|
||
) -> dict[str, Any]:
|
||
uid = str(uuid.uuid4())
|
||
ts = utc_now_iso()
|
||
self._app_users_repo().insert_user(
|
||
user_id=uid,
|
||
tenant_id=str(tenant_id),
|
||
username=str(username).strip(),
|
||
display_name=str(display_name).strip() or str(username).strip(),
|
||
role=str(role or "member"),
|
||
password_hash=str(password_hash or ""),
|
||
is_active=1 if is_active else 0,
|
||
created_at=ts,
|
||
avatar_attachment_id=None,
|
||
)
|
||
return self.get_user_by_id(tenant_id=tenant_id, user_id=uid) or {}
|
||
|
||
def update_user_account(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
display_name: str | None = None,
|
||
role: str | None = None,
|
||
is_active: bool | None = None,
|
||
password_hash: str | None = None,
|
||
avatar_attachment_id: str | None = None,
|
||
) -> bool:
|
||
return self._app_users_repo().update_user_account(
|
||
tenant_id=str(tenant_id),
|
||
user_id=str(user_id),
|
||
display_name=display_name,
|
||
role=role,
|
||
is_active=is_active,
|
||
password_hash=password_hash,
|
||
avatar_attachment_id=avatar_attachment_id,
|
||
)
|
||
|
||
def delete_user_account(self, *, tenant_id: str, user_id: str) -> int:
|
||
return self._app_users_repo().delete_user_account(tenant_id=str(tenant_id), user_id=str(user_id))
|
||
|
||
def upsert_user_permission(self, *, tenant_id: str, user_id: str, permission: str) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO user_permission (tenant_id, user_id, permission, created_at)
|
||
VALUES (?, ?, ?, ?)
|
||
""",
|
||
(str(tenant_id), str(user_id), str(permission), ts),
|
||
)
|
||
|
||
def delete_user_permission(self, *, tenant_id: str, user_id: str, permission: str) -> int:
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"DELETE FROM user_permission WHERE tenant_id = ? AND user_id = ? AND permission = ?",
|
||
(str(tenant_id), str(user_id), str(permission)),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
|
||
def list_user_permissions(self, *, tenant_id: str, user_id: str, role: str | None = None) -> list[str]:
|
||
out: set[str] = set()
|
||
role_name = str(role or "").strip()
|
||
with self._connect() as conn:
|
||
if role_name:
|
||
rows = conn.execute(
|
||
"SELECT permission FROM role_permission WHERE role = ? ORDER BY permission ASC",
|
||
(role_name,),
|
||
).fetchall()
|
||
out.update(str(r["permission"]) for r in rows)
|
||
rows2 = conn.execute(
|
||
"""
|
||
SELECT permission FROM user_permission
|
||
WHERE tenant_id = ? AND user_id = ?
|
||
ORDER BY permission ASC
|
||
""",
|
||
(str(tenant_id), str(user_id)),
|
||
).fetchall()
|
||
out.update(str(r["permission"]) for r in rows2)
|
||
return sorted(out)
|
||
|
||
def get_user_workspace_path_allowlist(self, *, tenant_id: str, user_id: str) -> dict[str, Any] | None:
|
||
tid = str(tenant_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
if not tid or not uid:
|
||
return None
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT tenant_id, user_id, extra_roots, allow_any_path, allow_high_risk_public_tools, updated_at
|
||
FROM user_workspace_path_allowlist
|
||
WHERE tenant_id = ? AND user_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(tid, uid),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
return {
|
||
"tenant_id": str(row["tenant_id"] or ""),
|
||
"user_id": str(row["user_id"] or ""),
|
||
"extra_roots": str(row["extra_roots"] or ""),
|
||
"allow_any_path": bool(int(row["allow_any_path"] or 0)),
|
||
"allow_high_risk_public_tools": bool(int(row["allow_high_risk_public_tools"] or 0)),
|
||
"updated_at": str(row["updated_at"] or ""),
|
||
}
|
||
|
||
def list_user_workspace_extra_roots_union(self) -> list[str]:
|
||
"""All ``|``-separated extra roots from every ``user_workspace_path_allowlist`` row.
|
||
|
||
Intended for **admin/diagnostics** only. MCP filesystem argv must use per-session policy
|
||
(see ``collect_filesystem_mcp_extra_roots(..., policy_session_id=...)``), not this union,
|
||
so one user cannot inherit another user's roots.
|
||
"""
|
||
out: list[str] = []
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"SELECT extra_roots FROM user_workspace_path_allowlist WHERE length(trim(extra_roots)) > 0"
|
||
).fetchall()
|
||
for r in rows:
|
||
raw = str(r["extra_roots"] or "")
|
||
for part in raw.split("|"):
|
||
s = part.strip().strip('"').strip("'")
|
||
if s:
|
||
out.append(s)
|
||
return out
|
||
|
||
def upsert_user_workspace_path_allowlist(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
extra_roots: str,
|
||
allow_any_path: bool,
|
||
allow_high_risk_public_tools: bool = False,
|
||
) -> None:
|
||
tid = str(tenant_id or "").strip()
|
||
uid = str(user_id or "").strip()
|
||
if not tid or not uid:
|
||
return
|
||
ts = utc_now_iso()
|
||
roots = str(extra_roots or "")
|
||
if len(roots) > 16000:
|
||
roots = roots[:16000]
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO user_workspace_path_allowlist (
|
||
tenant_id, user_id, extra_roots, allow_any_path, allow_high_risk_public_tools, updated_at
|
||
)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(tenant_id, user_id) DO UPDATE SET
|
||
extra_roots = excluded.extra_roots,
|
||
allow_any_path = excluded.allow_any_path,
|
||
allow_high_risk_public_tools = excluded.allow_high_risk_public_tools,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(
|
||
tid,
|
||
uid,
|
||
roots,
|
||
1 if allow_any_path else 0,
|
||
1 if allow_high_risk_public_tools else 0,
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def create_auth_session(
|
||
self,
|
||
*,
|
||
session_token_hash: str,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
role: str,
|
||
expires_at: str,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
self._auth_sessions_repo().insert_session(
|
||
session_token_hash=str(session_token_hash),
|
||
tenant_id=str(tenant_id),
|
||
user_id=str(user_id),
|
||
role=str(role),
|
||
created_at=ts,
|
||
expires_at=str(expires_at),
|
||
last_seen_at=ts,
|
||
)
|
||
|
||
def revoke_auth_session(self, *, session_token_hash: str) -> int:
|
||
ts = utc_now_iso()
|
||
return self._auth_sessions_repo().revoke_one(session_token_hash=str(session_token_hash), revoked_at=ts)
|
||
|
||
def revoke_all_auth_sessions(self) -> int:
|
||
ts = utc_now_iso()
|
||
return self._auth_sessions_repo().revoke_all_active(revoked_at=ts)
|
||
|
||
def get_auth_session(self, *, session_token_hash: str) -> dict[str, Any] | None:
|
||
return self._auth_sessions_repo().fetch_by_hash(session_token_hash=str(session_token_hash))
|
||
|
||
def touch_auth_session(self, *, session_token_hash: str) -> None:
|
||
ts = utc_now_iso()
|
||
self._auth_sessions_repo().touch(session_token_hash=str(session_token_hash), last_seen_at=ts)
|
||
|
||
def add_admin_audit_log(
|
||
self,
|
||
*,
|
||
actor_tenant_id: str,
|
||
actor_user_id: str,
|
||
action: str,
|
||
target_type: str,
|
||
target_id: str,
|
||
status: str,
|
||
detail: dict[str, Any] | None = None,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO admin_audit_log
|
||
(actor_tenant_id, actor_user_id, action, target_type, target_id, status, detail, timestamp)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
str(actor_tenant_id),
|
||
str(actor_user_id),
|
||
str(action),
|
||
str(target_type),
|
||
str(target_id),
|
||
str(status),
|
||
json.dumps(detail or {}, ensure_ascii=False, default=str),
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def list_admin_audit_logs(
|
||
self,
|
||
*,
|
||
tenant_id: str | None = None,
|
||
action: str | None = None,
|
||
actor_user_id: str | None = None,
|
||
status: str | None = None,
|
||
limit: int = 200,
|
||
offset: int = 0,
|
||
) -> list[dict[str, Any]]:
|
||
lim = max(1, min(int(limit), 500))
|
||
off = max(0, int(offset))
|
||
clauses: list[str] = []
|
||
params: list[Any] = []
|
||
tid = str(tenant_id or "").strip()
|
||
act = str(action or "").strip()
|
||
actor = str(actor_user_id or "").strip()
|
||
st = str(status or "").strip()
|
||
if tid:
|
||
clauses.append("l.actor_tenant_id = ?")
|
||
params.append(tid)
|
||
if act:
|
||
clauses.append("l.action = ?")
|
||
params.append(act)
|
||
if actor:
|
||
clauses.append("l.actor_user_id = ?")
|
||
params.append(actor)
|
||
if st:
|
||
clauses.append("l.status = ?")
|
||
params.append(st)
|
||
where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""
|
||
sql = f"""
|
||
SELECT l.actor_tenant_id, l.actor_user_id, l.action, l.target_type, l.target_id, l.status, l.detail,
|
||
l.timestamp, u.username AS actor_username, u.display_name AS actor_display_name
|
||
FROM admin_audit_log l
|
||
LEFT JOIN app_user u ON u.tenant_id = l.actor_tenant_id AND u.id = l.actor_user_id
|
||
{where_sql}
|
||
ORDER BY l.id DESC
|
||
LIMIT ?
|
||
OFFSET ?
|
||
"""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(sql, tuple(params + [lim, off])).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for r in rows:
|
||
try:
|
||
detail = json.loads(r["detail"])
|
||
except Exception:
|
||
detail = {}
|
||
out.append(
|
||
{
|
||
"actor_tenant_id": r["actor_tenant_id"],
|
||
"actor_user_id": r["actor_user_id"],
|
||
"actor_username": str(r["actor_username"] or "") if r["actor_username"] is not None else "",
|
||
"actor_display_name": str(r["actor_display_name"] or "") if r["actor_display_name"] is not None else "",
|
||
"action": r["action"],
|
||
"target_type": r["target_type"],
|
||
"target_id": r["target_id"],
|
||
"status": r["status"],
|
||
"detail": detail,
|
||
"timestamp": r["timestamp"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def count_admin_audit_logs(
|
||
self,
|
||
*,
|
||
tenant_id: str | None = None,
|
||
action: str | None = None,
|
||
actor_user_id: str | None = None,
|
||
status: str | None = None,
|
||
) -> int:
|
||
clauses: list[str] = []
|
||
params: list[Any] = []
|
||
tid = str(tenant_id or "").strip()
|
||
act = str(action or "").strip()
|
||
actor = str(actor_user_id or "").strip()
|
||
st = str(status or "").strip()
|
||
if tid:
|
||
clauses.append("actor_tenant_id = ?")
|
||
params.append(tid)
|
||
if act:
|
||
clauses.append("action = ?")
|
||
params.append(act)
|
||
if actor:
|
||
clauses.append("actor_user_id = ?")
|
||
params.append(actor)
|
||
if st:
|
||
clauses.append("status = ?")
|
||
params.append(st)
|
||
where_sql = f"WHERE {' AND '.join(clauses)}" if clauses else ""
|
||
sql = f"SELECT COUNT(1) AS c FROM admin_audit_log {where_sql}"
|
||
with self._connect() as conn:
|
||
row = conn.execute(sql, tuple(params)).fetchone()
|
||
return int((row["c"] if row and row["c"] is not None else 0) or 0)
|
||
|
||
def upsert_channel_identity(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
channel: str,
|
||
external_user_id: str,
|
||
user_id: str,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO channel_identity (tenant_id, channel, external_user_id, user_id, created_at)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
ON CONFLICT(tenant_id, channel, external_user_id) DO UPDATE SET
|
||
user_id = excluded.user_id
|
||
""",
|
||
(str(tenant_id), str(channel), str(external_user_id), str(user_id), ts),
|
||
)
|
||
|
||
def upsert_channel_identity_v2(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
channel: str,
|
||
account_id: str,
|
||
external_user_id: str,
|
||
user_id: str,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO channel_identity_v2 (tenant_id, channel, account_id, external_user_id, user_id, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(tenant_id, channel, account_id, external_user_id) DO UPDATE SET
|
||
user_id = excluded.user_id
|
||
""",
|
||
(str(tenant_id), str(channel), str(account_id), str(external_user_id), str(user_id), ts),
|
||
)
|
||
|
||
def resolve_user_by_channel_identity(
|
||
self, *, channel: str, external_user_id: str
|
||
) -> dict[str, Any] | None:
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT ci.tenant_id, ci.user_id, u.display_name, u.role
|
||
FROM channel_identity ci
|
||
JOIN app_user u ON u.id = ci.user_id
|
||
WHERE ci.channel = ? AND ci.external_user_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(str(channel), str(external_user_id)),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
return {
|
||
"tenant_id": row["tenant_id"],
|
||
"user_id": row["user_id"],
|
||
"display_name": row["display_name"],
|
||
"role": row["role"],
|
||
}
|
||
|
||
def resolve_user_by_channel_identity_v2(
|
||
self, *, channel: str, account_id: str, external_user_id: str
|
||
) -> dict[str, Any] | None:
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT ci.tenant_id, ci.user_id, u.display_name, u.role
|
||
FROM channel_identity_v2 ci
|
||
JOIN app_user u ON u.id = ci.user_id
|
||
WHERE ci.channel = ? AND ci.account_id = ? AND ci.external_user_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(str(channel), str(account_id), str(external_user_id)),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
return {
|
||
"tenant_id": row["tenant_id"],
|
||
"user_id": row["user_id"],
|
||
"display_name": row["display_name"],
|
||
"role": row["role"],
|
||
}
|
||
|
||
def upsert_user_channel_account(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
user_id: str,
|
||
channel: str,
|
||
account_id: str,
|
||
name: str | None = None,
|
||
config: dict[str, Any] | None = None,
|
||
is_active: bool = True,
|
||
) -> None:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO user_channel_account
|
||
(tenant_id, user_id, channel, account_id, name, config, is_active, created_at, updated_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(tenant_id, user_id, channel, account_id) DO UPDATE SET
|
||
name = excluded.name,
|
||
config = excluded.config,
|
||
is_active = excluded.is_active,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(
|
||
str(tenant_id),
|
||
str(user_id),
|
||
str(channel),
|
||
str(account_id),
|
||
str(name or "").strip(),
|
||
json.dumps(config or {}, ensure_ascii=False, default=str),
|
||
1 if is_active else 0,
|
||
ts,
|
||
ts,
|
||
),
|
||
)
|
||
|
||
def find_user_by_channel_account(self, *, channel: str, account_id: str) -> dict[str, Any] | None:
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT uca.tenant_id, uca.user_id, uca.channel, uca.account_id, uca.name, uca.config, uca.is_active,
|
||
u.display_name, u.role
|
||
FROM user_channel_account uca
|
||
JOIN app_user u ON u.id = uca.user_id
|
||
WHERE uca.channel = ? AND uca.account_id = ? AND uca.is_active = 1
|
||
ORDER BY uca.updated_at DESC
|
||
LIMIT 1
|
||
""",
|
||
(str(channel), str(account_id)),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
try:
|
||
cfg = json.loads(row["config"])
|
||
except Exception:
|
||
cfg = {}
|
||
return {
|
||
"tenant_id": row["tenant_id"],
|
||
"user_id": row["user_id"],
|
||
"channel": row["channel"],
|
||
"account_id": row["account_id"],
|
||
"name": str(row["name"] or "").strip(),
|
||
"display_name": row["display_name"],
|
||
"role": row["role"],
|
||
"config": cfg,
|
||
"is_active": bool(int(row["is_active"] or 0)),
|
||
}
|
||
|
||
def list_user_channel_accounts(
|
||
self, *, tenant_id: str, user_id: str, channel: str = "wecom", include_inactive: bool = True
|
||
) -> list[dict[str, Any]]:
|
||
where = ["uca.tenant_id = ?", "uca.user_id = ?", "uca.channel = ?"]
|
||
params: list[Any] = [str(tenant_id), str(user_id), str(channel)]
|
||
if not include_inactive:
|
||
where.append("uca.is_active = 1")
|
||
sql = f"""
|
||
SELECT uca.tenant_id, uca.user_id, uca.channel, uca.account_id, uca.name, uca.config, uca.is_active,
|
||
uca.created_at, uca.updated_at
|
||
FROM user_channel_account uca
|
||
WHERE {' AND '.join(where)}
|
||
ORDER BY uca.updated_at DESC, uca.account_id ASC
|
||
"""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(sql, params).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for row in rows:
|
||
try:
|
||
cfg = json.loads(row["config"])
|
||
except Exception:
|
||
cfg = {}
|
||
out.append(
|
||
{
|
||
"tenant_id": row["tenant_id"],
|
||
"user_id": row["user_id"],
|
||
"channel": row["channel"],
|
||
"account_id": row["account_id"],
|
||
"name": str(row["name"] or "").strip(),
|
||
"config": cfg,
|
||
"is_active": bool(int(row["is_active"] or 0)),
|
||
"created_at": row["created_at"],
|
||
"updated_at": row["updated_at"],
|
||
}
|
||
)
|
||
return out
|
||
|
||
def delete_user_channel_account(self, *, tenant_id: str, user_id: str, channel: str, account_id: str) -> int:
|
||
tid, uid, aid = str(tenant_id), str(user_id), str(account_id)
|
||
if str(channel) == "wecom":
|
||
self.delete_setting(f"wecom:bot_secret:{tid}:{uid}:{aid}")
|
||
self.delete_setting(f"wecom:agent_secret:{tid}:{uid}:{aid}")
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"""
|
||
DELETE FROM user_channel_account
|
||
WHERE tenant_id = ? AND user_id = ? AND channel = ? AND account_id = ?
|
||
""",
|
||
(tid, uid, str(channel), aid),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
|
||
def backfill_user_channel_account_names(self, *, channel: str = "wecom") -> int:
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE user_channel_account
|
||
SET name = account_id
|
||
WHERE channel = ? AND TRIM(COALESCE(name, '')) = ''
|
||
""",
|
||
(str(channel),),
|
||
)
|
||
return int(cur.rowcount or 0)
|
||
|
||
def create_bind_code(self, *, tenant_id: str, role: str, code: str) -> dict[str, Any]:
|
||
ts = utc_now_iso()
|
||
self._bind_code_repo().insert_bind_code(
|
||
code=str(code), tenant_id=str(tenant_id), role=str(role), created_at=ts
|
||
)
|
||
return {"code": code, "tenant_id": tenant_id, "role": role, "created_at": ts}
|
||
|
||
def consume_bind_code(
|
||
self, *, code: str, channel: str, external_user_id: str, display_name: str | None = None
|
||
) -> dict[str, Any] | None:
|
||
"""Consume a bind code and create/bind a new user. Returns binding info, or None if invalid/used."""
|
||
code = str(code or "").strip()
|
||
if not code:
|
||
return None
|
||
row = self._bind_code_repo().fetch_by_code(code=code)
|
||
if not row or row.get("used_at"):
|
||
return None
|
||
tenant_id = str(row["tenant_id"])
|
||
role = str(row["role"] or "member")
|
||
user = self.create_user(tenant_id=tenant_id, display_name=display_name or "User", role=role)
|
||
self.upsert_channel_identity(
|
||
tenant_id=tenant_id,
|
||
channel=channel,
|
||
external_user_id=external_user_id,
|
||
user_id=str(user["id"]),
|
||
)
|
||
ts = utc_now_iso()
|
||
self._bind_code_repo().mark_used(
|
||
code=code, used_at=ts, used_by_external_user_id=str(external_user_id)
|
||
)
|
||
return {"tenant_id": tenant_id, "user_id": user["id"], "role": role}
|
||
|
||
def list_bind_codes(self, *, tenant_id: str | None = None, limit: int = 200) -> list[dict[str, Any]]:
|
||
lim = max(1, int(limit))
|
||
rows = self._bind_code_repo().list_bind_codes(tenant_id=tenant_id, limit=lim)
|
||
return [
|
||
{
|
||
"code": r["code"],
|
||
"tenant_id": r["tenant_id"],
|
||
"role": r["role"],
|
||
"created_at": r["created_at"],
|
||
"used_at": r["used_at"],
|
||
"used_by_external_user_id": r["used_by_external_user_id"],
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def list_channel_identities(
|
||
self, *, tenant_id: str | None = None, channel: str | None = None, limit: int = 300
|
||
) -> list[dict[str, Any]]:
|
||
lim = max(1, int(limit))
|
||
where = []
|
||
params: list[Any] = []
|
||
if tenant_id:
|
||
where.append("ci.tenant_id = ?")
|
||
params.append(str(tenant_id))
|
||
if channel:
|
||
where.append("ci.channel = ?")
|
||
params.append(str(channel))
|
||
wsql = ("WHERE " + " AND ".join(where)) if where else ""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT ci.tenant_id, ci.channel, ci.external_user_id, ci.user_id, ci.created_at,
|
||
u.display_name, u.role
|
||
FROM channel_identity ci
|
||
JOIN app_user u ON u.id = ci.user_id
|
||
{wsql}
|
||
ORDER BY ci.created_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(*params, lim),
|
||
).fetchall()
|
||
return [
|
||
{
|
||
"tenant_id": r["tenant_id"],
|
||
"channel": r["channel"],
|
||
"external_user_id": r["external_user_id"],
|
||
"user_id": r["user_id"],
|
||
"display_name": r["display_name"],
|
||
"role": r["role"],
|
||
"created_at": r["created_at"],
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def list_channel_identities_v2(
|
||
self,
|
||
*,
|
||
tenant_id: str | None = None,
|
||
channel: str | None = None,
|
||
account_id: str | None = None,
|
||
user_id: str | None = None,
|
||
limit: int = 300,
|
||
) -> list[dict[str, Any]]:
|
||
lim = max(1, int(limit))
|
||
where = []
|
||
params: list[Any] = []
|
||
if tenant_id:
|
||
where.append("ci.tenant_id = ?")
|
||
params.append(str(tenant_id))
|
||
if channel:
|
||
where.append("ci.channel = ?")
|
||
params.append(str(channel))
|
||
if account_id:
|
||
where.append("ci.account_id = ?")
|
||
params.append(str(account_id))
|
||
if user_id:
|
||
where.append("ci.user_id = ?")
|
||
params.append(str(user_id))
|
||
wsql = ("WHERE " + " AND ".join(where)) if where else ""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT ci.tenant_id, ci.channel, ci.account_id, ci.external_user_id, ci.user_id, ci.created_at,
|
||
u.display_name, u.role,
|
||
COALESCE(uca.name, '') AS account_name
|
||
FROM channel_identity_v2 ci
|
||
JOIN app_user u ON u.id = ci.user_id
|
||
LEFT JOIN user_channel_account uca
|
||
ON uca.tenant_id = ci.tenant_id
|
||
AND uca.user_id = ci.user_id
|
||
AND uca.channel = ci.channel
|
||
AND uca.account_id = ci.account_id
|
||
{wsql}
|
||
ORDER BY ci.created_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(*params, lim),
|
||
).fetchall()
|
||
return [
|
||
{
|
||
"tenant_id": r["tenant_id"],
|
||
"channel": r["channel"],
|
||
"account_id": r["account_id"],
|
||
"account_name": str(r["account_name"] or "").strip(),
|
||
"external_user_id": r["external_user_id"],
|
||
"user_id": r["user_id"],
|
||
"display_name": r["display_name"],
|
||
"role": r["role"],
|
||
"created_at": r["created_at"],
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def get_or_create_channel_session(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
channel: str,
|
||
external_chat_id: str,
|
||
external_user_id: str,
|
||
session_title: str,
|
||
) -> str:
|
||
"""Return a stable session_id for a given channel chat context."""
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT session_id
|
||
FROM channel_session
|
||
WHERE tenant_id = ? AND channel = ? AND external_chat_id = ? AND external_user_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(str(tenant_id), str(channel), str(external_chat_id), str(external_user_id)),
|
||
).fetchone()
|
||
if row and row["session_id"]:
|
||
return str(row["session_id"])
|
||
sess = self.create_session(session_title)
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO channel_session
|
||
(tenant_id, channel, external_chat_id, external_user_id, session_id, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(str(tenant_id), str(channel), str(external_chat_id), str(external_user_id), str(sess.id), ts),
|
||
)
|
||
return str(sess.id)
|
||
|
||
def get_or_create_channel_session_v2(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
channel: str,
|
||
account_id: str,
|
||
external_chat_id: str,
|
||
external_user_id: str,
|
||
session_title: str,
|
||
) -> str:
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT session_id
|
||
FROM channel_session_v2
|
||
WHERE tenant_id = ? AND channel = ? AND account_id = ? AND external_chat_id = ? AND external_user_id = ?
|
||
LIMIT 1
|
||
""",
|
||
(str(tenant_id), str(channel), str(account_id), str(external_chat_id), str(external_user_id)),
|
||
).fetchone()
|
||
if row and row["session_id"]:
|
||
return str(row["session_id"])
|
||
sess = self.create_session(session_title)
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO channel_session_v2
|
||
(tenant_id, channel, account_id, external_chat_id, external_user_id, session_id, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(str(tenant_id), str(channel), str(account_id), str(external_chat_id), str(external_user_id), str(sess.id), ts),
|
||
)
|
||
return str(sess.id)
|
||
|
||
def backfill_ui_session_owner_from_channel_v2(self) -> int:
|
||
return self._ui_session_owner_repo().backfill_from_channel_v2(created_at=utc_now_iso())
|
||
|
||
# ----------------------------
|
||
# Oclaw tasks
|
||
# ----------------------------
|
||
def oclaw_task_create(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
session_id: str,
|
||
task_type: str = "async_turn",
|
||
payload: dict[str, Any] | None = None,
|
||
) -> OclawTask:
|
||
tid = str(uuid.uuid4())
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO oclaw_task
|
||
(id, tenant_id, session_id, task_type, status, payload, result, attempt_count, claimed_by, lease_expires_at, last_error, created_at, updated_at, finished_at)
|
||
VALUES (?, ?, ?, ?, 'queued', ?, '{}', 0, NULL, NULL, '', ?, ?, NULL)
|
||
""",
|
||
(
|
||
tid,
|
||
str(tenant_id),
|
||
str(session_id),
|
||
str(task_type or "async_turn"),
|
||
json.dumps(payload or {}, ensure_ascii=False),
|
||
ts,
|
||
ts,
|
||
),
|
||
)
|
||
got = self.oclaw_task_get(task_id=tid)
|
||
if not got:
|
||
raise RuntimeError("failed to create oclaw task")
|
||
return got
|
||
|
||
def oclaw_task_claim(
|
||
self,
|
||
*,
|
||
worker_id: str,
|
||
lease_seconds: int = 90,
|
||
task_type: str | None = None,
|
||
) -> OclawTask | None:
|
||
now = datetime.now(timezone.utc)
|
||
lease_expires = (now + timedelta(seconds=max(15, min(int(lease_seconds or 90), 1800)))).isoformat()
|
||
now_iso = now.isoformat()
|
||
with self._connect() as conn:
|
||
where = "(status = 'queued' OR (status = 'claimed' AND (lease_expires_at IS NULL OR lease_expires_at <= ?)))"
|
||
params: list[Any] = [now_iso]
|
||
tt = str(task_type or "").strip()
|
||
if tt:
|
||
where = f"{where} AND task_type = ?"
|
||
params.append(tt)
|
||
row = conn.execute(
|
||
f"""
|
||
SELECT id FROM oclaw_task
|
||
WHERE {where}
|
||
ORDER BY created_at ASC
|
||
LIMIT 1
|
||
""",
|
||
tuple(params),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
task_id = str(row["id"])
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE oclaw_task
|
||
SET status = 'claimed',
|
||
attempt_count = attempt_count + 1,
|
||
claimed_by = ?,
|
||
lease_expires_at = ?,
|
||
updated_at = ?
|
||
WHERE id = ?
|
||
AND (status = 'queued' OR (status = 'claimed' AND (lease_expires_at IS NULL OR lease_expires_at <= ?)))
|
||
""",
|
||
(str(worker_id), lease_expires, now_iso, task_id, now_iso),
|
||
)
|
||
if not cur.rowcount:
|
||
return None
|
||
return self.oclaw_task_get(task_id=task_id)
|
||
|
||
def oclaw_task_finish(self, *, task_id: str, result: dict[str, Any] | None = None) -> bool:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE oclaw_task
|
||
SET status = 'done',
|
||
result = ?,
|
||
lease_expires_at = NULL,
|
||
last_error = '',
|
||
updated_at = ?,
|
||
finished_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(json.dumps(result or {}, ensure_ascii=False), ts, ts, str(task_id)),
|
||
)
|
||
return bool(cur.rowcount and cur.rowcount > 0)
|
||
|
||
def oclaw_task_fail(self, *, task_id: str, error: str, result: dict[str, Any] | None = None) -> bool:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE oclaw_task
|
||
SET status = 'failed',
|
||
result = ?,
|
||
lease_expires_at = NULL,
|
||
last_error = ?,
|
||
updated_at = ?,
|
||
finished_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(json.dumps(result or {}, ensure_ascii=False), str(error or "")[:2000], ts, ts, str(task_id)),
|
||
)
|
||
return bool(cur.rowcount and cur.rowcount > 0)
|
||
|
||
def oclaw_task_get(self, *, task_id: str, tenant_id: str | None = None) -> OclawTask | None:
|
||
where = "id = ?"
|
||
params: list[Any] = [str(task_id)]
|
||
if tenant_id:
|
||
where += " AND tenant_id = ?"
|
||
params.append(str(tenant_id))
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT id, tenant_id, session_id, task_type, status, payload, result, attempt_count, claimed_by, lease_expires_at, last_error, created_at, updated_at, finished_at
|
||
FROM oclaw_task
|
||
WHERE """ + where + """
|
||
LIMIT 1
|
||
""",
|
||
tuple(params),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
return OclawTask(
|
||
id=str(row["id"]),
|
||
tenant_id=str(row["tenant_id"]),
|
||
session_id=str(row["session_id"]),
|
||
task_type=str(row["task_type"]),
|
||
status=str(row["status"]),
|
||
payload=str(row["payload"] or "{}"),
|
||
result=str(row["result"] or "{}"),
|
||
attempt_count=int(row["attempt_count"] or 0),
|
||
claimed_by=str(row["claimed_by"]) if row["claimed_by"] else None,
|
||
lease_expires_at=str(row["lease_expires_at"]) if row["lease_expires_at"] else None,
|
||
last_error=str(row["last_error"] or ""),
|
||
created_at=str(row["created_at"]),
|
||
updated_at=str(row["updated_at"]),
|
||
finished_at=str(row["finished_at"]) if row["finished_at"] else None,
|
||
)
|
||
|
||
def oclaw_task_list(
|
||
self,
|
||
*,
|
||
status: str | None = None,
|
||
limit: int = 50,
|
||
tenant_id: str | None = None,
|
||
session_id: str | None = None,
|
||
) -> list[OclawTask]:
|
||
lim = max(1, min(int(limit or 50), 500))
|
||
where = ""
|
||
params: list[Any] = []
|
||
if status:
|
||
where = "status = ?"
|
||
params.append(str(status))
|
||
if tenant_id:
|
||
where = f"{where} AND tenant_id = ?" if where else "tenant_id = ?"
|
||
params.append(str(tenant_id))
|
||
if session_id:
|
||
where = f"{where} AND session_id = ?" if where else "session_id = ?"
|
||
params.append(str(session_id))
|
||
where_clause = f"WHERE {where}" if where else ""
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT id, tenant_id, session_id, task_type, status, payload, result, attempt_count, claimed_by, lease_expires_at, last_error, created_at, updated_at, finished_at
|
||
FROM oclaw_task
|
||
{where_clause}
|
||
ORDER BY created_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(*params, lim),
|
||
).fetchall()
|
||
out: list[OclawTask] = []
|
||
for row in rows:
|
||
out.append(
|
||
OclawTask(
|
||
id=str(row["id"]),
|
||
tenant_id=str(row["tenant_id"]),
|
||
session_id=str(row["session_id"]),
|
||
task_type=str(row["task_type"]),
|
||
status=str(row["status"]),
|
||
payload=str(row["payload"] or "{}"),
|
||
result=str(row["result"] or "{}"),
|
||
attempt_count=int(row["attempt_count"] or 0),
|
||
claimed_by=str(row["claimed_by"]) if row["claimed_by"] else None,
|
||
lease_expires_at=str(row["lease_expires_at"]) if row["lease_expires_at"] else None,
|
||
last_error=str(row["last_error"] or ""),
|
||
created_at=str(row["created_at"]),
|
||
updated_at=str(row["updated_at"]),
|
||
finished_at=str(row["finished_at"]) if row["finished_at"] else None,
|
||
)
|
||
)
|
||
return out
|
||
|
||
def oclaw_run_upsert(
|
||
self,
|
||
*,
|
||
run_id: str,
|
||
tenant_id: str,
|
||
session_id: str,
|
||
status: str,
|
||
payload: dict[str, Any] | None = None,
|
||
) -> bool:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO oclaw_run(run_id, tenant_id, session_id, status, payload, created_at, updated_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(run_id) DO UPDATE SET
|
||
status = excluded.status,
|
||
payload = excluded.payload,
|
||
updated_at = excluded.updated_at
|
||
""",
|
||
(
|
||
str(run_id),
|
||
str(tenant_id),
|
||
str(session_id),
|
||
str(status),
|
||
json.dumps(payload or {}, ensure_ascii=False),
|
||
ts,
|
||
ts,
|
||
),
|
||
)
|
||
return True
|
||
|
||
def oclaw_attempt_append(
|
||
self,
|
||
*,
|
||
run_id: str,
|
||
tenant_id: str,
|
||
session_id: str,
|
||
attempt_no: int,
|
||
status: str,
|
||
reason: str = "",
|
||
payload: dict[str, Any] | None = None,
|
||
) -> int:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
if self._use_pg:
|
||
cur = conn.execute(
|
||
"""
|
||
INSERT INTO oclaw_attempt(run_id, tenant_id, session_id, attempt_no, status, reason, payload, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
RETURNING id
|
||
""",
|
||
(
|
||
str(run_id),
|
||
str(tenant_id),
|
||
str(session_id),
|
||
int(attempt_no),
|
||
str(status),
|
||
str(reason or ""),
|
||
json.dumps(payload or {}, ensure_ascii=False),
|
||
ts,
|
||
),
|
||
)
|
||
row = cur.fetchone()
|
||
return int(row["id"]) if row else 0
|
||
cur = conn.execute(
|
||
"""
|
||
INSERT INTO oclaw_attempt(run_id, tenant_id, session_id, attempt_no, status, reason, payload, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
str(run_id),
|
||
str(tenant_id),
|
||
str(session_id),
|
||
int(attempt_no),
|
||
str(status),
|
||
str(reason or ""),
|
||
json.dumps(payload or {}, ensure_ascii=False),
|
||
ts,
|
||
),
|
||
)
|
||
return int(cur.lastrowid or 0)
|
||
|
||
def oclaw_attempt_list(self, *, run_id: str, limit: int = 30) -> list[dict[str, Any]]:
|
||
lim = max(1, min(int(limit or 30), 200))
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT id, run_id, tenant_id, session_id, attempt_no, status, reason, payload, created_at
|
||
FROM oclaw_attempt
|
||
WHERE run_id = ?
|
||
ORDER BY attempt_no ASC
|
||
LIMIT ?
|
||
""",
|
||
(str(run_id), lim),
|
||
).fetchall()
|
||
out: list[dict[str, Any]] = []
|
||
for row in rows:
|
||
out.append(
|
||
{
|
||
"id": int(row["id"]),
|
||
"run_id": str(row["run_id"]),
|
||
"tenant_id": str(row["tenant_id"]),
|
||
"session_id": str(row["session_id"]),
|
||
"attempt_no": int(row["attempt_no"]),
|
||
"status": str(row["status"]),
|
||
"reason": str(row["reason"] or ""),
|
||
"payload": str(row["payload"] or "{}"),
|
||
"created_at": str(row["created_at"]),
|
||
}
|
||
)
|
||
return out
|
||
|
||
def oclaw_run_get(self, *, run_id: str, tenant_id: str | None = None) -> OclawRun | None:
|
||
where = "run_id = ?"
|
||
params: list[Any] = [str(run_id)]
|
||
if tenant_id:
|
||
where += " AND tenant_id = ?"
|
||
params.append(str(tenant_id))
|
||
with self._connect() as conn:
|
||
row = conn.execute(
|
||
f"""
|
||
SELECT run_id, tenant_id, session_id, status, payload, created_at, updated_at
|
||
FROM oclaw_run
|
||
WHERE {where}
|
||
LIMIT 1
|
||
""",
|
||
tuple(params),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
return OclawRun(
|
||
run_id=str(row["run_id"]),
|
||
tenant_id=str(row["tenant_id"]),
|
||
session_id=str(row["session_id"]),
|
||
status=str(row["status"]),
|
||
payload=str(row["payload"] or "{}"),
|
||
created_at=str(row["created_at"]),
|
||
updated_at=str(row["updated_at"]),
|
||
)
|
||
|
||
def oclaw_run_list(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
session_id: str | None = None,
|
||
status: str | None = None,
|
||
limit: int = 50,
|
||
) -> list[OclawRun]:
|
||
lim = max(1, min(int(limit or 50), 300))
|
||
where = "tenant_id = ?"
|
||
params: list[Any] = [str(tenant_id)]
|
||
if session_id:
|
||
where += " AND session_id = ?"
|
||
params.append(str(session_id))
|
||
if status:
|
||
where += " AND status = ?"
|
||
params.append(str(status))
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT run_id, tenant_id, session_id, status, payload, created_at, updated_at
|
||
FROM oclaw_run
|
||
WHERE {where}
|
||
ORDER BY updated_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(*params, lim),
|
||
).fetchall()
|
||
out: list[OclawRun] = []
|
||
for row in rows:
|
||
out.append(
|
||
OclawRun(
|
||
run_id=str(row["run_id"]),
|
||
tenant_id=str(row["tenant_id"]),
|
||
session_id=str(row["session_id"]),
|
||
status=str(row["status"]),
|
||
payload=str(row["payload"] or "{}"),
|
||
created_at=str(row["created_at"]),
|
||
updated_at=str(row["updated_at"]),
|
||
)
|
||
)
|
||
return out
|
||
|
||
# ----------------------------
|
||
# Todo items
|
||
# ----------------------------
|
||
def todo_create(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
owner_user_id: str,
|
||
title: str,
|
||
due_at: str | None = None,
|
||
assignee_user_id: str | None = None,
|
||
) -> dict[str, Any]:
|
||
tid = str(uuid.uuid4())
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO todo_item
|
||
(id, tenant_id, owner_user_id, assignee_user_id, title, due_at, status, created_at, updated_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
tid,
|
||
str(tenant_id),
|
||
str(owner_user_id),
|
||
str(assignee_user_id) if assignee_user_id else None,
|
||
str(title or "").strip(),
|
||
str(due_at) if due_at else None,
|
||
"open",
|
||
ts,
|
||
ts,
|
||
),
|
||
)
|
||
return {
|
||
"id": tid,
|
||
"tenant_id": tenant_id,
|
||
"owner_user_id": owner_user_id,
|
||
"assignee_user_id": assignee_user_id,
|
||
"title": title,
|
||
"due_at": due_at,
|
||
"status": "open",
|
||
"created_at": ts,
|
||
"updated_at": ts,
|
||
}
|
||
|
||
def todo_list(
|
||
self,
|
||
*,
|
||
tenant_id: str,
|
||
assignee_user_id: str | None = None,
|
||
status: str | None = "open",
|
||
limit: int = 50,
|
||
) -> list[dict[str, Any]]:
|
||
lim = max(1, min(int(limit), 200))
|
||
where = ["tenant_id = ?"]
|
||
params: list[Any] = [str(tenant_id)]
|
||
if assignee_user_id:
|
||
where.append("assignee_user_id = ?")
|
||
params.append(str(assignee_user_id))
|
||
if status:
|
||
where.append("status = ?")
|
||
params.append(str(status))
|
||
wsql = " AND ".join(where)
|
||
with self._connect() as conn:
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT id, tenant_id, owner_user_id, assignee_user_id, title, due_at, status, created_at, updated_at
|
||
FROM todo_item
|
||
WHERE {wsql}
|
||
ORDER BY updated_at DESC
|
||
LIMIT ?
|
||
""",
|
||
(*params, lim),
|
||
).fetchall()
|
||
return [
|
||
{
|
||
"id": r["id"],
|
||
"tenant_id": r["tenant_id"],
|
||
"owner_user_id": r["owner_user_id"],
|
||
"assignee_user_id": r["assignee_user_id"],
|
||
"title": r["title"],
|
||
"due_at": r["due_at"],
|
||
"status": r["status"],
|
||
"created_at": r["created_at"],
|
||
"updated_at": r["updated_at"],
|
||
}
|
||
for r in rows
|
||
]
|
||
|
||
def todo_set_status(self, *, tenant_id: str, todo_id: str, status: str) -> bool:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE todo_item
|
||
SET status = ?, updated_at = ?
|
||
WHERE tenant_id = ? AND id = ?
|
||
""",
|
||
(str(status), ts, str(tenant_id), str(todo_id)),
|
||
)
|
||
return bool(cur.rowcount and cur.rowcount > 0)
|
||
|
||
def todo_assign(self, *, tenant_id: str, todo_id: str, assignee_user_id: str) -> bool:
|
||
ts = utc_now_iso()
|
||
with self._connect() as conn:
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE todo_item
|
||
SET assignee_user_id = ?, updated_at = ?
|
||
WHERE tenant_id = ? AND id = ?
|
||
""",
|
||
(str(assignee_user_id), ts, str(tenant_id), str(todo_id)),
|
||
)
|
||
return bool(cur.rowcount and cur.rowcount > 0)
|