group_fqcd_jr/app/api/controllers/agent_runs.py

from collections.abc import AsyncIterator
from dataclasses import asdict

from fastapi import APIRouter, Depends, Request, status
from sqlalchemy.ext.asyncio import AsyncSession
from starlette.responses import StreamingResponse

from app.api.dependencies.auth import build_request_context
from app.api.dependencies.database import get_session
from app.api.dependencies.rate_limit import enforce_rate_limit
from app.api.schemas.agent_runs import (
    AgentRunAcceptedEnvelope,
    AgentRunAcceptedResponse,
    AgentRunCreateRequest,
    AgentRunStatusEnvelope,
    AgentRunStatusResponse,
)
from app.api.views.agent_run_sse import encode_events, recovery_events
from app.core.config import get_settings
from app.core.contracts import AgentRequest, RequestContext
from app.core.errors import SseNotAcceptableError
from app.service.agent_run_application_service import AgentRunApplicationService
from app.service.run_query_service import RunQueryService

router = APIRouter(prefix="/api/v1/agent-runs", tags=["agent-runs"],
                   dependencies=[Depends(enforce_rate_limit)])

SSE_MEDIA_TYPE = "text/event-stream"


def accepts_event_stream(accept: str | None) -> bool:
    """`Accept` 是否接受 `text/event-stream`（文档 §3.2：该头**非必填**）。

    - 未携带（`None` 或空串）→ 放行：文档写明"默认 `application/json`；SSE 为
      `text/event-stream`"，即由接口自身决定响应类型，不是客户端错误；
    - 携带 `text/event-stream`、`text/*` 或 `*/*` 且 `q != 0` → 放行；
    - 显式携带但只接受其他类型（如 `application/json`）→ 拒绝，由调用方转
      `406 SSE_NOT_ACCEPTABLE`（文档 §3.5/§6.4）。

    只做"是否可接受"的判定，不参与内容协商排序：SSE 端点只有一种表示。
    """
    if accept is None or not accept.strip():
        return True
    for entry in accept.split(","):
        parts = entry.split(";")
        media_type = parts[0].strip().lower()
        if media_type not in {SSE_MEDIA_TYPE, "text/*", "*/*"}:
            continue
        quality = 1.0
        for parameter in parts[1:]:
            name, _, value = parameter.partition("=")
            if name.strip().lower() == "q":
                try:
                    quality = float(value.strip())
                except ValueError:
                    quality = 0.0
        if quality > 0:
            return True
    return False


@router.post(
    "",
    response_model=AgentRunAcceptedEnvelope,
    status_code=status.HTTP_202_ACCEPTED,
)
async def create_agent_run(
    payload: AgentRunCreateRequest,
    request: Request,
    context: RequestContext = Depends(build_request_context),  # noqa: B008
    session: AsyncSession = Depends(get_session),  # noqa: B008
) -> AgentRunAcceptedEnvelope:
    request.state.request_context = context
    accepted = await AgentRunApplicationService(session).accept(
        AgentRequest(**payload.model_dump()), context)
    return AgentRunAcceptedEnvelope(
        data=AgentRunAcceptedResponse(
            run_id=accepted.run_id, trace_id=accepted.trace_id, status=accepted.status,
            status_url=f"/api/v1/agent-runs/{accepted.run_id}",
            events_url=f"/api/v1/agent-runs/{accepted.run_id}/events",
        ),
        meta={"trace_id": context.trace_id},
    )


@router.get("/{run_id}", response_model=AgentRunStatusEnvelope)
async def get_agent_run(
    run_id: str, context: RequestContext = Depends(build_request_context),  # noqa: B008
) -> AgentRunStatusEnvelope:
    """查询运行（文档 §6.3）。

    文档 §3.3 与 §6.3 都把成功响应定义为 `{data, meta:{trace_id}}` 信封；此前这里
    直接返回资源对象，客户端必须为这一个接口特判。**只改包装结构**：`data` 内的字段名
    与语义保持原样，`meta.trace_id` 用本次请求的 trace（`data.trace_id` 仍是运行自身的
    追踪标识，两者语义不同，不能互相替代）。
    """
    snapshot = await RunQueryService().get(run_id, context)
    return AgentRunStatusEnvelope(
        data=AgentRunStatusResponse(**asdict(snapshot)),
        meta={"trace_id": context.trace_id},
    )


@router.get("/{run_id}/events")
async def stream_agent_run_events(
    run_id: str,
    request: Request,
    context: RequestContext = Depends(build_request_context),  # noqa: B008
) -> StreamingResponse:
    query = RunQueryService()
    # 顺序按文档 §6.4 的主要错误列举：RUN_NOT_FOUND（含 AGENT_PERMISSION_DENIED 同级的
    # 可见性判定）在前、SSE_NOT_ACCEPTABLE 在后。可见性先行（auth 依赖已先于本函数执行）
    # 才能保证"运行是否存在"不因 Accept 头而异：否则用任意 run_id + 非法 Accept 探测，
    # 406 与 404 的差异就等价于一次存在性枚举。
    initial = await query.get(run_id, context)
    if not accepts_event_stream(request.headers.get("Accept")):
        raise SseNotAcceptableError("Accept 必须接受 text/event-stream")

    async def generate() -> AsyncIterator[str]:
        start_sent = False
        async for snapshot in query.watch(initial, context):
            if snapshot is None:
                yield ": heartbeat\n\n"
                continue
            result = snapshot.result or {}
            events = recovery_events(
                run_id=snapshot.run_id, trace_id=snapshot.trace_id, status=snapshot.status,
                error_code=snapshot.error_code, content=result.get("content"),
                tool_calls=result.get("tool_calls"),
                replay=initial.status in {"succeeded", "failed", "cancelled"},
                chunk_size=get_settings().sse_chunk_characters,
            )
            for encoded in encode_events(run_id, events[1:] if start_sent else events):
                yield encoded
            start_sent = True

    return StreamingResponse(generate(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache"})