Aether/tests/api/handlers/base/test_antigravity_v1internal.py

from __future__ import annotations

import json
from types import SimpleNamespace
from typing import cast
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from src.api.handlers.base.cli_protocol import CliHandlerProtocol
from src.api.handlers.base.stream_context import StreamContext
from src.api.handlers.gemini_cli.handler import GeminiCliMessageHandler
from src.services.provider.adapters.antigravity.envelope import (
    unwrap_v1internal_response,
    wrap_v1internal_request,
)


def _make_handler() -> GeminiCliMessageHandler:
    return GeminiCliMessageHandler(
        db=MagicMock(),  # type: ignore[arg-type]
        user=SimpleNamespace(id=1),  # type: ignore[arg-type]
        api_key=SimpleNamespace(id=1),  # type: ignore[arg-type]
        request_id="req_1",
        client_ip="127.0.0.1",
        user_agent="pytest",
        start_time=0.0,
    )


def test_wrap_v1internal_request_removes_inner_model() -> None:
    gemini_request = {
        "model": "gemini-2.0-flash",
        "contents": [{"role": "user", "parts": [{"text": "hi"}]}],
    }

    wrapped = wrap_v1internal_request(
        gemini_request,
        project_id="project-123",
        model="gemini-2.0-flash",
    )

    assert wrapped["project"] == "project-123"
    assert wrapped["model"] == "gemini-2.0-flash"
    assert "request" in wrapped
    assert "model" not in wrapped["request"]


def test_unwrap_v1internal_response() -> None:
    v1_resp = {
        "response": {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]},
        "responseId": "resp-123",
    }

    unwrapped = unwrap_v1internal_response(v1_resp)

    assert "response" not in unwrapped
    assert "candidates" in unwrapped
    assert unwrapped["_v1internal_response_id"] == "resp-123"


def test_convert_sse_line_unwraps_before_convert_stream_chunk() -> None:
    handler = _make_handler()

    ctx = StreamContext(model="gemini-2.0-flash", api_format="gemini:cli")
    ctx.provider_type = "antigravity"
    ctx.provider_api_format = "gemini:chat"
    ctx.client_api_format = "gemini:cli"

    v1_line = (
        'data: {"response": {"candidates": [{"content": {"parts": [{"text": "Hi"}]}}]},'
        ' "responseId": "123"}'
    )

    seen: dict[str, object] = {}

    class _DummyRegistry:
        def convert_stream_chunk(
            self, data_obj: object, *_args: object, **_kwargs: object
        ) -> list[str]:
            seen["data_obj"] = data_obj
            return []

    with patch(
        "src.api.handlers.base.cli_event_mixin.get_format_converter_registry",
        return_value=_DummyRegistry(),
    ):
        _lines, _events = handler._convert_sse_line(ctx, v1_line, [])

    assert isinstance(seen.get("data_obj"), dict)
    assert "response" not in seen["data_obj"]  # type: ignore[operator]
    assert "_v1internal_response_id" in seen["data_obj"]  # type: ignore[operator]


def test_handle_sse_event_unwraps_for_antigravity() -> None:
    handler = _make_handler()

    ctx = StreamContext(model="gemini-2.0-flash", api_format="gemini:cli")
    ctx.provider_type = "antigravity"

    v1_data = {
        "response": {
            "candidates": [
                {
                    "content": {"parts": [{"text": "Hello"}], "role": "model"},
                    "finishReason": "STOP",
                }
            ],
            "modelVersion": "gemini-2.0-flash-001",
        },
        "responseId": "123",
    }

    with patch.object(handler, "_process_event_data") as mock_process:
        cast(CliHandlerProtocol, handler)._handle_sse_event(
            ctx, None, json.dumps(v1_data), record_chunk=False
        )

    assert mock_process.call_count == 1
    passed_data = mock_process.call_args[0][2]
    assert isinstance(passed_data, dict)
    assert "response" not in passed_data
    assert "candidates" in passed_data


def test_handle_sse_event_caches_thought_signature_for_antigravity() -> None:
    from src.core.api_format.conversion.thinking_cache import signature_cache

    signature_cache.clear()

    handler = _make_handler()
    ctx = StreamContext(model="claude-sonnet-4-5", api_format="gemini:cli")
    ctx.provider_type = "antigravity"

    # 签名须 >= MIN_SIGNATURE_LENGTH(50)，否则会被忽略
    long_sig = "a" * 60

    payload = {
        "candidates": [
            {
                "content": {
                    "parts": [
                        {
                            "text": "t1",
                            "thought": True,
                            "thoughtSignature": long_sig,
                        }
                    ]
                }
            }
        ]
    }

    with patch.object(handler, "_process_event_data") as _mock_process:
        cast(CliHandlerProtocol, handler)._handle_sse_event(
            ctx, None, json.dumps(payload), record_chunk=False
        )

    assert signature_cache.get_or_dummy("claude-sonnet-4-5", "t1") == long_sig


def test_provider_type_drives_antigravity_usage_path() -> None:
    handler = _make_handler()

    event = {"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 2}}
    usage = handler._extract_usage_from_event(event, provider_type="antigravity")

    assert usage["input_tokens"] == 10
    assert usage["output_tokens"] == 2


@pytest.mark.asyncio
async def test_antigravity_forces_conversion_path_in_stream_with_prefetch() -> None:
    handler = _make_handler()

    ctx = StreamContext(model="gemini-2.0-flash", api_format="gemini:cli")
    ctx.provider_type = "antigravity"
    ctx.needs_conversion = False  # same-format case normally would passthrough

    class _AsyncIter:
        def __init__(self, items: list[bytes]) -> None:
            self._it = iter(items)

        def __aiter__(self) -> _AsyncIter:
            return self

        async def __anext__(self) -> bytes:
            try:
                return next(self._it)
            except StopIteration as e:
                raise StopAsyncIteration from e

    prefetched = [
        b'data: {"response": {"candidates": []}, "responseId": "1"}\n',
    ]
    byte_iter = _AsyncIter([])  # no more bytes after prefetch
    response_ctx = SimpleNamespace(__aexit__=AsyncMock(return_value=None))

    with patch.object(
        handler,
        "_convert_sse_line",
        return_value=(["data: {}"], []),
    ) as mock_convert:
        out = []
        async for chunk in handler._create_response_stream_with_prefetch(
            ctx, byte_iter, response_ctx, prefetched  # type: ignore[arg-type]
        ):
            out.append(chunk)

    assert mock_convert.call_count >= 1
    assert any(b"data: {}" in c for c in out)


def test_wrap_v1internal_request_injects_thought_signature_from_tool_cache() -> None:
    from src.core.api_format.conversion.thinking_cache import signature_cache

    signature_cache.clear()

    tool_id = "toolu_123"
    sig = "a" * 60
    signature_cache.cache_tool_signature(tool_id, sig)

    gemini_request = {
        "model": "gemini-2.0-flash",
        "contents": [
            {"role": "user", "parts": [{"text": "hi"}]},
            {"role": "model", "parts": [{"function_call": {"name": "do", "id": tool_id}}]},
        ],
    }

    wrapped = wrap_v1internal_request(
        gemini_request,
        project_id="project-123",
        model="gemini-2.0-flash",
    )

    model_turn = wrapped["request"]["contents"][1]
    part = model_turn["parts"][0]
    assert part["thoughtSignature"] == sig


def test_wrap_v1internal_request_injects_session_signature_when_tool_cache_missing() -> None:
    from src.core.api_format.conversion.thinking_cache import signature_cache

    signature_cache.clear()

    session_id = "sid-123"
    sig = "b" * 60
    signature_cache.cache_session_signature(session_id, sig, message_count=2)

    gemini_request = {
        "model": "gemini-2.0-flash",
        "sessionId": session_id,
        "contents": [
            {"role": "user", "parts": [{"text": "hi"}]},
            {
                "role": "model",
                "parts": [{"functionCall": {"name": "do", "id": "toolu_missing"}}],
            },
        ],
    }

    wrapped = wrap_v1internal_request(
        gemini_request,
        project_id="project-123",
        model="gemini-2.0-flash",
    )

    model_turn = wrapped["request"]["contents"][1]
    part = model_turn["parts"][0]
    assert part["thoughtSignature"] == sig