433 lines
15 KiB
Python
433 lines
15 KiB
Python
"""Admission helper shared by step-2 init and background job start.
|
|
|
|
The durable database lease is the source of truth. Gateway run inspection is
|
|
kept as a reconciliation fallback for legacy/untracked runs. When the limit is
|
|
enabled, missing room storage fails closed so a deployment error cannot bypass
|
|
the cap.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from fastapi import HTTPException, Request
|
|
from sqlalchemy import select
|
|
|
|
from deerflow.config.system_settings import get_roundtable_room_limit_settings
|
|
from deerflow.persistence.roundtable_jobs.sql import compute_active_dedupe_key
|
|
from deerflow.runtime.user_context import get_effective_user_id
|
|
|
|
|
|
def _room_key_from_payload(meta: Any, owner_user_id: str | None = None) -> str | None:
|
|
"""Resolve the live-execution key from thread/run metadata."""
|
|
if not isinstance(meta, dict):
|
|
return None
|
|
stamped = str(meta.get("roundtable_room_key") or "").strip()
|
|
if stamped:
|
|
return stamped
|
|
draft_id = str(meta.get("draft_id") or "").strip() or None
|
|
if not draft_id:
|
|
return None
|
|
task_id = str(meta.get("task_id") or meta.get("taskId") or "").strip() or None
|
|
return compute_active_dedupe_key(draft_id, owner_user_id, task_id)
|
|
|
|
|
|
def _room_scope_key_from_payload(
|
|
meta: Any,
|
|
owner_user_id: str | None = None,
|
|
) -> str | None:
|
|
"""Legacy persisted scope used only to de-duplicate fallback run evidence."""
|
|
if not isinstance(meta, dict):
|
|
return None
|
|
draft_id = str(meta.get("draft_id") or "").strip() or None
|
|
if not draft_id:
|
|
return None
|
|
task_id = str(meta.get("task_id") or meta.get("taskId") or "").strip() or None
|
|
owner = owner_user_id or str(meta.get("user_id") or "").strip() or None
|
|
return compute_active_dedupe_key(draft_id, owner, task_id)
|
|
|
|
|
|
def _thread_id_of(record: Any) -> str:
|
|
return str(getattr(record, "thread_id", "") or "").strip()
|
|
|
|
|
|
async def _all_active_run_records(request: Request) -> list[Any]:
|
|
"""Every in-flight Gateway run, regardless of which account started it."""
|
|
records: list[Any] = []
|
|
seen_ids: set[str] = set()
|
|
|
|
def _append(items: list[Any] | None) -> None:
|
|
for item in items or []:
|
|
task = getattr(item, "task", None)
|
|
if task is not None and callable(getattr(task, "done", None)) and task.done():
|
|
# A terminal asyncio task with a lagging pending/running status
|
|
# is cleanup residue, not a live room. This also prevents one
|
|
# exceptional worker-finalization path from pinning capacity.
|
|
continue
|
|
run_id = str(getattr(item, "run_id", "") or "")
|
|
thread_id = _thread_id_of(item)
|
|
marker = run_id or thread_id
|
|
if not marker or marker in seen_ids:
|
|
continue
|
|
seen_ids.add(marker)
|
|
records.append(item)
|
|
|
|
gate = getattr(request.app.state, "concurrency_gate", None)
|
|
if gate is not None and getattr(gate, "enabled", False):
|
|
try:
|
|
leases = await gate.list_active()
|
|
except Exception:
|
|
leases = None
|
|
if leases:
|
|
_append(list(leases))
|
|
|
|
run_mgr = getattr(request.app.state, "run_manager", None)
|
|
if run_mgr is not None:
|
|
try:
|
|
_append(await run_mgr.list_active())
|
|
except Exception:
|
|
pass
|
|
return records
|
|
|
|
|
|
async def _room_keys_from_thread_rows(
|
|
thread_ids: list[str],
|
|
) -> tuple[dict[str, str | None], set[str]]:
|
|
"""Load room keys for *all* owners. Returns (keys, unresolved thread ids)."""
|
|
remaining = {tid for tid in thread_ids if tid}
|
|
keys: dict[str, str | None] = {}
|
|
if not remaining:
|
|
return keys, remaining
|
|
try:
|
|
from deerflow.persistence.engine import get_session_factory
|
|
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
|
|
|
session_factory = get_session_factory()
|
|
if session_factory is None:
|
|
return keys, remaining
|
|
async with session_factory() as session:
|
|
rows = (
|
|
await session.execute(
|
|
select(
|
|
ThreadMetaRow.thread_id,
|
|
ThreadMetaRow.user_id,
|
|
ThreadMetaRow.metadata_json,
|
|
).where(ThreadMetaRow.thread_id.in_(list(remaining)))
|
|
)
|
|
).all()
|
|
for thread_id, user_id, metadata_json in rows:
|
|
remaining.discard(thread_id)
|
|
key = _room_key_from_payload(metadata_json, user_id)
|
|
if key:
|
|
keys[key] = _room_scope_key_from_payload(metadata_json, user_id)
|
|
except Exception:
|
|
return keys, remaining
|
|
return keys, remaining
|
|
|
|
|
|
async def extra_live_roundtable_keys(request: Request) -> dict[str, str | None]:
|
|
"""Room keys of step-2 sessions whose LangGraph runs are still in flight.
|
|
|
|
Scans the whole Gateway process: every user's runs are included. Several
|
|
seat runs in one session share ``roundtable_room_key``. Intent and
|
|
recommend threads have no such key and are ignored.
|
|
"""
|
|
records = await _all_active_run_records(request)
|
|
keys: dict[str, str | None] = {}
|
|
missing_tids: list[str] = []
|
|
seen_threads: set[str] = set()
|
|
for rec in records:
|
|
tid = _thread_id_of(rec)
|
|
if not tid or tid in seen_threads:
|
|
continue
|
|
seen_threads.add(tid)
|
|
meta = getattr(rec, "metadata", None)
|
|
owner = str(getattr(rec, "user_id", "") or "").strip() or None
|
|
if owner is None and isinstance(meta, dict):
|
|
owner = str(meta.get("user_id") or "").strip() or None
|
|
key = _room_key_from_payload(meta if isinstance(meta, dict) else None, owner)
|
|
if key:
|
|
keys[key] = _room_scope_key_from_payload(meta, owner)
|
|
continue
|
|
missing_tids.append(tid)
|
|
|
|
if not missing_tids:
|
|
return keys
|
|
|
|
from_rows, remaining = await _room_keys_from_thread_rows(missing_tids)
|
|
keys.update(from_rows)
|
|
|
|
thread_store = getattr(request.app.state, "thread_store", None)
|
|
if thread_store is None or not remaining:
|
|
return keys
|
|
for tid in remaining:
|
|
try:
|
|
row = await thread_store.get(tid, user_id=None)
|
|
except Exception:
|
|
continue
|
|
if not isinstance(row, dict):
|
|
continue
|
|
owner = str(row.get("user_id") or "").strip() or None
|
|
key = _room_key_from_payload(row.get("metadata"), owner)
|
|
if key:
|
|
keys[key] = _room_scope_key_from_payload(row.get("metadata"), owner)
|
|
return keys
|
|
|
|
|
|
async def backend_occupancy_keys(request: Request) -> dict[str, str | None]:
|
|
"""Live Gateway runs used to reconcile legacy leases, all accounts."""
|
|
return await extra_live_roundtable_keys(request)
|
|
|
|
|
|
def current_user_id(request: Request) -> str:
|
|
user = getattr(request.state, "user", None)
|
|
if user is not None and getattr(user, "id", None):
|
|
return str(user.id)
|
|
return get_effective_user_id()
|
|
|
|
|
|
def _store(request: Request):
|
|
return getattr(request.app.state, "roundtable_room_store", None)
|
|
|
|
|
|
async def admit_roundtable_room(
|
|
request: Request,
|
|
*,
|
|
draft_id: str | None,
|
|
task_id: str | None,
|
|
room_session_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
settings = get_roundtable_room_limit_settings()
|
|
store = _store(request)
|
|
extra_keys = await backend_occupancy_keys(request)
|
|
clean_room_session_id = (room_session_id or "").strip() or None
|
|
if settings.enabled and clean_room_session_id is None:
|
|
return {
|
|
"ok": False,
|
|
"enabled": True,
|
|
"max_rooms": settings.max_rooms,
|
|
"active_count": len(extra_keys),
|
|
"renewed": False,
|
|
"code": "ROOM_SESSION_REQUIRED",
|
|
"message": "会商会话标识缺失，请刷新页面后重试。",
|
|
}
|
|
if store is None:
|
|
return {
|
|
"ok": not settings.enabled,
|
|
"enabled": settings.enabled,
|
|
"max_rooms": settings.max_rooms,
|
|
"active_count": len(extra_keys),
|
|
"renewed": False,
|
|
"code": "" if not settings.enabled else "ROOM_STORE_UNAVAILABLE",
|
|
"message": "" if not settings.enabled else "会商并发服务暂不可用，请稍后再试。",
|
|
}
|
|
user_id = current_user_id(request)
|
|
clean_draft = (draft_id or "").strip() or None
|
|
clean_task = (task_id or "").strip() or None
|
|
return await store.admit(
|
|
draft_id=clean_draft,
|
|
user_id=user_id,
|
|
task_id=clean_task,
|
|
room_session_id=clean_room_session_id,
|
|
enabled=settings.enabled,
|
|
max_rooms=settings.max_rooms,
|
|
max_hold_minutes=getattr(settings, "max_hold_minutes", 30),
|
|
extra_keys=extra_keys,
|
|
)
|
|
|
|
|
|
async def renew_roundtable_room(
|
|
request: Request,
|
|
*,
|
|
draft_id: str | None,
|
|
task_id: str | None,
|
|
room_session_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Renew an existing browser lease; never admit or recreate a room."""
|
|
settings = get_roundtable_room_limit_settings()
|
|
store = _store(request)
|
|
extra_keys = await backend_occupancy_keys(request)
|
|
if store is None:
|
|
return {
|
|
"ok": not settings.enabled,
|
|
"enabled": settings.enabled,
|
|
"max_rooms": settings.max_rooms,
|
|
"active_count": len(extra_keys),
|
|
"renewed": False,
|
|
"code": "" if not settings.enabled else "ROOM_STORE_UNAVAILABLE",
|
|
"message": "" if not settings.enabled else "会商并发服务暂不可用，请稍后再试。",
|
|
}
|
|
return await store.renew(
|
|
draft_id=(draft_id or "").strip() or None,
|
|
user_id=current_user_id(request),
|
|
task_id=(task_id or "").strip() or None,
|
|
room_session_id=(room_session_id or "").strip() or None,
|
|
enabled=settings.enabled,
|
|
max_rooms=settings.max_rooms,
|
|
max_hold_minutes=getattr(settings, "max_hold_minutes", 30),
|
|
extra_keys=extra_keys,
|
|
)
|
|
|
|
|
|
async def release_roundtable_room(
|
|
request: Request,
|
|
*,
|
|
draft_id: str | None,
|
|
task_id: str | None,
|
|
room_session_id: str | None = None,
|
|
) -> bool:
|
|
"""Best-effort release used by both the API route and disconnect cleanup."""
|
|
store = _store(request)
|
|
if store is None:
|
|
return False
|
|
return await store.release(
|
|
draft_id=(draft_id or "").strip() or None,
|
|
user_id=current_user_id(request),
|
|
task_id=(task_id or "").strip() or None,
|
|
room_session_id=(room_session_id or "").strip() or None,
|
|
)
|
|
|
|
|
|
async def enforce_roundtable_room(
|
|
request: Request,
|
|
*,
|
|
draft_id: str | None,
|
|
task_id: str | None,
|
|
room_session_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Admit the room or raise 409 before any step-2 model work starts."""
|
|
decision = await admit_roundtable_room(
|
|
request,
|
|
draft_id=draft_id,
|
|
task_id=task_id,
|
|
room_session_id=room_session_id,
|
|
)
|
|
if not decision.get("ok", False):
|
|
raise HTTPException(status_code=409, detail=decision)
|
|
return decision
|
|
|
|
|
|
async def enforce_existing_roundtable_room(
|
|
request: Request,
|
|
*,
|
|
draft_id: str | None,
|
|
task_id: str | None,
|
|
room_session_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Require and renew an existing lease without performing admission."""
|
|
decision = await renew_roundtable_room(
|
|
request,
|
|
draft_id=draft_id,
|
|
task_id=task_id,
|
|
room_session_id=room_session_id,
|
|
)
|
|
if not decision.get("ok", False):
|
|
raise HTTPException(status_code=409, detail=decision)
|
|
return decision
|
|
|
|
|
|
async def _thread_room_scope(
|
|
request: Request,
|
|
thread_ids: list[str],
|
|
) -> tuple[str, str | None, str | None, str | None] | None:
|
|
"""Resolve ``(room_key, draft_id, owner_user_id, task_id)`` from metadata."""
|
|
wanted = [tid for tid in dict.fromkeys(thread_ids) if tid]
|
|
if not wanted:
|
|
return None
|
|
|
|
thread_store = getattr(request.app.state, "thread_store", None)
|
|
if thread_store is not None:
|
|
for tid in wanted:
|
|
try:
|
|
row = await thread_store.get(tid, user_id=None)
|
|
except Exception:
|
|
continue
|
|
if not isinstance(row, dict):
|
|
continue
|
|
meta = row.get("metadata")
|
|
owner = str(row.get("user_id") or "").strip() or None
|
|
key = _room_key_from_payload(meta, owner)
|
|
if key:
|
|
data = meta if isinstance(meta, dict) else {}
|
|
draft_id = str(data.get("draft_id") or "").strip() or None
|
|
task_id = str(data.get("task_id") or data.get("taskId") or "").strip() or None
|
|
return key, draft_id, owner, task_id
|
|
|
|
try:
|
|
from deerflow.persistence.engine import get_session_factory
|
|
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
|
|
|
session_factory = get_session_factory()
|
|
if session_factory is None:
|
|
return None
|
|
async with session_factory() as session:
|
|
rows = (
|
|
await session.execute(
|
|
select(
|
|
ThreadMetaRow.user_id,
|
|
ThreadMetaRow.metadata_json,
|
|
).where(ThreadMetaRow.thread_id.in_(wanted))
|
|
)
|
|
).all()
|
|
for owner, meta in rows:
|
|
key = _room_key_from_payload(meta, owner)
|
|
if key:
|
|
data = meta if isinstance(meta, dict) else {}
|
|
draft_id = str(data.get("draft_id") or "").strip() or None
|
|
task_id = str(data.get("task_id") or data.get("taskId") or "").strip() or None
|
|
return key, draft_id, owner, task_id
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
async def enforce_roundtable_run(
|
|
request: Request,
|
|
*,
|
|
draft_id: str | None,
|
|
task_id: str | None,
|
|
room_session_id: str | None = None,
|
|
thread_ids: list[str],
|
|
) -> dict[str, Any] | None:
|
|
"""Enforce the room cap at the final model-producing entry point.
|
|
|
|
Current clients send the draft scope on every run. A model run may only
|
|
renew the room that was admitted when Step 2 began; it must never create a
|
|
new room while the user resumes, answers a clarification, or intervenes.
|
|
Thread metadata is a compatibility fallback for initialized sessions.
|
|
"""
|
|
clean_draft = (draft_id or "").strip() or None
|
|
clean_task = (task_id or "").strip() or None
|
|
if clean_draft:
|
|
return await enforce_existing_roundtable_room(
|
|
request,
|
|
draft_id=clean_draft,
|
|
task_id=clean_task,
|
|
room_session_id=(room_session_id or "").strip() or None,
|
|
)
|
|
|
|
scope = await _thread_room_scope(request, thread_ids)
|
|
if scope is None:
|
|
return None
|
|
room_key, _stored_draft, _owner_user_id, _stored_task = scope
|
|
settings = get_roundtable_room_limit_settings()
|
|
store = _store(request)
|
|
if store is None:
|
|
if settings.enabled:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail="Roundtable room store not available",
|
|
)
|
|
return None
|
|
decision = await store.renew_room_key(
|
|
room_key=room_key,
|
|
holder_user_id=current_user_id(request),
|
|
enabled=settings.enabled,
|
|
max_rooms=settings.max_rooms,
|
|
max_hold_minutes=getattr(settings, "max_hold_minutes", 30),
|
|
extra_keys=await backend_occupancy_keys(request),
|
|
)
|
|
if not decision.get("ok", False):
|
|
raise HTTPException(status_code=409, detail=decision)
|
|
return decision
|