Aether/tests/api/internal_gateway/test_gateway_trace.py

"""Gateway internal report/finalize/trace route tests."""

import asyncio
import base64
import json
import time
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock

import pytest
from fastapi import BackgroundTasks, FastAPI
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from src.api.internal.gateway import (
    CONTROL_ACTION_HEADER,
    CONTROL_ACTION_PROXY_PUBLIC,
    CONTROL_EXECUTED_HEADER,
    GatewayAuthContext,
    GatewayExecutionDecisionResponse,
    GatewayExecuteRequest,
    GatewayExecutionPlanResponse,
    GatewayResolveRequest,
    GatewayStreamReportRequest,
    GatewaySyncReportRequest,
    _dispatch_gateway_sync_telemetry,
    _build_gateway_sync_error_payload,
    _build_gateway_sync_telemetry_writer,
    _run_gateway_stream_report_background,
    _run_gateway_sync_report_background,
    _stream_executor_requires_python_rewrite,
    _build_claude_chat_sync_decision,
    _build_claude_chat_stream_decision,
    _build_claude_cli_sync_decision,
    _build_claude_cli_stream_decision,
    _build_gemini_files_download_stream_decision,
    _build_gemini_files_proxy_sync_decision,
    _build_gemini_chat_sync_decision,
    _build_gemini_chat_stream_decision,
    _build_gemini_cli_sync_decision,
    _build_gemini_cli_stream_decision,
    _build_openai_chat_sync_decision,
    _build_openai_chat_stream_decision,
    _build_openai_cli_stream_decision,
    _build_openai_video_content_stream_decision,
    _extract_gateway_sync_error_message,
    _record_gateway_direct_candidate_graph,
    _resolve_gateway_sync_error_status_code,
    _build_openai_chat_sync_plan,
    _build_openai_cli_sync_decision,
    _build_openai_cli_stream_plan,
    _build_openai_cli_sync_plan,
    _is_streaming_sync_payload,
    _resolve_auth_context,
    _resolve_gateway_sync_adapter,
    classify_gateway_route,
    router,
)
from src.database import get_db
from src.models.database import Base, RequestCandidate
from src.services.orchestration.candidate_resolver import CandidateResolver
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
from src.services.request.execution_runtime_plan import ExecutionPlan, ExecutionPlanBody, PreparedExecutionPlan


def _wait_until(predicate: Any, *, timeout: float = 1.0, interval: float = 0.01) -> None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        if predicate():
            return
        time.sleep(interval)
    assert predicate()


def test_extract_gateway_sync_error_message_prefers_nested_error_message() -> None:
    payload = GatewaySyncReportRequest(
        trace_id="trace-error-message-123",
        report_kind="openai_chat_sync_finalize",
        status_code=429,
        headers={"content-type": "application/json"},
        body_json={"error": {"message": "rate limited"}},
    )

    assert _extract_gateway_sync_error_message(payload) == "rate limited"


def test_extract_gateway_sync_error_message_falls_back_to_body_base64() -> None:
    payload = GatewaySyncReportRequest(
        trace_id="trace-error-body-123",
        report_kind="openai_chat_sync_finalize",
        status_code=502,
        headers={"content-type": "text/plain"},
        body_base64=base64.b64encode(b"upstream exploded").decode("ascii"),
    )

    assert _extract_gateway_sync_error_message(payload) == "upstream exploded"


def test_build_gateway_sync_error_payload_returns_provider_error_when_no_conversion() -> None:
    payload = GatewaySyncReportRequest(
        trace_id="trace-error-payload-123",
        report_kind="openai_chat_sync_finalize",
        status_code=400,
        headers={"content-type": "application/json"},
        body_json={"error": {"message": "bad request"}},
    )

    assert _build_gateway_sync_error_payload(
        payload,
        client_api_format="openai:chat",
        provider_api_format="openai:chat",
        needs_conversion=False,
    ) == {"error": {"message": "bad request"}}


def test_build_gateway_sync_error_payload_builds_generic_client_error_for_non_json() -> None:
    payload = GatewaySyncReportRequest(
        trace_id="trace-error-generic-123",
        report_kind="openai_cli_sync_finalize",
        status_code=504,
        headers={"content-type": "text/plain"},
        body_base64=base64.b64encode(b"gateway timeout").decode("ascii"),
    )

    result = _build_gateway_sync_error_payload(
        payload,
        client_api_format="openai:cli",
        provider_api_format="openai:cli",
        needs_conversion=False,
    )

    assert result["error"]["message"] == "gateway timeout"


def test_resolve_gateway_sync_error_status_code_prefers_embedded_code() -> None:
    class _FakeParser:
        def parse_response(self, response: dict[str, object], status_code: int) -> SimpleNamespace:
            del response, status_code
            return SimpleNamespace(embedded_status_code=429)

    payload = GatewaySyncReportRequest(
        trace_id="trace-error-status-123",
        report_kind="openai_chat_sync_finalize",
        status_code=200,
        headers={"content-type": "application/json"},
        body_json={"error": {"message": "rate limited", "code": 429}},
    )

    assert _resolve_gateway_sync_error_status_code(payload, provider_parser=_FakeParser()) == 429


def test_resolve_gateway_sync_error_status_code_defaults_embedded_200_to_400() -> None:
    payload = GatewaySyncReportRequest(
        trace_id="trace-error-status-default-123",
        report_kind="openai_chat_sync_finalize",
        status_code=200,
        headers={"content-type": "application/json"},
        body_json={"error": {"message": "bad request"}},
    )

    assert _resolve_gateway_sync_error_status_code(payload) == 400


def test_record_gateway_direct_candidate_graph_creates_full_candidate_graph() -> None:
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(engine, tables=[RequestCandidate.__table__])
    SessionLocal = sessionmaker(bind=engine)

    with SessionLocal() as db:
        resolver = CandidateResolver(db=db, cache_scheduler=SimpleNamespace())
        user_api_key = SimpleNamespace(
            id="api-key-direct-123",
            name="client-key",
            user_id="user-direct-123",
            user=SimpleNamespace(id="user-direct-123", username="direct-user"),
        )
        candidates = [
            ProviderCandidate(
                provider=SimpleNamespace(id="provider-selected-123", name="openai", max_retries=1),
                endpoint=SimpleNamespace(id="endpoint-selected-123"),
                key=SimpleNamespace(id="key-selected-123"),
                provider_api_format="openai:chat",
            ),
            ProviderCandidate(
                provider=SimpleNamespace(id="provider-unused-123", name="openai", max_retries=1),
                endpoint=SimpleNamespace(id="endpoint-unused-123"),
                key=SimpleNamespace(id="key-unused-123"),
                provider_api_format="openai:chat",
            ),
            ProviderCandidate(
                provider=SimpleNamespace(id="provider-skipped-123", name="openai", max_retries=1),
                endpoint=SimpleNamespace(id="endpoint-skipped-123"),
                key=SimpleNamespace(id="key-skipped-123"),
                is_skipped=True,
                skip_reason="capability_miss",
                provider_api_format="openai:chat",
            ),
        ]

        _record_gateway_direct_candidate_graph(
            db=db,
            candidate_resolver=resolver,
            candidates=candidates,
            request_id="req-direct-candidate-graph-123",
            user_api_key=user_api_key,
            required_capabilities={"tools": True},
            selected_candidate_index=0,
        )

        rows = (
            db.query(RequestCandidate)
            .filter(RequestCandidate.request_id == "req-direct-candidate-graph-123")
            .order_by(RequestCandidate.candidate_index, RequestCandidate.retry_index)
            .all()
        )

        assert len(rows) == 3
        assert rows[0].status == "pending"
        assert rows[1].status == "unused"
        assert rows[2].status == "skipped"
        assert getattr(candidates[0], "request_candidate_id") == rows[0].id
        assert rows[0].started_at is not None
        assert rows[1].finished_at is not None


def test_record_gateway_direct_candidate_graph_uses_selected_pool_key_index() -> None:
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(engine, tables=[RequestCandidate.__table__])
    SessionLocal = sessionmaker(bind=engine)

    with SessionLocal() as db:
        resolver = CandidateResolver(db=db, cache_scheduler=SimpleNamespace())
        user_api_key = SimpleNamespace(
            id="api-key-pool-123",
            name="client-key",
            user_id="user-pool-123",
            user=SimpleNamespace(id="user-pool-123", username="pool-user"),
        )
        pool_keys = [
            SimpleNamespace(id="pool-key-0"),
            SimpleNamespace(id="pool-key-1"),
        ]
        candidate = PoolCandidate(
            provider=SimpleNamespace(id="provider-pool-123", name="claude", max_retries=1),
            endpoint=SimpleNamespace(id="endpoint-pool-123"),
            key=pool_keys[1],
            pool_keys=pool_keys,
            provider_api_format="claude:cli",
        )
        candidate._pool_key_index = 1

        _record_gateway_direct_candidate_graph(
            db=db,
            candidate_resolver=resolver,
            candidates=[candidate],
            request_id="req-direct-pool-graph-123",
            user_api_key=user_api_key,
            required_capabilities=None,
            selected_candidate_index=0,
        )

        rows = (
            db.query(RequestCandidate)
            .filter(RequestCandidate.request_id == "req-direct-pool-graph-123")
            .order_by(RequestCandidate.candidate_index, RequestCandidate.retry_index)
            .all()
        )

        assert len(rows) == 2
        assert rows[0].status == "unused"
        assert rows[1].status == "pending"
        assert getattr(candidate, "request_candidate_id") == rows[1].id


def test_record_gateway_direct_candidate_graph_is_idempotent_for_same_request_id() -> None:
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(engine, tables=[RequestCandidate.__table__])
    SessionLocal = sessionmaker(bind=engine)

    with SessionLocal() as db:
        resolver = CandidateResolver(db=db, cache_scheduler=SimpleNamespace())
        user_api_key = SimpleNamespace(
            id="api-key-direct-idempotent-123",
            name="client-key",
            user_id="user-direct-idempotent-123",
            user=SimpleNamespace(id="user-direct-idempotent-123", username="direct-user"),
        )
        request_id = "req-direct-candidate-idempotent-123"
        first_candidates = [
            ProviderCandidate(
                provider=SimpleNamespace(id="provider-selected-123", name="openai", max_retries=1),
                endpoint=SimpleNamespace(id="endpoint-selected-123"),
                key=SimpleNamespace(id="key-selected-123"),
                provider_api_format="openai:cli",
            ),
            ProviderCandidate(
                provider=SimpleNamespace(id="provider-unused-123", name="openai", max_retries=1),
                endpoint=SimpleNamespace(id="endpoint-unused-123"),
                key=SimpleNamespace(id="key-unused-123"),
                provider_api_format="openai:cli",
            ),
        ]

        _record_gateway_direct_candidate_graph(
            db=db,
            candidate_resolver=resolver,
            candidates=first_candidates,
            request_id=request_id,
            user_api_key=user_api_key,
            required_capabilities=None,
            selected_candidate_index=0,
        )
        first_selected_record_id = getattr(first_candidates[0], "request_candidate_id")

        second_candidates = [
            ProviderCandidate(
                provider=SimpleNamespace(id="provider-selected-123", name="openai", max_retries=1),
                endpoint=SimpleNamespace(id="endpoint-selected-123"),
                key=SimpleNamespace(id="key-selected-123"),
                provider_api_format="openai:cli",
            ),
            ProviderCandidate(
                provider=SimpleNamespace(id="provider-unused-123", name="openai", max_retries=1),
                endpoint=SimpleNamespace(id="endpoint-unused-123"),
                key=SimpleNamespace(id="key-unused-123"),
                provider_api_format="openai:cli",
            ),
        ]

        _record_gateway_direct_candidate_graph(
            db=db,
            candidate_resolver=resolver,
            candidates=second_candidates,
            request_id=request_id,
            user_api_key=user_api_key,
            required_capabilities=None,
            selected_candidate_index=0,
        )

        rows = (
            db.query(RequestCandidate)
            .filter(RequestCandidate.request_id == request_id)
            .order_by(RequestCandidate.candidate_index, RequestCandidate.retry_index)
            .all()
        )

        assert len(rows) == 2
        assert rows[0].status == "pending"
        assert rows[1].status == "unused"
        assert getattr(second_candidates[0], "request_candidate_id") == first_selected_record_id