ai-kefu-admin/app/platform/site_import/task.py

"""site_import 内存任务表（spec §3.3）——单进程 asyncio，(slug, task_id) 索引。

内存铁律：30 分钟惰性清理（查询/新建时顺带扫）；强引用防 GC（裸
create_task 会被回收——asyncio 官方坑，模块级 dict 存 task 对象）；
全平台并发 ≤2（复活任务占用同一名额）；频控 3 次/15 分钟（新建与
retry 同计）。异常类型：ImportRateLimited→429 / ConcurrentLimit→409 /
LookupError→404 / ValueError→409——API 层（Task 7）映射。

7 分钟总超时兜底：超时落在生成段（抓取产物已在）标 failed_phase=
"generate"——retry=generate 是 spec §3.3 的预期恢复路径（不重抓）；
落在抓取段则无产物，重开任务。
"""
from __future__ import annotations

import asyncio
import time
import uuid
from collections import defaultdict

from app.platform.site_import import crawler, generator
from app.platform.site_import.crawler import CrawlError, INSUFFICIENT_MESSAGE
from app.platform.site_import.generator import GenerateError
from app.platform.site_import.safe_fetch import FetchError

TASK_TTL = 30 * 60
TOTAL_TIMEOUT = 7 * 60
MAX_CONCURRENT = 2
IMPORT_WINDOW = 15 * 60
IMPORT_MAX = 3


class ImportRateLimited(Exception):
    """频控超限（15 分钟 3 次）——API 层映射 429。"""


class ConcurrentLimit(Exception):
    """全平台并发上限（≤2，复活任务同占名额）——API 层映射 409。"""


class _SlidingWindow:
    """(kind, slug) 键滑动窗计数——config_routes.SlidingWindowCounter 同款
    风格的独立实例（勿共用 email/push 测试端点的预算）。同款语义：超限后
    allow 守卫先于一切预检；计数只落在真正发起的这次（校验失败不消耗）。"""

    def __init__(self, window_seconds: float = IMPORT_WINDOW,
                 max_calls: int = IMPORT_MAX) -> None:
        self.window_seconds = window_seconds
        self.max_calls = max_calls
        self._calls: dict[tuple[str, str], list[float]] = defaultdict(list)

    def allow(self, key: tuple[str, str]) -> bool:
        self._prune(key, time.time())
        return len(self._calls.get(key, ())) < self.max_calls

    def record(self, key: tuple[str, str]) -> None:
        now = time.time()
        self._prune(key, now)
        self._calls[key].append(now)

    def _prune(self, key: tuple[str, str], now: float) -> None:
        calls = self._calls.get(key)
        if calls is None:
            return
        kept = [t for t in calls if t > now - self.window_seconds]
        if kept:
            self._calls[key] = kept
        else:
            self._calls.pop(key, None)


# 模块级单例：导入频控（键 ("import", slug)，全平台各租户独立预算）
_import_limiter = _SlidingWindow()

# 内存任务表：(slug, task_id) 联合索引——跨租户持他人 id 查询即 miss
_tasks: dict[tuple[str, str], "ImportTask"] = {}
_lock = asyncio.Lock()


class TaskState:
    PENDING = "pending"
    CRAWLING = "crawling"
    GENERATING = "generating"
    DONE = "done"
    FAILED = "failed"

    ACTIVE = (PENDING, CRAWLING, GENERATING)


class ImportTask:
    """一条导入任务：状态机 + 阶段产物 + 强引用的后台 aio task。

    failed_phase："generate" 可 retry=generate 复活（复用 crawl_result，
    不重抓）；"crawl"（SPA/内容不足/FetchError 网站打不开）与 None（超时
    落在抓取段/未知异常）不可复活——重开任务。
    """

    def __init__(self, slug: str, task_id: str, url: str | None) -> None:
        self.slug = slug
        self.task_id = task_id
        self.url = url
        self.state = TaskState.PENDING
        self.phase: str | None = None  # 运行中阶段（crawling/generating）
        self.pages_fetched = 0
        self.draft: dict | None = None
        self.error: str | None = None
        self.created_at = time.time()
        self.finished_at: float | None = None
        self.failed_phase: str | None = None
        self.crawl_result: dict | None = None  # retry=generate 复用
        self.need_crawl = True
        self._aio_task: asyncio.Task | None = None  # 强引用防 GC

    def fail(self, message: str, *, phase: str | None) -> None:
        self.state = TaskState.FAILED
        self.error = message
        self.failed_phase = phase
        self.phase = None
        self.finished_at = time.time()

    def _finish(self, draft: dict) -> None:
        self.draft = draft
        self.state = TaskState.DONE
        self.phase = None
        self.finished_at = time.time()


def _prune_expired(now: float | None = None) -> None:
    """30 分钟惰性清理（查询/新建顺带扫）。运行中任务防御性取消——
    7 分钟总超时下正常不会发生（30 >> 7），仅兜异常态。"""
    now = time.time() if now is None else now
    deadline = now - TASK_TTL
    for key in [k for k, t in _tasks.items() if t.created_at < deadline]:
        stale = _tasks.pop(key)
        if stale._aio_task is not None and not stale._aio_task.done():
            stale._aio_task.cancel()


def _count_active() -> int:
    return sum(1 for t in _tasks.values() if t.state in TaskState.ACTIVE)


def active_count() -> int:
    """全平台进行中任务数（pending/crawling/generating）。"""
    _prune_expired()
    return _count_active()


def get(slug: str, task_id: str) -> ImportTask | None:
    """(slug, task_id) 联合索引查询 + 惰性清理（跨租户 → None → 404）。"""
    _prune_expired()
    return _tasks.get((slug, task_id))


async def start(slug: str, url: str | None, *,
                 retry_task_id: str | None = None) -> str:
    """新建或复活（retry=generate）导入任务，返回 task_id。

    频控守卫最先（超限后 429 优先于一切预检）；record 只落在真正发起的
    这次——LookupError/ValueError/ConcurrentLimit 不消耗预算。
    """
    _prune_expired()
    key = ("import", slug)
    if not _import_limiter.allow(key):
        raise ImportRateLimited(
            f"导入太频繁，{IMPORT_WINDOW // 60} 分钟内最多 {IMPORT_MAX} 次，"
            "稍后再试")
    async with _lock:
        if retry_task_id is not None:
            target = _tasks.get((slug, retry_task_id))
            if target is None:
                raise LookupError(
                    f"task {retry_task_id!r} not found for tenant {slug!r}")
            if (target.state != TaskState.FAILED
                    or target.failed_phase != "generate"):
                raise ValueError("任务当前状态不可重试"
                                 "（仅 generate 阶段失败可复活）")
        # 并发名额对新建与 retry 复活同样生效（spec §3.3）
        if _count_active() >= MAX_CONCURRENT:
            raise ConcurrentLimit("已有导入任务进行中，稍后再试")
        if retry_task_id is not None:
            # 复活同 task_id：复用已抓产物，不重抓
            target.state = TaskState.PENDING
            target.error = None
            target.failed_phase = None
            target.finished_at = None
            target.need_crawl = False
        else:
            target = ImportTask(slug, uuid.uuid4().hex, url)
            _tasks[(slug, target.task_id)] = target
        _import_limiter.record(key)
        target._aio_task = asyncio.create_task(_pipeline(target))
        return target.task_id


async def _pipeline(task_obj: ImportTask) -> None:
    """后台管线：7 分钟总超时兜底 + 异常归档为白话失败。"""
    try:
        await asyncio.wait_for(_run(task_obj), TOTAL_TIMEOUT)
    except asyncio.TimeoutError:
        task_obj.fail(
            "导入超时，请重试",
            phase="generate" if task_obj.crawl_result is not None else None)
    except CrawlError as exc:
        task_obj.fail(str(exc), phase="crawl")  # SPA/内容不足——不可 retry
    except FetchError:
        # spec §6 行1：入口页不可达/超时是最高频真实失败路径——白话文案，
        # 不漏进通用兜底裸抛英文内部信息
        task_obj.fail("网站打不开或响应太慢——检查 URL", phase="crawl")
    except GenerateError as exc:
        task_obj.fail(str(exc), phase="generate")  # retry=generate 复活
    except Exception as exc:  # 防御：未知异常不裸抛（任务表对象完整性优先）
        task_obj.fail(f"导入失败：{exc}", phase=None)


async def _run(task_obj: ImportTask) -> None:
    if task_obj.need_crawl:
        task_obj.state = TaskState.CRAWLING
        task_obj.phase = "crawling"
        task_obj.crawl_result = await crawler.crawl(task_obj.url)
        task_obj.pages_fetched = int(
            task_obj.crawl_result.get("pages_fetched") or 0)
    task_obj.state = TaskState.GENERATING
    task_obj.phase = "generating"
    provider = generator.make_provider(None)  # resolver 三级链（平台设置/env）
    # source_text 传 analyze 实际消费的截断后文本（24K 同款——Task 5 备忘：
    # 否则截断段数字会被 build 的数字校验误剔除）
    source = ((task_obj.crawl_result or {}).get("pages_text")
              or "")[:generator.TOTAL_INPUT_CAP]
    analysis = await generator.analyze(provider, source)
    if not (analysis.get("product_lines") or []):
        # 内容不足三触发之一（调用 1 产品线为空）——同 crawl 类不可复活
        task_obj.fail(INSUFFICIENT_MESSAGE, phase="crawl")
        return
    gen = await generator.generate(provider, analysis)
    task_obj._finish(
        generator.build(gen, analysis, task_obj.crawl_result, source))


def _reset_for_test() -> None:
    """测试隔离（conftest autouse 前后各跑）：任务表清空（运行中 aio
    任务防御性取消）+ 频控窗计数/参数复位。"""
    for t in _tasks.values():
        if t._aio_task is not None and not t._aio_task.done():
            t._aio_task.cancel()
    _tasks.clear()
    _import_limiter._calls.clear()
    _import_limiter.window_seconds = IMPORT_WINDOW
    _import_limiter.max_calls = IMPORT_MAX