scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

fix(opencode): forward stable sessions automatically across inference routes (#48121)

Browse source 
Forward caller-supplied stable session metadata to recognized OpenCode destinations across Chat Completions, Responses, and native Messages. Preserve explicit headers, configured affinity headers, and existing transports.

Fixes #39503
This commit is contained in:
moe-berri 2026-10-23 17:38:47 -07:00 • committed by GitHub
parent f6b2086502
commit 9fab92672a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 507 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -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)

8
litellm/litellm_core_utils/provider_affinity.py
View file

@ -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)

27
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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,

10
litellm/llms/opencode/common_utils.py Normal file
View file

@ -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
)
)

6
litellm/llms/opencode/transformation.py Normal file
View file

@ -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

15
litellm/main.py
View file

@ -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(

1
litellm/messages/dispatch.py
View file

@ -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,

11
litellm/responses/main.py
View file

@ -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:

1
litellm/rust_bridge/catalog.py
View file

@ -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,

7
litellm/rust_bridge/messages/detectors.py
View file

@ -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

3
litellm/rust_bridge/responses/detectors.py
View file

@ -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:

9
tests/unit/litellm_core_utils/test_get_llm_provider_logic.py
View file

@ -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"

82
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -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()

90
tests/unit/llms/opencode/test_transformation.py Normal file
View file

@ -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"}

108
tests/unit/responses/test_responses_api_request_body.py
View file

@ -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"}

3
tests/unit/rust_bridge/test_detectors.py
View file

@ -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),

138
tests/unit/test_main.py
View file

@ -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 = {