From 54f26f09d1c4e6ce28592ee0c86fb608957a0449 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 03:30:00 -0700
Subject: [PATCH] fix(image variations): give each SDK client the session that
 matches it

The async image variations path builds its client params and hands them
to `AsyncOpenAI`, and the sync path hands its own to `OpenAI`, but the two
param blocks sit in the opposite order from the calls that consume them.
The previous commit read them in source order and put `aclient_session`
on the sync one, so a deployment that sets either session got a
`TypeError: Invalid http_client argument` on every image variation.

Each block now carries the session its SDK class accepts. The async path
also stops getting the sync session, which it had since before this
branch.

The new test drives both entry points against a loopback upstream with
both sessions set, and fails on either wiring being crossed.
---
 .../llms/openai/image_variations/handler.py   |  4 +-
 .../openai/image_variations/test_handler.py   | 83 +++++++++++++++++++
 2 files changed, 85 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_litellm/llms/openai/image_variations/test_handler.py

diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py
index fc9e644c5c0..ee33a5c6a14 100644
--- a/litellm/llms/openai/image_variations/handler.py
+++ b/litellm/llms/openai/image_variations/handler.py
@@ -63,7 +63,7 @@ class OpenAIImageVariationsHandler:
             init_client_params: Final = {
                 "api_key": api_key,
                 "base_url": api_base,
-                "http_client": litellm.client_session,
+                "http_client": litellm.aclient_session,
                 "default_headers": accept_encoding_header(),
                 "timeout": timeout,
                 "max_retries": max_retries,
@@ -181,7 +181,7 @@ class OpenAIImageVariationsHandler:
             init_client_params: Final = {
                 "api_key": api_key,
                 "base_url": api_base,
-                "http_client": litellm.aclient_session,
+                "http_client": litellm.client_session,
                 "default_headers": accept_encoding_header(),
                 "timeout": timeout,
                 "max_retries": max_retries,
diff --git a/tests/test_litellm/llms/openai/image_variations/test_handler.py b/tests/test_litellm/llms/openai/image_variations/test_handler.py
new file mode 100644
index 00000000000..501571f0962
--- /dev/null
+++ b/tests/test_litellm/llms/openai/image_variations/test_handler.py
@@ -0,0 +1,83 @@
+import io
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import httpx
+import pytest
+
+import litellm
+
+
+class _ImageVariationsUpstream(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def do_POST(self) -> None:
+        self.rfile.read(int(self.headers.get("Content-Length") or 0))
+        body = json.dumps({"created": 1, "data": [{"url": "https://example.invalid/variation.png"}]}).encode()
+        self.send_response(200)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def log_message(self, *args: object) -> None:
+        return
+
+
+@pytest.fixture
+def image_variations_upstream():
+    server = ThreadingHTTPServer(("127.0.0.1", 0), _ImageVariationsUpstream)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        yield f"http://127.0.0.1:{server.server_address[1]}"
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join(timeout=5)
+
+
+@pytest.fixture
+def shared_sessions(monkeypatch: pytest.MonkeyPatch):
+    sync_session = httpx.Client()
+    monkeypatch.setattr(litellm, "client_session", sync_session)
+    monkeypatch.setattr(litellm, "aclient_session", httpx.AsyncClient())
+    try:
+        yield
+    finally:
+        sync_session.close()
+
+
+def _square_png() -> io.BytesIO:
+    image = io.BytesIO(
+        b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00"
+        b"\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82"
+    )
+    image.name = "square.png"
+    return image
+
+
+def test_sync_image_variation_uses_the_sync_shared_session(image_variations_upstream, shared_sessions):
+    response = litellm.image_variation(
+        image=_square_png(),
+        model="dall-e-2",
+        api_key="sk-test",
+        api_base=image_variations_upstream,
+    )
+
+    assert response.data is not None
+    assert response.data[0].url == "https://example.invalid/variation.png"
+
+
+@pytest.mark.asyncio
+async def test_async_image_variation_uses_the_async_shared_session(image_variations_upstream, shared_sessions):
+    response = await litellm.aimage_variation(
+        image=_square_png(),
+        model="dall-e-2",
+        api_key="sk-test",
+        api_base=image_variations_upstream,
+    )
+
+    assert response.data is not None
+    assert response.data[0].url == "https://example.invalid/variation.png"