writer-work-flow/apps/api/ww_api/services/chain_runner.py

"""多章链长任务 runner（chain-workflow 设计 §5）——`run_chain_job` 在 jobs 壳里跑链图。

复用 `run_job` 的生命周期心智（set_running → work → complete/fail），但链多一态：
**interrupt 命中 → `awaiting_input`**（等作者裁决续跑）。故不直接复用 `run_job`，而是
仿其壳自建独立 session 写 job 状态，work 本体 = 驱动 LangGraph 链图。

纪律（守不变量）：
- **#1/#5**：图节点各自短事务落领域表（`chapters`/`chapter_reviews`/`chapter_digests`）；
  checkpointer 只存控制流；resume 重读领域表。本 runner 不持业务事务。
- **#3/#4**：`accept_op`（apps/api 注入）跑冲突 gate → 终稿提炼 digest → 单原子验收事务 →
  伏笔到期扫描，唯一正文写经此。
- **token 不泄漏**：job result/status/日志只记章号/计数/标志，绝不含 prompt/正文/token。
- **错误脱敏**：复用 `_classify_job_error`（AppError 透传安全文案，其余通用文案，P0-3）。

可测性：`session_factory`/`gateway`/`checkpointer`/`accept_op` 全可注入——单测直接 `await`
它，注 fake session 工厂 + mock 网关 + MemorySaver + fake accept_op（绝不联网/真 DB/真 LLM）。
"""

from __future__ import annotations

import uuid
from collections.abc import Awaitable, Callable, Sequence
from contextlib import AbstractAsyncContextManager
from typing import Any

import structlog
from langchain_core.runnables import RunnableConfig
from langgraph.types import Command
from ww_core.domain.chapter_repo import ChapterRepo, SqlChapterRepo
from ww_core.domain.digest_repo import SqlDigestAppendRepo
from ww_core.domain.job_repo import SqlJobRepo
from ww_core.domain.review_repo import SqlReviewRepo
from ww_core.memory import assemble
from ww_core.memory.sql_repositories import sql_memory_repos
from ww_core.orchestrator import REVIEW_SPECS, GatewayRun
from ww_core.orchestrator.chain import build_chain_graph, initial_chain_state
from ww_llm_gateway import Gateway

from ww_api.schemas.projects import ConflictDecision
from ww_api.services.accept_service import (
    assert_conflicts_resolved,
    run_accept_transaction,
)
from ww_api.services.digest_extraction import extract_digest_facts
from ww_api.services.foreshadow_scan import SessionFactory, run_overdue_scan
from ww_api.services.job_runner import _classify_job_error

log = structlog.get_logger(__name__)

# 链 job 的 kind（jobs.kind 自由 Text 列，零迁移复用；§7）。
JOB_KIND_CHAIN = "chain"


# 「按 session 建 digest 网关」缝：accept 节点在自建短事务里需 light 档网关跑终稿提炼。
DigestGatewayBuilder = Callable[[Any], Awaitable[Gateway]]


def build_accept_op(
    *,
    session_factory: SessionFactory,
    digest_gateway_builder: DigestGatewayBuilder,
    request_id: str | None = None,
) -> AcceptOp:
    """组装单章验收落库闭包（apps/api，守不变量 #3/#4）。

    链 `accept_chapter` 节点经此闭包把作者裁决落库——每章**独立短事务**（自建 session）：
    重读最近审稿 → 冲突 gate（缺判 → `CONFLICT_UNRESOLVED`，链失败）→ 终稿(=已存草稿)提炼
    digest（事务外）→ `run_accept_transaction` 单原子提交 → 提交后伏笔到期扫描（自建 session）。

    自动链无作者改稿：`final_text` = `chapters` 里该章草稿正文（write 节点所落）。
    `decisions` 为 dict 清单（interrupt 恢复值 / 无冲突 []），此处校验为 `ConflictDecision`。
    """

    async def accept_op(
        *,
        project_id: uuid.UUID,
        chapter_no: int,
        user_id: uuid.UUID,
        decisions: list[Any],
    ) -> None:
        parsed_decisions = [
            d if isinstance(d, ConflictDecision) else ConflictDecision.model_validate(d)
            for d in (decisions or [])
        ]
        async with session_factory() as session:
            chapter_repo: ChapterRepo = SqlChapterRepo(session)
            review_repo = SqlReviewRepo(session)
            digest_repo = SqlDigestAppendRepo(session)

            history = await review_repo.list_for_chapter(project_id, chapter_no)
            latest_review = history[0] if history else None
            # 冲突 gate（R5，事务前拦截，不写库；缺判 → CONFLICT_UNRESOLVED 抛，链失败）。
            assert_conflicts_resolved(latest_review, parsed_decisions)

            draft = await chapter_repo.get_draft(project_id, chapter_no)
            final_text = draft.content if draft is not None else ""

            gateway = await digest_gateway_builder(session)
            # R2：终稿 digest 提炼在开原子事务之前（别在持开事务里跨网络调 LLM）。
            digest_facts = await extract_digest_facts(
                gateway,
                final_text=final_text,
                user_id=user_id,
                project_id=project_id,
                chapter_no=chapter_no,
            )
            await run_accept_transaction(
                session=session,
                chapter_repo=chapter_repo,
                digest_repo=digest_repo,
                review_repo=review_repo,
                project_id=project_id,
                chapter_no=chapter_no,
                final_text=final_text,
                digest_facts=digest_facts,
                latest_review=latest_review,
                decisions=parsed_decisions,
            )
        # 验收提交后伏笔到期扫描（自建独立 session，§6.2；与单章端点同纪律）。
        await run_overdue_scan(
            session_factory,
            project_id=project_id,
            chapter_no=chapter_no,
            request_id=request_id,
        )

    return accept_op


# `()` → async-CM，进入得一个 langgraph checkpointer（Postgres 运行时 / MemorySaver 测试）。
CheckpointerCtx = Callable[[], AbstractAsyncContextManager[Any]]

#: apps/api 注入的「按 session 建 chapter draft repo」工厂（review 节点重读草稿用）。
ChapterRepoFactory = Callable[[Any], Any]
#: apps/api 注入的「按 session 建审稿留痕 repo」工厂（review/decide 节点用）。
ReviewRepoFactory = Callable[[Any], Any]
#: 单章验收落库闭包（apps/api 组装：冲突 gate + digest + 验收事务 + 到期扫描）。
AcceptOp = Callable[..., Awaitable[None]]


def _chain_result(state: dict[str, Any], *, awaiting_chapter: int | None) -> dict[str, Any]:
    """从图 state 取非密进度摘要（绝不含正文/token；§5）。

    `awaiting_chapter` 非 None = interrupt 命中该章待裁决；None 且 completed = 全跑完。
    """
    written = list(state.get("written") or [])
    return {
        "chain_key": state.get("chain_key"),
        "written": written,
        "completed": awaiting_chapter is None,
        "awaiting_chapter": awaiting_chapter,
    }


def _extract_awaiting_chapter(final: dict[str, Any]) -> int | None:
    """从图返回值判 interrupt：有 `__interrupt__` → 取暂停章号；否则 None（跑完）。

    interrupt 载荷 = `{"chapter_no": n}`（见链图 `_accept` 节点）。LangGraph 把它放在
    `final["__interrupt__"]`（Interrupt 对象列表）。容错读 `.value`/dict 两种形。
    """
    interrupts = final.get("__interrupt__")
    if not interrupts:
        return None
    first = interrupts[0]
    payload = getattr(first, "value", first)
    if isinstance(payload, dict):
        chapter_no = payload.get("chapter_no")
        return int(chapter_no) if chapter_no is not None else None
    return None


async def run_chain_job(
    session_factory: SessionFactory,
    job_id: uuid.UUID,
    *,
    project_id: uuid.UUID,
    user_id: uuid.UUID,
    chain_key: str,
    start_chapter_no: int,
    count: int,
    gateway: GatewayRun,
    checkpointer_ctx: CheckpointerCtx,
    accept_op: AcceptOp,
    chapter_repo_factory: ChapterRepoFactory,
    review_repo_factory: ReviewRepoFactory,
    resume_decisions: list[ConflictDecision] | None = None,
    review_specs: Sequence[Any] = REVIEW_SPECS,
    request_id: str | None = None,
) -> None:
    """跑/续一条多章链：set_running → 驱动链图 → 据返回置 done/awaiting_input/failed。

    `resume_decisions is None` → 初始 run（`graph.ainvoke(initial_state, config)`）；
    否则 resume（`graph.ainvoke(Command(resume=decisions), config)`），从 interrupt 续跑。
    `thread_id = str(job_id)`：同 job 的初始/续跑落同一检查点 thread。

    checkpointer 上下文（Postgres 连接 / MemorySaver）横跨整次 invoke；本 runner 在 task
    内打开/关闭它。异常一律被吞（后台任务边界），脱敏后置 job failed，不冒泡崩进程。
    """
    try:
        async with checkpointer_ctx() as checkpointer:
            await _set_running(session_factory, job_id)
            config: RunnableConfig = {"configurable": {"thread_id": str(job_id)}}
            graph = build_chain_graph(
                gateway,
                session_factory=session_factory,
                memory_repos_factory=sql_memory_repos,
                chapter_repo_factory=chapter_repo_factory,
                review_repo_factory=review_repo_factory,
                assemble=assemble,
                accept_op=accept_op,
                review_specs=review_specs,
                checkpointer=checkpointer,
            )
            if resume_decisions is None:
                initial = initial_chain_state(
                    project_id=project_id,
                    user_id=user_id,
                    start_chapter_no=start_chapter_no,
                    count=count,
                    chain_key=chain_key,
                )
                raw_final = await graph.ainvoke(initial, config=config)
            else:
                resume_value = [d.model_dump() for d in resume_decisions]
                raw_final = await graph.ainvoke(Command(resume=resume_value), config=config)

            final: dict[str, Any] = dict(raw_final)
            awaiting_chapter = _extract_awaiting_chapter(final)
            result = _chain_result(final, awaiting_chapter=awaiting_chapter)
            await _finish(session_factory, job_id, result, awaiting=awaiting_chapter is not None)
        log.info(
            "chain_job_settled",
            job_id=str(job_id),
            request_id=request_id,
            awaiting_chapter=awaiting_chapter,
            written_count=len(result["written"]),
        )
    except Exception as exc:  # noqa: BLE001 — 后台任务边界：记错误 + 置 job failed，不冒泡。
        log.error("chain_job_failed", job_id=str(job_id), request_id=request_id, error=str(exc))
        await _mark_failed(session_factory, job_id, _classify_job_error(exc), request_id)


async def _set_running(session_factory: SessionFactory, job_id: uuid.UUID) -> None:
    """独立短事务把 job 置 running（图驱动开始前）。"""
    async with session_factory() as session:
        await SqlJobRepo(session).set_running(job_id)
        await session.commit()


async def _finish(
    session_factory: SessionFactory,
    job_id: uuid.UUID,
    result: dict[str, Any],
    *,
    awaiting: bool,
) -> None:
    """据是否 interrupt 把 job 置 awaiting_input（待裁决）或 done（跑完），独立短事务。"""
    async with session_factory() as session:
        repo = SqlJobRepo(session)
        if awaiting:
            await repo.set_awaiting(job_id, result)
        else:
            await repo.complete(job_id, result)
        await session.commit()


async def _mark_failed(
    session_factory: SessionFactory,
    job_id: uuid.UUID,
    error: str,
    request_id: str | None,
) -> None:
    """在**全新** session 里把 job 置 failed（前一事务已因异常作废）。"""
    try:
        async with session_factory() as session:
            await SqlJobRepo(session).fail(job_id, error)
            await session.commit()
    except Exception as exc:  # noqa: BLE001 — 置失败态本身再炸只记日志，不冒泡。
        log.error(
            "chain_job_fail_mark_failed",
            job_id=str(job_id),
            request_id=request_id,
            error=str(exc),
        )


__all__ = [
    "JOB_KIND_CHAIN",
    "AcceptOp",
    "CheckpointerCtx",
    "DigestGatewayBuilder",
    "build_accept_op",
    "run_chain_job",
]