agent-Specialization/server/state.py

"""共享状态与常量，供各子模块使用。"""
from __future__ import annotations
import os
import threading
from collections import defaultdict, deque
from pathlib import Path
from typing import Dict, Any, Optional

from config import LOGS_DIR, PROJECT_MAX_STORAGE_BYTES, PROJECT_MAX_STORAGE_MB
from core.web_terminal import WebTerminal
from modules.custom_tool_registry import CustomToolRegistry
from modules.usage_tracker import UsageTracker
from modules.user_container_manager import UserContainerManager
from modules.user_manager import UserManager
from modules.api_user_manager import ApiUserManager

# 全局实例
user_manager = UserManager()
api_user_manager = ApiUserManager()
custom_tool_registry = CustomToolRegistry()
container_manager = UserContainerManager()
user_terminals: Dict[str, WebTerminal] = {}
terminal_rooms: Dict[str, set] = {}
connection_users: Dict[str, str] = {}
RECENT_UPLOAD_EVENT_LIMIT = 150
RECENT_UPLOAD_FEED_LIMIT = 60
stop_flags: Dict[str, Dict[str, Any]] = {}
active_polling_tasks: Dict[str, bool] = {}  # conversation_id -> is_polling

# 监控/限流/用量
MONITOR_FILE_TOOLS = {'write_file', 'edit_file'}
MONITOR_MEMORY_TOOLS = {'update_memory'}
MONITOR_SNAPSHOT_CHAR_LIMIT = 60000
MONITOR_MEMORY_ENTRY_LIMIT = 256
RATE_LIMIT_BUCKETS: Dict[str, deque] = defaultdict(deque)
FAILURE_TRACKERS: Dict[str, Dict[str, float]] = {}
pending_socket_tokens: Dict[str, Dict[str, Any]] = {}
usage_trackers: Dict[str, UsageTracker] = {}

MONITOR_SNAPSHOT_CACHE: Dict[str, Dict[str, Any]] = {}
MONITOR_SNAPSHOT_CACHE_LIMIT = 120
RECENT_UPLOAD_EVENT_LIMIT = 150
RECENT_UPLOAD_FEED_LIMIT = 60

# 路径与缓存设置（依赖项目配置）
PROJECT_STORAGE_CACHE: Dict[str, Dict[str, Any]] = {}
PROJECT_STORAGE_CACHE_TTL_SECONDS = float(os.environ.get("PROJECT_STORAGE_CACHE_TTL", "30"))

# 其他配置
DEFAULT_PORT = 8091
THINKING_FAILURE_KEYWORDS = ["⚠️", "🛑", "失败", "错误", "异常", "终止", "error", "failed", "未完成", "超时", "强制"]
CSRF_HEADER_NAME = "X-CSRF-Token"
CSRF_SESSION_KEY = "_csrf_token"
CSRF_SAFE_METHODS = {"GET", "HEAD", "OPTIONS", "TRACE"}
CSRF_PROTECTED_PATHS = {"/login", "/register", "/logout", "/host-login"}
CSRF_PROTECTED_PREFIXES = ("/api/",)
CSRF_EXEMPT_PATHS = {"/api/csrf-token"}
FAILED_LOGIN_LIMIT = 5
FAILED_LOGIN_LOCK_SECONDS = 300
SOCKET_TOKEN_TTL_SECONDS = 45
USER_IDLE_TIMEOUT_SECONDS = int(os.environ.get("USER_IDLE_TIMEOUT_SECONDS", "900"))
LAST_ACTIVE_FILE = Path(LOGS_DIR).expanduser().resolve() / "last_active.json"
_last_active_lock = threading.Lock()
_last_active_cache: Dict[str, float] = {}
_idle_reaper_started = False
TITLE_PROMPT_PATH = Path(__file__).resolve().parent.parent / "prompts" / "title_generation_prompt.txt"

# 项目存储限制常量也会被使用
PROJECT_MAX_STORAGE_BYTES = PROJECT_MAX_STORAGE_BYTES
PROJECT_MAX_STORAGE_MB = PROJECT_MAX_STORAGE_MB

__all__ = [
    "user_manager",
    "api_user_manager",
    "custom_tool_registry",
    "container_manager",
    "user_terminals",
    "terminal_rooms",
    "connection_users",
    "stop_flags",
    "MONITOR_FILE_TOOLS",
    "MONITOR_MEMORY_TOOLS",
    "MONITOR_SNAPSHOT_CHAR_LIMIT",
    "MONITOR_MEMORY_ENTRY_LIMIT",
    "RATE_LIMIT_BUCKETS",
    "FAILURE_TRACKERS",
    "pending_socket_tokens",
    "usage_trackers",
    "MONITOR_SNAPSHOT_CACHE",
    "MONITOR_SNAPSHOT_CACHE_LIMIT",
    "PROJECT_STORAGE_CACHE",
    "PROJECT_STORAGE_CACHE_TTL_SECONDS",
    "DEFAULT_PORT",
    "THINKING_FAILURE_KEYWORDS",
    "CSRF_HEADER_NAME",
    "CSRF_SESSION_KEY",
    "CSRF_SAFE_METHODS",
    "CSRF_PROTECTED_PATHS",
    "CSRF_PROTECTED_PREFIXES",
    "CSRF_EXEMPT_PATHS",
    "FAILED_LOGIN_LIMIT",
    "FAILED_LOGIN_LOCK_SECONDS",
    "SOCKET_TOKEN_TTL_SECONDS",
    "USER_IDLE_TIMEOUT_SECONDS",
    "LAST_ACTIVE_FILE",
    "_last_active_lock",
    "_last_active_cache",
    "_idle_reaper_started",
    "TITLE_PROMPT_PATH",
    "PROJECT_MAX_STORAGE_BYTES",
    "PROJECT_MAX_STORAGE_MB",
    "RECENT_UPLOAD_EVENT_LIMIT",
    "RECENT_UPLOAD_FEED_LIMIT",
    "make_stop_keys",
    "get_stop_flag",
    "set_stop_flag",
    "clear_stop_flag",
    "get_last_active_ts",
]


def get_last_active_ts(username: str, fallback: Optional[float] = None) -> Optional[float]:
    """
    返回最近活跃时间，优先使用缓存；当容器句柄中的时间更新、更晚时，自动刷新缓存。
    这样避免“缓存过旧导致刚触碰的容器被立即回收”的问题。
    """
    fallback_val: Optional[float]
    try:
        fallback_val = float(fallback) if fallback is not None else None
    except (TypeError, ValueError):
        fallback_val = None

    with _last_active_lock:
        cached = _last_active_cache.get(username)
        try:
            cached_val = float(cached) if cached is not None else None
        except (TypeError, ValueError):
            cached_val = None

        # 若没有缓存，或句柄时间更新、更晚，则刷新缓存
        if cached_val is None:
            if fallback_val is not None:
                _last_active_cache[username] = fallback_val
            return fallback_val

        if fallback_val is not None and fallback_val > cached_val:
            _last_active_cache[username] = fallback_val
            return fallback_val
    
        return cached_val


# ====== 停止标志辅助 ======
def make_stop_keys(client_sid: Optional[str] = None, username: Optional[str] = None):
    keys = []
    if client_sid:
        keys.append(client_sid)
    if username:
        keys.append(f"user:{username}")
    return keys


def set_stop_flag(client_sid: Optional[str], username: Optional[str], entry: Dict[str, Any]):
    for k in make_stop_keys(client_sid, username):
        stop_flags[k] = entry


def get_stop_flag(client_sid: Optional[str], username: Optional[str]) -> Optional[Dict[str, Any]]:
    for k in make_stop_keys(client_sid, username):
        val = stop_flags.get(k)
        if val:
            return val
    return None


def clear_stop_flag(client_sid: Optional[str], username: Optional[str]):
    for k in make_stop_keys(client_sid, username):
        stop_flags.pop(k, None)