social-app/backend/src/v1/agent/router.py

from __future__ import annotations

from collections.abc import AsyncIterator
import asyncio
from datetime import date
import re
import time
from typing import Annotated, Union

from ag_ui.core import RunAgentInput
from fastapi import APIRouter, Depends, Header, Query, Request, status, UploadFile
from fastapi import HTTPException
from fastapi.responses import JSONResponse, StreamingResponse

from core.agent.infrastructure.agui.stream import to_sse_event
from core.agent.domain.agui_input import (
    parse_run_input,
    validate_run_request_messages_contract,
)
from core.auth.models import CurrentUser
from services.base.redis import get_or_init_redis_client
from v1.agent.dependencies import get_agent_service
from v1.agent.schemas import AsrTranscribeResponse, TaskAcceptedResponse
from v1.agent.service import AgentService, asr_service
from v1.users.dependencies import get_current_user

router = APIRouter(prefix="/agent", tags=["agent"])
_LAST_EVENT_ID_RE = re.compile(r"^\d+-\d+$")
_RUNS_PER_MINUTE = 30
_MAX_SSE_CONNECTIONS_PER_USER = 3
_SSE_SLOT_TTL_SECONDS = 15 * 60


async def _allow_run_request(*, user_id: str) -> bool:
    try:
        redis = await get_or_init_redis_client()
        minute_bucket = int(time.time() // 60)
        key = f"agent:run-rate:{user_id}:{minute_bucket}"
        count = await redis.incr(key)
        if count == 1:
            await redis.expire(key, 70)
        return int(count) <= _RUNS_PER_MINUTE
    except Exception:  # noqa: BLE001
        return False


async def _acquire_sse_slot(*, user_id: str) -> bool:
    try:
        redis = await get_or_init_redis_client()
        key = f"agent:sse-active:{user_id}"
        count = await redis.incr(key)
        if count == 1:
            await redis.expire(key, _SSE_SLOT_TTL_SECONDS)
        if int(count) > _MAX_SSE_CONNECTIONS_PER_USER:
            await redis.decr(key)
            return False
        return True
    except Exception:  # noqa: BLE001
        return False


async def _release_sse_slot(*, user_id: str) -> None:
    try:
        redis = await get_or_init_redis_client()
        key = f"agent:sse-active:{user_id}"
        count = await redis.decr(key)
        if int(count) <= 0:
            await redis.delete(key)
    except Exception:  # noqa: BLE001
        return None


@router.post(
    "/runs", response_model=TaskAcceptedResponse, status_code=status.HTTP_202_ACCEPTED
)
async def enqueue_run(
    request: RunAgentInput,
    service: Annotated[AgentService, Depends(get_agent_service)],
    current_user: Annotated[CurrentUser, Depends(get_current_user)],
) -> TaskAcceptedResponse:
    try:
        normalized = parse_run_input(request.model_dump(mode="json", by_alias=True))
        validate_run_request_messages_contract(normalized)
    except ValueError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc
    allowed = await _allow_run_request(user_id=str(current_user.id))
    if not allowed:
        raise HTTPException(status_code=429, detail="Too many run requests")

    task = await service.enqueue_run(
        run_input=request,
        current_user=current_user,
    )
    return TaskAcceptedResponse(
        taskId=task.task_id,
        threadId=task.thread_id,
        runId=task.run_id,
        created=task.created,
    )


@router.post(
    "/runs/{thread_id}/resume",
    response_model=TaskAcceptedResponse,
    status_code=status.HTTP_202_ACCEPTED,
)
async def enqueue_resume(
    thread_id: str,
    request: RunAgentInput,
    service: Annotated[AgentService, Depends(get_agent_service)],
    current_user: Annotated[CurrentUser, Depends(get_current_user)],
) -> TaskAcceptedResponse:
    if request.thread_id != thread_id:
        raise HTTPException(status_code=422, detail="thread_id path/body mismatch")
    try:
        parse_run_input(request.model_dump(mode="json", by_alias=True))
    except ValueError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc
    task = await service.enqueue_resume(
        thread_id=thread_id,
        run_input=request,
        current_user=current_user,
    )
    return TaskAcceptedResponse(
        taskId=task.task_id,
        threadId=task.thread_id,
        runId=task.run_id,
        created=task.created,
    )


@router.get("/runs/{thread_id}/events")
async def stream_events(
    request: Request,
    thread_id: str,
    service: Annotated[AgentService, Depends(get_agent_service)],
    current_user: Annotated[CurrentUser, Depends(get_current_user)],
    last_event_id: str | None = Header(default=None, alias="Last-Event-ID"),
    idle_limit: int = Query(default=300, ge=1, le=3600),
) -> StreamingResponse:
    if last_event_id is not None and (
        len(last_event_id) > 32 or _LAST_EVENT_ID_RE.fullmatch(last_event_id) is None
    ):
        raise HTTPException(status_code=422, detail="Invalid Last-Event-ID")

    sse_slot_acquired = await _acquire_sse_slot(user_id=str(current_user.id))
    if not sse_slot_acquired:
        raise HTTPException(status_code=429, detail="Too many SSE connections")

    async def _event_iter() -> AsyncIterator[str]:
        cursor = last_event_id
        idle_polls = 0
        try:
            while not await request.is_disconnected() and idle_polls < idle_limit:
                rows = await service.stream_events(
                    thread_id=thread_id,
                    last_event_id=cursor,
                    current_user=current_user,
                )
                if not rows:
                    idle_polls += 1
                    yield ": keep-alive\n\n"
                    await asyncio.sleep(0.2)
                    continue

                idle_polls = 0
                for row in rows:
                    row_id = str(row.get("id", ""))
                    event = row.get("event")
                    if not row_id or not isinstance(event, dict):
                        continue
                    cursor = row_id
                    yield to_sse_event(row_id, event)
        finally:
            await _release_sse_slot(user_id=str(current_user.id))

    return StreamingResponse(
        _event_iter(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@router.get("/runs/{thread_id}/history")
async def get_history_snapshot(
    thread_id: str,
    service: Annotated[AgentService, Depends(get_agent_service)],
    current_user: Annotated[CurrentUser, Depends(get_current_user)],
    before: date | None = Query(default=None),
) -> dict[str, object]:
    return await service.get_history_snapshot(
        thread_id=thread_id,
        before=before,
        current_user=current_user,
    )


@router.get("/history")
async def get_user_history_snapshot(
    service: Annotated[AgentService, Depends(get_agent_service)],
    current_user: Annotated[CurrentUser, Depends(get_current_user)],
    thread_id: str | None = Query(default=None, alias="threadId"),
    before: date | None = Query(default=None),
) -> dict[str, object]:
    return await service.get_user_history_snapshot(
        current_user=current_user,
        thread_id=thread_id,
        before=before,
    )


@router.post(
    "/transcribe",
    response_model=AsrTranscribeResponse,
    status_code=status.HTTP_200_OK,
)
async def transcribe(
    audio: UploadFile,
    current_user: Annotated[CurrentUser, Depends(get_current_user)],
) -> Union[AsrTranscribeResponse, JSONResponse]:
    try:
        audio_data = await audio.read()
        if not audio_data:
            raise ValueError("Empty audio file")

        transcript = await asr_service.transcribe(
            audio_data, audio.filename or "unknown"
        )

        return AsrTranscribeResponse(transcript=transcript)

    except ValueError as exc:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={"detail": str(exc)},
        )
    except RuntimeError as exc:
        return JSONResponse(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            content={"detail": str(exc)},
        )