import json from collections.abc import AsyncIterator from typing import Any from unittest.mock import MagicMock, patch import pytest from fastapi.responses import JSONResponse, StreamingResponse from free_claude_code.api.handlers import ( MessagesHandler, ResponsesHandler, TokenCountHandler, ) from free_claude_code.application.errors import InvalidRequestError from free_claude_code.config.settings import Settings from free_claude_code.core.anthropic.models import ( Message, MessagesRequest, TokenCountRequest, ) from free_claude_code.core.anthropic.streaming import format_sse_event from free_claude_code.core.failures import ExecutionFailure, FailureKind from free_claude_code.core.openai_responses import OpenAIResponsesRequest _CLASSIFIER_SYSTEM = ( "You are a security monitor. Respond with yes or no." ) _CLASSIFIER_USER = ( "\nUser: review the repo\nWebFetch https://example.com: fetch\n" "\n immediately." ) class FakeProvider: def __init__(self, events: list[str] | None = None) -> None: self.preflight_calls: list[tuple[MessagesRequest, bool | None]] = [] self.requests: list[MessagesRequest] = [] self.stream_kwargs: list[dict[str, Any]] = [] self.events = events or [ 'event: message_start\ndata: {"type":"message_start"}\n\n', 'event: message_stop\ndata: {"type":"message_stop"}\n\n', ] def preflight_stream( self, request: MessagesRequest, *, thinking_enabled: bool | None = None ) -> None: self.preflight_calls.append((request, thinking_enabled)) async def cleanup(self) -> None: return None async def list_model_ids(self) -> frozenset[str]: return frozenset({"test-model"}) async def stream_response( self, request: MessagesRequest, input_tokens: int = 0, *, request_id: str | None = None, thinking_enabled: bool | None = None, ) -> AsyncIterator[str]: self.requests.append(request) self.stream_kwargs.append( { "input_tokens": input_tokens, "request_id": request_id, "thinking_enabled": thinking_enabled, } ) for event in self.events: yield event async def _streaming_body_text(response: StreamingResponse) -> str: parts: list[str] = [] async for chunk in response.body_iterator: if isinstance(chunk, bytes): parts.append(chunk.decode("utf-8")) else: parts.append(str(chunk)) return "".join(parts) def _json_response_content(response: JSONResponse) -> dict[str, Any]: content = json.loads(bytes(response.body).decode("utf-8")) assert isinstance(content, dict) return content def _trace_events(trace_mock: MagicMock, event: str) -> list[dict[str, Any]]: return [ dict(call.kwargs) for call in trace_mock.call_args_list if call.kwargs.get("event") == event ] @pytest.mark.asyncio async def test_messages_handler_passes_routed_request_and_stream_metadata() -> None: provider = FakeProvider() handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, messages=[Message(role="user", content="hi")], ) response = await handler.create(request) assert isinstance(response, StreamingResponse) body = await _streaming_body_text(response) assert "message_start" in body assert provider.requests[0].model == "test-model" assert provider.stream_kwargs[0]["input_tokens"] > 0 assert provider.stream_kwargs[0]["request_id"].startswith("req_") assert provider.stream_kwargs[0]["thinking_enabled"] is True assert len(provider.preflight_calls) == 1 @pytest.mark.asyncio @pytest.mark.parametrize("stream", [True, False]) async def test_messages_handler_preflight_invalid_request_stays_http_error( stream: bool, ) -> None: class RejectPreflightProvider(FakeProvider): def preflight_stream( self, request: MessagesRequest, *, thinking_enabled: bool | None = None, ) -> None: raise InvalidRequestError("bad tool shape") provider = RejectPreflightProvider() handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, messages=[Message(role="user", content="hi")], stream=stream, ) with pytest.raises(InvalidRequestError): await handler.create(request) @pytest.mark.asyncio async def test_messages_handler_aggregates_provider_stream_when_stream_false() -> None: provider = FakeProvider( [ format_sse_event( "message_start", { "type": "message_start", "message": { "id": "msg_test", "type": "message", "role": "assistant", "content": [], "model": "test-model", "stop_reason": None, "stop_sequence": None, "usage": {"input_tokens": 7, "output_tokens": 1}, }, }, ), format_sse_event( "content_block_start", { "type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}, }, ), format_sse_event( "content_block_delta", { "type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "OK"}, }, ), format_sse_event( "content_block_stop", {"type": "content_block_stop", "index": 0} ), format_sse_event( "message_delta", { "type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"input_tokens": 7, "output_tokens": 2}, }, ), format_sse_event("message_stop", {"type": "message_stop"}), ] ) handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, stream=False, messages=[Message(role="user", content="hi")], ) response = await handler.create(request) assert isinstance(response, JSONResponse) assert response.headers["content-type"].startswith("application/json") body = _json_response_content(response) assert body["id"] == "msg_test" assert body["type"] == "message" assert body["role"] == "assistant" assert body["model"] == "test-model" assert body["content"] == [{"type": "text", "text": "OK"}] assert body["stop_reason"] == "end_turn" assert body["usage"] == {"input_tokens": 7, "output_tokens": 2} @pytest.mark.asyncio async def test_messages_handler_returns_error_json_for_stream_false_sse_error() -> None: provider = FakeProvider( [ format_sse_event( "error", { "type": "error", "error": {"type": "api_error", "message": "upstream failed"}, }, ) ] ) handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, stream=False, messages=[Message(role="user", content="hi")], ) response = await handler.create(request) assert isinstance(response, JSONResponse) assert response.status_code == 500 assert response.headers["x-should-retry"] == "false" body = _json_response_content(response) assert body["type"] == "error" assert body["error"] == {"type": "api_error", "message": "upstream failed"} assert body["request_id"].startswith("req_") @pytest.mark.asyncio async def test_messages_handler_discards_partial_stream_false_output_on_error() -> None: provider = FakeProvider( [ format_sse_event( "message_start", { "type": "message_start", "message": { "id": "msg_partial", "type": "message", "role": "assistant", "content": [], "model": "test-model", "stop_reason": None, "stop_sequence": None, "usage": {"input_tokens": 1, "output_tokens": 1}, }, }, ), format_sse_event( "content_block_start", { "type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}, }, ), format_sse_event( "content_block_delta", { "type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "incomplete"}, }, ), format_sse_event( "error", { "type": "error", "error": { "type": "overloaded_error", "message": "provider overloaded", }, }, ), ] ) handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, stream=False, messages=[Message(role="user", content="hi")], ) response = await handler.create(request) assert isinstance(response, JSONResponse) assert response.status_code == 529 assert response.headers["x-should-retry"] == "false" body = _json_response_content(response) assert body["error"] == { "type": "overloaded_error", "message": "provider overloaded", } assert "content" not in body @pytest.mark.asyncio async def test_messages_handler_stream_false_provider_exception_keeps_status() -> None: class FailingProvider(FakeProvider): async def stream_response( self, request: Any, input_tokens: int = 0, *, request_id: str | None = None, thinking_enabled: bool | None = None, ) -> AsyncIterator[str]: self.requests.append(request) self.stream_kwargs.append( { "input_tokens": input_tokens, "request_id": request_id, "thinking_enabled": thinking_enabled, } ) raise ExecutionFailure( kind=FailureKind.RATE_LIMIT, status_code=429, message="upstream is busy", retryable=True, ) yield "unreachable" provider = FailingProvider() handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, stream=False, messages=[Message(role="user", content="hi")], ) response = await handler.create(request) assert isinstance(response, JSONResponse) assert response.status_code == 429 assert response.headers["x-should-retry"] == "false" body = _json_response_content(response) assert body["error"] == { "type": "rate_limit_error", "message": "upstream is busy", } @pytest.mark.asyncio async def test_messages_handler_forces_no_thinking_for_safety_classifier() -> None: provider = FakeProvider() handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, system=_CLASSIFIER_SYSTEM, messages=[Message(role="user", content=_CLASSIFIER_USER)], ) with patch("free_claude_code.api.handlers.messages.trace_event") as trace_mock: response = await handler.create(request) assert isinstance(response, StreamingResponse) await _streaming_body_text(response) assert provider.preflight_calls[0][1] is False assert provider.stream_kwargs[0]["thinking_enabled"] is False assert provider.requests[0].model == "test-model" assert provider.requests[0].system == _CLASSIFIER_SYSTEM assert _trace_events( trace_mock, "free_claude_code.api.optimization.safety_classifier_no_thinking" ) == [ { "stage": "routing", "event": "free_claude_code.api.optimization.safety_classifier_no_thinking", "source": "api", "model": "test-model", "changed": True, } ] @pytest.mark.asyncio async def test_messages_handler_preserves_thinking_for_non_classifier() -> None: provider = FakeProvider() handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, system="Explain XML formats.", messages=[ Message( role="user", content=( "Explain ... and a tag " "without making a verdict." ), ) ], ) with patch("free_claude_code.api.handlers.messages.trace_event") as trace_mock: response = await handler.create(request) assert isinstance(response, StreamingResponse) await _streaming_body_text(response) assert provider.preflight_calls[0][1] is True assert provider.stream_kwargs[0]["thinking_enabled"] is True assert ( _trace_events( trace_mock, "free_claude_code.api.optimization.safety_classifier_no_thinking", ) == [] ) @pytest.mark.asyncio async def test_messages_handler_keeps_existing_no_thinking_for_classifier() -> None: provider = FakeProvider() handler = MessagesHandler(Settings(), provider_resolver=lambda _: provider) request = MessagesRequest( model="claude-3-freecc-no-thinking/nvidia_nim/test-model", max_tokens=100, system=_CLASSIFIER_SYSTEM, messages=[Message(role="user", content=_CLASSIFIER_USER)], ) with patch("free_claude_code.api.handlers.messages.trace_event") as trace_mock: response = await handler.create(request) assert isinstance(response, StreamingResponse) await _streaming_body_text(response) assert provider.preflight_calls[0][1] is False assert provider.stream_kwargs[0]["thinking_enabled"] is False assert _trace_events( trace_mock, "free_claude_code.api.optimization.safety_classifier_no_thinking" ) == [ { "stage": "routing", "event": "free_claude_code.api.optimization.safety_classifier_no_thinking", "source": "api", "model": "test-model", "changed": False, } ] @pytest.mark.asyncio async def test_messages_handler_optimization_intercepts_before_provider_execution() -> ( None ): provider_resolver = MagicMock() handler = MessagesHandler(Settings(), provider_resolver=provider_resolver) request = MessagesRequest( model="nvidia_nim/test-model", max_tokens=100, messages=[Message(role="user", content="quota check")], ) optimized = object() with patch( "free_claude_code.api.handlers.messages.try_optimizations", return_value=optimized, ): assert await handler.create(request) is optimized provider_resolver.assert_not_called() @pytest.mark.asyncio async def test_responses_handler_bypasses_message_only_optimizations() -> None: provider = FakeProvider() handler = ResponsesHandler(Settings(), provider_resolver=lambda _: provider) with patch( "free_claude_code.api.handlers.messages.try_optimizations", side_effect=AssertionError("Responses must not use message optimizations"), ): response = await handler.create( OpenAIResponsesRequest( model="nvidia_nim/test-model", input="quota check", ) ) assert isinstance(response, StreamingResponse) body = await _streaming_body_text(response) assert "response.completed" in body assert provider.requests[0].messages[0].content == "quota check" @pytest.mark.asyncio async def test_responses_handler_does_not_apply_safety_classifier_policy() -> None: provider = FakeProvider() handler = ResponsesHandler(Settings(), provider_resolver=lambda _: provider) with patch("free_claude_code.api.handlers.messages.trace_event") as trace_mock: response = await handler.create( OpenAIResponsesRequest( model="nvidia_nim/test-model", input=_CLASSIFIER_USER, instructions=_CLASSIFIER_SYSTEM, ) ) assert isinstance(response, StreamingResponse) await _streaming_body_text(response) assert provider.preflight_calls[0][1] is True assert provider.stream_kwargs[0]["thinking_enabled"] is True assert ( _trace_events( trace_mock, "free_claude_code.api.optimization.safety_classifier_no_thinking", ) == [] ) def test_token_count_handler_routes_and_counts_tokens() -> None: handler = TokenCountHandler( Settings(), token_counter=lambda messages, system, tools: len(messages) + 41, ) with patch("free_claude_code.api.handlers.token_count.trace_event") as trace: response = handler.count( TokenCountRequest( model="nvidia_nim/test-model", messages=[Message(role="user", content="hi")], ), request_id="req_ingress", ) assert response.input_tokens == 42 assert all( call.kwargs["request_id"] == "req_ingress" for call in trace.call_args_list )