Mutual_Fund/api/chat/client_agent.py

"""Logged-in client Agent HTTP endpoints.

Authentication and customer-owned session enforcement are added in C2.
"""

from __future__ import annotations

import json

from fastapi import APIRouter, Depends, Request
from fastapi.responses import JSONResponse, StreamingResponse

from agent.client_agent.session import SessionOwnershipError
from api.deps import require_customer
from model.sys_user import SysUser
from service.customer_agent.chat import QueryTooLongError
from utils.request_id import get_request_id, new_request_id
from utils.response import fail, success


router = APIRouter()


def _runtime(request: Request):
    """Return the client Agent runtime installed by the application lifespan."""
    runtime = getattr(request.app.state, "client_agent_runtime", None)
    if runtime is None:
        raise RuntimeError("client agent runtime is not configured")
    return runtime


@router.post("/session/create")
async def create_session(
    request: Request,
    user: SysUser = Depends(require_customer),
):
    """Create a client Agent session during the architecture bootstrap phase."""
    runtime = _runtime(request)
    session_id = await runtime.session_service.create_session(user.id)
    return success({"session_id": session_id, "customer_id": user.id})


@router.post("/chat")
async def chat(
    request: Request,
    body: dict,
    user: SysUser = Depends(require_customer),
):
    """Reuse the existing customer Agent chat orchestration."""
    session_id = body.get("session_id", "")
    query = body.get("query", "")
    runtime = _runtime(request)
    trace_id = request.headers.get("X-Trace-Id") or get_request_id() or new_request_id()
    try:
        await runtime.session_service.verify_session_ownership(
            session_id, customer_id=user.id
        )
        retry_after = await runtime.session_service.consume_chat_quota(
            session_id, customer_id=user.id
        )
        if retry_after is not None:
            response = fail(429, "请求过于频繁", {"retry_after": retry_after})
            return JSONResponse(
                status_code=429,
                headers={"Retry-After": str(retry_after)},
                content=response.model_dump(),
            )
        result = await runtime.agent.handle(
            session_id,
            query,
            trace_id=trace_id,
            customer_id=user.id,
        )
    except SessionOwnershipError as exc:
        return JSONResponse(
            status_code=exc.code,
            content=fail(exc.code, exc.message).model_dump(),
        )
    except QueryTooLongError as exc:
        return JSONResponse(
            status_code=400,
            content=fail(400, str(exc)).model_dump(),
        )

    async def events():
        yield f"data: {json.dumps(result, ensure_ascii=False)}\n\n"

    return StreamingResponse(
        events(), media_type="text/event-stream", headers={"X-Trace-Id": trace_id}
    )


@router.post("/session/end")
async def end_session(
    request: Request,
    body: dict,
    user: SysUser = Depends(require_customer),
):
    """Close the client Agent session during the architecture bootstrap phase."""
    runtime = _runtime(request)
    session_id = body.get("session_id", "")
    try:
        await runtime.session_service.verify_session_ownership(
            session_id, customer_id=user.id
        )
    except SessionOwnershipError as exc:
        return JSONResponse(
            status_code=exc.code,
            content=fail(exc.code, exc.message).model_dump(),
        )
    warnings = await runtime.memory_service.close_session(
        customer_id=user.id,
        session_id=session_id,
    )
    if warnings:
        return success({"session_id": session_id, "archived": False, "warnings": warnings})
    if hasattr(runtime.redis, "delete"):
        await runtime.redis.delete(
            f"session:{session_id}",
            f"session:{session_id}:messages",
            f"rate:limit:client:{user.id}:{session_id}:chat",
        )
    return success({"session_id": session_id, "archived": True, "warnings": []})