Odysseus v1.0
This commit is contained in:
pewdiepie-archdaemon
53
core/__init__.py
Normal file
53
core/__init__.py
Normal file
@@ -0,0 +1,53 @@
|
||||
# core/__init__.py
|
||||
"""
|
||||
Chat Core — the essential chat experience.
|
||||
|
||||
This package contains only what's needed for:
|
||||
- Streaming LLM responses
|
||||
- Session management
|
||||
- Model routing
|
||||
- Authentication
|
||||
"""
|
||||
|
||||
from src.llm_core import (
|
||||
llm_call,
|
||||
llm_call_async,
|
||||
stream_llm,
|
||||
list_model_ids,
|
||||
normalize_model_id,
|
||||
LLMConfig,
|
||||
)
|
||||
from .auth import AuthManager
|
||||
from .constants import *
|
||||
from .middleware import SecurityHeadersMiddleware
|
||||
from .exceptions import (
|
||||
SessionNotFoundError,
|
||||
InvalidFileUploadError,
|
||||
LLMServiceError,
|
||||
WebSearchError,
|
||||
)
|
||||
from .models import Session, ChatMessage
|
||||
from .session_manager import SessionManager
|
||||
|
||||
__all__ = [
|
||||
# LLM
|
||||
"llm_call",
|
||||
"llm_call_async",
|
||||
"stream_llm",
|
||||
"list_model_ids",
|
||||
"normalize_model_id",
|
||||
"LLMConfig",
|
||||
# Auth
|
||||
"AuthManager",
|
||||
# Middleware
|
||||
"SecurityHeadersMiddleware",
|
||||
# Exceptions
|
||||
"SessionNotFoundError",
|
||||
"InvalidFileUploadError",
|
||||
"LLMServiceError",
|
||||
"WebSearchError",
|
||||
# Models
|
||||
"Session",
|
||||
"ChatMessage",
|
||||
"SessionManager",
|
||||
]
|
||||
43
core/atomic_io.py
Normal file
43
core/atomic_io.py
Normal file
@@ -0,0 +1,43 @@
|
||||
"""Atomic JSON file writes.
|
||||
|
||||
Use this everywhere a JSON config file is persisted. A plain `open("w") +
|
||||
json.dump` truncates the file on first write and only fills it with new
|
||||
content afterwards — a kill -9 / power loss / OOM in between produces a
|
||||
truncated or empty file. For password DBs (`auth.json`) and live state
|
||||
(`sessions.json`, `settings.json`, `integrations.json`, `cookbook_state.json`),
|
||||
that's a data-loss event.
|
||||
|
||||
`atomic_write_json` writes to a sibling tmp file, fsyncs, then `os.replace`s
|
||||
into place. On POSIX `os.replace` is atomic on the same filesystem.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
|
||||
def atomic_write_json(path: str, data: Any, *, indent: Optional[int] = None) -> None:
|
||||
"""Atomically persist `data` as JSON at `path`.
|
||||
|
||||
The temp file uses the live PID as a suffix so two processes saving the
|
||||
same file (e.g. unit tests) don't collide on the rename target.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
tmp = f"{path}.tmp.{os.getpid()}"
|
||||
with open(tmp, "w") as f:
|
||||
json.dump(data, f, indent=indent)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp, path)
|
||||
|
||||
|
||||
def atomic_write_text(path: str, text: str) -> None:
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
tmp = f"{path}.tmp.{os.getpid()}"
|
||||
with open(tmp, "w") as f:
|
||||
f.write(text)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp, path)
|
||||
426
core/auth.py
Normal file
426
core/auth.py
Normal file
@@ -0,0 +1,426 @@
|
||||
"""
|
||||
Authentication module — multi-user password hashing, session tokens, config persistence.
|
||||
Config stored in data/auth.json. Uses bcrypt directly.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
import bcrypt
|
||||
import pyotp
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
from core.atomic_io import atomic_write_json as _atomic_write_json # noqa: E402
|
||||
|
||||
DEFAULT_PRIVILEGES = {
|
||||
"can_use_agent": True,
|
||||
"can_use_browser": True,
|
||||
"can_use_bash": False,
|
||||
"can_use_documents": True,
|
||||
"can_use_research": True,
|
||||
"can_generate_images": True,
|
||||
"can_manage_memory": True,
|
||||
"max_messages_per_day": 0,
|
||||
"allowed_models": [],
|
||||
}
|
||||
|
||||
# Admins get everything
|
||||
ADMIN_PRIVILEGES = {k: (True if isinstance(v, bool) else (0 if isinstance(v, int) else [])) for k, v in DEFAULT_PRIVILEGES.items()}
|
||||
|
||||
DEFAULT_AUTH_PATH = os.path.join(
|
||||
Path(__file__).parent.parent, "data", "auth.json"
|
||||
)
|
||||
TOKEN_TTL = 60 * 60 * 24 * 7 # 7 days
|
||||
|
||||
|
||||
def _hash_password(password: str) -> str:
|
||||
return bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")
|
||||
|
||||
|
||||
def _verify_password(password: str, hashed: str) -> bool:
|
||||
return bcrypt.checkpw(password.encode("utf-8"), hashed.encode("utf-8"))
|
||||
|
||||
|
||||
class AuthManager:
|
||||
"""Manages multi-user password + session-token auth system."""
|
||||
|
||||
def __init__(self, auth_path: str = DEFAULT_AUTH_PATH):
|
||||
self.auth_path = auth_path
|
||||
self._sessions_path = os.path.join(os.path.dirname(auth_path), "sessions.json")
|
||||
self._config: Dict[str, Any] = {}
|
||||
self._sessions: Dict[str, Dict[str, Any]] = {} # token -> {username, expiry}
|
||||
# Guards mutations of self._sessions and the on-disk sessions.json.
|
||||
# Validate/create/revoke run concurrently from the FastAPI threadpool.
|
||||
self._sessions_lock = threading.RLock()
|
||||
self._load()
|
||||
self._load_sessions()
|
||||
self._migrate_single_user()
|
||||
self._migrate_legacy_admin_role()
|
||||
|
||||
def _load(self):
|
||||
try:
|
||||
if os.path.exists(self.auth_path):
|
||||
with open(self.auth_path, "r") as f:
|
||||
self._config = json.load(f)
|
||||
logger.info("Auth config loaded")
|
||||
else:
|
||||
self._config = {}
|
||||
logger.info("No auth config found — first-run setup required")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load auth config: {e}")
|
||||
self._config = {}
|
||||
|
||||
def _load_sessions(self):
|
||||
"""Load persisted session tokens from disk, pruning expired ones."""
|
||||
try:
|
||||
if os.path.exists(self._sessions_path):
|
||||
with open(self._sessions_path, "r") as f:
|
||||
data = json.load(f)
|
||||
now = time.time()
|
||||
self._sessions = {k: v for k, v in data.items() if v.get("expiry", 0) > now}
|
||||
pruned = len(data) - len(self._sessions)
|
||||
if pruned > 0:
|
||||
self._save_sessions()
|
||||
logger.info(f"Loaded {len(self._sessions)} session(s) from disk")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load sessions: {e}")
|
||||
self._sessions = {}
|
||||
|
||||
def _save_sessions(self):
|
||||
"""Persist session tokens to disk (atomic, lock-guarded)."""
|
||||
try:
|
||||
with self._sessions_lock:
|
||||
snapshot = dict(self._sessions)
|
||||
_atomic_write_json(self._sessions_path, snapshot)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save sessions: {e}")
|
||||
|
||||
def _migrate_single_user(self):
|
||||
"""Migrate old single-user format to multi-user format."""
|
||||
if "password_hash" in self._config and "users" not in self._config:
|
||||
old_user = self._config.get("username", "admin")
|
||||
old_hash = self._config["password_hash"]
|
||||
self._config = {
|
||||
"users": {
|
||||
old_user: {
|
||||
"password_hash": old_hash,
|
||||
"created": time.time(),
|
||||
"is_admin": True,
|
||||
}
|
||||
}
|
||||
}
|
||||
self._save()
|
||||
logger.info(f"Migrated single-user auth to multi-user (admin: {old_user})")
|
||||
|
||||
def _migrate_legacy_admin_role(self):
|
||||
"""Normalize setup.py's old role='admin' marker to is_admin=True."""
|
||||
changed = False
|
||||
for username, user in self.users.items():
|
||||
if user.get("role") == "admin" and "is_admin" not in user:
|
||||
user["is_admin"] = True
|
||||
changed = True
|
||||
logger.info(f"Migrated legacy admin role for '{username}'")
|
||||
if changed:
|
||||
self._save()
|
||||
|
||||
def _save(self):
|
||||
_atomic_write_json(self.auth_path, self._config, indent=2)
|
||||
|
||||
@property
|
||||
def users(self) -> Dict[str, Any]:
|
||||
return self._config.get("users", {})
|
||||
|
||||
@property
|
||||
def signup_enabled(self) -> bool:
|
||||
return self._config.get("signup_enabled", False)
|
||||
|
||||
@signup_enabled.setter
|
||||
def signup_enabled(self, value: bool):
|
||||
self._config["signup_enabled"] = value
|
||||
self._save()
|
||||
|
||||
@property
|
||||
def is_configured(self) -> bool:
|
||||
return len(self.users) > 0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Account management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def setup(self, username: str, password: str) -> bool:
|
||||
"""First-run admin setup. Only works if no users exist."""
|
||||
if self.is_configured:
|
||||
return False
|
||||
return self.create_user(username, password, is_admin=True)
|
||||
|
||||
def create_user(self, username: str, password: str, is_admin: bool = False) -> bool:
|
||||
"""Create a new user account."""
|
||||
username = username.strip().lower()
|
||||
if username in self.users:
|
||||
return False
|
||||
if "users" not in self._config:
|
||||
self._config["users"] = {}
|
||||
self._config["users"][username] = {
|
||||
"password_hash": _hash_password(password),
|
||||
"created": time.time(),
|
||||
"is_admin": is_admin,
|
||||
"privileges": dict(ADMIN_PRIVILEGES if is_admin else DEFAULT_PRIVILEGES),
|
||||
}
|
||||
self._save()
|
||||
logger.info(f"Created user '{username}' (admin={is_admin})")
|
||||
return True
|
||||
|
||||
def delete_user(self, username: str, requesting_user: str) -> bool:
|
||||
"""Delete a user. Only admins can delete, and can't delete themselves.
|
||||
|
||||
SECURITY: also revoke every active session token belonging to this
|
||||
user so any open browser tab they have gets kicked back to /login
|
||||
on the next request. Without this the user kept full access until
|
||||
their cookie expired naturally (default ~30 days).
|
||||
"""
|
||||
username = username.strip().lower()
|
||||
if username not in self.users:
|
||||
return False
|
||||
if username == requesting_user:
|
||||
return False
|
||||
if not self.users.get(requesting_user, {}).get("is_admin"):
|
||||
return False
|
||||
del self._config["users"][username]
|
||||
self._save()
|
||||
# Purge all sessions belonging to this user. validate_token doesn't
|
||||
# cross-check `self.users`, so without this step a deleted user's
|
||||
# cookie keeps authenticating.
|
||||
revoked = 0
|
||||
with self._sessions_lock:
|
||||
to_drop = [tok for tok, sess in self._sessions.items()
|
||||
if (sess or {}).get("username") == username]
|
||||
for tok in to_drop:
|
||||
self._sessions.pop(tok, None)
|
||||
revoked += 1
|
||||
if revoked:
|
||||
self._save_sessions()
|
||||
logger.info(f"Deleted user '{username}' (by {requesting_user}); revoked {revoked} active session(s)")
|
||||
return True
|
||||
|
||||
def is_admin(self, username: str) -> bool:
|
||||
return self.users.get(username, {}).get("is_admin", False)
|
||||
|
||||
def list_users(self) -> List[Dict[str, Any]]:
|
||||
return [
|
||||
{"username": u, "is_admin": d.get("is_admin", False), "privileges": self.get_privileges(u)}
|
||||
for u, d in self.users.items()
|
||||
]
|
||||
|
||||
def get_privileges(self, username: str) -> Dict[str, Any]:
|
||||
"""Get privileges for a user. Admins get all privileges."""
|
||||
user = self.users.get(username, {})
|
||||
if user.get("is_admin"):
|
||||
return dict(ADMIN_PRIVILEGES)
|
||||
# Merge stored privileges with defaults (in case new privileges were added)
|
||||
stored = user.get("privileges", {})
|
||||
return {**DEFAULT_PRIVILEGES, **stored}
|
||||
|
||||
def set_privileges(self, username: str, privileges: Dict[str, Any]) -> bool:
|
||||
"""Update privileges for a user. Can't modify admin privileges."""
|
||||
username = username.strip().lower()
|
||||
if username not in self.users:
|
||||
return False
|
||||
if self.users[username].get("is_admin"):
|
||||
return False # admins always have full access
|
||||
# Only allow known privilege keys
|
||||
current = self.get_privileges(username)
|
||||
for k, v in privileges.items():
|
||||
if k in DEFAULT_PRIVILEGES:
|
||||
current[k] = v
|
||||
self._config["users"][username]["privileges"] = current
|
||||
self._save()
|
||||
logger.info(f"Updated privileges for '{username}': {current}")
|
||||
return True
|
||||
|
||||
def change_password(self, username: str, current_password: str, new_password: str) -> bool:
|
||||
username = username.strip().lower()
|
||||
if username not in self.users:
|
||||
return False
|
||||
if not _verify_password(current_password, self.users[username]["password_hash"]):
|
||||
return False
|
||||
self._config["users"][username]["password_hash"] = _hash_password(new_password)
|
||||
self._save()
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# TOTP two-factor authentication
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def totp_enabled(self, username: str) -> bool:
|
||||
"""Check if 2FA is enabled for a user."""
|
||||
user = self.users.get(username.strip().lower(), {})
|
||||
return bool(user.get("totp_enabled"))
|
||||
|
||||
def totp_generate_secret(self, username: str) -> Optional[str]:
|
||||
"""Generate a new TOTP secret for a user. Returns the secret (not yet enabled)."""
|
||||
username = username.strip().lower()
|
||||
if username not in self.users:
|
||||
return None
|
||||
secret = pyotp.random_base32()
|
||||
self._config["users"][username]["totp_secret_pending"] = secret
|
||||
self._save()
|
||||
return secret
|
||||
|
||||
def totp_get_provisioning_uri(self, username: str, secret: str) -> str:
|
||||
"""Get the otpauth:// URI for QR code generation."""
|
||||
totp = pyotp.TOTP(secret)
|
||||
return totp.provisioning_uri(name=username, issuer_name="Odysseus")
|
||||
|
||||
def totp_confirm_enable(self, username: str, code: str) -> bool:
|
||||
"""Verify a TOTP code against the pending secret, then enable 2FA."""
|
||||
username = username.strip().lower()
|
||||
user = self.users.get(username, {})
|
||||
secret = user.get("totp_secret_pending")
|
||||
if not secret:
|
||||
return False
|
||||
totp = pyotp.TOTP(secret)
|
||||
if not totp.verify(code, valid_window=1):
|
||||
return False
|
||||
# Enable 2FA
|
||||
self._config["users"][username]["totp_secret"] = secret
|
||||
self._config["users"][username]["totp_enabled"] = True
|
||||
self._config["users"][username].pop("totp_secret_pending", None)
|
||||
# Generate backup codes
|
||||
backup = [secrets.token_hex(4) for _ in range(8)]
|
||||
self._config["users"][username]["totp_backup_codes"] = backup
|
||||
self._save()
|
||||
logger.info(f"2FA enabled for '{username}'")
|
||||
return True
|
||||
|
||||
def totp_verify(self, username: str, code: str) -> bool:
|
||||
"""Verify a TOTP code for login."""
|
||||
username = username.strip().lower()
|
||||
user = self.users.get(username, {})
|
||||
if not user.get("totp_enabled"):
|
||||
return True # 2FA not enabled, always pass
|
||||
secret = user.get("totp_secret")
|
||||
if not secret:
|
||||
return True
|
||||
# Check backup codes first
|
||||
backup = user.get("totp_backup_codes", [])
|
||||
if code in backup:
|
||||
backup.remove(code)
|
||||
self._config["users"][username]["totp_backup_codes"] = backup
|
||||
self._save()
|
||||
logger.info(f"Backup code used for '{username}' ({len(backup)} remaining)")
|
||||
return True
|
||||
totp = pyotp.TOTP(secret)
|
||||
return totp.verify(code, valid_window=1)
|
||||
|
||||
def totp_disable(self, username: str, password: str) -> bool:
|
||||
"""Disable 2FA for a user. Requires password confirmation."""
|
||||
username = username.strip().lower()
|
||||
if not self.verify_password(username, password):
|
||||
return False
|
||||
self._config["users"][username].pop("totp_secret", None)
|
||||
self._config["users"][username].pop("totp_secret_pending", None)
|
||||
self._config["users"][username].pop("totp_backup_codes", None)
|
||||
self._config["users"][username]["totp_enabled"] = False
|
||||
self._save()
|
||||
logger.info(f"2FA disabled for '{username}'")
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Login / logout / session tokens
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def verify_password(self, username: str, password: str) -> bool:
|
||||
username = username.strip().lower()
|
||||
if username not in self.users:
|
||||
return False
|
||||
return _verify_password(password, self.users[username]["password_hash"])
|
||||
|
||||
def create_session(self, username: str, password: str) -> Optional[str]:
|
||||
"""Verify credentials and return a session token, or None."""
|
||||
username = username.strip().lower()
|
||||
if not self.verify_password(username, password):
|
||||
return None
|
||||
token = secrets.token_hex(32)
|
||||
with self._sessions_lock:
|
||||
self._sessions[token] = {
|
||||
"username": username,
|
||||
"expiry": time.time() + TOKEN_TTL,
|
||||
}
|
||||
self._save_sessions()
|
||||
return token
|
||||
|
||||
def validate_token(self, token: Optional[str]) -> bool:
|
||||
if not token:
|
||||
return False
|
||||
expired = False
|
||||
deleted_user = False
|
||||
with self._sessions_lock:
|
||||
session = self._sessions.get(token)
|
||||
if session is None:
|
||||
return False
|
||||
if time.time() > session["expiry"]:
|
||||
self._sessions.pop(token, None)
|
||||
expired = True
|
||||
else:
|
||||
# SECURITY: if the user record has since been removed (admin
|
||||
# deleted them while their cookie was still valid), drop the
|
||||
# session so the next request kicks them out instead of
|
||||
# silently authenticating against a non-existent account.
|
||||
if session.get("username") not in self.users:
|
||||
self._sessions.pop(token, None)
|
||||
deleted_user = True
|
||||
if expired or deleted_user:
|
||||
self._save_sessions()
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_username_for_token(self, token: Optional[str]) -> Optional[str]:
|
||||
"""Return the username associated with a valid token."""
|
||||
if not token:
|
||||
return None
|
||||
expired = False
|
||||
deleted_user = False
|
||||
with self._sessions_lock:
|
||||
session = self._sessions.get(token)
|
||||
if session is None:
|
||||
return None
|
||||
if time.time() > session["expiry"]:
|
||||
self._sessions.pop(token, None)
|
||||
expired = True
|
||||
else:
|
||||
_u = session["username"]
|
||||
# SECURITY: orphan check — same rationale as validate_token.
|
||||
if _u not in self.users:
|
||||
self._sessions.pop(token, None)
|
||||
deleted_user = True
|
||||
else:
|
||||
return _u
|
||||
if expired or deleted_user:
|
||||
self._save_sessions()
|
||||
return None
|
||||
|
||||
def revoke_token(self, token: str):
|
||||
with self._sessions_lock:
|
||||
self._sessions.pop(token, None)
|
||||
self._save_sessions()
|
||||
|
||||
def status(self, token: Optional[str]) -> Dict[str, Any]:
|
||||
username = self.get_username_for_token(token)
|
||||
authenticated = username is not None
|
||||
result = {
|
||||
"configured": self.is_configured,
|
||||
"authenticated": authenticated,
|
||||
"username": username,
|
||||
"is_admin": self.is_admin(username) if username else False,
|
||||
}
|
||||
if authenticated:
|
||||
result["privileges"] = self.get_privileges(username)
|
||||
return result
|
||||
40
core/constants.py
Normal file
40
core/constants.py
Normal file
@@ -0,0 +1,40 @@
|
||||
# src/constants.py
|
||||
"""Application-wide constants and configuration values."""
|
||||
import os
|
||||
|
||||
APP_VERSION = "0.9.1"
|
||||
|
||||
# Base paths
|
||||
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/"
|
||||
STATIC_DIR = os.path.join(BASE_DIR, "static")
|
||||
DATA_DIR = os.path.join(BASE_DIR, "data")
|
||||
|
||||
# Data file paths
|
||||
SESSIONS_FILE = os.path.join(DATA_DIR, "sessions.json")
|
||||
MEMORY_FILE = os.path.join(DATA_DIR, "memory.json")
|
||||
MEMORY_DOC = os.path.join(DATA_DIR, "memory_doc.md")
|
||||
PERSONAL_DIR = os.path.join(DATA_DIR, "personal_docs")
|
||||
RUNBOOK_DIR = os.path.join(PERSONAL_DIR, "runbook")
|
||||
UPLOAD_DIR = os.path.join(DATA_DIR, "uploads")
|
||||
FEATURES_FILE = os.path.join(DATA_DIR, "features.json")
|
||||
SETTINGS_FILE = os.path.join(DATA_DIR, "settings.json")
|
||||
|
||||
# API Configuration
|
||||
MAX_CONTEXT_MESSAGES = 90
|
||||
REQUEST_TIMEOUT = 20
|
||||
OPENAI_COMPAT_PATH = "/v1/chat/completions"
|
||||
|
||||
# Environment variables with defaults
|
||||
DEFAULT_HOST = os.getenv("LLM_HOST", "localhost")
|
||||
LLM_HOSTS = [h.strip() for h in os.getenv("LLM_HOSTS", "").split(",") if h.strip()]
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
SEARXNG_INSTANCE = os.getenv('SEARXNG_INSTANCE', 'http://localhost:8080')
|
||||
|
||||
|
||||
# Cleanup configuration
|
||||
CLEANUP_ENABLED = os.getenv("CLEANUP_ENABLED", "True").lower() == "true"
|
||||
CLEANUP_INTERVAL_HOURS = int(os.getenv("CLEANUP_INTERVAL_HOURS", "24"))
|
||||
|
||||
# Default parameters
|
||||
DEFAULT_TEMPERATURE = 1.0
|
||||
DEFAULT_MAX_TOKENS = 0
|
||||
1776
core/database.py
Normal file
1776
core/database.py
Normal file
File diff suppressed because it is too large
Load Diff
29
core/exceptions.py
Normal file
29
core/exceptions.py
Normal file
@@ -0,0 +1,29 @@
|
||||
# src/exceptions.py
|
||||
"""Custom exceptions for the application."""
|
||||
|
||||
class SessionNotFoundError(Exception):
|
||||
"""Raised when a requested session is not found."""
|
||||
def __init__(self, session_id: str):
|
||||
self.session_id = session_id
|
||||
super().__init__(f"Session '{session_id}' not found")
|
||||
|
||||
class InvalidFileUploadError(Exception):
|
||||
"""Raised when a file upload fails validation."""
|
||||
def __init__(self, message: str, filename: str = None):
|
||||
self.filename = filename
|
||||
self.message = message
|
||||
super().__init__(message)
|
||||
|
||||
class LLMServiceError(Exception):
|
||||
"""Raised when there is an error communicating with the LLM service."""
|
||||
def __init__(self, message: str, endpoint: str = None):
|
||||
self.endpoint = endpoint
|
||||
self.message = message
|
||||
super().__init__(message)
|
||||
|
||||
class WebSearchError(Exception):
|
||||
"""Raised when there is an error with web search functionality."""
|
||||
def __init__(self, message: str, query: str = None):
|
||||
self.query = query
|
||||
self.message = message
|
||||
super().__init__(message)
|
||||
100
core/middleware.py
Normal file
100
core/middleware.py
Normal file
@@ -0,0 +1,100 @@
|
||||
# src/middleware.py
|
||||
# Shared middleware, decorators, and request helpers
|
||||
|
||||
import os
|
||||
import secrets
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.responses import Response
|
||||
|
||||
|
||||
# Per-process token that lets the in-app tool layer hit admin-gated
|
||||
# routes via HTTP loopback (the agent's tool calls don't carry the
|
||||
# admin user's session cookie). Set once at import; tools read the
|
||||
# same value from this module. Never persisted or exposed externally.
|
||||
INTERNAL_TOOL_TOKEN = os.environ.get("ODYSSEUS_INTERNAL_TOKEN") or secrets.token_hex(32)
|
||||
INTERNAL_TOOL_HEADER = "X-Odysseus-Internal-Token"
|
||||
|
||||
|
||||
def require_admin(request: Request):
|
||||
"""Raise 403 if the current user isn't an admin.
|
||||
Allows access when auth is explicitly disabled, or when the request carries
|
||||
the in-process internal-tool token used by loopback agent tools.
|
||||
"""
|
||||
# In-process bypass for tool-layer loopback calls. Two paths:
|
||||
# (a) header-direct (caller set X-Odysseus-Internal-Token), or
|
||||
# (b) the auth middleware already validated the token and stamped
|
||||
# request.state.current_user = "internal-tool".
|
||||
try:
|
||||
if request.headers.get(INTERNAL_TOOL_HEADER) == INTERNAL_TOOL_TOKEN:
|
||||
return
|
||||
if getattr(request.state, "current_user", None) == "internal-tool":
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
auth_mgr = getattr(request.app.state, "auth_manager", None)
|
||||
if os.getenv("AUTH_ENABLED", "true").lower() == "false":
|
||||
return
|
||||
if not auth_mgr or not auth_mgr.is_configured:
|
||||
raise HTTPException(403, "Admin only")
|
||||
user = getattr(request.state, "current_user", None)
|
||||
if not user or not auth_mgr.is_admin(user):
|
||||
raise HTTPException(403, "Admin only")
|
||||
|
||||
|
||||
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
|
||||
"""Add standard security headers to all responses."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next) -> Response:
|
||||
# Generate a per-request nonce for inline scripts
|
||||
nonce = secrets.token_hex(16)
|
||||
request.state.csp_nonce = nonce
|
||||
|
||||
response = await call_next(request)
|
||||
path = request.url.path
|
||||
|
||||
# Tool render endpoints are served inside iframes — allow framing by self
|
||||
is_tool_render = path.startswith("/api/tools/") and path.endswith("/render")
|
||||
# Visual report pages are self-contained HTML — need inline scripts + external images
|
||||
is_report = path.startswith("/api/research/report/")
|
||||
|
||||
response.headers["X-Content-Type-Options"] = "nosniff"
|
||||
response.headers["Referrer-Policy"] = "no-referrer"
|
||||
|
||||
if is_report:
|
||||
response.headers["Content-Security-Policy"] = (
|
||||
"default-src 'self'; "
|
||||
"script-src 'self' 'unsafe-inline'; "
|
||||
"style-src 'self' 'unsafe-inline'; "
|
||||
"font-src 'self'; "
|
||||
"img-src 'self' data: blob: https:; "
|
||||
"connect-src 'self'; "
|
||||
"frame-ancestors 'none'"
|
||||
)
|
||||
elif is_tool_render:
|
||||
# Tool iframe content: skip all framing headers — the iframe's
|
||||
# sandbox="allow-scripts" attribute provides isolation.
|
||||
# Don't overwrite the route's own restrictive CSP either.
|
||||
pass
|
||||
else:
|
||||
response.headers["X-Frame-Options"] = "DENY"
|
||||
# NOTE: `style-src 'unsafe-inline'` is intentionally retained.
|
||||
# `static/index.html` and `static/login.html` ship inline <style>
|
||||
# blocks, and several JS modules build runtime `style=""` attrs.
|
||||
# Migrating to nonce-only requires templating the HTML files +
|
||||
# auditing every JS-set style attribute. Since inline styles
|
||||
# don't execute script, the residual risk is visual-only.
|
||||
response.headers["Content-Security-Policy"] = (
|
||||
"default-src 'self'; "
|
||||
f"script-src 'self' 'nonce-{nonce}' https://cdn.jsdelivr.net; "
|
||||
"style-src 'self' 'unsafe-inline' https://cdn.jsdelivr.net; "
|
||||
"font-src 'self' https://cdn.jsdelivr.net; "
|
||||
"img-src 'self' data: blob:; "
|
||||
"media-src 'self' blob:; "
|
||||
"connect-src 'self'; "
|
||||
"frame-src 'self'; "
|
||||
"frame-ancestors 'none'"
|
||||
)
|
||||
return response
|
||||
84
core/models.py
Normal file
84
core/models.py
Normal file
@@ -0,0 +1,84 @@
|
||||
# core/models.py
|
||||
"""
|
||||
Pure data models — no database logic, no side effects.
|
||||
|
||||
These are simple datacontainers. All persistence is handled by SessionManager.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Any, Optional, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .session_manager import SessionManager
|
||||
|
||||
# Module-level session manager reference (set at app startup)
|
||||
_session_manager: Optional["SessionManager"] = None
|
||||
|
||||
|
||||
def set_session_manager(manager: "SessionManager"):
|
||||
"""Set the global session manager reference."""
|
||||
global _session_manager
|
||||
_session_manager = manager
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatMessage:
|
||||
"""A single chat message."""
|
||||
role: str
|
||||
content: str
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dict for API responses."""
|
||||
result = {"role": self.role, "content": self.content}
|
||||
if self.metadata:
|
||||
result["metadata"] = self.metadata
|
||||
return result
|
||||
|
||||
def get(self, key: str, default=None):
|
||||
"""Dict-like access for compatibility."""
|
||||
return getattr(self, key, default)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
"""A chat session — pure data container."""
|
||||
id: str
|
||||
name: str
|
||||
endpoint_url: str
|
||||
model: str
|
||||
rag: bool = False
|
||||
archived: bool = False
|
||||
headers: Optional[Dict[str, str]] = None
|
||||
history: List[ChatMessage] = None
|
||||
owner: Optional[str] = None
|
||||
is_important: bool = False
|
||||
message_count: int = 0
|
||||
|
||||
def __post_init__(self):
|
||||
if self.history is None:
|
||||
self.history = []
|
||||
if self.headers is None:
|
||||
self.headers = {}
|
||||
|
||||
def add_message(self, message: ChatMessage):
|
||||
"""
|
||||
Add a message to this session.
|
||||
|
||||
Delegates to SessionManager for persistence if available,
|
||||
otherwise just appends to history.
|
||||
"""
|
||||
self.history.append(message)
|
||||
self.message_count = len(self.history)
|
||||
|
||||
# Delegate to session manager for persistence
|
||||
if _session_manager:
|
||||
_session_manager._persist_message(self.id, message)
|
||||
|
||||
def get_context_messages(self) -> List[Dict[str, Any]]:
|
||||
"""Get messages in format for LLM API."""
|
||||
return [msg.to_dict() for msg in self.history]
|
||||
|
||||
def get(self, key: str, default=None):
|
||||
"""Dict-like access for compatibility."""
|
||||
return getattr(self, key, default)
|
||||
558
core/session_manager.py
Normal file
558
core/session_manager.py
Normal file
@@ -0,0 +1,558 @@
|
||||
# core/session_manager.py
|
||||
"""
|
||||
Session management — all session business logic and DB operations.
|
||||
|
||||
This is the single place that handles:
|
||||
- Loading/saving sessions to database
|
||||
- Adding messages to sessions
|
||||
- Session lifecycle (create, archive, delete)
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
import logging
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from typing import Dict, Optional
|
||||
|
||||
from .database import Session as DbSession, ChatMessage as DbChatMessage, Document as DbDocument, SessionLocal
|
||||
from .models import Session, ChatMessage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""
|
||||
Manages chat sessions with database persistence.
|
||||
|
||||
Usage:
|
||||
manager = SessionManager()
|
||||
session = manager.create_session(id, name, url, model)
|
||||
manager.add_message(session.id, ChatMessage("user", "hello"))
|
||||
session = manager.get_session(session_id)
|
||||
"""
|
||||
|
||||
def __init__(self, sessions_file: str = None):
|
||||
# sessions_file kept for backward compat, not used
|
||||
self.sessions: Dict[str, Session] = {}
|
||||
self.load_sessions()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Loading
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def load_sessions(self):
|
||||
"""Load recent session METADATA from the database — messages are
|
||||
hydrated on demand by `get_session`. Previously this walked every
|
||||
message of every session into RAM at boot, which on a long-running
|
||||
personal-server box could be tens of thousands of rows held forever
|
||||
in `self.sessions`.
|
||||
"""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_sessions = db.query(DbSession).filter(
|
||||
DbSession.archived == False,
|
||||
DbSession.message_count > 0,
|
||||
).order_by(DbSession.last_accessed.desc()).limit(100).all()
|
||||
|
||||
loaded_count = 0
|
||||
for db_session in db_sessions:
|
||||
try:
|
||||
session = self._db_to_session_meta(db_session)
|
||||
if session is not None:
|
||||
self.sessions[db_session.id] = session
|
||||
loaded_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading session {db_session.id}: {e}")
|
||||
continue
|
||||
|
||||
logger.info(f"Loaded {loaded_count} session(s) (metadata only)")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading sessions: {e}")
|
||||
self.sessions = {}
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def _db_to_session_meta(self, db_session: DbSession) -> Optional[Session]:
|
||||
"""Build a Session with empty history. `get_session` will hydrate
|
||||
messages from the DB on first read."""
|
||||
headers = db_session.headers
|
||||
if isinstance(headers, str):
|
||||
try:
|
||||
headers = json.loads(headers)
|
||||
except json.JSONDecodeError:
|
||||
headers = {}
|
||||
session = Session(
|
||||
id=db_session.id,
|
||||
name=db_session.name,
|
||||
endpoint_url=db_session.endpoint_url,
|
||||
model=db_session.model,
|
||||
rag=db_session.rag,
|
||||
archived=db_session.archived,
|
||||
headers=headers,
|
||||
history=[],
|
||||
owner=getattr(db_session, "owner", None),
|
||||
is_important=getattr(db_session, "is_important", False) or False,
|
||||
)
|
||||
session.message_count = getattr(db_session, "message_count", 0) or 0
|
||||
return session
|
||||
|
||||
def _db_to_session(self, db_session: DbSession, db) -> Optional[Session]:
|
||||
"""Convert a database session to a Session object."""
|
||||
history = []
|
||||
|
||||
# Try relationship first, then direct query
|
||||
if db_session.messages:
|
||||
for db_msg in db_session.messages:
|
||||
meta = json.loads(db_msg.meta_data) if db_msg.meta_data else {}
|
||||
if meta is None: meta = {}
|
||||
meta['_db_id'] = db_msg.id
|
||||
history.append(ChatMessage(
|
||||
role=db_msg.role,
|
||||
content=db_msg.content,
|
||||
metadata=meta,
|
||||
))
|
||||
else:
|
||||
db_messages = db.query(DbChatMessage).filter(
|
||||
DbChatMessage.session_id == db_session.id
|
||||
).order_by(DbChatMessage.timestamp).all()
|
||||
|
||||
for db_msg in db_messages:
|
||||
meta = json.loads(db_msg.meta_data) if db_msg.meta_data else {}
|
||||
if meta is None: meta = {}
|
||||
meta['_db_id'] = db_msg.id
|
||||
history.append(ChatMessage(
|
||||
role=db_msg.role,
|
||||
content=db_msg.content,
|
||||
metadata=meta,
|
||||
))
|
||||
|
||||
if not history:
|
||||
return None
|
||||
|
||||
# Parse headers
|
||||
headers = db_session.headers
|
||||
if isinstance(headers, str):
|
||||
try:
|
||||
headers = json.loads(headers)
|
||||
except json.JSONDecodeError:
|
||||
headers = {}
|
||||
|
||||
session = Session(
|
||||
id=db_session.id,
|
||||
name=db_session.name,
|
||||
endpoint_url=db_session.endpoint_url,
|
||||
model=db_session.model,
|
||||
rag=db_session.rag,
|
||||
archived=db_session.archived,
|
||||
headers=headers,
|
||||
history=history,
|
||||
owner=getattr(db_session, 'owner', None),
|
||||
is_important=getattr(db_session, 'is_important', False) or False,
|
||||
)
|
||||
|
||||
session.message_count = getattr(db_session, 'message_count', len(history))
|
||||
return session
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Message operations
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def add_message(self, session_id: str, message: ChatMessage):
|
||||
"""
|
||||
Add a message to a session and persist to database.
|
||||
|
||||
Args:
|
||||
session_id: Session ID
|
||||
message: ChatMessage to add
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
session.history.append(message)
|
||||
session.message_count = len(session.history)
|
||||
|
||||
self._persist_message(session_id, message)
|
||||
|
||||
def _persist_message(self, session_id: str, message: ChatMessage):
|
||||
"""Persist a single message to the database."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
msg_id = str(uuid.uuid4())
|
||||
db_message = DbChatMessage(
|
||||
id=msg_id,
|
||||
session_id=session_id,
|
||||
role=message.role,
|
||||
content=message.content,
|
||||
meta_data=json.dumps(message.metadata) if message.metadata else None
|
||||
)
|
||||
db.add(db_message)
|
||||
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.message_count = len(self.sessions.get(session_id, {}).history) if session_id in self.sessions else 0
|
||||
_now = datetime.now(timezone.utc)
|
||||
db_session.last_accessed = _now
|
||||
# Clean "last conversation" timestamp — only bumped here on a
|
||||
# real message persist, so it powers an accurate "Last active"
|
||||
# sort that ignores renames / model swaps / mere opens.
|
||||
db_session.last_message_at = _now
|
||||
|
||||
db.commit()
|
||||
|
||||
# Store DB ID on the in-memory message for edit/delete by ID
|
||||
if message.metadata is None:
|
||||
message.metadata = {}
|
||||
message.metadata['_db_id'] = msg_id
|
||||
|
||||
logger.debug(f"Persisted message to session {session_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error persisting message: {e}")
|
||||
db.rollback()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def truncate_messages(self, session_id: str, keep_count: int) -> bool:
|
||||
"""Truncate session history, keeping only the first `keep_count` messages."""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
if keep_count < 0:
|
||||
return False
|
||||
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_messages = db.query(DbChatMessage).filter(
|
||||
DbChatMessage.session_id == session_id
|
||||
).order_by(DbChatMessage.timestamp).all()
|
||||
|
||||
deleted = 0
|
||||
for msg in db_messages[keep_count:]:
|
||||
db.delete(msg)
|
||||
deleted += 1
|
||||
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.message_count = keep_count
|
||||
db_session.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
db.commit()
|
||||
|
||||
# Update in-memory
|
||||
session.history = session.history[:keep_count]
|
||||
|
||||
logger.info(f"Truncated session {session_id} to {keep_count} messages")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error truncating session: {e}")
|
||||
db.rollback()
|
||||
return False
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def replace_messages(self, session_id: str, messages: list) -> bool:
|
||||
"""Replace a session's persisted and in-memory history atomically."""
|
||||
session = self.get_session(session_id)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db.query(DbChatMessage).filter(DbChatMessage.session_id == session_id).delete()
|
||||
now = datetime.now(timezone.utc)
|
||||
for i, message in enumerate(messages):
|
||||
msg_id = str(uuid.uuid4())
|
||||
db_message = DbChatMessage(
|
||||
id=msg_id,
|
||||
session_id=session_id,
|
||||
role=message.role,
|
||||
content=message.content,
|
||||
meta_data=json.dumps(message.metadata) if message.metadata else None,
|
||||
timestamp=now + timedelta(microseconds=i),
|
||||
)
|
||||
db.add(db_message)
|
||||
if message.metadata is None:
|
||||
message.metadata = {}
|
||||
message.metadata["_db_id"] = msg_id
|
||||
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.message_count = len(messages)
|
||||
db_session.updated_at = now
|
||||
db_session.last_accessed = now
|
||||
db_session.last_message_at = now
|
||||
|
||||
db.commit()
|
||||
session.history = list(messages)
|
||||
session.message_count = len(messages)
|
||||
logger.info("Replaced session %s history with %d messages", session_id, len(messages))
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("Error replacing session history: %s", e)
|
||||
db.rollback()
|
||||
return False
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Session CRUD
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_session(self, session_id: str) -> Session:
|
||||
"""Get a session by ID, loading from DB if needed.
|
||||
|
||||
Sessions seeded by `load_sessions` start with empty history. The
|
||||
first read here hydrates them with the message rows.
|
||||
"""
|
||||
if session_id not in self.sessions:
|
||||
self._load_session_from_db(session_id)
|
||||
else:
|
||||
cached = self.sessions[session_id]
|
||||
# Lazy hydrate: metadata-only entries get their messages on first read.
|
||||
if not cached.history and getattr(cached, "message_count", 0) > 0:
|
||||
self._load_session_from_db(session_id)
|
||||
|
||||
# Update last_accessed
|
||||
self._touch_session(session_id)
|
||||
|
||||
return self.sessions[session_id]
|
||||
|
||||
def _load_session_from_db(self, session_id: str):
|
||||
"""Hydrate a single session (with messages) from the database."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session is None:
|
||||
raise KeyError(f"Session {session_id} not found")
|
||||
|
||||
session = self._db_to_session(db_session, db)
|
||||
if session:
|
||||
self.sessions[session_id] = session
|
||||
else:
|
||||
# No messages — fall back to metadata-only entry so callers
|
||||
# don't crash on KeyError for empty sessions.
|
||||
meta = self._db_to_session_meta(db_session)
|
||||
if meta is None:
|
||||
raise KeyError(f"Session {session_id} could not be loaded")
|
||||
self.sessions[session_id] = meta
|
||||
|
||||
except KeyError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading session {session_id}: {e}")
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def _touch_session(self, session_id: str):
|
||||
"""Update last_accessed timestamp."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.last_accessed = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating last_accessed: {e}")
|
||||
db.rollback()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def create_session(
|
||||
self,
|
||||
session_id: str,
|
||||
name: str,
|
||||
endpoint_url: str,
|
||||
model: str,
|
||||
rag: bool = False,
|
||||
owner: str = None
|
||||
) -> Session:
|
||||
"""Create a new session and save to database."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_session = DbSession(
|
||||
id=session_id,
|
||||
name=name,
|
||||
endpoint_url=endpoint_url,
|
||||
model=model,
|
||||
rag=rag,
|
||||
headers={},
|
||||
owner=owner,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc)
|
||||
)
|
||||
db.add(db_session)
|
||||
db.commit()
|
||||
|
||||
session = Session(
|
||||
id=session_id,
|
||||
name=name,
|
||||
endpoint_url=endpoint_url,
|
||||
model=model,
|
||||
rag=rag,
|
||||
headers={},
|
||||
owner=owner,
|
||||
)
|
||||
|
||||
self.sessions[session_id] = session
|
||||
return session
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"Error creating session: {e}")
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def delete_session(self, session_id: str) -> bool:
|
||||
"""Permanently delete a session and all its messages."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
# Detach documents so they survive as orphans in the library
|
||||
db.query(DbDocument).filter(DbDocument.session_id == session_id).update(
|
||||
{DbDocument.session_id: None}, synchronize_session=False
|
||||
)
|
||||
|
||||
# Delete messages
|
||||
db.query(DbChatMessage).filter(DbChatMessage.session_id == session_id).delete()
|
||||
|
||||
# Delete session
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db.delete(db_session)
|
||||
db.commit()
|
||||
|
||||
if session_id in self.sessions:
|
||||
del self.sessions[session_id]
|
||||
|
||||
logger.info(f"Deleted session {session_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting session: {e}")
|
||||
db.rollback()
|
||||
return False
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Session updates
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def update_session_name(self, session_id: str, name: str):
|
||||
"""Update session name."""
|
||||
if session_id not in self.sessions:
|
||||
return
|
||||
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.name = name
|
||||
db_session.updated_at = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
self.sessions[session_id].name = name
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"Error updating session name: {e}")
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def archive_session(self, session_id: str):
|
||||
"""Archive a session."""
|
||||
if session_id not in self.sessions:
|
||||
return
|
||||
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.archived = True
|
||||
db_session.updated_at = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
self.sessions[session_id].archived = True
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"Error archiving session: {e}")
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
def mark_important(self, session_id: str, important: bool = True):
|
||||
"""Mark session as important."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
|
||||
if db_session:
|
||||
db_session.is_important = important
|
||||
db_session.updated_at = datetime.now(timezone.utc)
|
||||
db.commit()
|
||||
|
||||
if session_id in self.sessions:
|
||||
self.sessions[session_id].is_important = important
|
||||
else:
|
||||
raise KeyError(f"Session {session_id} not found")
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"Error marking session important: {e}")
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Queries
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_sessions_for_user(self, username: Optional[str] = None) -> Dict[str, Session]:
|
||||
"""Return sessions for a specific user (or all if username is None)."""
|
||||
if username is None:
|
||||
return self.sessions
|
||||
return {
|
||||
sid: s for sid, s in self.sessions.items()
|
||||
if s.owner == username
|
||||
}
|
||||
|
||||
def save_sessions(self):
|
||||
"""No-op for DB compatibility."""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Cleanup
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def cleanup_empty_sessions(self, auto_archive_days: int = 30) -> dict:
|
||||
"""Clean up empty and old sessions."""
|
||||
db = SessionLocal()
|
||||
stats = {'deleted_empty': 0, 'archived_old': 0, 'total_checked': 0}
|
||||
|
||||
try:
|
||||
all_sessions = db.query(DbSession).all()
|
||||
cutoff_date = datetime.now(timezone.utc) - timedelta(days=auto_archive_days)
|
||||
|
||||
for db_session in all_sessions:
|
||||
stats['total_checked'] += 1
|
||||
|
||||
# Delete empty sessions
|
||||
if db_session.message_count == 0:
|
||||
if db_session.id in self.sessions:
|
||||
del self.sessions[db_session.id]
|
||||
db.delete(db_session)
|
||||
stats['deleted_empty'] += 1
|
||||
|
||||
# Archive old sessions
|
||||
elif (not db_session.archived and
|
||||
db_session.last_accessed and
|
||||
db_session.last_accessed < cutoff_date and
|
||||
db_session.message_count > 0 and
|
||||
not getattr(db_session, 'is_important', False)):
|
||||
db_session.archived = True
|
||||
stats['archived_old'] += 1
|
||||
|
||||
db.commit()
|
||||
logger.info(f"Cleanup: {stats['deleted_empty']} deleted, {stats['archived_old']} archived")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Cleanup error: {e}")
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
return stats
|
||||
Reference in New Issue
Block a user