优化全能创作

2026-09-11 19:11:27 +08:00
parent dbcac8a44b
commit 02e5dbfb60
25 changed files with 2398 additions and 462 deletions
@@ -5,6 +5,9 @@
 """
 from __future__ import annotations

+from datetime import timedelta
+
+from django.core.cache import cache
 from django.db import transaction
 from django.db.models import Max
 from django.utils import timezone
@@ -201,3 +204,186 @@ def pin_refs(conversation: CreationConversation, refs: list[dict]) -> list[dict]
        conversation.pinned_refs = merged
        conversation.save(update_fields=["pinned_refs", "updated_at"])
    return merged
+
+
+# ---------------------------------------------------------------- Agent 在途锁(一账号同时只整理一份方案)
+
+AGENT_LOCK_TTL_SECONDS = 20 * 60
+AGENT_STALE_MINUTES = 20
+AGENT_BUSY_DETAIL = "当前已有对话正在整理方案，请等待完成后再试"
+
+
+def agent_lock_key(team_id) -> str:
+    return f"omni:agent:{team_id}"
+
+
+def release_agent_lock(team_id, conversation_id=None) -> None:
+    """释放团队级 agent 锁。传 conversation_id 时只在锁仍指向该会话时删,避免误伤。"""
+    key = agent_lock_key(team_id)
+    if conversation_id is None:
+        cache.delete(key)
+        return
+    if str(cache.get(key) or "") == str(conversation_id):
+        cache.delete(key)
+
+
+def acquire_agent_lock(team_id, conversation_id) -> bool:
+    """cache.add 原子占坑。成功返回 True;已有占用返回 False。"""
+    return bool(cache.add(agent_lock_key(team_id), str(conversation_id), timeout=AGENT_LOCK_TTL_SECONDS))
+
+
+def cleanup_stale_agent_planning(team) -> int:
+    """把超时仍卡在 planning 的会话清回 idle,并尝试释放对应 Redis 锁。"""
+    from django.db.models import Q
+
+    cutoff = timezone.now() - timedelta(minutes=AGENT_STALE_MINUTES)
+    stale = list(
+        CreationConversation.objects.filter(
+            team=team,
+            agent_status=CreationConversation.AgentStatus.PLANNING,
+            is_deleted=False,
+        ).filter(
+            # started_at 为空也当过期(异常写入)
+            Q(agent_started_at__lt=cutoff) | Q(agent_started_at__isnull=True)
+        )[:20]
+    )
+    cleared = 0
+    now = timezone.now()
+    for conv in stale:
+        updated = CreationConversation.objects.filter(
+            pk=conv.pk,
+            agent_status=CreationConversation.AgentStatus.PLANNING,
+        ).update(
+            agent_status=CreationConversation.AgentStatus.IDLE,
+            agent_started_at=None,
+            updated_at=now,
+        )
+        if updated:
+            release_agent_lock(team.id, conv.id)
+            cleared += 1
+    return cleared
+
+
+def find_planning_conversation(team, *, exclude_id=None):
+    """返回团队里仍在 planning 的会话(已先清过期)。"""
+    cleanup_stale_agent_planning(team)
+    qs = CreationConversation.objects.filter(
+        team=team,
+        agent_status=CreationConversation.AgentStatus.PLANNING,
+        is_deleted=False,
+    )
+    if exclude_id is not None:
+        qs = qs.exclude(pk=exclude_id)
+    return qs.order_by("-agent_started_at", "-last_active_at").first()
+
+
+def team_agent_busy(team, *, exclude_id=None) -> bool:
+    """DB planning 或 Redis 锁任一占用 → 忙。exclude_id 用于同一会话续跑自检。"""
+    if find_planning_conversation(team, exclude_id=exclude_id) is not None:
+        return True
+    holder = cache.get(agent_lock_key(team.id))
+    if not holder:
+        return False
+    if exclude_id is not None and str(holder) == str(exclude_id):
+        return False
+    # 锁还在但持有会话已不在 planning → 孤儿锁,清掉放行
+    still = CreationConversation.objects.filter(
+        id=holder,
+        team=team,
+        agent_status=CreationConversation.AgentStatus.PLANNING,
+        is_deleted=False,
+    ).exists()
+    if not still:
+        release_agent_lock(team.id, holder)
+        return False
+    return True
+
+
+def begin_agent_planning(conversation: CreationConversation) -> bool:
+    """占 Redis 锁并把会话标成 planning。失败(冲突)返回 False,调用方回 409。"""
+    team_id = conversation.team_id
+    if team_agent_busy(conversation.team, exclude_id=conversation.id):
+        return False
+    if not acquire_agent_lock(team_id, conversation.id):
+        # 锁被别人占了;若锁指向自己(重入),允许续写状态
+        holder = cache.get(agent_lock_key(team_id))
+        if str(holder or "") != str(conversation.id):
+            return False
+    now = timezone.now()
+    CreationConversation.objects.filter(pk=conversation.pk).update(
+        agent_status=CreationConversation.AgentStatus.PLANNING,
+        agent_started_at=now,
+        last_active_at=now,
+        updated_at=now,
+    )
+    conversation.agent_status = CreationConversation.AgentStatus.PLANNING
+    conversation.agent_started_at = now
+    conversation.last_active_at = now
+    clear_agent_cancel(conversation.id)
+    return True
+
+
+def finish_agent_planning(
+    conversation: CreationConversation,
+    *,
+    awaiting_user: bool = False,
+) -> None:
+    """Celery finally:写回 agent_status 并释放锁。"""
+    status = (
+        CreationConversation.AgentStatus.AWAITING_USER
+        if awaiting_user
+        else CreationConversation.AgentStatus.IDLE
+    )
+    now = timezone.now()
+    CreationConversation.objects.filter(pk=conversation.pk).update(
+        agent_status=status,
+        agent_started_at=None,
+        last_active_at=now,
+        updated_at=now,
+    )
+    conversation.agent_status = status
+    conversation.agent_started_at = None
+    release_agent_lock(conversation.team_id, conversation.id)
+    # 正常结束也清掉取消标记,避免下一轮误判
+    clear_agent_cancel(conversation.id)
+
+
+def agent_cancel_key(conversation_id) -> str:
+    return f"omni:agent:cancel:{conversation_id}"
+
+
+def clear_agent_cancel(conversation_id) -> None:
+    cache.delete(agent_cancel_key(conversation_id))
+
+
+def is_agent_cancel_requested(conversation_id) -> bool:
+    """Celery tool round 之间读:用户点了终止则 True。"""
+    return bool(cache.get(agent_cancel_key(conversation_id)))
+
+
+def request_agent_cancel(conversation: CreationConversation) -> bool:
+    """用户终止整理方案:打 Redis 取消标,立刻 idle + 释放团队锁。
+
+    仅当当前仍是 planning 时成功。Celery 在下一轮工具间隙看到标后停跑,
+    已落库消息保留。返回 True=已受理, False=当时不是 planning(调用方回 409)。
+    """
+    # 先打标,再收态 —— worker 在长 LLM 调用回来后也能看见
+    cache.set(agent_cancel_key(conversation.id), "1", timeout=AGENT_LOCK_TTL_SECONDS)
+    updated = CreationConversation.objects.filter(
+        pk=conversation.pk,
+        agent_status=CreationConversation.AgentStatus.PLANNING,
+    ).update(
+        agent_status=CreationConversation.AgentStatus.IDLE,
+        agent_started_at=None,
+        last_active_at=timezone.now(),
+        updated_at=timezone.now(),
+    )
+    if not updated:
+        # 已经不是 planning(竞态完成/别人清了);仍留标一会儿无害,清掉避免脏状态
+        clear_agent_cancel(conversation.id)
+        conversation.refresh_from_db(fields=["agent_status", "agent_started_at"])
+        return False
+    conversation.agent_status = CreationConversation.AgentStatus.IDLE
+    conversation.agent_started_at = None
+    release_agent_lock(conversation.team_id, conversation.id)
+    return True