hunjing/backend/app/api/v1/endpoints/runs.py

"""Run 与 SSE API（Run 的启动契约见 kernel_runs.py）。"""

import asyncio
import json
from collections.abc import AsyncGenerator
from typing import Any

from fastapi import APIRouter, Header, Query
from sse_starlette.sse import EventSourceResponse

from app.core.config import settings
from app.core.errors import AppError
from app.core.time import to_iso, utc_now
from app.db.session import AsyncSessionLocal
from app.repositories import runtime as runtime_repo
from app.schemas.common import Page
from app.schemas.runtime import EventResponse, RunResponse
from app.services.event_bus import bus
from app.services.runtime import (
    cancel_run,
    get_run_response,
    list_run_event_responses,
    list_runs_response,
)

router = APIRouter(prefix="/runs", tags=["runs"])


@router.get("", response_model=Page[RunResponse])
async def list_runs(
    character_id: str | None = Query(default=None, description="按人物过滤。"),
    status: str | None = Query(default=None, description="按状态过滤。"),
    page: int = Query(default=1, ge=1),
    page_size: int = Query(default=20, ge=1, le=100),
) -> Page[RunResponse]:
    """Run 列表（按创建时间倒序）；归档页运行时间河数据源。"""
    return await list_runs_response(
        character_id=character_id,
        status=status,
        page=page,
        page_size=page_size,
    )


@router.get("/{run_id}", response_model=RunResponse)
async def get_run(run_id: str) -> RunResponse:
    return await get_run_response(run_id)


@router.get("/{run_id}/events", response_model=list[EventResponse])
async def list_run_events(
    run_id: str,
    limit: int = Query(default=200, ge=1, le=500),
    offset: int = Query(default=0, ge=0),
) -> list[EventResponse]:
    """Run 的事件时间线（按 sequence 升序）。"""
    return await list_run_event_responses(run_id, limit=limit, offset=offset)


@router.post("/{run_id}/cancel", response_model=RunResponse)
async def cancel_run_api(run_id: str) -> RunResponse:
    return await cancel_run(run_id)


@router.get("/{run_id}/stream")
async def stream_run(
    run_id: str,
    last_event_id: str | None = Header(default=None, alias="Last-Event-ID"),
) -> EventSourceResponse:
    """Run 事件流（SSE）：按 run_id 实时订阅事件增量。

    事件通过进程内 ``bus`` 实时推送（``append_event`` 自动 publish），
    DB 持久化事件供断线重连（``Last-Event-ID`` 回放）。
    """
    if not settings.enable_sse:
        raise AppError("invalid_request", "SSE 当前未启用。")
    async with AsyncSessionLocal() as session:
        run = await runtime_repo.get_run(session, run_id)
    if run is None:
        raise AppError("not_found", "Run 不存在。")
    return EventSourceResponse(_event_generator(run_id, last_event_id))


async def _event_generator(
    run_id: str, last_event_id: str | None
) -> AsyncGenerator[dict[str, Any], None]:
    # 断线重连：解析 Last-Event-ID 对应的 sequence 作为回放起点
    after_sequence = 0
    if last_event_id:
        async with AsyncSessionLocal() as session:
            previous_event = await runtime_repo.get_event(session, last_event_id)
            if previous_event and previous_event.run_id == run_id:
                after_sequence = previous_event.sequence

    last_sequence = after_sequence
    # 先订阅再回放，避免回放与订阅之间的事件落入间隙
    queue = bus.subscribe(run_id)
    try:
        for event in bus.list_after(run_id, after_sequence=last_sequence):
            last_sequence = event.sequence
            yield {
                "id": event.event_id,
                "event": event.type,
                "data": event.model_dump_json(),
            }
            if bus.is_terminal(event):
                return

        while True:
            try:
                event = await asyncio.wait_for(queue.get(), timeout=settings.sse_heartbeat_seconds)
            except TimeoutError:
                yield {
                    "event": "heartbeat",
                    "data": json.dumps(
                        {
                            "event_id": "heartbeat",
                            "run_id": run_id,
                            "type": "heartbeat",
                            "sequence": last_sequence,
                            "visibility": "public",
                            "payload": {},
                            "created_at": to_iso(utc_now()),
                        },
                        ensure_ascii=False,
                    ),
                }
                continue
            if event.sequence <= last_sequence:
                continue  # 回放阶段已转发，去重
            last_sequence = event.sequence
            yield {
                "id": event.event_id,
                "event": event.type,
                "data": event.model_dump_json(),
            }
            if bus.is_terminal(event):
                return
    finally:
        bus.unsubscribe(run_id, queue)