优化脚本就该模型

2026-08-21 15:08:24 +08:00
parent c2d3a786a3
commit 158c66ac25
12 changed files with 273 additions and 11 deletions
@@ -515,6 +515,61 @@ class AdminTaskMonitorTests(TestCase):
        task = self._mk_stuck_reserved("k-reap-perm")
        self.assertEqual(self.nc.post(f"/api/admin/tasks/{task.id}/reap/").status_code, 403)

+    def _mk_inflight_video(self, *, key: str, provider_task_id: str = "ark-poll-1"):
+        return self.AITask.objects.create(
+            team=self.team,
+            model_config=self.mc,
+            task_type=self.AITask.Type.FREE_VIDEO,
+            status=self.AITask.Status.POLLING,
+            estimated_cost="1.0",
+            actual_cost="0",
+            idempotency_key=key,
+            provider_task_id=provider_task_id,
+        )
+
+    def test_poll_inflight_free_video(self):
+        task = self._mk_inflight_video(key="k-poll-fv")
+
+        def fake_finalize(*, task):
+            task.status = self.AITask.Status.SUCCEEDED
+            task.save(update_fields=["status", "updated_at"])
+            return task
+
+        with patch("apps.ai.free_video.finalize_free_video", side_effect=fake_finalize) as mock_fin:
+            r = self.ac.post("/api/admin/tasks/poll/", {}, format="json")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.data["polled"], 1)
+        self.assertEqual(r.data["statuses"][str(task.id)], "succeeded")
+        mock_fin.assert_called_once()
+        task.refresh_from_db()
+        self.assertEqual(task.status, self.AITask.Status.SUCCEEDED)
+
+    def test_poll_skips_terminal_and_missing_provider_id(self):
+        self.AITask.objects.create(
+            team=self.team, model_config=self.mc, task_type=self.AITask.Type.FREE_VIDEO,
+            status=self.AITask.Status.POLLING, estimated_cost="1.0", actual_cost="0",
+            idempotency_key="k-poll-nopid",
+        )
+        with patch("apps.ai.free_video.finalize_free_video") as mock_fin:
+            r = self.ac.post("/api/admin/tasks/poll/", {}, format="json")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.data["polled"], 0)
+        mock_fin.assert_not_called()
+
+    def test_poll_task_ids_filter(self):
+        keep = self._mk_inflight_video(key="k-poll-keep", provider_task_id="ark-keep")
+        skip = self._mk_inflight_video(key="k-poll-skip", provider_task_id="ark-skip")
+        with patch("apps.ai.free_video.finalize_free_video", side_effect=lambda *, task: task) as mock_fin:
+            r = self.ac.post("/api/admin/tasks/poll/", {"task_ids": [str(keep.id)]}, format="json")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.data["polled"], 1)
+        self.assertIn(str(keep.id), r.data["statuses"])
+        self.assertNotIn(str(skip.id), r.data["statuses"])
+        mock_fin.assert_called_once()
+
+    def test_poll_requires_admin(self):
+        self.assertEqual(self.nc.post("/api/admin/tasks/poll/", {}, format="json").status_code, 403)
+

 class AdminBillingTests(TestCase):
    """Phase 7:计费审计(流水浏览/手动调额)+ 4 层额度策略(CRUD + 拦截生效)+ 权限。"""
@@ -21,6 +21,7 @@ from .views import (
    admin_task_reap,
    admin_task_retry,
    admin_tasks,
+    admin_tasks_poll,
    admin_prompt_template_detail,
    admin_prompt_templates,
    admin_quality_word_detail,
@@ -54,6 +55,7 @@ urlpatterns = [
    path("asset-reviews/submit/", admin_asset_reviews_submit, name="admin-asset-reviews-submit"),
    path("asset-reviews/poll/", admin_asset_reviews_poll, name="admin-asset-reviews-poll"),
    path("tasks/", admin_tasks, name="admin-tasks"),
+    path("tasks/poll/", admin_tasks_poll, name="admin-tasks-poll"),
    path("tasks/<uuid:task_id>/", admin_task_detail, name="admin-task-detail"),
    path("tasks/<uuid:task_id>/retry/", admin_task_retry, name="admin-task-retry"),
    path("tasks/<uuid:task_id>/reap/", admin_task_reap, name="admin-task-reap"),
@@ -1,5 +1,6 @@
 """平台超管后台 · 跨团队端点。所有视图统一挂 IsPlatformAdmin,非超管一律 403,写操作记审计。"""

+import logging
 from decimal import Decimal, ROUND_HALF_UP

 from django.db.models import Count, F, Q
@@ -39,6 +40,12 @@ from .serializers import (
    QualityWordSerializer,
 )

+logger = logging.getLogger(__name__)
+
+# 后台「刷新状态」只向供应商拉取已提交的异步视频任务;单次上限避免请求拖死。
+_ADMIN_TASK_POLL_LIMIT = 20
+_ADMIN_TASK_POLL_STATUSES = (AITask.Status.SUBMITTED, AITask.Status.POLLING)
+

 def _team_qs():
    return (
@@ -410,6 +417,31 @@ def admin_asset_reviews_poll(request):
    return Response({"polled": len(statuses), "statuses": statuses})


+def _refresh_inflight_task(task: AITask, *, operator) -> AITask:
+    """对单条在途任务向供应商拉一次最新态。无远端 ID / 非视频异步任务则原样返回。"""
+    if not task.provider_task_id:
+        return task
+    if task.task_type == AITask.Type.FREE_VIDEO:
+        from apps.ai.free_video import finalize_free_video
+
+        return finalize_free_video(task=task)
+    if task.task_type == AITask.Type.VIDEO_SEGMENT:
+        from apps.ai.services import poll_video_segment
+        from apps.projects.models import VideoSegment
+
+        segment_id = (task.request_payload or {}).get("video_segment_id")
+        if not segment_id:
+            return task
+        segment = VideoSegment.objects.filter(id=segment_id).first()
+        if segment is None:
+            return task
+        user = task.created_by or getattr(task.team, "owner", None) or operator
+        poll_video_segment(video_segment=segment, user=user)
+        task.refresh_from_db()
+        return task
+    return task
+
+
 # ─────────────────────────── AI 任务监控 + 成本异常 ───────────────────────────


@@ -440,6 +472,33 @@ def admin_tasks(request):
    return paginator.get_paginated_response(AdminTaskSerializer(page, many=True).data)


+@api_view(["POST"])
+@permission_classes([IsPlatformAdmin])
+def admin_tasks_poll(request):
+    """向供应商刷新在途视频任务状态。给 task_ids 则只拉这些,否则拉全平台 submitted/polling。
+
+    图片/脚本等同步任务没有远端 task id,点刷新后前端会再拉一次列表,worker 已落库的终态会一并更新。
+    """
+    ids = request.data.get("task_ids")
+    qs = (
+        AITask.objects.select_related("team", "model_config", "model_config__provider", "project")
+        .filter(status__in=_ADMIN_TASK_POLL_STATUSES)
+        .exclude(provider_task_id="")
+        .order_by("-updated_at")
+    )
+    if ids:
+        qs = qs.filter(id__in=ids)
+    statuses = {}
+    for task in qs[:_ADMIN_TASK_POLL_LIMIT]:
+        try:
+            refreshed = _refresh_inflight_task(task, operator=request.user)
+            statuses[str(refreshed.id)] = refreshed.status
+        except Exception:  # noqa: BLE001 — 单条失败不阻断整页刷新
+            logger.warning("admin poll task %s failed", task.id, exc_info=True)
+            statuses[str(task.id)] = task.status
+    return Response({"polled": len(statuses), "statuses": statuses})
+
+
@api_view(["GET"])
@permission_classes([IsPlatformAdmin])
 def admin_task_detail(request, task_id):