LLM_Proxy/tests/test_proxy_error_handling.py

from __future__ import annotations

import sys
import unittest
from pathlib import Path
from types import SimpleNamespace

import requests
import websocket
from flask import Flask

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src.application.app_context import AppContext
from src.external import LLMProvider, StaticUpstreamResponse
from src.external.stream_probe import BufferedUpstreamResponse
from src.hooks import BaseHook, HookContext, HookErrorType
from src.presentation.proxy_controller import ProxyController
from src.services.proxy_service import ProxyErrorInfo, ProxyService


class FakeLogger:
    def __init__(self) -> None:
        self.records: list[tuple[str, str]] = []

    def _log(self, level: str, msg: str, *args) -> None:
        rendered = msg % args if args else msg
        self.records.append((level, rendered))

    def info(self, msg: str, *args) -> None:
        self._log("info", msg, *args)

    def warning(self, msg: str, *args) -> None:
        self._log("warning", msg, *args)

    def error(self, msg: str, *args) -> None:
        self._log("error", msg, *args)

    def debug(self, msg: str, *args) -> None:
        self._log("debug", msg, *args)

    def messages(self, level: str) -> list[str]:
        return [message for record_level, message in self.records if record_level == level]


class FakeConfigManager:
    def __init__(self, whitelist_enabled: bool = False) -> None:
        self._whitelist_enabled = whitelist_enabled

    def is_chat_whitelist_enabled(self) -> bool:
        return self._whitelist_enabled


class FakeUserService:
    _UNSET = object()

    def __init__(self, *, user=_UNSET, accessible_models=None) -> None:
        self._user = {"username": "tester", "model_permissions": "*"} if user is self._UNSET else user
        self._accessible_models = None if accessible_models is None else list(accessible_models)

    def get_user_by_ip(self, ip_address: str, require_whitelist_access: bool = True):
        del ip_address, require_whitelist_access
        return self._user

    def can_user_access_model(self, user, model_name: str, available_models=None) -> bool:
        del user
        if self._accessible_models is None:
            return True
        return model_name in set(self._accessible_models)

    def get_accessible_models_for_user(self, user, available_models=None):
        del user
        if self._accessible_models is None:
            return list(available_models or [])
        if available_models is None:
            return list(self._accessible_models)
        available_set = set(available_models)
        return [model_name for model_name in self._accessible_models if model_name in available_set]


class FakeLogService:
    @staticmethod
    def log_request(**kwargs) -> None:
        del kwargs


class FakeProviderManager:
    def __init__(self, provider: LLMProvider) -> None:
        self._provider = provider

    def get_provider_for_model(self, model_name: str):
        if model_name in {f"{self._provider.name}/{self._provider.model_list[0]}", self._provider.model_list[0]}:
            return self._provider
        return None

    def list_model_names(self) -> tuple[str, ...]:
        return tuple(f"{self._provider.name}/{model}" for model in self._provider.model_list)

    def get_provider_view(self, provider_name: str):
        if provider_name != self._provider.name:
            return None
        return SimpleNamespace(
            name=self._provider.name,
            api=self._provider.api,
            transport=self._provider.transport,
            source_format=self._provider.source_format,
            target_format=self._provider.target_format,
            target_formats=self._provider.target_formats,
            model_list=self._provider.model_list,
            proxy=self._provider.proxy,
            timeout_seconds=self._provider.timeout_seconds,
            max_retries=self._provider.max_retries,
            verify_ssl=self._provider.verify_ssl,
            hook=self._provider.hook,
        )


class StubProxyService:
    def __init__(self, proxy_result) -> None:
        self._proxy_result = proxy_result

    def proxy_request(self, *args, **kwargs):
        del args, kwargs
        return self._proxy_result


class RecordingProxyService:
    def __init__(self, proxy_result) -> None:
        self._proxy_result = proxy_result
        self.last_args = None
        self.last_kwargs = None

    def proxy_request(self, *args, **kwargs):
        self.last_args = args
        self.last_kwargs = kwargs
        return self._proxy_result


class ContextRecordingHook(BaseHook):
    def __init__(self) -> None:
        self.contexts: list[HookContext] = []

    def header_hook(self, ctx: HookContext, headers: dict[str, str]) -> dict[str, str]:
        self.contexts.append(ctx)
        return headers


class ProxyControllerErrorFormatTests(unittest.TestCase):
    def test_list_models_includes_provider_protocol_metadata(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/responses",
            transport="http",
            source_format="openai_responses",
            target_format="codex",
            model_list=("gpt-5-codex",),
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService((None, 200, None)),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().get("/v1/models")

        self.assertEqual(200, response.status_code)
        self.assertEqual(
            {
                "object": "list",
                "data": [
                    {
                        "id": "demo/gpt-5-codex",
                        "object": "model",
                        "owned_by": "demo",
                        "provider_name": "demo",
                        "source_format": "openai_responses",
                        "target_formats": ["codex"],
                        "transport": "http",
                    }
                ],
            },
            response.get_json(),
        )

    def test_list_models_filters_to_models_allowed_for_whitelisted_user(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1", "gpt-4.1-mini"),
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(whitelist_enabled=True),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService((None, 200, None)),
            FakeUserService(accessible_models=["demo/gpt-4.1-mini"]),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().get("/v1/models", environ_base={"REMOTE_ADDR": "127.0.0.1"})

        self.assertEqual(200, response.status_code)
        payload = response.get_json()
        self.assertEqual(["demo/gpt-4.1-mini"], [item["id"] for item in payload["data"]])

    def test_list_models_rejects_non_whitelisted_ip_when_whitelist_enabled(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(whitelist_enabled=True),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService((None, 200, None)),
            FakeUserService(user=None),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().get("/v1/models", environ_base={"REMOTE_ADDR": "127.0.0.1"})

        self.assertEqual(403, response.status_code)
        self.assertEqual("ip_not_whitelisted", response.get_json()["error"]["code"])

    def test_chat_completions_rejects_model_not_allowed_for_whitelisted_user(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1", "gpt-4.1-mini"),
        )
        logger = FakeLogger()
        app = Flask(__name__)
        ctx = AppContext(
            logger=logger,
            config_manager=FakeConfigManager(whitelist_enabled=True),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService((None, 200, None)),
            FakeUserService(accessible_models=["demo/gpt-4.1-mini"]),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/chat/completions",
            json={"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(403, response.status_code)
        self.assertEqual("model_not_allowed", response.get_json()["error"]["code"])
        self.assertTrue(
            any("is not allowed to access model=demo/gpt-4.1" in msg for msg in logger.messages("warning"))
        )

    def test_chat_completions_returns_openai_style_error_payload_for_upstream_failures(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
        )
        logger = FakeLogger()
        app = Flask(__name__)
        ctx = AppContext(
            logger=logger,
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService(
                (
                    None,
                    502,
                    ProxyErrorInfo(
                        message="HTTP upstream request failed after 2 attempts: dial tcp timeout",
                        status_code=502,
                        error_type="upstream_error",
                        error_code="upstream_request_failed",
                    ),
                )
            ),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/chat/completions",
            json={"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(502, response.status_code)
        self.assertEqual(
            {
                "error": {
                    "message": "HTTP upstream request failed after 2 attempts: dial tcp timeout",
                    "type": "upstream_error",
                    "param": None,
                    "code": "upstream_request_failed",
                }
            },
            response.get_json(),
        )
        self.assertTrue(
            any("upstream_error=HTTP upstream request failed after 2 attempts: dial tcp timeout" in msg for msg in logger.messages("error"))
        )

    def test_responses_route_rejects_target_format_mismatch(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            target_format="openai_chat",
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService((None, 200, None)),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/responses",
            json={"model": "demo/gpt-4.1", "input": "hi"},
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(400, response.status_code)
        self.assertEqual(
            {
                "error": {
                    "message": "Model demo/gpt-4.1 is configured for downstream formats openai_chat, not one of openai_responses, codex",
                    "type": "invalid_request_error",
                    "param": None,
                    "code": "target_format_mismatch",
                }
            },
            response.get_json(),
        )

    def test_responses_route_allows_openai_responses_target(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            target_format="openai_responses",
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService(
                (
                    app.response_class(
                        '{"id":"resp_1","object":"response"}',
                        status=200,
                        mimetype="application/json",
                    ),
                    200,
                    None,
                )
            ),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/responses",
            json={"model": "demo/gpt-4.1", "input": "hi"},
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(200, response.status_code)
        self.assertEqual('{"id":"resp_1","object":"response"}', response.get_data(as_text=True))

    def test_chat_route_resolves_target_format_from_multi_target_provider(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            target_formats=("openai_chat", "claude_chat"),
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        proxy_service = RecordingProxyService(
            (
                app.response_class(
                    '{"id":"chatcmpl_1","object":"chat.completion"}',
                    status=200,
                    mimetype="application/json",
                ),
                200,
                None,
            )
        )
        ProxyController(
            ctx,
            proxy_service,
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/chat/completions",
            json={"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(200, response.status_code)
        self.assertEqual("openai_chat", proxy_service.last_kwargs["resolved_target_format"])

    def test_responses_route_allows_codex_target(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/responses",
            model_list=("gpt-5.2",),
            target_format="codex",
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService(
                (
                    app.response_class(
                        '{"id":"resp_1","object":"response"}',
                        status=200,
                        mimetype="application/json",
                    ),
                    200,
                    None,
                )
            ),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/responses",
            json={"model": "demo/gpt-5.2", "input": "hi"},
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(200, response.status_code)
        self.assertEqual('{"id":"resp_1","object":"response"}', response.get_data(as_text=True))

    def test_messages_route_rejects_target_format_mismatch(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            target_format="openai_chat",
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService((None, 200, None)),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/messages",
            json={"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]},
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(400, response.status_code)
        self.assertEqual(
            {
                "type": "error",
                "error": {
                    "type": "invalid_request_error",
                    "message": "Model demo/gpt-4.1 is configured for downstream formats openai_chat, not claude_chat",
                },
            },
            response.get_json(),
        )

    def test_messages_route_allows_claude_target(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/messages",
            model_list=("claude-sonnet-4-5",),
            target_format="claude_chat",
        )
        app = Flask(__name__)
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService(
                (
                    app.response_class(
                        '{"id":"msg_1","type":"message"}',
                        status=200,
                        mimetype="application/json",
                    ),
                    200,
                    None,
                )
            ),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/messages",
            json={
                "model": "demo/claude-sonnet-4-5",
                "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
            },
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(200, response.status_code)
        self.assertEqual('{"id":"msg_1","type":"message"}', response.get_data(as_text=True))

    def test_messages_route_returns_claude_style_error_payload_for_upstream_failures(self) -> None:
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/messages",
            model_list=("claude-sonnet-4-5",),
            target_format="claude_chat",
        )
        logger = FakeLogger()
        app = Flask(__name__)
        ctx = AppContext(
            logger=logger,
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=app,
        )
        ProxyController(
            ctx,
            StubProxyService(
                (
                    None,
                    502,
                    ProxyErrorInfo(
                        message="Upstream Claude request failed: overloaded",
                        status_code=502,
                        error_type="api_error",
                        error_code="upstream_overloaded",
                    ),
                )
            ),
            FakeUserService(),
            FakeLogService(),
            FakeProviderManager(provider),
        )

        response = app.test_client().post(
            "/v1/messages",
            json={
                "model": "demo/claude-sonnet-4-5",
                "max_tokens": 256,
                "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
            },
            environ_base={"REMOTE_ADDR": "127.0.0.1"},
        )

        self.assertEqual(502, response.status_code)
        self.assertEqual(
            {
                "type": "error",
                "error": {
                    "type": "api_error",
                    "message": "Upstream Claude request failed: overloaded",
                },
            },
            response.get_json(),
        )
        self.assertTrue(any("upstream_error=Upstream Claude request failed: overloaded" in msg for msg in logger.messages("error")))

class ProxyServiceErrorLoggingTests(unittest.TestCase):
    def test_proxy_service_logs_upstream_error_payload(self) -> None:
        logger = FakeLogger()
        ctx = AppContext(
            logger=logger,
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=Flask(__name__),
        )
        service = ProxyService(ctx)
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            max_retries=1,
        )
        upstream_response = BufferedUpstreamResponse(
            StaticUpstreamResponse(
                status_code=429,
                headers={"Content-Type": "application/json"},
            ),
            b'{"error":{"message":"Rate limit reached","type":"rate_limit_error","code":"rate_limit_exceeded"}}',
        )
        service._open_upstream_response = lambda *args, **kwargs: (upstream_response, False, 429)  # type: ignore[method-assign]

        response, status_code, failure_info = service.proxy_request(
            provider,
            {"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            {},
        )

        self.assertIsNone(failure_info)
        self.assertEqual(429, status_code)
        self.assertIsNotNone(response)
        self.assertIn("Rate limit reached", response.get_data(as_text=True))
        self.assertTrue(
            any(
                "Rate limit reached (type=rate_limit_error, code=rate_limit_exceeded)" in msg
                for msg in logger.messages("warning")
            )
        )


class ProxyServiceRetryHookContextTests(unittest.TestCase):
    def _build_service(self) -> ProxyService:
        ctx = AppContext(
            logger=FakeLogger(),
            config_manager=FakeConfigManager(),
            root_path=Path(__file__).resolve().parents[1],
            flask_app=Flask(__name__),
        )
        return ProxyService(ctx)

    @staticmethod
    def _chat_completion_response(status_code: int, body: bytes) -> tuple[BufferedUpstreamResponse, bool, int]:
        response = BufferedUpstreamResponse(
            StaticUpstreamResponse(
                status_code=status_code,
                headers={"Content-Type": "application/json"},
            ),
            body,
        )
        return response, False, status_code

    @staticmethod
    def _success_response() -> tuple[BufferedUpstreamResponse, bool, int]:
        return ProxyServiceRetryHookContextTests._chat_completion_response(
            200,
            (
                b'{"id":"chatcmpl_1","object":"chat.completion","created":123,"model":"gpt-4.1",'
                b'"choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],'
                b'"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}'
            ),
        )

    def test_hook_sees_previous_retryable_status_code_on_next_attempt(self) -> None:
        service = self._build_service()
        hook = ContextRecordingHook()
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            hook=hook,
            max_retries=2,
        )
        attempts = iter(
            [
                self._chat_completion_response(
                    429,
                    b'{"error":{"message":"Rate limit reached","type":"rate_limit_error","code":"rate_limit_exceeded"}}',
                ),
                self._success_response(),
            ]
        )
        service._open_upstream_response = lambda *args, **kwargs: next(attempts)  # type: ignore[method-assign]

        response, status_code, failure_info = service.proxy_request(
            provider,
            {"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            {},
        )

        self.assertIsNone(failure_info)
        self.assertEqual(200, status_code)
        self.assertIsNotNone(response)
        self.assertEqual(2, len(hook.contexts))
        self.assertIsNone(hook.contexts[0].last_status_code)
        self.assertIsNone(hook.contexts[0].last_error_type)
        self.assertEqual(429, hook.contexts[1].last_status_code)
        self.assertIsNone(hook.contexts[1].last_error_type)

    def test_hook_sees_timeout_error_type_on_next_attempt(self) -> None:
        service = self._build_service()
        hook = ContextRecordingHook()
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            hook=hook,
            max_retries=2,
        )
        calls = {"count": 0}

        def stub_open_upstream_response(*args, **kwargs):
            del args, kwargs
            if calls["count"] == 0:
                calls["count"] += 1
                raise requests.exceptions.Timeout("timed out")
            return self._success_response()

        service._open_upstream_response = stub_open_upstream_response  # type: ignore[method-assign]

        response, status_code, failure_info = service.proxy_request(
            provider,
            {"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            {},
        )

        self.assertIsNone(failure_info)
        self.assertEqual(200, status_code)
        self.assertIsNotNone(response)
        self.assertEqual(2, len(hook.contexts))
        self.assertIsNone(hook.contexts[1].last_status_code)
        self.assertEqual(HookErrorType.TIMEOUT, hook.contexts[1].last_error_type)

    def test_hook_sees_connection_error_type_on_next_attempt(self) -> None:
        service = self._build_service()
        hook = ContextRecordingHook()
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            hook=hook,
            max_retries=2,
        )
        calls = {"count": 0}

        def stub_open_upstream_response(*args, **kwargs):
            del args, kwargs
            if calls["count"] == 0:
                calls["count"] += 1
                raise requests.exceptions.ConnectionError("connection lost")
            return self._success_response()

        service._open_upstream_response = stub_open_upstream_response  # type: ignore[method-assign]

        response, status_code, failure_info = service.proxy_request(
            provider,
            {"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            {},
        )

        self.assertIsNone(failure_info)
        self.assertEqual(200, status_code)
        self.assertIsNotNone(response)
        self.assertEqual(2, len(hook.contexts))
        self.assertIsNone(hook.contexts[1].last_status_code)
        self.assertEqual(HookErrorType.CONNECTION_ERROR, hook.contexts[1].last_error_type)

    def test_hook_sees_websocket_error_type_on_next_attempt(self) -> None:
        service = self._build_service()
        hook = ContextRecordingHook()
        provider = LLMProvider(
            name="demo",
            api="wss://example.com/v1/chat/completions",
            transport="websocket",
            model_list=("gpt-4.1",),
            hook=hook,
            max_retries=2,
        )
        calls = {"count": 0}

        def stub_open_upstream_response(*args, **kwargs):
            del args, kwargs
            if calls["count"] == 0:
                calls["count"] += 1
                raise websocket.WebSocketException("socket closed")
            return self._success_response()

        service._open_upstream_response = stub_open_upstream_response  # type: ignore[method-assign]

        response, status_code, failure_info = service.proxy_request(
            provider,
            {"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            {},
        )

        self.assertIsNone(failure_info)
        self.assertEqual(200, status_code)
        self.assertIsNotNone(response)
        self.assertEqual(2, len(hook.contexts))
        self.assertIsNone(hook.contexts[1].last_status_code)
        self.assertEqual(HookErrorType.WEBSOCKET_ERROR, hook.contexts[1].last_error_type)

    def test_non_retryable_status_does_not_trigger_second_attempt(self) -> None:
        service = self._build_service()
        hook = ContextRecordingHook()
        provider = LLMProvider(
            name="demo",
            api="https://example.com/v1/chat/completions",
            model_list=("gpt-4.1",),
            hook=hook,
            max_retries=2,
        )
        attempts = iter(
            [
                self._chat_completion_response(
                    401,
                    b'{"error":{"message":"Unauthorized","type":"invalid_request_error","code":"invalid_api_key"}}',
                )
            ]
        )
        service._open_upstream_response = lambda *args, **kwargs: next(attempts)  # type: ignore[method-assign]

        response, status_code, failure_info = service.proxy_request(
            provider,
            {"model": "demo/gpt-4.1", "messages": [{"role": "user", "content": "hi"}]},
            {},
        )

        self.assertIsNone(failure_info)
        self.assertEqual(401, status_code)
        self.assertIsNotNone(response)
        self.assertEqual(1, len(hook.contexts))
        self.assertIsNone(hook.contexts[0].last_status_code)
        self.assertIsNone(hook.contexts[0].last_error_type)


if __name__ == "__main__":
    unittest.main()