feat(api): enhance timeout handling with swallow_exceptions option and update related logic

Author: sanjeevafk

api/routers/messages.py

@@ -647,7 +647,7 @@ async def send_message(request: Request, auth_data: dict = Depends(verify_token)
         history_messages = await _parse_snapshot_messages(snapshot_raw_messages, req.conversation_id)
         if not history_messages:
             db_start = time.perf_counter()
-            db_meta, db_messages = await with_timeout(
+            db_result = await with_timeout(
                 _load_conversation_from_db(
                     req.conversation_id,
                     user_id,
@@ -656,7 +656,12 @@ async def send_message(request: Request, auth_data: dict = Depends(verify_token)
                 timeout_seconds=CONTEXT_LOAD_TIMEOUTS["db_context"],
                 default=({}, []),
                 context_label="db_context_load",
+                swallow_exceptions=True,
             )
+            if db_result is None:
+                db_meta, db_messages = {}, []
+            else:
+                db_meta, db_messages = db_result
             db_ms = (time.perf_counter() - db_start) * 1000
             logger.info(
                 "timing_db_load",
@@ -665,12 +670,8 @@ async def send_message(request: Request, auth_data: dict = Depends(verify_token)
                 db_ms=round(db_ms, 2),
                 db_messages_count=len(db_messages),
             )
-            if not db_messages:
-                logger.warning(
-                    "db_context_timeout_fallback_to_empty",
-                    conversation_id=req.conversation_id,
-                    request_id=request_id,
-                )
+            # Note: with_timeout already logs timeout scenarios internally.
+            # Empty db_messages for new conversations is expected behavior.
             if db_meta:
                 snapshot_meta = db_meta
             if db_messages:
@@ -691,6 +692,7 @@ async def send_message(request: Request, auth_data: dict = Depends(verify_token)
             timeout_seconds=CONTEXT_LOAD_TIMEOUTS["intent_classify"],
             default=ConversationIntent(type="new_query", reason="intent_timeout_default"),
             context_label="intent_classification",
+            swallow_exceptions=True,
         )
         if intent is None:
             intent = ConversationIntent(type="new_query", reason="intent_none_default")
@@ -702,6 +704,7 @@ async def send_message(request: Request, auth_data: dict = Depends(verify_token)
         context_messages: list[ConversationMessage] = []
         context_signature = ""
         prompt_build_ms = 0.0
+        context_materialized = False
         socratic_context = build_socratic_context(history_messages)
 
         async def load_context_for_stream() -> tuple[list[ConversationMessage], str, float]:
@@ -724,6 +727,41 @@ async def load_context_for_stream() -> tuple[list[ConversationMessage], str, flo
             return loaded_messages, loaded_signature, local_prompt_build_ms
 
         context_messages_task = asyncio.create_task(load_context_for_stream())
+
+        async def ensure_context_materialized(
+            *, timeout_seconds: float, source: str
+        ) -> None:
+            nonlocal context_messages, context_signature, prompt_build_ms, context_materialized
+            if context_materialized:
+                return
+            try:
+                loaded_messages, loaded_signature, loaded_prompt_build_ms = await asyncio.wait_for(
+                    asyncio.shield(context_messages_task),
+                    timeout=timeout_seconds,
+                )
+                context_messages = loaded_messages
+                context_signature = loaded_signature
+                prompt_build_ms = loaded_prompt_build_ms
+            except (asyncio.TimeoutError, asyncio.CancelledError):
+                logger.warning(
+                    "context_load_timeout",
+                    request_id=request_id,
+                    timeout_seconds=timeout_seconds,
+                    source=source,
+                )
+                context_messages = []
+                context_signature = ""
+            except Exception as exc:
+                logger.warning(
+                    "context_load_error",
+                    request_id=request_id,
+                    source=source,
+                    error=str(exc),
+                )
+                context_messages = []
+                context_signature = ""
+            finally:
+                context_materialized = True
 
         last_three = history_messages[-3:]
         logger.info(
@@ -812,6 +850,7 @@ async def load_context_for_stream() -> tuple[list[ConversationMessage], str, flo
         system_prompt_bundle = "\n".join(
             [part for part in (system_prompt, mode_prompt, intent_prompt) if part]
         )
+        await ensure_context_materialized(timeout_seconds=1.0, source="pre_cache")
         cache_key = _message_cache_key(
             content=effective_content,
             mode=selected_mode,
@@ -994,38 +1033,9 @@ async def event_generator():
             user_sequence_id: int | None = None
             assistant_sequence_id: int | None = None
             redis_append_failed = False
-            context_load_timeout_seconds = 1.0
-            context_loaded_for_stream = False
 
             async def ensure_context_for_stream() -> None:
-                nonlocal context_messages, context_signature, prompt_build_ms, context_loaded_for_stream
-                if context_loaded_for_stream:
-                    return
-                context_loaded_for_stream = True
-                try:
-                    loaded_messages, loaded_signature, loaded_prompt_build_ms = await asyncio.wait_for(
-                        context_messages_task,
-                        timeout=context_load_timeout_seconds,
-                    )
-                    context_messages = loaded_messages
-                    context_signature = loaded_signature
-                    prompt_build_ms = loaded_prompt_build_ms
-                except (asyncio.TimeoutError, asyncio.CancelledError):
-                    logger.warning(
-                        "context_load_timeout_in_stream",
-                        request_id=request_id,
-                        timeout_seconds=context_load_timeout_seconds,
-                    )
-                    context_messages = []
-                    context_signature = ""
-                except Exception as exc:
-                    logger.warning(
-                        "context_load_error_in_stream",
-                        request_id=request_id,
-                        error=str(exc),
-                    )
-                    context_messages = []
-                    context_signature = ""
+                await ensure_context_materialized(timeout_seconds=1.0, source="stream")
 
             asyncio.create_task(
                 _capture_telemetry_async(
@@ -1712,13 +1722,7 @@ async def get_next_chunk():
                 elif status == "aborted":
                     error_type = "aborted"
                     error_message = "User aborted stream"
-                safe_user_id = safeNumber(user_id, default=None)
-                if safe_user_id is None:
-                    logger.warning(
-                        "messages_user_id_not_numeric",
-                        request_id=request_id,
-                        user_id_hash=user_id_hash,
-                    )
+                safe_user_id = user_id or None
                 payload = build_llm_request_payload(
                     request_id=request_id,
                     user_id=safe_user_id,
@@ -1755,12 +1759,13 @@ async def get_next_chunk():
             breakdown={
                 "snapshot_ms": round(snapshot_ms, 2),
                 "db_ms": round(db_ms, 2),
-                "context_build_ms": round(prompt_build_ms, 2),
             },
         )
         response_started = True
         return response
     finally:
+        if not response_started:
+            await _ingress_dedupe_clear(client_message_id)
         if not response_started and not lock_released:
             _release_conversation_lock(req.conversation_id)
             lock_released = True

api/routers/query.py

@@ -400,7 +400,7 @@ async def query_topic_stream(
         now_ts = int(time.time())
         idem_check = await check_idempotency_and_cache(
             idempotency_key=idempotency_key,
-            cache_key=cache_key,
+            cache_key=cache_key(topic, level, mode),
             now_ts=now_ts,
             idempotency_ttl=idempotency_ttl_seconds,
             idempotency_stale=idempotency_stale_seconds,

api/services/cache.py

@@ -74,6 +74,7 @@ async def _execute(self, *command: Any) -> Any:
             timeout_seconds=REDIS_REST_CALL_TIMEOUT_SECONDS,
             default=None,
             context_label=f"redis_call_{str(command[0]).lower() if command else 'unknown'}",
+            swallow_exceptions=True,
         )
         if response is None:
             raise RuntimeError("Redis REST call timed out or failed")
@@ -96,6 +97,7 @@ async def _execute_pipeline(self, commands: list[list[Any]]) -> list[Any]:
             timeout_seconds=REDIS_REST_CALL_TIMEOUT_SECONDS,
             default=None,
             context_label="redis_call_pipeline",
+            swallow_exceptions=True,
         )
         if response is None:
             raise RuntimeError("Redis REST pipeline call timed out or failed")

api/services/inference.py

@@ -1,10 +1,11 @@
 """Native multi-provider inference service."""
 
 import asyncio
+import inspect
 import json
 import re
 import time
-from typing import cast
+from typing import Any, cast
 import httpx
 import structlog
 from openai import APIConnectionError, APIStatusError, APITimeoutError
@@ -66,15 +67,31 @@
 
 
 class _SearchServiceShim:
-    async def get_search_context(self, topic: str):
-        return await _load_search_context(topic, mode=LEARNING_MODE)
+    async def get_search_context(self, topic: str, mode: str):
+        return await _load_search_context(topic, mode=mode)
 
     async def load_search_context(self, topic: str, *, mode: str):
         default_impl = getattr(self.get_search_context, "__func__", None) is _SearchServiceShim.get_search_context
         if default_impl:
             return await _load_search_context(topic, mode=mode)
+
+        get_context = cast(Any, self.get_search_context)
+        supports_mode = False
         try:
-            context = await self.get_search_context(topic)
+            signature = inspect.signature(get_context)
+            supports_mode = any(
+                param.kind is inspect.Parameter.VAR_KEYWORD or param.name == "mode"
+                for param in signature.parameters.values()
+            )
+        except (TypeError, ValueError):
+            # If signature inspection fails (e.g. dynamic callables), prefer the new contract.
+            supports_mode = True
+
+        try:
+            if supports_mode:
+                context = await get_context(topic, mode=mode)
+            else:
+                context = await get_context(topic)
         except Exception:
             return ""
         return _truncate_search_context(str(context or ""))

api/services/message_streaming.py

@@ -120,6 +120,8 @@ async def event_generator():
         telemetry_sink: dict[str, Any] = {}
         stream_failed = False
         last_progress_update = start_time
+        actual_context_messages = context_messages
+        actual_socratic_context = socratic_context
 
         async def update_progress() -> None:
             nonlocal last_progress_update
@@ -207,8 +209,6 @@ def emit(event: str, payload: dict[str, Any] | str) -> str:
                 return
 
             generation_start = time.perf_counter()
-            actual_context_messages = context_messages
-            actual_socratic_context = socratic_context
             if context_messages_task and not actual_context_messages:
                 try:
                     actual_context_messages = await asyncio.wait_for(
@@ -574,16 +574,18 @@ def emit(event: str, payload: dict[str, Any] | str) -> str:
                 sampled=True,
             )
             if assistant_message_id:
+                current_assistant_message_id = assistant_message_id
+
                 def _update_db():
                     try:
-                        ChatRepository.update_assistant_message(assistant_message_id, full_content)
+                        ChatRepository.update_assistant_message(current_assistant_message_id, full_content)
                     except Exception as exc:
                         logger.error(
                             "messages_assistant_update_failed",
                             error=str(exc),
                             request_id=request_id,
                             user_id_hash=user_id_hash,
-                            message_id=assistant_message_id,
+                            message_id=current_assistant_message_id,
                             retry=bool(req.regenerate),
                             sampled=False,
                         )

api/tests/test_inference.py

@@ -115,6 +115,25 @@ async def fake_call_model(_model, _prompt, **_kwargs):
     assert calls["count"] == 1
 
 
+@pytest.mark.asyncio
+async def test_search_service_load_search_context_forwards_mode_to_override(monkeypatch):
+    captured: dict[str, str] = {}
+
+    async def fake_search_context(_topic: str, *, mode: str):
+        captured["mode"] = mode
+        return f"context for {mode}"
+
+    monkeypatch.setattr(inference_module.search_service, "get_search_context", fake_search_context)
+
+    result = await inference_module.search_service.load_search_context(
+        "dns",
+        mode=inference_module.SOCRATIC_MODE,
+    )
+
+    assert result == "context for socratic"
+    assert captured["mode"] == inference_module.SOCRATIC_MODE
+
+
 @pytest.mark.asyncio
 async def test_learning_length_policy_default_adds_cue_when_trimmed(monkeypatch):
     async def fake_search_context(_topic: str):

api/tests/test_utils.py

@@ -1,6 +1,6 @@
 import pytest
 
-from utils import sanitize_topic, topic_cache_key
+from utils import sanitize_topic, topic_cache_key, with_timeout
 
 
 def test_sanitize_topic_valid():
@@ -25,3 +25,40 @@ def test_sanitize_topic_missing():
 def test_topic_cache_key_format():
     key = topic_cache_key("Hello World!", "eli5")
     assert key == "knowbear:hello_world:eli5"
+
+
+@pytest.mark.asyncio
+async def test_with_timeout_reraises_unexpected_exception_by_default():
+    async def boom():
+        raise ValueError("unexpected")
+
+    with pytest.raises(ValueError, match="unexpected"):
+        await with_timeout(
+            boom(),
+            timeout_seconds=0.05,
+            default="fallback",
+            context_label="unit_test_context",
+        )
+
+
+@pytest.mark.asyncio
+async def test_with_timeout_swallow_exceptions_returns_default_and_logs(caplog):
+    async def boom():
+        raise RuntimeError("network-ish failure")
+
+    with caplog.at_level("ERROR"):
+        result = await with_timeout(
+            boom(),
+            timeout_seconds=0.05,
+            default="fallback",
+            context_label="unit_test_context",
+            swallow_exceptions=True,
+        )
+
+    assert result == "fallback"
+    exception_logs = [
+        record for record in caplog.records if "timeout_wrapper_exception" in record.getMessage()
+    ]
+    assert exception_logs
+    assert any("unit_test_context" in record.getMessage() for record in exception_logs)
+    assert any(record.exc_info for record in exception_logs)

api/utils.py

@@ -170,6 +170,7 @@ async def with_timeout(
     error_on_timeout: bool = False,
     default: T | None = None,
     context_label: str = "operation",
+    swallow_exceptions: bool = False,
 ) -> T | None:
     """Run awaitable with a timeout and graceful fallback behavior."""
     try:
@@ -184,10 +185,9 @@ async def with_timeout(
         return default
     except asyncio.CancelledError:
         _logger.debug("timeout_wrapper_cancelled", extra={"context": context_label})
-        return default
-    except Exception as exc:
-        _logger.warning(
-            "timeout_wrapper_exception",
-            extra={"context": context_label, "error": str(exc)},
-        )
-        return default
+        raise
+    except Exception:
+        _logger.exception("timeout_wrapper_exception context=%s", context_label)
+        if swallow_exceptions:
+            return default
+        raise