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。
"""

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

import config
from agents import agent_registry, prompts
from api import concurrency
from api.models import ResumeRequest
from api.sse_core import ChatStream, StreamRequest, pending_interrupts, 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

_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],
    filename: str | None,
    saved_path: str | None,
) -> str:
    """按固定顺序拼装最终消息。

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

    if filename and saved_path:
        parts.append(
            f"用户上传了文件「{filename}」，文件路径：{saved_path}。"
            f"请用 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) -> 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"
    payload = await file.read()

    if len(payload) > MAX_UPLOAD_SIZE:
        raise HTTPException(
            status_code=413,
            detail={
                "code": "file_too_large",
                "message": f"文件超过上限 {MAX_UPLOAD_MB}MB",
                "max_size_mb": MAX_UPLOAD_MB,
            },
        )
    if not payload:
        raise HTTPException(
            status_code=400, detail={"code": "empty_file", "message": "上传的文件为空"}
        )

    target_dir = config.ROOT / "workspace" / "uploads"
    target_dir.mkdir(parents=True, exist_ok=True)
    target = target_dir / f"{uuid.uuid4().hex[:8]}_{safe_name}"
    target.write_bytes(payload)
    logger.info("已保存上传文件 path=%s size=%d", target, len(payload))
    return f"/workspace/uploads/{target.name}"


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

@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),
    user_name: str | None = Form(default=None),
    request: Request = None,
):
    """统一对话接口（SSE 流式）。

    请求体：multipart/form-data
    - message：用户消息（必填），支持 ``#skill:<name>``、``#session:<uuid>`` 前缀
    - session_id：可选，会话 ID（缺省则新建）
    - thread_id：可选，必须与 session_id 一致，否则 404
    - mode：可选，权限模式（plan | full）
    - file：可选，附件
    - user_name：可选，归属用户，缺省 default（2.0 新增，替代 1.5 的鉴权依赖）
    """
    return await _run_chat_pipeline(
        message=message,
        session_id=session_id,
        thread_id=thread_id,
        mode=mode,
        file=file,
        user_name=_resolve_user(user_name),
        billing_kind="chat",
        trace_id=_trace_id_of(request),
    )


@router.post("/chat/only")
async def chat_only(
    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),
    user_name: str | None = Form(default=None),
    request: Request = None,
):
    """纯对话入口。管线与 /chat 完全相同，只有计费类型不同（便于分账）。"""
    return await _run_chat_pipeline(
        message=message,
        session_id=session_id,
        thread_id=thread_id,
        mode=mode,
        file=file,
        user_name=_resolve_user(user_name),
        billing_kind="chat_only",
        trace_id=_trace_id_of(request),
    )


async def _run_chat_pipeline(
    message: str,
    session_id: str | None,
    thread_id: str | None,
    mode: str | None,
    file: UploadFile | None,
    user_name: str,
    billing_kind: str,
    trace_id: str = "",
):
    """/chat 与 /chat/only 的公共管线。

    步骤顺序与原由：
    0. 额度预检 —— 在开任何资源之前 fail-closed 拦掉超限用户
    1. 会话初始化 —— 校验归属、确定 session_id（它同时是 thread_id）
    2. 模式 —— 落库并算出 recursion_limit
    3. 标题 —— 首轮才生成
    4. 上传落盘 —— 存进 agent 可读的沙箱
    5. 消息拼装 —— 文件说明 + 引用 + 技能 + 偏好 + 模式指令
    6. 并发闸 + SSE 装配
    """
    started_at = db_conn.now()
    usage_day = db_conn.day_key(started_at)

    # 0) 额度预检
    try:
        guard.precheck(user_name, usage_day)
    except guard.QuotaExceededError as exc:
        raise HTTPException(
            status_code=429,
            detail={
                "code": "quota_exceeded",
                "message": str(exc),
                "limit_tokens": exc.quota.get("limit_tokens", 0),
                "used_tokens": exc.quota.get("used_tokens", 0),
            },
        ) from exc
    except guard.BillingConfigError as exc:
        # fail-closed：计费不可用时宁可拒绝服务，也不放行无法计费的请求
        raise HTTPException(
            status_code=500,
            detail={"code": "billing_unavailable", "message": str(exc)},
        ) from exc

    # 1) 会话
    session = _init_session(session_id, thread_id, user_name, mode or "full")

    # 2) 模式
    current_mode, graph_config = _setup_mode(session, mode)

    # 3) 标题
    _ensure_session_title(session, message)

    # 4) 上传
    filename = file.filename if file else None
    saved_path = await _save_upload(file) if file is not None else None

    # 5) 消息拼装
    ref = _resolve_refs_and_skill(message, user_name)
    final_message = _inject_all(
        ref["message"], user_name, current_mode, ref, filename, saved_path
    )

    # 6) 并发闸 + SSE
    slot = concurrency.acquire_task_slot(user_name)
    try:
        agent = agent_registry.get_agent()
        stream = ChatStream(
            StreamRequest(
                agent=agent,
                session_id=session["id"],
                thread_id=session["id"],
                config=graph_config,
                message=final_message,
                user_name=user_name,
                kind=billing_kind,
                model=config.MODEL,
                original_user_message=ref["original"] or message,
                trace_id=trace_id,
                slot=slot,
                started_at=started_at.timestamp(),
            )
        )
        return StreamingResponse(
            stream.generate(),
            media_type="text/event-stream",
            headers=sse_headers(session["id"]),
        )
    except HTTPException:
        slot.release()
        raise
    except Exception:
        # StreamingResponse 构造失败时生成器没机会跑 finally，这里兜底释放
        slot.release()
        raise


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

def _to_official_resume(payload: ResumeRequest, pending_count: int) -> dict[str, Any]:
    """把 1.5 的 ResumeRequest 翻译成官方 HITL 的 ``{"decisions": [...]}``。

    官方要求 decisions 与 action_requests **数量严格一一对应**，数量不符会
    抛错，所以这里按 ``pending_count`` 复制同一决策。

    翻译表：
    - ``approved=True`` → ``approve``
    - ``approved=False`` → ``reject``（带拒绝理由）
    - ``status=cancelled`` → ``reject``（用户取消等价于不执行）
    - ``status=error`` → ``reject``（前端恢复出错，不执行更安全）
    - 有 ``answers`` → ``respond``（把人的回答作为工具结果）
    """
    count = max(pending_count, 1)

    if payload.approved is True:
        decision: dict[str, Any] = {"type": "approve"}
    elif payload.approved is False:
        decision = {"type": "reject", "message": "用户拒绝了该操作，请勿重试，改为说明情况"}
    elif payload.status == "cancelled":
        decision = {"type": "reject", "message": "用户取消了该操作"}
    elif payload.status == "error":
        decision = {
            "type": "reject",
            "message": f"前端恢复出错，该操作未执行：{payload.error or '未提供详情'}",
        }
    elif payload.answers:
        decision = {"type": "respond", "message": "\n".join(payload.answers)}
    else:
        decision = {"type": "reject", "message": "未收到明确答复，该操作未执行"}

    return {"decisions": [dict(decision) for _ in range(count)]}


@router.post("/chat/resume")
async def chat_resume(payload: ResumeRequest, request: Request = None):
    """恢复被 interrupt 挂起的对话。"""
    user_name = _resolve_user(payload.user_name)
    session = _init_session(payload.session_id, payload.thread_id, user_name, "full")
    thread_id = session["id"]

    agent = agent_registry.get_agent()
    graph_config = {
        "configurable": {"thread_id": thread_id},
        "recursion_limit": RECURSION_LIMIT_PLAN
        if (session.get("mode") == "plan")
        else RECURSION_LIMIT_DEFAULT,
    }

    pending = await pending_interrupts(agent, graph_config)
    if not pending:
        # 没有挂起的中断却来恢复：多半是重复点击或刷新后的陈旧请求。
        # 409 而不是 500——这不是服务端故障。
        raise HTTPException(
            status_code=409,
            detail={"code": "no_pending_interrupt", "message": "该会话当前没有待处理的确认请求"},
        )

    action_count = 0
    for item in pending:
        if isinstance(item, dict) and isinstance(item.get("action_requests"), list):
            action_count = max(action_count, len(item["action_requests"]))

    resume_payload = _to_official_resume(payload, action_count)
    logger.info(
        "恢复中断 session=%s 决策=%s pending=%d",
        thread_id, resume_payload["decisions"], len(pending),
    )

    slot = concurrency.acquire_task_slot(user_name)
    try:
        stream = ChatStream(
            StreamRequest(
                agent=agent,
                session_id=thread_id,
                thread_id=thread_id,
                config=graph_config,
                message="",
                user_name=user_name,
                kind="resume",
                model=config.MODEL,
                trace_id=_trace_id_of(request),
                slot=slot,
                started_at=db_conn.now().timestamp(),
                resume_payload=resume_payload,
            )
        )
        return StreamingResponse(
            stream.generate(),
            media_type="text/event-stream",
            headers=sse_headers(thread_id),
        )
    except HTTPException:
        slot.release()
        raise
    except Exception:
        slot.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 api.sse_events import EventWriter

    writer = EventWriter(session_id)
    frames: list[dict[str, Any]] = []
    for item in pending:
        if isinstance(item, dict) and isinstance(item.get("action_requests"), list):
            for action in item["action_requests"]:
                frames.append(
                    {
                        "type": "interrupt",
                        "session_id": session_id,
                        "tool_approval": {
                            "tool": action.get("name"),
                            "args": action.get("args"),
                            "reason": action.get("description") or "",
                            "side_effect": "该操作需人工确认后才会执行",
                        },
                    }
                )
        else:
            frames.append({"type": "interrupt", "session_id": session_id, "message": "智能体暂停，等待输入"})
    return {"session_id": session_id, "interrupts": frames, "raw": pending}