614 lines
24 KiB
Python
614 lines
24 KiB
Python
"""全能创作 · 会话与消息的写入服务(契约 §1/§2)。
|
|
|
|
只放「怎么把一条消息安全落库」这类底座能力;agent 循环、工具执行、SSE 在
|
|
后续的 creation_agent.py 里,别混进来。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from datetime import timedelta
|
|
import uuid
|
|
|
|
from django.core.cache import cache
|
|
from django.db import transaction
|
|
from django.db.models import Max
|
|
from django.utils import timezone
|
|
|
|
from .models import CreationConversation, CreationMessage
|
|
|
|
|
|
@transaction.atomic
|
|
def append_message(
|
|
conversation: CreationConversation,
|
|
*,
|
|
role: str,
|
|
kind: str = CreationMessage.Kind.TEXT,
|
|
text: str = "",
|
|
payload: dict | None = None,
|
|
refs: list | None = None,
|
|
task=None,
|
|
) -> CreationMessage:
|
|
"""往会话尾部追加一条消息,并刷新 last_active_at。
|
|
|
|
seq 在事务里 select_for_update 锁住会话行再取 max+1 —— SSE 流式期间可能有
|
|
并发写(用户抢答 / 轮询回填),不锁会撞 uniq_creation_message_seq。
|
|
"""
|
|
locked = CreationConversation.objects.select_for_update().get(pk=conversation.pk)
|
|
next_seq = (locked.messages.aggregate(m=Max("seq"))["m"] or 0) + 1
|
|
message = CreationMessage.objects.create(
|
|
conversation=locked,
|
|
role=role,
|
|
kind=kind,
|
|
text=text,
|
|
payload=payload or {},
|
|
refs=refs or [],
|
|
task=task,
|
|
seq=next_seq,
|
|
)
|
|
locked.last_active_at = timezone.now()
|
|
locked.save(update_fields=["last_active_at", "updated_at"])
|
|
conversation.last_active_at = locked.last_active_at
|
|
return message
|
|
|
|
|
|
def _assets_from_task(task) -> list[dict]:
|
|
"""任务落库的资产 → 结果卡要用的 {id,url,cover,type}。
|
|
|
|
URL 必须走长期直链:预签名链 1 小时过期,写进消息 payload 第二天就打不开。
|
|
"""
|
|
from apps.ai.services import asset_stable_url
|
|
from apps.assets.models import Asset
|
|
|
|
items = []
|
|
assets = Asset.objects.filter(
|
|
origin_task=task, is_deleted=False, purged_at__isnull=True,
|
|
).prefetch_related("files")
|
|
for asset in assets:
|
|
url, cover = asset_stable_url(asset)
|
|
if not url:
|
|
continue
|
|
items.append({
|
|
"id": str(asset.id),
|
|
"url": url,
|
|
"cover": cover or url,
|
|
"type": "video" if asset.asset_type == Asset.Type.VIDEO else "image",
|
|
})
|
|
return items
|
|
|
|
|
|
def _meta_from_task(task, message: CreationMessage) -> dict:
|
|
payload = message.payload or {}
|
|
req = task.request_payload or {}
|
|
model = ""
|
|
if task.model_config_id:
|
|
model = task.model_config.display_name or task.model_config.name
|
|
return {
|
|
"model": model or payload.get("model") or req.get("model") or "",
|
|
"ratio": req.get("ratio") or payload.get("ratio") or "",
|
|
"prompt": payload.get("prompt") or req.get("prompt") or "",
|
|
}
|
|
|
|
|
|
def _message_task(message: CreationMessage):
|
|
"""GENERATING 消息挂的任务:优先 FK,payload.task_id 兜底(旧数据 / 序列化往返)。"""
|
|
if message.task_id:
|
|
return message.task
|
|
task_id = (message.payload or {}).get("task_id")
|
|
if not task_id:
|
|
return None
|
|
from .models import AITask
|
|
|
|
return AITask.objects.select_related("model_config").filter(id=task_id).first()
|
|
|
|
|
|
@transaction.atomic
|
|
def fail_generating_message(message: CreationMessage, error: str) -> CreationMessage:
|
|
"""生成失败:GENERATING 原地改成 ERROR,不另开一条,避免中间态刷屏。"""
|
|
message.kind = CreationMessage.Kind.ERROR
|
|
message.text = (error or "生成失败")[:500]
|
|
message.save(update_fields=["kind", "text", "updated_at"])
|
|
conversation = message.conversation
|
|
conversation.status = CreationConversation.Status.FAILED
|
|
conversation.last_active_at = timezone.now()
|
|
conversation.save(update_fields=["status", "last_active_at", "updated_at"])
|
|
return message
|
|
|
|
|
|
def sync_generating_message(message: CreationMessage) -> bool:
|
|
"""看挂着的 AITask 是否已经终态,是就把 GENERATING 改成 RESULT / ERROR。
|
|
|
|
出图/出片在 worker 里跑,agent 只提交。前端轮询 GET 会话时靠这个回填;
|
|
worker 结束时也会调一次,不用干等到下一次轮询。
|
|
返回是否改了这条消息。
|
|
"""
|
|
from .models import AITask
|
|
|
|
if message.kind != CreationMessage.Kind.GENERATING:
|
|
return False
|
|
payload = message.payload or {}
|
|
if payload.get("kind") == "video_segments":
|
|
return _sync_segmented_video_message(message)
|
|
task = _message_task(message)
|
|
if task is None:
|
|
return False
|
|
if payload.get("kind") == "video_merge" and task.status in (AITask.Status.SUBMITTED, AITask.Status.POLLING):
|
|
# Celery 不在线的本地环境同样可完成用户已确认的合并；此前从不在普通出片完成时调用。
|
|
run_segmented_video_merge(str(task.id))
|
|
task.refresh_from_db()
|
|
if task.status == AITask.Status.SUCCEEDED:
|
|
assets = _assets_from_task(task)
|
|
if not assets:
|
|
return False # 状态已成功但资产还没落(极端竞态),下轮再试
|
|
finish_generating_message(message, assets=assets, meta=_meta_from_task(task, message))
|
|
return True
|
|
if task.status in (AITask.Status.FAILED, AITask.Status.CANCELLED):
|
|
fail_generating_message(message, task.error_message or "生成失败")
|
|
return True
|
|
return False
|
|
|
|
|
|
def _sync_segmented_video_message(message: CreationMessage) -> bool:
|
|
"""同步全能创作的多段出片。
|
|
|
|
每段仍是普通 FREE_VIDEO 任务，故可沿用既有计费、轮询和资产落库；这里仅把它们聚合成
|
|
一张结果卡。任一片段先完成就先写回 GENERATING 卡供用户预览；所有分段完成后才转成
|
|
可合并的结果卡，绝不在此处触发 ffmpeg。
|
|
"""
|
|
from .free_video import finalize_free_video
|
|
from .models import AITask
|
|
|
|
payload = message.payload or {}
|
|
ids = [str(value) for value in payload.get("task_ids") or [] if value]
|
|
if not ids:
|
|
return False
|
|
by_id = {
|
|
str(task.id): task
|
|
for task in AITask.objects.filter(id__in=ids).select_related("model_config")
|
|
}
|
|
tasks = [by_id.get(task_id) for task_id in ids]
|
|
if any(task is None for task in tasks):
|
|
fail_generating_message(message, "分段视频任务不完整，请重新生成。")
|
|
return True
|
|
|
|
# 本地没有 worker 时，用户的会话轮询本身即可推进每个片段；已在 worker 中的任务则幂等返回。
|
|
refreshed = []
|
|
for task in tasks:
|
|
assert task is not None
|
|
if task.status in (AITask.Status.SUBMITTED, AITask.Status.POLLING):
|
|
try:
|
|
task = finalize_free_video(task=task)
|
|
except Exception: # noqa: BLE001 - 单段网络抖动不能使整组直接失败
|
|
task.refresh_from_db()
|
|
else:
|
|
task.refresh_from_db()
|
|
refreshed.append(task)
|
|
|
|
failed = next(
|
|
(task for task in refreshed if task.status in (AITask.Status.FAILED, AITask.Status.CANCELLED)),
|
|
None,
|
|
)
|
|
if failed is not None:
|
|
number = next((index + 1 for index, task in enumerate(refreshed) if task.id == failed.id), 1)
|
|
fail_generating_message(message, f"第 {number} 段生成失败：{failed.error_message or '请重试'}")
|
|
return True
|
|
assets: list[dict] = []
|
|
completed_segments = 0
|
|
for index, task in enumerate(refreshed, start=1):
|
|
if task.status != AITask.Status.SUCCEEDED:
|
|
continue
|
|
task_assets = _assets_from_task(task)
|
|
# 上游状态先成功、资产稍后才落库时，保留 GENERATING，下一轮再展示这段。
|
|
if not task_assets:
|
|
continue
|
|
completed_segments += 1
|
|
for asset in task_assets:
|
|
assets.append({**asset, "label": f"第 {index} 段", "segment_index": index})
|
|
|
|
# 先完成的片段必须立刻回到前端，不能等另一段慢任务一起完成才出现。
|
|
# 保持同一条 GENERATING 消息，避免把一个 60 秒视频拆成多条对话消息。
|
|
if not all(task.status == AITask.Status.SUCCEEDED for task in refreshed):
|
|
progress = {
|
|
**payload,
|
|
"assets": assets,
|
|
"completed_segment_count": completed_segments,
|
|
}
|
|
if progress != payload:
|
|
message.payload = progress
|
|
message.save(update_fields=["payload", "updated_at"])
|
|
return True
|
|
return False
|
|
|
|
# 任务都成功但有片段的资产还在落库，继续保持生成中，避免最终结果缺片。
|
|
if completed_segments != len(refreshed):
|
|
progress = {
|
|
**payload,
|
|
"assets": assets,
|
|
"completed_segment_count": completed_segments,
|
|
}
|
|
if progress != payload:
|
|
message.payload = progress
|
|
message.save(update_fields=["payload", "updated_at"])
|
|
return True
|
|
return False
|
|
|
|
first = refreshed[0]
|
|
finish_generating_message(
|
|
message,
|
|
assets=assets,
|
|
meta={
|
|
**_meta_from_task(first, message),
|
|
"kind": "video_segments",
|
|
"segment_count": len(refreshed),
|
|
"total_duration": payload.get("total_duration") or "",
|
|
"needs_merge": True,
|
|
"segments": payload.get("segments") or [],
|
|
},
|
|
)
|
|
return True
|
|
|
|
|
|
def start_segmented_video_merge(*, conversation: CreationConversation, message: CreationMessage, user):
|
|
"""用户明确点击后才建立合并任务；这之前绝不下载片段或调用 ffmpeg。"""
|
|
from .models import AITask
|
|
|
|
payload = dict(message.payload or {})
|
|
if message.kind != CreationMessage.Kind.RESULT or not payload.get("needs_merge"):
|
|
raise ValueError("这条结果不需要合并")
|
|
state = str(payload.get("merge_state") or "")
|
|
if state in {"queued", "processing", "completed"}:
|
|
raise ValueError("这组片段已经在合并中或已合并")
|
|
task_ids = [str(value) for value in payload.get("task_ids") or [] if value]
|
|
source_tasks = list(AITask.objects.filter(id__in=task_ids).select_related("model_config"))
|
|
if len(source_tasks) != len(task_ids):
|
|
raise ValueError("分段视频不完整，无法合并")
|
|
model_config = source_tasks[0].model_config
|
|
merge_task = AITask.objects.create(
|
|
team=conversation.team,
|
|
created_by=user,
|
|
project=None,
|
|
task_type=AITask.Type.EXPORT,
|
|
status=AITask.Status.SUBMITTED,
|
|
model_config=model_config,
|
|
idempotency_key=f"omni-video-merge:{conversation.id}:{message.id}:{uuid.uuid4()}",
|
|
request_payload={
|
|
"feature": "omni_video_merge",
|
|
"source_message_id": str(message.id),
|
|
"source_task_ids": task_ids,
|
|
"prompt": payload.get("prompt") or "",
|
|
},
|
|
)
|
|
payload["merge_state"] = "queued"
|
|
payload["merge_task_id"] = str(merge_task.id)
|
|
message.payload = payload
|
|
message.save(update_fields=["payload", "updated_at"])
|
|
generating = append_message(
|
|
conversation,
|
|
role="assistant",
|
|
kind=CreationMessage.Kind.GENERATING,
|
|
payload={"task_id": str(merge_task.id), "kind": "video_merge", "prompt": payload.get("prompt") or ""},
|
|
task=merge_task,
|
|
)
|
|
return merge_task, generating
|
|
|
|
|
|
def run_segmented_video_merge(task_id: str) -> None:
|
|
"""后台执行用户已确认的合并。任务创建不等于执行；仅此函数调用 ffmpeg。"""
|
|
from .free_video import _store_free_video_media
|
|
from .models import AITask
|
|
from .services import asset_stable_url
|
|
from .video_replace import _concat_shot_media
|
|
from apps.assets.models import Asset
|
|
|
|
task = AITask.objects.select_related("team", "created_by").filter(id=task_id).first()
|
|
if task is None or task.task_type != AITask.Type.EXPORT:
|
|
return
|
|
source_message_id = str((task.request_payload or {}).get("source_message_id") or "")
|
|
source_message = CreationMessage.objects.filter(id=source_message_id).first()
|
|
if source_message is None:
|
|
return
|
|
try:
|
|
with transaction.atomic():
|
|
locked = AITask.objects.select_for_update().get(id=task.id)
|
|
if locked.status not in (AITask.Status.SUBMITTED, AITask.Status.POLLING):
|
|
return
|
|
locked.status = AITask.Status.POSTPROCESSING
|
|
locked.save(update_fields=["status", "updated_at"])
|
|
urls: list[str] = []
|
|
for source_id in (locked.request_payload or {}).get("source_task_ids") or []:
|
|
source_asset = (
|
|
Asset.objects.filter(origin_task_id=source_id, is_deleted=False, purged_at__isnull=True)
|
|
.prefetch_related("files")
|
|
.first()
|
|
)
|
|
url, _cover = asset_stable_url(source_asset)
|
|
if not url:
|
|
raise ValueError("有片段尚未准备好，暂时不能合并")
|
|
urls.append(url)
|
|
merged_bytes = _concat_shot_media(urls)
|
|
_store_free_video_media(task=locked, video_bytes=merged_bytes)
|
|
with transaction.atomic():
|
|
locked = AITask.objects.select_for_update().get(id=task.id)
|
|
if locked.status != AITask.Status.POSTPROCESSING:
|
|
return
|
|
locked.status = AITask.Status.SUCCEEDED
|
|
locked.completed_at = timezone.now()
|
|
locked.save(update_fields=["status", "completed_at", "updated_at"])
|
|
source = CreationMessage.objects.select_for_update().get(id=source_message.id)
|
|
source_payload = dict(source.payload or {})
|
|
source_payload["merge_state"] = "completed"
|
|
source.payload = source_payload
|
|
source.save(update_fields=["payload", "updated_at"])
|
|
except Exception as exc: # noqa: BLE001 - 合并失败不影响已完成的分段视频
|
|
import logging
|
|
|
|
logging.getLogger(__name__).exception("omni video merge failed for %s", task_id)
|
|
with transaction.atomic():
|
|
locked = AITask.objects.select_for_update().get(id=task.id)
|
|
if locked.status == AITask.Status.POSTPROCESSING:
|
|
locked.status = AITask.Status.FAILED
|
|
locked.error_code = "MergeError"
|
|
locked.error_message = str(exc)[:500]
|
|
locked.completed_at = timezone.now()
|
|
locked.save(update_fields=["status", "error_code", "error_message", "completed_at", "updated_at"])
|
|
source = CreationMessage.objects.select_for_update().filter(id=source_message.id).first()
|
|
if source is not None:
|
|
source_payload = dict(source.payload or {})
|
|
source_payload["merge_state"] = "failed"
|
|
source.payload = source_payload
|
|
source.save(update_fields=["payload", "updated_at"])
|
|
finally:
|
|
# 复用既有「生成中 → 结果/错误」转换，生成的合成片会成为独立结果卡。
|
|
fresh = AITask.objects.filter(id=task_id).first()
|
|
if fresh is not None:
|
|
sync_generating_for_task(fresh)
|
|
|
|
|
|
def sync_generating_messages(conversation: CreationConversation) -> int:
|
|
"""把一条会话里所有已结束的 GENERATING 回填。返回改了几条。"""
|
|
pending = list(
|
|
conversation.messages.filter(kind=CreationMessage.Kind.GENERATING)
|
|
.select_related("task", "task__model_config")
|
|
)
|
|
return sum(1 for message in pending if sync_generating_message(message))
|
|
|
|
|
|
def sync_generating_for_task(task) -> int:
|
|
"""worker / poll 终态后:只扫挂在这个任务上的 GENERATING。失败不能向外抛。"""
|
|
try:
|
|
pending = list(
|
|
CreationMessage.objects.filter(
|
|
task=task, kind=CreationMessage.Kind.GENERATING,
|
|
).select_related("conversation", "task", "task__model_config")
|
|
)
|
|
return sum(1 for message in pending if sync_generating_message(message))
|
|
except Exception: # noqa: BLE001 — 回填失败不能把已经成功的出图任务打成失败
|
|
import logging
|
|
|
|
logging.getLogger(__name__).warning(
|
|
"omni create: sync generating for task %s failed", getattr(task, "id", "?"),
|
|
exc_info=True,
|
|
)
|
|
return 0
|
|
|
|
|
|
@transaction.atomic
|
|
def finish_generating_message(message: CreationMessage, *, assets: list[dict], meta: dict) -> CreationMessage:
|
|
"""把「生成中」原地改成「结果」(契约 §2:不新增消息,避免中间态刷屏)。
|
|
|
|
重生成是**新开一条** GENERATING → RESULT,所以对话流仍然是往下叠加;
|
|
这里改的只是同一次生成自己的中间态。
|
|
"""
|
|
message.kind = CreationMessage.Kind.RESULT
|
|
message.payload = {**(message.payload or {}), **meta, "assets": assets}
|
|
message.save(update_fields=["kind", "payload", "updated_at"])
|
|
conversation = message.conversation
|
|
conversation.status = CreationConversation.Status.COMPLETED
|
|
conversation.last_active_at = timezone.now()
|
|
conversation.save(update_fields=["status", "last_active_at", "updated_at"])
|
|
return message
|
|
|
|
|
|
def pin_refs(conversation: CreationConversation, refs: list[dict]) -> list[dict]:
|
|
"""把本轮引用的实体并进「实体锁定」。按 (type, id) 去重,保留首次出现的顺序 ——
|
|
每轮无条件带上它们,这是多次生成之间锁脸/锁商品的唯一手段(契约 §5)。
|
|
"""
|
|
merged = list(conversation.pinned_refs or [])
|
|
seen = {(item.get("type"), str(item.get("id"))) for item in merged}
|
|
changed = False
|
|
for ref in refs or []:
|
|
ref_type, ref_id = ref.get("type"), ref.get("id")
|
|
if not ref_type or not ref_id:
|
|
continue # 缺 type/id 的 ref 解析不出实体,直接丢
|
|
key = (ref_type, str(ref_id))
|
|
if key in seen:
|
|
continue
|
|
merged.append(ref)
|
|
seen.add(key)
|
|
changed = True
|
|
if changed:
|
|
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
|