604 lines
23 KiB
Python
604 lines
23 KiB
Python
"""Deployment-wide roundtable room leases and background-job occupancy.
|
||
|
||
A room is one live roundtable execution, not one saved draft, not one account,
|
||
and not one seat. Different pages running the same draft have distinct hidden
|
||
session ids; every seat and step of one execution shares one id. The active set
|
||
is the union of every user's:
|
||
|
||
- a durable foreground/session lease renewed by the browser or a run entry
|
||
- live ``extra_keys`` supplied by the Gateway as a reconciliation fallback
|
||
only when no durable lease/job already represents that persisted scope
|
||
- a background job that is executing or waiting for the user's Step-2
|
||
clarification (``queued`` / ``running`` / ``awaiting_input`` /
|
||
``cancel_requested``)
|
||
|
||
Historical jobs older than the hard room lifetime and merely viewing chat
|
||
history are not rooms. There is no per-user pool.
|
||
|
||
Admission and lease creation happen while holding one database-backed gate
|
||
row. Updating that row serializes contenders on SQLite, MySQL and PostgreSQL,
|
||
so four simultaneous requests cannot all observe the same stale count.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import hashlib
|
||
from collections.abc import Mapping
|
||
from datetime import UTC, datetime, timedelta
|
||
from typing import Any
|
||
|
||
from sqlalchemy import delete, select, update
|
||
from sqlalchemy.exc import IntegrityError
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||
|
||
from deerflow.persistence.roundtable_jobs.model import RoundtableJobRow
|
||
from deerflow.persistence.roundtable_jobs.sql import compute_active_dedupe_key
|
||
from deerflow.persistence.roundtable_rooms.model import RoundtableRoomLeaseRow
|
||
|
||
DEFAULT_LEASE_SECONDS = 5 * 60
|
||
MIN_LEASE_SECONDS = 60
|
||
MAX_LEASE_SECONDS = 60 * 60
|
||
JOB_LIVE_SECONDS = 15 * 60
|
||
OCCUPYING_JOB_STATUSES = ("queued", "running", "awaiting_input", "cancel_requested")
|
||
# A background job keeps running while Step 3 generates its artifacts, but the
|
||
# concurrency cap is intentionally scoped to Step 2. ``consensus`` is written
|
||
# before Step 3 begins, which releases the slot without waiting for reporting.
|
||
NON_OCCUPYING_JOB_PHASES = frozenset({
|
||
"consensus",
|
||
"summary_writing",
|
||
"report_drawing",
|
||
"report_done",
|
||
})
|
||
GATE_ROW_KEY = "__roundtable_room_gate__"
|
||
|
||
|
||
def compute_room_session_key(room_session_id: str | None) -> str | None:
|
||
"""Stable key for one live roundtable execution (not one saved draft)."""
|
||
clean = (room_session_id or "").strip()
|
||
if not clean:
|
||
return None
|
||
return hashlib.sha256(f"roundtable-session:{clean}".encode()).hexdigest()
|
||
|
||
|
||
def _aware(value: datetime | None) -> datetime | None:
|
||
if value is None:
|
||
return None
|
||
if value.tzinfo is None:
|
||
return value.replace(tzinfo=UTC)
|
||
return value
|
||
|
||
|
||
def _job_is_live(
|
||
lease_until: datetime | None,
|
||
updated_at: datetime | None,
|
||
now: datetime,
|
||
live_after: datetime,
|
||
) -> bool:
|
||
until = _aware(lease_until)
|
||
if until is not None and until > now:
|
||
return True
|
||
updated = _aware(updated_at)
|
||
return updated is not None and updated >= live_after
|
||
|
||
|
||
def normalize_lease_seconds(value: int | None) -> int:
|
||
"""Clamp a foreground room lease to the supported 1–60 minute window."""
|
||
try:
|
||
seconds = int(value) if value is not None else DEFAULT_LEASE_SECONDS
|
||
except (TypeError, ValueError):
|
||
seconds = DEFAULT_LEASE_SECONDS
|
||
return max(MIN_LEASE_SECONDS, min(MAX_LEASE_SECONDS, seconds))
|
||
|
||
|
||
def _hard_deadline_reached(
|
||
created_at: datetime | None,
|
||
now: datetime,
|
||
max_hold_minutes: int,
|
||
) -> bool:
|
||
created = _aware(created_at)
|
||
return created is not None and created + timedelta(minutes=max_hold_minutes) <= now
|
||
|
||
|
||
class RoundtableRoomRepository:
|
||
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
||
self._sf = session_factory
|
||
self._lock = asyncio.Lock()
|
||
|
||
async def admit(
|
||
self,
|
||
*,
|
||
draft_id: str | None,
|
||
user_id: str | None,
|
||
task_id: str | None,
|
||
room_session_id: str | None = None,
|
||
enabled: bool,
|
||
max_rooms: int,
|
||
max_hold_minutes: int = 30,
|
||
lease_seconds: int = DEFAULT_LEASE_SECONDS,
|
||
extra_keys: set[str] | Mapping[str, str | None] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""Atomically reserve/renew a room, or reject when the pool is full."""
|
||
key = compute_room_session_key(room_session_id) or compute_active_dedupe_key(
|
||
draft_id, user_id, task_id
|
||
)
|
||
return await self.admit_room_key(
|
||
room_key=key,
|
||
draft_id=draft_id,
|
||
user_id=user_id,
|
||
task_id=task_id,
|
||
holder_user_id=user_id,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
max_hold_minutes=max_hold_minutes,
|
||
lease_seconds=lease_seconds,
|
||
extra_keys=extra_keys,
|
||
)
|
||
|
||
async def admit_room_key(
|
||
self,
|
||
*,
|
||
room_key: str | None,
|
||
draft_id: str | None,
|
||
user_id: str | None,
|
||
task_id: str | None,
|
||
holder_user_id: str | None,
|
||
enabled: bool,
|
||
max_rooms: int,
|
||
max_hold_minutes: int = 30,
|
||
lease_seconds: int = DEFAULT_LEASE_SECONDS,
|
||
extra_keys: set[str] | Mapping[str, str | None] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""Atomically admit a pre-resolved room key.
|
||
|
||
``room_key`` is used by ``/run/stream`` when it resolves the room from
|
||
thread metadata. ``user_id`` remains the room owner while
|
||
``holder_user_id`` is the account currently renewing a shared room.
|
||
"""
|
||
async with self._lock:
|
||
await self._ensure_gate_row()
|
||
now = datetime.now(UTC)
|
||
async with self._sf() as session:
|
||
# This no-op-ish write is intentional: it is the portable DB
|
||
# mutex protecting the subsequent count + lease upsert.
|
||
await session.execute(
|
||
update(RoundtableRoomLeaseRow)
|
||
.where(RoundtableRoomLeaseRow.room_key == GATE_ROW_KEY)
|
||
.values(updated_at=now)
|
||
)
|
||
keys, durable_scopes, hard_expired_keys = await self._active_snapshot(
|
||
session,
|
||
now,
|
||
max_hold_minutes=max_hold_minutes,
|
||
)
|
||
self._merge_extra_keys(keys, durable_scopes, extra_keys)
|
||
if room_key is None:
|
||
await session.commit()
|
||
blocked = bool(enabled)
|
||
return self._decision(
|
||
ok=not blocked,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=len(keys),
|
||
renewed=False,
|
||
code="DRAFT_REQUIRED" if blocked else "",
|
||
message="" if not blocked else "进入第二步需要先保存会商草稿",
|
||
)
|
||
if room_key in hard_expired_keys:
|
||
await session.commit()
|
||
return self._decision(
|
||
ok=False,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=len(keys),
|
||
renewed=False,
|
||
code="ROOM_MAX_HOLD_EXCEEDED",
|
||
message="本次会商已达到最长占用时间，房间已自动释放。请返回第一步后重新发起。",
|
||
)
|
||
owns = room_key in keys
|
||
if enabled and not owns and len(keys) >= max_rooms:
|
||
await session.commit()
|
||
return self._decision(
|
||
ok=False,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=len(keys),
|
||
renewed=False,
|
||
code="ROUNDTABLE_ROOM_FULL",
|
||
message="当前会商房间暂时已满，请稍候片刻后再试。",
|
||
)
|
||
await self._upsert_lease(
|
||
session,
|
||
room_key=room_key,
|
||
draft_id=draft_id,
|
||
user_id=user_id,
|
||
task_id=task_id,
|
||
holder_user_id=holder_user_id,
|
||
now=now,
|
||
lease_seconds=lease_seconds,
|
||
)
|
||
await session.commit()
|
||
return self._decision(
|
||
ok=True,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=len(keys) if owns else len(keys) + 1,
|
||
renewed=owns,
|
||
code="",
|
||
message="",
|
||
)
|
||
|
||
async def release(
|
||
self,
|
||
*,
|
||
draft_id: str | None,
|
||
user_id: str | None,
|
||
task_id: str | None,
|
||
room_session_id: str | None = None,
|
||
) -> bool:
|
||
"""Drop a leftover foreground lease row if one still exists."""
|
||
key = compute_room_session_key(room_session_id) or compute_active_dedupe_key(
|
||
draft_id, user_id, task_id
|
||
)
|
||
if not key:
|
||
return False
|
||
async with self._lock:
|
||
await self._ensure_gate_row()
|
||
async with self._sf() as session:
|
||
await session.execute(
|
||
update(RoundtableRoomLeaseRow)
|
||
.where(RoundtableRoomLeaseRow.room_key == GATE_ROW_KEY)
|
||
.values(updated_at=datetime.now(UTC))
|
||
)
|
||
row = await session.get(RoundtableRoomLeaseRow, key)
|
||
if row is None or (user_id and row.holder_user_id not in (None, user_id)):
|
||
await session.commit()
|
||
return False
|
||
await session.delete(row)
|
||
await session.commit()
|
||
return True
|
||
|
||
async def started_at(
|
||
self,
|
||
*,
|
||
room_session_id: str | None,
|
||
holder_user_id: str | None,
|
||
) -> datetime | None:
|
||
"""Return the original admission time for a foreground-to-job handoff."""
|
||
key = compute_room_session_key(room_session_id)
|
||
if key is None:
|
||
return None
|
||
async with self._sf() as session:
|
||
row = await session.get(RoundtableRoomLeaseRow, key)
|
||
if row is None or (
|
||
holder_user_id
|
||
and row.holder_user_id not in (None, holder_user_id)
|
||
):
|
||
return None
|
||
return _aware(row.created_at)
|
||
|
||
async def renew(
|
||
self,
|
||
*,
|
||
draft_id: str | None,
|
||
user_id: str | None,
|
||
task_id: str | None,
|
||
room_session_id: str | None,
|
||
enabled: bool,
|
||
max_rooms: int,
|
||
max_hold_minutes: int = 30,
|
||
lease_seconds: int = DEFAULT_LEASE_SECONDS,
|
||
extra_keys: set[str] | Mapping[str, str | None] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""Renew one live foreground lease without ever creating it.
|
||
|
||
Heartbeats are deliberately weaker than admission. A delayed browser
|
||
callback must not resurrect a room that was released after a pause,
|
||
network disconnect or completed Step 2.
|
||
"""
|
||
key = compute_room_session_key(room_session_id)
|
||
if key is None:
|
||
return self._decision(
|
||
ok=not enabled,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=0,
|
||
renewed=False,
|
||
code="ROOM_SESSION_REQUIRED" if enabled else "",
|
||
message="会商会话标识缺失，请刷新页面后重试。" if enabled else "",
|
||
)
|
||
|
||
return await self.renew_room_key(
|
||
room_key=key,
|
||
holder_user_id=user_id,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
max_hold_minutes=max_hold_minutes,
|
||
lease_seconds=lease_seconds,
|
||
extra_keys=extra_keys,
|
||
)
|
||
|
||
async def renew_room_key(
|
||
self,
|
||
*,
|
||
room_key: str,
|
||
holder_user_id: str | None,
|
||
enabled: bool,
|
||
max_rooms: int,
|
||
max_hold_minutes: int = 30,
|
||
lease_seconds: int = DEFAULT_LEASE_SECONDS,
|
||
extra_keys: set[str] | Mapping[str, str | None] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""Renew a resolved persisted key without creating a missing lease."""
|
||
async with self._lock:
|
||
await self._ensure_gate_row()
|
||
now = datetime.now(UTC)
|
||
async with self._sf() as session:
|
||
await session.execute(
|
||
update(RoundtableRoomLeaseRow)
|
||
.where(RoundtableRoomLeaseRow.room_key == GATE_ROW_KEY)
|
||
.values(updated_at=now)
|
||
)
|
||
keys, durable_scopes, hard_expired_keys = await self._active_snapshot(
|
||
session,
|
||
now,
|
||
max_hold_minutes=max_hold_minutes,
|
||
)
|
||
self._merge_extra_keys(keys, durable_scopes, extra_keys)
|
||
row = await session.get(RoundtableRoomLeaseRow, room_key)
|
||
if row is not None and room_key in hard_expired_keys:
|
||
await session.commit()
|
||
return self._decision(
|
||
ok=False,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=len(keys),
|
||
renewed=False,
|
||
code="ROOM_MAX_HOLD_EXCEEDED",
|
||
message="本次会商已达到最长占用时间，房间已自动释放。请返回第一步后重新发起。",
|
||
)
|
||
owns = row is not None and room_key in keys and (
|
||
not holder_user_id or row.holder_user_id in (None, holder_user_id)
|
||
)
|
||
if not owns:
|
||
await session.commit()
|
||
return self._decision(
|
||
ok=not enabled,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=len(keys),
|
||
renewed=False,
|
||
code="ROOM_LEASE_EXPIRED" if enabled else "",
|
||
message="会商连接已断开，请重新恢复研讨。" if enabled else "",
|
||
)
|
||
|
||
row.lease_until = now + timedelta(seconds=normalize_lease_seconds(lease_seconds))
|
||
row.updated_at = now
|
||
await session.commit()
|
||
return self._decision(
|
||
ok=True,
|
||
enabled=enabled,
|
||
max_rooms=max_rooms,
|
||
active_count=len(keys),
|
||
renewed=True,
|
||
code="",
|
||
message="",
|
||
)
|
||
|
||
async def summary(
|
||
self,
|
||
extra_keys: set[str] | Mapping[str, str | None] | None = None,
|
||
*,
|
||
max_hold_minutes: int = 30,
|
||
) -> int:
|
||
"""Return the number of distinct occupied rooms across every user."""
|
||
async with self._lock:
|
||
await self._ensure_gate_row()
|
||
now = datetime.now(UTC)
|
||
async with self._sf() as session:
|
||
keys, durable_scopes, _hard_expired_keys = await self._active_snapshot(
|
||
session,
|
||
now,
|
||
max_hold_minutes=max_hold_minutes,
|
||
)
|
||
self._merge_extra_keys(keys, durable_scopes, extra_keys)
|
||
await session.commit()
|
||
return len(keys)
|
||
|
||
@staticmethod
|
||
def _merge_extra_keys(
|
||
keys: set[str],
|
||
durable_scopes: set[str],
|
||
extra_keys: set[str] | Mapping[str, str | None] | None,
|
||
) -> None:
|
||
"""Add run-manager fallbacks without double-counting durable occupancy.
|
||
|
||
The map value is the legacy draft/user-or-task scope. If a live page
|
||
lease or background job already owns that scope, its stamped run key is
|
||
merely evidence for the same room and must not become another room.
|
||
"""
|
||
if not extra_keys:
|
||
return
|
||
items = extra_keys.items() if isinstance(extra_keys, Mapping) else ((key, None) for key in extra_keys)
|
||
for key, scope_key in items:
|
||
if not key or (scope_key and scope_key in durable_scopes):
|
||
continue
|
||
keys.add(key)
|
||
|
||
async def _active_snapshot(
|
||
self,
|
||
session: AsyncSession,
|
||
now: datetime,
|
||
*,
|
||
max_hold_minutes: int,
|
||
) -> tuple[set[str], set[str], set[str]]:
|
||
lease_rows = (
|
||
await session.execute(
|
||
select(RoundtableRoomLeaseRow).where(
|
||
RoundtableRoomLeaseRow.room_key != GATE_ROW_KEY
|
||
)
|
||
)
|
||
).scalars().all()
|
||
expired: list[str] = []
|
||
keys: set[str] = set()
|
||
durable_scopes: set[str] = set()
|
||
hard_expired_keys: set[str] = set()
|
||
for row in lease_rows:
|
||
until = _aware(row.lease_until)
|
||
if until is None or until <= now:
|
||
expired.append(row.room_key)
|
||
continue
|
||
scope_key = compute_active_dedupe_key(row.draft_id, row.user_id, row.task_id)
|
||
if _hard_deadline_reached(row.created_at, now, max_hold_minutes):
|
||
# Keep the still-live row briefly as a non-occupying tombstone.
|
||
# Its scope suppresses a lagging RunManager fallback until the
|
||
# browser observes the rejected heartbeat and cancels the run.
|
||
hard_expired_keys.add(row.room_key)
|
||
if scope_key:
|
||
durable_scopes.add(scope_key)
|
||
continue
|
||
keys.add(row.room_key)
|
||
if scope_key:
|
||
durable_scopes.add(scope_key)
|
||
if expired:
|
||
await session.execute(
|
||
delete(RoundtableRoomLeaseRow).where(
|
||
RoundtableRoomLeaseRow.room_key.in_(expired)
|
||
)
|
||
)
|
||
|
||
live_after = now - timedelta(seconds=JOB_LIVE_SECONDS)
|
||
job_rows = (
|
||
await session.execute(
|
||
select(
|
||
RoundtableJobRow.id,
|
||
RoundtableJobRow.draft_id,
|
||
RoundtableJobRow.user_id,
|
||
RoundtableJobRow.task_id,
|
||
RoundtableJobRow.status,
|
||
RoundtableJobRow.phase,
|
||
RoundtableJobRow.lease_until,
|
||
RoundtableJobRow.updated_at,
|
||
RoundtableJobRow.created_at,
|
||
).where(RoundtableJobRow.status.in_(OCCUPYING_JOB_STATUSES))
|
||
)
|
||
).all()
|
||
for job_id, draft_id, user_id, task_id, status, phase, lease_until, updated_at, created_at in job_rows:
|
||
if phase in NON_OCCUPYING_JOB_PHASES:
|
||
continue
|
||
if _hard_deadline_reached(created_at, now, max_hold_minutes):
|
||
# The worker cancellation watcher will finalize this state.
|
||
# Marking the intent here also covers queued jobs and workers
|
||
# that stopped renewing but have not yet been reaped.
|
||
await session.execute(
|
||
update(RoundtableJobRow)
|
||
.where(
|
||
RoundtableJobRow.id == job_id,
|
||
RoundtableJobRow.status.in_(["queued", "running", "awaiting_input"]),
|
||
)
|
||
.values(
|
||
status="cancel_requested",
|
||
version=RoundtableJobRow.version + 1,
|
||
updated_at=now,
|
||
)
|
||
)
|
||
continue
|
||
if status != "awaiting_input" and not _job_is_live(
|
||
lease_until,
|
||
updated_at,
|
||
now,
|
||
live_after,
|
||
):
|
||
continue
|
||
job_key = compute_active_dedupe_key(draft_id, user_id, task_id)
|
||
if job_key and job_key in durable_scopes:
|
||
# Foreground → background handoff can overlap for a few
|
||
# milliseconds. Both represent one room for the same scope.
|
||
continue
|
||
if not job_key:
|
||
job_key = f"job:{job_id}"
|
||
keys.add(job_key)
|
||
scope_key = compute_active_dedupe_key(draft_id, user_id, task_id)
|
||
if scope_key:
|
||
durable_scopes.add(scope_key)
|
||
return keys, durable_scopes, hard_expired_keys
|
||
|
||
async def _ensure_gate_row(self) -> None:
|
||
"""Create the portable database mutex row once, tolerating races."""
|
||
now = datetime.now(UTC)
|
||
async with self._sf() as session:
|
||
if await session.get(RoundtableRoomLeaseRow, GATE_ROW_KEY) is not None:
|
||
return
|
||
session.add(
|
||
RoundtableRoomLeaseRow(
|
||
room_key=GATE_ROW_KEY,
|
||
draft_id=None,
|
||
user_id=None,
|
||
task_id=None,
|
||
holder_user_id=None,
|
||
lease_until=None,
|
||
created_at=now,
|
||
updated_at=now,
|
||
)
|
||
)
|
||
try:
|
||
await session.commit()
|
||
except IntegrityError:
|
||
await session.rollback()
|
||
|
||
async def _upsert_lease(
|
||
self,
|
||
session: AsyncSession,
|
||
*,
|
||
room_key: str,
|
||
draft_id: str | None,
|
||
user_id: str | None,
|
||
task_id: str | None,
|
||
holder_user_id: str | None,
|
||
now: datetime,
|
||
lease_seconds: int = DEFAULT_LEASE_SECONDS,
|
||
) -> None:
|
||
row = await session.get(RoundtableRoomLeaseRow, room_key)
|
||
lease_until = now + timedelta(seconds=normalize_lease_seconds(lease_seconds))
|
||
if row is None:
|
||
session.add(
|
||
RoundtableRoomLeaseRow(
|
||
room_key=room_key,
|
||
draft_id=draft_id,
|
||
user_id=user_id,
|
||
task_id=task_id,
|
||
holder_user_id=holder_user_id,
|
||
lease_until=lease_until,
|
||
created_at=now,
|
||
updated_at=now,
|
||
)
|
||
)
|
||
return
|
||
row.draft_id = draft_id
|
||
row.user_id = user_id
|
||
row.task_id = task_id
|
||
row.holder_user_id = holder_user_id
|
||
row.lease_until = lease_until
|
||
row.updated_at = now
|
||
|
||
@staticmethod
|
||
def _decision(
|
||
*,
|
||
ok: bool,
|
||
enabled: bool,
|
||
max_rooms: int,
|
||
active_count: int,
|
||
renewed: bool,
|
||
code: str,
|
||
message: str,
|
||
) -> dict[str, Any]:
|
||
return {
|
||
"ok": ok,
|
||
"enabled": enabled,
|
||
"max_rooms": max_rooms,
|
||
"active_count": active_count,
|
||
"renewed": renewed,
|
||
"code": code,
|
||
"message": message,
|
||
}
|