PlantAI-rag-new/bioflow_service/app/context.py

"""app/context.py —— 生信服务共享实例容器（db/events/kernel/catalog/runner/semaphore/planner）。

- 生产：app 启动时 init_context() 加载真实依赖（SQLite、stage_image_map.json、
  vstore FlowStore 本地实例（存在时）、真 docker 执行器）。
- 测试：build_context() 可注入 fake db/executor/flow_store（mock docker/LLM/vstore）。
- vstore 不存在/不可用：FlowStore 为 None，planner 自动降级 llm_only（FR-BIO-010 降级）。
- 并发信号量（BIOFLOW_CONCURRENCY，FR-BIO-031）每 context 一个。
"""
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field
from typing import Any

from .config import get_settings
from .db import Database, db as default_db
from .services.bioflow.catalog import StageCatalog, build_catalog
from .services.bioflow.events import EventBus
from .services.bioflow.kernel import Kernel, KernelConfig, load_kernel
from .services.bioflow.planner import Planner
from .services.bioflow.runner import Runner

logger = logging.getLogger(__name__)


@dataclass
class Context:
    db: Any = None
    events: EventBus | None = None
    kernel: Kernel | None = None
    catalog: StageCatalog | None = None
    runner: Runner | None = None
    semaphore: asyncio.Semaphore | None = None
    planner: Planner | None = None
    flow_store: Any = None
    executor: Any = None
    records: dict[str, Any] = field(default_factory=dict)  # job_id → JobRecord（内存）
    _tasks: list[asyncio.Task] = field(default_factory=list)

    def spawn(self, coro: Any) -> asyncio.Task:
        task = asyncio.create_task(coro)
        self._tasks.append(task)
        task.add_done_callback(lambda t: self._tasks.remove(t) if t in self._tasks else None)
        return task

    async def close(self) -> None:
        for task in self._tasks:
            task.cancel()
        if self._tasks:
            await asyncio.gather(*self._tasks, return_exceptions=True)
        self._tasks.clear()


# 全局 context（测试可整体替换）
context: Context = Context()


def init_events(jobs_dir: str) -> EventBus:
    from pathlib import Path

    return EventBus(Path(jobs_dir) / "events")


def init_flow_store(flow_docs_dir: str) -> Any | None:
    """vstore FlowStore 本地实例（存在则加载；失败/缺失返回 None → planner 降级）。"""
    try:
        import sys
        from pathlib import Path

        repo_root = Path(__file__).resolve().parent.parent.parent
        if str(repo_root) not in sys.path:
            sys.path.insert(0, str(repo_root))
        from vstore.flow import FlowStore  # P6 A4 交付；缺失则降级

        return FlowStore(flow_docs_dir)
    except Exception as exc:  # noqa: BLE001 - 缺失/损坏降级 llm_only，不阻断服务
        logger.warning("context: vstore FlowStore 不可用（planner 降级 llm_only）: %s", exc)
        return None


def build_context(
    *,
    db: Database | None = None,
    events: EventBus | None = None,
    kernel: Kernel | None = None,
    catalog: StageCatalog | None = None,
    executor: Any = None,
    semaphore: asyncio.Semaphore | None = None,
    flow_store: Any = None,
    runner_kwargs: dict[str, Any] | None = None,
) -> Context:
    """构造 context（生产与测试共用；各依赖可注入）。

    executor=None 时用 kernel.run_stage（真 docker）。
    """
    settings = get_settings()
    database = db or default_db
    bus = events or init_events(settings.jobs_dir_abs)
    kern = kernel or load_kernel()
    cat = catalog or build_catalog(kern.map)
    sem = semaphore or asyncio.Semaphore(settings.concurrency)
    exec_fn = executor or kern.run_stage
    runner = Runner(
        database,
        bus,
        exec_fn,
        cat,
        sem,
        **(runner_kwargs or {}),
    )
    planner = Planner(
        cat,
        flow_store,
        llm_plan_timeout=settings.llm_plan_timeout,
        plan_candidates=settings.plan_candidates,
        plan_max_rounds=settings.plan_max_rounds,
        archive_finder=_make_archive_finder(database),
    )
    return Context(
        db=database,
        events=bus,
        kernel=kern,
        catalog=cat,
        runner=runner,
        semaphore=sem,
        planner=planner,
        flow_store=flow_store,
        executor=exec_fn,
    )


def _make_archive_finder(database: Database):
    """planner 复用归档查找回调（async；store 查询）。"""
    from .services.bioflow import store

    async def find(name: str) -> dict[str, Any] | None:
        try:
            return await store.find_archive_by_name(database, name)
        except Exception:  # noqa: BLE001
            return None

    return find


async def init_context(db_path: str | None = None) -> Context:
    """应用启动：初始化全局 context（幂等；db_path 供测试替换）。

    BUG-R1-03 修复：生产启动链接线 FlowStore——此前 build_context 恒不传
    flow_store（默认 None），部署态 ask/chat_qa 引用块恒空（citations=[]）。
    现在按配置 flow_docs_dir 加载本地实例（vstore/data/flow_docs 已入库 724
    块）；缺失/失败自动降级 None（planner 走 llm_only，服务不阻断）。
    """
    global context
    from .db import init_db

    database = context.db
    if db_path is not None:
        database = await init_db(db_path)
    elif context.db is None:
        database = await init_db()
    flow_store = init_flow_store(get_settings().flow_docs_dir_abs)
    context = build_context(db=database, flow_store=flow_store)
    # TL 接线(2026-09-12):扩展执行器接 Dev-B extension_exec。
    # 失败(如 docker 不可用/配置缺失)静默降级——trial/run 返回 503 而非 500。
    try:
        from .services.bioflow import extensions as ext_mod
        # 2026-09-15 修复:裸 ExtensionExecConfig() 的 entry_script_dir 为空串,
        # build_command 用相对路径 run_extension.py 在 176 cwd 下找不到入口。
        # 改走 _default_exec_config()(经 _resolve_extension_entry_dir 解析部署树路径)。
        from .services.bioflow.extension_exec import _default_exec_config, ExtensionExecutor as RealExecutor
        cfg = _default_exec_config()
        real = RealExecutor(cfg)
        store = ext_mod.get_store()
        if store is not None:
            ext_mod.inject_real_executor(real, store)
            logger.info("extension executor wired (image suffix=%s)", cfg.docker_tag_suffix)
    except Exception as exc:  # noqa: BLE001
        logger.warning("extension executor wiring skipped: %s", exc)
    return context


async def shutdown_context() -> None:
    await context.close()


# ---- JobRecord 注册与恢复（路由层操作任务的内存记录）----


def register_record(record: Any) -> None:
    context.records[record.id] = record


def _record_from_job(job: dict[str, Any], stages: list[dict[str, Any]]) -> Any:
    """DB 行 → JobRecord（无运行中内存记录时恢复，跨重启安全）。"""
    from .services.bioflow.runner import JobRecord, params_dict

    record = JobRecord(
        id=job["id"],
        session_id=job["session_id"],
        title=job["title"],
        status=job["status"],
        stage_chain=list(job["stage_chain"] or []),
        params=params_dict(job.get("params")),
        auto_continue=bool(job.get("auto_continue")),
        archive_id=job.get("archive_id"),
    )
    # 登记路径恢复（server_registered.json；供执行容器自动挂载）
    try:
        import json as _json
        reg_file = record.job_dir / "server_registered.json"
        if reg_file.is_file():
            record.registered_server_paths = (
                _json.loads(reg_file.read_text(encoding="utf-8"))
                .get("registered_server_paths") or []
            )
    except Exception:  # noqa: BLE001 - 恢复失败不阻断（挂载缺失→阶段报错可查）
        pass
    for s in stages:
        record._stage_states[s["stage"]] = {
            "status": s["status"],
            "params": params_dict(s.get("params")),
            "attempt": int(s.get("attempt") or 0),
            "manifest_path": s.get("manifest_path"),
            "log_tail": s.get("log_tail") or "",
            "error": s.get("error"),
            "started_at": s.get("started_at"),
            "finished_at": s.get("finished_at"),
        }
    return record


async def load_record(job_id: str) -> Any | None:
    """取任务 JobRecord：内存优先，否则 DB 恢复（并注册回内存）。"""
    if job_id in context.records:
        return context.records[job_id]
    if context.db is None:
        return None
    from .services.bioflow import store

    job = await store.get_job(context.db, job_id)
    if job is None:
        return None
    stages = await store.list_stages(context.db, job_id)
    record = _record_from_job(job, stages)
    context.records[job_id] = record
    return record