From 09045c5162ac3db46ac05f1cb5543c1d2a217989 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 11 Sep 2026 09:05:02 -0700 Subject: [PATCH] refactor(ocr): resolve logging executor explicitly --- litellm/litellm_core_utils/call_completion.py | 1 + .../test_call_completion.py | 33 ++++++++++++++++++- tests/test_litellm_rust/conftest.py | 3 +- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/call_completion.py b/litellm/litellm_core_utils/call_completion.py index 7d65625a7a1..b4cb6e8b276 100644 --- a/litellm/litellm_core_utils/call_completion.py +++ b/litellm/litellm_core_utils/call_completion.py @@ -51,6 +51,7 @@ class CompletionExecutor(Protocol): def submit( self, function: Callable[..., object], + /, *args: object, ) -> Future[object]: ... diff --git a/tests/test_litellm/litellm_core_utils/test_call_completion.py b/tests/test_litellm/litellm_core_utils/test_call_completion.py index 41d7a1d07f0..78006070c3f 100644 --- a/tests/test_litellm/litellm_core_utils/test_call_completion.py +++ b/tests/test_litellm/litellm_core_utils/test_call_completion.py @@ -3,14 +3,16 @@ import contextvars import datetime import weakref from collections.abc import Callable, Coroutine -from concurrent.futures import Future +from concurrent.futures import Future, ThreadPoolExecutor from importlib import import_module +from threading import get_ident from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest import litellm +from litellm.litellm_core_utils import thread_pool_executor from litellm.litellm_core_utils.call_completion import CallCompletion, PythonCompletion from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import ocr as rust_ocr_bridge @@ -128,6 +130,35 @@ def test_python_completion_preserves_sync_context_and_response_identity() -> Non assert len(executor.submissions) == 1 +def test_sync_wrapper_dispatches_with_logging_executor_and_caller_context(monkeypatch: pytest.MonkeyPatch) -> None: + marker: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("wrapper-context", default="missing") + marker.set("request-context") + caller_thread: Final = get_ident() + response: Final = object() + observed: Final[list[tuple[object, str, int]]] = [] + logging_obj: Final = MagicMock() + + def record_success(result: object, start_time: datetime.datetime, end_time: datetime.datetime) -> None: + observed.append((result, marker.get(), get_ident())) + + def ocr(**kwargs: object) -> object: + return response + + logging_obj.success_handler.side_effect = record_success + monkeypatch.setattr("litellm.utils.function_setup", MagicMock(return_value=(logging_obj, {}))) + monkeypatch.setattr("litellm.utils.load_credentials_from_list", MagicMock()) + wrapped: Final = client(ocr) + with ThreadPoolExecutor(max_workers=1) as executor: + monkeypatch.setattr(thread_pool_executor, "executor", executor) + result: Final = wrapped() + + assert result is response + assert len(observed) == 1 + assert observed[0][0] is response + assert observed[0][1] == "request-context" + assert observed[0][2] != caller_thread + + @pytest.mark.asyncio async def test_call_completion_attaches_once_and_forwards_final_objects() -> None: python_completion: Final = RecordingCompletion() diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index b0c75d9d2f5..17312e4d9d5 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -11,7 +11,7 @@ import pytest_asyncio import litellm from litellm import utils -from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils import litellm_logging, thread_pool_executor from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivateUsage] # preserve raw configuration state in test isolation _CONFIGURATION, @@ -77,6 +77,7 @@ async def isolate_ocr_test_state() -> AsyncIterator[None]: stack.enter_context(_rebound(_CONFIGURATION, "override", None)) executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") stack.enter_context(_rebound(utils, "executor", executor)) + stack.enter_context(_rebound(thread_pool_executor, "executor", executor)) try: yield finally: