diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index a67c1bb7df8..3130d49a5e2 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -972,3 +972,9 @@ def _get_openai_compatible_provider_info( if dynamic_api_key is None and api_key is not None: dynamic_api_key = api_key return model, custom_llm_provider, dynamic_api_key, api_base + + +def get_provider_affinity_header(api_base: str | None) -> str | None: + from litellm.llms.opencode.transformation import OpenCodeProviderAdapter + + return OpenCodeProviderAdapter().get_provider_affinity_header(api_base) diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index 02c4cd69159..e022839c5c7 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -44,10 +44,10 @@ def _get_value(value: object, key: str) -> object | None: return getattr(value, key, None) -def _get_provider_affinity_header_name(litellm_params: object | None) -> str | None: +def _get_provider_affinity_header_name(litellm_params: object | None, default_header: str | None = None) -> str | None: header: Final = _get_value(litellm_params, "provider_affinity_header") if litellm_params is not None else None if header is None: - return None + return default_header if not isinstance(header, str): raise TypeError("provider_affinity_header must be a string") return validate_provider_affinity_header_name(header) @@ -84,9 +84,9 @@ def get_stable_session_id(litellm_params: object | None) -> str | None: def add_provider_affinity_header( - headers: Mapping[str, object], litellm_params: object | None + headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers - header_name: Final = _get_provider_affinity_header_name(litellm_params) + header_name: Final = _get_provider_affinity_header_name(litellm_params, default_header) if header_name is None or any(key.lower() == header_name.lower() for key in headers): return dict(headers) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a7ca0d5f49e..a692f3f964e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -52,7 +52,9 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.litellm_core_utils.realtime_errors import ( close_after_upstream_handshake_refusal, realtime_error_event, @@ -60,6 +62,7 @@ from litellm.litellm_core_utils.realtime_errors import ( ) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base +from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -2435,13 +2438,13 @@ class BaseLLMHTTPHandler: # Also check for extra_headers in kwargs (from config or direct calls) extra_headers_from_kwargs: Final = kwargs.get("extra_headers", None) # Merge all header sources: forwarded < extra_headers < provider_specific - merged_headers: Final = {} - if forwarded_headers: - merged_headers.update(forwarded_headers) - if extra_headers_from_kwargs: - merged_headers.update(extra_headers_from_kwargs) - if provider_specific_headers: - merged_headers.update(provider_specific_headers) + merged_headers: Final[Mapping[str, object]] = MappingProxyType( + { + **(forwarded_headers or {}), + **(extra_headers_from_kwargs or {}), + **(provider_specific_headers or {}), + } + ) request_params: Final = reduce( delete_nested_value, litellm_params.get("additional_drop_params") or (), @@ -2452,11 +2455,19 @@ class BaseLLMHTTPHandler: and anthropic_messages_provider_config.should_filter_anthropic_beta_headers() ): await aload_beta_headers_config() + affinity_api_base: Final = ( + AnthropicModelInfo.get_api_base(api_base) if custom_llm_provider == "anthropic" else api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) ( headers, api_base, ) = await anthropic_messages_provider_config.avalidate_anthropic_messages_environment( - headers=merged_headers or {}, + headers=( + add_provider_affinity_header(merged_headers, litellm_params, default_affinity_header) + if default_affinity_header is not None + else dict(merged_headers) + ), model=model, messages=messages, optional_params=request_params, diff --git a/litellm/llms/opencode/common_utils.py b/litellm/llms/opencode/common_utils.py new file mode 100644 index 00000000000..65a96411e52 --- /dev/null +++ b/litellm/llms/opencode/common_utils.py @@ -0,0 +1,10 @@ +import re + + +def is_opencode_api_base(api_base: str | None) -> bool: + return bool( + api_base + and re.fullmatch( + r"https://opencode\.ai/(?:zen(?:/go)?|inference/(?:go/)?(?:openai|anthropic))(?:/v1)?/?", api_base + ) + ) diff --git a/litellm/llms/opencode/transformation.py b/litellm/llms/opencode/transformation.py new file mode 100644 index 00000000000..9b62e7aca95 --- /dev/null +++ b/litellm/llms/opencode/transformation.py @@ -0,0 +1,6 @@ +from litellm.llms.opencode.common_utils import is_opencode_api_base + + +class OpenCodeProviderAdapter: + def get_provider_affinity_header(self, api_base: str | None) -> str | None: + return "x-opencode-session" if is_opencode_api_base(api_base) else None diff --git a/litellm/main.py b/litellm/main.py index 32a7c290bb8..327a7a1c4b5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -89,6 +89,7 @@ from litellm.litellm_core_utils.get_litellm_params import ( with_control_options, ) from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_provider_affinity_header, invalid_provider_routing_hints_message, ) from litellm.litellm_core_utils.get_provider_specific_headers import ( @@ -130,6 +131,7 @@ from litellm.llms.bedrock_mantle.chat.claude_transformation import bedrock_mantl from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, http2_enabled from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.vertex_ai.common_utils import ( VertexAIModelRoute, @@ -5990,19 +5992,28 @@ def completion( }, ) litellm_params: Final = with_control_options(requested_litellm_params, control_options) - if litellm_params.get("provider_affinity_header") is not None: + affinity_api_base: Final = ( + str(client.base_url) + if isinstance(client, openai.OpenAI | openai.AsyncOpenAI) + else get_openai_credentials(api_base=api_base).api_base + if custom_llm_provider == "openai" + else api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) + if litellm_params.get("provider_affinity_header") is not None or default_affinity_header is not None: try: headers = add_provider_affinity_header( headers=headers or litellm.headers or MappingProxyType({}), litellm_params=MappingProxyType( { - "provider_affinity_header": litellm_params["provider_affinity_header"], + "provider_affinity_header": litellm_params.get("provider_affinity_header"), "litellm_session_id": kwargs.get("litellm_session_id"), "session_id": kwargs.get("session_id"), "metadata": metadata, "litellm_metadata": kwargs.get("litellm_metadata"), } ), + default_header=default_affinity_header, ) except ValueError as affinity_error: raise litellm.BadRequestError( diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index cdb9a925403..2e9d6a77a62 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -103,6 +103,7 @@ def _target(request: NativeCall, model: str, provider: str | None) -> Target: _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **detectors.COMMON, + Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.response_cache, Feature.MODIFY_PARAMS: detectors.modify_params, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 09eea7cb323..0b526adf4fa 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -30,6 +30,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -39,6 +40,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig from litellm.responses.litellm_completion_transformation.compaction_items import ( with_compaction_items_as_summaries, @@ -1466,6 +1468,12 @@ def responses( # get llm provider logic litellm_params: Final = GenericLiteLLMParams(**kwargs) + affinity_api_base: Final = ( + get_openai_credentials(api_base=litellm_params.api_base).api_base + if custom_llm_provider == "openai" + else litellm_params.api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) try: effective_extra_headers: Final = ( add_provider_affinity_header( @@ -1479,8 +1487,9 @@ def responses( "litellm_metadata": kwargs.get("litellm_metadata"), } ), + default_header=default_affinity_header, ) - if litellm_params.provider_affinity_header is not None + if litellm_params.provider_affinity_header is not None or default_affinity_header is not None else extra_headers ) except ValueError as affinity_error: diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index ebb221e98be..74cd28161a4 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -165,6 +165,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { } MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { + Feature.PROVIDER_AFFINITY, Feature.SAFEGUARDS_CLASSIFIER, Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 691c04f2daa..b53e80805a5 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -5,10 +5,15 @@ from __future__ import annotations from typing import Final import litellm -from litellm.rust_bridge.detectors import overridden, present, streaming +from litellm.llms.opencode.common_utils import is_opencode_api_base +from litellm.rust_bridge.detectors import api_base, overridden, present, streaming from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str +def provider_affinity(request: NativeCall) -> bool: + return is_opencode_api_base(api_base(request)) + + def anthropic_workload_identity(request: NativeCall) -> bool: """Python mints an Anthropic federation token when no static credential is configured.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo diff --git a/litellm/rust_bridge/responses/detectors.py b/litellm/rust_bridge/responses/detectors.py index 392c11f48a8..763b63d6742 100644 --- a/litellm/rust_bridge/responses/detectors.py +++ b/litellm/rust_bridge/responses/detectors.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # tool lists stay from pydantic import TypeAdapter import litellm +from litellm.llms.opencode.common_utils import is_opencode_api_base from litellm.rust_bridge.detectors import api_base, call_callbacks, overridden, present from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -25,7 +26,7 @@ def client_headers(request: NativeCall) -> bool: def provider_affinity(request: NativeCall) -> bool: """Python validates and fills the affinity header from the session ID.""" - return present(request, "provider_affinity_header") + return present(request, "provider_affinity_header") or is_opencode_api_base(api_base(request)) def text_format(request: NativeCall) -> bool: diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index d7474936425..57e66805678 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -8,6 +8,7 @@ from litellm import CustomLLM from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.get_llm_provider_logic import ( get_llm_provider, + get_provider_affinity_header, inferred_provider, is_registered_custom_provider, ) @@ -706,3 +707,11 @@ class TestClaudeModelPatternMatching: set_fallback_generalizations([]) with pytest.raises(litellm.BadRequestError): litellm.get_llm_provider(model="claude-opus-4-9") + + +@pytest.mark.parametrize("prefix", ["openai", "anthropic", "custom_openai", "openai_like"]) +def test_opencode_affinity_resolution_preserves_existing_transport(prefix: str) -> None: + api_base: Final = "https://opencode.ai/zen/go/v1" + model, provider, _, destination = get_llm_provider(model=f"{prefix}/test-model", api_base=api_base) + assert (model, provider, destination) == ("test-model", prefix, api_base) + assert get_provider_affinity_header(destination) == "x-opencode-session" diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 63cc45b29df..8b32ab2e6b8 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5542,3 +5542,85 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ sent: Final = received[0] assert "safeguards" not in json.loads(sent.content) assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("environment_base", [None, "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) +@pytest.mark.parametrize( + ("api_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/inference/go/anthropic", None, "conversation-123"), + ("https://opencode.ai/inference/go/anthropic", "caller-override", "caller-override"), + ("https://other.example", None, None), + ], +) +async def test_messages_opencode_session_header( + stream: bool, + environment_base: str | None, + api_base: str, + explicit_header: str | None, + expected_session: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + if environment_base: + monkeypatch.setenv(environment_base, "https://opencode.ai/inference/go/anthropic") + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "OK"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"event: message_start\ndata: {json.dumps({'type': 'message_start', 'message': payload})}\n\n" + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ) + if stream + else httpx.Response(200, json=payload) + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + result: Final = await litellm.anthropic.messages.acreate( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": "hello"}], + max_tokens=128000, + api_base=None if environment_base and api_base == "https://opencode.ai/inference/go/anthropic" else api_base, + api_key="test-key", + stream=stream, + litellm_session_id="conversation-123", + client=upstream, + extra_headers={"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + ) + if stream: + chunks: Final = b"".join([chunk async for chunk in result]) + assert b'"text": "OK"' in chunks + assert b"message_stop" in chunks + else: + assert result["content"] == [{"type": "text", "text": "OK"}] + finally: + await upstream.close() + sent: Final = received.get_nowait() + assert str(sent.url) == f"{api_base}/v1/messages" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["x-api-key"] == "test-key" + assert json.loads(sent.content) == { + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 128000, + "stream": stream, + } + assert received.empty() diff --git a/tests/unit/llms/opencode/test_transformation.py b/tests/unit/llms/opencode/test_transformation.py new file mode 100644 index 00000000000..96203c23122 --- /dev/null +++ b/tests/unit/llms/opencode/test_transformation.py @@ -0,0 +1,90 @@ +from typing import Final + +import pytest + +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header +from litellm.llms.opencode.transformation import OpenCodeProviderAdapter + + +@pytest.mark.parametrize( + "api_base", + [ + "https://opencode.ai/zen/v1", + "https://opencode.ai/zen/go/v1/", + "https://opencode.ai/inference/openai/v1", + "https://opencode.ai/inference/go/openai/v1", + "https://opencode.ai/inference/go/anthropic/v1", + "https://opencode.ai/inference/anthropic/v1", + ], +) +def test_opencode_defaults_affinity_header_to_the_callers_session(api_base: str) -> None: + headers: Final = {"Authorization": "Bearer test-key"} + params: Final = {"metadata": {"session_id": "conversation-123"}} + assert add_provider_affinity_header( + headers, params, OpenCodeProviderAdapter().get_provider_affinity_header(api_base) + ) == { + **headers, + "x-opencode-session": "conversation-123", + } + assert headers == {"Authorization": "Bearer test-key"} + assert params == {"metadata": {"session_id": "conversation-123"}} + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "https://api.openai.com/v1", + "https://opencode.ai.example/zen/go/v1", + "https://opencode.ai@example.com/zen/go/v1", + "https://example.com/opencode.ai/zen/go/v1", + "https://example.com/v1?upstream=https://opencode.ai/zen/go/v1", + "https://opencode.ai/console/v1", + "http://opencode.ai/zen/go/v1", + ], +) +def test_other_destinations_do_not_receive_an_automatic_session_header(api_base: str | None) -> None: + headers: Final = {"Authorization": "Bearer test-key", "x-existing": "unchanged"} + assert ( + add_provider_affinity_header( + headers, + {"litellm_session_id": "conversation-123"}, + OpenCodeProviderAdapter().get_provider_affinity_header(api_base), + ) + == headers + ) + + +@pytest.mark.parametrize( + "params", + [ + {}, + {"litellm_trace_id": "request-123"}, + {"metadata": {"trace_id": "request-123"}}, + { + "litellm_session_id": "request-123", + "litellm_metadata": {"session_id": "request-123", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, + ], +) +def test_opencode_does_not_invent_a_conversation_id(params: dict[str, object]) -> None: + assert ( + add_provider_affinity_header( + {}, params, OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1") + ) + == {} + ) + + +def test_opencode_preserves_explicit_header_and_affinity_configuration() -> None: + assert add_provider_affinity_header( + {"X-OpenCode-Session": "explicit"}, + {"litellm_session_id": "conversation-123"}, + OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1"), + ) == {"X-OpenCode-Session": "explicit"} + assert add_provider_affinity_header( + {}, + {"litellm_session_id": "conversation-123", "provider_affinity_header": "x-custom-session"}, + OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1"), + ) == {"x-custom-session": "conversation-123"} diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index 436367bd82d..36b5d8f135b 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -4,6 +4,7 @@ over the wire and surface provider errors correctly. Expected JSON bodies are st in expected_responses_api_request/. """ +import asyncio import copy import json import logging @@ -1231,3 +1232,110 @@ def test_faked_stream_accepts_a_text_timeout(monkeypatch: pytest.MonkeyPatch) -> assert [event.response.output[0].content[0].text for event in response if event.type == "response.completed"] == [ "Done." ] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("api_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/inference/go/openai/v1", None, "conversation-123"), + ("https://opencode.ai/inference/go/openai/v1", "caller-override", "caller-override"), + ("https://other.example/v1", None, None), + ], +) +@pytest.mark.asyncio +async def test_responses_opencode_session_header( + asynchronous: bool, stream: bool, api_base: str, explicit_header: str | None, expected_session: str | None +) -> None: + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = _minimal_responses_api_payload("resp_test", "gpt-6-luna") + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': payload})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + if stream + else httpx.Response(200, json=payload) + ) + + args: Final = { + "model": "openai/gpt-6-luna", + "input": "hello", + "api_base": api_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + "stream": stream, + "extra_headers": {"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + } + if asynchronous: + async_client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + response: Final = await litellm.aresponses(**args, client=async_client) + if stream: + events: Final = tuple([event.type async for event in response]) + assert "response.completed" in events + else: + assert response.output[0].content[0].text == "Done." + finally: + await async_client.close() + else: + sync_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + try: + result: Final = litellm.responses(**args, client=sync_client) + if stream: + assert "response.completed" in tuple(event.type for event in result) + else: + assert result.output[0].content[0].text == "Done." + finally: + sync_client.close() + sent: Final = received.get_nowait() + assert str(sent.url) == f"{api_base}/responses" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "gpt-6-luna", "input": "hello", "stream": stream} + assert received.empty() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("source", ["OPENAI_BASE_URL", "OPENAI_API_BASE", "litellm.api_base"]) +@pytest.mark.parametrize("explicit_base", [None, "https://other.example/v1"]) +@pytest.mark.asyncio +async def test_responses_opencode_session_uses_effective_destination( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + asynchronous: bool, + source: str, + explicit_base: str | None, +) -> None: + default_base: Final = "https://opencode.ai/inference/go/openai/v1" + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", default_base if source == "litellm.api_base" else None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if source != "litellm.api_base": + monkeypatch.setenv(source, default_base) + endpoint: Final = respx_mock.post(f"{explicit_base or default_base}/responses").respond( + json=_minimal_responses_api_payload("resp_test", "gpt-6-luna") + ) + args: Final = { + "model": "openai/gpt-6-luna", + "input": "hello", + "api_base": explicit_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + } + response: Final = await litellm.aresponses(**args) if asynchronous else litellm.responses(**args) + assert response.output[0].content[0].text == "Done." + assert endpoint.call_count == 1 + sent: Final = endpoint.calls[0].request + assert sent.headers.get("x-opencode-session") == (None if explicit_base else "conversation-123") + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "gpt-6-luna", "input": "hello"} diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 726128386d5..78722458cf3 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -116,6 +116,9 @@ def _wrap_nested_level(value: object, depth: int) -> object: (responses_detectors.provider_affinity, {"provider_affinity_header": ""}, True), (responses_detectors.provider_affinity, {"provider_affinity_header": None, "session_id": "session-1"}, False), (responses_detectors.provider_affinity, {}, False), + (responses_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/openai/v1"}, True), + (messages_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/anthropic/v1"}, True), + (messages_detectors.provider_affinity, {"api_base": "https://other.example/v1"}, False), (detectors.streaming, {"stream": True}, True), (detectors.streaming, {"stream": False}, False), (detectors.timeout_variants, {"request_timeout": 5}, True), diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index c3d3575b8b3..70b268bf2a2 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -20,7 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import APITimeoutError, AsyncOpenAI +from openai import APITimeoutError, AsyncOpenAI, OpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion from pydantic import TypeAdapter @@ -95,6 +95,142 @@ def preserve_litellm_completion_state(monkeypatch: pytest.MonkeyPatch) -> None: WHITE_PNG: Final = (Path(__file__).parents[1] / "white_100x100.png").read_bytes() +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("client_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/zen/go/v1", None, "conversation-123"), + ("https://opencode.ai/zen/go/v1", "caller-override", "caller-override"), + ("https://other.example/v1", None, None), + ], +) +@pytest.mark.asyncio +async def test_completion_sends_session_only_to_the_actual_opencode_destination( + asynchronous: bool, + stream: bool, + client_base: str, + explicit_header: str | None, + expected_session: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://opencode.ai/zen/go/v1") + monkeypatch.setenv("OPENAI_API_BASE", "https://opencode.ai/zen/go/v1") + monkeypatch.setattr(litellm, "api_base", "https://opencode.ai/zen/go/v1") + requests: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + requests.put_nowait(request) + payload: Final = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk" if stream else "chat.completion", + "created": 1, + "model": "session-test", + "choices": [ + { + "index": 0, + "delta" if stream else "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + } + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(payload)}\n\ndata: [DONE]\n\n", + ) + if stream + else httpx.Response(200, json=payload) + ) + + args: Final = { + "model": "openai/session-test", + "api_base": "https://opencode.ai/zen/go/v1", + "messages": [{"role": "user", "content": "hello"}], + "litellm_session_id": "conversation-123", + "extra_headers": {"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + "stream": stream, + } + if asynchronous: + async with AsyncOpenAI( + api_key="test-key", + base_url=client_base, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) as async_client: + response: Final = await acompletion(**args, client=async_client) + content: Final = ( + "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + if stream + else response.choices[0].message.content + ) + assert content == "OK" + else: + with OpenAI( + api_key="test-key", base_url=client_base, http_client=httpx.Client(transport=httpx.MockTransport(respond)) + ) as sync_client: + result: Final = completion(**args, client=sync_client) + text: Final = ( + "".join(chunk.choices[0].delta.content or "" for chunk in result) + if stream + else result.choices[0].message.content + ) + assert text == "OK" + sent: Final = requests.get_nowait() + assert str(sent.url) == f"{client_base}/chat/completions" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == { + "model": "session-test", + "messages": args["messages"], + **({"stream": True} if stream else {}), + } + assert requests.empty() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("source", ["OPENAI_BASE_URL", "OPENAI_API_BASE", "litellm.api_base"]) +@pytest.mark.parametrize("explicit_base", [None, "https://other.example/v1"]) +@pytest.mark.asyncio +async def test_completion_opencode_session_uses_effective_destination( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + asynchronous: bool, + source: str, + explicit_base: str | None, +) -> None: + default_base: Final = "https://opencode.ai/inference/go/openai/v1" + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", default_base if source == "litellm.api_base" else None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if source != "litellm.api_base": + monkeypatch.setenv(source, default_base) + endpoint: Final = respx_mock.post(f"{explicit_base or default_base}/chat/completions").respond( + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "session-test", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + } + ) + args: Final = { + "model": "openai/session-test", + "messages": [{"role": "user", "content": "hello"}], + "api_base": explicit_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + } + response: Final = await litellm.acompletion(**args) if asynchronous else litellm.completion(**args) + assert response.choices[0].message.content == "OK" + assert endpoint.call_count == 1 + sent: Final = endpoint.calls[0].request + assert sent.headers.get("x-opencode-session") == (None if explicit_base else "conversation-123") + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "session-test", "messages": args["messages"]} + + @pytest.fixture def openai_api_response(): mock_response_data = {