Aether/tests/api/internal_gateway/test_gateway_reporting_success.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 _make_legacy_test_client(app: FastAPI) -> TestClient:
    return TestClient(
        app,
        base_url="http://127.0.0.1",
        headers={"x-aether-legacy-internal-gateway": "true"},
    )


def test_report_sync_route_records_openai_chat_sync_success(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_db] = lambda: object()
    monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
    record_mock = AsyncMock(return_value=None)
    monkeypatch.setattr(
        "src.api.internal.gateway._record_gateway_openai_chat_sync_success",
        record_mock,
    )

    client = _make_legacy_test_client(app)
    response = client.post(
        "/api/internal/gateway/report-sync",
        json={
            "trace_id": "trace-openai-chat-report",
            "report_kind": "openai_chat_sync_success",
            "report_context": {"user_id": "user-123", "api_key_id": "key-123"},
            "status_code": 200,
            "headers": {"content-type": "application/json"},
            "body_json": {
                "id": "chatcmpl-123",
                "object": "chat.completion",
                "choices": [],
            },
            "telemetry": {"elapsed_ms": 123},
        },
    )

    assert response.status_code == 200
    assert response.json() == {"ok": True}
    _wait_until(lambda: record_mock.await_count == 1)
    record_mock.assert_awaited_once()


def test_report_sync_route_records_openai_cli_sync_success(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_db] = lambda: object()
    monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
    record_mock = AsyncMock(return_value=None)
    monkeypatch.setattr(
        "src.api.internal.gateway._record_gateway_passthrough_cli_sync_success",
        record_mock,
    )

    client = _make_legacy_test_client(app)
    response = client.post(
        "/api/internal/gateway/report-sync",
        json={
            "trace_id": "trace-openai-cli-report",
            "report_kind": "openai_cli_sync_success",
            "report_context": {
                "user_id": "user-123",
                "api_key_id": "key-123",
                "client_api_format": "openai:cli",
            },
            "status_code": 200,
            "headers": {"content-type": "application/json"},
            "body_json": {
                "id": "resp_123",
                "object": "response",
                "output": [],
            },
            "telemetry": {"elapsed_ms": 91},
        },
    )

    assert response.status_code == 200
    assert response.json() == {"ok": True}
    _wait_until(lambda: record_mock.await_count == 1)
    record_mock.assert_awaited_once()


def test_report_sync_route_records_claude_cli_sync_success(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_db] = lambda: object()
    monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
    record_mock = AsyncMock(return_value=None)
    monkeypatch.setattr(
        "src.api.internal.gateway._record_gateway_passthrough_cli_sync_success",
        record_mock,
    )

    client = _make_legacy_test_client(app)
    response = client.post(
        "/api/internal/gateway/report-sync",
        json={
            "trace_id": "trace-claude-cli-report",
            "report_kind": "claude_cli_sync_success",
            "report_context": {
                "user_id": "user-123",
                "api_key_id": "key-123",
                "client_api_format": "claude:cli",
            },
            "status_code": 200,
            "headers": {"content-type": "application/json"},
            "body_json": {
                "id": "msg_123",
                "type": "message",
                "content": [],
            },
        },
    )

    assert response.status_code == 200
    assert response.json() == {"ok": True}
    _wait_until(lambda: record_mock.await_count == 1)
    record_mock.assert_awaited_once()


def test_report_sync_route_records_gemini_cli_sync_success(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_db] = lambda: object()
    monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
    record_mock = AsyncMock(return_value=None)
    monkeypatch.setattr(
        "src.api.internal.gateway._record_gateway_passthrough_cli_sync_success",
        record_mock,
    )

    client = _make_legacy_test_client(app)
    response = client.post(
        "/api/internal/gateway/report-sync",
        json={
            "trace_id": "trace-gemini-cli-report",
            "report_kind": "gemini_cli_sync_success",
            "report_context": {
                "user_id": "user-123",
                "api_key_id": "key-123",
                "client_api_format": "gemini:cli",
            },
            "status_code": 200,
            "headers": {"content-type": "application/json"},
            "body_json": {
                "candidates": [],
            },
        },
    )

    assert response.status_code == 200
    assert response.json() == {"ok": True}
    _wait_until(lambda: record_mock.await_count == 1)
    record_mock.assert_awaited_once()


@pytest.mark.parametrize(
    ("report_kind", "recorder_attr"),
    [
        ("openai_video_create_sync_success", "_record_gateway_openai_video_create_sync_success"),
        ("openai_video_remix_sync_success", "_record_gateway_openai_video_remix_sync_success"),
        ("openai_video_delete_sync_success", "_record_gateway_openai_video_delete_sync_success"),
        ("openai_video_cancel_sync_success", "_record_gateway_openai_video_cancel_sync_success"),
        ("gemini_video_create_sync_success", "_record_gateway_gemini_video_create_sync_success"),
        ("gemini_video_cancel_sync_success", "_record_gateway_gemini_video_cancel_sync_success"),
    ],
)
def test_report_sync_route_records_video_sync_success_variants(
    monkeypatch: pytest.MonkeyPatch,
    report_kind: str,
    recorder_attr: str,
) -> None:
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_db] = lambda: object()
    monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
    record_mock = AsyncMock(return_value=None)
    monkeypatch.setattr(f"src.api.internal.gateway.{recorder_attr}", record_mock)

    client = _make_legacy_test_client(app)
    response = client.post(
        "/api/internal/gateway/report-sync",
        json={
            "trace_id": f"trace-{report_kind}",
            "report_kind": report_kind,
            "report_context": {
                "user_id": "user-123",
                "api_key_id": "key-123",
                "task_id": "task-123",
            },
            "status_code": 200,
            "headers": {"content-type": "application/json"},
            "body_json": {},
            "telemetry": {"elapsed_ms": 33},
        },
    )

    assert response.status_code == 200
    assert response.json() == {"ok": True}
    _wait_until(lambda: record_mock.await_count == 1)
    record_mock.assert_awaited_once()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("recorder_attr", "finalize_attr"),
    [
        (
            "_record_gateway_openai_video_create_sync_success",
            "_finalize_gateway_openai_video_create_sync",
        ),
        (
            "_record_gateway_openai_video_remix_sync_success",
            "_finalize_gateway_openai_video_remix_sync",
        ),
        (
            "_record_gateway_gemini_video_create_sync_success",
            "_finalize_gateway_gemini_video_create_sync",
        ),
    ],
)
async def test_video_sync_success_recorders_delegate_to_finalize(
    monkeypatch: pytest.MonkeyPatch,
    recorder_attr: str,
    finalize_attr: str,
) -> None:
    import src.api.internal.gateway as gateway_module

    finalize_mock = AsyncMock(return_value=JSONResponse({"ok": True}))
    monkeypatch.setattr(f"src.api.internal.gateway.{finalize_attr}", finalize_mock)

    db = object()
    payload = GatewaySyncReportRequest(
        trace_id="trace-video-sync-success-123",
        report_kind="video-sync-success",
        report_context={"user_id": "user-123", "api_key_id": "key-123"},
        status_code=200,
        headers={"content-type": "application/json"},
        body_json={"id": "ext-123"},
    )

    recorder = getattr(gateway_module, recorder_attr)
    await recorder(payload, db=db)

    finalize_mock.assert_awaited_once_with(payload, db=db)


def test_report_sync_route_records_claude_chat_sync_success(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_db] = lambda: object()
    monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
    record_mock = AsyncMock(return_value=None)
    monkeypatch.setattr(
        "src.api.internal.gateway._record_gateway_passthrough_chat_sync_success",
        record_mock,
    )

    client = _make_legacy_test_client(app)
    response = client.post(
        "/api/internal/gateway/report-sync",
        json={
            "trace_id": "trace-claude-chat-report",
            "report_kind": "claude_chat_sync_success",
            "report_context": {
                "user_id": "user-123",
                "api_key_id": "key-123",
                "client_api_format": "claude:chat",
            },
            "status_code": 200,
            "headers": {"content-type": "application/json"},
            "body_json": {
                "id": "msg_123",
                "type": "message",
                "content": [],
            },
        },
    )

    assert response.status_code == 200
    assert response.json() == {"ok": True}
    _wait_until(lambda: record_mock.await_count == 1)
    record_mock.assert_awaited_once()


def test_report_sync_route_records_gemini_chat_sync_success(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_db] = lambda: object()
    monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
    record_mock = AsyncMock(return_value=None)
    monkeypatch.setattr(
        "src.api.internal.gateway._record_gateway_passthrough_chat_sync_success",
        record_mock,
    )

    client = _make_legacy_test_client(app)
    response = client.post(
        "/api/internal/gateway/report-sync",
        json={
            "trace_id": "trace-gemini-chat-report",
            "report_kind": "gemini_chat_sync_success",
            "report_context": {
                "user_id": "user-123",
                "api_key_id": "key-123",
                "client_api_format": "gemini:chat",
            },
            "status_code": 200,
            "headers": {"content-type": "application/json"},
            "body_json": {
                "candidates": [],
            },
        },
    )

    assert response.status_code == 200
    assert response.json() == {"ok": True}
    _wait_until(lambda: record_mock.await_count == 1)
    record_mock.assert_awaited_once()


@pytest.mark.asyncio


@pytest.mark.asyncio


@pytest.mark.asyncio
async def test_record_gateway_openai_chat_sync_success_uses_provider_and_client_bodies(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.api.internal import gateway as gateway_module

    record_success = AsyncMock(return_value=None)
    fake_handler = SimpleNamespace(
        _normalize_response=lambda response: response,
        _extract_usage=lambda response: {
            "input_tokens": response.get("usage", {}).get("prompt_tokens", 0),
            "output_tokens": response.get("usage", {}).get("completion_tokens", 0),
        },
        telemetry=SimpleNamespace(record_success=record_success),
    )

    class FakeOpenAIChatAdapter:
        API_FAMILY = SimpleNamespace(value="openai")
        ENDPOINT_KIND = SimpleNamespace(value="chat")

        def _create_handler(self, **kwargs: Any) -> Any:
            return fake_handler

    fake_user = object()
    fake_api_key = object()

    class FakeQuery:
        def __init__(self, value: Any) -> None:
            self.value = value

        def filter(self, *args: Any, **kwargs: Any) -> "FakeQuery":
            return self

        def first(self) -> Any:
            return self.value

    class FakeDb:
        def __init__(self) -> None:
            self._values = iter((fake_user, fake_api_key))

        def query(self, model: Any) -> FakeQuery:
            return FakeQuery(next(self._values))

    monkeypatch.setattr("src.api.handlers.openai.OpenAIChatAdapter", FakeOpenAIChatAdapter)

    provider_body = {
        "responseId": "resp-provider-123",
        "candidates": [
            {
                "content": {"parts": [{"text": "provider"}], "role": "model"},
                "finishReason": "STOP",
            }
        ],
        "usageMetadata": {
            "promptTokenCount": 2,
            "candidatesTokenCount": 3,
            "totalTokenCount": 5,
        },
    }
    client_body = {
        "id": "chatcmpl-client-123",
        "object": "chat.completion",
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "client"},
                "finish_reason": "stop",
            }
        ],
        "usage": {
            "prompt_tokens": 2,
            "completion_tokens": 3,
            "total_tokens": 5,
        },
    }
    payload = GatewaySyncReportRequest(
        trace_id="trace-openai-chat-conversion-123",
        report_kind="openai_chat_sync_success",
        report_context={
            "user_id": "user-123",
            "api_key_id": "key-123",
            "provider_name": "gemini",
            "provider_api_format": "gemini:chat",
            "client_api_format": "openai:chat",
            "model": "gpt-5",
            "mapped_model": "gpt-5",
            "request_id": "req-123",
        },
        status_code=200,
        headers={"content-type": "application/json"},
        body_json=provider_body,
        client_body_json=client_body,
        telemetry={"elapsed_ms": 31},
    )

    await gateway_module._record_gateway_openai_chat_sync_success(payload, db=FakeDb())

    kwargs = record_success.await_args.kwargs
    assert kwargs["response_body"] == provider_body
    assert kwargs["client_response_body"] == client_body
    assert kwargs["has_format_conversion"] is True


@pytest.mark.asyncio


@pytest.mark.asyncio


@pytest.mark.asyncio
async def test_record_gateway_passthrough_cli_sync_success_uses_provider_and_client_bodies(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.api.internal import gateway as gateway_module

    record_success = AsyncMock(return_value=None)

    class FakeCliHandler:
        def __init__(self, **kwargs: Any) -> None:
            self.parser = SimpleNamespace(
                extract_usage_from_response=lambda response: {
                    "input_tokens": response.get("usage", {}).get("input_tokens", 0),
                    "output_tokens": response.get("usage", {}).get("output_tokens", 0),
                    "cache_creation_tokens": 0,
                    "cache_read_tokens": 0,
                }
            )
            self.telemetry = SimpleNamespace(record_success=record_success)

        def _extract_response_metadata(self, response: dict[str, Any]) -> dict[str, Any]:
            return {"response_id": response.get("id")}

    class FakeOpenAICliAdapter:
        API_FAMILY = SimpleNamespace(value="openai")
        ENDPOINT_KIND = SimpleNamespace(value="cli")
        HANDLER_CLASS = FakeCliHandler
        allowed_api_formats = ["openai:cli"]

        @staticmethod
        def detect_capability_requirements(*args: Any, **kwargs: Any) -> dict[str, Any]:
            return {}

    fake_user = object()
    fake_api_key = object()

    class FakeQuery:
        def __init__(self, value: Any) -> None:
            self.value = value

        def filter(self, *args: Any, **kwargs: Any) -> "FakeQuery":
            return self

        def first(self) -> Any:
            return self.value

    class FakeDb:
        def __init__(self) -> None:
            self._values = iter((fake_user, fake_api_key))

        def query(self, model: Any) -> FakeQuery:
            return FakeQuery(next(self._values))

    monkeypatch.setattr("src.api.handlers.openai_cli.OpenAICliAdapter", FakeOpenAICliAdapter)

    provider_body = {
        "responseId": "resp-provider-cli-123",
        "candidates": [
            {
                "content": {"parts": [{"text": "provider"}], "role": "model"},
                "finishReason": "STOP",
            }
        ],
        "usageMetadata": {
            "promptTokenCount": 2,
            "candidatesTokenCount": 3,
            "totalTokenCount": 5,
        },
    }
    client_body = {
        "id": "resp-client-cli-123",
        "object": "response",
        "output": [],
        "usage": {
            "input_tokens": 2,
            "output_tokens": 3,
            "total_tokens": 5,
        },
    }
    payload = GatewaySyncReportRequest(
        trace_id="trace-openai-cli-conversion-123",
        report_kind="openai_cli_sync_success",
        report_context={
            "user_id": "user-123",
            "api_key_id": "key-123",
            "provider_name": "gemini",
            "provider_api_format": "gemini:cli",
            "client_api_format": "openai:cli",
            "model": "gpt-5",
            "mapped_model": "gpt-5",
            "request_id": "req-cli-123",
        },
        status_code=200,
        headers={"content-type": "application/json"},
        body_json=provider_body,
        client_body_json=client_body,
        telemetry={"elapsed_ms": 29},
    )

    await gateway_module._record_gateway_passthrough_cli_sync_success(
        payload, db=FakeDb()
    )

    kwargs = record_success.await_args.kwargs
    assert kwargs["response_body"] == provider_body
    assert kwargs["client_response_body"] == client_body
    assert kwargs["has_format_conversion"] is True


@pytest.mark.asyncio


@pytest.mark.asyncio


@pytest.mark.asyncio
async def test_record_gateway_passthrough_cli_sync_success_postprocesses_antigravity_response(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from src.api.internal import gateway as gateway_module

    record_success = AsyncMock(return_value=None)
    seen_cache: list[tuple[str, dict[str, Any]]] = []

    class FakeCliHandler:
        def __init__(self, **kwargs: Any) -> None:
            self.parser = SimpleNamespace(
                extract_usage_from_response=lambda response: {
                    "input_tokens": response.get("usageMetadata", {}).get("promptTokenCount", 0),
                    "output_tokens": response.get("usageMetadata", {}).get(
                        "candidatesTokenCount", 0
                    ),
                    "cache_creation_tokens": 0,
                    "cache_read_tokens": 0,
                }
            )
            self.telemetry = SimpleNamespace(record_success=record_success)

        def _extract_response_metadata(self, response: dict[str, Any]) -> dict[str, Any]:
            return {"response_id": response.get("_v1internal_response_id")}

    class FakeGeminiCliAdapter:
        API_FAMILY = SimpleNamespace(value="gemini")
        ENDPOINT_KIND = SimpleNamespace(value="cli")
        HANDLER_CLASS = FakeCliHandler
        allowed_api_formats = ["gemini:cli"]

        @staticmethod
        def detect_capability_requirements(*args: Any, **kwargs: Any) -> dict[str, Any]:
            return {}

    fake_user = object()
    fake_api_key = object()

    class FakeQuery:
        def __init__(self, value: Any) -> None:
            self.value = value

        def filter(self, *args: Any, **kwargs: Any) -> "FakeQuery":
            return self

        def first(self) -> Any:
            return self.value

    class FakeDb:
        def __init__(self) -> None:
            self._values = iter((fake_user, fake_api_key))

        def query(self, model: Any) -> FakeQuery:
            return FakeQuery(next(self._values))

    def _record_cache(model: str, response: dict[str, Any]) -> None:
        seen_cache.append((model, response))

    monkeypatch.setattr("src.api.handlers.gemini_cli.GeminiCliAdapter", FakeGeminiCliAdapter)
    monkeypatch.setattr(
        "src.services.provider.adapters.antigravity.envelope.cache_thought_signatures",
        _record_cache,
    )

    provider_body = {
        "_v1internal_response_id": "resp-antigravity-123",
        "candidates": [
            {
                "content": {
                    "parts": [
                        {
                            "text": "thinking",
                            "thoughtSignature": "a" * 60,
                        }
                    ],
                    "role": "model",
                },
                "finishReason": "STOP",
            }
        ],
        "usageMetadata": {
            "promptTokenCount": 2,
            "candidatesTokenCount": 3,
            "totalTokenCount": 5,
        },
    }
    payload = GatewaySyncReportRequest(
        trace_id="trace-antigravity-cli-report-123",
        report_kind="gemini_cli_sync_success",
        report_context={
            "user_id": "user-123",
            "api_key_id": "key-123",
            "provider_name": "antigravity",
            "provider_api_format": "gemini:cli",
            "client_api_format": "gemini:cli",
            "model": "claude-sonnet-4-5",
            "mapped_model": "claude-sonnet-4-5",
            "request_id": "req-antigravity-cli-123",
            "envelope_name": "antigravity:v1internal",
        },
        status_code=200,
        headers={"content-type": "application/json"},
        body_json=provider_body,
        telemetry={"elapsed_ms": 17},
    )

    await gateway_module._record_gateway_passthrough_cli_sync_success(payload, db=FakeDb())

    assert len(seen_cache) == 1
    assert seen_cache[0][0] == "claude-sonnet-4-5"
    assert seen_cache[0][1]["_v1internal_response_id"] == "resp-antigravity-123"