vent_agent_2.0/api/chat_routes.py

"""对话路由。

``POST /api/chat`` 是主 agent 的**唯一入口**，2.0 逐字保持 1.5 的对外契约
（multipart 入参、SSE 帧格式、thread_id ≡ session_id、错误码语义），这样 1.5
的前端不用改一行就能连上。

与 1.5 的差异有三处：

1. **不做鉴权**（实现层）：1.5 用 ``Depends(get_current_user)`` 校验
   ``x-access-token`` 并解析出用户；2.0 直接取 multipart 的可选字段
   ``user_name``（缺省 ``default``）。风险是客户端可自报身份，仅适用于
   内网/开发期——上线前把 ``_resolve_user`` 换成真实的 ``Depends`` 即可，
   其余代码不用动。
2. **没有意图识别**（实现层）：1.5 曾经按意图分流到不同的 agent，后来也
   收敛成了单一 dialog agent。2.0 直接就是单一 agent。
3. **删掉了 ``POST /api/chat/only``**（契约层，唯一一处收缩）：1.5 的这个
   端点与 ``/chat`` 除了流水上的 ``kind`` 标签外**行为完全相同**，而那个标签
   没有任何消费方（不聚合、不分账、不影响额度）。1.5 自己也为这份重复付过
   代价——见 AGENTS.md §六「为什么没有 /chat/only」。1.5 前端从未调用过它。
   若将来确有按渠道分账的需求，应当**先把聚合/额度按渠道拆开**，再引入入口，
   而不是先立一个行为相同的端点。

   保留 ``/chat`` 之外的另两个入口 ``/chat/resume`` 与 ``/chat/interrupt``：
   它们承载的是**不同能力**（中断恢复、挂起查询），不是同一能力的别名。
"""

from __future__ import annotations

import re
import uuid
from pathlib import Path
from typing import Any

from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile
from fastapi.responses import StreamingResponse
from starlette.background import BackgroundTask

import config
from agents import agent_registry, prompts
from runtime import execution
from runtime import concurrency
from api.models import ResumeRequest
from runtime.interrupts import pending_interrupts
from api.sse_core import stream_sse, sse_headers
from billing import guard
from db import chat_store, db_conn
from obs import get_logger

logger = get_logger(__name__)

router = APIRouter(tags=["对话"])

# 与 1.5 一致：plan 模式放宽递归上限（要先出计划再执行，步骤天然更多）
RECURSION_LIMIT_PLAN = 500
RECURSION_LIMIT_DEFAULT = 300

VALID_MODES = ("plan", "full")
# 标题自动生成的长度（骨架用截断而非 LLM 调用，见 _ensure_session_title）
TITLE_MAX_CHARS = 30
MAX_UPLOAD_MB = config.get_int("TFAGENTS_MAX_UPLOAD_MB", 20)
MAX_UPLOAD_SIZE = MAX_UPLOAD_MB * 1024 * 1024
# 一次请求的附件数量上限。风量研判至少要「配风计划 + 测风报表」两份，
# 所以不能只收一个；上限用于挡住误传整个目录。
MAX_UPLOAD_FILES = max(1, config.get_int("TFAGENTS_MAX_UPLOAD_FILES", 10))

_SAFE_NAME_RE = re.compile(r"[^\w\u4e00-\u9fff.\-]+")


# ============================================================
# 请求工具
# ============================================================

def _trace_id_of(request: Request | None) -> str:
    """取透传中间件解析出的 trace_id。"""
    if request is None:
        return ""
    return getattr(request.state, "trace_id", "") or ""


def _resolve_user(user_name: str | None) -> str:
    """解析归属用户。

    **这是 2.0 唯一的身份来源，也是它最大的安全缺口**：没有任何校验，
    客户端填谁就是谁。计费、会话归属、并发配额全都建立在这个值上，所以
    当前实现只适合内网或开发环境。接真实鉴权时把这里替换成依赖注入即可。
    """
    return (user_name or "default").strip() or "default"


# ============================================================
# 会话与模式
# ============================================================

def _init_session(
    session_id: str | None, thread_id: str | None, user_name: str, mode: str
) -> dict[str, Any]:
    """取得或新建会话。返回会话行。

    两条硬约束（沿用 1.5）：
    1. ``thread_id`` 必须等于 ``session_id``——checkpointer 是进程级共享的，
       允许客户端指定任意 thread_id 就等于允许读写别人的图状态。
    2. 会话必须归属当前用户，否则按 404 处理（不是 403：不泄露「这个 ID
       存在但不属于你」这个信息）。
    """
    if thread_id and session_id and thread_id != session_id:
        raise HTTPException(
            status_code=404, detail={"code": "thread_mismatch", "message": "会话不存在"}
        )

    if not session_id:
        return chat_store.create_session(user_name=user_name, mode=mode)

    session = chat_store.get_session(session_id)
    if session is None or session.get("user_name") != user_name:
        raise HTTPException(
            status_code=404, detail={"code": "session_not_found", "message": "会话不存在"}
        )
    return session


def _setup_mode(session: dict[str, Any], mode: str | None) -> tuple[str, dict[str, Any]]:
    """确定本轮权限模式并构造 LangGraph config。

    请求里带了合法 mode 就落库（会话级设置，后续轮次沿用）；没带就用会话
    里存的值。模式只影响「注入什么指令」与递归上限——2.0 不用
    interrupt_before/after 的静态中断，plan 模式的审批靠工具内 ``interrupt``。
    """
    current = session.get("mode") or "full"
    if mode:
        normalized = mode.strip().lower()
        if normalized not in VALID_MODES:
            raise HTTPException(
                status_code=400,
                detail={"code": "invalid_mode", "message": f"mode 只能是 {' | '.join(VALID_MODES)}"},
            )
        if normalized != current:
            chat_store.update_session_mode(session["id"], normalized)
        current = normalized

    recursion_limit = RECURSION_LIMIT_PLAN if current == "plan" else RECURSION_LIMIT_DEFAULT
    graph_config = {
        "configurable": {"thread_id": session["id"]},
        "recursion_limit": recursion_limit,
    }
    return current, graph_config


def _ensure_session_title(session: dict[str, Any], message: str) -> None:
    """首轮时给会话起个标题。

    骨架用**截断而非 LLM 生成**——标题不值得多花一次模型调用的钱和延迟。
    想换成 LLM 摘要只需替换本函数体，调用方不用动。
    """
    if session.get("title"):
        return
    title = (message or "").strip().split("\n", 1)[0][:TITLE_MAX_CHARS]
    if title:
        chat_store.update_session_title(session["id"], title)
        session["title"] = title


# ============================================================
# 消息注入
# ============================================================

def _resolve_refs_and_skill(
    original: str, user_name: str
) -> dict[str, Any]:
    """解析 ``#skill:`` 与 ``#session:`` 前缀。

    顺序是先技能后引用（与 1.5 一致）：两个前缀可以叠加，技能在前。
    技能名不做白名单校验——skill 目录里没有这个名字时，agent 读文件会失败
    并如实上报，不需要在路由层再维护一份名单（那份名单必然会和目录漂移）。
    """
    text = original or ""
    forced_skill: str | None = None
    reference: dict[str, Any] | None = None

    text, forced_skill = prompts.extract_skill_prefix(text)
    text, ref_id = prompts.extract_session_ref(text)

    if ref_id:
        ref_session = chat_store.get_session(ref_id)
        if ref_session is None or ref_session.get("user_name") != user_name:
            # 引用不存在的会话不报错，只在指令里说明——报错会让用户
            # 因为一个笔误丢掉整条消息
            reference = {"title": "", "excerpt": "", "error": "引用的会话不存在"}
        else:
            recent = chat_store.get_messages(ref_id, limit=6)
            excerpt = "\n".join(
                f"{'用户' if m.get('msg_role') == 'user' else '助手'}：{(m.get('content') or '')[:300]}"
                for m in recent
                if m.get("content")
            )
            reference = {"title": ref_session.get("title") or "未命名会话", "excerpt": excerpt}

    return {"message": text, "forced_skill": forced_skill, "reference": reference, "original": original}


def _inject_all(
    message: str,
    user_name: str,
    mode: str,
    ref: dict[str, Any],
    uploads: list[tuple[str, str]] | None = None,
) -> str:
    """按固定顺序拼装最终消息。

    顺序有意为之：文件说明在最前（它是本轮的事实前提），接着是引用材料，
    最后才是指令（技能、偏好、模式）——指令放最后是为了离模型最近，
    不容易被长材料淹没。
    """
    parts: list[str] = []

    uploads = uploads or []
    if len(uploads) == 1:
        filename, saved_path = uploads[0]
        parts.append(
            f"用户上传了文件「{filename}」，文件路径：{saved_path}。"
            f"请用 read_file 读取它，不要凭文件名猜测内容。"
        )
    elif uploads:
        listing = "\n".join(f"- 「{name}」→ {path}" for name, path in uploads)
        parts.append(
            f"用户上传了 {len(uploads)} 个文件：\n{listing}\n"
            "请用 read_file 逐个读取，不要凭文件名猜测内容，也不要在文件系统里搜索别处的资料。"
        )

    ref_instruction = prompts.build_ref_instruction(ref.get("reference"))
    if ref_instruction:
        parts.append(ref_instruction)

    if ref.get("forced_skill"):
        parts.append(prompts.build_skill_instruction(ref["forced_skill"]))

    pref_instruction = prompts.build_preference_instruction(
        chat_store.preferences_as_text(user_name)
    )
    if pref_instruction:
        parts.append(pref_instruction)

    parts.append(prompts.build_mode_instruction(mode))

    return f"{message}{''.join(parts)}"


# ============================================================
# 文件上传
# ============================================================

async def _save_upload(file: UploadFile, user_name: str, session_id: str) -> str:
    """把上传文件存进 agent 的沙箱目录，返回虚拟路径。

    存到 ``/workspace/uploads/`` 而不是别处，是因为权限规则里只有
    ``/workspace/**`` 是可写的——存到沙箱外 agent 自己都读不到。
    """
    raw_name = file.filename or "upload.bin"
    safe_name = _SAFE_NAME_RE.sub("_", Path(raw_name).name) or "upload.bin"
    from agents.workspace import session_root
    target_dir = session_root(user_name, session_id) / "workspace" / "uploads"
    target_dir.mkdir(parents=True, exist_ok=True)
    target = target_dir / f"{uuid.uuid4().hex}_{safe_name}"
    size = 0
    try:
        with target.open("xb") as output:
            while chunk := await file.read(64 * 1024):
                size += len(chunk)
                if size > MAX_UPLOAD_SIZE:
                    raise HTTPException(status_code=413, detail={"code": "file_too_large", "message": f"文件超过上限 {MAX_UPLOAD_MB}MB"})
                output.write(chunk)
        if not size:
            raise HTTPException(status_code=400, detail={"code": "empty_file", "message": "上传的文件为空"})
    except BaseException:
        target.unlink(missing_ok=True)
        raise
    logger.info("已保存上传文件 path=%s size=%d", target, size)
    return f"/workspace/uploads/{target.name}"


async def _save_uploads(files: list[UploadFile], user_name: str, session_id: str) -> list[tuple[str, str]]:
    """保存本次请求的全部附件，返回 [(原始文件名, 虚拟路径)]。

    校验顺序有意如此：先卡数量再落盘，避免误传目录时先写进去一堆文件再报错。
    """
    if len(files) > MAX_UPLOAD_FILES:
        raise HTTPException(status_code=400, detail={
            "code": "too_many_files",
            "message": f"一次最多上传 {MAX_UPLOAD_FILES} 个文件，本次 {len(files)} 个",
        })
    saved: list[tuple[str, str]] = []
    for item in files:
        saved.append((item.filename or "upload.bin", await _save_upload(item, user_name, session_id)))
    return saved


# ============================================================
# 主入口
# ============================================================

@router.post("/chat")
async def chat(
    message: str = Form(...),
    session_id: str | None = Form(default=None),
    thread_id: str | None = Form(default=None),
    mode: str | None = Form(default=None),
    file: UploadFile | None = File(default=None),
    files: list[UploadFile] = File(default_factory=list),
    user_name: str | None = Form(default=None),
    request: Request = None,
    request_id: str | None = Form(default=None),
):
    """统一对话接口（SSE 流式）。

    请求体：multipart/form-data
    - message：用户消息（必填），支持 ``#skill:<name>``、``#session:<uuid>`` 前缀
    - session_id：可选，会话 ID（缺省则新建）
    - thread_id：可选，必须与 session_id 一致，否则 404
    - mode：可选，权限模式（plan | full）
    - file：可选，单个附件（1.5 兼容字段，仍可用）
    - files：可选，多个附件（2.0 追加字段）。研判类任务一次要传多份资料，
      因此除兼容用的单份 file 外另开这一个；两者可同时使用，会合并处理
    - user_name：可选，归属用户，缺省 default（2.0 新增，替代 1.5 的鉴权依赖）
    """
    attachments = ([file] if file is not None else []) + list(files or [])
    return await _run_chat_pipeline(
        message=message,
        session_id=session_id,
        thread_id=thread_id,
        mode=mode,
        attachments=attachments,
        user_name=_resolve_user(user_name),
        trace_id=_trace_id_of(request),
        request_id=request_id,
    )


async def _run_chat_pipeline(*, request_id: str | None = None, **kwargs):
    """同一用户的请求键只接受一次；重试返回已有运行 ID。"""
    import asyncio
    import hashlib
    if request_id is None:
        return await _execute_chat_pipeline(**kwargs)
    if not request_id.strip() or len(request_id) > 128:
        raise HTTPException(status_code=400, detail={"code": "invalid_request_id", "message": "request_id 必须为 1 至 128 字符"})
    key = hashlib.sha256((kwargs["user_name"] + "\0" + request_id).encode()).hexdigest()
    with concurrency.session_guard(f"request:{key}"):
        existing = await asyncio.to_thread(db_conn.query_one, "SELECT * FROM request_keys WHERE id = ?", (key,))
        if existing:
            raise HTTPException(status_code=409, detail={"code": "request_already_processed", "message": "该请求已接受，请查看原会话和执行记录", "session_id": existing["session_id"], "run_id": existing["run_id"]})
        response = await _execute_chat_pipeline(**kwargs)
        try:
            await asyncio.to_thread(db_conn.execute,
                "INSERT INTO request_keys (id, session_id, run_id, created_at) VALUES (?, ?, ?, ?)",
                (key, response.headers["X-Session-Id"], response.headers["X-Run-Id"], db_conn.fmt_dt()))
        except BaseException:
            # 响应体尚未迭代时 aclose 不会运行生成器 finally。
            await response.background()
            raise
        return response


async def _execute_chat_pipeline(
    message: str,
    session_id: str | None,
    thread_id: str | None,
    mode: str | None,
    attachments: list[UploadFile] | None,
    user_name: str,
    trace_id: str = "",
):
    """先验证请求，再取得会话独占权，装配共用的执行流程。"""
    import asyncio
    if mode and mode.strip().lower() not in VALID_MODES:
        raise HTTPException(status_code=400, detail={"code": "invalid_mode", "message": "mode 只能是 plan | full"})
    if not message.strip():
        raise HTTPException(status_code=400, detail={"code": "empty_message", "message": "消息不能为空"})
    started_at = db_conn.now()
    # 新建会话前也预检，拒绝请求不会产生空会话。
    await asyncio.to_thread(guard.precheck, user_name)
    session = await asyncio.to_thread(_init_session, session_id, thread_id, user_name, mode or "full")
    admission = await execution.admit(user_name, session["id"])
    try:
        agent = await asyncio.to_thread(agent_registry.get_agent)
        graph_config = {"configurable": {"thread_id": session["id"]}}
        if await pending_interrupts(agent, graph_config):
            raise HTTPException(status_code=409, detail={"code": "pending_interrupt", "message": "请先处理当前待确认请求"})
        current_mode, graph_config = await asyncio.to_thread(_setup_mode, session, mode)
        await asyncio.to_thread(_ensure_session_title, session, message)
        saved_uploads = await _save_uploads(list(attachments or []), user_name, session["id"])
        ref = await asyncio.to_thread(_resolve_refs_and_skill, message, user_name)
        final_message = await asyncio.to_thread(_inject_all, ref["message"], user_name, current_mode, ref, saved_uploads)
        stream = await execution.build_execution(admission, agent=agent,
            session_id=session["id"], thread_id=session["id"], config=graph_config,
            message=final_message, user_name=user_name, kind="chat", model=config.MODEL,
            original_user_message=ref["original"] or message, trace_id=trace_id,
            started_at=started_at.timestamp())
        return StreamingResponse(stream_sse(stream), media_type="text/event-stream", background=BackgroundTask(admission.release),
                                 headers={**sse_headers(session["id"]), "X-Run-Id": stream.req.run_id})
    except BaseException:
        admission.release()
        raise


# ============================================================
# 中断恢复
# ============================================================

def _to_official_resume(payload: ResumeRequest, pending_count: int) -> dict[str, Any]:
    from api.interrupts import to_official_resume
    return to_official_resume(payload, pending_count)


@router.post("/chat/resume")
async def chat_resume(payload: ResumeRequest, request: Request = None):
    """在会话独占期间核对当前中断，再恢复图。"""
    import asyncio
    from api.interrupts import resume_value
    user_name = _resolve_user(payload.user_name)
    session = await asyncio.to_thread(_init_session, payload.session_id, payload.thread_id, user_name, "full")
    admission = await execution.admit(user_name, session["id"])
    try:
        agent = await asyncio.to_thread(agent_registry.get_agent)
        _, graph_config = _setup_mode(session, None)
        state = await agent.aget_state(graph_config)
        resume_payload = resume_value(payload, state)
        stream = await execution.build_execution(admission, agent=agent,
            session_id=session["id"], thread_id=session["id"], config=graph_config,
            message="", user_name=user_name, kind="resume", model=config.MODEL,
            trace_id=_trace_id_of(request), resume_payload=resume_payload)
        return StreamingResponse(stream_sse(stream), media_type="text/event-stream", background=BackgroundTask(admission.release),
                                 headers={**sse_headers(session["id"]), "X-Run-Id": stream.req.run_id})
    except BaseException:
        admission.release()
        raise


@router.get("/chat/interrupt")
async def get_interrupt(session_id: str, user_name: str | None = None):
    """查询会话当前挂起的中断。

    前端刷新页面后靠这个接口把提问卡/审批卡补渲染出来——否则用户会看到
    一个「什么都不动」的对话框，因为中断事件是流式的、刷新后就没了。

    有挂起返回 200 + 载荷；无挂起返回 204（无内容）。
    """
    from fastapi import Response

    user = _resolve_user(user_name)
    session = chat_store.get_session(session_id)
    if session is None or session.get("user_name") != user:
        raise HTTPException(
            status_code=404, detail={"code": "session_not_found", "message": "会话不存在"}
        )

    agent = agent_registry.get_agent()
    graph_config = {"configurable": {"thread_id": session_id}}
    pending = await pending_interrupts(agent, graph_config)
    if not pending:
        return Response(status_code=204)

    from runtime.interrupts import interrupt_event
    from runtime.events import EventBuilder
    writer = EventBuilder(session_id, encode=lambda event: event)
    frames = [interrupt_event(writer, item) for item in pending]
    return {"session_id": session_id, "interrupts": frames}