litellm/tests/integration/providers/test_databricks_chat_wire.py

import itertools
import json
import uuid
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final, Literal, TypeAlias

import anthropic
import openai
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from integration.translation.chat_completions.cross_model.contracts import assert_completion
from integration.translation.response_contracts import responses_structure
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter

_BACKEND: Final = "databricks-glm-5-2"
_API_KEY: Final = "synthetic-databricks-key"
_PROMPT: Final = "Summarise the cached briefing in one sentence."
_JSON_SCHEMA_BACKEND: Final = "databricks-qwen35-122b-a10b"
_CLAUDE_BACKEND: Final = "databricks-claude-haiku-4-5"
_JSON_SCHEMA_PROMPT: Final = "Return the requested JSON."
_JSON_CONTENT: Final = '{"p":{"name":"Ada"}}'
_PROVIDER_USAGE: Final[Mapping[str, JsonValue]] = {
    "prompt_tokens": 12011,
    "completion_tokens": 8,
    "total_tokens": 12019,
    "cache_read_input_tokens": 12002,
    "cache_creation_input_tokens": 0,
}
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_PERSON_SCHEMA: Final[Mapping[str, JsonValue]] = {
    "$defs": {
        "Person": {
            "type": "object",
            "properties": {"name": {"type": "string"}},
            "required": ["name"],
        }
    },
    "type": "object",
    "properties": {"p": {"$ref": "#/$defs/Person"}},
    "required": ["p"],
}
_JSON_SCHEMA_RESPONSE_FORMAT: Final[Mapping[str, JsonValue]] = {
    "type": "json_schema",
    "json_schema": {"name": "P", "strict": True, "schema": _PERSON_SCHEMA},
}
_MESSAGES_PERSON_SCHEMA: Final[Mapping[str, JsonValue]] = {
    "$defs": {
        "Person": {
            "type": "object",
            "properties": {"name": {"type": "string"}},
            "required": ["name"],
            "additionalProperties": False,
        }
    },
    "type": "object",
    "properties": {"p": {"$ref": "#/$defs/Person"}},
    "required": ["p"],
    "additionalProperties": False,
}
_MESSAGES_RESPONSE_FORMAT: Final[Mapping[str, JsonValue]] = {
    "type": "json_schema",
    "json_schema": {
        "name": "structured_output",
        "schema": _MESSAGES_PERSON_SCHEMA,
        "strict": True,
    },
}
_JSON_TOOL_CALL_MESSAGE: Final[dict[str, JsonValue]] = {
    "role": "assistant",
    "content": None,
    "tool_calls": [
        {
            "id": "json-tool-call",
            "type": "function",
            "function": {"name": "json_tool_call", "arguments": _JSON_CONTENT},
        }
    ],
}


def test_databricks_responses_reasoning_summary_keeps_reasoning_effort_a_string(gateway: Gateway) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST" and request.target == "/chat/completions"
        assert _JSON_OBJECT.validate_json(request.body) == {
            "model": _BACKEND,
            "messages": [{"role": "user", "content": "Say hello."}],
            "max_tokens": 64,
            "reasoning_effort": "low",
        }
        return Reply(body=_chat_completion_response(_BACKEND, {"role": "assistant", "content": "Hello."}, "stop"))

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        alias: Final = scenario.model(model=f"databricks/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/responses",
            {
                "model": alias,
                "input": "Say hello.",
                "max_output_tokens": 64,
                "reasoning": {"effort": "low", "summary": "auto"},
            },
        )
        assert response.status_code == 200, response.text
        body: Final = _JSON_OBJECT.validate_json(response.content)
        responses_structure(body)
        assert body["usage"] == {
            "input_tokens": 1,
            "output_tokens": 1,
            "total_tokens": 2,
            "input_tokens_details": None,
            "output_tokens_details": None,
            "cost": None,
        }
        assert len(wire.drain()) == 1


def _chat_completion_response(model: str, message: dict[str, JsonValue], finish_reason: str) -> bytes:
    response: Final = {
        "id": "databricks-json-schema-response",
        "object": "chat.completion",
        "created": 1,
        "model": model,
        "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
        "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
    }
    return json.dumps(response).encode()


class _PromptTokensDetails(BaseModel):
    model_config = ConfigDict(extra="ignore")
    cached_tokens: int | None = None


class _Usage(BaseModel):
    model_config = ConfigDict(extra="ignore")
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    prompt_tokens_details: _PromptTokensDetails | None = None


class _Delta(BaseModel):
    model_config = ConfigDict(extra="ignore")
    content: str | None = None


class _Choice(BaseModel):
    model_config = ConfigDict(extra="ignore")
    delta: _Delta


class _Chunk(BaseModel):
    model_config = ConfigDict(extra="ignore")
    id: str
    choices: tuple[_Choice, ...]
    usage: _Usage | None = None


def _frame(
    identity: str,
    choices: list[Mapping[str, object]],
    usage: Mapping[str, JsonValue] | None = None,
    model: str = _BACKEND,
) -> bytes:
    value: Final = {
        "id": identity,
        "object": "chat.completion.chunk",
        "created": 1,
        "model": model,
        "choices": choices,
        **({} if usage is None else {"usage": usage}),
    }
    return b"data: " + json.dumps(value).encode() + b"\n\n"


@pytest.mark.covers("other.provider_wire.databricks.stream_usage_and_cache_reads_reach_client_and_spend_log")
def test_databricks_stream_final_usage_chunk_reaches_client_and_spend_log(gateway: Gateway) -> None:
    identity: Final = f"databricks-stream-{uuid.uuid4().hex}"
    frames: Final = (
        _frame(
            identity, [{"index": 0, "delta": {"role": "assistant", "content": "The briefing "}, "finish_reason": None}]
        ),
        _frame(identity, [{"index": 0, "delta": {"content": "is short."}, "finish_reason": None}]),
        _frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}]),
        _frame(identity, [], usage=_PROVIDER_USAGE),
        b"data: [DONE]\n\n",
    )

    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/chat/completions"
        assert request.headers["authorization"] == f"Bearer {_API_KEY}"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["model"] == _BACKEND
        assert body["messages"] == [{"role": "user", "content": _PROMPT}]
        assert body["stream"] is True
        return Reply(content_type="text/event-stream", chunks=frames)

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        with gateway.client.stream(
            "POST",
            "/v1/chat/completions",
            json={
                "model": model,
                "messages": [{"role": "user", "content": _PROMPT}],
                "stream": True,
                "stream_options": {"include_usage": True},
            },
            headers={"Authorization": f"Bearer {gateway.key}"},
        ) as response:
            assert response.status_code == 200, response.read()
            lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
        assert lines[-1] == "data: [DONE]", lines
        chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
        assert {chunk.id for chunk in chunks} == {identity}
        assert (
            "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
            == "The briefing is short."
        )
        usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
        assert len(usages) == 1, lines
        assert (
            usages[0].prompt_tokens,
            usages[0].completion_tokens,
            usages[0].total_tokens,
            usages[0].prompt_tokens_details.cached_tokens if usages[0].prompt_tokens_details is not None else None,
        ) == (12011, 8, 12019, 12002), lines
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
        rows: Final = eventually(
            lambda: read_rows(
                'SELECT prompt_tokens, completion_tokens, total_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
                (identity,),
            ),
            lambda values: len(values) == 1,
            seconds=70,
        )
        assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["total_tokens"]) == (12011, 8, 12019)


def _shape_reply(identity: str, stream: bool) -> Reply:
    if stream:
        chunks: Final = (
            _frame(
                identity,
                [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}],
            ),
            _frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}]),
            b"data: [DONE]\n\n",
        )
        return Reply(content_type="text/event-stream", chunks=chunks)

    body: Final = {
        "id": identity,
        "object": "chat.completion",
        "created": 1,
        "model": _BACKEND,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "ok"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
    }
    return Reply(body=json.dumps(body).encode())


@pytest.mark.parametrize(
    "case,endpoint,provider_model,request_body,expected_messages,responses_mode,upstream_path",
    [
        pytest.param(
            "chat_nonstream",
            "/v1/chat/completions",
            "databricks-gemma-3-12b",
            {
                "messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
                "response_format": {"type": "json_object"},
            },
            [{"role": "user", "content": "Reply in JSON"}],
            False,
            "/chat/completions",
            id="A1-chat-nonstream",
        ),
        pytest.param(
            "chat_stream",
            "/v1/chat/completions",
            "databricks-gemma-3-12b",
            {
                "messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
                "response_format": {"type": "json_object"},
                "stream": True,
            },
            [{"role": "user", "content": "Reply in JSON"}],
            False,
            "/chat/completions",
            id="A2-chat-stream",
        ),
        pytest.param(
            "responses-mode-bridge",
            "/v1/chat/completions",
            "databricks-qwen3-next-80b-a3b-instruct",
            {
                "messages": [{"role": "user", "content": "Reply in JSON"}],
                "response_format": {"type": "json_object"},
            },
            [{"role": "user", "content": "Reply in JSON"}],
            True,
            "/chat/completions",
            id="A3-responses-mode",
        ),
        pytest.param(
            "responses-mode-bridge-unity-catalog",
            "/v1/chat/completions",
            "system.ai.deepseek-v4-1-flash",
            {
                "messages": [{"role": "user", "content": "Reply in JSON with key a. hi"}],
                "response_format": {"type": "json_object"},
            },
            [{"role": "user", "content": "Reply in JSON with key a. hi"}],
            True,
            "/ai-gateway/mlflow/v1/chat/completions",
            id="A3b-responses-mode-unity-catalog-model",
        ),
        pytest.param(
            "messages-endpoint",
            "/v1/messages",
            "databricks-gemma-3-12b",
            {
                "max_tokens": 32,
                "messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
            },
            [{"role": "user", "content": "Reply in JSON"}],
            False,
            "/chat/completions",
            id="A4-messages",
        ),
        pytest.param(
            "responses-endpoint",
            "/v1/responses",
            "databricks-gemma-3-12b",
            {"input": [{"role": "user", "content": [{"type": "input_text", "text": "Reply in JSON"}]}]},
            [{"role": "user", "content": "Reply in JSON"}],
            False,
            "/chat/completions",
            id="A5-responses-list-input",
        ),
    ],
)
def test_databricks_non_claude_single_text_block_reaches_upstream_as_string(
    gateway: Gateway,
    case: str,
    endpoint: str,
    provider_model: str,
    request_body: dict[str, JsonValue],
    expected_messages: list[JsonValue],
    responses_mode: bool,
    upstream_path: str,
) -> None:
    identity: Final = f"databricks-single-text-{case}-{uuid.uuid4().hex}"
    streaming: Final = request_body.get("stream") is True

    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["messages"] == expected_messages, body["messages"]
        return _shape_reply(identity, streaming)

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(
            model=f"databricks/{provider_model}",
            api_base=wire.url,
            api_key=_API_KEY,
            model_info={"mode": "responses"} if responses_mode else None,
        )
        headers: Final = {"anthropic-version": "2023-06-01"} if endpoint == "/v1/messages" else {}
        body: Final = {"model": model, **request_body}
        if streaming:
            with gateway.client.stream(
                "POST", endpoint, json=body, headers={"Authorization": f"Bearer {gateway.key}", **headers}
            ) as response:
                assert response.status_code == 200, response.read()
                lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
            assert lines[-1] == "data: [DONE]", lines
            assert any('"content":"ok"' in line for line in lines), lines
        else:
            response: Final = gateway.request("POST", endpoint, body, headers=headers)
            assert response.status_code == 200, response.text
            assert _JSON_OBJECT.validate_json(response.content)["id"]
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", upstream_path)]


@pytest.mark.parametrize(
    "case,provider_model,content,expected_content",
    [
        pytest.param(
            "cache-control",
            "databricks-gemma-3-12b",
            [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
            [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
            id="B1-cache-control-stays-list",
        ),
        pytest.param(
            "two-text-blocks",
            "databricks-gemma-3-12b",
            [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}],
            [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}],
            id="B2-two-text-blocks-stays-list",
        ),
        pytest.param(
            "text-and-image",
            "databricks-gemma-3-12b",
            [
                {"type": "text", "text": "Reply in JSON"},
                {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
            ],
            [
                {"type": "text", "text": "Reply in JSON"},
                {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
            ],
            id="B3-text-and-image-stays-list",
        ),
        pytest.param(
            "claude",
            "databricks-claude-opus-5-5",
            [{"type": "text", "text": "Reply in JSON"}],
            [{"type": "text", "text": "Reply in JSON"}],
            id="B4-claude-stays-list",
        ),
    ],
)
def test_databricks_list_content_stays_list(
    gateway: Gateway,
    case: str,
    provider_model: str,
    content: list[JsonValue],
    expected_content: list[JsonValue],
) -> None:
    identity: Final = f"databricks-list-content-{case}-{uuid.uuid4().hex}"
    expected_messages: Final = [{"role": "user", "content": expected_content}]

    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["messages"] == expected_messages, body["messages"]
        return _shape_reply(identity, False)

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{provider_model}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {"model": model, "messages": [{"role": "user", "content": content}]},
        )
        assert response.status_code == 200, response.text
        assert _JSON_OBJECT.validate_json(response.content)["id"]
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]


def test_databricks_json_schema_refs_are_preserved_in_chat_completions(gateway: Gateway) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/chat/completions"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["model"] == _JSON_SCHEMA_BACKEND
        assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
        assert body["stream"] is False
        assert body["response_format"] == _JSON_SCHEMA_RESPONSE_FORMAT
        return Reply(
            body=_chat_completion_response(
                _JSON_SCHEMA_BACKEND,
                {"role": "assistant", "content": _JSON_CONTENT},
                "stop",
            )
        )

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {
                "model": model,
                "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
                "response_format": _JSON_SCHEMA_RESPONSE_FORMAT,
            },
        )
        assert response.status_code == 200, response.text
        body: Final = _JSON_OBJECT.validate_json(response.content)
        assert_completion(body)
        assert body["choices"] == [
            {
                "index": 0,
                "message": {"role": "assistant", "content": _JSON_CONTENT},
                "finish_reason": "stop",
            }
        ]
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]


def test_databricks_streaming_json_schema_refs_are_preserved(gateway: Gateway) -> None:
    identity: Final = f"databricks-json-schema-stream-{uuid.uuid4().hex}"
    frames: Final = (
        _frame(
            identity,
            [{"index": 0, "delta": {"role": "assistant", "content": '{"p":'}, "finish_reason": None}],
            model=_JSON_SCHEMA_BACKEND,
        ),
        _frame(
            identity,
            [{"index": 0, "delta": {"content": '{"name":"Ada"}}'}, "finish_reason": None}],
            model=_JSON_SCHEMA_BACKEND,
        ),
        _frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}], model=_JSON_SCHEMA_BACKEND),
        b"data: [DONE]\n\n",
    )

    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/chat/completions"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["model"] == _JSON_SCHEMA_BACKEND
        assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
        assert body["stream"] is True
        assert body["response_format"] == _JSON_SCHEMA_RESPONSE_FORMAT
        return Reply(content_type="text/event-stream", chunks=frames)

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        with gateway.client.stream(
            "POST",
            "/v1/chat/completions",
            json={
                "model": model,
                "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
                "response_format": _JSON_SCHEMA_RESPONSE_FORMAT,
                "stream": True,
            },
            headers={"Authorization": f"Bearer {gateway.key}"},
        ) as response:
            assert response.status_code == 200, response.read()
            lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
        assert lines[-1] == "data: [DONE]", lines
        chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
        choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
        assert "".join(choice.delta.content or "" for choice in choices) == _JSON_CONTENT
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]


def test_databricks_claude_json_schema_refs_are_preserved_in_tool_parameters(gateway: Gateway) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/chat/completions"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["model"] == _CLAUDE_BACKEND
        assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
        assert "response_format" not in body
        assert body["tools"] == [
            {
                "type": "function",
                "function": {"name": "json_tool_call", "parameters": _PERSON_SCHEMA},
            }
        ]
        assert body["tool_choice"] == {"type": "function", "function": {"name": "json_tool_call"}}
        return Reply(body=_chat_completion_response(_CLAUDE_BACKEND, _JSON_TOOL_CALL_MESSAGE, "tool_calls"))

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {
                "model": model,
                "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
                "response_format": _JSON_SCHEMA_RESPONSE_FORMAT,
            },
        )
        assert response.status_code == 200, response.text
        body: Final = _JSON_OBJECT.validate_json(response.content)
        assert body["choices"] == [
            {
                "index": 0,
                "message": {"role": "assistant", "content": _JSON_CONTENT},
                "finish_reason": "stop",
            }
        ]
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]


def test_databricks_responses_json_schema_refs_are_preserved_in_chat_bridge(gateway: Gateway) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/chat/completions"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["model"] == _JSON_SCHEMA_BACKEND
        assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
        assert body["response_format"] == {
            "type": "json_schema",
            "json_schema": {"name": "P", "schema": _PERSON_SCHEMA, "strict": True},
        }
        return Reply(
            body=_chat_completion_response(
                _JSON_SCHEMA_BACKEND,
                {"role": "assistant", "content": _JSON_CONTENT},
                "stop",
            )
        )

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/responses",
            {
                "model": model,
                "input": _JSON_SCHEMA_PROMPT,
                "text": {
                    "format": {
                        "type": "json_schema",
                        "name": "P",
                        "strict": True,
                        "schema": _PERSON_SCHEMA,
                    }
                },
            },
        )
        assert response.status_code == 200, response.text
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]


def test_databricks_messages_json_schema_refs_are_preserved_for_qwen(gateway: Gateway) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/chat/completions"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["model"] == _JSON_SCHEMA_BACKEND
        assert body["response_format"] == _MESSAGES_RESPONSE_FORMAT
        return Reply(
            body=_chat_completion_response(
                _JSON_SCHEMA_BACKEND,
                {"role": "assistant", "content": _JSON_CONTENT},
                "stop",
            )
        )

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/messages",
            {
                "model": model,
                "max_tokens": 64,
                "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
                "output_format": {"type": "json_schema", "schema": _PERSON_SCHEMA},
            },
        )
        assert response.status_code == 200, response.text
        body: Final = _JSON_OBJECT.validate_json(response.content)
        assert body["content"] == [{"type": "text", "text": _JSON_CONTENT}]
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]


def test_databricks_messages_claude_json_schema_refs_are_preserved_in_tool_parameters(gateway: Gateway) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/chat/completions"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body["model"] == _CLAUDE_BACKEND
        assert "response_format" not in body
        assert body["tools"] == [
            {
                "type": "function",
                "function": {"name": "json_tool_call", "parameters": _MESSAGES_PERSON_SCHEMA},
            }
        ]
        assert body["tool_choice"] == {"type": "function", "function": {"name": "json_tool_call"}}
        return Reply(body=_chat_completion_response(_CLAUDE_BACKEND, _JSON_TOOL_CALL_MESSAGE, "tool_calls"))

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/messages",
            {
                "model": model,
                "max_tokens": 64,
                "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
                "output_format": {"type": "json_schema", "schema": _PERSON_SCHEMA},
            },
        )
        assert response.status_code == 200, response.text
        body: Final = _JSON_OBJECT.validate_json(response.content)
        assert body["content"] == [{"type": "text", "text": _JSON_CONTENT}]
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]


_Endpoint: TypeAlias = Literal["chat", "responses", "messages"]
_PATHS: Final[Mapping[_Endpoint, str]] = {
    "chat": "/v1/chat/completions",
    "responses": "/v1/responses",
    "messages": "/v1/messages",
}
_OUTAGE: Final = "upstream-outage"
_OUTAGE_MESSAGE: Final = "The service is temporarily unavailable"
_EXPECTED_JSON: Final[JsonValue] = _JSON_OBJECT.validate_json(_JSON_CONTENT)
_CONTENT_DELTAS: Final[tuple[Mapping[str, JsonValue], ...]] = (
    {"role": "assistant", "content": '{"p":'},
    {"content": '{"name":"Ada"}}'},
)
_TOOL_CALL_DELTAS: Final[tuple[Mapping[str, JsonValue], ...]] = (
    {
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {
                "index": 0,
                "id": "json-tool-call",
                "type": "function",
                "function": {"name": "json_tool_call", "arguments": ""},
            }
        ],
    },
    {"tool_calls": [{"index": 0, "function": {"arguments": _JSON_CONTENT}}]},
)


class _ChatMessage(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    content: str | None = None


class _ChatChoice(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    message: _ChatMessage


class _ChatCompletion(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    choices: tuple[_ChatChoice, ...]


class _OutputPart(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    type: str
    text: str | None = None


class _OutputItem(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    type: str
    content: tuple[_OutputPart, ...] = ()
    name: str | None = None
    arguments: str | None = None


class _ResponsesBody(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    output: tuple[_OutputItem, ...]


class _ResponsesEvent(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    type: str
    response: _ResponsesBody | None = None


class _MessagesBlock(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    type: str
    text: str | None = None


class _MessagesBody(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    content: tuple[_MessagesBlock, ...]


class _MessagesDelta(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    type: str | None = None
    text: str | None = None
    partial_json: str | None = None


class _MessagesEvent(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    type: str
    delta: _MessagesDelta | None = None


class _FunctionDelta(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    arguments: str | None = None


class _ToolCallDelta(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    function: _FunctionDelta | None = None


class _StreamDelta(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    content: str | None = None
    tool_calls: tuple[_ToolCallDelta, ...] | None = None


class _StreamChoice(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    delta: _StreamDelta


class _StreamChunk(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)
    choices: tuple[_StreamChoice, ...]


def _json_schema_format(schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
    return {"type": "json_schema", "json_schema": {"name": "P", "strict": True, "schema": dict(schema)}}


def _client_body(
    endpoint: _Endpoint,
    model: str,
    stream: bool,
    prompt: str = _JSON_SCHEMA_PROMPT,
) -> dict[str, JsonValue]:
    match endpoint:
        case "chat":
            return {
                "model": model,
                "messages": [{"role": "user", "content": prompt}],
                "response_format": _json_schema_format(_PERSON_SCHEMA),
                "stream": stream,
            }
        case "responses":
            return {
                "model": model,
                "input": prompt,
                "text": {
                    "format": {"type": "json_schema", "name": "P", "strict": True, "schema": dict(_PERSON_SCHEMA)}
                },
                "stream": stream,
            }
        case "messages":
            return {
                "model": model,
                "max_tokens": 64,
                "messages": [{"role": "user", "content": prompt}],
                "output_format": {"type": "json_schema", "schema": dict(_PERSON_SCHEMA)},
                "stream": stream,
            }


def _expected_outbound(endpoint: _Endpoint, backend: str) -> dict[str, JsonValue]:
    schema: Final = dict(_MESSAGES_PERSON_SCHEMA if endpoint == "messages" else _PERSON_SCHEMA)
    if "claude" in backend:
        return _JSON_OBJECT.validate_python(
            {
                "tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": schema}}],
                "tool_choice": {"type": "function", "function": {"name": "json_tool_call"}},
            }
        )
    name: Final = "structured_output" if endpoint == "messages" else "P"
    return _JSON_OBJECT.validate_python(
        {"response_format": {"type": "json_schema", "json_schema": {"name": name, "schema": schema, "strict": True}}}
    )


def _assert_outbound(request: Request, endpoint: _Endpoint, backend: str, stream: bool) -> None:
    assert (request.method, request.target) == ("POST", "/chat/completions")
    body: Final = _JSON_OBJECT.validate_json(request.body)
    expected: Final = _expected_outbound(endpoint, backend)
    assert {key: body.get(key) for key in expected} == expected
    assert (body["model"], body.get("stream") is True) == (backend, stream)
    assert "claude" not in backend or "response_format" not in body


def _upstream_reply(backend: str, stream: bool) -> Reply:
    claude: Final = "claude" in backend
    finish_reason: Final = "tool_calls" if claude else "stop"
    if not stream:
        message: Final[dict[str, JsonValue]] = (
            _JSON_TOOL_CALL_MESSAGE if claude else {"role": "assistant", "content": _JSON_CONTENT}
        )
        return Reply(body=_chat_completion_response(backend, message, finish_reason))
    identity: Final = f"databricks-json-schema-{uuid.uuid4().hex}"
    deltas: Final = _TOOL_CALL_DELTAS if claude else _CONTENT_DELTAS
    frames: Final = (
        *(_frame(identity, [{"index": 0, "delta": delta, "finish_reason": None}], model=backend) for delta in deltas),
        _frame(identity, [{"index": 0, "delta": {}, "finish_reason": finish_reason}], model=backend),
        b"data: [DONE]\n\n",
    )
    return Reply(content_type="text/event-stream", chunks=frames)


def _outage_or_reply(request: Request) -> Reply:
    if _OUTAGE.encode() in request.body:
        error: Final = {"error_code": "TEMPORARILY_UNAVAILABLE", "message": _OUTAGE_MESSAGE}
        return Reply(status=503, body=json.dumps(error).encode())
    body: Final = _JSON_OBJECT.validate_json(request.body)
    return _upstream_reply(str(body["model"]), body.get("stream") is True)


def _item_text(item: _OutputItem) -> str:
    match item.type:
        case "message":
            return "".join(part.text or "" for part in item.content)
        case "function_call" if item.name == "json_tool_call":
            return item.arguments or ""
        case _:
            return ""


def _responses_text(body: _ResponsesBody) -> str:
    return "".join(_item_text(item) for item in body.output)


def _chat_delta_text(delta: _StreamDelta) -> str:
    calls: Final = delta.tool_calls or ()
    return (delta.content or "") + "".join(call.function.arguments or "" for call in calls if call.function is not None)


def _messages_delta_text(delta: _MessagesDelta) -> str:
    match delta.type:
        case "text_delta":
            return delta.text or ""
        case "input_json_delta":
            return delta.partial_json or ""
        case _:
            return ""


def _final_text(endpoint: _Endpoint, content: bytes) -> str:
    match endpoint:
        case "chat":
            return _ChatCompletion.model_validate_json(content).choices[0].message.content or ""
        case "responses":
            return _responses_text(_ResponsesBody.model_validate_json(content))
        case "messages":
            return "".join(block.text or "" for block in _MessagesBody.model_validate_json(content).content)


def _streamed_text(endpoint: _Endpoint, data: tuple[str, ...]) -> str:
    match endpoint:
        case "chat":
            assert data[-1] == "[DONE]", data
            chunks: Final = tuple(_StreamChunk.model_validate_json(line) for line in data[:-1])
            choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
            return "".join(_chat_delta_text(choice.delta) for choice in choices)
        case "responses":
            events: Final = tuple(_ResponsesEvent.model_validate_json(line) for line in data if line != "[DONE]")
            completed: Final = tuple(event.response for event in events if event.type == "response.completed")
            assert len(completed) == 1 and completed[0] is not None, data
            return _responses_text(completed[0])
        case "messages":
            message_events: Final = tuple(_MessagesEvent.model_validate_json(line) for line in data)
            assert message_events[-1].type == "message_stop", data
            deltas: Final = tuple(event.delta for event in message_events if event.delta is not None)
            return "".join(_messages_delta_text(delta) for delta in deltas)


def _outcome(status: int, text: str) -> tuple[int, JsonValue]:
    if status == 200:
        return status, _JSON_OBJECT.validate_json(text)
    return status, _OUTAGE_MESSAGE if _OUTAGE_MESSAGE in text else text


def _call(gateway: Gateway, endpoint: _Endpoint, body: Mapping[str, JsonValue]) -> tuple[int, str]:
    if body["stream"] is not True:
        response: Final = gateway.request("POST", _PATHS[endpoint], body)
        return response.status_code, _final_text(endpoint, response.content) if response.is_success else response.text
    with gateway.client.stream(
        "POST", _PATHS[endpoint], json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"}
    ) as streamed:
        if not streamed.is_success:
            return streamed.status_code, streamed.read().decode()
        data: Final = tuple(line.removeprefix("data: ") for line in streamed.iter_lines() if line.startswith("data: "))
    return 200, _streamed_text(endpoint, data)


@pytest.mark.parametrize(
    ("endpoint", "backend", "stream"),
    [
        pytest.param("chat", _CLAUDE_BACKEND, True, id="chat-stream-claude"),
        pytest.param("responses", _CLAUDE_BACKEND, False, id="responses-claude"),
        pytest.param("responses", _JSON_SCHEMA_BACKEND, True, id="responses-stream-qwen"),
        pytest.param("responses", _CLAUDE_BACKEND, True, id="responses-stream-claude"),
        pytest.param("messages", _JSON_SCHEMA_BACKEND, True, id="messages-stream-qwen"),
        pytest.param("messages", _CLAUDE_BACKEND, True, id="messages-stream-claude"),
    ],
)
def test_databricks_json_schema_refs_are_preserved_on_every_endpoint(
    gateway: Gateway, endpoint: _Endpoint, backend: str, stream: bool
) -> None:
    with wire_server(lambda _request: _upstream_reply(backend, stream)) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{backend}", api_base=wire.url, api_key=_API_KEY)
        status, text = _call(gateway, endpoint, _client_body(endpoint, model, stream))
        received: Final = wire.drain()
    assert len(received) == 1, received
    _assert_outbound(received[0], endpoint, backend, stream)
    assert _outcome(status, text) == (200, _EXPECTED_JSON), text


@pytest.mark.parametrize(
    "response_format",
    [
        pytest.param(
            _json_schema_format(
                {
                    "definitions": {
                        "Person": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}
                    },
                    "type": "object",
                    "properties": {"p": {"$ref": "#/definitions/Person"}},
                    "required": ["p"],
                }
            ),
            id="definitions-ref",
        ),
        pytest.param(
            _json_schema_format(
                {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string"},
                        "children": {"type": "array", "items": {"$ref": "#"}},
                    },
                    "required": ["name", "children"],
                }
            ),
            id="recursive-root-ref",
        ),
        pytest.param(
            _json_schema_format({"type": "object", "properties": {"$ref": {"type": "string"}}, "required": ["$ref"]}),
            id="property-named-ref",
        ),
        pytest.param(
            _json_schema_format({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}),
            id="flat-schema",
        ),
        pytest.param({"type": "json_object"}, id="json-object"),
    ],
)
def test_databricks_response_format_reaches_upstream_as_sent(
    gateway: Gateway, response_format: dict[str, JsonValue]
) -> None:
    with (
        wire_server(lambda _request: _upstream_reply(_JSON_SCHEMA_BACKEND, False)) as wire,
        gateway.scenario() as scenario,
    ):
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {
                "model": model,
                "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
                "response_format": response_format,
            },
        )
        received: Final = wire.drain()
    assert response.status_code == 200, response.text
    assert _final_text("chat", response.content) == _JSON_CONTENT
    assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
    assert _JSON_OBJECT.validate_json(received[0].body)["response_format"] == response_format


def test_databricks_json_schema_without_schema_returns_the_upstream_error(gateway: Gateway) -> None:
    response_format: Final[dict[str, JsonValue]] = {"type": "json_schema", "json_schema": {"name": "P", "strict": True}}
    upstream_error: Final = {"error_code": "INVALID_PARAMETER_VALUE", "message": "json_schema.schema is required"}

    def respond(_request: Request) -> Reply:
        return Reply(status=400, body=json.dumps(upstream_error).encode())

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {
                "model": model,
                "messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
                "response_format": response_format,
            },
        )
        received: Final = wire.drain()
    assert (response.status_code, "json_schema.schema is required" in response.text) == (400, True), response.text
    assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
    assert _JSON_OBJECT.validate_json(received[0].body)["response_format"] == response_format


def _assert_burst_outbound(request: Request) -> None:
    body: Final = _JSON_OBJECT.validate_json(request.body)
    endpoint: Final[_Endpoint] = "messages" if body["response_format"] == _MESSAGES_RESPONSE_FORMAT else "chat"
    _assert_outbound(request, endpoint, _JSON_SCHEMA_BACKEND, body.get("stream") is True)


def _burst_request(index: int) -> tuple[_Endpoint, bool, bool]:
    endpoints: Final[tuple[_Endpoint, ...]] = ("chat", "responses", "messages")
    return endpoints[index % 3], index % 2 == 1, index % 5 == 2


def test_databricks_json_schema_burst_through_an_upstream_outage_keeps_every_ref(gateway: Gateway) -> None:
    plan: Final = tuple(_burst_request(index) for index in range(24))
    recovery: Final[tuple[_Endpoint, ...]] = ("chat", "responses", "messages")
    with wire_server(_outage_or_reply) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        endpoints: Final = tuple(endpoint for endpoint, _stream, _down in plan)
        bodies: Final = tuple(
            _client_body(endpoint, model, stream, f"{_JSON_SCHEMA_PROMPT} {index} {_OUTAGE if down else ''}")
            for index, (endpoint, stream, down) in enumerate(plan)
        )
        with ThreadPoolExecutor(max_workers=12) as pool:
            results: Final = tuple(pool.map(_call, itertools.repeat(gateway), endpoints, bodies))
        recovered: Final = tuple(
            _call(gateway, endpoint, _client_body(endpoint, model, False)) for endpoint in recovery
        )
        received: Final = wire.drain()
    assert len(received) == len(plan) + len(recovery), received
    for request in received:
        _assert_burst_outbound(request)
    expected: Final = tuple(
        (503, _OUTAGE_MESSAGE) if down else (200, _EXPECTED_JSON) for _endpoint, _stream, down in plan
    )
    assert tuple(_outcome(*result) for result in results) == expected
    assert tuple(_outcome(*result) for result in recovered) == ((200, _EXPECTED_JSON),) * len(recovery)


_GEMMA_BACKEND: Final = "databricks-gemma-3-12b"
_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True}


def _one_block(text: str) -> list[JsonValue]:
    return [{"type": "text", "text": text}]


def _plain_reply(_request: Request) -> Reply:
    return _shape_reply(f"databricks-plain-{uuid.uuid4().hex}", False)


def _sent_messages(request: Request) -> JsonValue:
    return _JSON_OBJECT.validate_json(request.body)["messages"]


def _tool_call(identity: str, city: str) -> dict[str, JsonValue]:
    return {
        "id": identity,
        "type": "function",
        "function": {"name": "get_weather", "arguments": json.dumps({"city": city})},
    }


@pytest.mark.parametrize(
    "messages,expected_messages",
    [
        pytest.param(
            [
                {"role": "system", "content": _one_block("Alpha rule.")},
                {"role": "system", "content": _one_block("Bravo rule.")},
                {"role": "user", "content": _one_block("Reply in JSON")},
            ],
            [
                {"role": "system", "content": "Alpha rule.\n\nBravo rule."},
                {"role": "user", "content": "Reply in JSON"},
            ],
            id="C1-two-one-block-system-messages",
        ),
        pytest.param(
            [
                {"role": "system", "content": _one_block("Alpha rule.")},
                {"role": "developer", "content": _one_block("Bravo rule.")},
                {"role": "user", "content": "Reply in JSON"},
            ],
            [
                {"role": "system", "content": "Alpha rule.\n\nBravo rule."},
                {"role": "user", "content": "Reply in JSON"},
            ],
            id="C2-system-then-developer-one-block",
        ),
        pytest.param(
            [
                {"role": "system", "content": _one_block("Alpha rule.")},
                {"role": "system", "content": "Bravo rule."},
                {"role": "user", "content": "Reply in JSON"},
            ],
            [
                {"role": "system", "content": "Alpha rule.\n\nBravo rule."},
                {"role": "user", "content": "Reply in JSON"},
            ],
            id="C3-one-block-then-string-system",
        ),
        pytest.param(
            [
                {"role": "system", "content": _one_block("Alpha rule.")},
                {
                    "role": "system",
                    "content": [{"type": "text", "text": "Bravo", "cache_control": {"type": "ephemeral"}}],
                },
                {"role": "user", "content": "Reply in JSON"},
            ],
            [
                {
                    "role": "system",
                    "content": [
                        {"type": "text", "text": "Alpha rule."},
                        {"type": "text", "text": "Bravo", "cache_control": {"type": "ephemeral"}},
                    ],
                },
                {"role": "user", "content": "Reply in JSON"},
            ],
            id="C4-one-block-then-cache-marked-system-stays-parts",
        ),
        pytest.param(
            [
                {"role": "user", "content": _one_block("Weather in Paris and Rome?")},
                {
                    "role": "assistant",
                    "content": _one_block("Checking both."),
                    "tool_calls": [_tool_call("call-paris", "Paris"), _tool_call("call-rome", "Rome")],
                },
                {"role": "tool", "tool_call_id": "call-paris", "content": _one_block("18C")},
                {"role": "tool", "tool_call_id": "call-rome", "content": _one_block("22C")},
            ],
            [
                {"role": "user", "content": "Weather in Paris and Rome?"},
                {"role": "assistant", "content": "Checking both.", "tool_calls": [_tool_call("call-paris", "Paris")]},
                {"role": "tool", "tool_call_id": "call-paris", "content": "18C"},
                {"role": "assistant", "tool_calls": [_tool_call("call-rome", "Rome")]},
                {"role": "tool", "tool_call_id": "call-rome", "content": "22C"},
            ],
            id="C5-parallel-tool-results-split-and-collapsed",
        ),
        pytest.param(
            [
                {"role": "user", "content": [{"type": "text", "text": "Reply in JSON", "cache_control": None}]},
                {"role": "assistant", "content": [{"type": "text", "text": "ok", "annotations": []}]},
                {"role": "user", "content": _one_block("   ")},
                {"role": "user", "content": _one_block("x" * 5120)},
            ],
            [
                {"role": "user", "content": [{"type": "text", "text": "Reply in JSON", "cache_control": None}]},
                {"role": "assistant", "content": [{"type": "text", "text": "ok", "annotations": []}]},
                {"role": "user"},
                {"role": "user", "content": "x" * 5120},
            ],
            id="F1-extra-keys-stay-list-blank-dropped-5kb-collapsed",
        ),
    ],
)
def test_databricks_non_claude_message_shapes_reach_upstream(
    gateway: Gateway, messages: list[JsonValue], expected_messages: list[JsonValue]
) -> None:
    with wire_server(_plain_reply) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {"model": model, "messages": messages, "cache": _NO_CACHE},
        )
        received: Final = wire.drain()
    assert response.status_code == 200, response.text
    assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
    assert _sent_messages(received[0]) == expected_messages


def test_databricks_responses_instructions_and_one_block_developer_reach_upstream_joined(gateway: Gateway) -> None:
    with wire_server(_plain_reply) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/responses",
            {
                "model": model,
                "instructions": "Alpha rule.",
                "input": [
                    {"role": "developer", "content": [{"type": "input_text", "text": "Bravo rule."}]},
                    {"role": "user", "content": [{"type": "input_text", "text": "Reply in JSON"}]},
                ],
            },
        )
        received: Final = wire.drain()
    assert response.status_code == 200, response.text
    assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
    assert _sent_messages(received[0]) == [
        {"role": "system", "content": "Alpha rule.\n\nBravo rule."},
        {"role": "user", "content": "Reply in JSON"},
    ]


def test_databricks_fallback_sends_the_collapsed_string_to_both_deployments(gateway: Gateway) -> None:
    marker: Final = uuid.uuid4().hex
    fallback_backend: Final = "databricks-qwen3-next-80b-a3b-instruct"

    def respond(request: Request) -> Reply:
        if _JSON_OBJECT.validate_json(request.body)["model"] == _GEMMA_BACKEND:
            return Reply(status=500, body=b'{"error_code": "INTERNAL_ERROR", "message": "primary deployment is down"}')
        return _shape_reply(f"databricks-fallback-{marker}", False)

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        primary: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        fallback: Final = scenario.model(model=f"databricks/{fallback_backend}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {
                "model": primary,
                "messages": [{"role": "user", "content": _one_block(f"Reply in JSON {marker}")}],
                "response_format": {"type": "json_object"},
                "fallbacks": [fallback],
                "num_retries": 0,
                "cache": _NO_CACHE,
            },
        )
        received: Final = wire.drain()
    assert response.status_code == 200, response.text
    assert _JSON_OBJECT.validate_json(response.content)["id"] == f"databricks-fallback-{marker}"
    assert [_JSON_OBJECT.validate_json(request.body)["model"] for request in received] == [
        _GEMMA_BACKEND,
        fallback_backend,
    ]
    assert [_sent_messages(request) for request in received] == [
        [{"role": "user", "content": f"Reply in JSON {marker}"}],
    ] * 2


def _proxy_root(gateway: Gateway) -> str:
    return str(gateway.client.base_url).rstrip("/")


def test_databricks_openai_sdk_sync_one_block_reaches_upstream_as_string_and_logs_spend(gateway: Gateway) -> None:
    identity: Final = f"databricks-openai-sdk-{uuid.uuid4().hex}"
    with (
        wire_server(lambda _request: _shape_reply(identity, False)) as wire,
        gateway.scenario() as scenario,
        openai.OpenAI(base_url=f"{_proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0) as client,
    ):
        model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        completion: Final = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
            response_format={"type": "json_object"},
            extra_body={"cache": _NO_CACHE},
        )
        received: Final = wire.drain()
    assert (completion.id, completion.choices[0].message.content) == (identity, "ok")
    assert [_sent_messages(request) for request in received] == [[{"role": "user", "content": "Reply in JSON"}]]
    rows: Final = eventually(
        lambda: read_rows('SELECT status, call_type FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)),
        lambda values: len(values) == 1,
        seconds=70,
    )
    assert (rows[0]["status"], rows[0]["call_type"]) == ("success", "acompletion")


async def test_databricks_openai_sdk_async_stream_one_block_reaches_upstream_as_string(gateway: Gateway) -> None:
    identity: Final = f"databricks-openai-async-{uuid.uuid4().hex}"
    with wire_server(lambda _request: _shape_reply(identity, True)) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        async with openai.AsyncOpenAI(
            base_url=f"{_proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0
        ) as client:
            stream: Final = await client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
                response_format={"type": "json_object"},
                stream=True,
                extra_body={"cache": _NO_CACHE},
            )
            chunks: Final = [chunk async for chunk in stream]
        received: Final = wire.drain()
    assert {chunk.id for chunk in chunks} == {identity}
    assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "ok"
    assert [_sent_messages(request) for request in received] == [[{"role": "user", "content": "Reply in JSON"}]]


async def test_databricks_anthropic_sdk_async_stream_one_block_system_and_user_reach_upstream_as_strings(
    gateway: Gateway,
) -> None:
    identity: Final = f"databricks-anthropic-async-{uuid.uuid4().hex}"
    with wire_server(lambda _request: _shape_reply(identity, True)) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        async with anthropic.AsyncAnthropic(
            base_url=_proxy_root(gateway), api_key=gateway.key, max_retries=0
        ) as client:
            async with client.messages.stream(
                model=model,
                max_tokens=32,
                system=[{"type": "text", "text": "Reply in JSON."}],
                messages=[{"role": "user", "content": [{"type": "text", "text": "What is 2+2?"}]}],
            ) as stream:
                final: Final = await stream.get_final_message()
        received: Final = wire.drain()
    assert [block.text for block in final.content if block.type == "text"] == ["ok"]
    assert [_sent_messages(request) for request in received] == [
        [{"role": "system", "content": "Reply in JSON."}, {"role": "user", "content": "What is 2+2?"}]
    ]


def test_databricks_one_block_burst_through_an_upstream_outage_reaches_upstream_as_strings(gateway: Gateway) -> None:
    plan: Final = tuple(_burst_request(index) for index in range(24))

    def respond(request: Request) -> Reply:
        if _OUTAGE.encode() in request.body:
            error: Final = {"error_code": "TEMPORARILY_UNAVAILABLE", "message": _OUTAGE_MESSAGE}
            return Reply(status=503, body=json.dumps(error).encode())
        body: Final = _JSON_OBJECT.validate_json(request.body)
        return _upstream_reply(_JSON_SCHEMA_BACKEND, body.get("stream") is True)

    def one_block_body(endpoint: _Endpoint, model: str, stream: bool, prompt: str) -> dict[str, JsonValue]:
        body: Final = _client_body(endpoint, model, stream, prompt)
        match endpoint:
            case "responses":
                return {**body, "input": [{"role": "user", "content": [{"type": "input_text", "text": prompt}]}]}
            case "chat" | "messages":
                return {**body, "messages": [{"role": "user", "content": _one_block(prompt)}]}

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        prompts: Final = tuple(
            f"{_JSON_SCHEMA_PROMPT} {index} {_OUTAGE if down else ''}".strip()
            for index, (_e, _s, down) in enumerate(plan)
        )
        bodies: Final = tuple(
            one_block_body(endpoint, model, stream, prompt) for (endpoint, stream, _down), prompt in zip(plan, prompts)
        )
        with ThreadPoolExecutor(max_workers=12) as pool:
            results: Final = tuple(pool.map(_call, itertools.repeat(gateway), (e for e, _s, _d in plan), bodies))
        received: Final = wire.drain()
    assert sorted(json.dumps(_sent_messages(request)) for request in received) == sorted(
        json.dumps([{"role": "user", "content": prompt}]) for prompt in prompts
    )
    expected: Final = tuple(
        (503, _OUTAGE_MESSAGE) if down else (200, _EXPECTED_JSON) for _endpoint, _stream, down in plan
    )
    assert tuple(_outcome(*result) for result in results) == expected


def test_databricks_transform_request_shows_the_collapsed_string(gateway: Gateway, tmp_path: Path) -> None:
    environment: Final = {
        "DATABRICKS_API_BASE": "https://databricks.invalid/serving-endpoints",
        "DATABRICKS_API_KEY": _API_KEY,
    }
    with owned_proxy_process(gateway, tmp_path, environment) as owned:
        response: Final = owned.gateway.request(
            "POST",
            "/utils/transform_request",
            {
                "call_type": "completion",
                "request_body": {
                    "model": f"databricks/{_GEMMA_BACKEND}",
                    "messages": [{"role": "user", "content": _one_block("Reply in JSON")}],
                },
            },
        )
    assert response.status_code == 200, response.text
    body: Final = _JSON_OBJECT.validate_json(response.content)
    assert body["error"] is None, body
    assert body["raw_request_api_base"] == "https://databricks.invalid/serving-endpoints/chat/completions"
    sent: Final = body["raw_request_body"]
    assert isinstance(sent, dict), body
    assert sent["messages"] == [{"role": "user", "content": "Reply in JSON"}]


def test_databricks_non_string_text_block_is_rejected_before_reaching_upstream(gateway: Gateway) -> None:
    with wire_server(_plain_reply) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
        response: Final = gateway.request(
            "POST",
            "/v1/chat/completions",
            {
                "model": model,
                "messages": [{"role": "user", "content": [{"type": "text", "text": 42}]}],
                "cache": _NO_CACHE,
            },
        )
        received: Final = wire.drain()
    assert response.status_code == 400, response.text
    assert "'int' object has no attribute 'strip'" in response.text
    assert received == ()