social-app/backend/tests/integration/v1/agent/test_sse_flow_live.py

from __future__ import annotations

import os
from datetime import datetime, timedelta, timezone
from uuid import UUID

import httpx
import jwt
import pytest
from sqlalchemy import select

from core.config import config
from core.db.session import AsyncSessionLocal
from models.agent_chat_message import AgentChatMessage
from models.agent_chat_session import AgentChatSession
from models.profile import Profile

BASE_URL = os.getenv("AGENT_LIVE_BASE_URL", "http://localhost:5775")


async def _owner_id() -> UUID:
    async with AsyncSessionLocal() as session:
        owner_id = (await session.execute(select(Profile.id).limit(1))).scalar_one_or_none()
        if owner_id is None:
            pytest.skip("profile owner not found")
        return owner_id


def _jwt_for(user_id: UUID) -> str:
    secret = config.supabase.jwt_secret
    if not secret:
        pytest.skip("JWT secret not configured")
    issuer = f"{config.supabase.public_url.rstrip('/')}/auth/v1"
    payload = {
        "sub": str(user_id),
        "role": "authenticated",
        "aud": "authenticated",
        "iss": issuer,
        "iat": datetime.now(timezone.utc),
        "exp": datetime.now(timezone.utc) + timedelta(minutes=30),
    }
    return jwt.encode(payload, secret, algorithm="HS256")


@pytest.mark.asyncio
@pytest.mark.live
async def test_agent_sse_closed_loop_live() -> None:
    if os.getenv("AGENT_LIVE_INTEGRATION") != "1":
        pytest.skip("set AGENT_LIVE_INTEGRATION=1 to run live integration test")

    owner_id = await _owner_id()
    token = _jwt_for(owner_id)
    headers = {"Authorization": f"Bearer {token}"}

    async with httpx.AsyncClient(timeout=30.0) as client:
        run_resp = await client.post(
            f"{BASE_URL}/api/v1/agent/runs",
            headers=headers,
            json={"prompt": "请用一句话介绍你自己"},
        )
        assert run_resp.status_code == 202

        accepted = run_resp.json()
        session_id = str(accepted["session_id"])
        assert session_id

        events_url = f"{BASE_URL}/api/v1/agent/runs/{session_id}/events"
        event_names: list[str] = []
        async with client.stream("GET", events_url, headers=headers, timeout=20.0) as sse_resp:
            assert sse_resp.status_code == 200
            assert sse_resp.headers.get("content-type", "").startswith("text/event-stream")
            async for line in sse_resp.aiter_lines():
                if line.startswith("event:"):
                    event_names.append(line.split(":", 1)[1].strip())

        assert "RUN_STARTED" in event_names
        assert "RUN_FINISHED" in event_names or "RUN_ERROR" in event_names

    async with AsyncSessionLocal() as session:
        session_row = await session.get(AgentChatSession, UUID(session_id))
        assert session_row is not None
        assert session_row.message_count >= 1
        assert session_row.total_tokens >= 0
        assert session_row.total_cost >= 0

        rows = await session.execute(
            select(AgentChatMessage).where(AgentChatMessage.session_id == UUID(session_id))
        )
        assert len(list(rows.scalars().all())) >= 1