from __future__ import annotations from collections import Counter, defaultdict from datetime import datetime, timezone from typing import Iterable, Sequence from fastapi import HTTPException from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.enums.generation_task import ( ChatGenerationPipelineStage, ChatGenerationTaskStatus, GenerationMode, ) from app.models.chat_generation_task import ChatGenerationTask from app.services.generation.pipeline.db_lock_service import execute_with_lock_timeout from app.services.operation_log_service import log_operation_event from app.services.video_upscale.guard_service import assert_no_recoverable_failed_upscale_tasks from app.services.resource_accounting_service import ( SOURCE_MODEL_CHAT_TASK, soft_delete_resources_by_source, ) ACTIVE_STAGES = { ChatGenerationPipelineStage.QUEUED.value, ChatGenerationPipelineStage.PREPARING.value, ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value, ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value, ChatGenerationPipelineStage.WAITING_REMOTE.value, ChatGenerationPipelineStage.POLLING.value, ChatGenerationPipelineStage.RESULT_READY.value, ChatGenerationPipelineStage.DOWNLOAD_QUEUED.value, ChatGenerationPipelineStage.DOWNLOADING.value, ChatGenerationPipelineStage.RETRY_WAITING.value, ChatGenerationPipelineStage.UPSCALE_QUEUED.value, ChatGenerationPipelineStage.UPSCALE_PROCESSING.value, ChatGenerationPipelineStage.UPSCALE_POLLING.value, ChatGenerationPipelineStage.UPSCALE_DOWNLOADING.value, ChatGenerationPipelineStage.UPSCALE_FINALIZING.value, ChatGenerationPipelineStage.UPSCALE_RETRY_WAITING.value, } def is_task_active(task: ChatGenerationTask) -> bool: return task.deleted_at is None and ( task.status == ChatGenerationTaskStatus.GENERATING.value or (task.pipeline_stage or "") in ACTIVE_STAGES ) def get_display_status(task: ChatGenerationTask) -> str: if task.deleted_at is not None: return "deleted" if task.pipeline_stage == ChatGenerationPipelineStage.DOWNLOAD_FAILED.value: return "download_failed" if task.pipeline_stage == ChatGenerationPipelineStage.UPSCALE_FAILED.value: return "failed" return task.status or ChatGenerationTaskStatus.PENDING.value async def load_children_map( db: AsyncSession, parent_ids: Sequence[str] | Iterable[str], *, include_deleted: bool = True, ) -> dict[str, list[ChatGenerationTask]]: ids = list(dict.fromkeys(str(item) for item in parent_ids if item)) if not ids: return {} query = select(ChatGenerationTask).where( ChatGenerationTask.parent_task_id.in_(ids), ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value, ) if not include_deleted: query = query.where(ChatGenerationTask.deleted_at.is_(None)) result = await db.execute( query.order_by( ChatGenerationTask.parent_task_id.asc(), ChatGenerationTask.generation_index.asc(), ChatGenerationTask.created_at.asc(), ) ) grouped: dict[str, list[ChatGenerationTask]] = defaultdict(list) for task in result.scalars().all(): if task.parent_task_id: grouped[task.parent_task_id].append(task) return dict(grouped) async def load_task_and_children( db: AsyncSession, *, task_id: str, user_id: str | None = None, include_deleted_children: bool = True, ) -> tuple[ChatGenerationTask | None, list[ChatGenerationTask]]: query = select(ChatGenerationTask).where(ChatGenerationTask.id == task_id) if user_id: query = query.where(ChatGenerationTask.user_id == user_id) result = await db.execute(query.limit(1)) task = result.scalar_one_or_none() if not task: return None, [] if task.generation_mode == GenerationMode.CHATAPI_CHILD.value and task.parent_task_id: parent_result = await db.execute( select(ChatGenerationTask).where(ChatGenerationTask.id == task.parent_task_id).limit(1) ) parent = parent_result.scalar_one_or_none() return parent or task, [task] if task.generation_mode != GenerationMode.CHATAPI_MAIN.value: return task, [] children_map = await load_children_map( db, [task.id], include_deleted=include_deleted_children, ) return task, children_map.get(task.id, []) def _generation_result_status(task: ChatGenerationTask) -> str: """返回任务真实生成结果，不受资源软删除影响。""" if task.pipeline_stage == ChatGenerationPipelineStage.DOWNLOAD_FAILED.value: return "download_failed" if task.status == ChatGenerationTaskStatus.FAILED.value or (task.pipeline_stage or "") in { ChatGenerationPipelineStage.FAILED.value, ChatGenerationPipelineStage.TIMEOUT.value, ChatGenerationPipelineStage.UPSCALE_FAILED.value, }: return "failed" if is_task_active(task): return "generating" if task.status == ChatGenerationTaskStatus.COMPLETED.value: return "completed" return task.status or "pending" def _build_summary(children: list[ChatGenerationTask]) -> str | None: if not children: return None result_counters: Counter[str] = Counter(_generation_result_status(child) for child in children) labels = { "completed": "完成", "failed": "生成失败", "download_failed": "下载失败", "generating": "生成中", "pending": "待处理", } parts = [f"{count}项{labels.get(status, status)}" for status, count in result_counters.items() if count] deleted_count = sum(1 for child in children if child.deleted_at is not None) if deleted_count: parts.append(f"{deleted_count}项资源已删除") return f"{len(children)}项中" + "，".join(parts) def _apply_main_task_status(main: ChatGenerationTask, children: list[ChatGenerationTask]) -> dict[str, object]: previous_status = main.status previous_stage = main.pipeline_stage active_children = [child for child in children if is_task_active(child)] failed_children = [ child for child in children if ( child.status == ChatGenerationTaskStatus.FAILED.value or (child.pipeline_stage or "") in { ChatGenerationPipelineStage.FAILED.value, ChatGenerationPipelineStage.TIMEOUT.value, ChatGenerationPipelineStage.DOWNLOAD_FAILED.value, ChatGenerationPipelineStage.UPSCALE_FAILED.value, } ) ] completed_children = [ child for child in children if child.status == ChatGenerationTaskStatus.COMPLETED.value ] if active_children: main.status = ChatGenerationTaskStatus.GENERATING.value main.pipeline_stage = active_children[0].pipeline_stage or ChatGenerationPipelineStage.QUEUED.value main.generated_at = None main.error_message = _build_summary(children) elif failed_children: main.status = ChatGenerationTaskStatus.FAILED.value if any(child.pipeline_stage == ChatGenerationPipelineStage.DOWNLOAD_FAILED.value for child in failed_children): main.pipeline_stage = ChatGenerationPipelineStage.DOWNLOAD_FAILED.value elif any(child.pipeline_stage == ChatGenerationPipelineStage.UPSCALE_FAILED.value for child in failed_children): main.pipeline_stage = ChatGenerationPipelineStage.UPSCALE_FAILED.value else: main.pipeline_stage = ChatGenerationPipelineStage.FAILED.value main.generated_at = max( (child.generated_at for child in completed_children if child.generated_at), default=datetime.now(timezone.utc), ) main.error_message = _build_summary(children) else: main.status = ChatGenerationTaskStatus.COMPLETED.value main.pipeline_stage = ChatGenerationPipelineStage.DONE.value main.generated_at = max( (child.generated_at for child in children if child.generated_at), default=main.generated_at or datetime.now(timezone.utc), ) main.error_message = None if main.gen_type == "video": main.credits_cost = round(sum(float(child.credits_cost or 0) for child in children), 2) main.text_credits_cost = round(sum(float(child.text_credits_cost or 0) for child in children), 2) main.text_tokens_used = sum(int(child.text_tokens_used or 0) for child in children) main.image_tokens_used = sum(int(child.image_tokens_used or 0) for child in children) main.video_tokens_used = sum(int(child.video_tokens_used or 0) for child in children) main.retry_count = int(main.manual_retry_count or 0) main.poll_count = sum(int(child.poll_count or 0) for child in children) main.poll_error_count = sum(int(child.poll_error_count or 0) for child in children) return { "before_status": previous_status, "before_stage": previous_stage, "after_status": main.status, "after_stage": main.pipeline_stage, "summary": _build_summary(children), } async def aggregate_main_tasks_status_batch( db: AsyncSession, *, parent_task_ids: Sequence[str] | Iterable[str], ) -> dict[str, ChatGenerationTask]: ids = list(dict.fromkeys(str(item) for item in parent_task_ids if item)) if not ids: return {} result = await execute_with_lock_timeout( db, select(ChatGenerationTask) .where( ChatGenerationTask.id.in_(ids), ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value, ) .order_by(ChatGenerationTask.id.asc()) .with_for_update(), ) mains = list(result.scalars().all()) children_map = await load_children_map(db, ids, include_deleted=True) log_snapshots: list[tuple[str, str | None, dict[str, object]]] = [] main_map: dict[str, ChatGenerationTask] = {} for main in mains: main_id = str(main.id) main_map[main_id] = main if main.deleted_at is not None: continue children = children_map.get(main_id, []) if not children: continue detail = _apply_main_task_status(main, children) log_snapshots.append((main_id, str(main.user_id) if main.user_id else None, detail)) await db.flush() for main_id, user_id, detail in log_snapshots: log_operation_event( domain="generation_ai_batch", event_type="MAIN_STATUS_AGGREGATED", event_status="success", source="service", user_id=user_id, group_id=main_id, task_id=main_id, detail=detail, ) return main_map async def aggregate_main_task_status( db: AsyncSession, *, parent_task_id: str, ) -> ChatGenerationTask | None: main_map = await aggregate_main_tasks_status_batch(db, parent_task_ids=[parent_task_id]) return main_map.get(str(parent_task_id)) async def aggregate_parent_for_child(db: AsyncSession, child: ChatGenerationTask | None) -> ChatGenerationTask | None: if not child or child.generation_mode != GenerationMode.CHATAPI_CHILD.value or not child.parent_task_id: return None # 项目关闭了 autoflush，先显式 flush 子任务的终态，确保聚合查询读取到本事务最新状态。 await db.flush() return await aggregate_main_task_status(db, parent_task_id=str(child.parent_task_id)) async def soft_delete_child_tasks_batch( db: AsyncSession, *, child_task_ids: Sequence[str] | Iterable[str], user_id: str, deleted_at: datetime | None = None, require_completed: bool = False, ) -> int: ids = list(dict.fromkeys(str(item) for item in child_task_ids if item)) if not ids: return 0 deleted_at = deleted_at or datetime.now(timezone.utc) result = await execute_with_lock_timeout( db, select(ChatGenerationTask) .where( ChatGenerationTask.id.in_(ids), ChatGenerationTask.user_id == user_id, ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value, ) .with_for_update() ) children = list(result.scalars().all()) found_ids = {str(child.id) for child in children} missing_ids = [item for item in ids if item not in found_ids] if missing_ids: raise HTTPException(status_code=404, detail=f"子任务不存在: {','.join(missing_ids)}") active_children = [child for child in children if child.deleted_at is None] running_ids = [child.id for child in active_children if is_task_active(child)] if running_ids: raise HTTPException(status_code=400, detail=f"仍有 {len(running_ids)} 个子任务生成中，暂不能删除") await assert_no_recoverable_failed_upscale_tasks(db, [str(child.id) for child in active_children]) if require_completed: invalid_ids = [ child.id for child in active_children if child.status != ChatGenerationTaskStatus.COMPLETED.value or child.generated_at is None ] if invalid_ids: raise HTTPException(status_code=409, detail=f"只有生成完成的资源才能从素材云删除: {','.join(invalid_ids)}") source_ids = [str(child.id) for child in active_children] freed_size = await soft_delete_resources_by_source( db, source_model=SOURCE_MODEL_CHAT_TASK, source_ids=source_ids, deleted_at=deleted_at, ) parent_ids = list(dict.fromkeys(str(child.parent_task_id) for child in active_children if child.parent_task_id)) for child in active_children: child.deleted_at = deleted_at await db.flush() for parent_id in parent_ids: await aggregate_main_task_status(db, parent_task_id=parent_id) return int(freed_size or 0) async def soft_delete_child_task( db: AsyncSession, *, child_task_id: str, user_id: str, deleted_at: datetime | None = None, ) -> int: deleted_at = deleted_at or datetime.now(timezone.utc) detail_result = await db.execute( select( ChatGenerationTask.parent_task_id, ChatGenerationTask.generation_index, ).where( ChatGenerationTask.id == child_task_id, ChatGenerationTask.user_id == user_id, ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value, ).limit(1) ) detail = detail_result.one_or_none() if not detail: raise HTTPException(status_code=404, detail="子任务不存在") log_operation_event( domain="generation_ai_batch", event_type="CHILD_RESOURCE_DELETE_START", event_status="started", source="service", user_id=user_id, group_id=detail.parent_task_id, task_id=child_task_id, detail={"generation_index": detail.generation_index}, ) freed_size = await soft_delete_child_tasks_batch( db, child_task_ids=[child_task_id], user_id=user_id, deleted_at=deleted_at, ) log_operation_event( domain="generation_ai_batch", event_type="CHILD_RESOURCE_DELETE_SUCCESS", event_status="success", source="service", user_id=user_id, group_id=detail.parent_task_id, task_id=child_task_id, detail={"generation_index": detail.generation_index, "freed_size_bytes": freed_size}, ) return freed_size async def soft_delete_top_level_task_group( db: AsyncSession, *, task_id: str, user_id: str, deleted_at: datetime | None = None, ) -> int: deleted_at = deleted_at or datetime.now(timezone.utc) result = await execute_with_lock_timeout( db, select(ChatGenerationTask) .where( ChatGenerationTask.id == task_id, ChatGenerationTask.user_id == user_id, ChatGenerationTask.generation_mode.in_( [GenerationMode.CHATAPI_ASYNC.value, GenerationMode.CHATAPI_MAIN.value] ), ChatGenerationTask.deleted_at.is_(None), ) .with_for_update() .limit(1) ) task = result.scalar_one_or_none() if not task: raise HTTPException(status_code=404, detail="任务不存在") if task.generation_mode == GenerationMode.CHATAPI_ASYNC.value: if is_task_active(task): raise HTTPException(status_code=400, detail="当前任务正在生成中，暂不能删除") await assert_no_recoverable_failed_upscale_tasks(db, [str(task.id)]) freed_size = await soft_delete_resources_by_source( db, source_model=SOURCE_MODEL_CHAT_TASK, source_ids=[task.id], deleted_at=deleted_at, ) task.deleted_at = deleted_at await db.flush() return int(freed_size or 0) children_map = await load_children_map(db, [task.id], include_deleted=True) children = children_map.get(task.id, []) active_ids = [child.id for child in children if is_task_active(child)] if active_ids: raise HTTPException(status_code=400, detail=f"任务组仍有 {len(active_ids)} 个子任务生成中，暂不能删除") active_children = [child for child in children if child.deleted_at is None] child_ids = [child.id for child in active_children] await assert_no_recoverable_failed_upscale_tasks(db, [str(child_id) for child_id in child_ids]) freed_size = await soft_delete_resources_by_source( db, source_model=SOURCE_MODEL_CHAT_TASK, source_ids=child_ids, deleted_at=deleted_at, ) for child in active_children: child.deleted_at = deleted_at task.deleted_at = deleted_at await db.flush() log_operation_event( domain="generation_ai_batch", event_type="BATCH_GROUP_DELETE_SUCCESS", event_status="success", source="service", user_id=user_id, group_id=task.id, task_id=task.id, detail={ "child_task_ids": child_ids, "freed_size_bytes": int(freed_size or 0), "physical_files_deleted": False, }, ) return int(freed_size or 0)