Aether/tests/services/test_kiro_envelope.py

from __future__ import annotations

from types import SimpleNamespace

import httpx
import pytest

from src.services.provider.adapters.kiro.context import (
    KiroRequestContext,
    get_kiro_request_context,
    set_kiro_request_context,
)
from src.services.provider.adapters.kiro.envelope import kiro_envelope
from src.services.provider.adapters.kiro.error_enhancer import (
    build_kiro_network_diagnostic,
    classify_kiro_http_status,
    enhance_kiro_http_error_text,
)
from src.services.provider.adapters.kiro.plugin import build_kiro_url
from src.services.provider.request_context import get_selected_base_url, set_selected_base_url

_REFRESH_TOKEN = "r" * 120


@pytest.fixture(autouse=True)
def _reset_kiro_context() -> None:  # type: ignore[misc]
    set_kiro_request_context(None)
    set_selected_base_url(None)
    yield
    set_kiro_request_context(None)
    set_selected_base_url(None)


def test_kiro_http_status_classification() -> None:
    assert classify_kiro_http_status(200) == "ok"
    assert classify_kiro_http_status(401) == "auth_error"
    assert classify_kiro_http_status(429) == "rate_limited"
    assert classify_kiro_http_status(503) == "upstream_server_error"


def test_kiro_envelope_records_http_status() -> None:
    set_kiro_request_context(KiroRequestContext(region="us-east-1", machine_id="mid-1"))

    kiro_envelope.on_http_status(base_url="https://q.us-east-1.amazonaws.com", status_code=429)

    ctx = get_kiro_request_context()
    assert ctx is not None
    assert ctx.last_http_status == 429
    assert ctx.last_http_error_category == "rate_limited"


def test_kiro_envelope_records_connection_error_summary() -> None:
    set_kiro_request_context(KiroRequestContext(region="us-east-1", machine_id="mid-1"))

    kiro_envelope.on_connection_error(
        base_url="https://q.us-east-1.amazonaws.com",
        exc=httpx.ConnectTimeout("dial timed out"),
    )

    ctx = get_kiro_request_context()
    assert ctx is not None
    assert ctx.last_connection_error_category == "connect_timeout"
    assert ctx.last_connection_error_summary is not None
    assert "ConnectTimeout" in ctx.last_connection_error_summary


def test_kiro_envelope_capture_selected_base_url() -> None:
    set_selected_base_url("https://q.us-west-2.amazonaws.com")
    assert kiro_envelope.capture_selected_base_url() == "https://q.us-west-2.amazonaws.com"


def test_build_kiro_url_sets_selected_base_url_and_applies_region() -> None:
    set_kiro_request_context(KiroRequestContext(region="eu-west-1", machine_id="mid-2"))
    endpoint = SimpleNamespace(base_url="https://q.{region}.amazonaws.com")

    url = build_kiro_url(
        endpoint,
        is_stream=True,
        effective_query_params={"alt": "sse"},
    )

    assert url.startswith("https://q.eu-west-1.amazonaws.com/generateAssistantResponse")
    assert get_selected_base_url() == "https://q.eu-west-1.amazonaws.com"


def test_build_kiro_network_diagnostic_prefers_connection_summary() -> None:
    diag = build_kiro_network_diagnostic(
        http_status=503,
        http_category="upstream_server_error",
        connection_summary="connect_timeout: ConnectTimeout: dial timed out",
    )
    assert diag == "network=connect_timeout: ConnectTimeout: dial timed out"


def test_kiro_envelope_wrap_request_adds_inference_config() -> None:
    wrapped, _ = kiro_envelope.wrap_request(
        {
            "model": "claude-sonnet-4-5",
            "max_tokens": 2048,
            "temperature": 0.3,
            "top_p": 0.8,
            "messages": [{"role": "user", "content": "hello"}],
        },
        model="claude-sonnet-4-5",
        url_model=None,
        decrypted_auth_config={
            "auth_method": "social",
            "refreshToken": _REFRESH_TOKEN,
            "profileArn": "arn:aws:iam::1:role/x",
        },
    )

    assert wrapped["conversationState"]["currentMessage"]["userInputMessage"]["modelId"] == (
        "claude-sonnet-4-5"
    )
    assert wrapped["inferenceConfig"] == {
        "maxTokens": 2048,
        "temperature": 0.3,
        "topP": 0.8,
    }
    assert wrapped["profileArn"] == "arn:aws:iam::1:role/x"


def test_kiro_envelope_omits_profile_arn_for_idc_auth() -> None:
    wrapped, _ = kiro_envelope.wrap_request(
        {
            "messages": [{"role": "user", "content": "hello"}],
        },
        model="claude-sonnet-4-5",
        url_model=None,
        decrypted_auth_config={
            "auth_method": "identity_center",
            "refreshToken": _REFRESH_TOKEN,
            "profileArn": "arn:aws:iam::1:role/x",
            "clientId": "cid",
            "clientSecret": "secret",
        },
    )

    assert "profileArn" not in wrapped


def test_enhance_kiro_http_error_text_maps_known_reason() -> None:
    text = enhance_kiro_http_error_text(
        '{"message":"Input is too long.","reason":"CONTENT_LENGTH_EXCEEDS_THRESHOLD"}',
        status_code=400,
    )

    assert "[CONTENT_LENGTH_EXCEEDS_THRESHOLD]" in text
    assert "输入超过模型上下文限制" in text