diff --git a/litellm/litellm_core_utils/openai_tool_name_mapping.py b/litellm/litellm_core_utils/openai_tool_name_mapping.py index dfce305ce87..0c55967ede5 100644 --- a/litellm/litellm_core_utils/openai_tool_name_mapping.py +++ b/litellm/litellm_core_utils/openai_tool_name_mapping.py @@ -50,10 +50,23 @@ def should_sanitize_openai_tool_names(litellm_provider: str) -> bool: return litellm_provider in _OPENAI_TOOL_NAME_VALIDATION_PROVIDERS -def begin_openai_tool_name_mapping_scope() -> None: - """Reset both mappings for this completion (call once at the start of completion()).""" - _CTX.set({}) - _CTX_REV.set({}) +def begin_openai_tool_name_mapping_scope(*, force_reset: bool = False) -> None: + """Initialise (or reset) both mapping dicts for one completion request. + + Idempotent by default: if the ContextVar already holds a dict this is a + no-op. This matters for async streaming where ``acompletion`` calls this + function in the *outer* async context before ``contextvars.copy_context()`` + so that both the executor and the stream consumer share the same dict + objects by reference. The inner ``completion()`` call then sees a non-None + dict and skips the reset, preserving the shared reference. + + Pass ``force_reset=True`` only when you explicitly want a clean slate + (e.g. in tests). + """ + if force_reset or _CTX.get() is None: + _CTX.set({}) + if force_reset or _CTX_REV.get() is None: + _CTX_REV.set({}) def _store(sanitized: str, original: str) -> None: diff --git a/litellm/main.py b/litellm/main.py index 3e471ac77c7..1cae4e4b3ff 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -611,6 +611,9 @@ async def acompletion( # noqa: PLR0915 # Use a partial function to pass your keyword arguments func = partial(completion, **completion_kwargs, **kwargs) + # Initialise the OpenAI tool-name mapping scope in the *outer* async + begin_openai_tool_name_mapping_scope() + # Add the context to the function ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 601650c3836..692f15d9df9 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4032,7 +4032,7 @@ def test_openai_tool_name_mapping_per_request_context(): ) from litellm.utils import validate_and_fix_openai_tools - begin_openai_tool_name_mapping_scope() + begin_openai_tool_name_mapping_scope(force_reset=True) validate_and_fix_openai_tools( tools=[ { @@ -4049,7 +4049,8 @@ def test_openai_tool_name_mapping_per_request_context(): restore_openai_tool_name_for_user("invalid_name_at_index_71") == "invalid.name.at.index.71" ) - begin_openai_tool_name_mapping_scope() + # Simulate a second request: force_reset=True clears the mapping. + begin_openai_tool_name_mapping_scope(force_reset=True) assert ( restore_openai_tool_name_for_user("invalid_name_at_index_71") == "invalid_name_at_index_71" @@ -4066,7 +4067,7 @@ def test_convert_to_model_response_restores_openai_tool_call_names_when_mapped() from litellm.types.utils import ModelResponse from litellm.utils import validate_and_fix_openai_tools - begin_openai_tool_name_mapping_scope() + begin_openai_tool_name_mapping_scope(force_reset=True) validate_and_fix_openai_tools( tools=[ { @@ -4121,7 +4122,7 @@ def test_tool_choice_function_name_sanitized_with_tools(): ) from litellm.utils import validate_and_fix_openai_tools - begin_openai_tool_name_mapping_scope() + begin_openai_tool_name_mapping_scope(force_reset=True) validate_and_fix_openai_tools( tools=[ { @@ -4152,7 +4153,7 @@ def test_tool_choice_function_name_uses_collision_suffix(): ) from litellm.utils import validate_and_fix_openai_tools - begin_openai_tool_name_mapping_scope() + begin_openai_tool_name_mapping_scope(force_reset=True) validate_and_fix_openai_tools( tools=[ { @@ -4175,3 +4176,41 @@ def test_tool_choice_function_name_uses_collision_suffix(): # "a.b" → "a_b"; "a/b" also sanitizes to "a_b" but gets suffix → "a_b_1" assert get_sanitized_tool_name("a.b") == "a_b" assert get_sanitized_tool_name("a/b") == "a_b_1" + + +def test_begin_scope_is_idempotent_shared_reference(): + """begin_openai_tool_name_mapping_scope must not replace the dict when already set. + + This simulates the async-streaming path: the outer acompletion() call + initialises the scope (sets _CTX = {}), copy_context() copies the + reference, and then the inner completion() call must not clobber it with a + new empty dict — otherwise mutations from the executor are invisible to the + stream consumer in the outer context. + """ + import contextvars + + from litellm.litellm_core_utils.openai_tool_name_mapping import ( + _CTX, + begin_openai_tool_name_mapping_scope, + restore_openai_tool_name_for_user, + store_openai_tool_name_mapping, + ) + + # Outer context initialises the scope (simulates acompletion pre-copy_context). + begin_openai_tool_name_mapping_scope(force_reset=True) + outer_dict = _CTX.get() + assert outer_dict is not None + + # Inner context (executor) calls begin_openai_tool_name_mapping_scope again + # and then writes a mapping — simulates what completion() does. + def inner(): + begin_openai_tool_name_mapping_scope() # must be a no-op (no force_reset) + inner_dict = _CTX.get() + assert inner_dict is outer_dict, "inner call must not replace the shared dict" + store_openai_tool_name_mapping("my_tool", "my.tool") + + ctx = contextvars.copy_context() + ctx.run(inner) + + # The outer context must see the mapping written by the inner context. + assert restore_openai_tool_name_for_user("my_tool") == "my.tool"