LQAgentPlatform/views/construction_write/task_cancel_views.py

# -*- coding: utf-8 -*-
"""
大纲生成任务取消接口
提供取消正在执行的大纲生成任务功能
"""

import json
import time
from typing import Optional, Dict, Any
from pydantic import BaseModel, Field
from fastapi import APIRouter, HTTPException
from foundation.observability.logger.loggering import write_logger as logger
from foundation.infrastructure.tracing import TraceContext, auto_trace
from core.base.workflow_manager import WorkflowManager
from core.base.sse_manager import unified_sse_manager

# 创建路由
task_cancel_router = APIRouter(prefix="/sgbx", tags=["施工方案编写"])

# 初始化工作流管理器
workflow_manager = WorkflowManager(
    max_concurrent_docs=3,
    max_concurrent_reviews=5
)


class TaskCancelRequest(BaseModel):
    """任务取消请求模型
    
    示例请求:
    {
        "task_id": "outline_abc123456789",
        "user_id": "user-001",
        "cancel_reason": "用户主动取消"
    }
    """
    task_id: str = Field(..., description="任务ID", example="outline_abc123456789")
    user_id: str = Field(..., description="用户ID", example="user-001")
    cancel_reason: Optional[str] = Field(
        default="用户主动取消",
        description="取消原因"
    )

    class Config:
        extra = "forbid"


class TaskCancelResponse(BaseModel):
    """任务取消响应模型"""
    code: int = Field(..., description="状态码")
    message: str = Field(..., description="状态消息")
    data: Optional[Dict[str, Any]] = Field(None, description="响应数据")


def validate_user_id(user_id: str) -> None:
    """验证用户标识"""
    supported_users = {'user-001', 'user-002', 'user-003'}
    if user_id not in supported_users:
        raise HTTPException(
            status_code=403,
            detail={
                "code": "CANCEL_001",
                "error_type": "INVALID_USER",
                "message": "用户标识未提供或无效"
            }
        )


@task_cancel_router.post("/task_cancel", response_model=TaskCancelResponse)
@auto_trace(generate_if_missing=True)
async def task_cancel(request: TaskCancelRequest):
    """
    取消大纲生成任务
    
    取消正在执行的大纲生成任务，支持取消本服务的生成任务。
    【修复】现在支持取消预注册状态（pending）的任务，即任务提交后、Worker 执行前的时间段。
    
    Args:
        request: 任务取消请求参数
        
    Returns:
        取消结果
        
    Example:
        POST /sgbx/task_cancel
        {
            "task_id": "outline_abc123456789",
            "user_id": "user-001",
            "cancel_reason": "用户主动取消"
        }
    """
    trace_id = f"cancel_{request.task_id}"
    TraceContext.set_trace_id(trace_id)
    
    try:
        logger.info(f"[{trace_id}] 接收任务取消请求: task_id={request.task_id}, user_id={request.user_id}")
        
        # 参数验证
        validate_user_id(request.user_id)
        
        # 检查任务是否存在
        try:
            task_info = await workflow_manager.get_outline_sgbx_task_info(request.task_id)
        except Exception as e:
            logger.warning(f"[{trace_id}] 获取任务信息异常: {e}")
            task_info = None
        
        if not task_info:
            return TaskCancelResponse(
                code=404,
                message="任务不存在或已完成",
                data={
                    "task_id": request.task_id,
                    "status": "not_found"
                }
            )
        
        # 检查任务状态
        task_status = task_info.get("status") or task_info.get("overall_task_status", "unknown")
        
        if task_status == "cancelled":
            return TaskCancelResponse(
                code=200,
                message="任务已处于取消状态",
                data={
                    "task_id": request.task_id,
                    "status": "already_cancelled",
                    "cancelled_at": task_info.get("cancelled_at")
                }
            )
        
        if task_status in ["completed", "failed"]:
            return TaskCancelResponse(
                code=400,
                message=f"任务已{task_status}，无法取消",
                data={
                    "task_id": request.task_id,
                    "current_status": task_status
                }
            )
        
        # 【修复】使用 workflow_manager 的 set_outline_terminate_signal 方法
        # 支持 pending（预注册）和 processing（执行中）两种状态
        cancelled_at = int(time.time())
        
        result = await workflow_manager.set_outline_terminate_signal(
            callback_task_id=request.task_id,
            operator=request.user_id
        )
        
        if not result.get("success"):
            logger.warning(f"[{trace_id}] 设置终止信号失败: {result.get('message')}")
            return TaskCancelResponse(
                code=400,
                message=result.get("message", "取消任务失败"),
                data={
                    "task_id": request.task_id,
                    "current_status": task_info.get("status")
                }
            )
        
        logger.info(f"[{trace_id}] 终止信号已设置: {request.task_id}")
        
        # 【修复】如果是预注册状态（pending），任务已被直接取消，无需其他操作
        is_pre_registered = task_info.get("is_pre_registered", False)
        if is_pre_registered or task_status == "pending":
            logger.info(f"[{trace_id}] 预注册任务已被取消: {request.task_id}")
            
            # 更新进度信息
            try:
                await workflow_manager.progress_manager.update_stage_progress(
                    callback_task_id=request.task_id,
                    overall_task_status="cancelled",
                    status="cancelled",
                    message=f"任务已被用户取消: {request.cancel_reason}"
                )
            except Exception as e:
                logger.warning(f"[{trace_id}] 更新进度信息失败: {e}")
            
            return TaskCancelResponse(
                code=200,
                message="任务已成功取消（未开始执行）",
                data={
                    "task_id": request.task_id,
                    "status": "cancelled",
                    "cancelled_at": cancelled_at,
                    "cancel_reason": request.cancel_reason,
                    "cancelled_by": request.user_id,
                    "is_pre_registered": True
                }
            )
        
        # 对于正在执行的任务（processing），设置额外的取消标志（兼容旧版）
        try:
            import redis.asyncio as redis_async
            from redis.asyncio.connection import ConnectionPool
            
            pool = ConnectionPool(
                host='127.0.0.1',
                port=6379,
                password='123456',
                db=0,
                decode_responses=True,
                max_connections=20,
                socket_connect_timeout=10,
                socket_timeout=10,
                retry_on_timeout=True,
                health_check_interval=30
            )
            
            redis_client = redis_async.Redis(connection_pool=pool)
            
            # 设置终止标志（兼容旧版 outline_workflow 的检查）
            terminate_data = json.dumps({
                "cancelled": True,
                "cancelled_by": request.user_id,
                "cancel_reason": request.cancel_reason,
                "cancelled_at": cancelled_at
            })
            
            await redis_client.set(f"terminate:{request.task_id}", terminate_data, ex=3600)
            
            await redis_client.close()
            await pool.disconnect()
            
        except Exception as e:
            logger.warning(f"[{trace_id}] 设置兼容终止标志失败: {e}")
            # 不影响主流程，继续执行
        
        # 尝试终止 Celery 任务
        celery_task_id = task_info.get("celery_task_id") or task_info.get("celery_id")
        if celery_task_id:
            try:
                from celery import current_app as celery_app
                celery_app.control.revoke(celery_task_id, terminate=True)
                logger.info(f"[{trace_id}] Celery 终止信号已发送: {celery_task_id}")
            except Exception as e:
                logger.warning(f"[{trace_id}] 终止 Celery 任务失败: {e}")
        
        # 发送取消事件到 SSE
        try:
            cancel_event = {
                "callback_task_id": request.task_id,
                "status": "cancelled",
                "overall_task_status": "cancelled",
                "message": f"任务已被用户取消: {request.cancel_reason}",
                "cancelled_at": cancelled_at,
                "cancelled_by": request.user_id
            }
            await unified_sse_manager.send_progress(request.task_id, cancel_event)
            await unified_sse_manager.close_connection(request.task_id)
        except Exception as e:
            logger.warning(f"[{trace_id}] 关闭 SSE 连接失败: {e}")
        
        return TaskCancelResponse(
            code=200,
            message="任务取消成功",
            data={
                "task_id": request.task_id,
                "status": "cancelled",
                "cancelled_at": cancelled_at,
                "cancel_reason": request.cancel_reason,
                "cancelled_by": request.user_id
            }
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"[{trace_id}] 取消任务异常: {str(e)}", exc_info=True)
        return TaskCancelResponse(
            code=500,
            message=f"取消任务失败: {str(e)}",
            data={"task_id": request.task_id}
        )


@task_cancel_router.get("/task_status")
@auto_trace(generate_if_missing=True)
async def get_task_status(
    task_id: str,
    user_id: str
):
    """
    获取任务状态
    
    Args:
        task_id: 任务ID
        user_id: 用户ID
        
    Returns:
        任务状态信息
    """
    try:
        logger.info(f"查询任务状态: task_id={task_id}")
        
        # 验证用户
        validate_user_id(user_id)
        
        # 获取任务信息
        task_info = await workflow_manager.get_outline_sgbx_task_info(task_id)
        
        if not task_info:
            return {
                "code": 404,
                "message": "任务不存在或已完成",
                "data": None
            }
        
        return {
            "code": 200,
            "message": "查询成功",
            "data": {
                "task_id": task_id,
                "status": task_info.get("status") or task_info.get("overall_task_status"),
                "progress": task_info.get("current", 0),
                "message": task_info.get("message", ""),
                "updated_at": task_info.get("updated_at")
            }
        }
        
    except Exception as e:
        logger.error(f"查询任务状态失败: {str(e)}", exc_info=True)
        return {
            "code": 500,
            "message": f"查询失败: {str(e)}",
            "data": None
        }