JOJO
b472f9fe5e
- 989 行单文件拆为 server/context/ 9 模块（identity/broadcast/personalization/usage/upload/conversation/resources/decorators/reaper）+ __init__.py 兼容 re-export，外部 import 路径不变 - 新增 RuntimeIdentity 显式身份快照（identity.py），get_user_resources 支持参数化身份驱动，为任务线程拆除 Flask 隐式上下文做准备 - 测试 patch 目标随迁至 server.context.personalization；新增 test_runtime_identity_resources 覆盖 web/host/api 身份路由矩阵
38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
"""配额追踪器（UsageTracker）的获取与广播。"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Optional
|
||
|
||
from modules.usage_tracker import UsageTracker
|
||
from server import state
|
||
|
||
|
||
def get_or_create_usage_tracker(username: Optional[str], workspace: Optional['modules.user_manager.UserWorkspace'] = None) -> Optional[UsageTracker]:
|
||
if not username:
|
||
return None
|
||
tracker = state.usage_trackers.get(username)
|
||
if tracker:
|
||
return tracker
|
||
from modules.user_manager import UserWorkspace # noqa: F401 # 保持类型引用兼容
|
||
if workspace is None:
|
||
workspace = state.user_manager.ensure_user_workspace(username)
|
||
record = state.user_manager.get_user(username)
|
||
role = getattr(record, "role", "user") if record else "user"
|
||
tracker = UsageTracker(str(workspace.data_dir), role=role or "user")
|
||
state.usage_trackers[username] = tracker
|
||
return tracker
|
||
|
||
|
||
def emit_user_quota_update(username: Optional[str]):
|
||
from server.extensions import socketio
|
||
if not username:
|
||
return
|
||
tracker = get_or_create_usage_tracker(username)
|
||
if not tracker:
|
||
return
|
||
try:
|
||
snapshot = tracker.get_quota_snapshot()
|
||
socketio.emit('quota_update', {'quotas': snapshot}, room=f"user_{username}")
|
||
except Exception:
|
||
pass
|