PlantAI-dev/tasks.py

"""
Celery任务定义
"""
import sys
import os
from pathlib import Path
from datetime import datetime

# 添加项目路径
PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT))

from celery_app import app


@app.task(
    bind=True,
    name="tasks.execute_lightweight_extension_operation",
    time_limit=None,
    soft_time_limit=None,
)
def execute_lightweight_extension_operation(self, payload: dict) -> dict:
    """Run one lightweight-extension operation without a Celery deadline."""
    from gfa_core.extension_stage_lightweight.api import execute_queued_operation

    return execute_queued_operation(payload)


def normalize_stage_params(stage: str, params: dict | None) -> tuple[dict, dict]:
    """Clamp per-stage CPU at the worker boundary and return resource metadata."""
    params = dict(params or {})
    try:
        from gfa_core.mas.stage_catalog import clamp_stage_cpu, get_stage_resource_profile

        requested_cpu = params.get("requested_cpu", params.get("cpu"))
        effective_cpu = clamp_stage_cpu(stage, params.get("cpu"))
        profile = get_stage_resource_profile(stage)
        if stage in {"rnaseq_qc", "rnaseq_align"}:
            from gfa_core.mas.resource_governor import recommend_rnaseq_parallelism, rnaseq_project_cpu_limit

            explicit_cpu = requested_cpu not in {None, "", 0, "0"}
            project_limit = rnaseq_project_cpu_limit(requested_cpu if explicit_cpu else None)
            explicit_threads = params.get("threads_per_sample") or params.get("threadsPerSample")
            plan = recommend_rnaseq_parallelism(
                total_cpu=project_limit,
                sample_parallel_jobs=params.get("sample_parallel_jobs") or params.get("sampleParallelJobs"),
                stage=stage,
            )
            effective_cpu = clamp_stage_cpu(stage, project_limit)
            params["sample_parallel_jobs"] = plan["sample_parallel_jobs"]
            if explicit_threads:
                params["threads_per_sample"] = max(1, int(explicit_threads))
            else:
                params.pop("threads_per_sample", None)
                params.pop("threadsPerSample", None)
            params["rnaseq_project_max_cpu"] = plan["project_cpu_limit"]
            params["resource_decision"] = {
                "stage": stage,
                "requested_cpu": requested_cpu,
                "effective_cpu": effective_cpu,
                "project_cpu_limit": plan["project_cpu_limit"],
                "sample_parallel_jobs": plan["sample_parallel_jobs"],
                "threads_per_sample": max(1, int(explicit_threads)) if explicit_threads else None,
                "reason": (
                    f"{stage} preserves the selected total CPU budget; per-sample threads "
                    "are derived after the real sample count is known"
                ),
            }
    except Exception as exc:
        requested_cpu = params.get("requested_cpu", params.get("cpu"))
        try:
            effective_cpu = max(1, min(int(requested_cpu), 4))
        except (TypeError, ValueError):
            effective_cpu = 4
        profile = {
            "queue": "medium",
            "resource_class": "medium",
            "default_cpu": effective_cpu,
            "max_cpu": effective_cpu,
            "profile_error": str(exc),
        }

    params["requested_cpu"] = requested_cpu
    params["cpu"] = effective_cpu
    _sync_member_set_id_with_scope(stage, params)
    return params, profile


def _sync_member_set_id_with_scope(stage: str, params: dict) -> None:
    """Keep a scope-derived member_set_id consistent with the execute-time scope_id.

    The execute layer can hand down a comp-polluted scope file together with a
    plan-time (no-comp) member_set_id, producing a self-contradictory manifest
    that postflight rejects as stage_identity_contract_mismatch.  Re-derive the
    scope token from the actual scope file so the CLI, the manifest and the
    postflight check all agree.  User-defined subsets (subfamily-xxx) are NOT
    scope-derived and are left untouched.
    """
    member_set_id = str(params.get("memberSetId") or params.get("member_set_id") or "").strip()
    scope_file = str(params.get("scopeFile") or params.get("scope_file") or "").strip()
    if not member_set_id or not scope_file:
        return
    if "subfamily-" in member_set_id:
        return
    parts = member_set_id.split(":")
    if len(parts) != 3 or parts[2] not in {"all_members", "domain_refined"}:
        return
    try:
        from gfa_core.gene_family_scope import read_scope, sanitize_token

        scope_path = Path(scope_file)
        if not scope_path.is_absolute():
            scope_path = PROJECT_ROOT / scope_path
        if not scope_path.exists():
            return
        scope_id = str(read_scope(scope_path).get("scope_id") or "").strip()
        if not scope_id:
            return
        expected_token = sanitize_token(scope_id)
        analysis_route = str(params.get("analysisRoute") or params.get("analysis_route") or parts[1] or "").strip()
        if analysis_route == "unscoped":
            analysis_route = "single_species" if stage == "identify" else ""
        expected_route = analysis_route or parts[1]
        expected_selector = "domain_refined" if expected_route == "subfamily_evidence" else parts[2]
        if parts[0] == expected_token and parts[1] == expected_route and parts[2] == expected_selector:
            return
        new_member_set_id = f"{expected_token}:{expected_route}:{expected_selector}"
        print(f"[MEMBER_SET SYNC] scope_id={scope_id} member_set_id {member_set_id} -> {new_member_set_id}")
        if params.get("memberSetId"):
            params["memberSetId"] = new_member_set_id
        if params.get("member_set_id"):
            params["member_set_id"] = new_member_set_id
        if not params.get("memberSetId") and not params.get("member_set_id"):
            params["memberSetId"] = new_member_set_id
    except Exception as exc:
        print(f"[MEMBER_SET SYNC] skipped: {exc}")


def _nonempty_prefix(value) -> str:
    text = str(value or "").strip()
    return text if text and text.lower() not in {"none", "null", "undefined", "nan", "-"} else ""


def _merge_manifest_identity_params(params: dict | None, manifest: dict | None) -> dict:
    """Prefer the identity actually written by the stage for postflight checks."""
    from gfa_core.mas.tools.stage_identity import merge_manifest_identity_params

    return merge_manifest_identity_params(params, manifest)


def _gene_family_from_config(config_path: Path):
    if not config_path.exists():
        return None
    try:
        import pandas as pd
        sheets = pd.read_excel(config_path, sheet_name=None)
        for sheet_name, df in sheets.items():
            if str(sheet_name).lower() in {"analysis", "meta", "metadata"}:
                for _, row in df.fillna("").iterrows():
                    cells = [str(v).strip() for v in row.tolist()]
                    for idx, cell in enumerate(cells[:-1]):
                        if cell.lower() in {"family", "gene_family", "analysis_family", "project_family"}:
                            value = _nonempty_prefix(cells[idx + 1])
                            if value:
                                return value
            if "Gene_Family" in df.columns:
                for value in df["Gene_Family"].tolist():
                    value = _nonempty_prefix(value)
                    if value:
                        return value
    except Exception as exc:
        print(f"[Celery] Failed to infer gene family from {config_path}: {exc}")
    return None


def _rnaseq_prefix_from_final(final_dir: Path) -> str:
    if not final_dir.exists():
        return ""
    preferred = [
        "*.rnaseq_wgcna.manifest.json",
        "*.rnaseq_fgsea.manifest.json",
        "*.rnaseq_candidate_family.manifest.json",
        "*.rnaseq_sample_qc.manifest.json",
        "*.rnaseq_de.manifest.json",
        "*.rnaseq_quant.manifest.json",
        "*.rnaseq_align.manifest.json",
        "*.rnaseq_qc.manifest.json",
        "*.annotation.manifest.json",
        "*.genome_prepare.manifest.json",
    ]
    for pattern in preferred:
        matches = sorted(final_dir.glob(pattern))
        if len(matches) == 1:
            suffix = pattern[1:]
            return matches[0].name[: -len(suffix)]
    return ""


def resolve_effective_prefix(job_id: str, prefix: str = "", stage: str = "") -> str:
    provided = _nonempty_prefix(prefix)
    gene_family_stages = {
        "identify", "phylogeny", "motif", "domain", "structure", "promoter", "tfbs",
        "synteny", "interspecies_synteny", "karyotype", "combo", "ppi",
        "pep", "summary", "reference_classification", "tree_plot",
    }

    def normalize_gene_family_prefix(value: str) -> str:
        value = _nonempty_prefix(value)
        if stage in gene_family_stages and value.startswith("Multi."):
            family = value.split(".", 1)[1].strip()
            if family:
                return f"Target.{family}"
        return value

    final_dir = PROJECT_ROOT / "workspace" / job_id / "final"
    if stage in {"rnaseq_enrich", "rnaseq_go_enrich", "rnaseq_kegg_enrich", "rnaseq_sample_qc", "rnaseq_fgsea", "rnaseq_wgcna", "rnaseq_candidate_family"}:
        rnaseq_prefix = _rnaseq_prefix_from_final(final_dir)
        if rnaseq_prefix and rnaseq_prefix != provided:
            print(f"[Celery] Overriding RNA-seq prefix {provided or '<empty>'} -> {rnaseq_prefix}")
            return rnaseq_prefix
    if provided:
        return normalize_gene_family_prefix(provided)
    input_dir = PROJECT_ROOT / "input" / job_id
    species_config = input_dir / "genome_prepare" / "species_config.json"
    if species_config.exists():
        try:
            import json
            data = json.loads(species_config.read_text(encoding="utf-8"))
            target = data.get("target") if isinstance(data.get("target"), dict) else None
            if target:
                value = _nonempty_prefix(target.get("prefix"))
                if value:
                    return normalize_gene_family_prefix(value)
            value = _nonempty_prefix(data.get("prefix"))
            if value:
                return normalize_gene_family_prefix(value)
        except Exception as exc:
            print(f"[Celery] Failed to read species_config prefix: {exc}")
    for base in [final_dir, input_dir / "final"]:
        if base.exists():
            for mf in sorted(base.glob("Target.*.genome_prepare.manifest.json")):
                name = mf.name
                return name.split(".genome_prepare.manifest", 1)[0]
            for mf in sorted(base.glob("Multi.*.genome_prepare.manifest.json")):
                name = mf.name
                return normalize_gene_family_prefix(name.split(".genome_prepare.manifest", 1)[0])
    family = _gene_family_from_config(input_dir / "config.xlsx")
    if family:
        return f"Target.{family}"
    for pattern in ("Target.*.final_members.pep.fa", "Multi.*.final_members.pep.fa"):
        for file in sorted(final_dir.glob(pattern)):
            return normalize_gene_family_prefix(file.name[:-len(".final_members.pep.fa")])
    return ""


def _split_prefix(value: str) -> tuple[str, str]:
    text = str(value or "").strip()
    if "." in text:
        left, right = text.split(".", 1)
        return left.strip(), right.strip()
    return "", text


def _list_param(value) -> list[str]:
    if isinstance(value, list):
        return [str(item).strip() for item in value if str(item).strip()]
    import re

    return [item.strip() for item in re.split(r"[,;，；\s]+", str(value or "")) if item.strip()]


def _resolved_artifacts_from_decisions(decisions: dict) -> dict:
    resolved = {}
    for key, decision in (decisions or {}).items():
        if not isinstance(decision, dict):
            continue
        selected = decision.get("selected") if isinstance(decision.get("selected"), dict) else None
        if selected and selected.get("path"):
            resolved[key] = {
                "path": selected.get("path"),
                "source": selected.get("source", ""),
                "kind": selected.get("kind", ""),
                "role": selected.get("role", ""),
                "species": selected.get("species", ""),
                "family": selected.get("family", ""),
            }
    return resolved


def attach_semantic_execution_context(job_id: str, stage: str, prefix: str, params: dict | None) -> dict:
    """Attach resolver decisions to params for execution and diagnosis.

    /api/execute already performs a semantic preflight, but Celery workers can
    also be called directly or run after code reloads.  Rebuilding here keeps
    execution, diagnosis, and logs on the same project-context view.
    """
    params = dict(params or {})
    if isinstance(params.get("semanticPreflight"), dict) and isinstance(params.get("resolvedArtifacts"), dict):
        return params
    try:
        from gfa_core.mas.artifact_resolver import ArtifactResolver
        from gfa_core.mas.context_index import build_project_context

        context = build_project_context(job_id=job_id, prefix=prefix, project_root=PROJECT_ROOT)
        resolver = ArtifactResolver(context)
        prefix_species, prefix_family = _split_prefix(prefix)
        family = str(params.get("geneFamily") or params.get("family") or prefix_family or context.family or "").strip()
        target_species = _list_param(params.get("targetSpecies") or params.get("target_species") or [])
        if target_species:
            species_list = target_species
        elif prefix_species and prefix_species.lower() not in {"multi", "target"}:
            species_list = [prefix_species]
        elif len(context.target_species) == 1:
            species_list = [context.target_species[0]]
        else:
            species_list = [""]

        decisions = {}
        for species in species_list:
            for role, decision in resolver.resolve_stage(stage, species=species, family=family).items():
                decisions[f"{species or 'auto'}:{role}"] = decision.to_dict()

        semantic_preflight = {
            "valid": True,
            "stage": stage,
            "prefix": prefix,
            "decisions": decisions,
            "context": {
                "family": context.family,
                "target_species": context.target_species,
                "reference_species": context.reference_species,
                "comparative_species": context.comparative_species,
                "config_path": context.config_path,
                "input_dir": context.input_dir,
                "final_dir": context.final_dir,
            },
        }
        params["semanticPreflight"] = semantic_preflight
        params["resolvedArtifacts"] = _resolved_artifacts_from_decisions(decisions)
    except Exception as exc:
        params["semanticContextError"] = str(exc)
    return params

# 初始化Orchestrator（每个worker进程初始化一次）
_orchestrator = None

def get_orchestrator():
    global _orchestrator
    if _orchestrator is None:
        from gfa_core.mas.config import config as mas_config
        from gfa_core.mas.llm import LLMManager
        from gfa_core.mas.memory import MemoryManager
        from gfa_core.mas.core.langgraph_orchestrator import LangGraphOrchestrator
        
        llm_config = mas_config.load_llm_config()
        llm_config["openai_api_key"] = mas_config.settings.openai_api_key
        llm_config["openai_base_url"] = mas_config.settings.openai_base_url
        llm_manager = LLMManager(config=llm_config)
        _orchestrator = LangGraphOrchestrator(llm_manager=llm_manager)
    return _orchestrator


def validate_single_stage_artifacts(job_id: str, prefix: str, stage: str, result: dict | None = None) -> dict | None:
    """Return a failure response if a single-stage task did not produce valid artifacts."""
    from gfa_core.mas.tools.stage_identity import get_stage_status_for_identity, resolve_stage_identity
    from gfa_core.mas.tools.stage_status import get_stage_status

    actual_stage = stage
    actual_prefix = prefix
    manifest = {}
    if isinstance(result, dict):
        final_output = result.get("final_output", {})
        if isinstance(final_output, dict) and isinstance(final_output.get("manifest"), dict):
            manifest = final_output.get("manifest") or {}
        elif isinstance(result.get("manifest"), dict):
            manifest = result.get("manifest") or {}
    if isinstance(manifest, dict) and manifest:
        actual_stage = _nonempty_prefix(manifest.get("stage")) or actual_stage
        actual_prefix = _nonempty_prefix(manifest.get("prefix")) or actual_prefix

    stage_context = manifest.get("stage_context") or manifest.get("stageContext") or {} if isinstance(manifest, dict) else {}
    identity_params = {
        "prefix": actual_prefix,
        "scopeFile": manifest.get("scope_file") or stage_context.get("scope_file") or "" if isinstance(manifest, dict) else "",
        "scopeId": manifest.get("scope_id") or stage_context.get("scope_id") or "" if isinstance(manifest, dict) else "",
        "scopeMode": manifest.get("scope_mode") or stage_context.get("scope_mode") or "" if isinstance(manifest, dict) else "",
        "analysisRoute": manifest.get("analysis_route") or stage_context.get("analysis_route") or "" if isinstance(manifest, dict) else "",
        "memberSetId": manifest.get("member_set_id") or stage_context.get("member_set_id") or "" if isinstance(manifest, dict) else "",
        "geneFamily": manifest.get("gene_family") or manifest.get("family") or stage_context.get("family") or "" if isinstance(manifest, dict) else "",
        "targetSpecies": manifest.get("target_species") or stage_context.get("target_species") or [] if isinstance(manifest, dict) else [],
        "referenceSpecies": manifest.get("reference_species") or stage_context.get("reference_species") or [] if isinstance(manifest, dict) else [],
        "comparativeSpecies": manifest.get("comparative_species") or stage_context.get("comparative_species") or [] if isinstance(manifest, dict) else [],
        "wgsSampleSetId": manifest.get("wgs_sample_set_id") if isinstance(manifest, dict) else "",
        "wgsReferenceId": manifest.get("wgs_reference_id") if isinstance(manifest, dict) else "",
        "wgsAnalysisId": manifest.get("wgs_analysis_id") if isinstance(manifest, dict) else "",
        "genotypeScopeId": manifest.get("genotype_scope_id") if isinstance(manifest, dict) else "",
        "phenotypeScopeId": manifest.get("phenotype_scope_id") if isinstance(manifest, dict) else "",
        "covariateScopeId": manifest.get("covariate_scope_id") if isinstance(manifest, dict) else "",
        "modelScopeId": manifest.get("model_scope_id") if isinstance(manifest, dict) else "",
        "sourceModelScopeId": manifest.get("source_model_scope_id") if isinstance(manifest, dict) else "",
        "analysisScopeId": manifest.get("analysis_scope_id") if isinstance(manifest, dict) else "",
        "runId": manifest.get("run_id") if isinstance(manifest, dict) else "",
    }
    try:
        identity = resolve_stage_identity(
            stage=actual_stage,
            prefix=actual_prefix,
            params=identity_params,
            job_id=job_id,
        )
        stage_status = get_stage_status_for_identity(
            job_id=job_id,
            identity=identity,
            strict_identity=bool(
                identity.scope_id
                or identity.analysis_route
                or identity.member_set_id
                or identity.wgs_sample_set_id
                or identity.wgs_reference_id
                or identity.wgs_analysis_id
                or identity.genotype_scope_id
                or identity.phenotype_scope_id
                or identity.covariate_scope_id
                or identity.model_scope_id
                or identity.source_model_scope_id
                or identity.analysis_scope_id
            ),
        )
    except Exception as exc:
        scope_id = str(identity_params.get("scopeId") or "")
        analysis_route = str(identity_params.get("analysisRoute") or "")
        member_set_id = str(identity_params.get("memberSetId") or "")
        print(f"[Celery] Identity-aware artifact validation fallback for {stage}: {exc}")
        stage_status = get_stage_status(
            job_id=job_id,
            prefix=actual_prefix,
            stage=actual_stage,
            scope_id=scope_id,
            analysis_route=analysis_route,
            member_set_id=member_set_id,
            strict_identity=bool(scope_id or analysis_route or member_set_id),
        )
    if actual_stage in {"rnaseq_go_enrich", "rnaseq_kegg_enrich"} and stage_status.completed:
        final_dir = PROJECT_ROOT / "workspace" / job_id / "final"
        suffix = ".GO_enrich.txt" if actual_stage == "rnaseq_go_enrich" else ".KEGG_enrich.txt"
        if any(final_dir.glob(f"*{suffix}")):
            return None

    if stage_status.completed:
        return None

    business_messages: list[str] = []
    if isinstance(result, dict):
        for error in result.get("errors", []) or []:
            if error:
                business_messages.append(str(error))
        final_output = result.get("final_output", {})
        if isinstance(final_output, dict):
            for key in ("response", "message", "error"):
                value = final_output.get(key)
                if value:
                    business_messages.append(str(value))

    business_error = next(
        (
            msg for msg in business_messages
            if "❌" in msg or "failed" in msg.lower() or "cannot choose a background" in msg.lower()
        ),
        "",
    )

    detail_parts = []
    if stage_status.errors:
        detail_parts.extend(stage_status.errors)
    if stage_status.missing_artifacts:
        detail_parts.append("缺失产物: " + "; ".join(stage_status.missing_artifacts))
    detail_msg = "；".join(detail_parts) or f"{stage} 未生成有效输出"
    error_msg = business_error or f"{stage} 执行失败：{detail_msg}"
    print(f"[Celery] 单阶段产物校验失败: {stage_status.to_dict()}")

    return {
        'status': 'failed',
        'success': False,
        'current_stage': stage,
        'progress': 0,
        'message': error_msg,
        'errors': [error_msg],
        'stage_status': stage_status.to_dict(),
        'result': (result or {}).get("final_output", {}) if isinstance(result, dict) else {},
        'logs': [
            {'time': datetime.now().isoformat(), 'level': 'error', 'message': error_msg}
        ]
    }


def _diagnose_single_stage_failure(
    job_id: str,
    stage: str,
    prefix: str = "",
    *,
    prepare_patch_workspace: bool = False,
) -> dict:
    """DiagnosisAgent is disabled; keep failures as plain execution failures."""
    return {
        "status": "diagnosis_disabled",
        "stage": stage,
        "job_id": job_id,
        "prefix": prefix,
        "issues": [],
    }
    try:
        from gfa_core.mas.agents.diagnosis import DiagnosisAgent

        diagnosis = DiagnosisAgent(None).diagnose_job_stage(
            job_id=job_id,
            stage=stage,
            prefix=prefix or "",
            project_root=PROJECT_ROOT,
            prepare_patch_workspace=prepare_patch_workspace,
        )
        return diagnosis if isinstance(diagnosis, dict) else {}
    except Exception as exc:
        print(f"[Celery] DiagnosisAgent failed for {job_id}/{stage}: {exc}")
        return {
            "status": "diagnosis_failed",
            "stage": stage,
            "job_id": job_id,
            "prefix": prefix,
            "error": str(exc),
            "policy": {
                "auto_modify_input": False,
                "auto_modify_global_code": False,
                "auto_rerun_stage": False,
            },
        }


def _diagnose_stage_gate(
    job_id: str,
    stage: str,
    prefix: str = "",
    *,
    gate: str,
    params: dict | None = None,
) -> dict:
    return {
        "status": "diagnosis_disabled",
        "stage": stage,
        "job_id": job_id,
        "prefix": prefix,
        "gate": gate,
        "issues": [],
    }
    try:
        from gfa_core.mas.agents.diagnosis import DiagnosisAgent

        diagnosis_prefix = "genome_prepare" if stage in {"genome_prepare", "genome_prepare_v2"} else (prefix or "")
        return DiagnosisAgent(None).diagnose_stage_gate(
            job_id=job_id,
            stage=stage,
            prefix=diagnosis_prefix,
            gate=gate,
            project_root=PROJECT_ROOT,
            params=params or {},
        )
    except Exception as exc:
        print(f"[Celery] Diagnosis gate failed for {job_id}/{stage}/{gate}: {exc}")
        return {
            "status": "diagnosis_failed",
            "stage": stage,
            "job_id": job_id,
            "prefix": prefix,
            "gate": gate,
            "error": str(exc),
            "issues": [],
        }


def _apply_runtime_repairs(
    job_id: str,
    stage: str,
    prefix: str = "",
    diagnosis: dict | None = None,
    *,
    execute_heavy: bool = False,
) -> dict:
    return {
        "success": False,
        "status": "diagnosis_disabled",
        "stage": stage,
        "job_id": job_id,
        "results": [],
    }
    try:
        from gfa_core.mas.agents.diagnosis import DiagnosisAgent

        return DiagnosisAgent(None).apply_runtime_repairs(
            job_id=job_id,
            stage=stage,
            prefix=prefix or "",
            diagnosis=diagnosis,
            project_root=PROJECT_ROOT,
            execute_heavy=execute_heavy,
        )
    except Exception as exc:
        print(f"[Celery] Runtime repair failed for {job_id}/{stage}: {exc}")
        return {
            "success": False,
            "status": "repair_failed",
            "stage": stage,
            "job_id": job_id,
            "error": str(exc),
            "results": [],
        }


def _attach_diagnosis_to_failure(response: dict, job_id: str, stage: str, prefix: str = "") -> dict:
    """Keep failed task responses free of DiagnosisAgent side effects."""
    if not isinstance(response, dict) or response.get("success") is not False:
        return response
    response.pop("pending_confirmation", None)
    response.pop("diagnosis", None)
    return response

    diagnosis = response.get("diagnosis")
    if not isinstance(diagnosis, dict) or not diagnosis:
        diagnosis = _diagnose_single_stage_failure(job_id=job_id, stage=stage, prefix=prefix)
        response["diagnosis"] = diagnosis

    session_id = str(response.get("session_id") or job_id or "").strip()
    if diagnosis.get("requires_user_confirmation"):
        response["pending_confirmation"] = {
            "status": "waiting",
            "type": "diagnosis_review",
            "session_id": session_id,
            "stage": stage,
            "job_id": job_id,
            "prefix": prefix,
            "diagnosis": diagnosis,
            "message": "检测到阶段失败，已生成诊断。请用户确认后再修订输入、执行副本修复或升级全局代码。",
            "allowed_actions": [
                "review_only",
                "prepare_job_patch",
                "apply_confirmed_global_fix_after_backup",
            ],
        }

    return response


@app.task(bind=True, max_retries=3)
def execute_agent_task(self, task_id: str, message: str, session_id: str, job_id: str = "", history: list = None, prefix: str = ""):
    """
    执行Agent分析任务
    """
    try:
        print(f"[Celery] Starting task {task_id}, job_id={job_id}")
        
        # 更新任务状态为运行中（使用self.request.id获取Celery任务ID）
        celery_task_id = self.request.id if self.request else task_id
        self.update_state(
            task_id=celery_task_id,
            state='RUNNING',
            meta={
                'current_stage': 'intent_parsing',
                'progress': 0,
                'message': message,
                'logs': [
                    {'time': datetime.now().isoformat(), 'level': 'info', 'message': '开始分析任务...'},
                    {'time': datetime.now().isoformat(), 'level': 'info', 'message': f'解析用户意图: {message}'}
                ]
            }
        )
        
        # 获取Orchestrator
        orchestrator = get_orchestrator()
        
        # 查询已完成的阶段
        completed_stages = []
        if job_id:
            from gfa_core.mas.tools.stage_status import list_completed_stages

            completed_stages = list_completed_stages(job_id=job_id, prefix=prefix or "")
        print(f"[tasks.py] Completed stages: {completed_stages}")
        
        # 构建初始状态
        from gfa_core.mas.core.state import GraphState
        initial_state = {
            "query": message,
            "session_id": session_id,
            "job_id": job_id,  # 传入 job_id 用于 auto_continue
            "files": [],
            "current_agent": "router",
            "status": "running",
            "errors": [],
            "final_output": {},
            "history": history or [],
            "messages": [],  # 添加 messages 字段
            "completed_stages": completed_stages,  # 添加已完成阶段
            "retry_count": 0,  # 添加重试计数
            "max_retries": 3,  # 添加最大重试次数
            "user_input": {"prefix": prefix}  # 传递 prefix 参数
        }
        
        print(f"[tasks.py] initial_state keys: {list(initial_state.keys())}")
        print(f"[tasks.py] initial_state['completed_stages']: {initial_state.get('completed_stages', 'NOT FOUND')}")
        
        # 运行Orchestrator
        import asyncio
        result = asyncio.run(orchestrator.run(initial_state))
        
        # 处理结果
        final_output = result.get("final_output", {})
        
        if result.get("status") == "completed":
            # ============ 生成累积解读报告（新增） ============
            try:
                import asyncio
                from gfa_core.mas.agents.result_parser import create_result_parser_agent
                from gfa_core.mas.llm import LLMManager

                # 从 result 中提取 prefix（如果传入的 prefix 为空）
                if not prefix:
                    prefix = result.get("final_output", {}).get("user_input", {}).get("prefix", "")
                if not prefix:
                    print("[Celery] 未解析到输出标签，跳过累积解读报告生成")
                    raise RuntimeError("missing_prefix_for_interpretation")

                # 获取已完成的阶段列表（统一 manifest/artifacts 校验）
                from gfa_core.mas.tools.stage_status import list_completed_stages

                completed_stages = list_completed_stages(job_id=job_id, prefix=prefix)

                # 按顺序排序
                stage_order = ["identify", "phylogeny", "structure", "motif", "domain", "promoter", "synteny", "interspecies_synteny", "karyotype", "expression", "kaks", "ppi"]
                completed_stages = [s for s in stage_order if s in completed_stages] + [s for s in completed_stages if s not in stage_order]

                print(f"[Celery] 生成累积解读，已完成阶段: {completed_stages}")

                # 创建 ResultParserAgent（不需要 LLMManager）
                agent = create_result_parser_agent(None)

                # 生成累积解读报告
                report = asyncio.run(agent.generate_cumulative_report(
                    job_id=job_id,
                    prefix=prefix,
                    completed_stages=completed_stages
                ))

                # 保存解读报告
                save_success = agent.save_interpretation(job_id, prefix, report)

                if save_success:
                    print(f"[Celery] 解读报告已保存到 interpretation.json")
                else:
                    print(f"[Celery] 保存解读报告失败")

            except Exception as e:
                # 解读失败不阻塞主流程
                print(f"[Celery] 生成解读报告失败（非阻塞）: {e}")
                import traceback
                traceback.print_exc()

            return {
                'status': 'completed',
                'current_stage': result.get("current_agent", "unknown"),
                'progress': 100,
                'message': message,
                'result': {
                    'response': final_output.get("response", ""),
                    'analysis_plan': final_output.get("analysis_plan", {}),
                    'execution_result': final_output.get("execution_result", {})
                },
                'logs': [
                    {'time': datetime.now().isoformat(), 'level': 'success', 'message': '分析任务完成'}
                ]
            }
        else:
            return {
                'status': 'failed',
                'current_stage': result.get("current_agent", "unknown"),
                'progress': 0,
                'message': message,
                'result': {},
                'logs': [
                    {'time': datetime.now().isoformat(), 'level': 'error', 'message': f"分析失败: {result.get('errors', ['Unknown error'])}"}
                ]
            }
            
    except Exception as e:
        import traceback
        error_msg = f"任务异常: {str(e)}\n{traceback.format_exc()}"
        print(f"[Celery] Task {task_id} failed: {error_msg}")
        celery_task_id = self.request.id if self.request else task_id
        self.update_state(
            task_id=celery_task_id,
            state='FAILURE',
            meta={
                'exc_type': type(e).__name__,
                'exc_message': str(e),
            }
        )
        raise


@app.task(bind=True, max_retries=3)
def execute_single_stage_task(self, task_id: str, stage: str, job_id: str, prefix: str, 
                               params: dict = None, session_id: str = None):
    """
    执行单个分析阶段任务（用于/api/execute统一架构）
    
    参数:
    - task_id: 任务ID
    - stage: 分析阶段（如identify, phylogeny等）
    - job_id: 任务目录ID
    - prefix: 分析前缀（例如 At.WRKY；基因家族流程通常由 config/scope 自动生成）
    - params: 执行参数
    - session_id: 会话ID
    """
    resource_acquired = False
    resource_outcome = "failed"
    resource_lease = None
    try:
        print(f"[Celery] Starting single stage task {task_id}, stage={stage}, job_id={job_id}")
        params, resource_profile = normalize_stage_params(stage, params)
        print(
            "[Celery] Resource profile "
            f"stage={stage}, queue={resource_profile.get('queue')}, "
            f"class={resource_profile.get('resource_class')}, "
            f"requested_cpu={params.get('requested_cpu')}, effective_cpu={params.get('cpu')}"
        )
        
        celery_task_id = self.request.id if self.request else task_id

        try:
            from gfa_core.mas.resource_governor import ResourceLease
            from gfa_core.mas.stage_catalog import get_stage_queue

            resource_queue = get_stage_queue(stage)
            resource_lease = ResourceLease(
                task_id=task_id,
                stage=stage,
                job_id=job_id,
                prefix=prefix,
                cpu=int(params.get("cpu") or 1),
                queue=resource_queue,
                resource_class=str(resource_profile.get("resource_class") or ""),
                concurrency_group=str(resource_profile.get("concurrency_group") or ""),
                celery_task_id=celery_task_id,
            )

            def _update_waiting_resource(wait_result: dict):
                usage = wait_result.get("usage", {}) if isinstance(wait_result, dict) else {}
                reason = wait_result.get("reason", "等待资源可用") if isinstance(wait_result, dict) else "等待资源可用"
                self.update_state(
                    task_id=celery_task_id,
                    state="WAITING_RESOURCE",
                    meta={
                        "status": "waiting_resource",
                        "current_stage": stage,
                        "progress": 0,
                        "message": reason,
                        "resource_usage": usage,
                        "resource_profile": resource_profile,
                        "logs": [
                            {
                                "time": datetime.now().isoformat(),
                                "level": "info",
                                "message": reason,
                            }
                        ],
                    },
                )

            acquire_result = resource_lease.wait(update_state=_update_waiting_resource)
            resource_acquired = True
            print(f"[Celery] Resource lease granted for {task_id}: {acquire_result.get('usage')}")
        except Exception as exc:
            print(f"[Celery] Resource governor unavailable, continuing without lease: {exc}")
        
        # 更新任务状态
        self.update_state(
            task_id=celery_task_id,
            state='RUNNING',
            meta={
                'current_stage': stage,
                'progress': 10,
                'message': (
                    f"开始执行 {stage}（请求CPU {params.get('requested_cpu') or params.get('cpu')}，"
                    f"调度CPU {params.get('cpu')}，{resource_profile.get('resource_class')}队列）..."
                ),
                'logs': [
                    {
                        'time': datetime.now().isoformat(),
                        'level': 'info',
                        'message': (
                            f"开始执行 {stage}; requested_cpu={params.get('requested_cpu')}, "
                            f"effective_cpu={params.get('cpu')}, queue_class={resource_profile.get('resource_class')}"
                        )
                    }
                ]
            }
        )

        # /api/execute 是“用户已明确选择阶段”的入口。这里直接调用
        # StageToolRegistry，避免再走自然语言规划/PreFlight，把
        # rnaseq_enrich 误展开成 rnaseq_qc 并错误要求 FASTQ。
        try:
            import asyncio
            from gfa_core.mas.tools.stage_tools import stage_registry

            prefix = resolve_effective_prefix(job_id, prefix, stage)
            if params is not None:
                params["prefix"] = prefix
            params = attach_semantic_execution_context(job_id, stage, prefix, params)
            print(f"[Celery] Effective prefix for {stage}: {prefix or '<empty>'}")
            if params.get("resolvedArtifacts"):
                print(f"[Celery] Resolved artifacts for {stage}: {params.get('resolvedArtifacts')}")
            qa_preflight = _diagnose_stage_gate(job_id, stage, prefix, gate="preflight", params=params)
            qa_repair = {}
            if qa_preflight.get("runtime_repair_candidates"):
                execute_heavy = bool((params or {}).get("diagnosisHeavyRepair", True))
                qa_repair = _apply_runtime_repairs(
                    job_id,
                    stage,
                    prefix,
                    qa_preflight,
                    execute_heavy=execute_heavy,
                )
                qa_preflight = _diagnose_stage_gate(job_id, stage, prefix, gate="preflight", params=params)
            blocking_preflight = [
                issue for issue in qa_preflight.get("issues", [])
                if issue.get("severity") == "blocking" or issue.get("blocks_downstream")
            ]
            if blocking_preflight:
                error_msg = "❌ 诊断门控阻断执行：发现输入/上游产物存在会导致阶段失败的问题"
                return _attach_diagnosis_to_failure({
                    'status': 'failed',
                    'success': False,
                    'current_stage': stage,
                    'progress': 0,
                    'session_id': session_id or job_id,
                    'message': error_msg,
                    'errors': [error_msg],
                    'diagnosis': qa_preflight,
                    'qa_preflight': qa_preflight,
                    'qa_runtime_repair': qa_repair,
                    'semantic_preflight': params.get("semanticPreflight"),
                    'resolved_artifacts': params.get("resolvedArtifacts", {}),
                    'resource_profile': resource_profile,
                    'logs': [
                        {'time': datetime.now().isoformat(), 'level': 'error', 'message': error_msg}
                    ]
                }, job_id, stage, prefix)
            direct_result = asyncio.run(stage_registry.execute(
                stage_id=stage,
                job_id=job_id,
                prefix=prefix,
                extra_params=params or {},
            ))
            print(f"[Celery] Direct single-stage result: {direct_result}")

            success = direct_result.get("status") == "success"
            if success:
                result_manifest = direct_result.get("manifest") if isinstance(direct_result.get("manifest"), dict) else {}
                result_stage = _nonempty_prefix(result_manifest.get("stage")) or stage
                result_prefix = _nonempty_prefix(result_manifest.get("prefix")) or prefix
                artifact_failure = validate_single_stage_artifacts(job_id, result_prefix, result_stage, {
                    "final_output": {
                        "message": direct_result.get("message", ""),
                        "manifest": result_manifest,
                        "manifest_path": direct_result.get("manifest_path"),
                    }
                })
                if artifact_failure:
                    artifact_failure["session_id"] = session_id or job_id
                    return _attach_diagnosis_to_failure(artifact_failure, job_id, result_stage, result_prefix)
                postflight_params = _merge_manifest_identity_params(params, result_manifest)
                qa_postflight = _diagnose_stage_gate(job_id, result_stage, result_prefix, gate="postflight", params=postflight_params)
                qa_post_repair = {}
                if qa_postflight.get("runtime_repair_candidates"):
                    qa_post_repair = _apply_runtime_repairs(
                        job_id,
                        result_stage,
                        result_prefix,
                        qa_postflight,
                        execute_heavy=bool((params or {}).get("diagnosisHeavyRepair", True)),
                    )
                    qa_postflight = _diagnose_stage_gate(job_id, result_stage, result_prefix, gate="postflight", params=postflight_params)
                blocking_postflight = [
                    issue for issue in qa_postflight.get("issues", [])
                    if issue.get("severity") == "blocking" or issue.get("blocks_downstream")
                ]
                if blocking_postflight:
                    error_msg = "❌ 阶段已运行，但诊断 postflight 发现关键产物不可安全下游使用"
                    return _attach_diagnosis_to_failure({
                        'status': 'failed',
                        'success': False,
                        'current_stage': result_stage,
                        'progress': 100,
                        'session_id': session_id or job_id,
                        'message': error_msg,
                        'errors': [error_msg],
                        'result': direct_result,
                        'diagnosis': qa_postflight,
                        'qa_preflight': qa_preflight,
                        'qa_runtime_repair': qa_repair,
                        'qa_postflight': qa_postflight,
                        'qa_postflight_repair': qa_post_repair,
                        'semantic_preflight': params.get("semanticPreflight"),
                        'resolved_artifacts': params.get("resolvedArtifacts", {}),
                        'resource_profile': resource_profile,
                        'logs': [
                            {'time': datetime.now().isoformat(), 'level': 'error', 'message': error_msg}
                        ]
                    }, job_id, result_stage, result_prefix)
                resource_outcome = "completed"
                return {
                    'status': 'completed',
                    'success': True,
                    'current_stage': result_stage,
                    'progress': 100,
                    'message': f'{result_stage} 执行完成',
                    'result': direct_result,
                    'qa_preflight': qa_preflight,
                    'qa_runtime_repair': qa_repair,
                    'qa_postflight': qa_postflight,
                    'qa_postflight_repair': qa_post_repair,
                    'semantic_preflight': params.get("semanticPreflight"),
                    'resolved_artifacts': params.get("resolvedArtifacts", {}),
                    'resource_profile': resource_profile,
                    'logs': [
                        {'time': datetime.now().isoformat(), 'level': 'success', 'message': direct_result.get("message", f'{stage} 执行成功')}
                    ]
                }

            error_msg = direct_result.get("error") or direct_result.get("message") or f"{stage} 执行失败"
            qa_failure = _diagnose_stage_gate(job_id, stage, prefix, gate="failure", params=params)
            return _attach_diagnosis_to_failure({
                'status': 'failed',
                'success': False,
                'current_stage': stage,
                'progress': 0,
                'session_id': session_id or job_id,
                'message': error_msg,
                'errors': [error_msg],
                'result': direct_result,
                'diagnosis': qa_failure,
                'qa_preflight': qa_preflight,
                'qa_runtime_repair': qa_repair,
                'qa_failure': qa_failure,
                'semantic_preflight': params.get("semanticPreflight"),
                'resolved_artifacts': params.get("resolvedArtifacts", {}),
                'resource_profile': resource_profile,
                'logs': [
                    {'time': datetime.now().isoformat(), 'level': 'error', 'message': error_msg}
                ]
            }, job_id, stage, prefix)
        except Exception as direct_exc:
            print(f"[Celery] Direct single-stage execution failed, falling back to orchestrator: {direct_exc}")
        
        # 获取Orchestrator
        orchestrator = get_orchestrator()
        
        # 构建初始状态（单stage模式）
        from gfa_core.mas.core.langgraph_orchestrator import GraphState
        initial_state: GraphState = {
            "query": f"执行{stage}分析",
            "session_id": session_id or task_id,
            "job_id": job_id,
            "files": [],
            "current_agent": "router",
            "status": "running",
            "errors": [],
            "final_output": {},
            "messages": [],
            "retry_count": 0,
            "max_retries": 3,
            "user_input": {
                "prefix": prefix,
                "stage": stage,  # 指定要执行的stage
                "params": params or {}
            },
            "single_stage_mode": True,  # 标记为单stage模式
            "target_stage": stage  # 目标stage
        }
        
        print(f"[tasks.py] 单stage模式: target={stage}, prefix={prefix}")
        
        # 运行LangGraph
        import asyncio
        result = asyncio.run(orchestrator.run(initial_state))
        
        # 处理结果
        if result.get("status") in ["completed", "success"]:
            # 单阶段任务的最终成功标准必须是真实产物，而不是仅依赖
            # LangGraph/Celery 的运行状态。
            artifact_failure = validate_single_stage_artifacts(job_id, prefix, stage, result)
            if artifact_failure:
                artifact_failure["session_id"] = session_id or job_id
                return _attach_diagnosis_to_failure(artifact_failure, job_id, stage, prefix)

            # ============ 生成累积解读报告（新增） ============
            try:
                import asyncio
                from gfa_core.mas.agents.result_parser import create_result_parser_agent

                # 获取已完成的阶段列表（统一 manifest/artifacts 校验）
                from gfa_core.mas.tools.stage_status import list_completed_stages

                completed_stages = list_completed_stages(job_id=job_id, prefix=prefix)
                print(f"[Celery DEBUG] Unified completed_stages: {completed_stages}")

                # 按顺序排序
                stage_order = ["identify", "phylogeny", "structure", "motif", "domain", "promoter", "synteny", "interspecies_synteny", "karyotype", "expression", "kaks", "ppi"]
                completed_stages = [s for s in stage_order if s in completed_stages] + [s for s in completed_stages if s not in stage_order]

                print(f"[Celery] 生成累积解读，已完成阶段: {completed_stages}")

                # 创建 ResultParserAgent（不需要 LLMManager）
                agent = create_result_parser_agent(None)

                # 生成累积解读报告
                report = asyncio.run(agent.generate_cumulative_report(
                    job_id=job_id,
                    prefix=prefix,
                    completed_stages=completed_stages
                ))

                # 保存解读报告
                save_success = agent.save_interpretation(job_id, prefix, report)

                if save_success:
                    print(f"[Celery] 解读报告已保存到 interpretation.json")
                else:
                    print(f"[Celery] 保存解读报告失败")

            except Exception as e:
                # 解读失败不阻塞主流程
                print(f"[Celery] 生成解读报告失败（非阻塞）: {e}")
                import traceback
                traceback.print_exc()

            resource_outcome = "completed"
            return {
                'status': 'completed',
                'success': True,  # 添加 success 字段，供前端判断
                'current_stage': stage,
                'progress': 100,
                'message': f'{stage} 执行完成',
                'result': result.get("final_output", {}),
                'resource_profile': resource_profile,
                'logs': [
                    {'time': datetime.now().isoformat(), 'level': 'success', 'message': f'{stage} 执行成功'}
                ]
            }
        else:
            # 检查是否有等待确认的状态（优先检查 result.status）
            if result.get("status") == "waiting_confirmation":
                pending_confirmation = result.get("pending_confirmation", {})
                # 将确认状态存储到Celery任务meta中
                return {
                    'status': 'waiting_confirmation',
                    'current_stage': stage,
                    'progress': 50,
                    'message': '等待用户确认修复方案',
                    'pending_confirmation': pending_confirmation,
                    'logs': [
                        {'time': datetime.now().isoformat(), 'level': 'warning', 'message': '检测到错误，等待用户确认修复方案'}
                    ]
                }
            
            # 旧逻辑：检查 pending_confirmation.status == "waiting"（兼容）
            pending_confirmation = result.get("pending_confirmation")
            if pending_confirmation and pending_confirmation.get("status") == "waiting":
                # 将确认状态存储到Celery任务meta中
                return {
                    'status': 'waiting_confirmation',
                    'current_stage': stage,
                    'progress': 50,
                    'message': '等待用户确认修复方案',
                    'pending_confirmation': pending_confirmation,
                    'logs': [
                        {'time': datetime.now().isoformat(), 'level': 'warning', 'message': '检测到错误，等待用户确认修复方案'}
                    ]
                }
            
            # 构建详细错误消息：优先保留 run_stage/stage_tools 返回的业务错误，
            # 避免被后续 manifest 校验错误覆盖成 "Manifest not found"。
            final_output = result.get("final_output", {}) or {}
            execution_results = result.get("execution_results", []) or []
            business_error = (
                final_output.get("response")
                or final_output.get("message")
                or final_output.get("error")
                or result.get("execution_error")
            )
            if not business_error:
                for item in execution_results:
                    stage_result = item.get("result", {}) if isinstance(item, dict) else {}
                    candidate = stage_result.get("error") or stage_result.get("message")
                    if candidate:
                        business_error = candidate
                        break

            if isinstance(business_error, str) and (
                "❌" in business_error
                or "failed" in business_error.lower()
                or "无法" in business_error
                or "缺少" in business_error
            ):
                error_msg = business_error
            else:
                error_msg = f'{stage} 执行失败'
            diagnosis = result.get("diagnosis_result", {})
            if diagnosis:
                suggestion = diagnosis.get("suggestion", "")
                error_type = diagnosis.get("error_type", "")
                if suggestion:
                    error_msg += f"\n\n诊断: {error_type}\n建议: {suggestion}"
            
            return _attach_diagnosis_to_failure({
                'status': 'failed',
                'success': False,  # 添加 success 字段，供前端判断
                'current_stage': stage,
                'progress': 0,
                'session_id': session_id or job_id,
                'message': error_msg,
                'errors': result.get("errors", ["Unknown error"]),
                'diagnosis': diagnosis,
                'logs': [
                    {'time': datetime.now().isoformat(), 'level': 'error', 'message': error_msg}
                ]
            }, job_id, stage, prefix)
            
    except Exception as e:
        import traceback
        error_msg = f"任务异常: {str(e)}\n{traceback.format_exc()}"
        print(f"[Celery] Single stage task {task_id} failed: {error_msg}")
        celery_task_id = self.request.id if self.request else task_id
        self.update_state(
            task_id=celery_task_id,
            state='FAILURE',
            meta={
                'exc_type': type(e).__name__,
                'exc_message': str(e),
            }
        )
        raise
    finally:
        if resource_acquired and resource_lease is not None:
            try:
                resource_lease.release(status=resource_outcome, message=f"{stage} {resource_outcome}")
            except Exception as release_exc:
                print(f"[Celery] Failed to release resource lease for {task_id}: {release_exc}")