From 6e3fc60f0443abe056d019969a3577acee73e9c1 Mon Sep 17 00:00:00 2001
From: abhishekacharya200 <abhishekacharya200@users.noreply.github.com>
Date: Mon, 19 Oct 2026 22:46:25 +1100
Subject: [PATCH] fix(codestral): honor timeout on FIM completions

The sync, sync streaming and async streaming paths dropped the caller's
timeout and used the module level client default (6000s). Only async
non streaming passed it through.
---
 litellm/llms/codestral/completion/handler.py  |  6 ++-
 .../test_codestral_completion_handler.py      | 46 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py
index 6b172e77295..ec1d8622ad0 100644
--- a/litellm/llms/codestral/completion/handler.py
+++ b/litellm/llms/codestral/completion/handler.py
@@ -104,8 +104,9 @@ async def make_call(
     model: str,
     messages: list,
     logging_obj,
+    timeout: float | httpx.Timeout | None = None,
 ):
-    response: Final = await client.post(api_base, headers=headers, data=data, stream=True)
+    response: Final = await client.post(api_base, headers=headers, data=data, stream=True, timeout=timeout)
 
     if response.status_code != 200:
         raise TextCompletionCodestralError(status_code=response.status_code, message=response.text)
@@ -342,6 +343,7 @@ class CodestralTextCompletion:
                 headers=headers,
                 data=json.dumps({**data, "stream": True}),
                 stream=stream,
+                timeout=timeout,
             )
             _response: Final = CustomStreamWrapper(
                 iter_sse_lines(response.iter_text()),
@@ -356,6 +358,7 @@ class CodestralTextCompletion:
                 url=completion_url,
                 headers=headers,
                 data=json.dumps(data),
+                timeout=timeout,
             )
         return self.process_text_completion_response(
             model=model,
@@ -449,6 +452,7 @@ class CodestralTextCompletion:
                 model=model,
                 messages=messages,
                 logging_obj=logging_obj,
+                timeout=timeout,
             ),
             model=model,
             custom_llm_provider="text-completion-codestral",
diff --git a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py
index 9ce50d8ceca..dac670eacb0 100644
--- a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py
+++ b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py
@@ -7,7 +7,17 @@ import respx
 
 import litellm
 
-_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
+CODESTRAL_FIM_URL: Final = "https://codestral.mistral.ai/v1/fim/completions"
+_LINE_SEPARATED_TEXT: Final = "before
between
after"
+
+COMPLETION_BODY: Final = {
+    "id": "cmpl-1",
+    "object": "chat.completion",
+    "created": 1,
+    "model": "codestral-latest",
+    "choices": [{"index": 0, "message": {"role": "assistant", "content": "return 1"}, "finish_reason": "stop"}],
+    "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6},
+}
 
 
 def _fim_stream_body(text: str) -> bytes:
@@ -40,7 +50,7 @@ async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta(
     sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
-    route: Final = respx_mock.post("https://codestral.mistral.ai/v1/fim/completions").mock(
+    route: Final = respx_mock.post(CODESTRAL_FIM_URL).mock(
         return_value=httpx.Response(200, content=_fim_stream_body(_LINE_SEPARATED_TEXT))
     )
 
@@ -48,3 +58,35 @@ async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta(
 
     assert texts == [_LINE_SEPARATED_TEXT, None]
     assert json.loads(route.calls.last.request.content)["stream"] is True
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_codestral_text_completion_sends_the_callers_timeout(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool, stream: bool
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    response: Final = (
+        httpx.Response(200, content=_fim_stream_body("return 1"))
+        if stream
+        else httpx.Response(200, json=COMPLETION_BODY)
+    )
+    route: Final = respx_mock.post(CODESTRAL_FIM_URL).mock(return_value=response)
+
+    kwargs: Final = {
+        "model": "text-completion-codestral/codestral-latest",
+        "prompt": "def one():",
+        "api_key": "test-codestral-key",
+        "timeout": 7.0,
+        "stream": stream,
+    }
+    if sync_mode:
+        result = litellm.text_completion(**kwargs)
+        text = "".join(chunk.choices[0].text or "" for chunk in result) if stream else result.choices[0].text
+    else:
+        result = await litellm.atext_completion(**kwargs)
+        text = "".join([chunk.choices[0].text or "" async for chunk in result]) if stream else result.choices[0].text
+
+    assert text == "return 1"
+    assert route.calls[0].request.extensions["timeout"] == {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0}