diff --git a/apps/api/app/api/routes/generation_tasks.py b/apps/api/app/api/routes/generation_tasks.py index 9af9d0ecb..fe24672a2 100755 --- a/apps/api/app/api/routes/generation_tasks.py +++ b/apps/api/app/api/routes/generation_tasks.py @@ -358,6 +358,49 @@ def create_generation_task( ) ) try: + # 兜底关联编辑计划:前端未传 source_edit_plan_id 时, + # 通过 template_id + user_id 在 DB 层直接查找最新的 plan。 + # 必须在 enqueue 之前执行,避免 worker 读取时 source_edit_plan_id 为空(竞态条件) + if not task.source_edit_plan_id and request.template_id: + try: + from packages.adapters.sqlalchemy_impl.models import EditPlanModel + + _plan_model = ( + db.query(EditPlanModel) + .filter( + EditPlanModel.template_id == request.template_id, + EditPlanModel.created_by_user_id == user_id, + ) + .order_by(EditPlanModel.created_at.desc()) + .first() + ) + if _plan_model: + task.source_edit_plan_id = _plan_model.id + generation_task_repository.update(task) + logger.info( + "[生成任务] 自动关联编辑计划: task_id=%s plan_id=%s", + task.id, + _plan_model.id, + ) + except Exception: + logger.warning( + "[生成任务] 查找关联编辑计划失败(不影响主流程): task_id=%s", + task.id, + exc_info=True, + ) + + # 回写 plan.config:必须在 enqueue 之前执行, + # 确保 worker 读取 plan 时 config 中已包含 generation_task_id。 + # 只在首个任务时回写一次,避免批量生成时循环覆盖。 + _effective_plan_id = task.source_edit_plan_id + if _effective_plan_id and len(created_tasks) == 0: + _writeback_edit_plan_config( + plan_id=_effective_plan_id, + task_id=task.id, + title_config=request.title_config, + db=db, + ) + if safe_enqueue_generation_task( task, generation_task_repository, @@ -366,45 +409,6 @@ def create_generation_task( log_task_status=True, ): created_tasks.append(task) - # 只在首个成功任务时回写一次 plan.config, - # 避免批量生成时循环覆盖 generation_task_id - if request.source_edit_plan_id and len(created_tasks) == 1: - _writeback_edit_plan_config( - plan_id=request.source_edit_plan_id, - task_id=task.id, - title_config=request.title_config, - db=db, - ) - - # 兜底关联编辑计划:前端未传 source_edit_plan_id 时, - # 通过 template_id + user_id 在 DB 层直接查找最新的 plan - if not task.source_edit_plan_id and request.template_id: - try: - from packages.adapters.sqlalchemy_impl.models import EditPlanModel - - _plan_model = ( - db.query(EditPlanModel) - .filter( - EditPlanModel.template_id == request.template_id, - EditPlanModel.created_by_user_id == user_id, - ) - .order_by(EditPlanModel.created_at.desc()) - .first() - ) - if _plan_model: - task.source_edit_plan_id = _plan_model.id - generation_task_repository.update(task) - logger.info( - "[生成任务] 自动关联编辑计划: task_id=%s plan_id=%s", - task.id, - _plan_model.id, - ) - except Exception: - logger.warning( - "[生成任务] 查找关联编辑计划失败(不影响主流程): task_id=%s", - task.id, - exc_info=True, - ) else: failed_tasks.append(task) except UserPendingLimitExceeded as _e: diff --git a/tests/unit/test_race_condition_fallback_plan.py b/tests/unit/test_race_condition_fallback_plan.py new file mode 100644 index 000000000..5690f22f5 --- /dev/null +++ b/tests/unit/test_race_condition_fallback_plan.py @@ -0,0 +1,98 @@ +""" +回归测试:验证 generation_tasks.py 中兜底关联 edit plan 在 enqueue 之前执行。 + +根因(PR #1481 后续修复):兜底关联逻辑原来在 safe_enqueue_generation_task 之后执行, +导致 worker 在 enqueue 后立即读取 task 时,source_edit_plan_id 仍为空(竞态条件)。 +""" + +import ast +import textwrap + + +def _get_function_source(filepath, func_name): + """提取函数源码""" + with open(filepath) as f: + source = f.read() + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: + lines = source.splitlines() + start = node.lineno - 1 + end = node.end_lineno + return textwrap.dedent("\n".join(lines[start:end])) + return None + + +def _find_try_block_source(func_source): + """在函数源码中找到包含 safe_enqueue_generation_task 的 try 块""" + tree = ast.parse(textwrap.dedent(func_source)) + for node in ast.walk(tree): + if isinstance(node, ast.Try): + # 检查 try 块中是否包含 safe_enqueue_generation_task + block_lines = func_source.splitlines() + block_text = "\n".join(block_lines[node.lineno - 1 : node.end_lineno]) + if "safe_enqueue_generation_task" in block_text: + return block_text + return None + + +def test_fallback_before_enqueue(): + """验证兜底关联 edit plan 的代码在 safe_enqueue_generation_task 调用之前""" + filepath = "apps/api/app/api/routes/generation_tasks.py" + func_source = _get_function_source(filepath, "create_generation_task") + assert func_source is not None, "create_generation_task function not found" + + try_block = _find_try_block_source(func_source) + assert try_block is not None, "try block with safe_enqueue_generation_task not found" + + # 定位关键标记在 try 块中的行号 + lines = try_block.splitlines() + + fallback_line = None + enqueue_line = None + writeback_line = None + + for i, line in enumerate(lines): + if "not task.source_edit_plan_id and request.template_id" in line and fallback_line is None: + fallback_line = i + if "safe_enqueue_generation_task(" in line and enqueue_line is None: + enqueue_line = i + if "_writeback_edit_plan_config(" in line and "def " not in line and writeback_line is None: + writeback_line = i + + assert fallback_line is not None, "兜底关联逻辑 not found in try block" + assert enqueue_line is not None, "safe_enqueue_generation_task call not found in try block" + assert writeback_line is not None, "_writeback_edit_plan_config call not found in try block" + + # 核心断言:兜底关联和回写都在 enqueue 之前 + assert fallback_line < enqueue_line, f"兜底关联(行{fallback_line})应在 enqueue(行{enqueue_line})之前" + assert writeback_line < enqueue_line, f"回写 config(行{writeback_line})应在 enqueue(行{enqueue_line})之前" + + +def test_fallback_sets_source_edit_plan_id(): + """验证兜底关联逻辑会设置 task.source_edit_plan_id""" + filepath = "apps/api/app/api/routes/generation_tasks.py" + func_source = _get_function_source(filepath, "create_generation_task") + assert func_source is not None + + try_block = _find_try_block_source(func_source) + assert try_block is not None + + # 验证兜底逻辑包含赋值语句 + assert "task.source_edit_plan_id = _plan_model.id" in try_block, "兜底关联逻辑应设置 task.source_edit_plan_id" + assert "generation_task_repository.update(task)" in try_block, "兜底关联后应持久化 task 到 DB" + + +def test_writeback_uses_effective_plan_id(): + """验证回写 config 使用的是 effective_plan_id(包含兜底结果),而非仅 request.source_edit_plan_id""" + filepath = "apps/api/app/api/routes/generation_tasks.py" + func_source = _get_function_source(filepath, "create_generation_task") + assert func_source is not None + + try_block = _find_try_block_source(func_source) + assert try_block is not None + + # 验证使用了 _effective_plan_id 或 task.source_edit_plan_id,而非仅 request.source_edit_plan_id + # 修复前用的是 request.source_edit_plan_id,修复后应该用 task.source_edit_plan_id + uses_effective = "_effective_plan_id" in try_block or "task.source_edit_plan_id" in try_block + assert uses_effective, "回写 config 应使用包含兜底结果的有效 plan_id"